allow to configure host header

Author: sschepens

boomer/boomer.go

@@ -18,10 +18,10 @@ package boomer
 import (
 	"crypto/tls"
 	"math"
+	"net"
 	"runtime"
 	"sync"
 	"time"
-	"net"
 
 	"github.com/Clever/leakybucket"
 	"github.com/Clever/leakybucket/memory"
@@ -40,12 +40,13 @@ type Result struct {
 type Boomer struct {
 	// Request is the request to be made.
 	Request *fasthttp.Request
+	Addr    string
 
 	// Timeout in seconds.
-	Timeout time.Duration
+	Timeout        time.Duration
 	ConnectTimeout time.Duration
-	ReadTimeout time.Duration
-	WriteTimeout time.Duration
+	ReadTimeout    time.Duration
+	WriteTimeout   time.Duration
 
 	// C is the concurrency level, the number of concurrent workers to run.
 	C uint
@@ -66,13 +67,14 @@ type Boomer struct {
 	jobs     chan *fasthttp.Request
 	running  bool
 	wg       *sync.WaitGroup
-	client   *fasthttp.Client
+	client   *fasthttp.HostClient
 }
 
 // NewBoomer returns a new instance of Boomer for the specified request.
-func NewBoomer(req *fasthttp.Request) *Boomer {
+func NewBoomer(addr string, req *fasthttp.Request) *Boomer {
 	return &Boomer{
 		C:       uint(runtime.NumCPU()),
+		Addr:    addr,
 		Request: req,
 		results: make(chan Result),
 		stop:    make(chan struct{}),
@@ -169,15 +171,16 @@ func (b *Boomer) Run() {
 	if b.running {
 		return
 	}
-	b.client = &fasthttp.Client{
+	b.client = &fasthttp.HostClient{
+		Addr: b.Addr,
 		Dial: func(addr string) (net.Conn, error) {
 			return fasthttp.DialTimeout(addr, b.ConnectTimeout)
 		},
 		TLSConfig: &tls.Config{
 			InsecureSkipVerify: true,
 		},
-		MaxConnsPerHost: math.MaxInt32,
-		ReadTimeout: b.ReadTimeout,
+		MaxConns:     math.MaxInt32,
+		ReadTimeout:  b.ReadTimeout,
 		WriteTimeout: b.WriteTimeout,
 	}
 	b.running = true

boomer/boomer_test.go

@@ -38,7 +38,7 @@ func TestN(t *testing.T) {
 	req := fasthttp.AcquireRequest()
 	req.SetRequestURI(server.URL)
 	req.Header.SetMethod("GET")
-	boomer := NewBoomer(req).
+	boomer := NewBoomer(string(req.Host()), req).
 		WithAmount(20).
 		WithConcurrency(2)
 	go func() {
@@ -64,7 +64,7 @@ func TestQPS(t *testing.T) {
 	req := fasthttp.AcquireRequest()
 	req.SetRequestURI(server.URL)
 	req.Header.SetMethod("GET")
-	boomer := NewBoomer(req).
+	boomer := NewBoomer(string(req.Host()), req).
 		WithAmount(20).
 		WithConcurrency(2).
 		WithRateLimit(1, time.Second)
@@ -103,7 +103,7 @@ func TestRequest(t *testing.T) {
 	req.Header.Set("Content-Type", "text/html")
 	req.Header.Set("X-some", "value")
 	req.Header.Set("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte("username:password")))
-	boomer := NewBoomer(req).
+	boomer := NewBoomer(string(req.Host()), req).
 		WithAmount(1).
 		WithConcurrency(1)
 	go func() {
@@ -141,7 +141,7 @@ func TestBody(t *testing.T) {
 	req.SetRequestURI(server.URL)
 	req.Header.SetMethod("POST")
 	req.SetBody([]byte("Body"))
-	boomer := NewBoomer(req).
+	boomer := NewBoomer(string(req.Host()), req).
 		WithAmount(10).
 		WithConcurrency(1)
 	go func() {

pla.go

@@ -107,6 +107,10 @@ func main() {
 			usageAndExit("invalid url ''" + req.URI().String() + "'', unable to detect host")
 		}
 	}
+	addr := string(req.URI().Host())
+	if !strings.Contains(addr, ":") {
+		addr = addr + ":80"
+	}
 	req.Header.SetMethod(method)
 	req.SetBodyString(*body)
 	req.Header.SetContentLength(len(req.Body()))
@@ -120,7 +124,11 @@ func main() {
 		if err != nil {
 			usageAndExit(err.Error())
 		}
-		req.Header.Set(match[1], match[2])
+		if match[1] == "Host" || match[1] == "host" {
+			req.SetHost(match[2])
+		} else {
+			req.Header.Set(match[1], match[2])
+		}
 	}
 
 	if !*disableCompression {
@@ -132,7 +140,7 @@ func main() {
 	}
 
 	ui = interfaces.NewBasicInterface()
-	boomerInstance = boomer.NewBoomer(req).
+	boomerInstance = boomer.NewBoomer(addr, req).
 		WithAmount(*n).
 		WithConcurrency(*c).
 		WithDuration(*duration).