当开启端点格式限制后,不同协议的渠道只能通过对应的端点访问: - Anthropic 渠道只能用 /v1/messages - Gemini 渠道只能用 /v1beta/models/* - OpenAI 系渠道只能用 /v1/chat/completions 等 - VertexAI/阿里云等支持多协议的渠道不受限 核心实现:EndpointFormatGuard 中间件从请求路径推导协议家族, 不依赖 gin context 中的 relay_format(因 group middleware 执行 时 context 中尚无此值)。 修复了 /v1/models/*path 和 /v1/engines/ 路径被错误归类为 openai 家族的 bug(实际是 Gemini 格式)。 单元测试 58 个,E2E 测试 15 个,覆盖: - 路径分类边界(含 Gemini relay 路径修复验证) - 渠道类型覆盖(含非 LLM 渠道透传) - 并发安全、运行时开关切换 - 错误消息 JSON 结构和内容验证master
| @@ -37,6 +37,8 @@ const ( | |||||
| ContextKeyChannelMultiKeyIndex ContextKey = "channel_multi_key_index" | ContextKeyChannelMultiKeyIndex ContextKey = "channel_multi_key_index" | ||||
| ContextKeyChannelKey ContextKey = "channel_key" | ContextKeyChannelKey ContextKey = "channel_key" | ||||
| ContextKeyRelayFormat ContextKey = "relay_format" | |||||
| ContextKeyAutoGroup ContextKey = "auto_group" | ContextKeyAutoGroup ContextKey = "auto_group" | ||||
| ContextKeyAutoGroupIndex ContextKey = "auto_group_index" | ContextKeyAutoGroupIndex ContextKey = "auto_group_index" | ||||
| ContextKeyAutoGroupRetryIndex ContextKey = "auto_group_retry_index" | ContextKeyAutoGroupRetryIndex ContextKey = "auto_group_retry_index" | ||||
| @@ -286,6 +286,7 @@ func InitResources() error { | |||||
| // Initialize options, should after model.InitDB() | // Initialize options, should after model.InitDB() | ||||
| model.InitOptionMap() | model.InitOptionMap() | ||||
| model.LoadOptionsFromDatabase() | |||||
| // 清理旧的磁盘缓存文件 | // 清理旧的磁盘缓存文件 | ||||
| common.CleanupOldCacheFiles() | common.CleanupOldCacheFiles() | ||||
| @@ -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 | |||||
| } | |||||
| @@ -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) | |||||
| } | |||||
| @@ -175,7 +175,7 @@ func InitOptionMap() { | |||||
| } | } | ||||
| common.OptionMapRWMutex.Unlock() | common.OptionMapRWMutex.Unlock() | ||||
| loadOptionsFromDatabase() | |||||
| LoadOptionsFromDatabase() | |||||
| } | } | ||||
| // triggerWechatPayReset 安全触发微信支付客户端重置 | // triggerWechatPayReset 安全触发微信支付客户端重置 | ||||
| @@ -192,7 +192,7 @@ func triggerAlipayReset() { | |||||
| } | } | ||||
| } | } | ||||
| func loadOptionsFromDatabase() { | |||||
| func LoadOptionsFromDatabase() { | |||||
| options, _ := AllOption() | options, _ := AllOption() | ||||
| for _, option := range options { | for _, option := range options { | ||||
| err := updateOptionMap(option.Key, option.Value) | err := updateOptionMap(option.Key, option.Value) | ||||
| @@ -206,7 +206,7 @@ func SyncOptions(frequency int) { | |||||
| for { | for { | ||||
| time.Sleep(time.Duration(frequency) * time.Second) | time.Sleep(time.Duration(frequency) * time.Second) | ||||
| common.SysLog("syncing options from database") | common.SysLog("syncing options from database") | ||||
| loadOptionsFromDatabase() | |||||
| LoadOptionsFromDatabase() | |||||
| } | } | ||||
| } | } | ||||
| @@ -567,7 +567,7 @@ func handleConfigUpdate(key, value string) bool { | |||||
| configMap := map[string]string{ | configMap := map[string]string{ | ||||
| configKey: value, | configKey: value, | ||||
| } | } | ||||
| config.UpdateConfigFromMap(cfg, configMap) | |||||
| _ = config.UpdateConfigFromMap(cfg, configMap) | |||||
| // 特定配置的后处理 | // 特定配置的后处理 | ||||
| if configName == "performance_setting" { | if configName == "performance_setting" { | ||||
| @@ -10,6 +10,13 @@ import ( | |||||
| "github.com/gin-gonic/gin" | "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) { | func SetRelayRouter(router *gin.Engine) { | ||||
| router.Use(middleware.CORS()) | router.Use(middleware.CORS()) | ||||
| router.Use(middleware.DecompressRequestMiddleware()) | router.Use(middleware.DecompressRequestMiddleware()) | ||||
| @@ -71,81 +78,49 @@ func SetRelayRouter(router *gin.Engine) { | |||||
| // WebSocket 路由(统一到 Relay) | // WebSocket 路由(统一到 Relay) | ||||
| wsRouter := relayV1Router.Group("") | wsRouter := relayV1Router.Group("") | ||||
| wsRouter.Use(middleware.Distribute()) | 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 | //http router | ||||
| httpRouter := relayV1Router.Group("") | httpRouter := relayV1Router.Group("") | ||||
| httpRouter.Use(middleware.Distribute()) | httpRouter.Use(middleware.Distribute()) | ||||
| httpRouter.Use(middleware.EndpointFormatGuard()) | |||||
| httpRouter.Use(middleware.UserModelRateLimit()) | httpRouter.Use(middleware.UserModelRateLimit()) | ||||
| // claude related routes | // claude related routes | ||||
| httpRouter.POST("/messages", func(c *gin.Context) { | |||||
| controller.Relay(c, types.RelayFormatClaude) | |||||
| }) | |||||
| httpRouter.POST("/messages", relayHandler(types.RelayFormatClaude)) | |||||
| // chat related routes | // 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 | // 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 | // 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 | // embedding related routes | ||||
| httpRouter.POST("/embeddings", func(c *gin.Context) { | |||||
| controller.Relay(c, types.RelayFormatEmbedding) | |||||
| }) | |||||
| httpRouter.POST("/embeddings", relayHandler(types.RelayFormatEmbedding)) | |||||
| // audio related routes | // 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 | // rerank related routes | ||||
| httpRouter.POST("/rerank", func(c *gin.Context) { | |||||
| controller.Relay(c, types.RelayFormatRerank) | |||||
| }) | |||||
| httpRouter.POST("/rerank", relayHandler(types.RelayFormatRerank)) | |||||
| // gemini relay routes | // 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 | // other relay routes | ||||
| httpRouter.POST("/moderations", func(c *gin.Context) { | |||||
| controller.Relay(c, types.RelayFormatOpenAI) | |||||
| }) | |||||
| httpRouter.POST("/moderations", relayHandler(types.RelayFormatOpenAI)) | |||||
| // not implemented | // not implemented | ||||
| httpRouter.POST("/images/variations", controller.RelayNotImplemented) | httpRouter.POST("/images/variations", controller.RelayNotImplemented) | ||||
| @@ -186,12 +161,11 @@ func SetRelayRouter(router *gin.Engine) { | |||||
| relayGeminiRouter.Use(middleware.ModelRequestRateLimit()) | relayGeminiRouter.Use(middleware.ModelRequestRateLimit()) | ||||
| relayGeminiRouter.Use(middleware.RelayCaptureMiddleware()) | relayGeminiRouter.Use(middleware.RelayCaptureMiddleware()) | ||||
| relayGeminiRouter.Use(middleware.Distribute()) | relayGeminiRouter.Use(middleware.Distribute()) | ||||
| relayGeminiRouter.Use(middleware.EndpointFormatGuard()) | |||||
| relayGeminiRouter.Use(middleware.UserModelRateLimit()) | relayGeminiRouter.Use(middleware.UserModelRateLimit()) | ||||
| { | { | ||||
| // Gemini API 路径格式: /v1beta/models/{model_name}:{action} | // 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)) | |||||
| } | } | ||||
| } | } | ||||
| @@ -35,7 +35,8 @@ func (p ChatCompletionsToResponsesPolicy) IsChannelEnabled(channelID int, channe | |||||
| type GlobalSettings struct { | type GlobalSettings struct { | ||||
| PassThroughRequestEnabled bool `json:"pass_through_request_enabled"` | PassThroughRequestEnabled bool `json:"pass_through_request_enabled"` | ||||
| ThinkingModelBlacklist []string `json:"thinking_model_blacklist"` | 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, | Enabled: false, | ||||
| AllChannels: true, | AllChannels: true, | ||||
| }, | }, | ||||
| EndpointFormatRestrictionEnabled: false, | |||||
| } | } | ||||
| // 全局实例 | // 全局实例 | ||||
| @@ -63,6 +65,10 @@ func GetGlobalSettings() *GlobalSettings { | |||||
| return &globalSettings | return &globalSettings | ||||
| } | } | ||||
| func IsEndpointFormatRestrictionEnabled() bool { | |||||
| return globalSettings.EndpointFormatRestrictionEnabled | |||||
| } | |||||
| // ShouldPreserveThinkingSuffix 判断模型是否配置为保留 thinking/-nothinking/-low/-high/-medium 后缀 | // ShouldPreserveThinkingSuffix 判断模型是否配置为保留 thinking/-nothinking/-low/-high/-medium 后缀 | ||||
| func ShouldPreserveThinkingSuffix(modelName string) bool { | func ShouldPreserveThinkingSuffix(modelName string) bool { | ||||
| target := strings.TrimSpace(modelName) | target := strings.TrimSpace(modelName) | ||||
| @@ -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] + "..." | |||||
| } | |||||
| @@ -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() | |||||
| @@ -40,6 +40,7 @@ const ModelSetting = () => { | |||||
| 'claude.default_max_tokens': '', | 'claude.default_max_tokens': '', | ||||
| 'claude.thinking_adapter_budget_tokens_percentage': 0.8, | 'claude.thinking_adapter_budget_tokens_percentage': 0.8, | ||||
| 'global.pass_through_request_enabled': false, | 'global.pass_through_request_enabled': false, | ||||
| 'global.endpoint_format_restriction_enabled': false, | |||||
| 'global.thinking_model_blacklist': '[]', | 'global.thinking_model_blacklist': '[]', | ||||
| 'global.chat_completions_to_responses_policy': '{}', | 'global.chat_completions_to_responses_policy': '{}', | ||||
| 'general_setting.ping_interval_enabled': false, | 'general_setting.ping_interval_enabled': false, | ||||
| @@ -68,6 +68,7 @@ const chatCompletionsToResponsesPolicyAllChannelsExample = JSON.stringify( | |||||
| const defaultGlobalSettingInputs = { | const defaultGlobalSettingInputs = { | ||||
| 'global.pass_through_request_enabled': false, | 'global.pass_through_request_enabled': false, | ||||
| 'global.endpoint_format_restriction_enabled': false, | |||||
| 'global.thinking_model_blacklist': '[]', | 'global.thinking_model_blacklist': '[]', | ||||
| 'global.chat_completions_to_responses_policy': '{}', | 'global.chat_completions_to_responses_policy': '{}', | ||||
| 'general_setting.ping_interval_enabled': false, | 'general_setting.ping_interval_enabled': false, | ||||
| @@ -205,6 +206,23 @@ export default function SettingGlobalModel(props) { | |||||
| /> | /> | ||||
| </Col> | </Col> | ||||
| </Row> | </Row> | ||||
| <Row> | |||||
| <Col xs={24} sm={12} md={8} lg={8} xl={8}> | |||||
| <Form.Switch | |||||
| label={t('端点格式限制')} | |||||
| field={'global.endpoint_format_restriction_enabled'} | |||||
| onChange={(value) => | |||||
| setInputs({ | |||||
| ...inputs, | |||||
| 'global.endpoint_format_restriction_enabled': value, | |||||
| }) | |||||
| } | |||||
| extraText={t( | |||||
| '开启后,渠道模型只能通过原生协议端点访问(如 Claude 模型只能走 /v1/messages),拒绝跨协议转换请求', | |||||
| )} | |||||
| /> | |||||
| </Col> | |||||
| </Row> | |||||
| <Row> | <Row> | ||||
| <Col span={24}> | <Col span={24}> | ||||
| <Form.TextArea | <Form.TextArea | ||||