Harden account group routing controls

Author: DeliciousBuding

admin/account_groups.go

@@ -201,9 +201,37 @@ func (h *Handler) DeleteAccountGroup(c *gin.Context) {
 			h.store.ApplyAccountGroups(acc.DBID, groups)
 		}
 	}
+	if err := h.removeDeletedGroupFromAPIKeyScopes(ctx, id); err != nil {
+		writeInternalError(c, err)
+		return
+	}
 	writeMessage(c, http.StatusOK, "分组已删除")
 }
 
+func (h *Handler) removeDeletedGroupFromAPIKeyScopes(ctx context.Context, groupID int64) error {
+	if h == nil || h.db == nil || groupID <= 0 {
+		return nil
+	}
+	keys, err := h.db.ListAPIKeys(ctx)
+	if err != nil {
+		return err
+	}
+	for _, key := range keys {
+		if key == nil || !containsInt64(key.AllowedGroupIDs, groupID) {
+			continue
+		}
+		next := removeInt64(key.AllowedGroupIDs, groupID)
+		if err := h.db.UpdateAPIKeyAllowedGroupIDs(ctx, key.ID, next); err != nil {
+			return err
+		}
+		if h.store != nil {
+			h.store.SetAPIKeyAllowedGroups(key.ID, next)
+		}
+		h.invalidateAPIKeyRuntimeCaches(ctx, key.Key)
+	}
+	return nil
+}
+
 func sanitizeAccountGroupName(raw string) (string, error) {
 	name := strings.TrimSpace(raw)
 	if name == "" {
@@ -230,6 +258,15 @@ func removeInt64(slice []int64, target int64) []int64 {
 	return out
 }
 
+func containsInt64(slice []int64, target int64) bool {
+	for _, v := range slice {
+		if v == target {
+			return true
+		}
+	}
+	return false
+}
+
 func dedupeInt64(ids []int64) []int64 {
 	if len(ids) == 0 {
 		return nil

admin/handler.go

@@ -3361,6 +3361,9 @@ func (h *Handler) DeleteAPIKey(c *gin.Context) {
 		writeError(c, http.StatusInternalServerError, "删除失败: "+err.Error())
 		return
 	}
+	if h.store != nil {
+		h.store.SetAPIKeyAllowedGroups(id, nil)
+	}
 	h.invalidateAPIKeyRuntimeCaches(ctx, keyToInvalidate)
 	writeMessage(c, http.StatusOK, "已删除")
 }

auth/fast_scheduler_test.go

@@ -147,6 +147,29 @@ func TestStoreNextExcludingRespectsAPIKeyWhitelist(t *testing.T) {
 	}
 }
 
+func TestStoreNextExcludingRespectsAPIKeyAllowedGroups(t *testing.T) {
+	restricted := newFastSchedulerTestAccount(1, HealthTierHealthy, 120, 1)
+	restricted.GroupIDs = []int64{10}
+	fallback := newFastSchedulerTestAccount(2, HealthTierHealthy, 80, 1)
+	fallback.GroupIDs = []int64{20}
+
+	store := &Store{
+		accounts:       []*Account{restricted, fallback},
+		maxConcurrency: 1,
+	}
+	store.SetAPIKeyAllowedGroups(1, []int64{20})
+
+	got := store.NextExcluding(1, nil)
+	if got == nil {
+		t.Fatal("NextExcluding() returned nil")
+	}
+	defer store.Release(got)
+
+	if got.DBID != 2 {
+		t.Fatalf("NextExcluding() picked dbID=%d, want 2", got.DBID)
+	}
+}
+
 func TestStoreNextSkipsDispatchPausedAccount(t *testing.T) {
 	paused := newFastSchedulerTestAccount(1, HealthTierHealthy, 120, 1)
 	atomic.StoreInt32(&paused.DispatchPaused, 1)

auth/session_affinity_test.go

@@ -316,3 +316,26 @@ func TestNextForSessionFallsBackWhenAPIKeyNotAllowed(t *testing.T) {
 		t.Fatalf("proxyURL = %q, want empty fallback proxy", proxyURL)
 	}
 }
+
+func TestNextForSessionFallsBackWhenAPIKeyGroupNotAllowed(t *testing.T) {
+	store := &Store{
+		accounts: []*Account{
+			{DBID: 1, AccessToken: "tok-1", GroupIDs: []int64{20}},
+			{DBID: 2, AccessToken: "tok-2", GroupIDs: []int64{10}},
+		},
+		maxConcurrency: 2,
+	}
+	store.SetAPIKeyAllowedGroups(1, []int64{20})
+	store.bindSessionAffinity("session-1", store.accounts[1], "http://proxy-2")
+
+	acc, proxyURL := store.NextForSession("session-1", 1, nil)
+	if acc == nil {
+		t.Fatal("expected fallback account")
+	}
+	if acc.DBID != 1 {
+		t.Fatalf("account DBID = %d, want %d", acc.DBID, 1)
+	}
+	if proxyURL != "" {
+		t.Fatalf("proxyURL = %q, want empty fallback proxy", proxyURL)
+	}
+}

auth/store.go

@@ -2515,7 +2515,7 @@ func (s *Store) NextExcludingWithFilter(apiKeyID int64, exclude map[int64]bool,
 			if !acc.IsAvailable() {
 				continue
 			}
-			if !acc.AllowsAPIKey(apiKeyID) {
+			if !s.accountAllowedForAPIKey(acc, apiKeyID) {
 				continue
 			}
 			if filter != nil && !filter(acc) {
@@ -2709,7 +2709,7 @@ func (s *Store) takeByIDExcluding(id int64, apiKeyID int64, exclude map[int64]bo
 	if s.accountHasCachedCooldown(target) {
 		return nil
 	}
-	if !target.AllowsAPIKey(apiKeyID) {
+	if !s.accountAllowedForAPIKey(target, apiKeyID) {
 		return nil
 	}
 	if filter != nil && !filter(target) {
@@ -2761,7 +2761,7 @@ func (s *Store) hasDispatchCandidateWithFilter(apiKeyID int64, exclude map[int64
 		if s.accountHasCachedCooldown(acc) {
 			continue
 		}
-		if !acc.AllowsAPIKey(apiKeyID) {
+		if !s.accountAllowedForAPIKey(acc, apiKeyID) {
 			continue
 		}
 		if filter != nil && !filter(acc) {
@@ -3146,6 +3146,13 @@ func (s *Store) APIKeyAllowsAccount(apiKeyID int64, acc *Account) bool {
 	return false
 }
 
+func (s *Store) accountAllowedForAPIKey(acc *Account, apiKeyID int64) bool {
+	if acc == nil {
+		return false
+	}
+	return acc.AllowsAPIKey(apiKeyID) && s.APIKeyAllowsAccount(apiKeyID, acc)
+}
+
 func (s *Store) ApplyOpenAIResponsesConfig(dbID int64, baseURL, apiKey string, models []string, proxyURL string) bool {
 	acc := s.FindByID(dbID)
 	if acc == nil {

database/postgres.go

@@ -868,7 +868,7 @@ func (row *APIKeyRow) IsQuotaExhausted() bool {
 }
 
 func (row *APIKeyRow) HasAccessConstraints() bool {
-	return row != nil && (row.QuotaLimit > 0 || row.ExpiresAt.Valid)
+	return row != nil && (row.QuotaLimit > 0 || row.ExpiresAt.Valid || len(row.AllowedGroupIDs) > 0)
 }
 
 // UpdateAPIKeyName updates the display name of an API key without changing the key value.

frontend/src/locales/en.json

@@ -404,13 +404,17 @@
     "groupManageTitle": "Account Group Management",
     "groupCreate": "Create Group",
     "groupCreateTitle": "Create Account Group",
+    "groupEdit": "Edit Group",
+    "groupEditTitle": "Edit Account Group",
     "groupName": "Group Name",
+    "groupNameRequired": "Enter a group name",
     "groupNamePlaceholder": "e.g. Pro Pool",
     "groupDescription": "Description",
     "groupDescriptionPlaceholder": "Optional description",
     "groupColor": "Color",
     "groupColorPlaceholder": "#2563eb",
     "groupMembers": "Members",
+    "groupNoDescription": "No description",
     "groupEmpty": "No groups yet",
     "groupEmptyDesc": "Create groups, then assign them in account editing.",
     "groupCreated": "Group created",

frontend/src/locales/zh.json

@@ -404,13 +404,17 @@
     "groupManageTitle": "账号分组管理",
     "groupCreate": "创建分组",
     "groupCreateTitle": "创建账号分组",
+    "groupEdit": "编辑分组",
+    "groupEditTitle": "编辑账号分组",
     "groupName": "分组名称",
+    "groupNameRequired": "请输入分组名称",
     "groupNamePlaceholder": "例如 Pro 池",
     "groupDescription": "分组说明",
     "groupDescriptionPlaceholder": "可选说明",
     "groupColor": "颜色",
     "groupColorPlaceholder": "#2563eb",
     "groupMembers": "成员",
+    "groupNoDescription": "暂无说明",
     "groupEmpty": "暂无分组",
     "groupEmptyDesc": "创建分组后可在账号编辑中分配。",
     "groupCreated": "分组已创建",