144 lines
3.8 KiB
Go
144 lines
3.8 KiB
Go
package middlewares
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
var (
|
|
limit = 1
|
|
burst = 60
|
|
)
|
|
|
|
func TestRateLimitMiddleware(t *testing.T) {
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
// Create a new rate limiter that allows 2 requests per second with a burst of 4
|
|
|
|
// Create a new Gin router with the rate limit middleware
|
|
router := setupRouter()
|
|
|
|
// Test cases
|
|
tests := []struct {
|
|
name string
|
|
requests int
|
|
expectedStatus int
|
|
sleep time.Duration
|
|
}{
|
|
{"Allow first request", 1, http.StatusOK, 0},
|
|
{"Allow up to burst limit", burst - 1, http.StatusOK, 0},
|
|
{"Block after burst limit", burst + 20, http.StatusTooManyRequests, 0},
|
|
{"Allow after rate limit replenishes", 1, http.StatusOK, time.Second},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
time.Sleep(tt.sleep)
|
|
var status int
|
|
for i := 0; i < tt.requests; i++ {
|
|
status = makeRequest(router, "192.168.0.2")
|
|
|
|
}
|
|
if status != tt.expectedStatus {
|
|
t.Errorf("Expected status %d, got %d", tt.expectedStatus, status)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestIPRateLimiter(t *testing.T) {
|
|
limiter := NewIPRateLimiter(limit, burst)
|
|
|
|
limiter1 := limiter.GetLimiter("127.0.0.1")
|
|
limiter2 := limiter.GetLimiter("192.168.0.1")
|
|
|
|
if limiter1 == limiter2 {
|
|
t.Error("Expected different limiters for different IPs")
|
|
}
|
|
|
|
limiter3 := limiter.GetLimiter("127.0.0.1")
|
|
if limiter1 != limiter3 {
|
|
t.Error("Expected the same limiter for the same IP")
|
|
}
|
|
}
|
|
|
|
func TestDifferentRateLimits(t *testing.T) {
|
|
testCases := []struct {
|
|
name string
|
|
duration time.Duration
|
|
requests int
|
|
expectedRequests int
|
|
ip string
|
|
}{
|
|
{"Low rate", 5 * time.Second, burst + 5*limit, burst + 5*limit, "192.168.23.3"},
|
|
{"Low rate with limiting", 4 * time.Second, burst + 4*limit + 4, burst + 4*limit, "192.168.23.4"},
|
|
{"High rate", 1 * time.Second, burst + 5*limit, burst + limit, "192.168.23.5"},
|
|
{"Fractional rate", 10 * time.Second, (burst + limit) / 2, (burst + limit) / 2, "192.168.23.6"},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
router := setupRouter()
|
|
testRateLimit(t, router, tc.duration, tc.requests, tc.expectedRequests, tc.ip)
|
|
})
|
|
}
|
|
}
|
|
|
|
func testRateLimit(t *testing.T, router *gin.Engine, duration time.Duration, requests int, expectedRequests int, ip string) {
|
|
start := time.Now()
|
|
successCount := 0
|
|
totalRequests := 0
|
|
|
|
t.Logf("Sleeping for: %v", time.Duration(duration.Nanoseconds()/int64(requests)))
|
|
for time.Since(start) < duration {
|
|
status := makeRequest(router, ip)
|
|
totalRequests++
|
|
|
|
if status == http.StatusOK {
|
|
successCount++
|
|
}
|
|
|
|
time.Sleep(time.Duration(duration.Nanoseconds() / int64(requests)))
|
|
}
|
|
|
|
actualDuration := time.Since(start)
|
|
|
|
t.Logf("limit: %v, burst: %v", limit, burst)
|
|
t.Logf("Test duration: %v", actualDuration)
|
|
t.Logf("Successful requests: %d", successCount)
|
|
t.Logf("Expected successful requests: %d", expectedRequests)
|
|
t.Logf("Total requests: %d", totalRequests)
|
|
|
|
if successCount < int(expectedRequests)-4 || successCount > int(expectedRequests)+4 {
|
|
t.Errorf("Expected around %d successful requests, got %d", expectedRequests, successCount)
|
|
}
|
|
|
|
if requests-expectedRequests != 0 && totalRequests <= successCount {
|
|
t.Errorf("Expected some requests to be rate limited")
|
|
}
|
|
}
|
|
|
|
func setupRouter() *gin.Engine {
|
|
gin.SetMode(gin.TestMode)
|
|
router := gin.New()
|
|
limiter := NewIPRateLimiter(limit, burst)
|
|
router.Use(RateLimitMiddleware(limiter))
|
|
router.GET("/test", func(c *gin.Context) {
|
|
c.String(http.StatusOK, "success")
|
|
})
|
|
return router
|
|
}
|
|
|
|
func makeRequest(router *gin.Engine, ip string) int {
|
|
w := httptest.NewRecorder()
|
|
req, _ := http.NewRequest("GET", "/test", nil)
|
|
req.Header.Set("X-Forwarded-For", ip) // Set a consistent IP
|
|
router.ServeHTTP(w, req)
|
|
return w.Code
|
|
}
|