fix(apikey):修复api key异步模型越权调用问题
This commit is contained in:
chengqian committed
1 parent
908dce79f7
commit
3256d84c30
3 files changed
+230
-31
No files matched your search
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
Reference in new issue
Block a user