seewe/sub2api
1
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-10-14 12:02:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): constrain account model route ownership

Browse Source
This commit is contained in:
wucm667
2026-09-23 15:29:11 +08:00
parent 13fcfc0f3b
commit e878ab710f
2 changed files with 119 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/service/openai_account_scheduler.go
+5
View File
@@ -1776,6 +1776,11 @@ func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatibleReason(ctx con
if account == nil {
return false, "account_nil"
}
if source, ok := CompositeRouteSourceFromContext(ctx); ok && source == CompositeRouteSourceAccount {
if publicModel, modelOK := RequestedPublicModelFromContext(ctx); modelOK && !explicitModelMappingClaims(*account, publicModel) {
return false, "account_model_not_owned"
}
}
if req.RequirePrivacySet && !account.IsPrivacySet() {
return false, "privacy_not_set"
}
backend/internal/service/openai_account_scheduler_test.go
+114
View File
@@ -44,6 +44,10 @@ func (r schedulerTestOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx
return result, nil
}
func (r schedulerTestOpenAIAccountRepo) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) {
return append([]Account(nil), r.accounts...), nil
}
func (r schedulerTestOpenAIAccountRepo) ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) {
var result []Account
for _, acc := range r.accounts {
@@ -1257,6 +1261,116 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedSessionIn
}
}
func TestOpenAIGatewayService_SelectAccountWithScheduler_AccountModelRouteOnlyUsesClaimingAccount(t *testing.T) {
resetOpenAIAdvancedSchedulerSettingCacheForTest()
ctx := context.Background()
groupID := int64(1010711)
owner := Account{
ID: 371011, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10,
GroupIDs: []int64{groupID},
Credentials: map[string]any{"model_mapping": map[string]any{"team-alias": "gpt-5.1"}},
}
nonOwner := Account{
ID: 371012, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0,
GroupIDs: []int64{groupID},
}
repo := schedulerTestOpenAIAccountRepo{accounts: []Account{owner, nonOwner}}
ownership, err := (&GatewayService{accountRepo: repo}).resolveCompositeModelOwnership(ctx, groupID, "team-alias")
require.NoError(t, err)
require.Equal(t, CompositeModelOwnership{TargetPlatform: PlatformOpenAI, Matched: true}, ownership)
ctx = WithCompositeRouteDecision(ctx, CompositeRouteDecision{
Matched: true, Source: CompositeRouteSourceAccount, GroupID: groupID,
PublicModel: "team-alias", TargetPlatform: ownership.TargetPlatform, UpstreamModel: "team-alias",
})
for _, tt := range []struct {
name string
sessionHash string
sessionBindings map[string]int64
}{
{name: "initial filter excludes non owner"},
{name: "non owner sticky hit is rejected", sessionHash: "non-owner", sessionBindings: map[string]int64{"openai:non-owner": nonOwner.ID}},
} {
t.Run(tt.name, func(t *testing.T) {
cfg := &config.Config{}
cfg.Gateway.Scheduling.LoadBatchEnabled = false
svc := &OpenAIGatewayService{
accountRepo: repo,
cache: &schedulerTestGatewayCache{sessionBindings: tt.sessionBindings},
cfg: cfg,
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
}
selection, _, selectErr := svc.SelectAccountWithScheduler(
ctx, &groupID, "", tt.sessionHash, "team-alias", nil, OpenAIUpstreamTransportAny, false,
)
require.NoError(t, selectErr)
require.NotNil(t, selection)
require.NotNil(t, selection.Account)
require.Equal(t, owner.ID, selection.Account.ID)
if selection.ReleaseFunc != nil {
selection.ReleaseFunc()
}
})
}
}
func TestOpenAIGatewayService_SelectAccountWithScheduler_AccountModelRouteDBRecheckFailsOverToOwner(t *testing.T) {
ctx := context.Background()
groupID := int64(1010712)
mapping := map[string]any{"model_mapping": map[string]any{"team-alias": "gpt-5.1"}}
stalePrimary := &Account{
ID: 371021, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0,
GroupIDs: []int64{groupID}, Credentials: mapping,
}
staleBackup := &Account{
ID: 371022, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10,
GroupIDs: []int64{groupID}, Credentials: mapping,
}
dbPrimary := *stalePrimary
dbPrimary.Credentials = nil
dbBackup := *staleBackup
repo := schedulerTestOpenAIAccountRepo{accounts: []Account{dbPrimary, dbBackup}}
ctx = WithCompositeRouteDecision(ctx, CompositeRouteDecision{
Matched: true, Source: CompositeRouteSourceAccount, GroupID: groupID,
PublicModel: "team-alias", TargetPlatform: PlatformOpenAI, UpstreamModel: "team-alias",
})
acquiredIDs, releasedIDs := []int64{}, []int64{}
svc := &OpenAIGatewayService{
accountRepo: repo,
cfg: &config.Config{RunMode: config.RunModeStandard},
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
schedulerSnapshot: &SchedulerSnapshotService{cache: &openAISnapshotCacheStub{
snapshotAccounts: []*Account{stalePrimary, staleBackup},
accountsByID: map[int64]*Account{stalePrimary.ID: stalePrimary, staleBackup.ID: staleBackup},
}},
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{
acquiredIDs: &acquiredIDs, releasedIDs: &releasedIDs,
}),
}
selection, _, err := svc.SelectAccountWithScheduler(
ctx, &groupID, "", "", "team-alias", nil, OpenAIUpstreamTransportAny, false,
)
require.NoError(t, err)
require.NotNil(t, selection)
require.NotNil(t, selection.Account)
require.Equal(t, dbBackup.ID, selection.Account.ID)
if containsInt64(acquiredIDs, dbPrimary.ID) {
require.Contains(t, releasedIDs, dbPrimary.ID)
}
require.Contains(t, acquiredIDs, dbBackup.ID)
if selection.ReleaseFunc != nil {
selection.ReleaseFunc()
}
}
func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedPreviousRequiresMovableContext(t *testing.T) {
resetOpenAIAdvancedSchedulerSettingCacheForTest()