feat: Add rate limiting

Adds a rate limiting mechanism, that will send HTTP/429 responses once a
defined limit is reached (Token Bucket)

Signed-off-by: Manuel Rüger <[email protected]>

Author: mrueg

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/prometheus/common v0.64.0
 	golang.org/x/crypto v0.39.0
 	golang.org/x/sync v0.15.0
+	golang.org/x/time v0.12.0
 	gopkg.in/yaml.v2 v2.4.0
 )
 

go.sum

@@ -61,6 +61,8 @@ golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw=
 golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
 golang.org/x/text v0.26.0 h1:P42AVeLghgTYr4+xUnTRKDMqpar+PtX7KWuNQL21L8M=
 golang.org/x/text v0.26.0/go.mod h1:QK15LZJUUQVJxhz7wXgxSy/CJaTFjd0G+YLonydOVQA=
+golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE=
+golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg=
 google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY=
 google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

web/handler.go

@@ -24,6 +24,7 @@ import (
 	"sync"
 
 	"golang.org/x/crypto/bcrypt"
+	"golang.org/x/time/rate"
 )
 
 // extraHTTPHeaders is a map of HTTP headers that can be added to HTTP
@@ -80,6 +81,7 @@ type webHandler struct {
 	handler       http.Handler
 	logger        *slog.Logger
 	cache         *cache
+	limiter       *rate.Limiter
 	// bcryptMtx is there to ensure that bcrypt.CompareHashAndPassword is run
 	// only once in parallel as this is CPU intensive.
 	bcryptMtx sync.Mutex
@@ -93,6 +95,11 @@ func (u *webHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	if u.limiter != nil && !u.limiter.Allow() {
+		http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
+		return
+	}
+
 	// Configure http headers.
 	for k, v := range c.HTTPConfig.Header {
 		w.Header().Set(k, v)

web/testdata/web_config_rate_limiter_capacity_one.yaml

@@ -0,0 +1,3 @@
+rate_limiter_config:
+  rate: 1
+  burst: 1

web/testdata/web_config_rate_limiter_nonblocking.yaml

@@ -0,0 +1,3 @@
+rate_limiter_config:
+  rate: 0
+  burst: 0

web/tls_config.go

@@ -31,6 +31,7 @@ import (
 	"github.com/mdlayher/vsock"
 	config_util "github.com/prometheus/common/config"
 	"golang.org/x/sync/errgroup"
+	"golang.org/x/time/rate"
 	"gopkg.in/yaml.v2"
 )
 
@@ -40,9 +41,10 @@ var (
 )
 
 type Config struct {
-	TLSConfig  TLSConfig                     `yaml:"tls_server_config"`
-	HTTPConfig HTTPConfig                    `yaml:"http_server_config"`
-	Users      map[string]config_util.Secret `yaml:"basic_auth_users"`
+	TLSConfig         TLSConfig                     `yaml:"tls_server_config"`
+	HTTPConfig        HTTPConfig                    `yaml:"http_server_config"`
+	RateLimiterConfig RateLimiterConfig             `yaml:"rate_limiter_config"`
+	Users             map[string]config_util.Secret `yaml:"basic_auth_users"`
 }
 
 type TLSConfig struct {
@@ -109,6 +111,11 @@ type HTTPConfig struct {
 	Header map[string]string `yaml:"headers,omitempty"`
 }
 
+type RateLimiterConfig struct {
+	Burst int `yaml:"burst"`
+	Rate  int `yaml:"rate"`
+}
+
 func getConfig(configPath string) (*Config, error) {
 	content, err := os.ReadFile(configPath)
 	if err != nil {
@@ -366,11 +373,19 @@ func Serve(l net.Listener, server *http.Server, flags *FlagConfig, logger *slog.
 		return err
 	}
 
+	var limiter *rate.Limiter
+	// Setup Rate Limiter
+	if c.RateLimiterConfig.Rate != 0 && c.RateLimiterConfig.Burst != 0 {
+		limiter = rate.NewLimiter(rate.Limit(c.RateLimiterConfig.Rate), c.RateLimiterConfig.Burst)
+		logger.Info("Rate Limiter is enabled.", "burst", c.RateLimiterConfig.Burst, "rate", c.RateLimiterConfig.Rate)
+	}
+
 	server.Handler = &webHandler{
 		tlsConfigPath: tlsConfigPath,
 		logger:        logger,
 		handler:       handler,
 		cache:         newCache(),
+		limiter:       limiter,
 	}
 
 	config, err := ConfigToTLSConfig(&c.TLSConfig)

web/tls_config_test.go

@@ -72,6 +72,7 @@ var (
 		// Introduced in Go 1.21
 		"Certificate required": regexp.MustCompile(`certificate required`),
 		"Unknown CA":           regexp.MustCompile(`unknown certificate authority`),
+		"Too Many Requests":    regexp.MustCompile(`Too Many Requests`),
 	}
 )
 
@@ -98,6 +99,7 @@ type TestInputs struct {
 	Username            string
 	Password            string
 	ClientCertificate   string
+	Requests            int
 }
 
 func TestYAMLFiles(t *testing.T) {
@@ -364,6 +366,19 @@ func TestServerBehaviour(t *testing.T) {
 			ClientCertificate: "client2_selfsigned",
 			ExpectedError:     ErrorMap["Invalid client cert"],
 		},
+		{
+			Name:           "valid rate limiter that doesn't block",
+			YAMLConfigPath: "testdata/web_config_rate_limiter_nonblocking.yaml",
+			UseTLSClient:   false,
+			ExpectedError:  nil,
+		},
+		{
+			Name:           "valid rate limiter with a capacity of one",
+			YAMLConfigPath: "testdata/web_config_rate_limiter_capacity_one.yaml",
+			UseTLSClient:   false,
+			Requests:       10,
+			ExpectedError:  ErrorMap["Too Many Requests"],
+		},
 	}
 	for _, testInputs := range testTables {
 		t.Run(testInputs.Name, testInputs.Test)
@@ -511,35 +526,41 @@ func (test *TestInputs) Test(t *testing.T) {
 		if test.Username != "" {
 			req.SetBasicAuth(test.Username, test.Password)
 		}
+
 		return client.Do(req)
 	}
 	go func() {
 		time.Sleep(250 * time.Millisecond)
-		r, err := ClientConnection()
-		if err != nil {
-			recordConnectionError(err)
-			return
-		}
 
-		if test.ActualCipher != 0 {
-			if r.TLS.CipherSuite != test.ActualCipher {
-				recordConnectionError(
-					fmt.Errorf("bad cipher suite selected. Expected: %s, got: %s",
-						tls.CipherSuiteName(test.ActualCipher),
-						tls.CipherSuiteName(r.TLS.CipherSuite),
-					),
-				)
+		for req := 0; req <= test.Requests; req++ {
+
+			r, err := ClientConnection()
+
+			if err != nil {
+				recordConnectionError(err)
+				return
 			}
-		}
 
-		body, err := io.ReadAll(r.Body)
-		if err != nil {
-			recordConnectionError(err)
-			return
-		}
-		if string(body) != "Hello World!" {
-			recordConnectionError(errors.New(string(body)))
-			return
+			if test.ActualCipher != 0 {
+				if r.TLS.CipherSuite != test.ActualCipher {
+					recordConnectionError(
+						fmt.Errorf("bad cipher suite selected. Expected: %s, got: %s",
+							tls.CipherSuiteName(test.ActualCipher),
+							tls.CipherSuiteName(r.TLS.CipherSuite),
+						),
+					)
+				}
+			}
+
+			body, err := io.ReadAll(r.Body)
+			if err != nil {
+				recordConnectionError(err)
+				return
+			}
+			if string(body) != "Hello World!" {
+				recordConnectionError(errors.New(string(body)))
+				return
+			}
 		}
 		recordConnectionError(nil)
 	}()