Add rate limiting

README.md

@@ -7,14 +7,15 @@ GSS (Go serve SPA) is a containerized web server for single-page applications wr
 - Optimized for single-page apps.
 - Automatically serves pre-compressed brotli and gzip files if available.
 - Sensible default cache configuration.
+- Configurable rate limiter.
 - Optional out-of-the-box metrics.
 - Docker-based.
 - Configurable via YAML.
 - Lightweight.
 
 ## Usage
 
-GSS works as a Docker image. By default it serves a directory in the container named `dist` at port `8080`, and exposes a metrics endpoint at `:9090/metrics` if enabled.
+GSS works as a Docker image. By default it serves a directory in the container named `dist` at port `8080`.
 
 ### Running container directly
 
@@ -106,6 +107,20 @@ Enables metrics collection and exposes an endpoint at `:9090/metrics`. Collected
 > metrics: true
 > ```
 
+### Rate limit per minute: `rateLimit`
+
+##### string: integer
+
+Enables rate limiting per minute per IP using a memory store. 15 by default.
+
+> Example:
+>
+> ```yaml
+> # gss.yaml
+>
+> rateLimit: 10
+> ```
+
 ## Contributing
 
 This project started as a way to learn and to solve a need I had. If you think it can be improved in any way, you are very welcome to contribute!

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/felixge/httpsnoop v1.0.3
 	github.com/prometheus/client_golang v1.15.1
 	github.com/rs/zerolog v1.29.1
+	github.com/sethvargo/go-limiter v0.7.2
 	github.com/stretchr/testify v1.8.2
 	gopkg.in/yaml.v2 v2.4.0
 )
@@ -15,13 +16,15 @@ require (
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/golang/protobuf v1.5.3 // indirect
+	github.com/kr/text v0.2.0 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.18 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/prometheus/client_model v0.3.0 // indirect
 	github.com/prometheus/common v0.42.0 // indirect
 	github.com/prometheus/procfs v0.9.0 // indirect
+	github.com/rogpeppe/go-internal v1.10.0 // indirect
 	golang.org/x/sys v0.7.0 // indirect
 	google.golang.org/protobuf v1.30.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect

gss.go

@@ -14,6 +14,8 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
+	"github.com/sethvargo/go-limiter/httplimit"
+	"github.com/sethvargo/go-limiter/memorystore"
 	"gopkg.in/yaml.v2"
 )
 
@@ -22,7 +24,7 @@ func main() {
 
 	var metrics *metrics
 	cfg := newConfig().withYAML()
-	if cfg.Metrics {
+	if cfg.MetricsEnabled {
 		metrics = registerMetrics()
 		internalServer := newInternalServer(metrics)
 		go func() {
@@ -48,12 +50,20 @@ func setUpLogger() {
 }
 
 type config struct {
-	Headers map[string]string `yaml:"headers,omitempty"`
-	Metrics bool              `yaml:"metrics,omitempty"`
+	ResponseHeaders    map[string]string `yaml:"headers,omitempty"`
+	MetricsEnabled     bool              `yaml:"metrics,omitempty"`
+	RateLimitPerMinute int               `yaml:"rateLimit,omitempty"`
 }
 
 func newConfig() *config {
-	return &config{}
+	return &config{
+		// Default values
+		ResponseHeaders: map[string]string{
+			"Server": "GSS",
+		},
+		MetricsEnabled:     false,
+		RateLimitPerMinute: 15,
+	}
 }
 
 func (c *config) withYAML() *config {
@@ -95,10 +105,23 @@ func newFileServer(cfg *config, metrics *metrics) *fileServer {
 }
 
 func (f *fileServer) init() *fileServer {
-	if f.Config.Metrics {
-		f.Server.Handler = metricsMiddleware(f.Metrics)(f.setHeaders((f.serveSPA())))
+	store, err := memorystore.New(&memorystore.Config{
+		Tokens:   uint64(f.Config.RateLimitPerMinute),
+		Interval: time.Minute,
+	})
+	if err != nil {
+		log.Fatal().Msgf("Error creating rate limit store: %v", err)
+	}
+
+	rateLimit, err := httplimit.NewMiddleware(store, httplimit.IPKeyFunc())
+	if err != nil {
+		log.Fatal().Msgf("Error creating rate limit middleware: %v", err)
+	}
+
+	if f.Config.MetricsEnabled {
+		f.Server.Handler = metricsMiddleware(f.Metrics)(rateLimit.Handle(f.setHeaders((f.serveSPA()))))
 	} else {
-		f.Server.Handler = f.setHeaders((f.serveSPA()))
+		f.Server.Handler = rateLimit.Handle(f.setHeaders((f.serveSPA())))
 	}
 
 	return f
@@ -112,7 +135,7 @@ func (f *fileServer) setHeaders(h http.Handler) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Vary", "Accept-Encoding")
 
-		for k, v := range f.Config.Headers {
+		for k, v := range f.Config.ResponseHeaders {
 			w.Header().Set(k, v)
 		}
 

gss_test.go

@@ -15,7 +15,7 @@ func TestGSS(t *testing.T) {
 		t.Parallel()
 
 		cfg := &config{
-			Headers: map[string]string{
+			ResponseHeaders: map[string]string{
 				"X-Test": "test",
 			},
 		}
@@ -28,6 +28,21 @@ func TestGSS(t *testing.T) {
 		assert.Equal(t, "test", w.Header().Get("X-Test"))
 	})
 
+	t.Run("rate limits requests", func(t *testing.T) {
+		t.Parallel()
+
+		cfg := &config{
+			RateLimitPerMinute: 10,
+		}
+		fileServer := newFileServer(cfg, metrics).init()
+		w := httptest.NewRecorder()
+		r := httptest.NewRequest(http.MethodGet, "/", nil)
+
+		fileServer.Server.Handler.ServeHTTP(w, r)
+
+		assert.Equal(t, "10", w.Header().Get("X-Ratelimit-Limit"))
+	})
+
 	t.Run("redirects index correctly", func(t *testing.T) {
 		t.Parallel()
 

test/gss.yaml

@@ -1,5 +1,5 @@
 headers:
   Referrer-Policy: "strict-origin-when-cross-origin"
-  Server: "GSS"
   Strict-Transport-Security: "max-age=63072000; includeSubDomains; preload"
 metrics: true
+rateLimit: 10