feat: Add Rate Limiting

docs/web-config.yml

@@ -10,3 +10,8 @@ tls_server_config:
 basic_auth_users:
   alice: $2y$10$mDwo.lAisC94iLAyP81MCesa29IzH37oigHC/42V2pdJlUprsJPze
   bob: $2y$10$hLqFl9jSjoAAy95Z/zw8Ye8wkdMBM8c5Bn1ptYqP/AXyV0.oy0S8m
+
+# Rate limiting requests on the endpoint using a token bucket
+rate_limit:
+  interval: "1s" # time interval between two requests, set to 0 to disable rate limiter
+  burst: 20 # and permits a burst of up to 20 requests.

docs/web-configuration.md

@@ -20,6 +20,7 @@ Generic placeholders are defined as follows:
 * `<filename>`: a valid path in the current working directory
 * `<secret>`: a regular string that is a secret, such as a password
 * `<string>`: a regular string
+* `<int>`: a regular integer
 
 ```
 tls_server_config:
@@ -125,6 +126,12 @@ http_server_config:
 # required. Passwords are hashed with bcrypt.
 basic_auth_users:
   [ <string>: <secret> ... ]
+
+
+# Rate limiting requests on the endpoint using a token bucket
+rate_limit:
+  interval: <duration> # time interval between two requests, set to 0 to disable rate limiter
+  burst: <int> # and permits a burst of <int> requests.
 ```
 
 [A sample configuration file](web-config.yml) is provided.

go.mod

@@ -10,6 +10,7 @@ require (
 	go.yaml.in/yaml/v2 v2.4.2
 	golang.org/x/crypto v0.41.0
 	golang.org/x/sync v0.16.0
+	golang.org/x/time v0.12.0
 )
 
 require (

go.sum

@@ -62,6 +62,8 @@ golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI=
 golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
 golang.org/x/text v0.28.0 h1:rhazDwis8INMIwQ4tpjLDzUhx6RlXqZNPEM0huQojng=
 golang.org/x/text v0.28.0/go.mod h1:U8nCwOR8jO/marOQ0QbDiOngZVEBB7MAiitBuMjXiNU=
+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.8 h1:xHScyCOEuuwZEc6UtSOvPbAT4zRh0xcNRYekJwfqyMc=
 google.golang.org/protobuf v1.36.8/go.mod h1:fuxRtAxBytpl4zzqUh6/eyUujkJdNiuEkXntxiD/uRU=
 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_nonblocking.yaml

@@ -0,0 +1,3 @@
+rate_limit:
+  interval: 0
+  burst: 0

web/testdata/web_config_rate_limiter_one_second.yaml

@@ -0,0 +1,3 @@
+rate_limit:
+  interval: "1s"
+  burst: 0

web/tls_config.go

@@ -26,12 +26,14 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/coreos/go-systemd/v22/activation"
 	"github.com/mdlayher/vsock"
 	config_util "github.com/prometheus/common/config"
 	"go.yaml.in/yaml/v2"
 	"golang.org/x/sync/errgroup"
+	"golang.org/x/time/rate"
 )
 
 var (
@@ -40,9 +42,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_limit"`
+	Users             map[string]config_util.Secret `yaml:"basic_auth_users"`
 }
 
 type TLSConfig struct {
@@ -109,6 +112,11 @@ type HTTPConfig struct {
 	Header map[string]string `yaml:"headers,omitempty"`
 }
 
+type RateLimiterConfig struct {
+	Burst    int           `yaml:"burst"`
+	Interval time.Duration `yaml:"interval"`
+}
+
 func getConfig(configPath string) (*Config, error) {
 	content, err := os.ReadFile(configPath)
 	if err != nil {
@@ -365,11 +373,18 @@ func Serve(l net.Listener, server *http.Server, flags *FlagConfig, logger *slog.
 		return err
 	}
 
+	var limiter *rate.Limiter
+	if c.RateLimiterConfig.Interval != 0 {
+		limiter = rate.NewLimiter(rate.Every(c.RateLimiterConfig.Interval), c.RateLimiterConfig.Burst)
+		logger.Info("Rate Limiter is enabled.", "burst", c.RateLimiterConfig.Burst, "interval", c.RateLimiterConfig.Interval)
+	}
+
 	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,20 @@ func TestServerBehaviour(t *testing.T) {
 			ClientCertificate: "client2_selfsigned",
 			ExpectedError:     ErrorMap["Invalid client cert"],
 		},
+		{
+			Name:           "valid rate limiter (no rate limiter set up) that doesn't block",
+			YAMLConfigPath: "testdata/web_config_rate_limiter_nonblocking.yaml",
+			UseTLSClient:   false,
+			Requests:       10,
+			ExpectedError:  nil,
+		},
+		{
+			Name:           "valid rate limiter with an interval of one second",
+			YAMLConfigPath: "testdata/web_config_rate_limiter_one_second.yaml",
+			UseTLSClient:   false,
+			Requests:       10,
+			ExpectedError:  ErrorMap["Too Many Requests"],
+		},
 	}
 	for _, testInputs := range testTables {
 		t.Run(testInputs.Name, testInputs.Test)
@@ -511,35 +527,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)
 	}()