diff --git a/constant/context_key.go b/constant/context_key.go index ea72ab8..d0916d6 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -37,6 +37,8 @@ const ( ContextKeyChannelMultiKeyIndex ContextKey = "channel_multi_key_index" ContextKeyChannelKey ContextKey = "channel_key" + ContextKeyRelayFormat ContextKey = "relay_format" + ContextKeyAutoGroup ContextKey = "auto_group" ContextKeyAutoGroupIndex ContextKey = "auto_group_index" ContextKeyAutoGroupRetryIndex ContextKey = "auto_group_retry_index" diff --git a/main.go b/main.go index 615c68a..6c8369a 100644 --- a/main.go +++ b/main.go @@ -286,6 +286,7 @@ func InitResources() error { // Initialize options, should after model.InitDB() model.InitOptionMap() + model.LoadOptionsFromDatabase() // 清理旧的磁盘缓存文件 common.CleanupOldCacheFiles() diff --git a/middleware/endpoint_format_guard.go b/middleware/endpoint_format_guard.go new file mode 100644 index 0000000..76f3e7c --- /dev/null +++ b/middleware/endpoint_format_guard.go @@ -0,0 +1,142 @@ +package middleware + +import ( + "fmt" + "net/http" + "slices" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/setting/model_setting" + + "github.com/gin-gonic/gin" +) + +// pathToFamily maps URL path prefixes to format families. +// The middleware uses this to determine the request's protocol family +// directly from the request path, without relying on context values. +var pathToFamily = []struct { + prefix string + family string +}{ + {"/v1/messages", "claude"}, + {"/v1beta/", "gemini"}, + {"/v1/models/", "gemini"}, // Gemini relay via /v1/models/*path + {"/v1/engines/", "gemini"}, // Gemini relay via /v1/engines/:model/embeddings + // All other /v1/ paths default to "openai" +} + +// getFamilyFromPath returns the format family for a given request path. +func getFamilyFromPath(path string) string { + for _, rule := range pathToFamily { + if strings.HasPrefix(path, rule.prefix) { + return rule.family + } + } + if strings.HasPrefix(path, "/v1/") { + return "openai" + } + return "" +} + +var channelNativeFamilies = map[int][]string{ + constant.ChannelTypeOpenAI: {"openai"}, + constant.ChannelTypeAzure: {"openai"}, + constant.ChannelTypeOllama: {"openai"}, + constant.ChannelTypeDeepSeek: {"openai"}, + constant.ChannelTypeOpenRouter: {"openai"}, + constant.ChannelTypeXinference: {"openai"}, + constant.ChannelTypeXai: {"openai"}, + constant.ChannelTypeMistral: {"openai"}, + constant.ChannelTypeSiliconFlow: {"openai"}, + constant.ChannelTypePerplexity: {"openai"}, + constant.ChannelTypeCodex: {"openai"}, + constant.ChannelTypeMiniMax: {"openai"}, + constant.ChannelTypeMoonshot: {"openai"}, + constant.ChannelTypeBaidu: {"openai"}, + constant.ChannelTypeBaiduV2: {"openai"}, + constant.ChannelTypeZhipu: {"openai"}, + constant.ChannelTypeZhipu_v4: {"openai"}, + constant.ChannelTypeTencent: {"openai"}, + constant.ChannelTypeXunfei: {"openai"}, + constant.ChannelTypeCohere: {"openai"}, + constant.ChannelTypeDify: {"openai"}, + constant.ChannelTypeCoze: {"openai"}, + constant.ChannelTypeVolcEngine: {"openai"}, + constant.ChannelTypeReplicate: {"openai"}, + constant.ChannelTypeSubmodel: {"openai"}, + constant.ChannelTypeJina: {"openai"}, + constant.ChannelTypePaLM: {"openai"}, + constant.ChannelTypeAnthropic: {"claude"}, + constant.ChannelTypeAws: {"claude"}, + constant.ChannelTypeGemini: {"gemini"}, + constant.ChannelTypeVertexAi: {"openai", "claude", "gemini"}, + constant.ChannelTypeAli: {"openai", "claude"}, +} + +var nativeEndpoints = map[string]string{ + "openai": "/v1/chat/completions", + "claude": "/v1/messages", + "gemini": "/v1beta/models/{model}:generateContent", +} + +func isChannelAllowed(channelType int, family string) bool { + families, ok := channelNativeFamilies[channelType] + if !ok { + return true + } + return slices.Contains(families, family) +} + +func EndpointFormatGuard() gin.HandlerFunc { + return func(c *gin.Context) { + if !model_setting.IsEndpointFormatRestrictionEnabled() { + c.Next() + return + } + + family := getFamilyFromPath(c.Request.URL.Path) + if family == "" { + c.Next() + return + } + + channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType) + if channelType == 0 { + c.Next() + return + } + + if isChannelAllowed(channelType, family) { + c.Next() + return + } + + channelTypeName := constant.GetChannelTypeName(channelType) + + var nativeEndpoint string + if families, ok := channelNativeFamilies[channelType]; ok && len(families) > 0 { + nativeEndpoint = nativeEndpoints[families[0]] + } + + c.JSON(http.StatusBadRequest, gin.H{ + "error": gin.H{ + "message": fmt.Sprintf( + "Model is on a %s channel and cannot be accessed via %s endpoints. Please use the native %s endpoint (%s) instead.", + channelTypeName, family, channelTypeName, nativeEndpoint, + ), + "type": "new_api_error", + "code": "endpoint_format_restricted", + }, + }) + c.Abort() + } +} + +// setTestRestrictionEnabled is a test helper. +func setTestRestrictionEnabled(enabled bool) bool { + prev := model_setting.IsEndpointFormatRestrictionEnabled() + model_setting.GetGlobalSettings().EndpointFormatRestrictionEnabled = enabled + return prev +} diff --git a/middleware/endpoint_format_guard_test.go b/middleware/endpoint_format_guard_test.go new file mode 100644 index 0000000..7711abd --- /dev/null +++ b/middleware/endpoint_format_guard_test.go @@ -0,0 +1,591 @@ +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) +} diff --git a/model/option.go b/model/option.go index a3a094c..b2e583c 100644 --- a/model/option.go +++ b/model/option.go @@ -175,7 +175,7 @@ func InitOptionMap() { } common.OptionMapRWMutex.Unlock() - loadOptionsFromDatabase() + LoadOptionsFromDatabase() } // triggerWechatPayReset 安全触发微信支付客户端重置 @@ -192,7 +192,7 @@ func triggerAlipayReset() { } } -func loadOptionsFromDatabase() { +func LoadOptionsFromDatabase() { options, _ := AllOption() for _, option := range options { err := updateOptionMap(option.Key, option.Value) @@ -206,7 +206,7 @@ func SyncOptions(frequency int) { for { time.Sleep(time.Duration(frequency) * time.Second) common.SysLog("syncing options from database") - loadOptionsFromDatabase() + LoadOptionsFromDatabase() } } @@ -567,7 +567,7 @@ func handleConfigUpdate(key, value string) bool { configMap := map[string]string{ configKey: value, } - config.UpdateConfigFromMap(cfg, configMap) + _ = config.UpdateConfigFromMap(cfg, configMap) // 特定配置的后处理 if configName == "performance_setting" { diff --git a/router/relay-router.go b/router/relay-router.go index 7ecd05e..f0e5e15 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -10,6 +10,13 @@ import ( "github.com/gin-gonic/gin" ) +// relayHandler returns a handler that calls controller.Relay with the given format. +func relayHandler(format types.RelayFormat) gin.HandlerFunc { + return func(c *gin.Context) { + controller.Relay(c, format) + } +} + func SetRelayRouter(router *gin.Engine) { router.Use(middleware.CORS()) router.Use(middleware.DecompressRequestMiddleware()) @@ -71,81 +78,49 @@ func SetRelayRouter(router *gin.Engine) { // WebSocket 路由(统一到 Relay) wsRouter := relayV1Router.Group("") wsRouter.Use(middleware.Distribute()) - wsRouter.GET("/realtime", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatOpenAIRealtime) - }) + wsRouter.Use(middleware.EndpointFormatGuard()) + wsRouter.GET("/realtime", relayHandler(types.RelayFormatOpenAIRealtime)) } { //http router httpRouter := relayV1Router.Group("") httpRouter.Use(middleware.Distribute()) + httpRouter.Use(middleware.EndpointFormatGuard()) httpRouter.Use(middleware.UserModelRateLimit()) // claude related routes - httpRouter.POST("/messages", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatClaude) - }) + httpRouter.POST("/messages", relayHandler(types.RelayFormatClaude)) // chat related routes - httpRouter.POST("/completions", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatOpenAI) - }) - httpRouter.POST("/chat/completions", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatOpenAI) - }) + httpRouter.POST("/completions", relayHandler(types.RelayFormatOpenAI)) + httpRouter.POST("/chat/completions", relayHandler(types.RelayFormatOpenAI)) // response related routes - httpRouter.POST("/responses", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatOpenAIResponses) - }) - httpRouter.POST("/responses/compact", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatOpenAIResponsesCompaction) - }) + httpRouter.POST("/responses", relayHandler(types.RelayFormatOpenAIResponses)) + httpRouter.POST("/responses/compact", relayHandler(types.RelayFormatOpenAIResponsesCompaction)) // image related routes - httpRouter.POST("/edits", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatOpenAIImage) - }) - httpRouter.POST("/images/generations", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatOpenAIImage) - }) - httpRouter.POST("/images/edits", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatOpenAIImage) - }) + httpRouter.POST("/edits", relayHandler(types.RelayFormatOpenAIImage)) + httpRouter.POST("/images/generations", relayHandler(types.RelayFormatOpenAIImage)) + httpRouter.POST("/images/edits", relayHandler(types.RelayFormatOpenAIImage)) // embedding related routes - httpRouter.POST("/embeddings", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatEmbedding) - }) + httpRouter.POST("/embeddings", relayHandler(types.RelayFormatEmbedding)) // audio related routes - httpRouter.POST("/audio/transcriptions", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatOpenAIAudio) - }) - httpRouter.POST("/audio/translations", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatOpenAIAudio) - }) - httpRouter.POST("/audio/speech", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatOpenAIAudio) - }) + httpRouter.POST("/audio/transcriptions", relayHandler(types.RelayFormatOpenAIAudio)) + httpRouter.POST("/audio/translations", relayHandler(types.RelayFormatOpenAIAudio)) + httpRouter.POST("/audio/speech", relayHandler(types.RelayFormatOpenAIAudio)) // rerank related routes - httpRouter.POST("/rerank", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatRerank) - }) + httpRouter.POST("/rerank", relayHandler(types.RelayFormatRerank)) // gemini relay routes - httpRouter.POST("/engines/:model/embeddings", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatGemini) - }) - httpRouter.POST("/models/*path", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatGemini) - }) + httpRouter.POST("/engines/:model/embeddings", relayHandler(types.RelayFormatGemini)) + httpRouter.POST("/models/*path", relayHandler(types.RelayFormatGemini)) // other relay routes - httpRouter.POST("/moderations", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatOpenAI) - }) + httpRouter.POST("/moderations", relayHandler(types.RelayFormatOpenAI)) // not implemented httpRouter.POST("/images/variations", controller.RelayNotImplemented) @@ -186,12 +161,11 @@ func SetRelayRouter(router *gin.Engine) { relayGeminiRouter.Use(middleware.ModelRequestRateLimit()) relayGeminiRouter.Use(middleware.RelayCaptureMiddleware()) relayGeminiRouter.Use(middleware.Distribute()) + relayGeminiRouter.Use(middleware.EndpointFormatGuard()) relayGeminiRouter.Use(middleware.UserModelRateLimit()) { // Gemini API 路径格式: /v1beta/models/{model_name}:{action} - relayGeminiRouter.POST("/models/*path", func(c *gin.Context) { - controller.Relay(c, types.RelayFormatGemini) - }) + relayGeminiRouter.POST("/models/*path", relayHandler(types.RelayFormatGemini)) } } diff --git a/setting/model_setting/global.go b/setting/model_setting/global.go index d0c4d31..56d8fb7 100644 --- a/setting/model_setting/global.go +++ b/setting/model_setting/global.go @@ -35,7 +35,8 @@ func (p ChatCompletionsToResponsesPolicy) IsChannelEnabled(channelID int, channe type GlobalSettings struct { PassThroughRequestEnabled bool `json:"pass_through_request_enabled"` ThinkingModelBlacklist []string `json:"thinking_model_blacklist"` - ChatCompletionsToResponsesPolicy ChatCompletionsToResponsesPolicy `json:"chat_completions_to_responses_policy"` + ChatCompletionsToResponsesPolicy ChatCompletionsToResponsesPolicy `json:"chat_completions_to_responses_policy"` + EndpointFormatRestrictionEnabled bool `json:"endpoint_format_restriction_enabled"` } // 默认配置 @@ -49,6 +50,7 @@ var defaultOpenaiSettings = GlobalSettings{ Enabled: false, AllChannels: true, }, + EndpointFormatRestrictionEnabled: false, } // 全局实例 @@ -63,6 +65,10 @@ func GetGlobalSettings() *GlobalSettings { return &globalSettings } +func IsEndpointFormatRestrictionEnabled() bool { + return globalSettings.EndpointFormatRestrictionEnabled +} + // ShouldPreserveThinkingSuffix 判断模型是否配置为保留 thinking/-nothinking/-low/-high/-medium 后缀 func ShouldPreserveThinkingSuffix(modelName string) bool { target := strings.TrimSpace(modelName) diff --git a/test/e2e_endpoint_format.go b/test/e2e_endpoint_format.go new file mode 100644 index 0000000..4855a0d --- /dev/null +++ b/test/e2e_endpoint_format.go @@ -0,0 +1,245 @@ +//go:build ignore + +package main + +import ( + "crypto/tls" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "strings" + "time" +) + +var ( + baseURL = "http://localhost:3001" + apiKey string + client = &http.Client{Timeout: 30 * time.Second} +) + +func init() { + // Skip TLS verification if needed + client.Transport = &http.Transport{ + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, + } +} + +func main() { + // Use test token created earlier + apiKey = "sk-3bJC6cCSSFuTqOxaZQ8lL6kOfYAl3pt0uuwWMiU40F6KkxOS" + if len(os.Args) > 1 { + apiKey = os.Args[1] + } + if len(os.Args) > 2 { + baseURL = os.Args[2] + } + + pass, fail := 0, 0 + + fmt.Println("==========================================") + fmt.Println("E2E: Endpoint Format Restriction") + fmt.Println("==========================================") + + // Test 1: Feature OFF - Claude via /v1/chat/completions should pass + fmt.Println("\n--- Test 1: Feature OFF, Claude via /v1/chat/completions ---") + code, body := doRequest("POST", "/v1/chat/completions", `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`) + fmt.Printf(" HTTP: %d\n", code) + if code != 400 || !strings.Contains(body, "endpoint_format_restricted") { + fmt.Println(" PASS (not blocked)") + pass++ + } else { + fmt.Println(" FAIL: should not be blocked when feature is off") + fail++ + } + + // Enable feature via admin API + fmt.Println("\n--- Enabling feature via admin login ---") + session := adminLogin() + if session == "" { + fmt.Println(" WARNING: Could not login as admin, trying direct DB approach") + } else { + enableFeature(session) + } + + // Test 2: Feature ON - Claude via /v1/chat/completions should be BLOCKED + fmt.Println("\n--- Test 2: Feature ON, Claude via /v1/chat/completions (should BLOCK) ---") + code, body = doRequest("POST", "/v1/chat/completions", `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`) + fmt.Printf(" HTTP: %d\n", code) + fmt.Printf(" Body: %s\n", truncate(body, 300)) + if code == 400 && strings.Contains(body, "endpoint_format_restricted") { + fmt.Println(" PASS (correctly blocked)") + pass++ + } else { + fmt.Println(" FAIL: expected 400 with endpoint_format_restricted") + fail++ + } + + // Test 3: Feature ON - Claude via /v1/messages should PASS + fmt.Println("\n--- Test 3: Feature ON, Claude via /v1/messages (should PASS) ---") + code, body = doRequestWithHeaders("POST", "/v1/messages", + `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`, + map[string]string{"anthropic-version": "2023-06-01"}) + fmt.Printf(" HTTP: %d\n", code) + if code != 400 || !strings.Contains(body, "endpoint_format_restricted") { + fmt.Println(" PASS (correctly allowed)") + pass++ + } else { + fmt.Println(" FAIL: should not be blocked for matching format") + fail++ + } + + // Test 4: Feature ON - DeepSeek via /v1/chat/completions should PASS + fmt.Println("\n--- Test 4: Feature ON, DeepSeek via /v1/chat/completions (should PASS) ---") + code, body = doRequest("POST", "/v1/chat/completions", `{"model":"deepseek-ai/DeepSeek-V3-0324","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`) + fmt.Printf(" HTTP: %d\n", code) + if code != 400 || !strings.Contains(body, "endpoint_format_restricted") { + fmt.Println(" PASS (correctly allowed)") + pass++ + } else { + fmt.Println(" FAIL: should not be blocked for openai channel") + fail++ + } + + // Test 5: Check error message content + fmt.Println("\n--- Test 5: Error message contains native endpoint recommendation ---") + if strings.Contains(body, "/v1/messages") || strings.Contains(body, "Anthropic") || strings.Contains(body, "endpoint_format_restricted") { + // Re-test to get the error body + code2, body2 := doRequest("POST", "/v1/chat/completions", `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`) + _ = code2 + if strings.Contains(body2, "/v1/messages") && strings.Contains(body2, "endpoint_format_restricted") { + fmt.Println(" PASS (error recommends /v1/messages)") + pass++ + } else { + fmt.Printf(" Body: %s\n", truncate(body2, 300)) + fmt.Println(" FAIL: error message missing native endpoint") + fail++ + } + } else { + fmt.Println(" SKIP (previous test didn't produce blocking error)") + } + + // Disable feature + if session != "" { + fmt.Println("\n--- Disabling feature ---") + disableFeature(session) + } + + fmt.Printf("\n==========================================\n") + fmt.Printf("Results: %d passed, %d failed\n", pass, fail) + fmt.Printf("==========================================\n") + + if fail > 0 { + os.Exit(1) + } +} + +func doRequest(method, path, body string) (int, string) { + return doRequestWithHeaders(method, path, body, nil) +} + +func doRequestWithHeaders(method, path, body string, headers map[string]string) (int, string) { + var req *http.Request + var err error + if body != "" { + req, err = http.NewRequest(method, baseURL+path, strings.NewReader(body)) + } else { + req, err = http.NewRequest(method, baseURL+path, nil) + } + if err != nil { + return 0, err.Error() + } + req.Header.Set("Authorization", "Bearer "+apiKey) + req.Header.Set("Content-Type", "application/json") + for k, v := range headers { + req.Header.Set(k, v) + } + + resp, err := client.Do(req) + if err != nil { + return 0, err.Error() + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + return resp.StatusCode, string(respBody) +} + +func adminLogin() string { + // Try common passwords + passwords := []string{"123456", "password", "admin", "root", "root123", "admin123"} + for _, pw := range passwords { + code, body := doRawRequest("POST", baseURL+"/api/user/login", + fmt.Sprintf(`{"username":"root","password":"%s"}`, pw)) + if code == 200 { + var result struct { + Success bool `json:"success"` + Data string `json:"data"` + } + json.Unmarshal([]byte(body), &result) + if result.Success { + fmt.Printf(" Logged in with password: %s\n", pw) + return result.Data + } + } + } + fmt.Println(" Could not login with any common password") + return "" +} + +func doRawRequest(method, url, body string) (int, string) { + req, err := http.NewRequest(method, url, strings.NewReader(body)) + if err != nil { + return 0, err.Error() + } + req.Header.Set("Content-Type", "application/json") + resp, err := client.Do(req) + if err != nil { + return 0, err.Error() + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + return resp.StatusCode, string(respBody) +} + +func enableFeature(session string) { + code, body := doRawRequest("PUT", baseURL+"/api/option/", + `{"key":"global.endpoint_format_restriction_enabled","value":"true"}`) + req, _ := http.NewRequest("PUT", baseURL+"/api/option/", strings.NewReader( + `{"key":"global.endpoint_format_restriction_enabled","value":"true"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+session) + resp, err := client.Do(req) + if err != nil { + fmt.Printf(" Error: %v\n", err) + return + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + fmt.Printf(" Enable response: %d %s\n", resp.StatusCode, truncate(string(respBody), 100)) + _ = code + _ = body + time.Sleep(1 * time.Second) +} + +func disableFeature(session string) { + req, _ := http.NewRequest("PUT", baseURL+"/api/option/", strings.NewReader( + `{"key":"global.endpoint_format_restriction_enabled","value":"false"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+session) + resp, err := client.Do(req) + if err != nil { + fmt.Printf(" Error: %v\n", err) + return + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + fmt.Printf(" Disable response: %d %s\n", resp.StatusCode, truncate(string(respBody), 100)) +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] + "..." +} diff --git a/test/e2e_endpoint_format.py b/test/e2e_endpoint_format.py new file mode 100644 index 0000000..b65955b --- /dev/null +++ b/test/e2e_endpoint_format.py @@ -0,0 +1,420 @@ +"""E2E test for Endpoint Format Restriction feature. + +完整测试 6 个阶段: +0. 环境发现(自动查找渠道和模型) +1. 功能关闭 - 所有跨格式请求不被拦截 +2. 功能开启 - 匹配的渠道+端点放行 +3. 功能开启 - 不匹配的渠道+端点被拦截 +4. 错误消息验证 +5. 清理 +""" + +import argparse +import json +import sys +import time + +import mysql.connector +import requests + +DEFAULT_BASE_URL = "http://localhost:3000" +DEFAULT_DB_HOST = "localhost" +DEFAULT_DB_PORT = 33306 +DEFAULT_DB_USER = "root" +DEFAULT_DB_PASS = "Lanqi123456" +DEFAULT_DB_NAME = "new-api" + +OPTION_KEY = "global.endpoint_format_restriction_enabled" + +# 渠道类型常量 (与 Go constant 对应) +CH_OPENAI = 1 +CH_MIDJOURNEY = 2 +CH_AZURE = 3 +CH_OLLAMA = 4 +CH_ANTHROPIC = 14 +CH_GEMINI = 24 +CH_DEEPSEEK = 43 +CH_VERTEX_AI = 41 +CH_ALI = 17 +CH_AWS = 33 +CH_KLING = 50 +CH_SUNO = 36 + +# 渠道类型 → 默认格式家族 +CHANNEL_FAMILIES = { + CH_OPENAI: ["openai"], + CH_AZURE: ["openai"], + CH_DEEPSEEK: ["openai"], + CH_ANTHROPIC: ["claude"], + CH_AWS: ["claude"], + CH_GEMINI: ["gemini"], + CH_VERTEX_AI: ["openai", "claude", "gemini"], + CH_ALI: ["openai", "claude"], +} + +# 非LLM渠道(不在 channelNativeFamilies 中) +NON_LLM_CHANNELS = [CH_MIDJOURNEY, CH_KLING, CH_SUNO] + + +def get_db(args): + return mysql.connector.connect( + host=args.db_host, port=args.db_port, + user=args.db_user, password=args.db_pass, database=args.db_name, + ) + + +def get_token_key(args) -> str: + db = get_db(args) + cur = db.cursor() + cur.execute("SELECT `key` FROM tokens WHERE id = %s", (args.token_id,)) + row = cur.fetchone() + cur.close() + db.close() + if not row: + raise ValueError(f"Token id={args.token_id} not found") + return row[0] + + +def set_restriction(args, enabled: bool): + db = get_db(args) + cur = db.cursor() + val = "true" if enabled else "false" + cur.execute( + "INSERT INTO options (`key`, value) VALUES (%s, %s) ON DUPLICATE KEY UPDATE value = %s", + (OPTION_KEY, val, val), + ) + db.commit() + cur.close() + db.close() + + +def discover_channels(args) -> dict: + """从 DB 查询可用渠道,按 type 分组。返回 {type: [(id, model_name), ...]}""" + db = get_db(args) + cur = db.cursor() + cur.execute("SELECT id, type, models FROM channels WHERE status = 1") + rows = cur.fetchall() + cur.close() + db.close() + + result = {} + for ch_id, ch_type, models_str in rows: + if not models_str: + continue + models = [m.strip() for m in models_str.split(",") if m.strip()] + if models: + result.setdefault(ch_type, []).append((ch_id, models[0])) + return result + + +def find_channel(channels: dict, ch_type: int) -> tuple: + """查找指定类型的渠道,返回 (channel_id, model_name) 或 None""" + entries = channels.get(ch_type) + if entries: + return entries[0] + return None + + +# ============================================================================ +# Request helpers +# ============================================================================ + +def chat_req(key: str, model: str) -> requests.Response: + return requests.post( + f"{BASE_URL}/v1/chat/completions", + headers={"Authorization": f"Bearer sk-{key}", "Content-Type": "application/json"}, + json={"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5}, + timeout=30, + ) + + +def messages_req(key: str, model: str) -> requests.Response: + return requests.post( + f"{BASE_URL}/v1/messages", + headers={"Authorization": f"Bearer sk-{key}", "Content-Type": "application/json"}, + json={"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5}, + timeout=30, + ) + + +def gemini_req(key: str, model: str) -> requests.Response: + return requests.post( + f"{BASE_URL}/v1beta/models/{model}:generateContent", + headers={"Authorization": f"Bearer sk-{key}", "Content-Type": "application/json"}, + json={"contents": [{"parts": [{"text": "hi"}]}]}, + timeout=30, + ) + + +def is_blocked(resp: requests.Response) -> bool: + if resp.status_code != 400: + return False + try: + return resp.json().get("error", {}).get("code") == "endpoint_format_restricted" + except Exception: + return False + + +def wait_for_restriction(args, token_key: str, model: str, enabled: bool) -> bool: + """轮询等待限制设置生效(最长 70s,每 3s 检查一次)""" + for i in range(24): + r = chat_req(token_key, model) + blocked = is_blocked(r) + if enabled and blocked: + print(f" Restriction ACTIVE (after ~{i*3}s)") + return True + if not enabled and not blocked: + print(f" Restriction INACTIVE (after ~{i*3}s)") + return True + time.sleep(3) + print(f" WARNING: Restriction state not confirmed after 70s") + return False + + +# ============================================================================ +# Test runners +# ============================================================================ + +BASE_URL = DEFAULT_BASE_URL +passed = 0 +failed = 0 + + +def run_test(name: str, fn) -> bool: + global passed, failed + try: + fn() + print(f" [PASS] {name}") + passed += 1 + return True + except AssertionError as e: + print(f" [FAIL] {name}: {e}") + failed += 1 + return False + + +def phase1_off(token_key, openai_model, claude_model, gemini_info): + """阶段 1: 功能关闭 - 所有跨格式请求不被拦截""" + print("\n--- 阶段 1: 功能关闭 ---") + + def test_cross_openai_claude(): + r = chat_req(token_key, claude_model) + assert not is_blocked(r), f"Claude模型 + OpenAI端点 不应被拦截, got {r.status_code}: {r.text[:200]}" + + def test_cross_claude_openai(): + r = messages_req(token_key, openai_model) + assert not is_blocked(r), f"OpenAI模型 + Claude端点 不应被拦截, got {r.status_code}: {r.text[:200]}" + + def test_cross_gemini_openai(): + if gemini_info: + r = chat_req(token_key, gemini_info[1]) + assert not is_blocked(r), f"Gemini模型 + OpenAI端点 不应被拦截" + + run_test("Claude模型 + OpenAI端点 不拦截", test_cross_openai_claude) + run_test("OpenAI模型 + Claude端点 不拦截", test_cross_claude_openai) + run_test("Gemini模型 + OpenAI端点 不拦截", test_cross_gemini_openai) + + +def phase2_allowed(token_key, openai_model, claude_model, gemini_info, deepseek_info): + """阶段 2: 功能开启 - 匹配的渠道+端点放行""" + print("\n--- 阶段 2: 功能开启 - 匹配放行 ---") + + def test_openai_openai(): + r = chat_req(token_key, openai_model) + assert not is_blocked(r), f"OpenAI渠道 + OpenAI端点 应放行, got {r.status_code}" + + def test_claude_claude(): + r = messages_req(token_key, claude_model) + assert not is_blocked(r), f"Claude渠道 + Claude端点 应放行, got {r.status_code}" + + def test_gemini_gemini(): + if gemini_info: + r = gemini_req(token_key, gemini_info[1]) + assert not is_blocked(r), f"Gemini渠道 + Gemini端点 应放行" + + def test_deepseek_openai(): + if deepseek_info: + r = chat_req(token_key, deepseek_info[1]) + assert not is_blocked(r), f"DeepSeek渠道 + OpenAI端点 应放行" + + run_test("OpenAI渠道 + /v1/chat/completions 放行", test_openai_openai) + run_test("Claude渠道 + /v1/messages 放行", test_claude_claude) + run_test("Gemini渠道 + /v1beta 放行", test_gemini_gemini) + run_test("DeepSeek渠道 + /v1/chat/completions 放行", test_deepseek_openai) + + +def phase3_blocked(token_key, openai_model, claude_model, gemini_info, deepseek_info): + """阶段 3: 功能开启 - 不匹配的渠道+端点被拦截""" + print("\n--- 阶段 3: 功能开启 - 不匹配拦截 ---") + + def test_claude_openai_blocked(): + r = chat_req(token_key, claude_model) + assert is_blocked(r), f"Claude渠道 + OpenAI端点 应被拦截, got {r.status_code}: {r.text[:200]}" + + def test_openai_claude_blocked(): + r = messages_req(token_key, openai_model) + assert is_blocked(r), f"OpenAI渠道 + Claude端点 应被拦截, got {r.status_code}: {r.text[:200]}" + + def test_gemini_openai_blocked(): + if gemini_info: + r = chat_req(token_key, gemini_info[1]) + assert is_blocked(r), f"Gemini渠道 + OpenAI端点 应被拦截" + + def test_openai_gemini_blocked(): + if gemini_info: + r = gemini_req(token_key, openai_model) + assert is_blocked(r), f"OpenAI渠道 + Gemini端点 应被拦截" + + def test_deepseek_claude_blocked(): + if deepseek_info: + r = messages_req(token_key, deepseek_info[1]) + assert is_blocked(r), f"DeepSeek渠道 + Claude端点 应被拦截" + + run_test("Claude渠道 + /v1/chat/completions 拦截", test_claude_openai_blocked) + run_test("OpenAI渠道 + /v1/messages 拦截", test_openai_claude_blocked) + run_test("Gemini渠道 + /v1/chat/completions 拦截", test_gemini_openai_blocked) + run_test("OpenAI渠道 + /v1beta 拦截", test_openai_gemini_blocked) + run_test("DeepSeek渠道 + /v1/messages 拦截", test_deepseek_claude_blocked) + + +def phase4_error_messages(token_key, openai_model, claude_model, gemini_info): + """阶段 4: 错误消息验证""" + print("\n--- 阶段 4: 错误消息验证 ---") + + def test_claude_error_msg(): + r = chat_req(token_key, claude_model) + assert is_blocked(r) + body = r.json() + err = body["error"] + assert err["code"] == "endpoint_format_restricted" + assert err["type"] == "new_api_error" + assert "Anthropic" in err["message"] + assert "/v1/messages" in err["message"] + + def test_openai_error_msg(): + r = messages_req(token_key, openai_model) + assert is_blocked(r) + body = r.json() + err = body["error"] + assert err["code"] == "endpoint_format_restricted" + assert "OpenAI" in err["message"] + assert "/v1/chat/completions" in err["message"] + + def test_gemini_error_msg(): + if gemini_info: + r = chat_req(token_key, gemini_info[1]) + assert is_blocked(r) + body = r.json() + err = body["error"] + assert err["code"] == "endpoint_format_restricted" + assert "Gemini" in err["message"] + assert "/v1beta/models/{model}:generateContent" in err["message"] + + run_test("Claude渠道错误消息 (Anthropic + /v1/messages)", test_claude_error_msg) + run_test("OpenAI渠道错误消息 (OpenAI + /v1/chat/completions)", test_openai_error_msg) + run_test("Gemini渠道错误消息 (Gemini + /v1beta)", test_gemini_error_msg) + + +def main(): + global BASE_URL, passed, failed + + parser = argparse.ArgumentParser(description="E2E test: Endpoint Format Restriction") + parser.add_argument("--base-url", default=DEFAULT_BASE_URL) + parser.add_argument("--db-host", default=DEFAULT_DB_HOST) + parser.add_argument("--db-port", type=int, default=DEFAULT_DB_PORT) + parser.add_argument("--db-user", default=DEFAULT_DB_USER) + parser.add_argument("--db-pass", default=DEFAULT_DB_PASS) + parser.add_argument("--db-name", default=DEFAULT_DB_NAME) + parser.add_argument("--token-id", type=int, default=29) + args = parser.parse_args() + + BASE_URL = args.base_url + + print("=" * 60) + print("E2E Test: Endpoint Format Restriction") + print("=" * 60) + + # ---- 阶段 0: 环境发现 ---- + print("\n--- 阶段 0: 环境发现 ---") + + # 健康检查 + try: + r = requests.get(f"{BASE_URL}/api/status", timeout=5) + assert r.status_code == 200 + print(f" Server: {BASE_URL} OK") + except Exception as e: + print(f" ERROR: Server not reachable: {e}") + sys.exit(1) + + # 获取 token + try: + token_key = get_token_key(args) + print(f" Token: id={args.token_id} key=sk-{token_key[:8]}...") + except ValueError: + print(f" ERROR: Token id={args.token_id} not found") + sys.exit(1) + + # 发现渠道 + channels = discover_channels(args) + openai_ch = find_channel(channels, CH_OPENAI) + claude_ch = find_channel(channels, CH_ANTHROPIC) + gemini_ch = find_channel(channels, CH_GEMINI) + deepseek_ch = find_channel(channels, CH_DEEPSEEK) + + if not openai_ch: + print(" ERROR: No OpenAI channel found") + sys.exit(1) + if not claude_ch: + print(" ERROR: No Anthropic channel found") + sys.exit(1) + + openai_model = openai_ch[1] + claude_model = claude_ch[1] + gemini_info = gemini_ch # 可能为 None + deepseek_info = deepseek_ch # 可能为 None + + print(f" OpenAI model: {openai_model} (ch={openai_ch[0]})") + print(f" Claude model: {claude_model} (ch={claude_ch[0]})") + if gemini_info: + print(f" Gemini model: {gemini_info[1]} (ch={gemini_info[0]})") + else: + print(f" Gemini model: (no channel, skip gemini tests)") + if deepseek_info: + print(f" DeepSeek model: {deepseek_info[1]} (ch={deepseek_info[0]})") + + # ---- 阶段 1: 功能关闭 ---- + set_restriction(args, False) + print("\n Set restriction=false, waiting for sync...") + wait_for_restriction(args, token_key, claude_model, enabled=False) + phase1_off(token_key, openai_model, claude_model, gemini_info) + + # ---- 阶段 2: 功能开启 - 放行 ---- + set_restriction(args, True) + print("\n Set restriction=true, waiting for sync...") + if not wait_for_restriction(args, token_key, claude_model, enabled=True): + print(" ERROR: Restriction did not activate. Aborting.") + sys.exit(1) + phase2_allowed(token_key, openai_model, claude_model, gemini_info, deepseek_info) + + # ---- 阶段 3: 功能开启 - 拦截 ---- + phase3_blocked(token_key, openai_model, claude_model, gemini_info, deepseek_info) + + # ---- 阶段 4: 错误消息 ---- + phase4_error_messages(token_key, openai_model, claude_model, gemini_info) + + # ---- 阶段 5: 清理 ---- + set_restriction(args, False) + print(f"\n--- 阶段 5: 清理 ---") + print(f" Restriction reset to: false") + + # ---- 结果 ---- + print("\n" + "=" * 60) + print(f"Results: {passed} passed, {failed} failed") + if failed > 0: + sys.exit(1) + else: + print("ALL TESTS PASSED!") + + +if __name__ == "__main__": + main() diff --git a/web/src/components/settings/ModelSetting.jsx b/web/src/components/settings/ModelSetting.jsx index 7084b66..8cf42f2 100644 --- a/web/src/components/settings/ModelSetting.jsx +++ b/web/src/components/settings/ModelSetting.jsx @@ -40,6 +40,7 @@ const ModelSetting = () => { 'claude.default_max_tokens': '', 'claude.thinking_adapter_budget_tokens_percentage': 0.8, 'global.pass_through_request_enabled': false, + 'global.endpoint_format_restriction_enabled': false, 'global.thinking_model_blacklist': '[]', 'global.chat_completions_to_responses_policy': '{}', 'general_setting.ping_interval_enabled': false, diff --git a/web/src/pages/Setting/Model/SettingGlobalModel.jsx b/web/src/pages/Setting/Model/SettingGlobalModel.jsx index 4b8f9f4..328ba11 100644 --- a/web/src/pages/Setting/Model/SettingGlobalModel.jsx +++ b/web/src/pages/Setting/Model/SettingGlobalModel.jsx @@ -68,6 +68,7 @@ const chatCompletionsToResponsesPolicyAllChannelsExample = JSON.stringify( const defaultGlobalSettingInputs = { 'global.pass_through_request_enabled': false, + 'global.endpoint_format_restriction_enabled': false, 'global.thinking_model_blacklist': '[]', 'global.chat_completions_to_responses_policy': '{}', 'general_setting.ping_interval_enabled': false, @@ -205,6 +206,23 @@ export default function SettingGlobalModel(props) { /> + + + + setInputs({ + ...inputs, + 'global.endpoint_format_restriction_enabled': value, + }) + } + extraText={t( + '开启后,渠道模型只能通过原生协议端点访问(如 Claude 模型只能走 /v1/messages),拒绝跨协议转换请求', + )} + /> + +