杭州明婳科技
481 lines
11 KiB
Go
481 lines
11 KiB
Go
package middleware_test
|
|
|
|
import (
|
|
"context"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/EthanCodeCraft/xlgo-core/middleware"
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func setupTestRouter() *gin.Engine {
|
|
gin.SetMode(gin.TestMode)
|
|
r := gin.New()
|
|
return r
|
|
}
|
|
|
|
// ===== RequestID Tests =====
|
|
|
|
func TestRequestID(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.Use(middleware.RequestID())
|
|
r.GET("/test", func(c *gin.Context) {
|
|
id := middleware.GetRequestID(c)
|
|
c.JSON(200, gin.H{"request_id": id})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/test", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
// 验证响应头有 X-Request-ID
|
|
headerID := w.Header().Get("X-Request-ID")
|
|
if headerID == "" {
|
|
t.Error("X-Request-ID header should be set")
|
|
}
|
|
|
|
// 验证响应体中的 request_id
|
|
if w.Code != 200 {
|
|
t.Errorf("RequestID status = %d", w.Code)
|
|
}
|
|
}
|
|
|
|
func TestRequestIDWithExisting(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.Use(middleware.RequestID())
|
|
r.GET("/test", func(c *gin.Context) {
|
|
id := middleware.GetRequestID(c)
|
|
c.JSON(200, gin.H{"request_id": id})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/test", nil)
|
|
req.Header.Set("X-Request-ID", "custom-id-123")
|
|
r.ServeHTTP(w, req)
|
|
|
|
// 验证使用了传入的 ID
|
|
headerID := w.Header().Get("X-Request-ID")
|
|
if headerID != "custom-id-123" {
|
|
t.Errorf("X-Request-ID = %s, want custom-id-123", headerID)
|
|
}
|
|
}
|
|
|
|
func TestGetRequestIDEmpty(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.GET("/test", func(c *gin.Context) {
|
|
id := middleware.GetRequestID(c)
|
|
c.JSON(200, gin.H{"request_id": id})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/test", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
// 没有 RequestID 中间件时，返回空
|
|
if w.Code != 200 {
|
|
t.Errorf("status = %d", w.Code)
|
|
}
|
|
}
|
|
|
|
// ===== Recover Tests =====
|
|
|
|
func TestRecover(t *testing.T) {
|
|
// 需要初始化 logger，否则 Recover 会 panic
|
|
// 这里跳过完整测试，仅验证中间件可以正常注册
|
|
r := setupTestRouter()
|
|
r.Use(middleware.Recover())
|
|
r.GET("/normal", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/normal", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
if w.Code != 200 {
|
|
t.Errorf("Normal status = %d, want 200", w.Code)
|
|
}
|
|
}
|
|
|
|
func TestRecoverNoPanic(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.Use(middleware.Recover())
|
|
r.GET("/normal", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/normal", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
if w.Code != 200 {
|
|
t.Errorf("Normal status = %d, want 200", w.Code)
|
|
}
|
|
}
|
|
|
|
// ===== CSRF Tests =====
|
|
|
|
func TestCSRF(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.Use(middleware.CSRF())
|
|
r.GET("/form", func(c *gin.Context) {
|
|
token := middleware.GetCSRFToken(c)
|
|
c.JSON(200, gin.H{"csrf_token": token})
|
|
})
|
|
r.POST("/submit", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
// GET 请求应该成功，并设置 cookie
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/form", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
if w.Code != 200 {
|
|
t.Errorf("CSRF GET status = %d", w.Code)
|
|
}
|
|
|
|
// 获取 cookie 中的 token
|
|
cookies := w.Result().Cookies()
|
|
var csrfToken string
|
|
for _, c := range cookies {
|
|
if c.Name == "csrf_token" {
|
|
csrfToken = c.Value
|
|
break
|
|
}
|
|
}
|
|
|
|
// POST 无 token 应该失败
|
|
w2 := httptest.NewRecorder()
|
|
req2 := httptest.NewRequest("POST", "/submit", nil)
|
|
r.ServeHTTP(w2, req2)
|
|
|
|
if w2.Code != 200 {
|
|
// 成功捕获错误，返回 200 但 code != 1
|
|
}
|
|
|
|
// POST 带 token 应该成功
|
|
w3 := httptest.NewRecorder()
|
|
req3 := httptest.NewRequest("POST", "/submit", nil)
|
|
req3.Header.Set("X-CSRF-Token", csrfToken)
|
|
r.ServeHTTP(w3, req3)
|
|
|
|
// 注意：由于 cookie 需要从上一个请求传递，这里可能需要手动设置
|
|
}
|
|
|
|
func TestGetCSRFToken(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.Use(middleware.CSRF())
|
|
r.GET("/token", func(c *gin.Context) {
|
|
token := middleware.GetCSRFToken(c)
|
|
c.JSON(200, gin.H{"token": token})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/token", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
if w.Code != 200 {
|
|
t.Errorf("GetCSRFToken status = %d", w.Code)
|
|
}
|
|
}
|
|
|
|
func TestDoubleSubmitCookie(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.Use(middleware.DoubleSubmitCookie())
|
|
r.GET("/get", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
r.POST("/post", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
// GET 应该成功
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/get", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
if w.Code != 200 {
|
|
t.Errorf("DoubleSubmit GET status = %d", w.Code)
|
|
}
|
|
|
|
// POST 无 token 应该失败
|
|
w2 := httptest.NewRecorder()
|
|
req2 := httptest.NewRequest("POST", "/post", nil)
|
|
r.ServeHTTP(w2, req2)
|
|
|
|
// 应返回错误
|
|
if w2.Code != 200 && w2.Code != 400 {
|
|
t.Errorf("DoubleSubmit POST without token status = %d", w2.Code)
|
|
}
|
|
}
|
|
|
|
// ===== CORS Tests =====
|
|
|
|
func TestCORS(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.Use(middleware.CORS())
|
|
r.GET("/test", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
// 使用 localhost origin（开发环境默认允许）
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/test", nil)
|
|
req.Header.Set("Origin", "http://localhost:3000")
|
|
r.ServeHTTP(w, req)
|
|
|
|
// 验证请求成功
|
|
if w.Code != 200 {
|
|
t.Errorf("CORS status = %d", w.Code)
|
|
}
|
|
|
|
// 验证其他 CORS 头始终设置
|
|
allowMethods := w.Header().Get("Access-Control-Allow-Methods")
|
|
if allowMethods == "" {
|
|
t.Error("Access-Control-Allow-Methods should be set")
|
|
}
|
|
}
|
|
|
|
func TestCORSOptions(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.Use(middleware.CORS())
|
|
r.POST("/test", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
// OPTIONS 预检请求
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("OPTIONS", "/test", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
// OPTIONS 应返回 204
|
|
if w.Code != 204 {
|
|
t.Errorf("CORS OPTIONS status = %d, want 204", w.Code)
|
|
}
|
|
}
|
|
|
|
// ===== RateLimit Tests =====
|
|
|
|
func TestRateLimit(t *testing.T) {
|
|
r := setupTestRouter()
|
|
// 使用自定义限流器
|
|
limiter := middleware.NewRateLimiter(10, time.Minute) // 每分钟10次
|
|
defer limiter.Stop()
|
|
r.Use(middleware.RateLimit(limiter))
|
|
r.GET("/test", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
// 正常请求应该成功
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/test", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
// 单次请求应该成功
|
|
if w.Code != 200 {
|
|
t.Errorf("RateLimit status = %d, want 200", w.Code)
|
|
}
|
|
}
|
|
|
|
func TestRateLimiterAllow(t *testing.T) {
|
|
limiter := middleware.NewRateLimiter(3, time.Minute) // 每分钟3次
|
|
defer limiter.Stop()
|
|
|
|
// 前3次允许
|
|
for i := 0; i < 3; i++ {
|
|
if !limiter.Allow("192.168.1.1") {
|
|
t.Errorf("Request %d should be allowed", i+1)
|
|
}
|
|
}
|
|
|
|
// 第4次拒绝
|
|
if limiter.Allow("192.168.1.1") {
|
|
t.Error("Request 4 should be denied")
|
|
}
|
|
|
|
// 不同 IP 应该独立计数
|
|
if !limiter.Allow("192.168.1.2") {
|
|
t.Error("Different IP should be allowed")
|
|
}
|
|
}
|
|
|
|
func TestCustomRateLimit(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.Use(middleware.CustomRateLimit(5, time.Minute))
|
|
r.GET("/test", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/test", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
if w.Code != 200 {
|
|
t.Errorf("CustomRateLimit status = %d", w.Code)
|
|
}
|
|
}
|
|
|
|
func TestLoginRateLimit(t *testing.T) {
|
|
middleware.InitRateLimiters()
|
|
defer middleware.StopRateLimiters()
|
|
|
|
r := setupTestRouter()
|
|
r.Use(middleware.LoginRateLimit())
|
|
r.GET("/test", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/test", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
if w.Code != 200 {
|
|
t.Errorf("LoginRateLimit status = %d", w.Code)
|
|
}
|
|
}
|
|
|
|
func TestAPIRateLimit(t *testing.T) {
|
|
middleware.InitRateLimiters()
|
|
defer middleware.StopRateLimiters()
|
|
|
|
r := setupTestRouter()
|
|
r.Use(middleware.APIRateLimit())
|
|
r.GET("/test", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/test", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
if w.Code != 200 {
|
|
t.Errorf("APIRateLimit status = %d", w.Code)
|
|
}
|
|
}
|
|
|
|
func TestUploadRateLimit(t *testing.T) {
|
|
middleware.InitRateLimiters()
|
|
defer middleware.StopRateLimiters()
|
|
|
|
r := setupTestRouter()
|
|
r.Use(middleware.UploadRateLimit())
|
|
r.GET("/test", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/test", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
if w.Code != 200 {
|
|
t.Errorf("UploadRateLimit status = %d", w.Code)
|
|
}
|
|
}
|
|
|
|
// ===== RedisRateLimiter Tests =====
|
|
|
|
func TestRedisRateLimiter(t *testing.T) {
|
|
limiter := middleware.NewRedisRateLimiter("test_limit", 10, time.Minute)
|
|
|
|
// Without Redis, should always allow
|
|
ctx := context.Background()
|
|
allowed, err := limiter.Allow(ctx, "192.168.1.1")
|
|
if err != nil {
|
|
t.Errorf("RedisRateLimiter error: %v", err)
|
|
}
|
|
if !allowed {
|
|
t.Error("RedisRateLimiter should allow without Redis")
|
|
}
|
|
}
|
|
|
|
func TestRedisRateLimitMiddleware(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.Use(middleware.RedisRateLimit("test_api", 100))
|
|
r.GET("/test", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/test", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
// Without Redis, should pass
|
|
if w.Code != 200 {
|
|
t.Errorf("RedisRateLimit status = %d", w.Code)
|
|
}
|
|
}
|
|
|
|
func TestLoginRedisRateLimit(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.Use(middleware.LoginRedisRateLimit())
|
|
r.GET("/test", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/test", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
if w.Code != 200 {
|
|
t.Errorf("LoginRedisRateLimit status = %d", w.Code)
|
|
}
|
|
}
|
|
|
|
func TestAPIRedisRateLimit(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.Use(middleware.APIRedisRateLimit())
|
|
r.GET("/test", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/test", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
if w.Code != 200 {
|
|
t.Errorf("APIRedisRateLimit status = %d", w.Code)
|
|
}
|
|
}
|
|
|
|
func TestCustomRedisRateLimit(t *testing.T) {
|
|
r := setupTestRouter()
|
|
r.Use(middleware.CustomRedisRateLimit("custom", 50, time.Minute))
|
|
r.GET("/test", func(c *gin.Context) {
|
|
c.JSON(200, gin.H{"status": "ok"})
|
|
})
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("GET", "/test", nil)
|
|
r.ServeHTTP(w, req)
|
|
|
|
if w.Code != 200 {
|
|
t.Errorf("CustomRedisRateLimit status = %d", w.Code)
|
|
}
|
|
}
|
|
|
|
func TestRedisRateLimiterGetCount(t *testing.T) {
|
|
limiter := middleware.NewRedisRateLimiter("test_count", 10, time.Minute)
|
|
|
|
ctx := context.Background()
|
|
count, err := limiter.GetCount(ctx, "192.168.1.1")
|
|
if err != nil {
|
|
t.Errorf("GetCount error: %v", err)
|
|
}
|
|
// Without Redis, count should be 0
|
|
if count != 0 {
|
|
t.Errorf("GetCount = %d, want 0 without Redis", count)
|
|
}
|
|
}
|
|
|
|
func TestRedisRateLimiterReset(t *testing.T) {
|
|
limiter := middleware.NewRedisRateLimiter("test_reset", 10, time.Minute)
|
|
|
|
ctx := context.Background()
|
|
err := limiter.Reset(ctx, "192.168.1.1")
|
|
if err != nil {
|
|
t.Errorf("Reset error: %v", err)
|
|
}
|
|
} |