Feature/add random dir

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.

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,8 +247,6 @@ func main() {
 	req.Header = header
 
 	w := &requester.Work{
-		Request:            req,
-		RequestBody:        bodyAll,
 		N:                  num,
 		C:                  conc,
 		QPS:                q,
@@ -235,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)

requester/requester.go

@@ -20,6 +20,7 @@ import (
 	"crypto/tls"
 	"io"
 	"io/ioutil"
+	"math/rand"
 	"net/http"
 	"net/http/httptrace"
 	"net/url"
@@ -48,14 +49,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 +146,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 +230,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 +253,25 @@ 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)
+	}
+}
+
+// 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 {
@@ -275,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
 }

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,
 	}
@@ -132,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")
+	}
+}