From 3256d84c306e31f7a9e638e5273c13876917b052 Mon Sep 17 00:00:00 2001 From: chengqian Date: Fri, 7 Aug 2026 17:41:16 +0800 Subject: [PATCH] =?UTF-8?q?fix(apikey):=E4=BF=AE=E5=A4=8Dapi=20key?= =?UTF-8?q?=E5=BC=82=E6=AD=A5=E6=A8=A1=E5=9E=8B=E8=B6=8A=E6=9D=83=E8=B0=83?= =?UTF-8?q?=E7=94=A8=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/handler/relay/task.go | 90 ++++++++++----- internal/handler/relay/task_test.go | 168 ++++++++++++++++++++++++++++ relay/common/provider.go | 3 + 3 files changed, 230 insertions(+), 31 deletions(-) create mode 100644 internal/handler/relay/task_test.go diff --git a/internal/handler/relay/task.go b/internal/handler/relay/task.go index 8c50c69..84d6826 100644 --- a/internal/handler/relay/task.go +++ b/internal/handler/relay/task.go @@ -3,6 +3,9 @@ package relay import ( "context" "encoding/json" + "errors" + "fmt" + "net/http" "strings" "time" @@ -58,27 +61,7 @@ func HandleTaskSubmit(r *ghttp.Request) { // 选择渠道 channelMeta, err := selectTaskChannel(r, body) if err != nil { - statusCode := 503 - errType := "server_error" - errMsg := err.Error() - - if err == relay_common.ErrTenantModelNotEnabled { - statusCode = 403 - errType = "permission_denied" - errMsg = "当前租户未启用该模型，请联系管理员" - } else if err == relay_common.ErrMemberModelNotAllowed { - statusCode = 403 - errType = "permission_denied" - errMsg = "当前成员无权使用该模型" - } else if err == relay_common.ErrChannelUnavailable { - statusCode = 503 - errType = "server_error" - errMsg = "该模型暂无可用的渠道，请稍后重试或联系管理员" - } - - r.Response.WriteStatus(statusCode, g.Map{ - "error": g.Map{"type": errType, "message": errMsg}, - }) + writeTaskChannelSelectionError(r, err) return } @@ -166,6 +149,58 @@ func extractTaskID(r *ghttp.Request) string { return r.Get("task_id").String() } +type taskModelAccessChecker interface { + CheckMemberModelAccess(ctx context.Context, tenantID, userID int64, modelName string) (bool, error) + CheckApiKeyModelAccess(ctx context.Context, apiKeyID int64, modelName string) (bool, error) +} + +func checkTaskModelAccess( + ctx context.Context, + checker taskModelAccessChecker, + tenantID, userID, apiKeyID int64, + modelName string, +) error { + allowed, err := checker.CheckMemberModelAccess(ctx, tenantID, userID, modelName) + if err != nil { + return fmt.Errorf("check member model access: %w", err) + } + if !allowed { + return relay_common.ErrMemberModelNotAllowed + } + + allowed, err = checker.CheckApiKeyModelAccess(ctx, apiKeyID, modelName) + if err != nil { + return fmt.Errorf("check API key model access: %w", err) + } + if !allowed { + return relay_common.ErrApiKeyModelNotAllowed + } + + return nil +} + +func taskChannelSelectionError(err error) (statusCode int, errType, message string) { + switch { + case errors.Is(err, relay_common.ErrTenantModelNotEnabled): + return http.StatusForbidden, "permission_denied", "当前租户未启用该模型，请联系管理员" + case errors.Is(err, relay_common.ErrMemberModelNotAllowed): + return http.StatusForbidden, "permission_denied", "当前成员无权使用该模型" + case errors.Is(err, relay_common.ErrApiKeyModelNotAllowed): + return http.StatusForbidden, "permission_denied", "当前 API Key 无权使用该模型" + case errors.Is(err, relay_common.ErrChannelUnavailable): + return http.StatusServiceUnavailable, "server_error", "该模型暂无可用的渠道，请稍后重试或联系管理员" + default: + return http.StatusServiceUnavailable, "server_error", err.Error() + } +} + +func writeTaskChannelSelectionError(r *ghttp.Request, err error) { + statusCode, errType, message := taskChannelSelectionError(err) + r.Response.WriteStatus(statusCode, g.Map{ + "error": g.Map{"type": errType, "message": message}, + }) +} + // selectTaskChannel 为异步任务选择渠道（新调度引擎单次选择）。 // 任务提交是瞬时动作：选定并物化后立即结束调度会话释放租约； // 会话绑定已由调度器在选择时写入，提交结果不额外上报健康（与旧行为一致）。 @@ -183,11 +218,8 @@ func selectTaskChannel(r *ghttp.Request, body []byte) (*relay_common.ChannelMeta userID := middleware.GetUserID(ctx) apiKeyID := middleware.GetApiKeyID(ctx) - // 检查成员模型范围 - if allowed, err := relayDataProvider.CheckMemberModelAccess(ctx, tenantID, userID, req.Model); err != nil { + if err := checkTaskModelAccess(ctx, relayDataProvider, tenantID, userID, apiKeyID, req.Model); err != nil { return nil, err - } else if !allowed { - return nil, relay_common.ErrMemberModelNotAllowed } enabled, channelScope, err := relayDataProvider.CheckTenantModelAccess(ctx, tenantID, req.Model) @@ -267,9 +299,7 @@ func HandleMjSubmit(r *ghttp.Request) { channelMeta, err := selectTaskChannel(r, body) if err != nil { - r.Response.WriteStatus(503, g.Map{ - "error": g.Map{"type": "server_error", "message": err.Error()}, - }) + writeTaskChannelSelectionError(r, err) return } @@ -322,9 +352,7 @@ func HandleAliImageSubmit(r *ghttp.Request) { channelMeta, err := selectTaskChannel(r, body) if err != nil { - r.Response.WriteStatus(503, g.Map{ - "error": g.Map{"type": "server_error", "message": err.Error()}, - }) + writeTaskChannelSelectionError(r, err) return } diff --git a/internal/handler/relay/task_test.go b/internal/handler/relay/task_test.go new file mode 100644 index 0000000..2f658cd --- /dev/null +++ b/internal/handler/relay/task_test.go @@ -0,0 +1,168 @@ +package relay + +import ( + "context" + "errors" + "net/http/httptest" + "testing" + + "github.com/gogf/gf/v2/net/ghttp" + + "github.com/qianfree/team-api/internal/middleware" + relay_common "github.com/qianfree/team-api/relay/common" +) + +type taskModelAccessCheckerStub struct { + relay_common.DataProvider + + memberAllowed bool + apiKeyAllowed bool + memberErr error + apiKeyErr error + apiKeyChecks int + tenantChecks int +} + +func (s *taskModelAccessCheckerStub) CheckMemberModelAccess( + context.Context, + int64, + int64, + string, +) (bool, error) { + return s.memberAllowed, s.memberErr +} + +func (s *taskModelAccessCheckerStub) CheckApiKeyModelAccess( + context.Context, + int64, + string, +) (bool, error) { + s.apiKeyChecks++ + return s.apiKeyAllowed, s.apiKeyErr +} + +func (s *taskModelAccessCheckerStub) CheckTenantModelAccess( + context.Context, + int64, + string, +) (bool, []int64, error) { + s.tenantChecks++ + return false, nil, errors.New("tenant model check should not be reached") +} + +func TestCheckTaskModelAccess(t *testing.T) { + errMemberCheck := errors.New("member check failed") + errApiKeyCheck := errors.New("api key check failed") + + tests := []struct { + name string + checker *taskModelAccessCheckerStub + wantErr error + wantApiKeyChecks int + }{ + { + name: "allows model when both scopes allow it", + checker: &taskModelAccessCheckerStub{ + memberAllowed: true, + apiKeyAllowed: true, + }, + wantApiKeyChecks: 1, + }, + { + name: "rejects member before checking API key", + checker: &taskModelAccessCheckerStub{ + memberAllowed: false, + apiKeyAllowed: true, + }, + wantErr: relay_common.ErrMemberModelNotAllowed, + }, + { + name: "rejects model outside API key scope", + checker: &taskModelAccessCheckerStub{ + memberAllowed: true, + apiKeyAllowed: false, + }, + wantErr: relay_common.ErrApiKeyModelNotAllowed, + wantApiKeyChecks: 1, + }, + { + name: "propagates member scope lookup error", + checker: &taskModelAccessCheckerStub{ + memberErr: errMemberCheck, + }, + wantErr: errMemberCheck, + }, + { + name: "propagates API key scope lookup error", + checker: &taskModelAccessCheckerStub{ + memberAllowed: true, + apiKeyErr: errApiKeyCheck, + }, + wantErr: errApiKeyCheck, + wantApiKeyChecks: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := checkTaskModelAccess( + context.Background(), + tt.checker, + 1, + 2, + 3, + "video-model", + ) + if !errors.Is(err, tt.wantErr) { + t.Fatalf("checkTaskModelAccess() error = %v, want %v", err, tt.wantErr) + } + if tt.checker.apiKeyChecks != tt.wantApiKeyChecks { + t.Fatalf("API key checks = %d, want %d", tt.checker.apiKeyChecks, tt.wantApiKeyChecks) + } + }) + } +} + +func TestSelectTaskChannelRejectsModelOutsideApiKeyScope(t *testing.T) { + checker := &taskModelAccessCheckerStub{ + memberAllowed: true, + apiKeyAllowed: false, + } + originalProvider := relayDataProvider + relayDataProvider = checker + t.Cleanup(func() { + relayDataProvider = originalProvider + }) + + r := &ghttp.Request{ + Request: httptest.NewRequest("POST", "/v1/video/generations", nil), + } + r.SetCtxVar(middleware.CtxKeyTenantID, int64(1)) + r.SetCtxVar(middleware.CtxKeyUserID, int64(2)) + r.SetCtxVar(middleware.CtxKeyApiKeyID, int64(3)) + + _, err := selectTaskChannel(r, []byte(`{"model":"video-model"}`)) + if !errors.Is(err, relay_common.ErrApiKeyModelNotAllowed) { + t.Fatalf("selectTaskChannel() error = %v, want %v", err, relay_common.ErrApiKeyModelNotAllowed) + } + if checker.apiKeyChecks != 1 { + t.Fatalf("API key checks = %d, want 1", checker.apiKeyChecks) + } + if checker.tenantChecks != 0 { + t.Fatalf("tenant model checks = %d, want 0 after authorization denial", checker.tenantChecks) + } +} + +func TestTaskChannelSelectionError_ApiKeyModelDenied(t *testing.T) { + statusCode, errType, message := taskChannelSelectionError(relay_common.ErrApiKeyModelNotAllowed) + + if statusCode != 403 { + t.Fatalf("status code = %d, want 403", statusCode) + } + if errType != "permission_denied" { + t.Fatalf("error type = %q, want permission_denied", errType) + } + if message != "当前 API Key 无权使用该模型" { + t.Fatalf("message = %q, want API Key model denial message", message) + } +} diff --git a/relay/common/provider.go b/relay/common/provider.go index d539931..d5c0240 100644 --- a/relay/common/provider.go +++ b/relay/common/provider.go @@ -21,6 +21,9 @@ var ErrTenantModelNotEnabled = errors.New("model not enabled for this tenant") // ErrMemberModelNotAllowed 成员无权使用该模型 var ErrMemberModelNotAllowed = errors.New("model not allowed for this member") +// ErrApiKeyModelNotAllowed API Key 无权使用该模型 +var ErrApiKeyModelNotAllowed = errors.New("model not allowed for this API key") + // DeprecationInfo 模型弃用信息 type DeprecationInfo struct { Deprecated bool