xlgo-core/middleware/middleware_test.go

package middleware_test

import (
	"context"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/EthanCodeCraft/xlgo-core/config"
	"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)
	}
}

// CORS 规范要求：AllowCredentials=false 时不应发送 Access-Control-Allow-Credentials 头。
// 历史 bug：旧实现在 if/else 两个分支都设了 "true"，相当于永远允许凭证。
func TestCORSAllowCredentialsDefault(t *testing.T) {
	r := setupTestRouter()
	r.Use(middleware.CORSWithConfig(&config.CORSConfig{
		AllowedOrigins:   []string{"https://example.com"},
		AllowCredentials: false,
	}))
	r.GET("/test", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "ok"})
	})

	w := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/test", nil)
	req.Header.Set("Origin", "https://example.com")
	r.ServeHTTP(w, req)

	if got := w.Header().Get("Access-Control-Allow-Credentials"); got != "" {
		t.Errorf("Access-Control-Allow-Credentials = %q, want empty when AllowCredentials=false", got)
	}
}

// AllowCredentials=true 且 Origin 在白名单内：发送凭证头并回显具体 Origin。
func TestCORSAllowCredentialsExplicitOrigin(t *testing.T) {
	r := setupTestRouter()
	r.Use(middleware.CORSWithConfig(&config.CORSConfig{
		AllowedOrigins:   []string{"https://example.com"},
		AllowCredentials: true,
	}))
	r.GET("/test", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "ok"})
	})

	w := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/test", nil)
	req.Header.Set("Origin", "https://example.com")
	r.ServeHTTP(w, req)

	if got := w.Header().Get("Access-Control-Allow-Credentials"); got != "true" {
		t.Errorf("Access-Control-Allow-Credentials = %q, want \"true\"", got)
	}
	if got := w.Header().Get("Access-Control-Allow-Origin"); got != "https://example.com" {
		t.Errorf("Access-Control-Allow-Origin = %q, want \"https://example.com\"", got)
	}
	if got := w.Header().Get("Vary"); got != "Origin" {
		t.Errorf("Vary = %q, want \"Origin\"", got)
	}
}

// CORS 规范：AllowCredentials=true 时禁止使用 "*"，必须回显具体 Origin。
func TestCORSWildcardWithCredentials(t *testing.T) {
	r := setupTestRouter()
	r.Use(middleware.CORSWithConfig(&config.CORSConfig{
		AllowedOrigins:   []string{"*"},
		AllowCredentials: true,
	}))
	r.GET("/test", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "ok"})
	})

	w := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/test", nil)
	req.Header.Set("Origin", "https://anywhere.example")
	r.ServeHTTP(w, req)

	if got := w.Header().Get("Access-Control-Allow-Origin"); got == "*" {
		t.Errorf("Access-Control-Allow-Origin must NOT be \"*\" when credentials are allowed, got %q", got)
	}
	if got := w.Header().Get("Access-Control-Allow-Origin"); got != "https://anywhere.example" {
		t.Errorf("Access-Control-Allow-Origin = %q, want echoed origin", got)
	}
	if got := w.Header().Get("Access-Control-Allow-Credentials"); got != "true" {
		t.Errorf("Access-Control-Allow-Credentials = %q, want \"true\"", got)
	}
}

// AllowedOrigins=["*"] 且 AllowCredentials=false：保持 "*" 通配符语义。
func TestCORSWildcardWithoutCredentials(t *testing.T) {
	r := setupTestRouter()
	r.Use(middleware.CORSWithConfig(&config.CORSConfig{
		AllowedOrigins:   []string{"*"},
		AllowCredentials: false,
	}))
	r.GET("/test", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "ok"})
	})

	w := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/test", nil)
	req.Header.Set("Origin", "https://anywhere.example")
	r.ServeHTTP(w, req)

	if got := w.Header().Get("Access-Control-Allow-Origin"); got != "*" {
		t.Errorf("Access-Control-Allow-Origin = %q, want \"*\"", got)
	}
	if got := w.Header().Get("Access-Control-Allow-Credentials"); got != "" {
		t.Errorf("Access-Control-Allow-Credentials = %q, want empty (AllowCredentials=false)", got)
	}
}

// Origin 不在白名单时不应回显，避免反射型 CORS 漏洞。
func TestCORSOriginNotAllowed(t *testing.T) {
	r := setupTestRouter()
	r.Use(middleware.CORSWithConfig(&config.CORSConfig{
		AllowedOrigins:   []string{"https://allowed.example"},
		AllowCredentials: true,
	}))
	r.GET("/test", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "ok"})
	})

	w := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/test", nil)
	req.Header.Set("Origin", "https://evil.example")
	r.ServeHTTP(w, req)

	if got := w.Header().Get("Access-Control-Allow-Origin"); got != "" {
		t.Errorf("Access-Control-Allow-Origin = %q, want empty for non-whitelisted origin", got)
	}
	if got := w.Header().Get("Access-Control-Allow-Credentials"); got != "" {
		t.Errorf("Access-Control-Allow-Credentials = %q, want empty for non-whitelisted origin", got)
	}
}

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