feat: add client rate limits (#281)

* feat: add client rate limits

Signed-off-by: Felipe Zipitria <felipe.zipitria@owasp.org>

* Apply suggestions from code review

Co-authored-by: Max Leske <maxleske@gmail.com>

* fix: update naming after code review merge

---------

Signed-off-by: Felipe Zipitria <felipe.zipitria@owasp.org>
Co-authored-by: Max Leske <maxleske@gmail.com>

cmd/run.go

@@ -54,6 +54,7 @@ func NewRunCommand() *cobra.Command {
 	runCmd.Flags().Duration("wait-for-connection-timeout", http.DefaultConnectionTimeout, "Http connection timeout, The timeout includes connection time, any redirects, and reading the response body.")
 	runCmd.Flags().Bool("wait-for-insecure-skip-tls-verify", http.DefaultInsecureSkipTLSVerify, "Skips tls certificate checks for the HTTPS request.")
 	runCmd.Flags().Bool("wait-for-no-redirect", http.DefaultNoRedirect, "Do not follow HTTP 3xx redirects.")
+	runCmd.Flags().DurationP("rate-limit", "r", 0, "Limit the request rate to the server to 1 request per specified duration. 0 is the default, and disables rate limiting.")
 
 	return runCmd
 }
@@ -83,6 +84,7 @@ func runE(cmd *cobra.Command, args []string) error {
 	connectionTimeout, _ := cmd.Flags().GetDuration("wait-for-connection-timeout")
 	insecureSkipTLSVerify, _ := cmd.Flags().GetBool("wait-for-insecure-skip-tls-verify")
 	noRedirect, _ := cmd.Flags().GetBool("wait-for-no-redirect")
+	rateLimit, _ := cmd.Flags().GetDuration("rate-limit")
 
 	if exclude != "" && include != "" {
 		cmd.SilenceUsage = false
@@ -161,6 +163,7 @@ func runE(cmd *cobra.Command, args []string) error {
 		ShowOnlyFailed: showOnlyFailed,
 		ConnectTimeout: connectTimeout,
 		ReadTimeout:    readTimeout,
+		RateLimit:      rateLimit,
 	}, out)
 
 	if err != nil {

ftwhttp/client.go

@@ -4,9 +4,11 @@
 package ftwhttp
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
+	"golang.org/x/time/rate"
 	"net"
 	"net/http/cookiejar"
 	"strings"
@@ -21,6 +23,7 @@ func NewClientConfig() ClientConfig {
 	return ClientConfig{
 		ConnectTimeout: 3 * time.Second,
 		ReadTimeout:    1 * time.Second,
+		RateLimiter:    rate.NewLimiter(rate.Inf, 1),
 	}
 }
 
@@ -44,6 +47,11 @@ func (c *Client) SetRootCAs(cas *x509.CertPool) {
 	c.config.RootCAs = cas
 }
 
+// SetRateLimiter sets the rate limiter for the client.
+func (c *Client) SetRateLimiter(limiter *rate.Limiter) {
+	c.config.RateLimiter = limiter
+}
+
 // NewConnection creates a new Connection based on a Destination
 func (c *Client) NewConnection(d Destination) error {
 	if c.Transport != nil && c.Transport.connection != nil {
@@ -110,7 +118,12 @@ func (c *Client) dial(d Destination) (net.Conn, error) {
 func (c *Client) Do(req Request) (*Response, error) {
 	var response *Response
 
-	err := c.Transport.Request(&req)
+	err := c.config.RateLimiter.Wait(context.Background()) // This is a blocking call. Honors the rate limit
+	if err != nil {
+		log.Error().Msgf("http/client: error waiting on rate limiter: %s\n", err.Error())
+		return response, err
+	}
+	err = c.Transport.Request(&req)
 
 	if err != nil {
 		log.Error().Msgf("http/client: error sending request: %s\n", err.Error())

ftwhttp/client_test.go

@@ -6,9 +6,11 @@ package ftwhttp
 import (
 	"bytes"
 	"fmt"
+	"golang.org/x/time/rate"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/suite"
 )
@@ -32,6 +34,7 @@ func (s *clientTestSuite) SetupTest() {
 	var err error
 	s.client, err = NewClient(NewClientConfig())
 	s.Require().NoError(err)
+	s.Require().Equal(s.client.config.RateLimiter, rate.NewLimiter(rate.Inf, 1))
 	s.Nil(s.client.Transport, "Transport not expected to be initialized yet")
 }
 
@@ -74,6 +77,18 @@ func (s *clientTestSuite) TestNewClient() {
 	s.NotNil(s.client.Jar, "Error creating Client")
 }
 
+func (s *clientTestSuite) TestSetRootCAs() {
+	s.client.SetRootCAs(nil)
+	s.Nil(s.client.config.RootCAs, "Error setting RootCAs")
+}
+
+func (s *clientTestSuite) TestSetRateLimiter() {
+	newRateLimiter := rate.NewLimiter(rate.Every(10*time.Second), 100)
+	s.client.SetRateLimiter(newRateLimiter)
+	rl := s.client.config.RateLimiter
+	s.Require().Equal(newRateLimiter, rl, "Error setting RateLimiter")
+}
+
 func (s *clientTestSuite) TestConnectDestinationHTTPS() {
 	s.httpTestServer(secureServer)
 	d, err := DestinationFromString(s.ts.URL)
@@ -211,3 +226,33 @@ func (s *clientTestSuite) TestNewOrReusedConnectionReusesTransport() {
 
 	s.Equal(begin, s.client.Transport.duration.begin, "Transport must not be reinitialized when reusing connection")
 }
+
+// TestClientRateLimits tests the rate limiter functionality of the client. Test should take at least 3 seconds to run.
+func (s *clientTestSuite) TestClientRateLimits() {
+	waitTime := 3 * time.Second
+	s.httpTestServer(insecureServer)
+	d, err := DestinationFromString(s.ts.URL)
+	s.Require().NoError(err, "Failed to construct destination from test server")
+
+	newRateLimiter := rate.NewLimiter(rate.Every(waitTime), 1)
+	s.client.SetRateLimiter(newRateLimiter)
+	err = s.client.NewOrReusedConnection(*d)
+	s.Require().NoError(err, "Failed to create new or to reuse connection")
+
+	rl := &RequestLine{
+		Method:  "GET",
+		URI:     "/get",
+		Version: "HTTP/1.1",
+	}
+
+	h := Header{"Accept": "*/*", "User-Agent": "go-ftw test agent", "Host": "localhost"}
+	req := NewRequest(rl, h, nil, true)
+
+	// We need to do at least 2 calls so there is a wait between both.
+	before := time.Now()
+	_, err = s.client.Do(*req)
+	_, err = s.client.Do(*req)
+	after := time.Now()
+
+	s.GreaterOrEqual(after.Sub(before), waitTime, "Rate limiter did not work as expected")
+}

ftwhttp/types.go

@@ -8,6 +8,8 @@ import (
 	"net"
 	"net/http"
 	"time"
+
+	"golang.org/x/time/rate"
 )
 
 // ClientConfig provides configuration options for the HTTP client.
@@ -18,6 +20,8 @@ type ClientConfig struct {
 	ReadTimeout time.Duration
 	// RootCAs is the set of root CA certificates that is used to verify server
 	RootCAs *x509.CertPool
+	// RateLimiter is the rate limiter to use for requests.
+	RateLimiter *rate.Limiter
 }
 
 // Client is the top level abstraction in http

go.mod

@@ -21,6 +21,7 @@ require (
 	github.com/stretchr/testify v1.9.0
 	github.com/yargevad/filepathx v1.0.0
 	golang.org/x/net v0.22.0
+	golang.org/x/time v0.5.0
 	wait4x.dev/v2 v2.14.0
 )
 

go.sum

@@ -42,10 +42,6 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/huandu/xstrings v1.3.3 h1:/Gcsuc1x8JVbJ9/rlye4xZnVAbEkGauT8lbebqcQws4=
 github.com/huandu/xstrings v1.3.3/go.mod h1:y5/lhBue+AyNmUVz9RLU9xbLR0o4KIIExikq4ovT0aE=
-github.com/icza/backscanner v0.0.0-20230330133933-bf6beb754c70 h1:xrd41BUTgqxyYFfFwGdt/bnwS8KNYzPraj8WgvJ5NWk=
-github.com/icza/backscanner v0.0.0-20230330133933-bf6beb754c70/go.mod h1:GYeBD1CF7AqnKZK+UCytLcY3G+UKo0ByXX/3xfdNyqQ=
-github.com/icza/backscanner v0.0.0-20240221180550-0d5d27a5d977 h1:XCXNe6YufH4flQt4eM+D771E+Rc3w+wcE34YMS/DpzI=
-github.com/icza/backscanner v0.0.0-20240221180550-0d5d27a5d977/go.mod h1:GYeBD1CF7AqnKZK+UCytLcY3G+UKo0ByXX/3xfdNyqQ=
 github.com/icza/backscanner v0.0.0-20240221180818-f23e3ba0e79f h1:EKPpaKkARuHjoV/ZKzk3vqbSJXULRSivDCQhL+tF77Y=
 github.com/icza/backscanner v0.0.0-20240221180818-f23e3ba0e79f/go.mod h1:GYeBD1CF7AqnKZK+UCytLcY3G+UKo0ByXX/3xfdNyqQ=
 github.com/icza/mighty v0.0.0-20180919140131-cfd07d671de6 h1:8UsGZ2rr2ksmEru6lToqnXgA8Mz1DP11X4zSJ159C3k=
@@ -96,16 +92,9 @@ github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyh
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
-github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c=
-github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
 github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
 github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
-github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
-github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
-github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/tidwall/gjson v1.14.4 h1:uo0p8EbA09J7RQaflQ1aBRffTR7xedD2bcIVSYxLnkM=
@@ -122,8 +111,6 @@ github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5t
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
-golang.org/x/crypto v0.19.0 h1:ENy+Az/9Y1vSrlrvBSyna3PITt4tiZLf7sgCjZBX7Wo=
-golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
 golang.org/x/crypto v0.21.0 h1:X31++rzVUdKhX5sWmSOFZxx8UW/ldWx55cbf08iNAMA=
 golang.org/x/crypto v0.21.0/go.mod h1:0BP7YvVV9gBbVKyeTG0Gyn+gZm94bibOW5BjDEYAOMs=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
@@ -132,8 +119,6 @@ golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLL
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
 golang.org/x/net v0.5.0/go.mod h1:DivGGAXEgPSlEBzxGzZI+ZLohi+xUj054jfeKui00ws=
-golang.org/x/net v0.21.0 h1:AQyQV4dYCvJ7vGmJyKki9+PBdyvhkSd8EIx/qb0AYv4=
-golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44=
 golang.org/x/net v0.22.0 h1:9sGLhx7iRIHEiX0oAJ3MRZMUCElJgy7Br1nO+AMN3Tc=
 golang.org/x/net v0.22.0/go.mod h1:JKghWKKOSdJwpW2GEx0Ja7fmaKnMsbu+MWVZTokSYmg=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
@@ -151,8 +136,6 @@ golang.org/x/sys v0.0.0-20220908164124-27713097b956/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.4.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.17.0 h1:25cE3gD+tdBA7lp7QfhuV+rJiE9YXTcS3VG1SqssI/Y=
-golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4=
 golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
@@ -165,6 +148,8 @@ golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
 golang.org/x/text v0.6.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
 golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
+golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
+golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
@@ -175,7 +160,6 @@ golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2/go.mod h1:K8+ghG5WaK9qNq
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
-gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.0/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

runner/run.go

@@ -6,6 +6,7 @@ package runner
 import (
 	"errors"
 	"fmt"
+	"golang.org/x/time/rate"
 	"regexp"
 	"time"
 
@@ -40,6 +41,9 @@ func Run(cfg *config.FTWConfiguration, tests []*test.FTWTest, c RunnerConfig, ou
 	if c.ReadTimeout != 0 {
 		conf.ReadTimeout = c.ReadTimeout
 	}
+	if c.RateLimit != 0 {
+		conf.RateLimiter = rate.NewLimiter(rate.Every(c.RateLimit), 1)
+	}
 	client, err := ftwhttp.NewClient(conf)
 	if err != nil {
 		return &TestRunContext{}, err

runner/types.go

@@ -29,6 +29,8 @@ type RunnerConfig struct {
 	ConnectTimeout time.Duration
 	// ReadTimeout is the timeout for receiving responses during test execution.
 	ReadTimeout time.Duration
+	// RateLimit is the rate limit for requests to the server. 0 is unlimited.
+	RateLimit time.Duration
 }
 
 // TestRunContext carries information about the current test run.