package middleware import ( "bytes" "net/http" "net/http/httptest" "strconv" "testing" "github.com/gin-gonic/gin" "github.com/user-management-system/internal/config" ) func performRateLimitedRequest(router *gin.Engine, path string, userID int64) *httptest.ResponseRecorder { recorder := httptest.NewRecorder() req := httptest.NewRequest(http.MethodGet, path, nil) req.RemoteAddr = "127.0.0.1:12345" req.Header.Set("X-Test-User-ID", strconv.FormatInt(userID, 10)) router.ServeHTTP(recorder, req) return recorder } func performRefreshRateLimitedRequestWithCookie(router *gin.Engine, refreshToken string) *httptest.ResponseRecorder { recorder := httptest.NewRecorder() req := httptest.NewRequest(http.MethodPost, "/auth/refresh", nil) req.RemoteAddr = "127.0.0.1:12345" if refreshToken != "" { req.AddCookie(&http.Cookie{Name: "ums_refresh_token", Value: refreshToken}) } router.ServeHTTP(recorder, req) return recorder } func performRefreshRateLimitedRequestWithBody(router *gin.Engine, refreshToken string) *httptest.ResponseRecorder { recorder := httptest.NewRecorder() body := bytes.NewBufferString(`{"refresh_token":"` + refreshToken + `"}`) req := httptest.NewRequest(http.MethodPost, "/auth/refresh", body) req.RemoteAddr = "127.0.0.1:12345" req.Header.Set("Content-Type", "application/json") router.ServeHTTP(recorder, req) return recorder } func TestRateLimitMiddleware_API_ScopesBudgetByRouteForAuthenticatedUser(t *testing.T) { gin.SetMode(gin.TestMode) rateLimitMiddleware := NewRateLimitMiddleware(config.RateLimitConfig{}) router := gin.New() router.Use(func(c *gin.Context) { rawUserID := c.GetHeader("X-Test-User-ID") if rawUserID != "" { userID, err := strconv.ParseInt(rawUserID, 10, 64) if err == nil { c.Set("user_id", userID) } } c.Next() }) protected := router.Group("") protected.Use(rateLimitMiddleware.API()) protected.GET("/users", func(c *gin.Context) { c.Status(http.StatusOK) }) protected.GET("/roles", func(c *gin.Context) { c.Status(http.StatusOK) }) for i := 0; i < 100; i++ { recorder := performRateLimitedRequest(router, "/users", 1) if recorder.Code != http.StatusOK { t.Fatalf("request %d to /users returned %d, want %d", i+1, recorder.Code, http.StatusOK) } } sameRouteOverflow := performRateLimitedRequest(router, "/users", 1) if sameRouteOverflow.Code != http.StatusTooManyRequests { t.Fatalf("overflow request to /users returned %d, want %d", sameRouteOverflow.Code, http.StatusTooManyRequests) } differentRoute := performRateLimitedRequest(router, "/roles", 1) if differentRoute.Code != http.StatusOK { t.Fatalf("request to /roles after exhausting /users budget returned %d, want %d", differentRoute.Code, http.StatusOK) } } func TestRateLimitMiddleware_Refresh_ScopesBudgetByRefreshCookie(t *testing.T) { gin.SetMode(gin.TestMode) rateLimitMiddleware := NewRateLimitMiddleware(config.RateLimitConfig{}) router := gin.New() router.POST("/auth/refresh", rateLimitMiddleware.Refresh(), func(c *gin.Context) { c.Status(http.StatusOK) }) for i := 0; i < 10; i++ { recorder := performRefreshRateLimitedRequestWithCookie(router, "refresh-token-a") if recorder.Code != http.StatusOK { t.Fatalf("request %d for refresh-token-a returned %d, want %d", i+1, recorder.Code, http.StatusOK) } } sameTokenOverflow := performRefreshRateLimitedRequestWithCookie(router, "refresh-token-a") if sameTokenOverflow.Code != http.StatusTooManyRequests { t.Fatalf("overflow request for refresh-token-a returned %d, want %d", sameTokenOverflow.Code, http.StatusTooManyRequests) } differentToken := performRefreshRateLimitedRequestWithCookie(router, "refresh-token-b") if differentToken.Code != http.StatusOK { t.Fatalf("request for refresh-token-b after exhausting refresh-token-a budget returned %d, want %d", differentToken.Code, http.StatusOK) } } func TestRateLimitMiddleware_Refresh_ScopesBudgetByRefreshTokenBody(t *testing.T) { gin.SetMode(gin.TestMode) rateLimitMiddleware := NewRateLimitMiddleware(config.RateLimitConfig{}) router := gin.New() router.POST("/auth/refresh", rateLimitMiddleware.Refresh(), func(c *gin.Context) { c.Status(http.StatusOK) }) for i := 0; i < 10; i++ { recorder := performRefreshRateLimitedRequestWithBody(router, "refresh-token-a") if recorder.Code != http.StatusOK { t.Fatalf("request %d for refresh-token-a body returned %d, want %d", i+1, recorder.Code, http.StatusOK) } } sameTokenOverflow := performRefreshRateLimitedRequestWithBody(router, "refresh-token-a") if sameTokenOverflow.Code != http.StatusTooManyRequests { t.Fatalf("overflow request for refresh-token-a body returned %d, want %d", sameTokenOverflow.Code, http.StatusTooManyRequests) } differentToken := performRefreshRateLimitedRequestWithBody(router, "refresh-token-b") if differentToken.Code != http.StatusOK { t.Fatalf("request for refresh-token-b body after exhausting refresh-token-a budget returned %d, want %d", differentToken.Code, http.StatusOK) } }