mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-11-03 00:33:13 +00:00
fix(groups): merge main and resolve model allowlist conflicts
This commit is contained in:
252 files changed
+11371
-1598
No files matched your search
@@ -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))
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
Reference in new issue
Block a user