package security import ( "net/http/httptest" "testing" "tyc-server/app/main/api/internal/config" ) // 创建测试配置 func createTestConfig() *config.SecurityConfig { return &config.SecurityConfig{ RateLimit: struct { Enabled bool `json:"enabled" yaml:"enabled"` WindowSize int64 `json:"windowSize" yaml:"windowSize"` MaxRequests int64 `json:"maxRequests" yaml:"maxRequests"` TriggerThreshold int64 `json:"triggerThreshold" yaml:"triggerThreshold"` TriggerWindow int64 `json:"triggerWindow" yaml:"triggerWindow"` }{ Enabled: true, WindowSize: 60, MaxRequests: 5, TriggerThreshold: 5, TriggerWindow: 24, }, IPBlacklist: struct { Enabled bool `json:"enabled" yaml:"enabled"` }{ Enabled: true, }, UserBlacklist: struct { Enabled bool `json:"enabled" yaml:"enabled"` }{ Enabled: true, }, AnomalyDetection: struct { Enabled bool `json:"enabled" yaml:"enabled"` }{ Enabled: true, }, BurstAttack: struct { Enabled bool `json:"enabled" yaml:"enabled"` TimeWindow int64 `json:"timeWindow" yaml:"timeWindow"` MaxConcurrent int64 `json:"maxConcurrent" yaml:"maxConcurrent"` }{ Enabled: true, TimeWindow: 1, MaxConcurrent: 20, }, } } // 测试客户端标识生成 func TestClientIDGeneration(t *testing.T) { config := createTestConfig() // 使用nil Redis进行测试,只测试不依赖Redis的逻辑 middleware := NewSecurityMiddleware(config, nil) // 测试IP标识 req := httptest.NewRequest("GET", "/test", nil) req.Header.Set("X-Real-IP", "192.168.1.100") clientID := middleware.getClientID(req) expected := "ip:192.168.1.100" if clientID != expected { t.Errorf("期望客户端标识 %s,但得到了 %s", expected, clientID) } } // 测试真实IP获取 func TestRealIPExtraction(t *testing.T) { config := createTestConfig() middleware := NewSecurityMiddleware(config, nil) // 测试X-Forwarded-For req := httptest.NewRequest("GET", "/test", nil) req.Header.Set("X-Forwarded-For", "203.0.113.1, 192.168.1.1") ip := middleware.getClientIP(req) expected := "203.0.113.1" if ip != expected { t.Errorf("期望IP %s,但得到了 %s", expected, ip) } // 测试X-Real-IP(创建新的请求对象) req2 := httptest.NewRequest("GET", "/test", nil) req2.Header.Set("X-Real-IP", "198.51.100.1") ip = middleware.getClientIP(req2) expected = "198.51.100.1" if ip != expected { t.Errorf("期望IP %s,但得到了 %s", expected, ip) } // 测试直接连接 req3 := httptest.NewRequest("GET", "/test", nil) req3.RemoteAddr = "192.168.1.50:12345" ip = middleware.getClientIP(req3) expected = "192.168.1.50" if ip != expected { t.Errorf("期望IP %s,但得到了 %s", expected, ip) } // 测试优先级:X-Forwarded-For 应该优先于 X-Real-IP req4 := httptest.NewRequest("GET", "/test", nil) req4.Header.Set("X-Forwarded-For", "10.0.0.1, 10.0.0.2") req4.Header.Set("X-Real-IP", "10.0.0.3") ip = middleware.getClientIP(req4) expected = "10.0.0.1" // X-Forwarded-For 应该优先 if ip != expected { t.Errorf("优先级测试失败:期望IP %s,但得到了 %s", expected, ip) } } // 测试中间件创建 func TestNewSecurityMiddleware(t *testing.T) { config := createTestConfig() middleware := NewSecurityMiddleware(config, nil) if middleware == nil { t.Error("中间件创建失败") } if middleware.config != config { t.Error("配置设置失败") } } // 测试配置验证 func TestConfigValidation(t *testing.T) { config := createTestConfig() if !config.RateLimit.Enabled { t.Error("频率限制应该启用") } if config.RateLimit.MaxRequests != 5 { t.Error("最大请求数设置错误") } if !config.IPBlacklist.Enabled { t.Error("IP黑名单应该启用") } if !config.UserBlacklist.Enabled { t.Error("用户黑名单应该启用") } if !config.AnomalyDetection.Enabled { t.Error("异常检测应该启用") } }