Skip to content

Commit 530a162

Browse files
author
QTom
committed
fix(gateway): 分组隔离 — 禁止未分组账号被跨组调度
当 API Key 无分组时,调度仅从未分组账号池中选取。 修复 isAccountInGroup 在 groupID==nil 时的逻辑, 同时补全 scheduler_snapshot_service 和 gemini_compat_service 中的 SimpleMode 保护,确保分组隔离在所有调度路径生效。 新增 ListSchedulableUngroupedByPlatform/s 方法, 使用 Ent 的 Not(HasAccountGroups()) 谓词实现未分组账号隔离。 新增 17 个单元和端到端隔离测试,覆盖所有分支和边界条件。
1 parent 9792b17 commit 530a162

14 files changed

Lines changed: 475 additions & 10 deletions

backend/internal/handler/sora_client_handler_test.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2089,6 +2089,12 @@ func (r *stubAccountRepoForHandler) ListSchedulableByPlatforms(context.Context,
20892089
func (r *stubAccountRepoForHandler) ListSchedulableByGroupIDAndPlatforms(context.Context, int64, []string) ([]service.Account, error) {
20902090
return r.accounts, nil
20912091
}
2092+
func (r *stubAccountRepoForHandler) ListSchedulableUngroupedByPlatform(_ context.Context, _ string) ([]service.Account, error) {
2093+
return r.accounts, nil
2094+
}
2095+
func (r *stubAccountRepoForHandler) ListSchedulableUngroupedByPlatforms(_ context.Context, _ []string) ([]service.Account, error) {
2096+
return r.accounts, nil
2097+
}
20922098
func (r *stubAccountRepoForHandler) SetRateLimited(context.Context, int64, time.Time) error {
20932099
return nil
20942100
}

backend/internal/handler/sora_gateway_handler_test.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,12 @@ func (r *stubAccountRepo) ListSchedulableByPlatforms(ctx context.Context, platfo
182182
func (r *stubAccountRepo) ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]service.Account, error) {
183183
return r.ListSchedulableByPlatforms(ctx, platforms)
184184
}
185+
func (r *stubAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
186+
return r.ListSchedulableByPlatform(ctx, platform)
187+
}
188+
func (r *stubAccountRepo) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]service.Account, error) {
189+
return r.ListSchedulableByPlatforms(ctx, platforms)
190+
}
185191
func (r *stubAccountRepo) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
186192
return nil
187193
}

backend/internal/repository/account_repo.go

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -829,6 +829,51 @@ func (r *accountRepository) ListSchedulableByPlatforms(ctx context.Context, plat
829829
return r.accountsToService(ctx, accounts)
830830
}
831831

832+
func (r *accountRepository) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
833+
now := time.Now()
834+
accounts, err := r.client.Account.Query().
835+
Where(
836+
dbaccount.PlatformEQ(platform),
837+
dbaccount.StatusEQ(service.StatusActive),
838+
dbaccount.SchedulableEQ(true),
839+
dbaccount.Not(dbaccount.HasAccountGroups()),
840+
tempUnschedulablePredicate(),
841+
notExpiredPredicate(now),
842+
dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
843+
dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
844+
).
845+
Order(dbent.Asc(dbaccount.FieldPriority)).
846+
All(ctx)
847+
if err != nil {
848+
return nil, err
849+
}
850+
return r.accountsToService(ctx, accounts)
851+
}
852+
853+
func (r *accountRepository) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]service.Account, error) {
854+
if len(platforms) == 0 {
855+
return nil, nil
856+
}
857+
now := time.Now()
858+
accounts, err := r.client.Account.Query().
859+
Where(
860+
dbaccount.PlatformIn(platforms...),
861+
dbaccount.StatusEQ(service.StatusActive),
862+
dbaccount.SchedulableEQ(true),
863+
dbaccount.Not(dbaccount.HasAccountGroups()),
864+
tempUnschedulablePredicate(),
865+
notExpiredPredicate(now),
866+
dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
867+
dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
868+
).
869+
Order(dbent.Asc(dbaccount.FieldPriority)).
870+
All(ctx)
871+
if err != nil {
872+
return nil, err
873+
}
874+
return r.accountsToService(ctx, accounts)
875+
}
876+
832877
func (r *accountRepository) ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]service.Account, error) {
833878
if len(platforms) == 0 {
834879
return nil, nil

backend/internal/server/api_contract_test.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1027,6 +1027,14 @@ func (s *stubAccountRepo) ListSchedulableByGroupIDAndPlatforms(ctx context.Conte
10271027
return nil, errors.New("not implemented")
10281028
}
10291029

1030+
func (s *stubAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
1031+
return nil, errors.New("not implemented")
1032+
}
1033+
1034+
func (s *stubAccountRepo) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]service.Account, error) {
1035+
return nil, errors.New("not implemented")
1036+
}
1037+
10301038
func (s *stubAccountRepo) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
10311039
return errors.New("not implemented")
10321040
}

backend/internal/service/account_service.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,8 @@ type AccountRepository interface {
5454
ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error)
5555
ListSchedulableByPlatforms(ctx context.Context, platforms []string) ([]Account, error)
5656
ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]Account, error)
57+
ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error)
58+
ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]Account, error)
5759

5860
SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error
5961
SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error

backend/internal/service/account_service_delete_test.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,14 @@ func (s *accountRepoStub) ListSchedulableByGroupIDAndPlatforms(ctx context.Conte
147147
panic("unexpected ListSchedulableByGroupIDAndPlatforms call")
148148
}
149149

150+
func (s *accountRepoStub) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
151+
panic("unexpected ListSchedulableUngroupedByPlatform call")
152+
}
153+
154+
func (s *accountRepoStub) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]Account, error) {
155+
panic("unexpected ListSchedulableUngroupedByPlatforms call")
156+
}
157+
150158
func (s *accountRepoStub) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
151159
panic("unexpected SetRateLimited call")
152160
}

0 commit comments

Comments
 (0)