From 12b2d74b9fe3e432ff1b193b5c22ac0833fb8a33 Mon Sep 17 00:00:00 2001
From: Mark Adams <mark@markadams.me>
Date: Wed, 31 May 2017 12:33:00 -0500
Subject: [PATCH] Add -disableRedirects flag to disable following of redirects
 (#41)

---
 hey.go                 | 3 +++
 requester/requester.go | 9 +++++++++
 2 files changed, 12 insertions(+)

diff --git a/hey.go b/hey.go
index 7001d3c1..71b59c7b 100644
--- a/hey.go
+++ b/hey.go
@@ -69,6 +69,7 @@ var (
 
 	disableCompression = flag.Bool("disable-compression", false, "")
 	disableKeepAlives  = flag.Bool("disable-keepalive", false, "")
+	disableRedirects   = flag.Bool("disable-redirects", false, "")
 	proxyAddr          = flag.String("x", "", "")
 
 	enableTrace = flag.Bool("more", false, "")
@@ -102,6 +103,7 @@ Options:
   -disable-compression  Disable compression.
   -disable-keepalive    Disable keep-alive, prevents re-use of TCP
                         connections between different HTTP requests.
+  -disable-redirects    Disable following of HTTP redirects
   -cpus                 Number of used cpu cores.
                         (default for current machine is %d cores)
   -more                 Provides information on DNS lookup, dialup, request and
@@ -214,6 +216,7 @@ func main() {
 		Timeout:            *t,
 		DisableCompression: *disableCompression,
 		DisableKeepAlives:  *disableKeepAlives,
+		DisableRedirects:   *disableRedirects,
 		H2:                 *h2,
 		ProxyAddr:          proxyURL,
 		Output:             *output,
diff --git a/requester/requester.go b/requester/requester.go
index 6f05cdd0..ec6eb5d4 100644
--- a/requester/requester.go
+++ b/requester/requester.go
@@ -76,6 +76,9 @@ type Work struct {
 	// DisableKeepAlives is an option to prevents re-use of TCP connections between different HTTP requests
 	DisableKeepAlives bool
 
+	// DisableRedirects is an option to prevent the following of HTTP redirects
+	DisableRedirects bool
+
 	// Output represents the output type. If "csv" is provided, the
 	// output will be dumped as a csv stream.
 	Output string
@@ -232,7 +235,13 @@ func (b *Work) runWorker(n int) {
 	} else {
 		tr.TLSNextProto = make(map[string]func(string, *tls.Conn) http.RoundTripper)
 	}
+
 	client := &http.Client{Transport: tr, Timeout: time.Duration(b.Timeout) * time.Second}
+	if b.DisableRedirects {
+		client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
+			return http.ErrUseLastResponse
+		}
+	}
 	for i := 0; i < n; i++ {
 		if b.QPS > 0 {
 			<-throttle