package middleware import ( "encoding/json" "net/http/httptest" "sync" "sync/atomic" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) // ============================================================================ // Group A: getFamilyFromPath 纯函数测试 // ============================================================================ func TestGetFamilyFromPath_Claude(t *testing.T) { assert.Equal(t, "claude", getFamilyFromPath("/v1/messages")) } func TestGetFamilyFromPath_ClaudeSubpath(t *testing.T) { assert.Equal(t, "claude", getFamilyFromPath("/v1/messages/")) } func TestGetFamilyFromPath_ClaudeBatchPath(t *testing.T) { assert.Equal(t, "claude", getFamilyFromPath("/v1/messages/batches")) } func TestGetFamilyFromPath_Gemini(t *testing.T) { assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/models/gemini-pro:generateContent")) assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/")) } func TestGetFamilyFromPath_GeminiCompatPath(t *testing.T) { assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/openai/models")) } func TestGetFamilyFromPath_GeminiStreamPath(t *testing.T) { assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/models/gemini-2.0-flash:streamGenerateContent")) } func TestGetFamilyFromPath_V1ModelsPath(t *testing.T) { // /v1/models/*path 用于 Gemini relay,应归类为 gemini 家族 assert.Equal(t, "gemini", getFamilyFromPath("/v1/models/gemini-pro:generateContent")) assert.Equal(t, "gemini", getFamilyFromPath("/v1/models/gemini-2.0-flash:streamGenerateContent")) } func TestGetFamilyFromPath_V1EnginesPath(t *testing.T) { // /v1/engines/:model/embeddings 用于 Gemini relay,应归类为 gemini 家族 assert.Equal(t, "gemini", getFamilyFromPath("/v1/engines/text-embedding-004/embeddings")) } func TestGetFamilyFromPath_OpenAI(t *testing.T) { assert.Equal(t, "openai", getFamilyFromPath("/v1/chat/completions")) assert.Equal(t, "openai", getFamilyFromPath("/v1/completions")) assert.Equal(t, "openai", getFamilyFromPath("/v1/embeddings")) assert.Equal(t, "openai", getFamilyFromPath("/v1/audio/speech")) assert.Equal(t, "openai", getFamilyFromPath("/v1/images/generations")) assert.Equal(t, "openai", getFamilyFromPath("/v1/rerank")) assert.Equal(t, "openai", getFamilyFromPath("/v1/moderations")) } func TestGetFamilyFromPath_Realtime(t *testing.T) { assert.Equal(t, "openai", getFamilyFromPath("/v1/realtime")) } func TestGetFamilyFromPath_Responses(t *testing.T) { assert.Equal(t, "openai", getFamilyFromPath("/v1/responses")) assert.Equal(t, "openai", getFamilyFromPath("/v1/responses/compact")) } func TestGetFamilyFromPath_Unknown(t *testing.T) { assert.Equal(t, "", getFamilyFromPath("/api/status")) assert.Equal(t, "", getFamilyFromPath("/pg/chat/completions")) assert.Equal(t, "", getFamilyFromPath("/")) assert.Equal(t, "", getFamilyFromPath("/v1")) } // ============================================================================ // Group B: isChannelAllowed 纯函数测试 // ============================================================================ func TestIsChannelAllowed_OpenAIChannel_OpenAIFormat(t *testing.T) { assert.True(t, isChannelAllowed(constant.ChannelTypeOpenAI, "openai")) } func TestIsChannelAllowed_OpenAIChannel_ClaudeFormat(t *testing.T) { assert.False(t, isChannelAllowed(constant.ChannelTypeOpenAI, "claude")) } func TestIsChannelAllowed_OpenAIChannel_GeminiFormat(t *testing.T) { assert.False(t, isChannelAllowed(constant.ChannelTypeOpenAI, "gemini")) } func TestIsChannelAllowed_ClaudeChannel_ClaudeFormat(t *testing.T) { assert.True(t, isChannelAllowed(constant.ChannelTypeAnthropic, "claude")) } func TestIsChannelAllowed_ClaudeChannel_OpenAIFormat(t *testing.T) { assert.False(t, isChannelAllowed(constant.ChannelTypeAnthropic, "openai")) } func TestIsChannelAllowed_GeminiChannel(t *testing.T) { assert.True(t, isChannelAllowed(constant.ChannelTypeGemini, "gemini")) assert.False(t, isChannelAllowed(constant.ChannelTypeGemini, "openai")) assert.False(t, isChannelAllowed(constant.ChannelTypeGemini, "claude")) } func TestIsChannelAllowed_AzureChannel(t *testing.T) { assert.True(t, isChannelAllowed(constant.ChannelTypeAzure, "openai")) assert.False(t, isChannelAllowed(constant.ChannelTypeAzure, "claude")) assert.False(t, isChannelAllowed(constant.ChannelTypeAzure, "gemini")) } func TestIsChannelAllowed_DeepSeekChannel(t *testing.T) { assert.True(t, isChannelAllowed(constant.ChannelTypeDeepSeek, "openai")) assert.False(t, isChannelAllowed(constant.ChannelTypeDeepSeek, "claude")) assert.False(t, isChannelAllowed(constant.ChannelTypeDeepSeek, "gemini")) } func TestIsChannelAllowed_VertexAI_AllFormats(t *testing.T) { assert.True(t, isChannelAllowed(constant.ChannelTypeVertexAi, "openai")) assert.True(t, isChannelAllowed(constant.ChannelTypeVertexAi, "claude")) assert.True(t, isChannelAllowed(constant.ChannelTypeVertexAi, "gemini")) } func TestIsChannelAllowed_Ali(t *testing.T) { assert.True(t, isChannelAllowed(constant.ChannelTypeAli, "openai")) assert.True(t, isChannelAllowed(constant.ChannelTypeAli, "claude")) assert.False(t, isChannelAllowed(constant.ChannelTypeAli, "gemini")) } func TestIsChannelAllowed_Midjourney_PassThrough(t *testing.T) { // Midjourney 不在 channelNativeFamilies 中,所有格式都应放行 assert.True(t, isChannelAllowed(constant.ChannelTypeMidjourney, "openai")) assert.True(t, isChannelAllowed(constant.ChannelTypeMidjourney, "claude")) assert.True(t, isChannelAllowed(constant.ChannelTypeMidjourney, "gemini")) } func TestIsChannelAllowed_Kling_PassThrough(t *testing.T) { // Kling 不在 channelNativeFamilies 中,所有格式都应放行 assert.True(t, isChannelAllowed(constant.ChannelTypeKling, "openai")) assert.True(t, isChannelAllowed(constant.ChannelTypeKling, "claude")) } func TestIsChannelAllowed_Suno_PassThrough(t *testing.T) { // Suno 不在 channelNativeFamilies 中 assert.True(t, isChannelAllowed(constant.ChannelTypeSunoAPI, "openai")) } func TestIsChannelAllowed_UnknownChannel(t *testing.T) { assert.True(t, isChannelAllowed(999, "openai")) assert.True(t, isChannelAllowed(999, "claude")) assert.True(t, isChannelAllowed(999, "gemini")) } func TestIsChannelAllowed_EmptyFamily(t *testing.T) { // 空 family 表示路径不匹配任何已知格式,guard 会直接放行 // isChannelAllowed 不会被空 family 调用,但验证行为:openai 在 families 中不含 "" assert.False(t, isChannelAllowed(constant.ChannelTypeOpenAI, "")) } // ============================================================================ // Helper: runGuardTest // ============================================================================ func runGuardTestWithMethod(method, path string, channelType int) (bool, *httptest.ResponseRecorder) { called := false gin.SetMode(gin.TestMode) r := gin.New() r.Use(func(c *gin.Context) { if channelType >= 0 { c.Set(string(constant.ContextKeyChannelType), channelType) } c.Next() }) r.Use(EndpointFormatGuard()) // 注册所有相关路由 r.POST("/v1/chat/completions", func(ctx *gin.Context) { called = true }) r.POST("/v1/completions", func(ctx *gin.Context) { called = true }) r.POST("/v1/messages", func(ctx *gin.Context) { called = true }) r.POST("/v1/embeddings", func(ctx *gin.Context) { called = true }) r.POST("/v1/rerank", func(ctx *gin.Context) { called = true }) r.POST("/v1/responses", func(ctx *gin.Context) { called = true }) r.POST("/v1/responses/compact", func(ctx *gin.Context) { called = true }) r.POST("/v1/moderations", func(ctx *gin.Context) { called = true }) r.POST("/v1/audio/speech", func(ctx *gin.Context) { called = true }) r.POST("/v1/images/generations", func(ctx *gin.Context) { called = true }) r.POST("/v1/models/*path", func(ctx *gin.Context) { called = true }) r.POST("/v1/engines/:model/embeddings", func(ctx *gin.Context) { called = true }) r.GET("/v1/realtime", func(ctx *gin.Context) { called = true }) r.Any("/v1beta/models/*path", func(ctx *gin.Context) { called = true }) w := httptest.NewRecorder() req := httptest.NewRequest(method, path, nil) r.ServeHTTP(w, req) return called, w } func runGuardTest(path string, channelType int) (bool, *httptest.ResponseRecorder) { return runGuardTestWithMethod("POST", path, channelType) } // ============================================================================ // Group C: 中间件集成测试 // ============================================================================ func TestEndpointFormatGuard_Disabled(t *testing.T) { orig := setTestRestrictionEnabled(false) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic) assert.True(t, called, "handler should be called when restriction is disabled") } func TestEndpointFormatGuard_Disabled_MultipleChannelTypes(t *testing.T) { orig := setTestRestrictionEnabled(false) defer setTestRestrictionEnabled(orig) for _, ch := range []int{constant.ChannelTypeOpenAI, constant.ChannelTypeAnthropic, constant.ChannelTypeGemini, constant.ChannelTypeMidjourney} { called, _ := runGuardTest("/v1/chat/completions", ch) assert.True(t, called, "all channel types should pass when restriction is disabled") } } func TestEndpointFormatGuard_Enabled_Matched(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeOpenAI) assert.True(t, called, "handler should be called when format matches channel type") } func TestEndpointFormatGuard_Enabled_Mismatched(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic) assert.False(t, called, "handler should NOT be called when format mismatches") assert.Equal(t, 400, w.Code) } func TestEndpointFormatGuard_Enabled_ClaudeEndpoint_ClaudeChannel(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1/messages", constant.ChannelTypeAnthropic) assert.True(t, called, "claude endpoint + claude channel should be allowed") } func TestEndpointFormatGuard_Enabled_ClaudeEndpoint_OpenAIChannel(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, w := runGuardTest("/v1/messages", constant.ChannelTypeOpenAI) assert.False(t, called, "claude endpoint + openai channel should be blocked") assert.Equal(t, 400, w.Code) } func TestEndpointFormatGuard_Enabled_GeminiEndpoint_GeminiChannel(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeGemini) assert.True(t, called, "gemini endpoint + gemini channel should be allowed") } func TestEndpointFormatGuard_Enabled_GeminiEndpoint_OpenAIChannel(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, w := runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeOpenAI) assert.False(t, called, "gemini endpoint + openai channel should be blocked") assert.Equal(t, 400, w.Code) } func TestEndpointFormatGuard_Enabled_GeminiChannel_OpenAIEndpoint(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeGemini) assert.False(t, called, "gemini channel should be blocked on openai endpoint") assert.Equal(t, 400, w.Code) } func TestEndpointFormatGuard_Enabled_GeminiChannel_V1ModelsPath(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1/models/gemini-pro:generateContent", constant.ChannelTypeGemini) assert.True(t, called, "gemini channel should be allowed on /v1/models/* path (gemini family)") } func TestEndpointFormatGuard_Enabled_DeepSeekAllowedOnOpenAI(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeDeepSeek) assert.True(t, called, "deepseek channel should be allowed on openai endpoint") } func TestEndpointFormatGuard_Enabled_DeepSeekBlockedOnClaude(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, w := runGuardTest("/v1/messages", constant.ChannelTypeDeepSeek) assert.False(t, called, "deepseek channel should be blocked on claude endpoint") assert.Equal(t, 400, w.Code) } func TestEndpointFormatGuard_Enabled_MidjourneyPassThrough(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeMidjourney) assert.True(t, called, "midjourney channel should bypass restriction (not LLM)") } func TestEndpointFormatGuard_Enabled_KlingPassThrough(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeKling) assert.True(t, called, "kling channel should bypass restriction (not LLM)") } func TestEndpointFormatGuard_Enabled_RealtimeWebSocket(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTestWithMethod("GET", "/v1/realtime", constant.ChannelTypeOpenAI) assert.True(t, called, "openai channel should be allowed on realtime websocket") } func TestEndpointFormatGuard_Enabled_RealtimeWebSocket_Blocked(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, w := runGuardTestWithMethod("GET", "/v1/realtime", constant.ChannelTypeAnthropic) assert.False(t, called, "anthropic channel should be blocked on realtime websocket") assert.Equal(t, 400, w.Code) } func TestEndpointFormatGuard_Enabled_ResponsesEndpoint(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1/responses", constant.ChannelTypeOpenAI) assert.True(t, called, "openai channel should be allowed on /v1/responses") called, w := runGuardTest("/v1/responses", constant.ChannelTypeAnthropic) assert.False(t, called, "anthropic channel should be blocked on /v1/responses") assert.Equal(t, 400, w.Code) } func TestEndpointFormatGuard_Enabled_AWS_ClaudeAllowed(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1/messages", constant.ChannelTypeAws) assert.True(t, called, "AWS channel should be allowed for claude endpoint") } func TestEndpointFormatGuard_Enabled_AWS_OpenAIBlocked(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAws) assert.False(t, called, "AWS channel should be blocked for openai endpoint") assert.Equal(t, 400, w.Code) } func TestEndpointFormatGuard_Enabled_VertexAI_AllEndpoints(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeVertexAi) assert.True(t, called, "vertexai should be allowed on openai endpoint") called, _ = runGuardTest("/v1/messages", constant.ChannelTypeVertexAi) assert.True(t, called, "vertexai should be allowed on claude endpoint") called, _ = runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeVertexAi) assert.True(t, called, "vertexai should be allowed on gemini endpoint") } func TestEndpointFormatGuard_Enabled_Ali_GeminiBlocked(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeAli) assert.True(t, called, "ali should be allowed on openai endpoint") called, _ = runGuardTest("/v1/messages", constant.ChannelTypeAli) assert.True(t, called, "ali should be allowed on claude endpoint") called, w := runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeAli) assert.False(t, called, "ali should be blocked on gemini endpoint") assert.Equal(t, 400, w.Code) } func TestEndpointFormatGuard_Enabled_NoChannelType(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1/chat/completions", -1) assert.True(t, called, "handler should be called when no channelType in context") } func TestEndpointFormatGuard_Enabled_ChannelTypeZero(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) called, _ := runGuardTest("/v1/chat/completions", 0) assert.True(t, called, "handler should be called when channelType is 0") } func TestEndpointFormatGuard_NonV1Path(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) // /api/status 不在 guard 路由组中,无法通过 runGuardTest 测试 // 改用直接调用 getFamilyFromPath 验证返回空 family := getFamilyFromPath("/api/status") assert.Equal(t, "", family, "non-v1 paths should have empty family") } // ============================================================================ // Group D: 错误消息验证 // ============================================================================ func TestEndpointFormatGuard_ErrorMessage_ClaudeChannel(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) _, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic) require.Equal(t, 400, w.Code) var body map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body)) errObj := body["error"].(map[string]interface{}) assert.Equal(t, "endpoint_format_restricted", errObj["code"]) assert.Equal(t, "new_api_error", errObj["type"]) msg := errObj["message"].(string) assert.Contains(t, msg, "Anthropic") assert.Contains(t, msg, "/v1/messages") assert.Contains(t, msg, "openai endpoints") } func TestEndpointFormatGuard_ErrorMessage_OpenAIChannel(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) _, w := runGuardTest("/v1/messages", constant.ChannelTypeOpenAI) require.Equal(t, 400, w.Code) var body map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body)) errObj := body["error"].(map[string]interface{}) assert.Equal(t, "endpoint_format_restricted", errObj["code"]) msg := errObj["message"].(string) assert.Contains(t, msg, "OpenAI") assert.Contains(t, msg, "/v1/chat/completions") assert.Contains(t, msg, "claude endpoints") } func TestEndpointFormatGuard_ErrorMessage_GeminiChannel(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) _, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeGemini) require.Equal(t, 400, w.Code) var body map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body)) errObj := body["error"].(map[string]interface{}) assert.Equal(t, "endpoint_format_restricted", errObj["code"]) msg := errObj["message"].(string) assert.Contains(t, msg, "Gemini") assert.Contains(t, msg, "/v1beta/models/{model}:generateContent") assert.Contains(t, msg, "openai endpoints") } func TestEndpointFormatGuard_ErrorMessage_DeepSeekChannel(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) _, w := runGuardTest("/v1/messages", constant.ChannelTypeDeepSeek) require.Equal(t, 400, w.Code) var body map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body)) errObj := body["error"].(map[string]interface{}) assert.Equal(t, "endpoint_format_restricted", errObj["code"]) msg := errObj["message"].(string) assert.Contains(t, msg, "DeepSeek") assert.Contains(t, msg, "/v1/chat/completions") } // ============================================================================ // Group E: 并发测试 // ============================================================================ func TestEndpointFormatGuard_ConcurrentRequests(t *testing.T) { orig := setTestRestrictionEnabled(true) defer setTestRestrictionEnabled(orig) var allowedCount, blockedCount int64 var wg sync.WaitGroup for i := 0; i < 20; i++ { wg.Add(1) go func(idx int) { defer wg.Done() var ch int if idx%2 == 0 { ch = constant.ChannelTypeOpenAI } else { ch = constant.ChannelTypeAnthropic } called, w := runGuardTest("/v1/chat/completions", ch) if idx%2 == 0 { if called && w.Code == 200 { atomic.AddInt64(&allowedCount, 1) } } else { if !called && w.Code == 400 { atomic.AddInt64(&blockedCount, 1) } } }(i) } wg.Wait() assert.Equal(t, int64(10), atomic.LoadInt64(&allowedCount), "10 个 OpenAI 请求应放行") assert.Equal(t, int64(10), atomic.LoadInt64(&blockedCount), "10 个 Anthropic 请求应被拦截") } // ============================================================================ // Group F: 功能开关测试 // ============================================================================ func TestEndpointFormatGuard_RuntimeToggle(t *testing.T) { orig := setTestRestrictionEnabled(false) defer setTestRestrictionEnabled(orig) // 阶段 1: 关闭 → 应放行 called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic) assert.True(t, called, "should pass when restriction is OFF") // 阶段 2: 开启 → 应拦截 setTestRestrictionEnabled(true) called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic) assert.False(t, called, "should be blocked when restriction is ON") assert.Equal(t, 400, w.Code) // 阶段 3: 再次关闭 → 应放行 setTestRestrictionEnabled(false) called, _ = runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic) assert.True(t, called, "should pass again when restriction is turned OFF") } // ============================================================================ // 辅助测试 // ============================================================================ func TestEndpointFormatGuard_ContextKeyInt(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil) c.Set(string(constant.ContextKeyChannelType), 14) val := common.GetContextKeyInt(c, constant.ContextKeyChannelType) assert.Equal(t, 14, val) }