mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
017bbcb901 | ||
|
|
9ecb340822 | ||
|
|
a0f41f95a0 | ||
|
|
a60a29549f | ||
|
|
b31b435091 |
@@ -1 +1 @@
|
||||
0.2.9
|
||||
0.2.10
|
||||
|
||||
@@ -89,6 +89,7 @@ type Config struct {
|
||||
Pricing PricingConfig `mapstructure:"pricing"`
|
||||
Gateway GatewayConfig `mapstructure:"gateway"`
|
||||
APIKeyAuth APIKeyAuthCacheConfig `mapstructure:"api_key_auth_cache"`
|
||||
APIKeyCreate APIKeyCreateConfig `mapstructure:"api_key_create"`
|
||||
SubscriptionCache SubscriptionCacheConfig `mapstructure:"subscription_cache"`
|
||||
SubscriptionMaintenance SubscriptionMaintenanceConfig `mapstructure:"subscription_maintenance"`
|
||||
Dashboard DashboardCacheConfig `mapstructure:"dashboard_cache"`
|
||||
@@ -1696,6 +1697,14 @@ type APIKeyAuthCacheConfig struct {
|
||||
InvalidAbuse InvalidAuthAbuseConfig `mapstructure:"invalid_abuse"`
|
||||
}
|
||||
|
||||
// APIKeyCreateConfig 用户创建 API Key 的防滥用限制(0 表示不限制)
|
||||
type APIKeyCreateConfig struct {
|
||||
// MaxActivePerUser 单个用户同时存在(未删除)的 API Key 上限
|
||||
MaxActivePerUser int `mapstructure:"max_active_per_user"`
|
||||
// MaxPerUserPerHour 单个用户每小时可创建的 API Key 次数(删除不返还次数)
|
||||
MaxPerUserPerHour int `mapstructure:"max_per_user_per_hour"`
|
||||
}
|
||||
|
||||
type InvalidAuthAbuseConfig struct {
|
||||
Enabled bool `mapstructure:"enabled"`
|
||||
Threshold int `mapstructure:"threshold"`
|
||||
@@ -2332,6 +2341,8 @@ func setDefaults() {
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.window_seconds", 60)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.block_seconds", 60)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.capacity", 16384)
|
||||
viper.SetDefault("api_key_create.max_active_per_user", 200)
|
||||
viper.SetDefault("api_key_create.max_per_user_per_hour", 60)
|
||||
|
||||
// Subscription auth L1 cache
|
||||
viper.SetDefault("subscription_cache.l1_size", 16384)
|
||||
@@ -2711,6 +2722,12 @@ func (c *Config) Validate() error {
|
||||
return fmt.Errorf("server.h2c.max_upload_buffer_per_stream must be positive")
|
||||
}
|
||||
}
|
||||
if c.APIKeyCreate.MaxActivePerUser < 0 {
|
||||
return fmt.Errorf("api_key_create.max_active_per_user must be non-negative")
|
||||
}
|
||||
if c.APIKeyCreate.MaxPerUserPerHour < 0 {
|
||||
return fmt.Errorf("api_key_create.max_per_user_per_hour must be non-negative")
|
||||
}
|
||||
if c.APIKeyAuth.InvalidAbuse.Enabled {
|
||||
if c.APIKeyAuth.InvalidAbuse.Threshold < 10 {
|
||||
return fmt.Errorf("api_key_auth_cache.invalid_abuse.threshold must be at least 10")
|
||||
|
||||
@@ -101,8 +101,11 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
// 解析渠道级模型映射
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, reqModel)
|
||||
|
||||
// Claude Code only restriction
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly {
|
||||
// Claude Code only restriction: /v1/chat/completions is never a Claude Code
|
||||
// endpoint. With a fallback group the request continues and account selection
|
||||
// (checkClaudeCodeRestriction) schedules it in the fallback group; without one
|
||||
// it is rejected here.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly && apiKey.Group.FallbackGroupID == nil {
|
||||
h.chatCompletionsErrorResponse(c, http.StatusForbidden, "permission_error",
|
||||
"This group is restricted to Claude Code clients (/v1/messages only)")
|
||||
return
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// groupScopedSchedulerCache 按桶的分组只返回该分组的成员账号,并记录被查询过的分组。
|
||||
type groupScopedSchedulerCache struct {
|
||||
*fakeSchedulerCache
|
||||
mu sync.Mutex
|
||||
groupIDs []int64
|
||||
}
|
||||
|
||||
func (c *groupScopedSchedulerCache) GetSnapshot(_ context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) {
|
||||
c.mu.Lock()
|
||||
c.groupIDs = append(c.groupIDs, bucket.GroupID)
|
||||
c.mu.Unlock()
|
||||
var members []*service.Account
|
||||
for _, account := range c.accounts {
|
||||
for _, ag := range account.AccountGroups {
|
||||
if ag.GroupID == bucket.GroupID {
|
||||
members = append(members, account)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return members, true, nil
|
||||
}
|
||||
|
||||
func (c *groupScopedSchedulerCache) queriedGroupIDs() []int64 {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return append([]int64(nil), c.groupIDs...)
|
||||
}
|
||||
|
||||
type groupMapRepo struct {
|
||||
*fakeGroupRepo
|
||||
groups map[int64]*service.Group
|
||||
}
|
||||
|
||||
func (r *groupMapRepo) GetByID(_ context.Context, id int64) (*service.Group, error) {
|
||||
if group, ok := r.groups[id]; ok {
|
||||
return group, nil
|
||||
}
|
||||
return nil, service.ErrGroupNotFound
|
||||
}
|
||||
|
||||
func (r *groupMapRepo) GetByIDLite(ctx context.Context, id int64) (*service.Group, error) {
|
||||
return r.GetByID(ctx, id)
|
||||
}
|
||||
|
||||
func TestGatewayOpenAICompatibleHandlersClaudeCodeOnlyFallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
const (
|
||||
primaryGroupID = int64(9300)
|
||||
fallbackGroupID = int64(9301)
|
||||
primaryAccountID = int64(9310)
|
||||
fallbackAccountID = int64(9311)
|
||||
)
|
||||
|
||||
// 账号不带 api_key:转发在取令牌时失败,请求不会触达上游,断言只看选号结果。
|
||||
newAccount := func(id, groupID int64) *service.Account {
|
||||
return &service.Account{
|
||||
ID: id, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive, Schedulable: true, Concurrency: 1,
|
||||
AccountGroups: []service.AccountGroup{{AccountID: id, GroupID: groupID}},
|
||||
}
|
||||
}
|
||||
|
||||
endpoints := []struct {
|
||||
name string
|
||||
path string
|
||||
body string
|
||||
call func(*GatewayHandler, *gin.Context)
|
||||
}{
|
||||
{
|
||||
name: "responses", path: "/v1/responses",
|
||||
body: `{"model":"claude-sonnet-4-5","input":"hello","stream":false}`,
|
||||
call: (*GatewayHandler).Responses,
|
||||
},
|
||||
{
|
||||
name: "chat completions", path: "/v1/chat/completions",
|
||||
body: `{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hello"}],"stream":false}`,
|
||||
call: (*GatewayHandler).ChatCompletions,
|
||||
},
|
||||
}
|
||||
|
||||
for _, ep := range endpoints {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
hasFallback bool
|
||||
}{
|
||||
{name: "with fallback group", hasFallback: true},
|
||||
{name: "without fallback group", hasFallback: false},
|
||||
} {
|
||||
t.Run(ep.name+"/"+tc.name, func(t *testing.T) {
|
||||
primary := &service.Group{
|
||||
ID: primaryGroupID, Hydrated: true, Platform: service.PlatformAnthropic,
|
||||
Status: service.StatusActive, ClaudeCodeOnly: true,
|
||||
}
|
||||
if tc.hasFallback {
|
||||
fallbackID := fallbackGroupID
|
||||
primary.FallbackGroupID = &fallbackID
|
||||
}
|
||||
fallback := &service.Group{
|
||||
ID: fallbackGroupID, Hydrated: true, Platform: service.PlatformAnthropic,
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
|
||||
schedulerCache := &groupScopedSchedulerCache{fakeSchedulerCache: &fakeSchedulerCache{accounts: []*service.Account{
|
||||
newAccount(primaryAccountID, primaryGroupID),
|
||||
newAccount(fallbackAccountID, fallbackGroupID),
|
||||
}}}
|
||||
gatewayService := service.NewGatewayService(
|
||||
nil, &groupMapRepo{fakeGroupRepo: &fakeGroupRepo{}, groups: map[int64]*service.Group{
|
||||
primaryGroupID: primary,
|
||||
fallbackGroupID: fallback,
|
||||
}}, nil, nil, nil, nil, nil, nil, nil,
|
||||
service.NewSchedulerSnapshotService(schedulerCache, nil, nil, nil, nil),
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
billingCacheService := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billingCacheService.Stop)
|
||||
h := &GatewayHandler{
|
||||
gatewayService: gatewayService,
|
||||
billingCacheService: billingCacheService,
|
||||
concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&fakeConcurrencyCache{}), SSEPingFormatClaude, 0),
|
||||
maxAccountSwitches: 1,
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
primaryGroupIDRef := primaryGroupID
|
||||
apiKey := &service.APIKey{
|
||||
ID: 9320, UserID: 9330, GroupID: &primaryGroupIDRef, Group: primary, Status: service.StatusActive,
|
||||
User: &service.User{ID: 9330, Concurrency: 10, Balance: 100},
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
ctx := context.WithValue(context.Background(), ctxkey.Group, primary)
|
||||
req := httptest.NewRequest(http.MethodPost, ep.path, bytes.NewBufferString(ep.body)).WithContext(ctx)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.Request = req
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10})
|
||||
|
||||
ep.call(h, c)
|
||||
|
||||
selected, reachedSelection := c.Get(opsAccountIDKey)
|
||||
if !tc.hasFallback {
|
||||
require.Equal(t, http.StatusForbidden, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "This group is restricted to Claude Code clients")
|
||||
require.False(t, reachedSelection, "a Claude Code only group without fallback must be rejected before account selection")
|
||||
require.Empty(t, schedulerCache.queriedGroupIDs())
|
||||
return
|
||||
}
|
||||
require.NotContains(t, recorder.Body.String(), "restricted to Claude Code clients")
|
||||
require.True(t, reachedSelection, "a Claude Code only group with fallback must reach account selection")
|
||||
require.Equal(t, fallbackAccountID, selected)
|
||||
queried := schedulerCache.queriedGroupIDs()
|
||||
require.Contains(t, queried, fallbackGroupID)
|
||||
require.NotContains(t, queried, primaryGroupID)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -105,13 +105,11 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
// 解析渠道级模型映射
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(requestCtx, apiKey.GroupID, reqModel)
|
||||
|
||||
// Claude Code only restriction:
|
||||
// /v1/responses is never a Claude Code endpoint.
|
||||
// When claude_code_only is enabled, this endpoint is rejected.
|
||||
// The existing service-layer checkClaudeCodeRestriction handles degradation
|
||||
// to fallback groups when the Forward path calls SelectAccountForModelWithExclusions.
|
||||
// Here we just reject at handler level since /v1/responses clients can't be Claude Code.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly {
|
||||
// Claude Code only restriction: /v1/responses is never a Claude Code
|
||||
// endpoint. With a fallback group the request continues and account selection
|
||||
// (checkClaudeCodeRestriction) schedules it in the fallback group; without one
|
||||
// it is rejected here.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly && apiKey.Group.FallbackGroupID == nil {
|
||||
h.responsesErrorResponse(c, http.StatusForbidden, "permission_error",
|
||||
"This group is restricted to Claude Code clients (/v1/messages only)")
|
||||
return
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
const (
|
||||
apiKeyRateLimitKeyPrefix = "apikey:ratelimit:"
|
||||
apiKeyRateLimitDuration = 24 * time.Hour
|
||||
apiKeyCreateCountKeyPrefix = "apikey:create_count:"
|
||||
apiKeyAuthCachePrefix = "apikey:auth:"
|
||||
authCacheInvalidateChannel = "auth:cache:invalidate"
|
||||
)
|
||||
@@ -24,6 +25,11 @@ func apiKeyRateLimitKey(userID int64) string {
|
||||
return fmt.Sprintf("%s%d", apiKeyRateLimitKeyPrefix, userID)
|
||||
}
|
||||
|
||||
// apiKeyCreateCountKey generates the Redis key for per-user API key creation counting.
|
||||
func apiKeyCreateCountKey(userID int64) string {
|
||||
return fmt.Sprintf("%s%d", apiKeyCreateCountKeyPrefix, userID)
|
||||
}
|
||||
|
||||
func apiKeyAuthCacheKey(key string) string {
|
||||
return fmt.Sprintf("%s%s", apiKeyAuthCachePrefix, key)
|
||||
}
|
||||
@@ -54,9 +60,18 @@ func (c *apiKeyCache) IncrementCreateAttemptCount(ctx context.Context, userID in
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *apiKeyCache) DeleteCreateAttemptCount(ctx context.Context, userID int64) error {
|
||||
key := apiKeyRateLimitKey(userID)
|
||||
return c.rdb.Del(ctx, key).Err()
|
||||
// IncrementCreateCount 在固定窗口内累加创建次数并返回累加后的值。
|
||||
// ExpireNX 只在首次创建计数键时设置过期,后续创建不会延长窗口;
|
||||
// MULTI 保证 INCR 与 EXPIRE 同时生效,避免计数键丢失 TTL 后永久封禁。
|
||||
func (c *apiKeyCache) IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error) {
|
||||
key := apiKeyCreateCountKey(userID)
|
||||
pipe := c.rdb.TxPipeline()
|
||||
incr := pipe.Incr(ctx, key)
|
||||
pipe.ExpireNX(ctx, key, window)
|
||||
if _, err := pipe.Exec(ctx); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return incr.Val(), nil
|
||||
}
|
||||
|
||||
func (c *apiKeyCache) IncrementDailyUsage(ctx context.Context, apiKey string) error {
|
||||
|
||||
@@ -51,19 +51,6 @@ func (s *ApiKeyCacheSuite) TestCreateAttemptCount() {
|
||||
s.AssertTTLWithin(ttl, 1*time.Second, apiKeyRateLimitDuration)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "delete_removes_key",
|
||||
fn: func(ctx context.Context, rdb *redis.Client, cache *apiKeyCache) {
|
||||
userID := int64(1)
|
||||
|
||||
require.NoError(s.T(), cache.IncrementCreateAttemptCount(ctx, userID))
|
||||
require.NoError(s.T(), cache.DeleteCreateAttemptCount(ctx, userID), "DeleteCreateAttemptCount")
|
||||
|
||||
count, err := cache.GetCreateAttemptCount(ctx, userID)
|
||||
require.NoError(s.T(), err, "expected nil error after delete")
|
||||
require.Equal(s.T(), 0, count, "expected zero count after delete")
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAPIKeyCacheIncrementCreateCountUsesFixedWindow(t *testing.T) {
|
||||
server := miniredis.RunT(t)
|
||||
client := redis.NewClient(&redis.Options{Addr: server.Addr()})
|
||||
defer func() { _ = client.Close() }()
|
||||
cache := NewAPIKeyCache(client)
|
||||
ctx := context.Background()
|
||||
key := apiKeyCreateCountKey(7)
|
||||
|
||||
count, err := cache.IncrementCreateCount(ctx, 7, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), count)
|
||||
require.Equal(t, time.Hour, server.TTL(key))
|
||||
|
||||
// 后续创建不得延长窗口,否则持续创建会让计数永不过期。
|
||||
server.FastForward(40 * time.Minute)
|
||||
count, err = cache.IncrementCreateCount(ctx, 7, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(2), count)
|
||||
require.Equal(t, 20*time.Minute, server.TTL(key))
|
||||
|
||||
// 其他用户独立计数。
|
||||
count, err = cache.IncrementCreateCount(ctx, 8, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), count)
|
||||
|
||||
// 窗口到期后重新计数。
|
||||
server.FastForward(21 * time.Minute)
|
||||
count, err = cache.IncrementCreateCount(ctx, 7, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), count)
|
||||
}
|
||||
@@ -1737,8 +1737,8 @@ func (stubApiKeyCache) IncrementCreateAttemptCount(ctx context.Context, userID i
|
||||
return nil
|
||||
}
|
||||
|
||||
func (stubApiKeyCache) DeleteCreateAttemptCount(ctx context.Context, userID int64) error {
|
||||
return nil
|
||||
func (stubApiKeyCache) IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (stubApiKeyCache) IncrementDailyUsage(ctx context.Context, apiKey string) error {
|
||||
|
||||
@@ -30,6 +30,8 @@ var (
|
||||
ErrAPIKeyTooShort = infraerrors.BadRequest("API_KEY_TOO_SHORT", "api key must be at least 16 characters")
|
||||
ErrAPIKeyInvalidChars = infraerrors.BadRequest("API_KEY_INVALID_CHARS", "api key can only contain letters, numbers, underscores, and hyphens")
|
||||
ErrAPIKeyRateLimited = infraerrors.TooManyRequests("API_KEY_RATE_LIMITED", "too many failed attempts, please try again later")
|
||||
ErrAPIKeyCreateLimited = infraerrors.TooManyRequests("API_KEY_CREATE_RATE_LIMITED", "too many api keys created recently, please try again later")
|
||||
ErrAPIKeyCountExceeded = infraerrors.Forbidden("API_KEY_COUNT_EXCEEDED", "api key count limit reached, please delete unused keys first")
|
||||
ErrAPIKeyAuthOverloaded = infraerrors.ServiceUnavailable("API_KEY_AUTH_OVERLOADED", "api key authentication is temporarily overloaded")
|
||||
ErrInvalidIPPattern = infraerrors.BadRequest("INVALID_IP_PATTERN", "invalid IP or CIDR pattern")
|
||||
// ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key has expired")
|
||||
@@ -48,6 +50,7 @@ const (
|
||||
defaultAuthLookupConcurrency = 64
|
||||
defaultNegativeAuthCacheSize = 16384
|
||||
apiKeyMaxErrorsPerHour = 20
|
||||
apiKeyCreateCountWindow = time.Hour
|
||||
apiKeyLastUsedMinTouch = 30 * time.Second
|
||||
apiKeySortCurrentConcurrency = "current_concurrency"
|
||||
// DB 写失败后的短退避,避免请求路径持续同步重试造成写风暴与高延迟。
|
||||
@@ -172,7 +175,7 @@ type APIKeyQuotaUsageState struct {
|
||||
type APIKeyCache interface {
|
||||
GetCreateAttemptCount(ctx context.Context, userID int64) (int, error)
|
||||
IncrementCreateAttemptCount(ctx context.Context, userID int64) error
|
||||
DeleteCreateAttemptCount(ctx context.Context, userID int64) error
|
||||
IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error)
|
||||
|
||||
IncrementDailyUsage(ctx context.Context, apiKey string) error
|
||||
SetDailyUsageExpiry(ctx context.Context, apiKey string, ttl time.Duration) error
|
||||
@@ -435,6 +438,34 @@ func (s *APIKeyService) checkAPIKeyRateLimit(ctx context.Context, userID int64)
|
||||
return nil
|
||||
}
|
||||
|
||||
// checkAPIKeyCreateLimits 校验创建 API Key 的防滥用限制(对自定义与自动生成的 Key 一视同仁)。
|
||||
// 数量上限按未删除的 Key 计;创建次数按固定窗口累计,删除 Key 不返还次数,
|
||||
// 以阻断"删除后反复新建"的循环。Redis 出错时与自定义 Key 限流一致,不阻止用户操作。
|
||||
func (s *APIKeyService) checkAPIKeyCreateLimits(ctx context.Context, userID int64) error {
|
||||
if s.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
if maxActive := s.cfg.APIKeyCreate.MaxActivePerUser; maxActive > 0 {
|
||||
count, err := s.apiKeyRepo.CountByUserID(ctx, userID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("count api keys: %w", err)
|
||||
}
|
||||
if count >= int64(maxActive) {
|
||||
return ErrAPIKeyCountExceeded
|
||||
}
|
||||
}
|
||||
if maxPerHour := s.cfg.APIKeyCreate.MaxPerUserPerHour; maxPerHour > 0 && s.cache != nil {
|
||||
count, err := s.cache.IncrementCreateCount(ctx, userID, apiKeyCreateCountWindow)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
if count > int64(maxPerHour) {
|
||||
return ErrAPIKeyCreateLimited
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// incrementAPIKeyErrorCount 增加用户创建自定义Key的错误计数
|
||||
func (s *APIKeyService) incrementAPIKeyErrorCount(ctx context.Context, userID int64) {
|
||||
if s.cache == nil {
|
||||
@@ -530,6 +561,10 @@ func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIK
|
||||
}
|
||||
}
|
||||
|
||||
if err := s.checkAPIKeyCreateLimits(ctx, userID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 创建API Key记录
|
||||
apiKey := &APIKey{
|
||||
UserID: userID,
|
||||
@@ -820,10 +855,6 @@ func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req
|
||||
if req.Status != nil {
|
||||
apiKey.Status = *req.Status
|
||||
fields.Status = true
|
||||
// 如果状态改变,清除Redis缓存
|
||||
if s.cache != nil {
|
||||
_ = s.cache.DeleteCreateAttemptCount(ctx, apiKey.UserID)
|
||||
}
|
||||
}
|
||||
|
||||
// Update quota fields
|
||||
@@ -931,9 +962,7 @@ func (s *APIKeyService) Delete(ctx context.Context, id int64, userID int64) erro
|
||||
}
|
||||
|
||||
// 删除成功后再清理缓存,避免"缓存已清但删除失败"的竞态。
|
||||
if s.cache != nil {
|
||||
_ = s.cache.DeleteCreateAttemptCount(ctx, userID)
|
||||
}
|
||||
// 注意:不清零创建相关计数,否则"删除后反复新建"即可绕过创建限流。
|
||||
s.InvalidateAuthCacheByKey(ctx, key)
|
||||
s.lastUsedTouchL1.Delete(id)
|
||||
|
||||
|
||||
@@ -138,8 +138,8 @@ func (s *authCacheStub) IncrementCreateAttemptCount(ctx context.Context, userID
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *authCacheStub) DeleteCreateAttemptCount(ctx context.Context, userID int64) error {
|
||||
return nil
|
||||
func (s *authCacheStub) IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (s *authCacheStub) IncrementDailyUsage(ctx context.Context, apiKey string) error {
|
||||
|
||||
@@ -0,0 +1,158 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type createLimitAPIKeyRepoStub struct {
|
||||
*apiKeyRepoStub
|
||||
activeCount int64
|
||||
countErr error
|
||||
created []*APIKey
|
||||
}
|
||||
|
||||
func (s *createLimitAPIKeyRepoStub) CountByUserID(ctx context.Context, userID int64) (int64, error) {
|
||||
return s.activeCount, s.countErr
|
||||
}
|
||||
|
||||
func (s *createLimitAPIKeyRepoStub) ExistsByKey(ctx context.Context, key string) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func (s *createLimitAPIKeyRepoStub) Create(ctx context.Context, key *APIKey) error {
|
||||
s.created = append(s.created, key)
|
||||
return nil
|
||||
}
|
||||
|
||||
type createLimitCacheStub struct {
|
||||
*apiKeyCacheStub
|
||||
createCounts map[int64]int64
|
||||
incrErr error
|
||||
windows []time.Duration
|
||||
}
|
||||
|
||||
func (s *createLimitCacheStub) IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error) {
|
||||
s.windows = append(s.windows, window)
|
||||
if s.incrErr != nil {
|
||||
return 0, s.incrErr
|
||||
}
|
||||
s.createCounts[userID]++
|
||||
return s.createCounts[userID], nil
|
||||
}
|
||||
|
||||
func newCreateLimitService(repo *createLimitAPIKeyRepoStub, cache *createLimitCacheStub, maxActive, maxPerHour int) *APIKeyService {
|
||||
cfg := &config.Config{}
|
||||
cfg.APIKeyCreate.MaxActivePerUser = maxActive
|
||||
cfg.APIKeyCreate.MaxPerUserPerHour = maxPerHour
|
||||
return &APIKeyService{
|
||||
apiKeyRepo: repo,
|
||||
userRepo: &userRepoStub{user: &User{ID: 7}},
|
||||
cache: cache,
|
||||
cfg: cfg,
|
||||
}
|
||||
}
|
||||
|
||||
func newCreateLimitStubs() (*createLimitAPIKeyRepoStub, *createLimitCacheStub) {
|
||||
return &createLimitAPIKeyRepoStub{apiKeyRepoStub: &apiKeyRepoStub{}},
|
||||
&createLimitCacheStub{apiKeyCacheStub: &apiKeyCacheStub{}, createCounts: map[int64]int64{}}
|
||||
}
|
||||
|
||||
func TestAPIKeyServiceCreate_RateLimitsGeneratedKeys(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
svc := newCreateLimitService(repo, cache, 0, 3)
|
||||
|
||||
for i := 0; i < 3; i++ {
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.ErrorIs(t, err, ErrAPIKeyCreateLimited)
|
||||
require.Len(t, repo.created, 3)
|
||||
require.Equal(t, apiKeyCreateCountWindow, cache.windows[0])
|
||||
}
|
||||
|
||||
func TestAPIKeyServiceCreate_RateLimitsCustomKeys(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
svc := newCreateLimitService(repo, cache, 0, 1)
|
||||
custom := "sk-custom-key-0123456789"
|
||||
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k", CustomKey: &custom})
|
||||
require.NoError(t, err)
|
||||
_, err = svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k", CustomKey: &custom})
|
||||
require.ErrorIs(t, err, ErrAPIKeyCreateLimited)
|
||||
require.Len(t, repo.created, 1)
|
||||
}
|
||||
|
||||
// 删除 Key 释放数量名额,但不返还创建次数,删建循环仍受每小时次数限制。
|
||||
func TestAPIKeyServiceCreate_DeleteDoesNotRefundCreateCount(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
repo.apiKey = &APIKey{ID: 1, UserID: 7, Key: "k"}
|
||||
svc := newCreateLimitService(repo, cache, 1, 2)
|
||||
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.NoError(t, err)
|
||||
repo.activeCount = 1
|
||||
_, err = svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.ErrorIs(t, err, ErrAPIKeyCountExceeded)
|
||||
|
||||
require.NoError(t, svc.Delete(context.Background(), 1, 7))
|
||||
repo.activeCount = 0
|
||||
_, err = svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.NoError(t, err)
|
||||
|
||||
repo.activeCount = 0
|
||||
_, err = svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.ErrorIs(t, err, ErrAPIKeyCreateLimited)
|
||||
require.Len(t, repo.created, 2)
|
||||
}
|
||||
|
||||
func TestAPIKeyServiceCreate_ActiveLimitDoesNotConsumeCreateCount(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
repo.activeCount = 5
|
||||
svc := newCreateLimitService(repo, cache, 5, 10)
|
||||
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.ErrorIs(t, err, ErrAPIKeyCountExceeded)
|
||||
require.Zero(t, cache.createCounts[7])
|
||||
}
|
||||
|
||||
func TestAPIKeyServiceCreate_ZeroLimitsDisableChecks(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
repo.countErr = errors.New("must not count")
|
||||
svc := newCreateLimitService(repo, cache, 0, 0)
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.Empty(t, cache.windows)
|
||||
}
|
||||
|
||||
func TestAPIKeyServiceCreate_CountErrorFailsClosed(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
repo.countErr = errors.New("db down")
|
||||
svc := newCreateLimitService(repo, cache, 5, 0)
|
||||
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.Error(t, err)
|
||||
require.Empty(t, repo.created)
|
||||
}
|
||||
|
||||
func TestAPIKeyServiceCreate_RedisErrorFailsOpen(t *testing.T) {
|
||||
repo, cache := newCreateLimitStubs()
|
||||
cache.incrErr = errors.New("redis down")
|
||||
svc := newCreateLimitService(repo, cache, 0, 1)
|
||||
|
||||
for i := 0; i < 3; i++ {
|
||||
_, err := svc.Create(context.Background(), 7, CreateAPIKeyRequest{Name: "k"})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
}
|
||||
@@ -234,11 +234,7 @@ func (s *apiKeyRepoStub) GetRateLimitData(ctx context.Context, id int64) (*APIKe
|
||||
|
||||
// apiKeyCacheStub 是 APIKeyCache 接口的测试桩实现。
|
||||
// 用于验证删除操作时缓存清理逻辑是否被正确调用。
|
||||
//
|
||||
// 设计说明:
|
||||
// - invalidated: 记录被清除缓存的用户 ID 列表
|
||||
type apiKeyCacheStub struct {
|
||||
invalidated []int64 // 记录调用 DeleteCreateAttemptCount 时传入的用户 ID
|
||||
deleteAuthKeys []string // 记录调用 DeleteAuthCache 时传入的缓存 key
|
||||
}
|
||||
|
||||
@@ -252,11 +248,9 @@ func (s *apiKeyCacheStub) IncrementCreateAttemptCount(ctx context.Context, userI
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteCreateAttemptCount 记录被清除缓存的用户 ID。
|
||||
// 删除 API Key 时会调用此方法清除用户的创建尝试计数缓存。
|
||||
func (s *apiKeyCacheStub) DeleteCreateAttemptCount(ctx context.Context, userID int64) error {
|
||||
s.invalidated = append(s.invalidated, userID)
|
||||
return nil
|
||||
// IncrementCreateCount 空实现,本测试不验证此行为
|
||||
func (s *apiKeyCacheStub) IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// IncrementDailyUsage 空实现,本测试不验证此行为
|
||||
@@ -306,8 +300,7 @@ func TestApiKeyService_Delete_OwnerMismatch(t *testing.T) {
|
||||
|
||||
err := svc.Delete(context.Background(), 10, 2) // API Key ID=10, 调用者 userID=2
|
||||
require.ErrorIs(t, err, ErrInsufficientPerms)
|
||||
require.Empty(t, repo.deletedIDs) // 验证删除操作未被调用
|
||||
require.Empty(t, cache.invalidated) // 验证缓存未被清除
|
||||
require.Empty(t, repo.deletedIDs) // 验证删除操作未被调用
|
||||
require.Empty(t, cache.deleteAuthKeys)
|
||||
}
|
||||
|
||||
@@ -316,7 +309,7 @@ func TestApiKeyService_Delete_OwnerMismatch(t *testing.T) {
|
||||
// - GetKeyAndOwnerID 返回所有者 ID 为 7
|
||||
// - 调用者 userID 为 7(匹配)
|
||||
// - Delete 成功执行
|
||||
// - 缓存被正确清除(使用 ownerID)
|
||||
// - 认证缓存被正确清除
|
||||
// - 返回 nil 错误
|
||||
func TestApiKeyService_Delete_Success(t *testing.T) {
|
||||
repo := &apiKeyRepoStub{
|
||||
@@ -328,8 +321,7 @@ func TestApiKeyService_Delete_Success(t *testing.T) {
|
||||
|
||||
err := svc.Delete(context.Background(), 42, 7) // API Key ID=42, 调用者 userID=7
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []int64{42}, repo.deletedIDs) // 验证正确的 API Key 被删除
|
||||
require.Equal(t, []int64{7}, cache.invalidated) // 验证所有者的缓存被清除
|
||||
require.Equal(t, []int64{42}, repo.deletedIDs) // 验证正确的 API Key 被删除
|
||||
require.Equal(t, []string{svc.authCacheKey("k")}, cache.deleteAuthKeys)
|
||||
_, exists := svc.lastUsedTouchL1.Load(int64(42))
|
||||
require.False(t, exists, "delete should clear touch debounce cache")
|
||||
@@ -349,7 +341,6 @@ func TestApiKeyService_Delete_NotFound(t *testing.T) {
|
||||
err := svc.Delete(context.Background(), 99, 1)
|
||||
require.ErrorIs(t, err, ErrAPIKeyNotFound)
|
||||
require.Empty(t, repo.deletedIDs)
|
||||
require.Empty(t, cache.invalidated)
|
||||
require.Empty(t, cache.deleteAuthKeys)
|
||||
}
|
||||
|
||||
@@ -496,6 +487,5 @@ func TestApiKeyService_Delete_DeleteFails(t *testing.T) {
|
||||
require.Error(t, err)
|
||||
require.ErrorContains(t, err, "delete api key")
|
||||
require.Equal(t, []int64{3}, repo.deletedIDs) // 验证 DeleteWithAudit 被调用
|
||||
require.Empty(t, cache.invalidated) // 验证删除失败时缓存未被清除(新顺序:先删后清)
|
||||
require.Empty(t, cache.deleteAuthKeys) // 验证删除失败时 auth 缓存未被清除
|
||||
}
|
||||
|
||||
@@ -42,8 +42,8 @@ func (s *quotaStateCacheStub) IncrementCreateAttemptCount(context.Context, int64
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *quotaStateCacheStub) DeleteCreateAttemptCount(context.Context, int64) error {
|
||||
return nil
|
||||
func (s *quotaStateCacheStub) IncrementCreateCount(context.Context, int64, time.Duration) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (s *quotaStateCacheStub) IncrementDailyUsage(context.Context, string) error {
|
||||
|
||||
@@ -67,8 +67,10 @@ func (*authInvalidationCacheStub) GetCreateAttemptCount(context.Context, int64)
|
||||
func (*authInvalidationCacheStub) IncrementCreateAttemptCount(context.Context, int64) error {
|
||||
return nil
|
||||
}
|
||||
func (*authInvalidationCacheStub) DeleteCreateAttemptCount(context.Context, int64) error { return nil }
|
||||
func (*authInvalidationCacheStub) IncrementDailyUsage(context.Context, string) error { return nil }
|
||||
func (*authInvalidationCacheStub) IncrementCreateCount(context.Context, int64, time.Duration) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (*authInvalidationCacheStub) IncrementDailyUsage(context.Context, string) error { return nil }
|
||||
func (*authInvalidationCacheStub) SetDailyUsageExpiry(context.Context, string, time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -717,6 +717,19 @@ api_key_auth_cache:
|
||||
# 每进程最多跟踪的客户端身份数量,确保内存有界。
|
||||
capacity: 16384
|
||||
|
||||
# =============================================================================
|
||||
# API Key Creation Limits
|
||||
# 用户创建 API Key 的防滥用限制(0 表示不限制)
|
||||
# =============================================================================
|
||||
api_key_create:
|
||||
# Maximum non-deleted API keys per user
|
||||
# 单个用户同时存在(未删除)的 API Key 上限
|
||||
max_active_per_user: 200
|
||||
# Maximum API key creations per user per hour (custom and generated keys).
|
||||
# Deleting a key does not refund the count, which blocks delete/recreate loops.
|
||||
# 单个用户每小时可创建的 API Key 次数(自定义与自动生成均计入);删除 Key 不返还次数,阻断删建循环
|
||||
max_per_user_per_hour: 60
|
||||
|
||||
# =============================================================================
|
||||
# Dashboard Cache Configuration
|
||||
# 仪表盘缓存配置
|
||||
|
||||
Reference in New Issue
Block a user