fix: harden handler context and rate limit isolation

internal/api/middleware/ratelimit_test.go

@@ -0,0 +1,107 @@
+package middleware
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/user-management-system/internal/config"
+)
+
+func init() {
+	gin.SetMode(gin.TestMode)
+}
+
+func newRateLimitTestEngine(mw gin.HandlerFunc) *gin.Engine {
+	engine := gin.New()
+	engine.Use(mw)
+	engine.GET("/ping", func(c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{"ok": true})
+	})
+	return engine
+}
+
+func performRateLimitRequest(engine *gin.Engine, remoteAddr string, setup func(*http.Request)) int {
+	req := httptest.NewRequest(http.MethodGet, "/ping", nil)
+	req.RemoteAddr = remoteAddr
+	if setup != nil {
+		setup(req)
+	}
+	w := httptest.NewRecorder()
+	engine.ServeHTTP(w, req)
+	return w.Code
+}
+
+func TestRateLimitMiddleware_LoginUsesIndependentIPBuckets(t *testing.T) {
+	mw := NewRateLimitMiddleware(config.RateLimitConfig{})
+	engine := newRateLimitTestEngine(mw.Login())
+
+	for i := 0; i < 5; i++ {
+		if code := performRateLimitRequest(engine, "1.1.1.1:1234", nil); code != http.StatusOK {
+			t.Fatalf("ip1 request %d expected 200, got %d", i+1, code)
+		}
+	}
+	if code := performRateLimitRequest(engine, "1.1.1.1:1234", nil); code != http.StatusTooManyRequests {
+		t.Fatalf("ip1 sixth request expected 429, got %d", code)
+	}
+
+	if code := performRateLimitRequest(engine, "2.2.2.2:1234", nil); code != http.StatusOK {
+		t.Fatalf("independent ip should not be throttled, got %d", code)
+	}
+}
+
+func TestRateLimitMiddleware_APIPrefersUserIDOverSharedIP(t *testing.T) {
+	mw := NewRateLimitMiddleware(config.RateLimitConfig{})
+	engine := gin.New()
+	engine.Use(func(c *gin.Context) {
+		if userID := c.GetHeader("X-Test-User-ID"); userID != "" {
+			c.Set("user_id", userID)
+		}
+		c.Next()
+	})
+	engine.Use(mw.limitForKey("api-test", 60, 1))
+	engine.GET("/ping", func(c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{"ok": true})
+	})
+
+	setupUser1 := func(req *http.Request) {
+		req.Header.Set("X-Test-User-ID", "101")
+	}
+	setupUser2 := func(req *http.Request) {
+		req.Header.Set("X-Test-User-ID", "202")
+	}
+
+	if code := performRateLimitRequest(engine, "9.9.9.9:1234", setupUser1); code != http.StatusOK {
+		t.Fatalf("user1 first request expected 200, got %d", code)
+	}
+	if code := performRateLimitRequest(engine, "9.9.9.9:1234", setupUser1); code != http.StatusTooManyRequests {
+		t.Fatalf("user1 second request expected 429, got %d", code)
+	}
+	if code := performRateLimitRequest(engine, "9.9.9.9:1234", setupUser2); code != http.StatusOK {
+		t.Fatalf("user2 should have independent bucket on shared ip, got %d", code)
+	}
+}
+
+func TestRateLimitMiddleware_CleansUpIdleLimiters(t *testing.T) {
+	mw := NewRateLimitMiddleware(config.RateLimitConfig{})
+	mw.cleanupInt = 10 * time.Millisecond
+	engine := newRateLimitTestEngine(mw.limitForKey("cleanup", 1, 2))
+
+	if code := performRateLimitRequest(engine, "3.3.3.3:1234", nil); code != http.StatusOK {
+		t.Fatalf("seed request expected 200, got %d", code)
+	}
+	if got := len(mw.limiters); got != 1 {
+		t.Fatalf("expected 1 limiter after seed request, got %d", got)
+	}
+
+	time.Sleep(1100 * time.Millisecond)
+	if code := performRateLimitRequest(engine, "4.4.4.4:1234", nil); code != http.StatusOK {
+		t.Fatalf("cleanup trigger request expected 200, got %d", code)
+	}
+
+	if got := len(mw.limiters); got != 1 {
+		t.Fatalf("expected stale limiter to be cleaned up, got %d entries", got)
+	}
+}