mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
shaw
80a229bce5
修复 PR #3768 批量图像 MVP 合并后审计报告中的全部问题： 结算与计费（高危）： - 所有 SETTLEMENT_* 失败（超冻结/计数非法/manifest 冲突/定价缺失/扣费失败） 统一计入 retry_count 并在耗尽时释放冻结转 failed，消灭 settling 无限 requeue 导致的冻结余额永久锁死 - 耗尽出口的释放指纹统一为 RequestHash，与 processor/Cancel/recovery 一致； release 遇同 request id 指纹冲突视为幂等成功，治愈历史毒消息 - 管理端校验 hold_multiplier >= discount_multiplier，定价快照对存量脏数据钳制 - 释放前校验 per-job hold claim（dedup+归档表），杜绝幻影释放 索引对账（高危）： - provider 输出与提交 custom_id 集对账：未知条目丢弃并记事件， 漏项补 PROVIDER_RESULT_MISSING 失败行，保证 success+fail == item_count 提交与恢复（高危）： - 提交前转 uploading 并在 provider.Submit 期间心跳刷新 updated_at； 恢复扫描改为原子复核（FailStaleUnsubmittedBatchImageJob）， 消灭慢提交被误杀退款而上游任务照常计费的孤儿场景 - 上游任务创建成功但本地状态推进失败时，尽力取消上游并清理输入 - recovery 释放失败时入队交由 worker releaseTerminalHold 兜底重试 队列与并发（中危）： - Enqueue（SetNX+LPush）与 Reserve（BRPop+ZAdd）均改为 Lua 原子脚本， 消灭崩溃窗口导致 job 脱离队列、被 7 天 inflight 键锁死 - 锁冲突按 LockConflictDelay 重新入队（原直接丢弃需等 10 分钟 stale 恢复） - 处理期间心跳：active zset 续期（ZAddXX 防幽灵成员）+ 锁 TTL 续期 - ReplaceBatchImageItemsForJob 增加 indexing 状态守卫，防掉队 worker 重写账目 存量回归（中危）： - image-only 定价条目（仅图片价无 token 价）恢复 token 计费 fail-closed， 不再按 $0 计费；图片计费路径不受影响 - 鉴权余额门槛恢复 balance <= 0 语义，MinimumBalanceReserve 不再作硬 403 加固： - ZIP max_items 钳制到管理员上限；Submit 补齐 Platform==Gemini 校验； gemini downloadUri 跟随前做 host 白名单校验 - 批量客户端改用共享 httpclient（拨号/TLS/响应头超时有界） - 审计点名的忽略错误（MarkDownloaded/SettlementFailed/AppendEvent 等）改为记日志
832 lines
28 KiB
Go
832 lines
28 KiB
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type fakeAPIKeyRepo struct {
|
|
getByKey func(ctx context.Context, key string) (*service.APIKey, error)
|
|
updateLastUsed func(ctx context.Context, id int64, usedAt time.Time) error
|
|
}
|
|
|
|
type fakeGoogleSubscriptionRepo struct {
|
|
getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error)
|
|
updateStatus func(ctx context.Context, subscriptionID int64, status string) error
|
|
activateWindow func(ctx context.Context, id int64, start time.Time) error
|
|
resetDaily func(ctx context.Context, id int64, start time.Time) error
|
|
resetWeekly func(ctx context.Context, id int64, start time.Time) error
|
|
resetMonthly func(ctx context.Context, id int64, start time.Time) error
|
|
}
|
|
|
|
func (f fakeAPIKeyRepo) Create(ctx context.Context, key *service.APIKey) error {
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) GetByID(ctx context.Context, id int64) (*service.APIKey, error) {
|
|
return nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) GetKeyAndOwnerID(ctx context.Context, id int64) (string, int64, error) {
|
|
return "", 0, errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) GetByKey(ctx context.Context, key string) (*service.APIKey, error) {
|
|
if f.getByKey == nil {
|
|
return nil, errors.New("unexpected call")
|
|
}
|
|
return f.getByKey(ctx, key)
|
|
}
|
|
func (f fakeAPIKeyRepo) GetByKeyForAuth(ctx context.Context, key string) (*service.APIKey, error) {
|
|
return f.GetByKey(ctx, key)
|
|
}
|
|
func (f fakeAPIKeyRepo) Update(ctx context.Context, key *service.APIKey) error {
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) Delete(ctx context.Context, id int64) error {
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) DeleteWithAudit(ctx context.Context, id int64) error {
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams, _ service.APIKeyListFilters) ([]service.APIKey, *pagination.PaginationResult, error) {
|
|
return nil, nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
|
|
return nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) CountByUserID(ctx context.Context, userID int64) (int64, error) {
|
|
return 0, errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) ExistsByKey(ctx context.Context, key string) (bool, error) {
|
|
return false, errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.APIKey, *pagination.PaginationResult, error) {
|
|
return nil, nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) SearchAPIKeys(ctx context.Context, userID int64, keyword string, limit int) ([]service.APIKey, error) {
|
|
return nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) ClearGroupIDByGroupID(ctx context.Context, groupID int64) (int64, error) {
|
|
return 0, errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) CountByGroupID(ctx context.Context, groupID int64) (int64, error) {
|
|
return 0, errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) ListKeysByUserID(ctx context.Context, userID int64) ([]string, error) {
|
|
return nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) {
|
|
return nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error) {
|
|
return 0, errors.New("not implemented")
|
|
}
|
|
func (f fakeAPIKeyRepo) UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error {
|
|
if f.updateLastUsed != nil {
|
|
return f.updateLastUsed(ctx, id, usedAt)
|
|
}
|
|
return nil
|
|
}
|
|
func (f fakeAPIKeyRepo) IncrementRateLimitUsage(ctx context.Context, id int64, cost float64) error {
|
|
return nil
|
|
}
|
|
func (f fakeAPIKeyRepo) ResetRateLimitWindows(ctx context.Context, id int64) error {
|
|
return nil
|
|
}
|
|
func (f fakeAPIKeyRepo) GetRateLimitData(ctx context.Context, id int64) (*service.APIKeyRateLimitData, error) {
|
|
return &service.APIKeyRateLimitData{}, nil
|
|
}
|
|
func (f fakeAPIKeyRepo) UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (int64, error) {
|
|
return 0, errors.New("not implemented")
|
|
}
|
|
|
|
func (f fakeGoogleSubscriptionRepo) Create(ctx context.Context, sub *service.UserSubscription) error {
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) GetByID(ctx context.Context, id int64) (*service.UserSubscription, error) {
|
|
return nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.UserSubscription, error) {
|
|
return nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) GetByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
|
|
return nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) GetActiveByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
|
|
if f.getActive != nil {
|
|
return f.getActive(ctx, userID, groupID)
|
|
}
|
|
return nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) Update(ctx context.Context, sub *service.UserSubscription) error {
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) Delete(ctx context.Context, id int64) error {
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) Restore(ctx context.Context, subscriptionID int64, restoredStatus string) (*service.UserSubscription, error) {
|
|
return nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) ListByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
|
|
return nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) ListActiveByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
|
|
return nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
|
return nil, nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, platform, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
|
return nil, nil, errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) ExistsByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
|
|
return false, errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) ExistsActiveByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
|
|
return false, errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) ExtendExpiry(ctx context.Context, subscriptionID int64, newExpiresAt time.Time) error {
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) UpdateStatus(ctx context.Context, subscriptionID int64, status string) error {
|
|
if f.updateStatus != nil {
|
|
return f.updateStatus(ctx, subscriptionID, status)
|
|
}
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) UpdateNotes(ctx context.Context, subscriptionID int64, notes string) error {
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, start time.Time) error {
|
|
if f.activateWindow != nil {
|
|
return f.activateWindow(ctx, id, start)
|
|
}
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, start time.Time) error {
|
|
if f.resetDaily != nil {
|
|
return f.resetDaily(ctx, id, start)
|
|
}
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) ResetWeeklyUsage(ctx context.Context, id int64, start time.Time) error {
|
|
if f.resetWeekly != nil {
|
|
return f.resetWeekly(ctx, id, start)
|
|
}
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) ResetMonthlyUsage(ctx context.Context, id int64, start time.Time) error {
|
|
if f.resetMonthly != nil {
|
|
return f.resetMonthly(ctx, id, start)
|
|
}
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) IncrementUsage(ctx context.Context, id int64, costUSD float64) error {
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) BatchUpdateExpiredStatus(ctx context.Context) (int64, error) {
|
|
return 0, errors.New("not implemented")
|
|
}
|
|
|
|
type googleErrorResponse struct {
|
|
Error struct {
|
|
Code int `json:"code"`
|
|
Message string `json:"message"`
|
|
Status string `json:"status"`
|
|
} `json:"error"`
|
|
}
|
|
|
|
func newTestAPIKeyService(repo service.APIKeyRepository) *service.APIKeyService {
|
|
return service.NewAPIKeyService(
|
|
repo,
|
|
nil, // userRepo (unused in GetByKey)
|
|
nil, // groupRepo
|
|
nil, // userSubRepo
|
|
nil, // userGroupRateRepo
|
|
nil, // cache
|
|
&config.Config{},
|
|
)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_MissingKey(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
r := gin.New()
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
return nil, errors.New("should not be called")
|
|
},
|
|
})
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusUnauthorized, rec.Code)
|
|
var resp googleErrorResponse
|
|
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
|
require.Equal(t, http.StatusUnauthorized, resp.Error.Code)
|
|
require.Equal(t, "API key is required", resp.Error.Message)
|
|
require.Equal(t, "UNAUTHENTICATED", resp.Error.Status)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_QueryApiKeyRejected(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
r := gin.New()
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
return nil, errors.New("should not be called")
|
|
},
|
|
})
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test?api_key=legacy", nil)
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusBadRequest, rec.Code)
|
|
var resp googleErrorResponse
|
|
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
|
require.Equal(t, http.StatusBadRequest, resp.Error.Code)
|
|
require.Equal(t, "Query parameter api_key is deprecated. Use Authorization header or key instead.", resp.Error.Message)
|
|
require.Equal(t, "INVALID_ARGUMENT", resp.Error.Status)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogleSetsGroupContext(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
group := &service.Group{
|
|
ID: 99,
|
|
Name: "g1",
|
|
Status: service.StatusActive,
|
|
Platform: service.PlatformGemini,
|
|
Hydrated: true,
|
|
}
|
|
user := &service.User{
|
|
ID: 7,
|
|
Role: service.RoleUser,
|
|
Status: service.StatusActive,
|
|
Balance: 10,
|
|
Concurrency: 3,
|
|
}
|
|
apiKey := &service.APIKey{
|
|
ID: 100,
|
|
UserID: user.ID,
|
|
Key: "test-key",
|
|
Status: service.StatusActive,
|
|
User: user,
|
|
Group: group,
|
|
}
|
|
apiKey.GroupID = &group.ID
|
|
|
|
apiKeyService := service.NewAPIKeyService(
|
|
fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
if key != apiKey.Key {
|
|
return nil, service.ErrAPIKeyNotFound
|
|
}
|
|
clone := *apiKey
|
|
return &clone, nil
|
|
},
|
|
},
|
|
nil,
|
|
nil,
|
|
nil,
|
|
nil,
|
|
nil,
|
|
&config.Config{RunMode: config.RunModeSimple},
|
|
)
|
|
|
|
cfg := &config.Config{RunMode: config.RunModeSimple}
|
|
r := gin.New()
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
|
r.GET("/v1beta/test", func(c *gin.Context) {
|
|
groupFromCtx, ok := c.Request.Context().Value(ctxkey.Group).(*service.Group)
|
|
if !ok || groupFromCtx == nil || groupFromCtx.ID != group.ID {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"ok": false})
|
|
return
|
|
}
|
|
c.JSON(http.StatusOK, gin.H{"ok": true})
|
|
})
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
req.Header.Set("x-api-key", apiKey.Key)
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusOK, rec.Code)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_QueryKeyAllowedOnV1Beta(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
r := gin.New()
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
return &service.APIKey{
|
|
ID: 1,
|
|
Key: key,
|
|
Status: service.StatusActive,
|
|
User: &service.User{
|
|
ID: 123,
|
|
Status: service.StatusActive,
|
|
},
|
|
}, nil
|
|
},
|
|
})
|
|
cfg := &config.Config{RunMode: config.RunModeSimple}
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test?key=valid", nil)
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusOK, rec.Code)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_InvalidKey(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
r := gin.New()
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
return nil, service.ErrAPIKeyNotFound
|
|
},
|
|
})
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
req.Header.Set("Authorization", "Bearer invalid")
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusUnauthorized, rec.Code)
|
|
var resp googleErrorResponse
|
|
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
|
require.Equal(t, http.StatusUnauthorized, resp.Error.Code)
|
|
require.Equal(t, "Invalid API key", resp.Error.Message)
|
|
require.Equal(t, "UNAUTHENTICATED", resp.Error.Status)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_MarksUnavailableGroupBusinessLimited(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
groupID := int64(101)
|
|
user := &service.User{
|
|
ID: 7,
|
|
Role: service.RoleUser,
|
|
Status: service.StatusActive,
|
|
Balance: 10,
|
|
Concurrency: 3,
|
|
}
|
|
apiKey := &service.APIKey{
|
|
ID: 100,
|
|
UserID: user.ID,
|
|
GroupID: &groupID,
|
|
Key: "google-group-deleted",
|
|
Status: service.StatusActive,
|
|
User: user,
|
|
Group: &service.Group{
|
|
ID: groupID,
|
|
Name: "deleted",
|
|
Status: "deleted",
|
|
Platform: service.PlatformGemini,
|
|
Hydrated: true,
|
|
},
|
|
}
|
|
|
|
r := gin.New()
|
|
var markedBusinessLimited bool
|
|
var businessLimitedReason string
|
|
r.Use(func(c *gin.Context) {
|
|
c.Next()
|
|
markedBusinessLimited = service.HasOpsClientBusinessLimited(c)
|
|
if v, ok := c.Get(service.OpsClientBusinessLimitedReasonKey); ok {
|
|
businessLimitedReason, _ = v.(string)
|
|
}
|
|
})
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
if key != apiKey.Key {
|
|
return nil, service.ErrAPIKeyNotFound
|
|
}
|
|
clone := *apiKey
|
|
return &clone, nil
|
|
},
|
|
})
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{RunMode: config.RunModeSimple}))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
req.Header.Set("x-goog-api-key", apiKey.Key)
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusForbidden, rec.Code)
|
|
var resp googleErrorResponse
|
|
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
|
require.Equal(t, "API Key 所属分组已删除", resp.Error.Message)
|
|
require.True(t, markedBusinessLimited)
|
|
require.Equal(t, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable, businessLimitedReason)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_RepoError(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
r := gin.New()
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
return nil, errors.New("db down")
|
|
},
|
|
})
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
req.Header.Set("Authorization", "Bearer any")
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusInternalServerError, rec.Code)
|
|
var resp googleErrorResponse
|
|
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
|
require.Equal(t, http.StatusInternalServerError, resp.Error.Code)
|
|
require.Equal(t, "Failed to validate API key", resp.Error.Message)
|
|
require.Equal(t, "INTERNAL", resp.Error.Status)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_DisabledKey(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
r := gin.New()
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
return &service.APIKey{
|
|
ID: 1,
|
|
Key: key,
|
|
Status: service.StatusDisabled,
|
|
User: &service.User{
|
|
ID: 123,
|
|
Status: service.StatusActive,
|
|
},
|
|
}, nil
|
|
},
|
|
})
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
req.Header.Set("Authorization", "Bearer disabled")
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusUnauthorized, rec.Code)
|
|
var resp googleErrorResponse
|
|
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
|
require.Equal(t, http.StatusUnauthorized, resp.Error.Code)
|
|
require.Equal(t, "API key is disabled", resp.Error.Message)
|
|
require.Equal(t, "UNAUTHENTICATED", resp.Error.Status)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_InsufficientBalance(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
r := gin.New()
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
return &service.APIKey{
|
|
ID: 1,
|
|
Key: key,
|
|
Status: service.StatusActive,
|
|
User: &service.User{
|
|
ID: 123,
|
|
Status: service.StatusActive,
|
|
Balance: 0,
|
|
},
|
|
}, nil
|
|
},
|
|
})
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
req.Header.Set("Authorization", "Bearer ok")
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusForbidden, rec.Code)
|
|
var resp googleErrorResponse
|
|
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
|
require.Equal(t, http.StatusForbidden, resp.Error.Code)
|
|
require.Equal(t, "Insufficient account balance", resp.Error.Message)
|
|
require.Equal(t, "PERMISSION_DENIED", resp.Error.Status)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_BalanceBelowMinimumReserve(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
// 鉴权层保持历史语义：MinimumBalanceReserve 只用于 billing-cache 预检，
|
|
// 0 < balance < reserve 的用户不得在鉴权中间件被硬 403。
|
|
r := gin.New()
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
return &service.APIKey{
|
|
ID: 1,
|
|
Key: key,
|
|
Status: service.StatusActive,
|
|
User: &service.User{
|
|
ID: 123,
|
|
Status: service.StatusActive,
|
|
Balance: 0.005,
|
|
},
|
|
}, nil
|
|
},
|
|
})
|
|
cfg := &config.Config{}
|
|
cfg.Billing.MinimumBalanceReserve = 0.01
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
req.Header.Set("Authorization", "Bearer ok")
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusOK, rec.Code)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_RejectsExhaustedBalance(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
r := gin.New()
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
return &service.APIKey{
|
|
ID: 1,
|
|
Key: key,
|
|
Status: service.StatusActive,
|
|
User: &service.User{
|
|
ID: 123,
|
|
Status: service.StatusActive,
|
|
Balance: 0,
|
|
},
|
|
}, nil
|
|
},
|
|
})
|
|
cfg := &config.Config{}
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
req.Header.Set("Authorization", "Bearer ok")
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusForbidden, rec.Code)
|
|
var resp googleErrorResponse
|
|
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
|
require.Equal(t, http.StatusForbidden, resp.Error.Code)
|
|
require.Equal(t, "Insufficient account balance", resp.Error.Message)
|
|
require.Equal(t, "PERMISSION_DENIED", resp.Error.Status)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_TouchesLastUsedOnSuccess(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
user := &service.User{
|
|
ID: 11,
|
|
Role: service.RoleUser,
|
|
Status: service.StatusActive,
|
|
Balance: 10,
|
|
Concurrency: 3,
|
|
}
|
|
apiKey := &service.APIKey{
|
|
ID: 201,
|
|
UserID: user.ID,
|
|
Key: "google-touch-ok",
|
|
Status: service.StatusActive,
|
|
User: user,
|
|
}
|
|
|
|
var touchedID int64
|
|
var touchedAt time.Time
|
|
r := gin.New()
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
if key != apiKey.Key {
|
|
return nil, service.ErrAPIKeyNotFound
|
|
}
|
|
clone := *apiKey
|
|
return &clone, nil
|
|
},
|
|
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
|
|
touchedID = id
|
|
touchedAt = usedAt
|
|
return nil
|
|
},
|
|
})
|
|
cfg := &config.Config{RunMode: config.RunModeSimple}
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
req.Header.Set("x-goog-api-key", apiKey.Key)
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusOK, rec.Code)
|
|
require.Equal(t, apiKey.ID, touchedID)
|
|
require.False(t, touchedAt.IsZero())
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_TouchFailureDoesNotBlock(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
user := &service.User{
|
|
ID: 12,
|
|
Role: service.RoleUser,
|
|
Status: service.StatusActive,
|
|
Balance: 10,
|
|
Concurrency: 3,
|
|
}
|
|
apiKey := &service.APIKey{
|
|
ID: 202,
|
|
UserID: user.ID,
|
|
Key: "google-touch-fail",
|
|
Status: service.StatusActive,
|
|
User: user,
|
|
}
|
|
|
|
touchCalls := 0
|
|
r := gin.New()
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
if key != apiKey.Key {
|
|
return nil, service.ErrAPIKeyNotFound
|
|
}
|
|
clone := *apiKey
|
|
return &clone, nil
|
|
},
|
|
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
|
|
touchCalls++
|
|
return errors.New("write failed")
|
|
},
|
|
})
|
|
cfg := &config.Config{RunMode: config.RunModeSimple}
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
req.Header.Set("x-goog-api-key", apiKey.Key)
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusOK, rec.Code)
|
|
require.Equal(t, 1, touchCalls)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_TouchesLastUsedInStandardMode(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
user := &service.User{
|
|
ID: 13,
|
|
Role: service.RoleUser,
|
|
Status: service.StatusActive,
|
|
Balance: 10,
|
|
Concurrency: 3,
|
|
}
|
|
apiKey := &service.APIKey{
|
|
ID: 203,
|
|
UserID: user.ID,
|
|
Key: "google-touch-standard",
|
|
Status: service.StatusActive,
|
|
User: user,
|
|
}
|
|
|
|
touchCalls := 0
|
|
r := gin.New()
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
if key != apiKey.Key {
|
|
return nil, service.ErrAPIKeyNotFound
|
|
}
|
|
clone := *apiKey
|
|
return &clone, nil
|
|
},
|
|
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
|
|
touchCalls++
|
|
return nil
|
|
},
|
|
})
|
|
cfg := &config.Config{RunMode: config.RunModeStandard}
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
req.Header.Set("Authorization", "Bearer "+apiKey.Key)
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusOK, rec.Code)
|
|
require.Equal(t, 1, touchCalls)
|
|
}
|
|
|
|
func TestApiKeyAuthWithSubscriptionGoogle_SubscriptionLimitExceededReturns429(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
limit := 1.0
|
|
group := &service.Group{
|
|
ID: 77,
|
|
Name: "gemini-sub",
|
|
Status: service.StatusActive,
|
|
Platform: service.PlatformGemini,
|
|
Hydrated: true,
|
|
SubscriptionType: service.SubscriptionTypeSubscription,
|
|
DailyLimitUSD: &limit,
|
|
}
|
|
user := &service.User{
|
|
ID: 999,
|
|
Role: service.RoleUser,
|
|
Status: service.StatusActive,
|
|
Balance: 10,
|
|
Concurrency: 3,
|
|
}
|
|
apiKey := &service.APIKey{
|
|
ID: 501,
|
|
UserID: user.ID,
|
|
Key: "google-sub-limit",
|
|
Status: service.StatusActive,
|
|
User: user,
|
|
Group: group,
|
|
}
|
|
apiKey.GroupID = &group.ID
|
|
|
|
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
|
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
|
if key != apiKey.Key {
|
|
return nil, service.ErrAPIKeyNotFound
|
|
}
|
|
clone := *apiKey
|
|
return &clone, nil
|
|
},
|
|
})
|
|
|
|
now := time.Now()
|
|
sub := &service.UserSubscription{
|
|
ID: 601,
|
|
UserID: user.ID,
|
|
GroupID: group.ID,
|
|
Status: service.SubscriptionStatusActive,
|
|
ExpiresAt: now.Add(24 * time.Hour),
|
|
DailyWindowStart: &now,
|
|
DailyUsageUSD: 10,
|
|
}
|
|
subscriptionService := service.NewSubscriptionService(nil, fakeGoogleSubscriptionRepo{
|
|
getActive: func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
|
|
if userID != user.ID || groupID != group.ID {
|
|
return nil, service.ErrSubscriptionNotFound
|
|
}
|
|
clone := *sub
|
|
return &clone, nil
|
|
},
|
|
updateStatus: func(ctx context.Context, subscriptionID int64, status string) error { return nil },
|
|
activateWindow: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
|
resetDaily: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
|
resetWeekly: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
|
resetMonthly: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
|
}, nil, nil, &config.Config{RunMode: config.RunModeStandard})
|
|
|
|
r := gin.New()
|
|
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, &config.Config{RunMode: config.RunModeStandard}))
|
|
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
|
req.Header.Set("x-goog-api-key", apiKey.Key)
|
|
rec := httptest.NewRecorder()
|
|
r.ServeHTTP(rec, req)
|
|
|
|
require.Equal(t, http.StatusTooManyRequests, rec.Code)
|
|
var resp googleErrorResponse
|
|
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
|
require.Equal(t, http.StatusTooManyRequests, resp.Error.Code)
|
|
require.Equal(t, "RESOURCE_EXHAUSTED", resp.Error.Status)
|
|
require.Contains(t, resp.Error.Message, "daily usage limit exceeded")
|
|
}
|