Add possibility to set an User Agent per request from file

README.md

@@ -54,6 +54,10 @@ Options:
   -disable-redirects    Disable following of HTTP redirects
   -cpus                 Number of used cpu cores.
                         (default for current machine is 8 cores)
+
+  -user-agent-feed      Set a User Agent for each request from file.
+						Expected on User Agent per line.
+						For example, /home/user/ua.txt or ./ua.txt.
 ```
 
 Previously known as [github.com/rakyll/boom](https://github.com/rakyll/boom).

hey.go

@@ -18,6 +18,7 @@ package main
 import (
 	"flag"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"math"
 	"net/http"
@@ -63,6 +64,7 @@ var (
 	disableKeepAlives  = flag.Bool("disable-keepalive", false, "")
 	disableRedirects   = flag.Bool("disable-redirects", false, "")
 	proxyAddr          = flag.String("x", "", "")
+	userAgentFeedFile  = flag.String("user-agent-feed", "", "")
 )
 
 var usage = `Usage: hey [options...] <url>
@@ -99,6 +101,10 @@ Options:
   -disable-redirects    Disable following of HTTP redirects
   -cpus                 Number of used cpu cores.
                         (default for current machine is %d cores)
+
+  -user-agent-feed      Set a User Agent for each request from file.
+                        Expected on User Agent per line.
+                        For example, /home/user/ua.txt or ./ua.txt.
 `
 
 func main() {
@@ -212,6 +218,14 @@ func main() {
 	header.Set("User-Agent", ua)
 	req.Header = header
 
+	var userAgentFeed io.ReadSeeker
+	if *userAgentFeedFile != "" {
+		userAgentFeed, err = os.Open(*userAgentFeedFile)
+		if err != nil {
+			usageAndExit(err.Error())
+		}
+	}
+
 	w := &requester.Work{
 		Request:            req,
 		RequestBody:        bodyAll,
@@ -225,6 +239,7 @@ func main() {
 		H2:                 *h2,
 		ProxyAddr:          proxyURL,
 		Output:             *output,
+		UserAgentFeed:      userAgentFeed,
 	}
 	w.Init()
 

requester/requester.go

@@ -95,6 +95,11 @@ type Work struct {
 
 	report *report
 	mutex  sync.Mutex
+
+	// UserAgentFeed is where the UserAgentFeeder will read to set a user agent
+	// for each request. Optional.
+	UserAgentFeed   io.ReadSeeker
+	userAgentFeeder *userAgentFeeder
 }
 
 func (b *Work) writer() io.Writer {
@@ -109,6 +114,9 @@ func (b *Work) Init() {
 	b.initOnce.Do(func() {
 		b.results = make(chan *result, min(b.C*1000, maxResult))
 		b.stopCh = make(chan struct{}, b.C)
+		if b.UserAgentFeed != nil {
+			b.userAgentFeeder = newUserAgentFeeder(b.UserAgentFeed)
+		}
 	})
 }
 
@@ -148,6 +156,12 @@ func (b *Work) makeRequest(c *http.Client) {
 	var dnsStart, connStart, resStart, reqStart, delayStart time.Duration
 	var dnsDuration, connDuration, resDuration, reqDuration, delayDuration time.Duration
 	req := cloneRequest(b.Request, b.RequestBody)
+
+	// Set the request user agent from the feeder, if any
+	if b.userAgentFeeder != nil {
+		b.userAgentFeeder.Feed(req)
+	}
+
 	trace := &httptrace.ClientTrace{
 		DNSStart: func(info httptrace.DNSStartInfo) {
 			dnsStart = now()
@@ -178,6 +192,7 @@ func (b *Work) makeRequest(c *http.Client) {
 		},
 	}
 	req = req.WithContext(httptrace.WithClientTrace(req.Context(), trace))
+
 	resp, err := c.Do(req)
 	if err == nil {
 		size = resp.ContentLength

requester/requester_test.go

@@ -21,6 +21,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -159,3 +160,50 @@ func TestRaceConditionDNSLookup(t *testing.T) {
 		t.Errorf("Expected to send 5000 requests, found %v", count)
 	}
 }
+
+func TestUserAgentFeeder(t *testing.T) {
+	count := 0
+	userAgentCount := make(map[string]int)
+	mutex := &sync.Mutex{}
+	handler := func(w http.ResponseWriter, r *http.Request) {
+		mutex.Lock()
+		defer mutex.Unlock()
+		count++
+		userAgentCount[r.UserAgent()]++
+	}
+	server := httptest.NewServer(http.HandlerFunc(handler))
+	defer server.Close()
+
+	// userAgentFeed contains also an empty feed that should fallback to
+	// userAgentDefault
+	userAgentFeed := []string{"hey/1.0.0", "", "hey/1.1.0"}
+	userAgentDefault := "hey/0.0.1"
+
+	uaf := strings.NewReader(strings.Join(userAgentFeed, "\n"))
+
+	req, _ := http.NewRequest("GET", server.URL, nil)
+	req.Header.Set("User-Agent", userAgentDefault)
+	w := &Work{
+		Request:       req,
+		N:             6,
+		C:             2,
+		UserAgentFeed: uaf,
+	}
+	w.Run()
+	if count != 6 {
+		t.Errorf("Expected to send 6 requests, found %v", count)
+	}
+
+	for _, ua := range userAgentFeed {
+		if ua == "" {
+			ua = userAgentDefault
+		}
+		v, ok := userAgentCount[ua]
+		if !ok {
+			t.Errorf("Expected to send requests with user agent %s but no one was sent", ua)
+		}
+		if v != 2 {
+			t.Errorf("Expected to send 2 requests with user agent %s, found %v", ua, v)
+		}
+	}
+}

requester/ua_feeder.go

@@ -0,0 +1,65 @@
+// Copyright 2019 ScientiaMobile Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// userAgentFeeder allow to set a User Agent per request reading from a io.ReadSeeker feed.
+
+package requester
+
+import (
+	"bufio"
+	"io"
+	"net/http"
+	"sync"
+)
+
+func newUserAgentFeeder(reader io.ReadSeeker) *userAgentFeeder {
+	return &userAgentFeeder{
+		reader:  reader,
+		scanner: bufio.NewScanner(reader),
+	}
+}
+
+type userAgentFeeder struct {
+	reader  io.ReadSeeker
+	scanner *bufio.Scanner
+	mutex   sync.Mutex
+}
+
+// userAgent returns a User Agent from the feed. When all the feed is read
+// it will restart to read from the beginning
+func (uaf *userAgentFeeder) userAgent() (string, error) {
+	uaf.mutex.Lock()
+	if uaf.scanner.Scan() == true {
+		ua := uaf.scanner.Text()
+		uaf.mutex.Unlock()
+		return ua, nil
+	}
+	if err := uaf.scanner.Err(); err != nil {
+		uaf.mutex.Unlock()
+		return "", err
+	}
+	uaf.mutex.Unlock()
+	uaf.reader.Seek(0, io.SeekStart)
+	uaf.scanner = bufio.NewScanner(uaf.reader)
+	return uaf.userAgent()
+}
+
+// Feed sets the request User Agent header with the feed entry. If feed entry is
+// an empty string the UserAgent won't be updated.
+func (uaf *userAgentFeeder) Feed(req *http.Request) {
+	userAgent, _ := uaf.userAgent()
+	if userAgent != "" {
+		req.Header.Set("User-Agent", userAgent)
+	}
+}

requester/ua_feeder_test.go

@@ -0,0 +1,82 @@
+// Copyright 2019 ScientiaMobile Inc. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package requester
+
+import (
+	"fmt"
+	"net/http"
+	"strings"
+	"testing"
+)
+
+func Test_userAgentFeeder_userAgent(t *testing.T) {
+	userAgent1 := "hey/1.0"
+	userAgent2 := "hey/1.1"
+
+	reader := strings.NewReader(fmt.Sprintf("%s\n%s\n", userAgent1, userAgent2))
+	uaf := newUserAgentFeeder(reader)
+
+	got, err := uaf.userAgent()
+	if err != nil {
+		t.Errorf("userAgentFeeder.userAgent() error = %v, want nil", err)
+	}
+	expected := userAgent1
+	if got != expected {
+		t.Errorf("userAgentFeeder.userAgent() = %v, want %v", got, expected)
+	}
+
+	got, err = uaf.userAgent()
+	if err != nil {
+		t.Errorf("userAgentFeeder.userAgent() error = %v, want nil", err)
+	}
+	expected = userAgent2
+	if got != expected {
+		t.Errorf("userAgentFeeder.userAgent() = %v, want %v", got, expected)
+	}
+
+	// Lines ended should start from beginning
+	got, err = uaf.userAgent()
+	if err != nil {
+		t.Errorf("userAgentFeeder.userAgent() error = %v, want nil", err)
+	}
+	expected = userAgent1
+	if got != expected {
+		t.Errorf("userAgentFeeder.userAgent() = %v, want %v", got, expected)
+	}
+}
+
+func Test_userAgentFeeder_Feed(t *testing.T) {
+	// userAgentFeed contains also an empty feed that should fallback to
+	// userAgentDefault
+	userAgentFeed := []string{"hey/1.0.0", "", "hey/1.1.0"}
+	userAgentDefault := "hey/0.0.1"
+
+	reader := strings.NewReader(strings.Join(userAgentFeed, "\n"))
+	uaf := newUserAgentFeeder(reader)
+
+	req, _ := http.NewRequest("GET", "", nil)
+	req.Header.Set("User-Agent", userAgentDefault)
+
+	for _, ua := range userAgentFeed {
+		r := cloneRequest(req, nil)
+		uaf.Feed(r)
+		if ua == "" {
+			ua = userAgentDefault
+		}
+		if ua != r.UserAgent() {
+			t.Errorf("userAgentFeeder.Feed(req) = %v, want %v", r.UserAgent(), ua)
+		}
+	}
+}