fix(groups): merge main and resolve model allowlist conflicts

This commit is contained in:
Zero Clover committed 2026-09-07 18:33:39 +08:00
commit 81ff9384cb
252 files changed
+11371 -1598

No files matched your search

+28 -3
View File
@@ -221,16 +221,21 @@ func (r *accountRepository) CreateWithAccountGroups(ctx context.Context, account
// Reuse a caller-owned transaction when this repository is already transactional.
txClient = r.client
}
groupIDs := make([]int64, 0, len(groups))
for i := range groups {
groupIDs = append(groupIDs, groups[i].GroupID)
}
if err := lockLiveGroups(ctx, txClient, groupIDs); err != nil {
return err
}
if err := createAccountRecord(ctx, txClient, account); err != nil {
return err
}
groupIDs := make([]int64, 0, len(groups))
if len(groups) > 0 {
builders := make([]*dbent.AccountGroupCreate, 0, len(groups))
for i := range groups {
groups[i].AccountID = account.ID
groupIDs = append(groupIDs, groups[i].GroupID)
builders = append(builders, txClient.AccountGroup.Create().
SetAccountID(account.ID).
SetGroupID(groups[i].GroupID).
@@ -1759,7 +1764,19 @@ func (r *accountRepository) ClearError(ctx context.Context, id int64) error {
}
func (r *accountRepository) AddToGroup(ctx context.Context, accountID, groupID int64, priority int) error {
_, err := r.client.AccountGroup.Create().
tx, err := r.client.Tx(ctx)
if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
return err
}
client := r.client
if tx != nil {
defer func() { _ = tx.Rollback() }()
client = tx.Client()
}
if err := lockLiveGroups(ctx, client, []int64{groupID}); err != nil {
return err
}
_, err = client.AccountGroup.Create().
SetAccountID(accountID).
SetGroupID(groupID).
SetPriority(priority).
@@ -1767,6 +1784,11 @@ func (r *accountRepository) AddToGroup(ctx context.Context, accountID, groupID i
if err != nil {
return err
}
if tx != nil {
if err := tx.Commit(); err != nil {
return err
}
}
payload := buildSchedulerGroupPayload([]int64{groupID})
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountGroupsChanged, &accountID, nil, payload); err != nil {
logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue add to group failed: account=%d group=%d err=%v", accountID, groupID, err)
@@ -1827,6 +1849,9 @@ func (r *accountRepository) BindGroups(ctx context.Context, accountID int64, gro
// 已处于外部事务中(ErrTxStarted),复用当前 client
txClient = r.client
}
if err := lockLiveGroups(ctx, txClient, groupIDs); err != nil {
return err
}
if _, err := txClient.AccountGroup.Delete().Where(dbaccountgroup.AccountIDEQ(accountID)).Exec(ctx); err != nil {
return err
@@ -16,7 +16,7 @@ import (
func (r *channelRepository) ListModelPricing(ctx context.Context, channelID int64) ([]service.ChannelModelPricing, error) {
rows, err := r.db.QueryContext(ctx,
`SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, image_input_price, image_output_price, per_request_price, time_pricing, created_at, updated_at
`SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, max_reasoning_effort_multiplier, image_input_price, image_output_price, per_request_price, time_pricing, created_at, updated_at
FROM channel_model_pricing WHERE channel_id = $1 ORDER BY id`, channelID,
)
if err != nil {
@@ -61,10 +61,10 @@ func (r *channelRepository) UpdateModelPricing(ctx context.Context, pricing *ser
}
result, err := r.db.ExecContext(ctx,
`UPDATE channel_model_pricing
SET models = $1, billing_mode = $2, input_price = $3, output_price = $4, cache_write_price = $5, cache_write_1h_price = $6, cache_read_price = $7, fast_multiplier = $8, flex_multiplier = $9, image_input_price = $10, image_output_price = $11, per_request_price = $12, time_pricing = $13, platform = $14, updated_at = NOW()
WHERE id = $15`,
SET models = $1, billing_mode = $2, input_price = $3, output_price = $4, cache_write_price = $5, cache_write_1h_price = $6, cache_read_price = $7, fast_multiplier = $8, flex_multiplier = $9, max_reasoning_effort_multiplier = $10, image_input_price = $11, image_output_price = $12, per_request_price = $13, time_pricing = $14, platform = $15, updated_at = NOW()
WHERE id = $16`,
modelsJSON, billingMode, pricing.InputPrice, pricing.OutputPrice, pricing.CacheWritePrice, pricing.CacheWrite1hPrice, pricing.CacheReadPrice,
pricing.FastMultiplier, pricing.FlexMultiplier, pricing.ImageInputPrice, pricing.ImageOutputPrice, pricing.PerRequestPrice,
pricing.FastMultiplier, pricing.FlexMultiplier, pricing.MaxReasoningEffortMultiplier, pricing.ImageInputPrice, pricing.ImageOutputPrice, pricing.PerRequestPrice,
timePricingJSON, pricing.Platform, pricing.ID,
)
if err != nil {
@@ -96,7 +96,7 @@ func (r *channelRepository) ReplaceModelPricing(ctx context.Context, channelID i
// batchLoadModelPricing 批量加载多个渠道的模型定价(含区间)
func (r *channelRepository) batchLoadModelPricing(ctx context.Context, channelIDs []int64) (map[int64][]service.ChannelModelPricing, error) {
rows, err := r.db.QueryContext(ctx,
`SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, image_input_price, image_output_price, per_request_price, time_pricing, created_at, updated_at
`SELECT id, channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, max_reasoning_effort_multiplier, image_input_price, image_output_price, per_request_price, time_pricing, created_at, updated_at
FROM channel_model_pricing WHERE channel_id = ANY($1) ORDER BY channel_id, id`,
pq.Array(channelIDs),
)
@@ -180,7 +180,7 @@ func scanModelPricingRows(rows *sql.Rows) ([]service.ChannelModelPricing, []int6
if err := rows.Scan(
&p.ID, &p.ChannelID, &p.Platform, &modelsJSON, &p.BillingMode,
&p.InputPrice, &p.OutputPrice, &p.CacheWritePrice, &p.CacheWrite1hPrice, &p.CacheReadPrice,
&p.FastMultiplier, &p.FlexMultiplier,
&p.FastMultiplier, &p.FlexMultiplier, &p.MaxReasoningEffortMultiplier,
&p.ImageInputPrice, &p.ImageOutputPrice, &p.PerRequestPrice, &timePricingJSON, &p.CreatedAt, &p.UpdatedAt,
); err != nil {
return nil, nil, fmt.Errorf("scan model pricing: %w", err)
@@ -247,11 +247,11 @@ func createModelPricingExec(ctx context.Context, exec dbExec, pricing *service.C
platform = "anthropic"
}
err = exec.QueryRowContext(ctx,
`INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, image_input_price, image_output_price, per_request_price, time_pricing)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15) RETURNING id, created_at, updated_at`,
`INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, max_reasoning_effort_multiplier, image_input_price, image_output_price, per_request_price, time_pricing)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16) RETURNING id, created_at, updated_at`,
pricing.ChannelID, platform, modelsJSON, billingMode,
pricing.InputPrice, pricing.OutputPrice, pricing.CacheWritePrice, pricing.CacheWrite1hPrice, pricing.CacheReadPrice,
pricing.FastMultiplier, pricing.FlexMultiplier, pricing.ImageInputPrice, pricing.ImageOutputPrice,
pricing.FastMultiplier, pricing.FlexMultiplier, pricing.MaxReasoningEffortMultiplier, pricing.ImageInputPrice, pricing.ImageOutputPrice,
pricing.PerRequestPrice, timePricingJSON,
).Scan(&pricing.ID, &pricing.CreatedAt, &pricing.UpdatedAt)
if err != nil {
@@ -16,7 +16,7 @@ import (
var channelModelPricingTimePricingColumns = []string{
"id", "channel_id", "platform", "models", "billing_mode", "input_price", "output_price",
"cache_write_price", "cache_write_1h_price", "cache_read_price", "fast_multiplier", "flex_multiplier", "image_input_price", "image_output_price",
"cache_write_price", "cache_write_1h_price", "cache_read_price", "fast_multiplier", "flex_multiplier", "max_reasoning_effort_multiplier", "image_input_price", "image_output_price",
"per_request_price", "time_pricing", "created_at", "updated_at",
}
@@ -33,7 +33,7 @@ func newChannelModelPricingTimePricingRepo(t *testing.T) (*channelRepository, sq
func modelPricingTimePricingRow(timePricing any) *sqlmock.Rows {
return sqlmock.NewRows(channelModelPricingTimePricingColumns).AddRow(
int64(11), int64(7), "openai", `["gpt-5"]`, service.BillingModeToken,
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, timePricing,
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, timePricing,
time.Date(2026, 8, 17, 0, 0, 0, 0, time.UTC), time.Date(2026, 8, 17, 1, 0, 0, 0, time.UTC),
)
}
@@ -105,10 +105,10 @@ func TestChannelModelPricingTimePricingCreateAndUpdateRoundTrip(t *testing.T) {
t.Run("create writes JSON", func(t *testing.T) {
repo, mock := newChannelModelPricingTimePricingRepo(t)
mock.ExpectQuery(regexp.QuoteMeta("INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, image_input_price, image_output_price, per_request_price, time_pricing)")).
mock.ExpectQuery(regexp.QuoteMeta("INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, max_reasoning_effort_multiplier, image_input_price, image_output_price, per_request_price, time_pricing)")).
WithArgs(
int64(7), "openai", []byte(`["gpt-5"]`), service.BillingModeToken,
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, channelModelPricingTimePricingJSON,
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, channelModelPricingTimePricingJSON,
).
WillReturnRows(sqlmock.NewRows([]string{"id", "created_at", "updated_at"}).AddRow(int64(11), time.Time{}, time.Time{}))
@@ -118,10 +118,10 @@ func TestChannelModelPricingTimePricingCreateAndUpdateRoundTrip(t *testing.T) {
t.Run("update writes JSON and entry ID", func(t *testing.T) {
repo, mock := newChannelModelPricingTimePricingRepo(t)
mock.ExpectExec(`(?s)UPDATE channel_model_pricing.*per_request_price = \$12, time_pricing = \$13, platform = \$14.*WHERE id = \$15`).
mock.ExpectExec(`(?s)UPDATE channel_model_pricing.*per_request_price = \$13, time_pricing = \$14, platform = \$15.*WHERE id = \$16`).
WithArgs(
[]byte(`["gpt-5"]`), service.BillingModeToken,
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, channelModelPricingTimePricingJSON, "openai", int64(11),
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, channelModelPricingTimePricingJSON, "openai", int64(11),
).
WillReturnResult(sqlmock.NewResult(0, 1))
@@ -153,10 +153,10 @@ func TestChannelModelPricingTimePricingCreateAndUpdateWriteNullWhenDisabled(t *t
t.Run("create writes SQL NULL", func(t *testing.T) {
repo, mock := newChannelModelPricingTimePricingRepo(t)
mock.ExpectQuery(regexp.QuoteMeta("INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, image_input_price, image_output_price, per_request_price, time_pricing)")).
mock.ExpectQuery(regexp.QuoteMeta("INSERT INTO channel_model_pricing (channel_id, platform, models, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, fast_multiplier, flex_multiplier, max_reasoning_effort_multiplier, image_input_price, image_output_price, per_request_price, time_pricing)")).
WithArgs(
int64(7), "openai", []byte(`["gpt-5"]`), service.BillingModeToken,
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
).
WillReturnRows(sqlmock.NewRows([]string{"id", "created_at", "updated_at"}).AddRow(int64(11), time.Time{}, time.Time{}))
@@ -166,10 +166,10 @@ func TestChannelModelPricingTimePricingCreateAndUpdateWriteNullWhenDisabled(t *t
t.Run("update writes SQL NULL", func(t *testing.T) {
repo, mock := newChannelModelPricingTimePricingRepo(t)
mock.ExpectExec(`(?s)UPDATE channel_model_pricing.*per_request_price = \$12, time_pricing = \$13, platform = \$14.*WHERE id = \$15`).
mock.ExpectExec(`(?s)UPDATE channel_model_pricing.*per_request_price = \$13, time_pricing = \$14, platform = \$15.*WHERE id = \$16`).
WithArgs(
[]byte(`["gpt-5"]`), service.BillingModeToken,
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, "openai", int64(11),
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, "openai", int64(11),
).
WillReturnResult(sqlmock.NewResult(0, 1))
+85 -9
View File
@@ -29,6 +29,39 @@ type groupRepository struct {
sql sqlExecutor
}
// lockLiveGroups makes account-group inserts participate in the same row-lock
// protocol as guarded group deletion. FOR SHARE conflicts with the deleter's
// FOR UPDATE lock, and READ COMMITTED rechecks deleted_at after any wait.
func lockLiveGroups(ctx context.Context, exec sqlExecutor, groupIDs []int64) error {
if len(groupIDs) == 0 {
return nil
}
unique := make(map[int64]struct{}, len(groupIDs))
for _, id := range groupIDs {
unique[id] = struct{}{}
}
rows, err := exec.QueryContext(ctx, `/* account_group_live_group_lock */
SELECT id FROM groups
WHERE id = ANY($1) AND deleted_at IS NULL
ORDER BY id
FOR SHARE`, pq.Array(groupIDs))
if err != nil {
return err
}
defer func() { _ = rows.Close() }()
locked := 0
for rows.Next() {
locked++
}
if err := rows.Err(); err != nil {
return err
}
if locked != len(unique) {
return service.ErrGroupNotFound
}
return nil
}
func NewGroupRepository(client *dbent.Client, sqlDB *sql.DB) service.GroupRepository {
return newGroupRepositoryWithSQL(client, sqlDB)
}
@@ -426,6 +459,15 @@ func (r *groupRepository) List(ctx context.Context, params pagination.Pagination
func (r *groupRepository) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, status, search string, isExclusive *bool) ([]service.Group, *pagination.PaginationResult, error) {
q := r.client.Group.Query()
return r.listWithFiltersQuery(ctx, q, params, platform, status, search, isExclusive)
}
func (r *groupRepository) ListBindableWithFilters(ctx context.Context, params pagination.PaginationParams, platform, status, search string, isExclusive *bool) ([]service.Group, *pagination.PaginationResult, error) {
q := r.client.Group.Query().Where(group.PlatformNEQ(service.PlatformComposite))
return r.listWithFiltersQuery(ctx, q, params, platform, status, search, isExclusive)
}
func (r *groupRepository) listWithFiltersQuery(ctx context.Context, q *dbent.GroupQuery, params pagination.PaginationParams, platform, status, search string, isExclusive *bool) ([]service.Group, *pagination.PaginationResult, error) {
if platform != "" {
q = q.Where(group.PlatformEQ(platform))
@@ -807,12 +849,14 @@ func (r *groupRepository) DeleteAccountGroupsByGroupID(ctx context.Context, grou
}
func (r *groupRepository) DeleteCascade(ctx context.Context, id int64) ([]int64, error) {
g, err := r.client.Group.Query().Where(group.IDEQ(id)).Only(ctx)
if err != nil {
return nil, translatePersistenceError(err, service.ErrGroupNotFound, nil)
}
groupSvc := groupEntityToService(g)
return r.deleteCascade(ctx, id, false)
}
func (r *groupRepository) DeleteCascadeIfEmpty(ctx context.Context, id int64) ([]int64, error) {
return r.deleteCascade(ctx, id, true)
}
func (r *groupRepository) deleteCascade(ctx context.Context, id int64, requireEmpty bool) ([]int64, error) {
// 使用 ent 事务统一包裹:避免手工基于 *sql.Tx 构造 ent client 带来的驱动断言问题,
// 同时保证级联删除的原子性。
tx, err := r.client.Tx(ctx)
@@ -830,13 +874,14 @@ func (r *groupRepository) DeleteCascade(ctx context.Context, id int64) ([]int64,
// Lock the group row to avoid concurrent writes while we cascade.
// 这里使用 exec.QueryContext 手动扫描,确保同一事务内加锁并能区分"未找到"与其他错误。
rows, err := exec.QueryContext(ctx, "SELECT id FROM groups WHERE id = $1 AND deleted_at IS NULL FOR UPDATE", id)
rows, err := exec.QueryContext(ctx, "SELECT id, subscription_type FROM groups WHERE id = $1 AND deleted_at IS NULL FOR UPDATE", id)
if err != nil {
return nil, err
}
var lockedID int64
var subscriptionType string
if rows.Next() {
if err := rows.Scan(&lockedID); err != nil {
if err := rows.Scan(&lockedID, &subscriptionType); err != nil {
_ = rows.Close()
return nil, err
}
@@ -850,9 +895,22 @@ func (r *groupRepository) DeleteCascade(ctx context.Context, id int64) ([]int64,
if lockedID == 0 {
return nil, service.ErrGroupNotFound
}
if requireEmpty {
var hasAccount bool
if err := scanSingleRow(ctx, exec, `SELECT EXISTS (
SELECT 1 FROM account_groups ag
JOIN accounts a ON a.id = ag.account_id
WHERE ag.group_id = $1 AND a.deleted_at IS NULL
)`, []any{id}, &hasAccount); err != nil {
return nil, err
}
if hasAccount {
return nil, service.ErrGroupNotEmpty
}
}
var affectedUserIDs []int64
if groupSvc.IsSubscriptionType() {
if subscriptionType == service.SubscriptionTypeSubscription {
// 只查询未软删除的订阅,避免通知已取消订阅的用户
rows, err := exec.QueryContext(ctx, "SELECT user_id FROM user_subscriptions WHERE group_id = $1 AND deleted_at IS NULL", id)
if err != nil {
@@ -1020,8 +1078,21 @@ func (r *groupRepository) BindAccountsToGroup(ctx context.Context, groupID int64
return nil
}
tx, err := r.client.Tx(ctx)
if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
return err
}
exec := sqlExecutor(r.client)
if tx != nil {
defer func() { _ = tx.Rollback() }()
exec = tx.Client()
}
if err := lockLiveGroups(ctx, exec, []int64{groupID}); err != nil {
return err
}
// 使用 INSERT ... ON CONFLICT DO NOTHING 忽略已存在的绑定
_, err := r.sql.ExecContext(
_, err = exec.ExecContext(
ctx,
`INSERT INTO account_groups (account_id, group_id, priority, created_at)
SELECT unnest($1::bigint[]), $2, 50, NOW()
@@ -1032,6 +1103,11 @@ func (r *groupRepository) BindAccountsToGroup(ctx context.Context, groupID int64
if err != nil {
return err
}
if tx != nil {
if err := tx.Commit(); err != nil {
return err
}
}
// 发送调度器事件
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupID, nil); err != nil {
@@ -8,10 +8,12 @@ import (
"errors"
"strings"
"testing"
"time"
dbent "github.com/Wei-Shaw/sub2api/ent"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
)
@@ -245,6 +247,103 @@ func (s *GroupRepoSuite) TestDelete() {
s.Require().ErrorIs(err, service.ErrGroupNotFound)
}
func (s *GroupRepoSuite) TestDeleteCascadeIfEmptyRejectsGroupWithNonDeletedAccount() {
group := &service.Group{Name: "guarded-non-empty", Platform: service.PlatformAnthropic, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard}
s.Require().NoError(s.repo.Create(s.ctx, group))
var accountID int64
s.Require().NoError(scanSingleRow(s.ctx, s.tx,
"INSERT INTO accounts (name, platform, type) VALUES ($1, $2, $3) RETURNING id",
[]any{"guarded-account", service.PlatformAnthropic, service.AccountTypeOAuth}, &accountID))
_, err := s.tx.ExecContext(s.ctx, "INSERT INTO account_groups (account_id, group_id, priority, created_at) VALUES ($1, $2, 1, NOW())", accountID, group.ID)
s.Require().NoError(err)
_, err = s.repo.DeleteCascadeIfEmpty(s.ctx, group.ID)
s.Require().ErrorIs(err, service.ErrGroupNotEmpty)
_, err = s.repo.GetByID(s.ctx, group.ID)
s.Require().NoError(err)
var bindings int
s.Require().NoError(scanSingleRow(s.ctx, s.tx, "SELECT COUNT(*) FROM account_groups WHERE group_id = $1", []any{group.ID}, &bindings))
s.Require().Equal(1, bindings)
}
func (s *GroupRepoSuite) TestDeleteCascadeIfEmptyDeletesEmptyGroup() {
group := &service.Group{Name: "guarded-empty", Platform: service.PlatformAnthropic, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard}
s.Require().NoError(s.repo.Create(s.ctx, group))
_, err := s.repo.DeleteCascadeIfEmpty(s.ctx, group.ID)
s.Require().NoError(err)
_, err = s.repo.GetByID(s.ctx, group.ID)
s.Require().ErrorIs(err, service.ErrGroupNotFound)
}
func (s *GroupRepoSuite) TestDeleteCascadeIfEmptyIgnoresBindingsToSoftDeletedAccounts() {
group := &service.Group{Name: "guarded-soft-deleted-account", Platform: service.PlatformAnthropic, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard}
s.Require().NoError(s.repo.Create(s.ctx, group))
var accountID int64
s.Require().NoError(scanSingleRow(s.ctx, s.tx,
"INSERT INTO accounts (name, platform, type, deleted_at) VALUES ($1, $2, $3, NOW()) RETURNING id",
[]any{"guarded-deleted-account", service.PlatformAnthropic, service.AccountTypeOAuth}, &accountID))
_, err := s.tx.ExecContext(s.ctx, "INSERT INTO account_groups (account_id, group_id, priority, created_at) VALUES ($1, $2, 1, NOW())", accountID, group.ID)
s.Require().NoError(err)
_, err = s.repo.DeleteCascadeIfEmpty(s.ctx, group.ID)
s.Require().NoError(err)
_, err = s.repo.GetByID(s.ctx, group.ID)
s.Require().ErrorIs(err, service.ErrGroupNotFound)
}
func TestBindAccountsToGroupWaitingBehindGuardedDeleteCannotCommit(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
var groupID, accountID int64
require.NoError(t, scanSingleRow(ctx, integrationDB,
"INSERT INTO groups (name, platform, rate_multiplier, status, subscription_type) VALUES ($1, $2, 1, $3, $4) RETURNING id",
[]any{"guarded-delete-bind-race", service.PlatformAnthropic, service.StatusActive, service.SubscriptionTypeStandard}, &groupID))
require.NoError(t, scanSingleRow(ctx, integrationDB,
"INSERT INTO accounts (name, platform, type) VALUES ($1, $2, $3) RETURNING id",
[]any{"guarded-delete-bind-race-account", service.PlatformAnthropic, service.AccountTypeOAuth}, &accountID))
t.Cleanup(func() {
_, _ = integrationDB.ExecContext(context.Background(), "DELETE FROM account_groups WHERE group_id = $1 OR account_id = $2", groupID, accountID)
_, _ = integrationDB.ExecContext(context.Background(), "DELETE FROM accounts WHERE id = $1", accountID)
_, _ = integrationDB.ExecContext(context.Background(), "DELETE FROM groups WHERE id = $1", groupID)
})
deleteTx, err := integrationEntClient.Tx(ctx)
require.NoError(t, err)
defer func() { _ = deleteTx.Rollback() }()
rows, err := deleteTx.Client().QueryContext(ctx, "SELECT id FROM groups WHERE id = $1 FOR UPDATE", groupID)
require.NoError(t, err)
require.True(t, rows.Next())
require.NoError(t, rows.Close())
bindDone := make(chan error, 1)
go func() {
bindDone <- newGroupRepositoryWithSQL(integrationEntClient, integrationDB).
BindAccountsToGroup(ctx, groupID, []int64{accountID})
}()
require.Eventually(t, func() bool {
var waiting bool
err := scanSingleRow(ctx, integrationDB, `SELECT EXISTS (
SELECT 1 FROM pg_stat_activity
WHERE query LIKE '/* account_group_live_group_lock */%'
AND wait_event_type = 'Lock'
)`, nil, &waiting)
return err == nil && waiting
}, 5*time.Second, 10*time.Millisecond, "binder did not wait on the guarded deletion lock")
deleteRepo := newGroupRepositoryWithSQL(deleteTx.Client(), deleteTx)
_, err = deleteRepo.DeleteCascadeIfEmpty(ctx, groupID)
require.NoError(t, err)
require.NoError(t, deleteTx.Commit())
require.ErrorIs(t, <-bindDone, service.ErrGroupNotFound)
var bindings int
require.NoError(t, scanSingleRow(ctx, integrationDB, "SELECT COUNT(*) FROM account_groups WHERE group_id = $1", []any{groupID}, &bindings))
require.Zero(t, bindings)
}
// --- List / ListWithFilters ---
func (s *GroupRepoSuite) TestList() {
@@ -619,6 +718,22 @@ func (s *GroupRepoSuite) TestListActive() {
s.Require().True(found, "active1 group should be in results")
}
func (s *GroupRepoSuite) TestListBindableWithFiltersAppliesFilterBeforePagination() {
for _, g := range []*service.Group{
{Name: "bindable-a", Platform: service.PlatformAnthropic, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard},
{Name: "composite-hidden", Platform: service.PlatformComposite, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard},
{Name: "bindable-b", Platform: service.PlatformOpenAI, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard},
} {
s.Require().NoError(s.repo.Create(s.ctx, g))
}
groups, page, err := s.repo.ListBindableWithFilters(s.ctx, pagination.PaginationParams{Page: 2, PageSize: 1, SortBy: "name", SortOrder: "asc"}, "", service.StatusActive, "bindable-", nil)
s.Require().NoError(err)
s.Require().Equal(int64(2), page.Total)
s.Require().Len(groups, 1)
s.Require().NotEqual(service.PlatformComposite, groups[0].Platform)
}
func (s *GroupRepoSuite) TestListActiveByPlatform() {
s.Require().NoError(s.repo.Create(s.ctx, &service.Group{
Name: "g1",
+24 -9
View File
@@ -213,7 +213,7 @@ func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID i
}
// 执行请求
client := httpClientForUpstreamRequest(entry.client, req)
client := s.httpClientForUpstreamRequest(entry.client, req)
client = httpClientWithGrokAccessDeniedFallback(client)
resp, err := servertiming.Do(client, req)
if err != nil {
@@ -277,7 +277,7 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco
return nil, err
}
client := httpClientForUpstreamRequest(entry.client, req)
client := s.httpClientForUpstreamRequest(entry.client, req)
client = httpClientWithGrokAccessDeniedFallback(client)
resp, err := servertiming.Do(client, req)
if err != nil {
@@ -297,15 +297,27 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco
return resp, nil
}
func httpClientForUpstreamRequest(client *http.Client, req *http.Request) *http.Client {
if client == nil || req == nil || !service.HTTPUpstreamRedirectsDisabled(req.Context()) {
// httpClientForUpstreamRequest 按请求上下文的标记派生客户端:禁用重定向,或对重定向的每一跳做主机校验。
// 派生的克隆与缓存客户端共享 Transport;未打标记时原样返回。
func (s *httpUpstreamService) httpClientForUpstreamRequest(client *http.Client, req *http.Request) *http.Client {
if client == nil || req == nil {
return client
}
clone := *client
clone.CheckRedirect = func(*http.Request, []*http.Request) error {
return http.ErrUseLastResponse
ctx := req.Context()
switch {
case service.HTTPUpstreamRedirectsDisabled(ctx):
clone := *client
clone.CheckRedirect = func(*http.Request, []*http.Request) error {
return http.ErrUseLastResponse
}
return &clone
case service.HTTPUpstreamPublicHostsOnly(ctx) && client.CheckRedirect == nil:
clone := *client
clone.CheckRedirect = s.redirectChecker
return &clone
default:
return client
}
return &clone
}
// grokAccessDeniedFallbackTransport preserves the subscription CLI proxy as
@@ -587,8 +599,11 @@ func (s *httpUpstreamService) shouldValidateResolvedIP() bool {
return !s.cfg.Security.URLAllowlist.AllowPrivateHosts
}
// validateRequestHost 校验请求主机的解析结果不落在回环、私网、链路本地或未指定地址。
// 是否全局启用由 security.url_allowlist 决定;带 WithHTTPUpstreamPublicHostsOnly 标记的请求无论配置如何都校验。
func (s *httpUpstreamService) validateRequestHost(req *http.Request) error {
if !s.shouldValidateResolvedIP() {
publicHostsOnly := req != nil && service.HTTPUpstreamPublicHostsOnly(req.Context())
if !s.shouldValidateResolvedIP() && !publicHostsOnly {
return nil
}
if req == nil || req.URL == nil {
@@ -961,3 +961,63 @@ func hasEntry(svc *httpUpstreamService, target *upstreamClientEntry) bool {
}
return false
}
func TestHTTPUpstreamDoPublicHostsOnlyRejectsPrivateDestinationBeforeConnecting(t *testing.T) {
var calls atomic.Int64
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
calls.Add(1)
w.WriteHeader(http.StatusOK)
}))
t.Cleanup(target.Close)
upstream := NewHTTPUpstream(nil)
plain, err := http.NewRequestWithContext(t.Context(), http.MethodGet, target.URL, nil)
require.NoError(t, err)
resp, err := upstream.Do(plain, "", 1, 1)
require.NoError(t, err)
require.NoError(t, resp.Body.Close())
require.Equal(t, int64(1), calls.Load(), "loopback stays reachable for requests without the marker")
guarded, err := http.NewRequestWithContext(service.WithHTTPUpstreamPublicHostsOnly(t.Context()), http.MethodGet, target.URL, nil)
require.NoError(t, err)
resp, err = upstream.Do(guarded, "", 1, 1)
require.Error(t, err)
require.Nil(t, resp)
require.Contains(t, err.Error(), "not allowed")
require.Equal(t, int64(1), calls.Load(), "marked request must be rejected before any connection is made")
}
func TestHTTPUpstreamPublicHostsOnlyValidatesEveryRedirectHop(t *testing.T) {
upstream, ok := NewHTTPUpstream(nil).(*httpUpstreamService)
require.True(t, ok)
base := &http.Client{}
plain, err := http.NewRequestWithContext(t.Context(), http.MethodGet, "https://cdn.example.com/a.png", nil)
require.NoError(t, err)
require.Same(t, base, upstream.httpClientForUpstreamRequest(base, plain))
guarded, err := http.NewRequestWithContext(service.WithHTTPUpstreamPublicHostsOnly(t.Context()), http.MethodGet, "https://cdn.example.com/a.png", nil)
require.NoError(t, err)
client := upstream.httpClientForUpstreamRequest(base, guarded)
require.NotSame(t, base, client)
require.NotNil(t, client.CheckRedirect)
require.Nil(t, base.CheckRedirect, "the cached client must stay untouched")
via := []*http.Request{guarded}
for _, hop := range []string{
"http://127.0.0.1:8080/a.png",
"http://[::1]:8080/a.png",
"http://10.0.0.8/a.png",
"http://169.254.169.254/latest/meta-data/",
"http://0.0.0.0/a.png",
} {
hopReq, err := http.NewRequestWithContext(guarded.Context(), http.MethodGet, hop, nil)
require.NoError(t, err)
require.Error(t, client.CheckRedirect(hopReq, via), "hop=%s", hop)
}
publicHop, err := http.NewRequestWithContext(guarded.Context(), http.MethodGet, "http://93.184.216.34/a.png", nil)
require.NoError(t, err)
require.NoError(t, client.CheckRedirect(publicHop, via))
require.Error(t, client.CheckRedirect(publicHop, make([]*http.Request, 10)), "redirect chain stays capped")
}
@@ -216,6 +216,15 @@ func (c *sessionLimitCache) RegisterSession(ctx context.Context, accountID int64
return result == 1, nil
}
// UnregisterSession 立即移除会话注册(不等待空闲超时)
// 请求最终失败时调用:上游从未服务该会话,继续占槽会卡住 max_sessions 受限的账号
func (c *sessionLimitCache) UnregisterSession(ctx context.Context, accountID int64, sessionUUID string) error {
if sessionUUID == "" {
return nil
}
return c.rdb.ZRem(ctx, sessionLimitKey(accountID), sessionUUID).Err()
}
// RefreshSession 刷新会话时间戳
func (c *sessionLimitCache) RefreshSession(ctx context.Context, accountID int64, sessionUUID string, idleTimeout time.Duration) error {
if sessionUUID == "" {
@@ -0,0 +1,92 @@
//go:build integration
package repository
import (
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/suite"
)
type SessionLimitCacheSuite struct {
IntegrationRedisSuite
cache service.SessionLimitCache
}
func (s *SessionLimitCacheSuite) SetupTest() {
s.IntegrationRedisSuite.SetupTest()
s.cache = NewSessionLimitCache(s.rdb, 5)
}
// TestUnregisterSession_FreesSlotForNewSession 复现并验证修复场景:
// max_sessions=1 的账号上,会话 A 注册后若请求转发失败,会话槽会被 A 占住
// 整个空闲超时窗口(默认 5 分钟),导致窗口内新会话 B 全部被拒。
// UnregisterSession 必须立即释放槽位,让 B 无需等待空闲超时即可注册。
func (s *SessionLimitCacheSuite) TestUnregisterSession_FreesSlotForNewSession() {
const accountID = int64(101)
maxSessions := 1
idleTimeout := 5 * time.Minute
// 会话 A 注册成功(请求开始转发)
allowed, err := s.cache.RegisterSession(s.ctx, accountID, "session-A", maxSessions, idleTimeout)
s.RequireNoError(err, "RegisterSession A")
s.True(allowed, "首个会话应被允许")
// 转发失败后的新会话 B 在空闲超时窗口内被拒(修复前的行为)
allowed, err = s.cache.RegisterSession(s.ctx, accountID, "session-B", maxSessions, idleTimeout)
s.RequireNoError(err, "RegisterSession B before release")
s.False(allowed, "槽位被 A 占用时新会话应被拒绝")
// 请求最终失败:立即释放 A 的会话注册(不等待空闲超时)
s.RequireNoError(s.cache.UnregisterSession(s.ctx, accountID, "session-A"), "UnregisterSession A")
// A 释放后 B 立即可注册,无需等待空闲超时
allowed, err = s.cache.RegisterSession(s.ctx, accountID, "session-B", maxSessions, idleTimeout)
s.RequireNoError(err, "RegisterSession B after release")
s.True(allowed, "释放后新会话应立即可注册")
count, err := s.cache.GetActiveSessionCount(s.ctx, accountID)
s.RequireNoError(err, "GetActiveSessionCount")
s.Equal(1, count, "释放后活跃会话数应为 1(仅 B)")
}
// TestUnregisterSession_IdempotentAndMissing 验证释放操作幂等:
// 重复释放与释放不存在的会话均不应报错。
func (s *SessionLimitCacheSuite) TestUnregisterSession_IdempotentAndMissing() {
const accountID = int64(102)
s.RequireNoError(s.cache.UnregisterSession(s.ctx, accountID, "missing"), "释放不存在的会话")
s.RequireNoError(s.cache.UnregisterSession(s.ctx, accountID, "missing"), "重复释放")
s.RequireNoError(s.cache.UnregisterSession(s.ctx, accountID, ""), "空 sessionUUID 应为 no-op")
}
// TestUnregisterSession_OnlyTargetsSpecifiedSession 验证释放只影响目标会话,
// 不影响同账号上的其他活跃会话。
func (s *SessionLimitCacheSuite) TestUnregisterSession_OnlyTargetsSpecifiedSession() {
const accountID = int64(103)
maxSessions := 2
idleTimeout := 5 * time.Minute
for _, sid := range []string{"session-A", "session-B"} {
allowed, err := s.cache.RegisterSession(s.ctx, accountID, sid, maxSessions, idleTimeout)
s.RequireNoError(err, "RegisterSession "+sid)
s.True(allowed)
}
s.RequireNoError(s.cache.UnregisterSession(s.ctx, accountID, "session-A"), "UnregisterSession A")
active, err := s.cache.IsSessionActive(s.ctx, accountID, "session-A")
s.RequireNoError(err, "IsSessionActive A")
s.False(active, "A 应已释放")
active, err = s.cache.IsSessionActive(s.ctx, accountID, "session-B")
s.RequireNoError(err, "IsSessionActive B")
s.True(active, "B 不应受影响")
}
func TestSessionLimitCacheSuite(t *testing.T) {
suite.Run(t, new(SessionLimitCacheSuite))
}