fix: always rate limit, fall back to client IP when header absent

api/middleware.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/json"
 	"io"
+	"net"
 	"net/http"
 
 	"github.com/didip/tollbooth/v5"
@@ -142,17 +143,32 @@ func (a *API) loadInstanceConfig(w http.ResponseWriter, r *http.Request) (contex
 func (a *API) limitHandler(lmt *limiter.Limiter) middlewareHandler {
 	return func(w http.ResponseWriter, req *http.Request) (context.Context, error) {
 		c := req.Context()
-		if limitHeader := a.config.RateLimitHeader; limitHeader != "" {
-			key := req.Header.Get(a.config.RateLimitHeader)
-			err := tollbooth.LimitByKeys(lmt, []string{key})
-			if err != nil {
-				return c, httpError(http.StatusTooManyRequests, "Rate limit exceeded")
-			}
+		key := a.rateLimitKey(req)
+		if err := tollbooth.LimitByKeys(lmt, []string{key}); err != nil {
+			return c, tooManyRequestsError("Rate limit exceeded")
 		}
 		return c, nil
 	}
 }
 
+// rateLimitKey returns the value to bucket rate limits by. It prefers the
+// configured header (typically set by an upstream proxy) and falls back to the
+// client IP so requests are always rate limited even when the header is
+// missing or unset. chi/middleware.RealIP rewrites RemoteAddr from
+// X-Forwarded-For, so this remains correct behind a proxy.
+func (a *API) rateLimitKey(req *http.Request) string {
+	if h := a.config.RateLimitHeader; h != "" {
+		if v := req.Header.Get(h); v != "" {
+			return v
+		}
+	}
+	host, _, err := net.SplitHostPort(req.RemoteAddr)
+	if err != nil {
+		return req.RemoteAddr
+	}
+	return host
+}
+
 func (a *API) verifyOperatorRequest(w http.ResponseWriter, req *http.Request) (context.Context, error) {
 	c, _, err := a.extractOperatorRequest(w, req)
 	return c, err

api/middleware_ratelimit_test.go

@@ -0,0 +1,42 @@
+package api
+
+import (
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/netlify/gotrue/conf"
+)
+
+func TestRateLimitKey(t *testing.T) {
+	tests := []struct {
+		name       string
+		headerName string
+		headerVal  string
+		remoteAddr string
+		want       string
+	}{
+		{"configured header present", "X-Forwarded-For", "1.2.3.4", "10.0.0.1:54321", "1.2.3.4"},
+		{"configured header empty falls back to ip", "X-Forwarded-For", "", "10.0.0.1:54321", "10.0.0.1"},
+		{"no header configured falls back to ip", "", "", "10.0.0.1:54321", "10.0.0.1"},
+		{"remoteAddr without port falls back as-is", "", "", "10.0.0.1", "10.0.0.1"},
+		{"ipv6 remoteAddr", "", "", "[::1]:12345", "::1"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			config := &conf.GlobalConfiguration{}
+			config.RateLimitHeader = tt.headerName
+			a := &API{config: config}
+
+			req := httptest.NewRequest("POST", "/token", nil)
+			req.RemoteAddr = tt.remoteAddr
+			if tt.headerName != "" && tt.headerVal != "" {
+				req.Header.Set(tt.headerName, tt.headerVal)
+			}
+
+			assert.Equal(t, tt.want, a.rateLimitKey(req))
+		})
+	}
+}