feat(api-key): 增加 API Key 上次使用时间并补齐测试
This commit is contained in:
@@ -0,0 +1,40 @@
|
||||
package dto
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAPIKeyFromService_MapsLastUsedAt(t *testing.T) {
|
||||
lastUsed := time.Now().UTC().Truncate(time.Second)
|
||||
src := &service.APIKey{
|
||||
ID: 1,
|
||||
UserID: 2,
|
||||
Key: "sk-map-last-used",
|
||||
Name: "Mapper",
|
||||
Status: service.StatusActive,
|
||||
LastUsedAt: &lastUsed,
|
||||
}
|
||||
|
||||
out := APIKeyFromService(src)
|
||||
require.NotNil(t, out)
|
||||
require.NotNil(t, out.LastUsedAt)
|
||||
require.WithinDuration(t, lastUsed, *out.LastUsedAt, time.Second)
|
||||
}
|
||||
|
||||
func TestAPIKeyFromService_MapsNilLastUsedAt(t *testing.T) {
|
||||
src := &service.APIKey{
|
||||
ID: 1,
|
||||
UserID: 2,
|
||||
Key: "sk-map-last-used-nil",
|
||||
Name: "MapperNil",
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
|
||||
out := APIKeyFromService(src)
|
||||
require.NotNil(t, out)
|
||||
require.Nil(t, out.LastUsedAt)
|
||||
}
|
||||
@@ -77,6 +77,7 @@ func APIKeyFromService(k *service.APIKey) *APIKey {
|
||||
Status: k.Status,
|
||||
IPWhitelist: k.IPWhitelist,
|
||||
IPBlacklist: k.IPBlacklist,
|
||||
LastUsedAt: k.LastUsedAt,
|
||||
Quota: k.Quota,
|
||||
QuotaUsed: k.QuotaUsed,
|
||||
ExpiresAt: k.ExpiresAt,
|
||||
|
||||
@@ -38,6 +38,7 @@ type APIKey struct {
|
||||
Status string `json:"status"`
|
||||
IPWhitelist []string `json:"ip_whitelist"`
|
||||
IPBlacklist []string `json:"ip_blacklist"`
|
||||
LastUsedAt *time.Time `json:"last_used_at"`
|
||||
Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited)
|
||||
QuotaUsed float64 `json:"quota_used"` // Used quota amount in USD
|
||||
ExpiresAt *time.Time `json:"expires_at"` // Expiration time (nil = never expires)
|
||||
|
||||
@@ -34,6 +34,7 @@ func (r *apiKeyRepository) Create(ctx context.Context, key *service.APIKey) erro
|
||||
SetName(key.Name).
|
||||
SetStatus(key.Status).
|
||||
SetNillableGroupID(key.GroupID).
|
||||
SetNillableLastUsedAt(key.LastUsedAt).
|
||||
SetQuota(key.Quota).
|
||||
SetQuotaUsed(key.QuotaUsed).
|
||||
SetNillableExpiresAt(key.ExpiresAt)
|
||||
@@ -48,6 +49,7 @@ func (r *apiKeyRepository) Create(ctx context.Context, key *service.APIKey) erro
|
||||
created, err := builder.Save(ctx)
|
||||
if err == nil {
|
||||
key.ID = created.ID
|
||||
key.LastUsedAt = created.LastUsedAt
|
||||
key.CreatedAt = created.CreatedAt
|
||||
key.UpdatedAt = created.UpdatedAt
|
||||
}
|
||||
@@ -394,6 +396,21 @@ func (r *apiKeyRepository) IncrementQuotaUsed(ctx context.Context, id int64, amo
|
||||
return updated.QuotaUsed, nil
|
||||
}
|
||||
|
||||
func (r *apiKeyRepository) UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error {
|
||||
affected, err := r.client.APIKey.Update().
|
||||
Where(apikey.IDEQ(id), apikey.DeletedAtIsNil()).
|
||||
SetLastUsedAt(usedAt).
|
||||
SetUpdatedAt(usedAt).
|
||||
Save(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if affected == 0 {
|
||||
return service.ErrAPIKeyNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func apiKeyEntityToService(m *dbent.APIKey) *service.APIKey {
|
||||
if m == nil {
|
||||
return nil
|
||||
@@ -406,6 +423,7 @@ func apiKeyEntityToService(m *dbent.APIKey) *service.APIKey {
|
||||
Status: m.Status,
|
||||
IPWhitelist: m.IPWhitelist,
|
||||
IPBlacklist: m.IPBlacklist,
|
||||
LastUsedAt: m.LastUsedAt,
|
||||
CreatedAt: m.CreatedAt,
|
||||
UpdatedAt: m.UpdatedAt,
|
||||
GroupID: m.GroupID,
|
||||
|
||||
156
backend/internal/repository/api_key_repo_last_used_unit_test.go
Normal file
156
backend/internal/repository/api_key_repo_last_used_unit_test.go
Normal file
@@ -0,0 +1,156 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/ent/enttest"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"entgo.io/ent/dialect"
|
||||
entsql "entgo.io/ent/dialect/sql"
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
func newAPIKeyRepoSQLite(t *testing.T) (*apiKeyRepository, *dbent.Client) {
|
||||
t.Helper()
|
||||
|
||||
db, err := sql.Open("sqlite", "file:api_key_repo_last_used?mode=memory&cache=shared")
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
_, err = db.Exec("PRAGMA foreign_keys = ON")
|
||||
require.NoError(t, err)
|
||||
|
||||
drv := entsql.OpenDB(dialect.SQLite, db)
|
||||
client := enttest.NewClient(t, enttest.WithOptions(dbent.Driver(drv)))
|
||||
t.Cleanup(func() { _ = client.Close() })
|
||||
|
||||
return &apiKeyRepository{client: client}, client
|
||||
}
|
||||
|
||||
func mustCreateAPIKeyRepoUser(t *testing.T, ctx context.Context, client *dbent.Client, email string) *service.User {
|
||||
t.Helper()
|
||||
u, err := client.User.Create().
|
||||
SetEmail(email).
|
||||
SetPasswordHash("test-password-hash").
|
||||
SetRole(service.RoleUser).
|
||||
SetStatus(service.StatusActive).
|
||||
Save(ctx)
|
||||
require.NoError(t, err)
|
||||
return userEntityToService(u)
|
||||
}
|
||||
|
||||
func TestAPIKeyRepository_CreateWithLastUsedAt(t *testing.T) {
|
||||
repo, client := newAPIKeyRepoSQLite(t)
|
||||
ctx := context.Background()
|
||||
user := mustCreateAPIKeyRepoUser(t, ctx, client, "create-last-used@test.com")
|
||||
|
||||
lastUsed := time.Now().UTC().Add(-time.Hour).Truncate(time.Second)
|
||||
key := &service.APIKey{
|
||||
UserID: user.ID,
|
||||
Key: "sk-create-last-used",
|
||||
Name: "CreateWithLastUsed",
|
||||
Status: service.StatusActive,
|
||||
LastUsedAt: &lastUsed,
|
||||
}
|
||||
|
||||
require.NoError(t, repo.Create(ctx, key))
|
||||
require.NotNil(t, key.LastUsedAt)
|
||||
require.WithinDuration(t, lastUsed, *key.LastUsedAt, time.Second)
|
||||
|
||||
got, err := repo.GetByID(ctx, key.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, got.LastUsedAt)
|
||||
require.WithinDuration(t, lastUsed, *got.LastUsedAt, time.Second)
|
||||
}
|
||||
|
||||
func TestAPIKeyRepository_UpdateLastUsed(t *testing.T) {
|
||||
repo, client := newAPIKeyRepoSQLite(t)
|
||||
ctx := context.Background()
|
||||
user := mustCreateAPIKeyRepoUser(t, ctx, client, "update-last-used@test.com")
|
||||
|
||||
key := &service.APIKey{
|
||||
UserID: user.ID,
|
||||
Key: "sk-update-last-used",
|
||||
Name: "UpdateLastUsed",
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
require.NoError(t, repo.Create(ctx, key))
|
||||
|
||||
before, err := repo.GetByID(ctx, key.ID)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, before.LastUsedAt)
|
||||
|
||||
target := time.Now().UTC().Add(2 * time.Minute).Truncate(time.Second)
|
||||
require.NoError(t, repo.UpdateLastUsed(ctx, key.ID, target))
|
||||
|
||||
after, err := repo.GetByID(ctx, key.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, after.LastUsedAt)
|
||||
require.WithinDuration(t, target, *after.LastUsedAt, time.Second)
|
||||
require.WithinDuration(t, target, after.UpdatedAt, time.Second)
|
||||
}
|
||||
|
||||
func TestAPIKeyRepository_UpdateLastUsedDeletedKey(t *testing.T) {
|
||||
repo, client := newAPIKeyRepoSQLite(t)
|
||||
ctx := context.Background()
|
||||
user := mustCreateAPIKeyRepoUser(t, ctx, client, "deleted-last-used@test.com")
|
||||
|
||||
key := &service.APIKey{
|
||||
UserID: user.ID,
|
||||
Key: "sk-update-last-used-deleted",
|
||||
Name: "UpdateLastUsedDeleted",
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
require.NoError(t, repo.Create(ctx, key))
|
||||
require.NoError(t, repo.Delete(ctx, key.ID))
|
||||
|
||||
err := repo.UpdateLastUsed(ctx, key.ID, time.Now().UTC())
|
||||
require.ErrorIs(t, err, service.ErrAPIKeyNotFound)
|
||||
}
|
||||
|
||||
func TestAPIKeyRepository_UpdateLastUsedDBError(t *testing.T) {
|
||||
repo, client := newAPIKeyRepoSQLite(t)
|
||||
ctx := context.Background()
|
||||
user := mustCreateAPIKeyRepoUser(t, ctx, client, "db-error-last-used@test.com")
|
||||
|
||||
key := &service.APIKey{
|
||||
UserID: user.ID,
|
||||
Key: "sk-update-last-used-db-error",
|
||||
Name: "UpdateLastUsedDBError",
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
require.NoError(t, repo.Create(ctx, key))
|
||||
|
||||
require.NoError(t, client.Close())
|
||||
err := repo.UpdateLastUsed(ctx, key.ID, time.Now().UTC())
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestAPIKeyRepository_CreateDuplicateKey(t *testing.T) {
|
||||
repo, client := newAPIKeyRepoSQLite(t)
|
||||
ctx := context.Background()
|
||||
user := mustCreateAPIKeyRepoUser(t, ctx, client, "duplicate-key@test.com")
|
||||
|
||||
first := &service.APIKey{
|
||||
UserID: user.ID,
|
||||
Key: "sk-duplicate",
|
||||
Name: "first",
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
second := &service.APIKey{
|
||||
UserID: user.ID,
|
||||
Key: "sk-duplicate",
|
||||
Name: "second",
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
|
||||
require.NoError(t, repo.Create(ctx, first))
|
||||
err := repo.Create(ctx, second)
|
||||
require.ErrorIs(t, err, service.ErrAPIKeyExists)
|
||||
}
|
||||
@@ -83,6 +83,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"status": "active",
|
||||
"ip_whitelist": null,
|
||||
"ip_blacklist": null,
|
||||
"last_used_at": null,
|
||||
"quota": 0,
|
||||
"quota_used": 0,
|
||||
"expires_at": null,
|
||||
@@ -122,6 +123,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"status": "active",
|
||||
"ip_whitelist": null,
|
||||
"ip_blacklist": null,
|
||||
"last_used_at": null,
|
||||
"quota": 0,
|
||||
"quota_used": 0,
|
||||
"expires_at": null,
|
||||
@@ -1471,6 +1473,20 @@ func (r *stubApiKeyRepo) IncrementQuotaUsed(ctx context.Context, id int64, amoun
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error {
|
||||
key, ok := r.byID[id]
|
||||
if !ok {
|
||||
return service.ErrAPIKeyNotFound
|
||||
}
|
||||
ts := usedAt
|
||||
key.LastUsedAt = &ts
|
||||
key.UpdatedAt = usedAt
|
||||
clone := *key
|
||||
r.byID[id] = &clone
|
||||
r.byKey[clone.Key] = &clone
|
||||
return nil
|
||||
}
|
||||
|
||||
type stubUsageLogRepo struct {
|
||||
userLogs map[int64][]service.UsageLog
|
||||
}
|
||||
|
||||
@@ -125,6 +125,7 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
|
||||
setGroupContext(c, apiKey.Group)
|
||||
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
@@ -184,6 +185,7 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
|
||||
setGroupContext(c, apiKey.Group)
|
||||
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
|
||||
|
||||
c.Next()
|
||||
}
|
||||
|
||||
@@ -64,6 +64,7 @@ func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subs
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
|
||||
setGroupContext(c, apiKey.Group)
|
||||
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
@@ -104,6 +105,7 @@ func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subs
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
|
||||
setGroupContext(c, apiKey.Group)
|
||||
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
@@ -18,7 +19,8 @@ import (
|
||||
)
|
||||
|
||||
type fakeAPIKeyRepo struct {
|
||||
getByKey func(ctx context.Context, key string) (*service.APIKey, error)
|
||||
getByKey func(ctx context.Context, key string) (*service.APIKey, error)
|
||||
updateLastUsed func(ctx context.Context, id int64, usedAt time.Time) error
|
||||
}
|
||||
|
||||
func (f fakeAPIKeyRepo) Create(ctx context.Context, key *service.APIKey) error {
|
||||
@@ -78,6 +80,12 @@ func (f fakeAPIKeyRepo) ListKeysByGroupID(ctx context.Context, groupID int64) ([
|
||||
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
|
||||
}
|
||||
|
||||
type googleErrorResponse struct {
|
||||
Error struct {
|
||||
@@ -356,3 +364,144 @@ func TestApiKeyAuthWithSubscriptionGoogle_InsufficientBalance(t *testing.T) {
|
||||
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)
|
||||
}
|
||||
|
||||
@@ -351,6 +351,147 @@ func TestAPIKeyAuthIPRestrictionDoesNotTrustSpoofedForwardHeaders(t *testing.T)
|
||||
require.Contains(t, w.Body.String(), "ACCESS_DENIED")
|
||||
}
|
||||
|
||||
func TestAPIKeyAuthTouchesLastUsedOnSuccess(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
user := &service.User{
|
||||
ID: 7,
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 10,
|
||||
Concurrency: 3,
|
||||
}
|
||||
apiKey := &service.APIKey{
|
||||
ID: 100,
|
||||
UserID: user.ID,
|
||||
Key: "touch-ok",
|
||||
Status: service.StatusActive,
|
||||
User: user,
|
||||
}
|
||||
|
||||
var touchedID int64
|
||||
var touchedAt time.Time
|
||||
apiKeyRepo := &stubApiKeyRepo{
|
||||
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}
|
||||
apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
|
||||
router := newAuthTestRouter(apiKeyService, nil, cfg)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
req.Header.Set("x-api-key", apiKey.Key)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, apiKey.ID, touchedID)
|
||||
require.False(t, touchedAt.IsZero(), "expected touch timestamp")
|
||||
}
|
||||
|
||||
func TestAPIKeyAuthTouchLastUsedFailureDoesNotBlock(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
user := &service.User{
|
||||
ID: 8,
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 10,
|
||||
Concurrency: 3,
|
||||
}
|
||||
apiKey := &service.APIKey{
|
||||
ID: 101,
|
||||
UserID: user.ID,
|
||||
Key: "touch-fail",
|
||||
Status: service.StatusActive,
|
||||
User: user,
|
||||
}
|
||||
|
||||
touchCalls := 0
|
||||
apiKeyRepo := &stubApiKeyRepo{
|
||||
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("db unavailable")
|
||||
},
|
||||
}
|
||||
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
|
||||
router := newAuthTestRouter(apiKeyService, nil, cfg)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
req.Header.Set("x-api-key", apiKey.Key)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code, "touch failure should not block request")
|
||||
require.Equal(t, 1, touchCalls)
|
||||
}
|
||||
|
||||
func TestAPIKeyAuthTouchesLastUsedInStandardMode(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
user := &service.User{
|
||||
ID: 9,
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 10,
|
||||
Concurrency: 3,
|
||||
}
|
||||
apiKey := &service.APIKey{
|
||||
ID: 102,
|
||||
UserID: user.ID,
|
||||
Key: "touch-standard",
|
||||
Status: service.StatusActive,
|
||||
User: user,
|
||||
}
|
||||
|
||||
touchCalls := 0
|
||||
apiKeyRepo := &stubApiKeyRepo{
|
||||
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}
|
||||
apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
|
||||
router := newAuthTestRouter(apiKeyService, nil, cfg)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
req.Header.Set("x-api-key", apiKey.Key)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, 1, touchCalls)
|
||||
}
|
||||
|
||||
func newAuthTestRouter(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) *gin.Engine {
|
||||
router := gin.New()
|
||||
router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, subscriptionService, cfg)))
|
||||
@@ -361,7 +502,8 @@ func newAuthTestRouter(apiKeyService *service.APIKeyService, subscriptionService
|
||||
}
|
||||
|
||||
type stubApiKeyRepo struct {
|
||||
getByKey func(ctx context.Context, key string) (*service.APIKey, error)
|
||||
getByKey func(ctx context.Context, key string) (*service.APIKey, error)
|
||||
updateLastUsed func(ctx context.Context, id int64, usedAt time.Time) error
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) Create(ctx context.Context, key *service.APIKey) error {
|
||||
@@ -439,6 +581,13 @@ func (r *stubApiKeyRepo) IncrementQuotaUsed(ctx context.Context, id int64, amoun
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubApiKeyRepo) UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error {
|
||||
if r.updateLastUsed != nil {
|
||||
return r.updateLastUsed(ctx, id, usedAt)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type stubUserSubscriptionRepo struct {
|
||||
getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error)
|
||||
updateStatus func(ctx context.Context, subscriptionID int64, status string) error
|
||||
|
||||
@@ -19,6 +19,7 @@ type APIKey struct {
|
||||
Status string
|
||||
IPWhitelist []string
|
||||
IPBlacklist []string
|
||||
LastUsedAt *time.Time
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
User *User
|
||||
|
||||
@@ -5,6 +5,8 @@ import (
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
@@ -32,6 +34,7 @@ var (
|
||||
|
||||
const (
|
||||
apiKeyMaxErrorsPerHour = 20
|
||||
apiKeyLastUsedMinTouch = 30 * time.Second
|
||||
)
|
||||
|
||||
type APIKeyRepository interface {
|
||||
@@ -58,6 +61,7 @@ type APIKeyRepository interface {
|
||||
|
||||
// Quota methods
|
||||
IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error)
|
||||
UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error
|
||||
}
|
||||
|
||||
// APIKeyCache defines cache operations for API key service
|
||||
@@ -125,6 +129,8 @@ type APIKeyService struct {
|
||||
authCacheL1 *ristretto.Cache
|
||||
authCfg apiKeyAuthCacheConfig
|
||||
authGroup singleflight.Group
|
||||
lastUsedTouchL1 sync.Map // keyID -> time.Time
|
||||
lastUsedTouchSF singleflight.Group
|
||||
}
|
||||
|
||||
// NewAPIKeyService 创建API Key服务实例
|
||||
@@ -527,6 +533,7 @@ func (s *APIKeyService) Delete(ctx context.Context, id int64, userID int64) erro
|
||||
if err := s.apiKeyRepo.Delete(ctx, id); err != nil {
|
||||
return fmt.Errorf("delete api key: %w", err)
|
||||
}
|
||||
s.lastUsedTouchL1.Delete(id)
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -558,6 +565,37 @@ func (s *APIKeyService) ValidateKey(ctx context.Context, key string) (*APIKey, *
|
||||
return apiKey, user, nil
|
||||
}
|
||||
|
||||
// TouchLastUsed 通过防抖更新 api_keys.last_used_at,减少高频写放大。
|
||||
// 该操作为尽力而为,不应阻塞主请求链路。
|
||||
func (s *APIKeyService) TouchLastUsed(ctx context.Context, keyID int64) error {
|
||||
if keyID <= 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
if v, ok := s.lastUsedTouchL1.Load(keyID); ok {
|
||||
if last, ok := v.(time.Time); ok && now.Sub(last) < apiKeyLastUsedMinTouch {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
_, err, _ := s.lastUsedTouchSF.Do(strconv.FormatInt(keyID, 10), func() (any, error) {
|
||||
latest := time.Now()
|
||||
if v, ok := s.lastUsedTouchL1.Load(keyID); ok {
|
||||
if last, ok := v.(time.Time); ok && latest.Sub(last) < apiKeyLastUsedMinTouch {
|
||||
return nil, nil
|
||||
}
|
||||
}
|
||||
|
||||
if err := s.apiKeyRepo.UpdateLastUsed(ctx, keyID, latest); err != nil {
|
||||
return nil, fmt.Errorf("touch api key last used: %w", err)
|
||||
}
|
||||
s.lastUsedTouchL1.Store(keyID, latest)
|
||||
return nil, nil
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
// IncrementUsage 增加API Key使用次数(可选:用于统计)
|
||||
func (s *APIKeyService) IncrementUsage(ctx context.Context, keyID int64) error {
|
||||
// 使用Redis计数器
|
||||
|
||||
@@ -103,6 +103,10 @@ func (s *authRepoStub) IncrementQuotaUsed(ctx context.Context, id int64, amount
|
||||
panic("unexpected IncrementQuotaUsed call")
|
||||
}
|
||||
|
||||
func (s *authRepoStub) UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error {
|
||||
panic("unexpected UpdateLastUsed call")
|
||||
}
|
||||
|
||||
type authCacheStub struct {
|
||||
getAuthCache func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error)
|
||||
setAuthKeys []string
|
||||
|
||||
@@ -24,10 +24,13 @@ import (
|
||||
// - deleteErr: 模拟 Delete 返回的错误
|
||||
// - deletedIDs: 记录被调用删除的 API Key ID,用于断言验证
|
||||
type apiKeyRepoStub struct {
|
||||
apiKey *APIKey // GetKeyAndOwnerID 的返回值
|
||||
getByIDErr error // GetKeyAndOwnerID 的错误返回值
|
||||
deleteErr error // Delete 的错误返回值
|
||||
deletedIDs []int64 // 记录已删除的 API Key ID 列表
|
||||
apiKey *APIKey // GetKeyAndOwnerID 的返回值
|
||||
getByIDErr error // GetKeyAndOwnerID 的错误返回值
|
||||
deleteErr error // Delete 的错误返回值
|
||||
deletedIDs []int64 // 记录已删除的 API Key ID 列表
|
||||
updateLastUsed func(ctx context.Context, id int64, usedAt time.Time) error
|
||||
touchedIDs []int64
|
||||
touchedUsedAts []time.Time
|
||||
}
|
||||
|
||||
// 以下方法在本测试中不应被调用,使用 panic 确保测试失败时能快速定位问题
|
||||
@@ -122,6 +125,15 @@ func (s *apiKeyRepoStub) IncrementQuotaUsed(ctx context.Context, id int64, amoun
|
||||
panic("unexpected IncrementQuotaUsed call")
|
||||
}
|
||||
|
||||
func (s *apiKeyRepoStub) UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error {
|
||||
s.touchedIDs = append(s.touchedIDs, id)
|
||||
s.touchedUsedAts = append(s.touchedUsedAts, usedAt)
|
||||
if s.updateLastUsed != nil {
|
||||
return s.updateLastUsed(ctx, id, usedAt)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// apiKeyCacheStub 是 APIKeyCache 接口的测试桩实现。
|
||||
// 用于验证删除操作时缓存清理逻辑是否被正确调用。
|
||||
//
|
||||
@@ -214,12 +226,15 @@ func TestApiKeyService_Delete_Success(t *testing.T) {
|
||||
}
|
||||
cache := &apiKeyCacheStub{}
|
||||
svc := &APIKeyService{apiKeyRepo: repo, cache: cache}
|
||||
svc.lastUsedTouchL1.Store(int64(42), time.Now())
|
||||
|
||||
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, []string{svc.authCacheKey("k")}, cache.deleteAuthKeys)
|
||||
_, exists := svc.lastUsedTouchL1.Load(int64(42))
|
||||
require.False(t, exists, "delete should clear touch debounce cache")
|
||||
}
|
||||
|
||||
// TestApiKeyService_Delete_NotFound 测试删除不存在的 API Key 时返回正确的错误。
|
||||
|
||||
141
backend/internal/service/api_key_service_touch_last_used_test.go
Normal file
141
backend/internal/service/api_key_service_touch_last_used_test.go
Normal file
@@ -0,0 +1,141 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAPIKeyService_TouchLastUsed_InvalidKeyID(t *testing.T) {
|
||||
repo := &apiKeyRepoStub{
|
||||
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
|
||||
return errors.New("should not be called")
|
||||
},
|
||||
}
|
||||
svc := &APIKeyService{apiKeyRepo: repo}
|
||||
|
||||
require.NoError(t, svc.TouchLastUsed(context.Background(), 0))
|
||||
require.NoError(t, svc.TouchLastUsed(context.Background(), -1))
|
||||
require.Empty(t, repo.touchedIDs)
|
||||
}
|
||||
|
||||
func TestAPIKeyService_TouchLastUsed_FirstTouchSucceeds(t *testing.T) {
|
||||
repo := &apiKeyRepoStub{}
|
||||
svc := &APIKeyService{apiKeyRepo: repo}
|
||||
|
||||
err := svc.TouchLastUsed(context.Background(), 123)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []int64{123}, repo.touchedIDs)
|
||||
require.Len(t, repo.touchedUsedAts, 1)
|
||||
require.False(t, repo.touchedUsedAts[0].IsZero())
|
||||
|
||||
cached, ok := svc.lastUsedTouchL1.Load(int64(123))
|
||||
require.True(t, ok, "successful touch should update debounce cache")
|
||||
_, isTime := cached.(time.Time)
|
||||
require.True(t, isTime)
|
||||
}
|
||||
|
||||
func TestAPIKeyService_TouchLastUsed_DebouncedWithinWindow(t *testing.T) {
|
||||
repo := &apiKeyRepoStub{}
|
||||
svc := &APIKeyService{apiKeyRepo: repo}
|
||||
|
||||
require.NoError(t, svc.TouchLastUsed(context.Background(), 123))
|
||||
require.NoError(t, svc.TouchLastUsed(context.Background(), 123))
|
||||
|
||||
require.Equal(t, []int64{123}, repo.touchedIDs, "second touch within debounce window should not hit repository")
|
||||
}
|
||||
|
||||
func TestAPIKeyService_TouchLastUsed_ExpiredDebounceTouchesAgain(t *testing.T) {
|
||||
repo := &apiKeyRepoStub{}
|
||||
svc := &APIKeyService{apiKeyRepo: repo}
|
||||
|
||||
require.NoError(t, svc.TouchLastUsed(context.Background(), 123))
|
||||
|
||||
// 强制将 debounce 时间回拨到窗口之外,触发第二次写库。
|
||||
svc.lastUsedTouchL1.Store(int64(123), time.Now().Add(-apiKeyLastUsedMinTouch-time.Second))
|
||||
|
||||
require.NoError(t, svc.TouchLastUsed(context.Background(), 123))
|
||||
require.Len(t, repo.touchedIDs, 2)
|
||||
require.Equal(t, int64(123), repo.touchedIDs[0])
|
||||
require.Equal(t, int64(123), repo.touchedIDs[1])
|
||||
}
|
||||
|
||||
func TestAPIKeyService_TouchLastUsed_RepoError(t *testing.T) {
|
||||
repo := &apiKeyRepoStub{
|
||||
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
|
||||
return errors.New("db write failed")
|
||||
},
|
||||
}
|
||||
svc := &APIKeyService{apiKeyRepo: repo}
|
||||
|
||||
err := svc.TouchLastUsed(context.Background(), 123)
|
||||
require.Error(t, err)
|
||||
require.ErrorContains(t, err, "touch api key last used")
|
||||
require.Equal(t, []int64{123}, repo.touchedIDs)
|
||||
|
||||
_, ok := svc.lastUsedTouchL1.Load(int64(123))
|
||||
require.False(t, ok, "failed touch should not update debounce cache")
|
||||
}
|
||||
|
||||
type touchSingleflightRepo struct {
|
||||
*apiKeyRepoStub
|
||||
mu sync.Mutex
|
||||
calls int
|
||||
blockCh chan struct{}
|
||||
}
|
||||
|
||||
func (r *touchSingleflightRepo) UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error {
|
||||
r.mu.Lock()
|
||||
r.calls++
|
||||
r.mu.Unlock()
|
||||
<-r.blockCh
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestAPIKeyService_TouchLastUsed_ConcurrentFirstTouchDeduplicated(t *testing.T) {
|
||||
repo := &touchSingleflightRepo{
|
||||
apiKeyRepoStub: &apiKeyRepoStub{},
|
||||
blockCh: make(chan struct{}),
|
||||
}
|
||||
svc := &APIKeyService{apiKeyRepo: repo}
|
||||
|
||||
const workers = 20
|
||||
startCh := make(chan struct{})
|
||||
errCh := make(chan error, workers)
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for i := 0; i < workers; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-startCh
|
||||
errCh <- svc.TouchLastUsed(context.Background(), 321)
|
||||
}()
|
||||
}
|
||||
|
||||
close(startCh)
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
repo.mu.Lock()
|
||||
defer repo.mu.Unlock()
|
||||
return repo.calls >= 1
|
||||
}, time.Second, 10*time.Millisecond)
|
||||
|
||||
close(repo.blockCh)
|
||||
wg.Wait()
|
||||
close(errCh)
|
||||
|
||||
for err := range errCh {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
repo.mu.Lock()
|
||||
defer repo.mu.Unlock()
|
||||
require.Equal(t, 1, repo.calls, "并发首次 touch 只应写库一次")
|
||||
}
|
||||
Reference in New Issue
Block a user