user-system/internal/api/middleware/ratelimit_test.go

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)
	}
}