shaw
86fb4781f4
UserRepository.Update and APIKeyRepository.Update rewrote the whole row on every call, regardless of which fields the caller meant to change. Several columns on those tables are maintained by dedicated atomic paths (balance deduction, quota and rate-limit counters, limit adjustments, activity timestamps), so a caller holding a slightly older snapshot could silently roll them back - a lost update. Both methods now take an explicit column mask and persist only the columns the caller declares; everything else keeps its current database value. - All user and API-key call sites declare exactly what they mutate, which turns admin edits and profile saves into genuine partial updates. - Email uniqueness locking/lookup and allowed_groups sync only run when those fields are part of the update. - UserUpdateFields deliberately has no balance/total_recharged members, so Update cannot touch them. New AdjustBalance/SetBalance apply the change in a single statement and return before/after values; admin balance adjustment uses them instead of read-modify-write. - promo_codes.used_count is no longer written by Update; it is only ever incremented by the redemption path. - The billing hot path that marks an API key quota-exhausted writes only status. - Dropped a no-op row write in RevokeAllUserTokens: users has no token_version column, so it persisted nothing while still overwriting concurrently-updated columns. Adds integration coverage that a stale snapshot cannot revert concurrent atomic writes, and unit coverage pinning the column set each entry point declares.
907 lines
31 KiB
Go
907 lines
31 KiB
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync/atomic"
|
|
"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"
|
|
)
|
|
|
|
func TestGoogleAPIKeyAuthRejectsOversizedCredentialsBeforeLookup(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
var calls atomic.Int32
|
|
repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
|
|
calls.Add(1)
|
|
return nil, service.ErrAPIKeyNotFound
|
|
}}
|
|
cfg := &config.Config{RunMode: config.RunModeSimple}
|
|
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
|
|
r := gin.New()
|
|
var reason IngressRejectReason
|
|
var rejected bool
|
|
r.Use(func(c *gin.Context) {
|
|
c.Next()
|
|
reason, rejected = GetIngressRejectReason(c)
|
|
})
|
|
r.Use(APIKeyAuthGoogle(svc, cfg))
|
|
r.GET("/v1beta/models", func(c *gin.Context) { c.Status(http.StatusOK) })
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
|
|
req.Header.Set("x-goog-api-key", strings.Repeat("x", service.MaxAPIKeyCredentialBytes+1))
|
|
r.ServeHTTP(w, req)
|
|
require.Equal(t, http.StatusUnauthorized, w.Code)
|
|
require.Zero(t, calls.Load())
|
|
require.True(t, rejected)
|
|
require.Equal(t, IngressRejectInvalidAPIKey, reason)
|
|
}
|
|
|
|
func TestGoogleAPIKeyAuthMarksLookupBulkheadRejection(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
|
|
return nil, service.ErrAPIKeyAuthOverloaded
|
|
}}
|
|
cfg := &config.Config{RunMode: config.RunModeSimple}
|
|
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
|
|
r := gin.New()
|
|
var reason IngressRejectReason
|
|
var rejected bool
|
|
r.Use(func(c *gin.Context) {
|
|
c.Next()
|
|
reason, rejected = GetIngressRejectReason(c)
|
|
})
|
|
r.Use(APIKeyAuthGoogle(svc, cfg))
|
|
r.GET("/v1beta/models", func(c *gin.Context) { c.Status(http.StatusOK) })
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
|
|
req.Header.Set("x-goog-api-key", "valid-shape")
|
|
r.ServeHTTP(w, req)
|
|
require.Equal(t, http.StatusServiceUnavailable, w.Code)
|
|
require.True(t, rejected)
|
|
require.Equal(t, IngressRejectAPIKeyAuthOverloaded, reason)
|
|
}
|
|
|
|
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 {
|
|
getByID func(ctx context.Context, id int64) (*service.UserSubscription, error)
|
|
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, _ service.APIKeyUpdateFields) 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) {
|
|
if f.getByID != nil {
|
|
return f.getByID(ctx, id)
|
|
}
|
|
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) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time) error {
|
|
return errors.New("not implemented")
|
|
}
|
|
func (f fakeGoogleSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, _ *time.Time, 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, _ *time.Time, 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, _ *time.Time, 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
|
|
},
|
|
})
|
|
var rejectReason IngressRejectReason
|
|
var rejected bool
|
|
r.Use(func(c *gin.Context) {
|
|
c.Next()
|
|
rejectReason, rejected = GetIngressRejectReason(c)
|
|
})
|
|
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)
|
|
require.True(t, rejected)
|
|
require.Equal(t, IngressRejectInvalidAPIKey, rejectReason)
|
|
}
|
|
|
|
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
|
|
var rejectReason IngressRejectReason
|
|
var rejected bool
|
|
r.Use(func(c *gin.Context) {
|
|
c.Next()
|
|
markedBusinessLimited = service.HasOpsClientBusinessLimited(c)
|
|
rejectReason, rejected = GetIngressRejectReason(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)
|
|
require.True(t, rejected)
|
|
require.Equal(t, IngressRejectGroupDeleted, rejectReason)
|
|
}
|
|
|
|
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")
|
|
}
|