diff --git a/README.md b/README.md index 917c903e..227cf083 100644 --- a/README.md +++ b/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). diff --git a/hey.go b/hey.go index d96455b9..7a1fcb9f 100644 --- a/hey.go +++ b/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...] @@ -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() diff --git a/requester/requester.go b/requester/requester.go index 35c2a43b..010369d3 100644 --- a/requester/requester.go +++ b/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 diff --git a/requester/requester_test.go b/requester/requester_test.go index 917df235..aa9d569a 100644 --- a/requester/requester_test.go +++ b/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) + } + } +} diff --git a/requester/ua_feeder.go b/requester/ua_feeder.go new file mode 100644 index 00000000..86a7de5f --- /dev/null +++ b/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) + } +} diff --git a/requester/ua_feeder_test.go b/requester/ua_feeder_test.go new file mode 100644 index 00000000..e3af89fb --- /dev/null +++ b/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) + } + } +}