From 97ef9d022d3155bd6c70786a4132b63ae84d1354 Mon Sep 17 00:00:00 2001 From: Martin Magakian Date: Mon, 27 Sep 2021 13:17:39 +0200 Subject: [PATCH 1/2] Rework request creation to send in order to allow different request strategies --- hey.go | 3 +-- requester/requester.go | 27 ++++++++++++--------------- requester/requester_test.go | 23 +++++++++++------------ 3 files changed, 24 insertions(+), 29 deletions(-) diff --git a/hey.go b/hey.go index f727e26b..4b07ef3f 100644 --- a/hey.go +++ b/hey.go @@ -222,8 +222,7 @@ func main() { req.Header = header w := &requester.Work{ - Request: req, - RequestBody: bodyAll, + NextRequest: requester.DuplicateNextRequest(req, bodyAll), N: num, C: conc, QPS: q, diff --git a/requester/requester.go b/requester/requester.go index fd7277e7..93dd362e 100644 --- a/requester/requester.go +++ b/requester/requester.go @@ -48,14 +48,9 @@ type result struct { } type Work struct { - // Request is the request to be made. - Request *http.Request - - RequestBody []byte - - // RequestFunc is a function to generate requests. If it is nil, then - // Request and RequestData are cloned for each request. - RequestFunc func() *http.Request + // NextRequest is a function to generate requests. + // It allow setting different request strategies for different use case. + NextRequest func() *http.Request // N is the total number of requests to make. N int @@ -150,12 +145,8 @@ func (b *Work) makeRequest(c *http.Client) { var code int var dnsStart, connStart, resStart, reqStart, delayStart time.Duration var dnsDuration, connDuration, resDuration, reqDuration, delayDuration time.Duration - var req *http.Request - if b.RequestFunc != nil { - req = b.RequestFunc() - } else { - req = cloneRequest(b.Request, b.RequestBody) - } + + req := b.NextRequest() trace := &httptrace.ClientTrace{ DNSStart: func(info httptrace.DNSStartInfo) { dnsStart = now() @@ -238,7 +229,6 @@ func (b *Work) runWorkers() { tr := &http.Transport{ TLSClientConfig: &tls.Config{ InsecureSkipVerify: true, - ServerName: b.Request.Host, }, MaxIdleConnsPerHost: min(b.C, maxIdleConn), DisableCompression: b.DisableCompression, @@ -262,6 +252,13 @@ func (b *Work) runWorkers() { wg.Wait() } +// DuplicateNextRequest returns a func that duplicate the request and body everytime it's called. +func DuplicateNextRequest(request *http.Request, body []byte) func() *http.Request { + return func() *http.Request { + return cloneRequest(request, body) + } +} + // cloneRequest returns a clone of the provided *http.Request. // The clone is a shallow copy of the struct and its Header map. func cloneRequest(r *http.Request, body []byte) *http.Request { diff --git a/requester/requester_test.go b/requester/requester_test.go index 9188f77f..311757ed 100644 --- a/requester/requester_test.go +++ b/requester/requester_test.go @@ -35,9 +35,9 @@ func TestN(t *testing.T) { req, _ := http.NewRequest("GET", server.URL, nil) w := &Work{ - Request: req, - N: 20, - C: 2, + NextRequest: DuplicateNextRequest(req, []byte{}), + N: 20, + C: 2, } w.Run() if count != 20 { @@ -56,10 +56,10 @@ func TestQps(t *testing.T) { req, _ := http.NewRequest("GET", server.URL, nil) w := &Work{ - Request: req, - N: 20, - C: 2, - QPS: 1, + NextRequest: DuplicateNextRequest(req, []byte{}), + N: 20, + C: 2, + QPS: 1, } wg.Add(1) time.AfterFunc(time.Second, func() { @@ -90,9 +90,9 @@ func TestRequest(t *testing.T) { req.Header = header req.SetBasicAuth("username", "password") w := &Work{ - Request: req, - N: 1, - C: 1, + NextRequest: DuplicateNextRequest(req, []byte{}), + N: 1, + C: 1, } w.Run() if uri != "/" { @@ -122,8 +122,7 @@ func TestBody(t *testing.T) { req, _ := http.NewRequest("POST", server.URL, bytes.NewBuffer([]byte("Body"))) w := &Work{ - Request: req, - RequestBody: []byte("Body"), + NextRequest: DuplicateNextRequest(req, []byte("Body")), N: 10, C: 1, } From c9376e8b774edb205b865cde88b57e2cdb9880eb Mon Sep 17 00:00:00 2001 From: Martin Magakian Date: Mon, 27 Sep 2021 14:38:45 +0200 Subject: [PATCH 2/2] Enable passing a directory as data sources --- README.md | 1 + hey.go | 31 ++++++++++++++++++++++++++++++- requester/requester.go | 14 ++++++++++++++ requester/requester_test.go | 28 ++++++++++++++++++++++++++++ 4 files changed, 73 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index d7c7c369..df8d34c6 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,7 @@ Options: -A HTTP Accept header. -d HTTP request body. -D HTTP request body from file. For example, /home/user/file.txt or ./file.txt. + -R HTTP request body from files in directory. For example, ./dir. -T Content-type, defaults to "text/html". -a Basic authentication, username:password. -x HTTP Proxy address as host:port. diff --git a/hey.go b/hey.go index 4b07ef3f..526b507f 100644 --- a/hey.go +++ b/hey.go @@ -24,6 +24,7 @@ import ( gourl "net/url" "os" "os/signal" + "path" "regexp" "runtime" "strings" @@ -43,6 +44,7 @@ var ( headers = flag.String("h", "", "") body = flag.String("d", "", "") bodyFile = flag.String("D", "", "") + bodyDir = flag.String("R", "", "") accept = flag.String("A", "", "") contentType = flag.String("T", "text/html", "") authHeader = flag.String("a", "", "") @@ -87,6 +89,7 @@ Options: -A HTTP Accept header. -d HTTP request body. -D HTTP request body from file. For example, /home/user/file.txt or ./file.txt. + -R HTTP request body from files in directory. For example, ./dir. -T Content-type, defaults to "text/html". -U User-Agent, defaults to version "hey/0.0.1". -a Basic authentication, username:password. @@ -181,6 +184,28 @@ func main() { } bodyAll = slurp } + var bodyAllDir [][]byte + if *bodyDir != "" { + if bodyAll != nil { + errAndExit("Cannot use both body and directory.") + } + + files, err := ioutil.ReadDir(*bodyDir) + if err != nil { + errAndExit(err.Error()) + } + + bodyAllDir = make([][]byte, len(files)) + for i, f := range files { + if !f.IsDir() { + slurp, err := ioutil.ReadFile(path.Join(*bodyDir, f.Name())) + if err != nil { + errAndExit(err.Error()) + } + bodyAllDir[i] = slurp + } + } + } var proxyURL *gourl.URL if *proxyAddr != "" { @@ -222,7 +247,6 @@ func main() { req.Header = header w := &requester.Work{ - NextRequest: requester.DuplicateNextRequest(req, bodyAll), N: num, C: conc, QPS: q, @@ -234,6 +258,11 @@ func main() { ProxyAddr: proxyURL, Output: *output, } + if *bodyDir != "" { + w.NextRequest = requester.DuplicateNextRequest(req, bodyAll) + } else { + w.NextRequest = requester.DuplicateNextRequestWithRandomBody(req, bodyAllDir) + } w.Init() c := make(chan os.Signal, 1) diff --git a/requester/requester.go b/requester/requester.go index 93dd362e..7a3f0d6a 100644 --- a/requester/requester.go +++ b/requester/requester.go @@ -20,6 +20,7 @@ import ( "crypto/tls" "io" "io/ioutil" + "math/rand" "net/http" "net/http/httptrace" "net/url" @@ -259,6 +260,18 @@ func DuplicateNextRequest(request *http.Request, body []byte) func() *http.Reque } } +// DuplicateNextRequestWithRandomBody returns a func that duplicate the request and +// pick a random body everytime it's called. +func DuplicateNextRequestWithRandomBody(request *http.Request, bodies [][]byte) func() *http.Request { + return func() *http.Request { + if len(bodies) != 0 { + randomBody := bodies[rand.Intn(len(bodies))] + return cloneRequest(request, randomBody) + } + return cloneRequest(request, []byte{}) + } +} + // cloneRequest returns a clone of the provided *http.Request. // The clone is a shallow copy of the struct and its Header map. func cloneRequest(r *http.Request, body []byte) *http.Request { @@ -272,6 +285,7 @@ func cloneRequest(r *http.Request, body []byte) *http.Request { } if len(body) > 0 { r2.Body = ioutil.NopCloser(bytes.NewReader(body)) + r2.ContentLength = int64(len(body)) } return r2 } diff --git a/requester/requester_test.go b/requester/requester_test.go index 311757ed..59a0fd92 100644 --- a/requester/requester_test.go +++ b/requester/requester_test.go @@ -131,3 +131,31 @@ func TestBody(t *testing.T) { t.Errorf("Expected to work 10 times, found %v", count) } } + +func TestRandomBody(t *testing.T) { + var countBody1 int64 + var countBody2 int64 + handler := func(w http.ResponseWriter, r *http.Request) { + body, _ := ioutil.ReadAll(r.Body) + if string(body) == "Body1" { + atomic.AddInt64(&countBody1, 1) + } + if string(body) == "Body2" { + atomic.AddInt64(&countBody2, 1) + } + } + server := httptest.NewServer(http.HandlerFunc(handler)) + defer server.Close() + + req, _ := http.NewRequest("POST", server.URL, bytes.NewBuffer([]byte("Body"))) + w := &Work{ + NextRequest: DuplicateNextRequestWithRandomBody(req, [][]byte{[]byte("Body1"), []byte("Body2")}), + N: 100, + C: 1, + } + w.Run() + if (10 >= countBody1 && countBody1 <= 50) || + (10 >= countBody2 && countBody2 <= 50) { + t.Errorf("Unexpected random statistic found") + } +}