api_key_auth_test.go 55 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794
  1. //go:build unit
  2. package middleware
  3. import (
  4. "context"
  5. "encoding/json"
  6. "errors"
  7. "net/http"
  8. "net/http/httptest"
  9. "strings"
  10. "sync/atomic"
  11. "testing"
  12. "time"
  13. "github.com/Wei-Shaw/sub2api/internal/config"
  14. "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
  15. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  16. "github.com/Wei-Shaw/sub2api/internal/service"
  17. "github.com/gin-gonic/gin"
  18. "github.com/stretchr/testify/require"
  19. )
  20. func TestAPIKeyAuthRejectsOversizedCredentialsBeforeLookup(t *testing.T) {
  21. gin.SetMode(gin.TestMode)
  22. var calls atomic.Int32
  23. repo := &stubApiKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
  24. calls.Add(1)
  25. return nil, service.ErrAPIKeyNotFound
  26. }}
  27. cfg := &config.Config{RunMode: config.RunModeSimple}
  28. svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
  29. for _, headers := range []map[string]string{
  30. {"x-api-key": strings.Repeat("x", service.MaxAPIKeyCredentialBytes+1)},
  31. {"Authorization": "Bearer " + strings.Repeat("x", service.MaxAPIKeyCredentialBytes+1)},
  32. {"Authorization": strings.Repeat("x", maxAPIKeyAuthorizationHeaderBytes+1)},
  33. } {
  34. r := gin.New()
  35. r.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(svc, nil, cfg)))
  36. r.GET("/t", func(c *gin.Context) { c.Status(http.StatusOK) })
  37. w := httptest.NewRecorder()
  38. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  39. for name, value := range headers {
  40. req.Header.Set(name, value)
  41. }
  42. r.ServeHTTP(w, req)
  43. require.Equal(t, http.StatusUnauthorized, w.Code)
  44. }
  45. require.Zero(t, calls.Load())
  46. }
  47. func TestSimpleModeBypassesQuotaCheck(t *testing.T) {
  48. gin.SetMode(gin.TestMode)
  49. limit := 1.0
  50. group := &service.Group{
  51. ID: 42,
  52. Name: "sub",
  53. Status: service.StatusActive,
  54. Hydrated: true,
  55. SubscriptionType: service.SubscriptionTypeSubscription,
  56. DailyLimitUSD: &limit,
  57. }
  58. user := &service.User{
  59. ID: 7,
  60. Role: service.RoleUser,
  61. Status: service.StatusActive,
  62. Balance: 10,
  63. Concurrency: 3,
  64. }
  65. apiKey := &service.APIKey{
  66. ID: 100,
  67. UserID: user.ID,
  68. Key: "test-key",
  69. Status: service.StatusActive,
  70. User: user,
  71. Group: group,
  72. }
  73. apiKey.GroupID = &group.ID
  74. apiKeyRepo := &stubApiKeyRepo{
  75. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  76. if key != apiKey.Key {
  77. return nil, service.ErrAPIKeyNotFound
  78. }
  79. clone := *apiKey
  80. return &clone, nil
  81. },
  82. }
  83. t.Run("standard_mode_completes_maintenance_before_request", func(t *testing.T) {
  84. cfg := &config.Config{RunMode: config.RunModeStandard}
  85. cfg.SubscriptionMaintenance.WorkerCount = 1
  86. cfg.SubscriptionMaintenance.QueueSize = 1
  87. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  88. past := time.Now().Add(-48 * time.Hour)
  89. sub := &service.UserSubscription{
  90. ID: 55,
  91. UserID: user.ID,
  92. GroupID: group.ID,
  93. Status: service.SubscriptionStatusActive,
  94. ExpiresAt: time.Now().Add(24 * time.Hour),
  95. DailyWindowStart: &past,
  96. WeeklyWindowStart: &past,
  97. MonthlyWindowStart: &past,
  98. DailyUsageUSD: 0,
  99. }
  100. maintenanceCalled := make(chan struct{}, 1)
  101. subscriptionRepo := &stubUserSubscriptionRepo{
  102. getByID: func(ctx context.Context, id int64) (*service.UserSubscription, error) {
  103. clone := *sub
  104. return &clone, nil
  105. },
  106. getActive: func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
  107. clone := *sub
  108. return &clone, nil
  109. },
  110. updateStatus: func(ctx context.Context, subscriptionID int64, status string) error { return nil },
  111. activateWindow: func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error { return nil },
  112. resetDaily: func(ctx context.Context, id int64, start time.Time) error {
  113. sub.DailyWindowStart = &start
  114. sub.DailyUsageUSD = 0
  115. maintenanceCalled <- struct{}{}
  116. return nil
  117. },
  118. resetWeekly: func(ctx context.Context, id int64, start time.Time) error {
  119. sub.WeeklyWindowStart = &start
  120. return nil
  121. },
  122. resetMonthly: func(ctx context.Context, id int64, start time.Time) error {
  123. sub.MonthlyWindowStart = &start
  124. return nil
  125. },
  126. }
  127. subscriptionService := service.NewSubscriptionService(nil, subscriptionRepo, nil, nil, cfg)
  128. t.Cleanup(subscriptionService.Stop)
  129. router := newAuthTestRouter(apiKeyService, subscriptionService, cfg)
  130. w := httptest.NewRecorder()
  131. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  132. req.Header.Set("x-api-key", apiKey.Key)
  133. router.ServeHTTP(w, req)
  134. require.Equal(t, http.StatusOK, w.Code)
  135. select {
  136. case <-maintenanceCalled:
  137. // ok
  138. case <-time.After(time.Second):
  139. t.Fatalf("expected maintenance to complete before response")
  140. }
  141. })
  142. t.Run("standard_mode_revalidates_cas_loser_from_database", func(t *testing.T) {
  143. cfg := &config.Config{RunMode: config.RunModeStandard}
  144. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  145. past := time.Now().Add(-48 * time.Hour)
  146. current := time.Now()
  147. stale := &service.UserSubscription{
  148. ID: 56,
  149. UserID: user.ID,
  150. GroupID: group.ID,
  151. Status: service.SubscriptionStatusActive,
  152. ExpiresAt: current.Add(24 * time.Hour),
  153. DailyWindowStart: &past,
  154. WeeklyWindowStart: &past,
  155. MonthlyWindowStart: &past,
  156. DailyUsageUSD: 10,
  157. }
  158. fresh := *stale
  159. fresh.DailyWindowStart = &current
  160. fresh.WeeklyWindowStart = &current
  161. fresh.MonthlyWindowStart = &current
  162. fresh.DailyUsageUSD = 2
  163. subscriptionRepo := &stubUserSubscriptionRepo{
  164. getActive: func(context.Context, int64, int64) (*service.UserSubscription, error) {
  165. clone := *stale
  166. return &clone, nil
  167. },
  168. getByID: func(context.Context, int64) (*service.UserSubscription, error) {
  169. clone := fresh
  170. return &clone, nil
  171. },
  172. resetDaily: func(context.Context, int64, time.Time) error { return nil },
  173. resetWeekly: func(context.Context, int64, time.Time) error { return nil },
  174. resetMonthly: func(context.Context, int64, time.Time) error { return nil },
  175. }
  176. subscriptionService := service.NewSubscriptionService(nil, subscriptionRepo, nil, nil, cfg)
  177. router := newAuthTestRouter(apiKeyService, subscriptionService, cfg)
  178. w := httptest.NewRecorder()
  179. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  180. req.Header.Set("x-api-key", apiKey.Key)
  181. router.ServeHTTP(w, req)
  182. require.Equal(t, http.StatusTooManyRequests, w.Code)
  183. })
  184. t.Run("simple_mode_bypasses_quota_check", func(t *testing.T) {
  185. cfg := &config.Config{RunMode: config.RunModeSimple}
  186. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  187. subscriptionService := service.NewSubscriptionService(nil, &stubUserSubscriptionRepo{}, nil, nil, cfg)
  188. router := newAuthTestRouter(apiKeyService, subscriptionService, cfg)
  189. w := httptest.NewRecorder()
  190. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  191. req.Header.Set("x-api-key", apiKey.Key)
  192. router.ServeHTTP(w, req)
  193. require.Equal(t, http.StatusOK, w.Code)
  194. })
  195. t.Run("simple_mode_accepts_lowercase_bearer", func(t *testing.T) {
  196. cfg := &config.Config{RunMode: config.RunModeSimple}
  197. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  198. subscriptionService := service.NewSubscriptionService(nil, &stubUserSubscriptionRepo{}, nil, nil, cfg)
  199. router := newAuthTestRouter(apiKeyService, subscriptionService, cfg)
  200. w := httptest.NewRecorder()
  201. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  202. req.Header.Set("Authorization", "bearer "+apiKey.Key)
  203. router.ServeHTTP(w, req)
  204. require.Equal(t, http.StatusOK, w.Code)
  205. })
  206. t.Run("standard_mode_enforces_quota_check", func(t *testing.T) {
  207. cfg := &config.Config{RunMode: config.RunModeStandard}
  208. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  209. now := time.Now()
  210. sub := &service.UserSubscription{
  211. ID: 55,
  212. UserID: user.ID,
  213. GroupID: group.ID,
  214. Status: service.SubscriptionStatusActive,
  215. ExpiresAt: now.Add(24 * time.Hour),
  216. DailyWindowStart: &now,
  217. DailyUsageUSD: 10,
  218. }
  219. subscriptionRepo := &stubUserSubscriptionRepo{
  220. getActive: func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
  221. if userID != sub.UserID || groupID != sub.GroupID {
  222. return nil, service.ErrSubscriptionNotFound
  223. }
  224. clone := *sub
  225. return &clone, nil
  226. },
  227. updateStatus: func(ctx context.Context, subscriptionID int64, status string) error { return nil },
  228. activateWindow: func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error { return nil },
  229. resetDaily: func(ctx context.Context, id int64, start time.Time) error { return nil },
  230. resetWeekly: func(ctx context.Context, id int64, start time.Time) error { return nil },
  231. resetMonthly: func(ctx context.Context, id int64, start time.Time) error { return nil },
  232. }
  233. subscriptionService := service.NewSubscriptionService(nil, subscriptionRepo, nil, nil, cfg)
  234. router := newAuthTestRouter(apiKeyService, subscriptionService, cfg)
  235. w := httptest.NewRecorder()
  236. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  237. req.Header.Set("x-api-key", apiKey.Key)
  238. router.ServeHTTP(w, req)
  239. require.Equal(t, http.StatusTooManyRequests, w.Code)
  240. require.Contains(t, w.Body.String(), "USAGE_LIMIT_EXCEEDED")
  241. })
  242. }
  243. func TestAPIKeyAuthSetsGroupContext(t *testing.T) {
  244. gin.SetMode(gin.TestMode)
  245. group := &service.Group{
  246. ID: 101,
  247. Name: "g1",
  248. Status: service.StatusActive,
  249. Platform: service.PlatformAnthropic,
  250. Hydrated: true,
  251. }
  252. user := &service.User{
  253. ID: 7,
  254. Role: service.RoleUser,
  255. Status: service.StatusActive,
  256. Balance: 10,
  257. Concurrency: 3,
  258. }
  259. apiKey := &service.APIKey{
  260. ID: 100,
  261. UserID: user.ID,
  262. Key: "test-key",
  263. Status: service.StatusActive,
  264. User: user,
  265. Group: group,
  266. }
  267. apiKey.GroupID = &group.ID
  268. apiKeyRepo := &stubApiKeyRepo{
  269. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  270. if key != apiKey.Key {
  271. return nil, service.ErrAPIKeyNotFound
  272. }
  273. clone := *apiKey
  274. return &clone, nil
  275. },
  276. }
  277. cfg := &config.Config{RunMode: config.RunModeSimple}
  278. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  279. router := gin.New()
  280. router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg)))
  281. router.GET("/t", func(c *gin.Context) {
  282. groupFromCtx, ok := c.Request.Context().Value(ctxkey.Group).(*service.Group)
  283. if !ok || groupFromCtx == nil || groupFromCtx.ID != group.ID {
  284. c.JSON(http.StatusInternalServerError, gin.H{"ok": false})
  285. return
  286. }
  287. userIDFromCtx, ok := c.Request.Context().Value(ctxkey.UserID).(int64)
  288. if !ok || userIDFromCtx != user.ID {
  289. c.JSON(http.StatusInternalServerError, gin.H{"ok": false})
  290. return
  291. }
  292. c.JSON(http.StatusOK, gin.H{"ok": true})
  293. })
  294. w := httptest.NewRecorder()
  295. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  296. req.Header.Set("x-api-key", apiKey.Key)
  297. router.ServeHTTP(w, req)
  298. require.Equal(t, http.StatusOK, w.Code)
  299. }
  300. func TestAPIKeyAuthRejectsExclusiveGroupWhenUserNoLongerAllowed(t *testing.T) {
  301. gin.SetMode(gin.TestMode)
  302. group := &service.Group{
  303. ID: 202,
  304. Name: "exclusive",
  305. Status: service.StatusActive,
  306. IsExclusive: true,
  307. Hydrated: true,
  308. }
  309. user := &service.User{
  310. ID: 7,
  311. Role: service.RoleUser,
  312. Status: service.StatusActive,
  313. Balance: 10,
  314. Concurrency: 3,
  315. AllowedGroups: []int64{},
  316. }
  317. apiKey := &service.APIKey{
  318. ID: 100,
  319. UserID: user.ID,
  320. Key: "test-key",
  321. Status: service.StatusActive,
  322. User: user,
  323. Group: group,
  324. }
  325. apiKey.GroupID = &group.ID
  326. apiKeyRepo := &stubApiKeyRepo{
  327. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  328. if key != apiKey.Key {
  329. return nil, service.ErrAPIKeyNotFound
  330. }
  331. clone := *apiKey
  332. return &clone, nil
  333. },
  334. }
  335. cfg := &config.Config{RunMode: config.RunModeSimple}
  336. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  337. router := newAuthTestRouter(apiKeyService, nil, cfg)
  338. w := httptest.NewRecorder()
  339. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  340. req.Header.Set("x-api-key", apiKey.Key)
  341. router.ServeHTTP(w, req)
  342. require.Equal(t, http.StatusForbidden, w.Code)
  343. require.Contains(t, w.Body.String(), "GROUP_NOT_ALLOWED")
  344. }
  345. func TestAPIKeyAuthOverwritesInvalidContextGroup(t *testing.T) {
  346. gin.SetMode(gin.TestMode)
  347. group := &service.Group{
  348. ID: 101,
  349. Name: "g1",
  350. Status: service.StatusActive,
  351. Platform: service.PlatformAnthropic,
  352. Hydrated: true,
  353. }
  354. user := &service.User{
  355. ID: 7,
  356. Role: service.RoleUser,
  357. Status: service.StatusActive,
  358. Balance: 10,
  359. Concurrency: 3,
  360. }
  361. apiKey := &service.APIKey{
  362. ID: 100,
  363. UserID: user.ID,
  364. Key: "test-key",
  365. Status: service.StatusActive,
  366. User: user,
  367. Group: group,
  368. }
  369. apiKey.GroupID = &group.ID
  370. apiKeyRepo := &stubApiKeyRepo{
  371. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  372. if key != apiKey.Key {
  373. return nil, service.ErrAPIKeyNotFound
  374. }
  375. clone := *apiKey
  376. return &clone, nil
  377. },
  378. }
  379. cfg := &config.Config{RunMode: config.RunModeSimple}
  380. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  381. router := gin.New()
  382. router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg)))
  383. invalidGroup := &service.Group{
  384. ID: group.ID,
  385. Platform: group.Platform,
  386. Status: group.Status,
  387. }
  388. router.GET("/t", func(c *gin.Context) {
  389. groupFromCtx, ok := c.Request.Context().Value(ctxkey.Group).(*service.Group)
  390. if !ok || groupFromCtx == nil || groupFromCtx.ID != group.ID || !groupFromCtx.Hydrated || groupFromCtx == invalidGroup {
  391. c.JSON(http.StatusInternalServerError, gin.H{"ok": false})
  392. return
  393. }
  394. c.JSON(http.StatusOK, gin.H{"ok": true})
  395. })
  396. w := httptest.NewRecorder()
  397. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  398. req.Header.Set("x-api-key", apiKey.Key)
  399. req = req.WithContext(context.WithValue(req.Context(), ctxkey.Group, invalidGroup))
  400. router.ServeHTTP(w, req)
  401. require.Equal(t, http.StatusOK, w.Code)
  402. }
  403. func TestAPIKeyAuthRejectsUnavailableGroup(t *testing.T) {
  404. gin.SetMode(gin.TestMode)
  405. groupID := int64(101)
  406. user := &service.User{
  407. ID: 7,
  408. Role: service.RoleUser,
  409. Status: service.StatusActive,
  410. Balance: 10,
  411. Concurrency: 3,
  412. }
  413. tests := []struct {
  414. name string
  415. group *service.Group
  416. wantStatus int
  417. wantCode string
  418. wantMarked bool
  419. wantReject IngressRejectReason
  420. }{
  421. {
  422. name: "active group passes",
  423. group: &service.Group{
  424. ID: groupID,
  425. Name: "active",
  426. Status: service.StatusActive,
  427. Platform: service.PlatformAnthropic,
  428. Hydrated: true,
  429. },
  430. wantStatus: http.StatusOK,
  431. },
  432. {
  433. name: "disabled group is forbidden",
  434. group: &service.Group{
  435. ID: groupID,
  436. Name: "disabled",
  437. Status: service.StatusDisabled,
  438. Platform: service.PlatformAnthropic,
  439. Hydrated: true,
  440. },
  441. wantStatus: http.StatusForbidden,
  442. wantCode: "GROUP_DISABLED",
  443. wantMarked: true,
  444. wantReject: IngressRejectGroupDisabled,
  445. },
  446. {
  447. name: "deleted status group is forbidden",
  448. group: &service.Group{
  449. ID: groupID,
  450. Name: "deleted",
  451. Status: "deleted",
  452. Platform: service.PlatformAnthropic,
  453. Hydrated: true,
  454. },
  455. wantStatus: http.StatusForbidden,
  456. wantCode: "GROUP_DELETED",
  457. wantMarked: true,
  458. wantReject: IngressRejectGroupDeleted,
  459. },
  460. {
  461. name: "missing group edge is forbidden",
  462. group: nil,
  463. wantStatus: http.StatusForbidden,
  464. wantCode: "GROUP_DELETED",
  465. wantMarked: true,
  466. wantReject: IngressRejectGroupDeleted,
  467. },
  468. }
  469. for _, tt := range tests {
  470. t.Run(tt.name, func(t *testing.T) {
  471. apiKey := &service.APIKey{
  472. ID: 100,
  473. UserID: user.ID,
  474. GroupID: &groupID,
  475. Key: "test-key",
  476. Status: service.StatusActive,
  477. User: user,
  478. Group: tt.group,
  479. }
  480. apiKeyRepo := &stubApiKeyRepo{
  481. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  482. if key != apiKey.Key {
  483. return nil, service.ErrAPIKeyNotFound
  484. }
  485. clone := *apiKey
  486. return &clone, nil
  487. },
  488. }
  489. cfg := &config.Config{RunMode: config.RunModeStandard}
  490. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  491. router := gin.New()
  492. var markedBusinessLimited bool
  493. var businessLimitedReason string
  494. var rejectReason IngressRejectReason
  495. var rejected bool
  496. router.Use(func(c *gin.Context) {
  497. c.Next()
  498. markedBusinessLimited = service.HasOpsClientBusinessLimited(c)
  499. rejectReason, rejected = GetIngressRejectReason(c)
  500. if v, ok := c.Get(service.OpsClientBusinessLimitedReasonKey); ok {
  501. businessLimitedReason, _ = v.(string)
  502. }
  503. })
  504. router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg)))
  505. router.GET("/t", func(c *gin.Context) {
  506. c.JSON(http.StatusOK, gin.H{"ok": true})
  507. })
  508. w := httptest.NewRecorder()
  509. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  510. req.Header.Set("x-api-key", apiKey.Key)
  511. router.ServeHTTP(w, req)
  512. require.Equal(t, tt.wantStatus, w.Code)
  513. if tt.wantCode != "" {
  514. require.Contains(t, w.Body.String(), tt.wantCode)
  515. }
  516. require.Equal(t, tt.wantMarked, markedBusinessLimited)
  517. require.Equal(t, tt.wantReject != "", rejected)
  518. require.Equal(t, tt.wantReject, rejectReason)
  519. if tt.wantMarked {
  520. require.Equal(t, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable, businessLimitedReason)
  521. }
  522. })
  523. }
  524. }
  525. func TestAPIKeyAuthMarksOnlyExpectedIngressRejections(t *testing.T) {
  526. gin.SetMode(gin.TestMode)
  527. tests := []struct {
  528. name string
  529. path string
  530. key string
  531. authHeader string
  532. repoErr error
  533. wantStatus int
  534. wantCode string
  535. wantReason IngressRejectReason
  536. }{
  537. {
  538. name: "query key deprecated",
  539. path: "/t?key=legacy",
  540. wantStatus: http.StatusBadRequest,
  541. wantCode: "api_key_in_query_deprecated",
  542. wantReason: IngressRejectQueryAPIKeyDeprecated,
  543. },
  544. {
  545. name: "missing key",
  546. path: "/t",
  547. wantStatus: http.StatusUnauthorized,
  548. wantCode: "API_KEY_REQUIRED",
  549. wantReason: IngressRejectAPIKeyRequired,
  550. },
  551. {
  552. name: "malformed authorization",
  553. path: "/t",
  554. authHeader: "Basic not-a-bearer-key",
  555. wantStatus: http.StatusUnauthorized,
  556. wantCode: "API_KEY_REQUIRED",
  557. wantReason: IngressRejectInvalidAPIKey,
  558. },
  559. {
  560. name: "oversized key",
  561. path: "/t",
  562. key: strings.Repeat("x", service.MaxAPIKeyCredentialBytes+1),
  563. wantStatus: http.StatusUnauthorized,
  564. wantCode: "INVALID_API_KEY",
  565. wantReason: IngressRejectInvalidAPIKey,
  566. },
  567. {
  568. name: "invalid key",
  569. path: "/t",
  570. key: "invalid",
  571. repoErr: service.ErrAPIKeyNotFound,
  572. wantStatus: http.StatusUnauthorized,
  573. wantCode: "INVALID_API_KEY",
  574. wantReason: IngressRejectInvalidAPIKey,
  575. },
  576. {
  577. name: "repository failure remains operational error",
  578. path: "/t",
  579. key: "valid-shape",
  580. repoErr: errors.New("database unavailable"),
  581. wantStatus: http.StatusInternalServerError,
  582. wantCode: "INTERNAL_ERROR",
  583. },
  584. {
  585. name: "auth lookup bulkhead rejection is an admission rejection",
  586. path: "/t",
  587. key: "valid-shape",
  588. repoErr: service.ErrAPIKeyAuthOverloaded,
  589. wantStatus: http.StatusServiceUnavailable,
  590. wantCode: "API_KEY_AUTH_OVERLOADED",
  591. wantReason: IngressRejectAPIKeyAuthOverloaded,
  592. },
  593. }
  594. for _, tt := range tests {
  595. t.Run(tt.name, func(t *testing.T) {
  596. repo := &stubApiKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
  597. return nil, tt.repoErr
  598. }}
  599. cfg := &config.Config{RunMode: config.RunModeSimple}
  600. apiKeyService := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
  601. router := gin.New()
  602. var reason IngressRejectReason
  603. var rejected bool
  604. router.Use(func(c *gin.Context) {
  605. c.Next()
  606. reason, rejected = GetIngressRejectReason(c)
  607. })
  608. router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg)))
  609. router.GET("/t", func(c *gin.Context) { c.Status(http.StatusOK) })
  610. w := httptest.NewRecorder()
  611. req := httptest.NewRequest(http.MethodGet, tt.path, nil)
  612. if tt.key != "" {
  613. req.Header.Set("x-api-key", tt.key)
  614. }
  615. if tt.authHeader != "" {
  616. req.Header.Set("Authorization", tt.authHeader)
  617. }
  618. router.ServeHTTP(w, req)
  619. require.Equal(t, tt.wantStatus, w.Code)
  620. require.Contains(t, w.Body.String(), tt.wantCode)
  621. require.Equal(t, tt.wantReason != "", rejected)
  622. require.Equal(t, tt.wantReason, reason)
  623. })
  624. }
  625. }
  626. func TestAPIKeyAuthSetsOpsFallbackKeyOnEarlyAbort(t *testing.T) {
  627. gin.SetMode(gin.TestMode)
  628. groupID := int64(101)
  629. user := &service.User{
  630. ID: 7,
  631. Role: service.RoleUser,
  632. Status: service.StatusActive,
  633. Balance: 10,
  634. Concurrency: 3,
  635. }
  636. apiKey := &service.APIKey{
  637. ID: 100,
  638. UserID: user.ID,
  639. GroupID: &groupID,
  640. Key: "test-key",
  641. Status: service.StatusActive,
  642. User: user,
  643. Group: &service.Group{
  644. ID: groupID,
  645. Name: "disabled",
  646. Status: service.StatusDisabled,
  647. Platform: service.PlatformAnthropic,
  648. Hydrated: true,
  649. },
  650. }
  651. apiKeyRepo := &stubApiKeyRepo{
  652. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  653. if key != apiKey.Key {
  654. return nil, service.ErrAPIKeyNotFound
  655. }
  656. clone := *apiKey
  657. return &clone, nil
  658. },
  659. }
  660. cfg := &config.Config{RunMode: config.RunModeStandard}
  661. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  662. router := gin.New()
  663. var fallback *service.APIKey
  664. var fallbackOK bool
  665. router.Use(func(c *gin.Context) {
  666. c.Next()
  667. fallback, fallbackOK = GetOpsFallbackAPIKey(c)
  668. })
  669. router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg)))
  670. router.GET("/t", func(c *gin.Context) {
  671. c.JSON(http.StatusOK, gin.H{"ok": true})
  672. })
  673. w := httptest.NewRecorder()
  674. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  675. req.Header.Set("x-api-key", apiKey.Key)
  676. router.ServeHTTP(w, req)
  677. // 分组停用 → 早退中断,但 ops fallback key 仍应写入,含 user/group/platform。
  678. require.Equal(t, http.StatusForbidden, w.Code)
  679. require.Contains(t, w.Body.String(), "GROUP_DISABLED")
  680. require.True(t, fallbackOK, "鉴权早退时也应写入 ops fallback api key")
  681. require.NotNil(t, fallback)
  682. require.Equal(t, apiKey.ID, fallback.ID)
  683. require.NotNil(t, fallback.User)
  684. require.Equal(t, user.ID, fallback.User.ID)
  685. require.NotNil(t, fallback.GroupID)
  686. require.Equal(t, groupID, *fallback.GroupID)
  687. require.NotNil(t, fallback.Group)
  688. require.Equal(t, service.PlatformAnthropic, fallback.Group.Platform)
  689. }
  690. func TestAPIKeyAuthGoogleSetsOpsFallbackKeyOnEarlyAbort(t *testing.T) {
  691. gin.SetMode(gin.TestMode)
  692. groupID := int64(202)
  693. user := &service.User{
  694. ID: 9,
  695. Role: service.RoleUser,
  696. Status: service.StatusActive,
  697. Balance: 10,
  698. Concurrency: 3,
  699. }
  700. apiKey := &service.APIKey{
  701. ID: 200,
  702. UserID: user.ID,
  703. GroupID: &groupID,
  704. Key: "g-key",
  705. Status: service.StatusActive,
  706. User: user,
  707. Group: &service.Group{
  708. ID: groupID,
  709. Name: "disabled",
  710. Status: service.StatusDisabled,
  711. Platform: service.PlatformGemini,
  712. Hydrated: true,
  713. },
  714. }
  715. apiKeyRepo := &stubApiKeyRepo{
  716. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  717. if key != apiKey.Key {
  718. return nil, service.ErrAPIKeyNotFound
  719. }
  720. clone := *apiKey
  721. return &clone, nil
  722. },
  723. }
  724. cfg := &config.Config{RunMode: config.RunModeStandard}
  725. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  726. router := gin.New()
  727. var fallback *service.APIKey
  728. var fallbackOK bool
  729. router.Use(func(c *gin.Context) {
  730. c.Next()
  731. fallback, fallbackOK = GetOpsFallbackAPIKey(c)
  732. })
  733. router.Use(gin.HandlerFunc(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg)))
  734. router.GET("/t", func(c *gin.Context) {
  735. c.JSON(http.StatusOK, gin.H{"ok": true})
  736. })
  737. w := httptest.NewRecorder()
  738. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  739. req.Header.Set("x-goog-api-key", apiKey.Key)
  740. router.ServeHTTP(w, req)
  741. require.Equal(t, http.StatusForbidden, w.Code)
  742. require.True(t, fallbackOK, "Google 鉴权早退时也应写入 ops fallback api key")
  743. require.NotNil(t, fallback)
  744. require.Equal(t, apiKey.ID, fallback.ID)
  745. require.NotNil(t, fallback.User)
  746. require.Equal(t, user.ID, fallback.User.ID)
  747. }
  748. func TestRequireGroupAssignmentMarksUngroupedKeyBusinessLimited(t *testing.T) {
  749. gin.SetMode(gin.TestMode)
  750. settingService := service.NewSettingService(fakeSettingRepo{
  751. values: map[string]string{
  752. service.SettingKeyAllowUngroupedKeyScheduling: "false",
  753. },
  754. }, &config.Config{})
  755. apiKey := &service.APIKey{
  756. ID: 100,
  757. Key: "ungrouped-key",
  758. Status: service.StatusActive,
  759. }
  760. router := gin.New()
  761. var markedBusinessLimited bool
  762. var businessLimitedReason string
  763. var rejectReason IngressRejectReason
  764. var rejected bool
  765. router.Use(func(c *gin.Context) {
  766. c.Next()
  767. markedBusinessLimited = service.HasOpsClientBusinessLimited(c)
  768. rejectReason, rejected = GetIngressRejectReason(c)
  769. if v, ok := c.Get(service.OpsClientBusinessLimitedReasonKey); ok {
  770. businessLimitedReason, _ = v.(string)
  771. }
  772. })
  773. router.Use(func(c *gin.Context) {
  774. c.Set(string(ContextKeyAPIKey), apiKey)
  775. c.Next()
  776. })
  777. router.Use(RequireGroupAssignment(settingService, AnthropicErrorWriter))
  778. router.GET("/t", func(c *gin.Context) {
  779. c.JSON(http.StatusOK, gin.H{"ok": true})
  780. })
  781. w := httptest.NewRecorder()
  782. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  783. router.ServeHTTP(w, req)
  784. require.Equal(t, http.StatusForbidden, w.Code)
  785. require.Contains(t, w.Body.String(), "not assigned to any group")
  786. require.True(t, rejected)
  787. require.Equal(t, IngressRejectGroupUnassigned, rejectReason)
  788. require.True(t, markedBusinessLimited)
  789. require.Equal(t, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnassigned, businessLimitedReason)
  790. }
  791. func TestAPIKeyAuthIPRestrictionUsesTrustedPathWhenSwitchDisabled(t *testing.T) {
  792. gin.SetMode(gin.TestMode)
  793. user := &service.User{
  794. ID: 7,
  795. Role: service.RoleUser,
  796. Status: service.StatusActive,
  797. Balance: 10,
  798. Concurrency: 3,
  799. }
  800. apiKey := &service.APIKey{
  801. ID: 100,
  802. UserID: user.ID,
  803. Key: "test-key",
  804. Status: service.StatusActive,
  805. User: user,
  806. IPWhitelist: []string{"1.2.3.4"},
  807. }
  808. apiKeyRepo := &stubApiKeyRepo{
  809. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  810. if key != apiKey.Key {
  811. return nil, service.ErrAPIKeyNotFound
  812. }
  813. clone := *apiKey
  814. return &clone, nil
  815. },
  816. }
  817. cfg := &config.Config{RunMode: config.RunModeSimple}
  818. cfg.SetTrustForwardedIPForAPIKeyACL(false)
  819. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  820. router := gin.New()
  821. require.NoError(t, router.SetTrustedProxies(nil))
  822. var markedBusinessLimited bool
  823. var businessLimitedReason string
  824. router.Use(func(c *gin.Context) {
  825. c.Next()
  826. markedBusinessLimited = service.HasOpsClientBusinessLimited(c)
  827. if v, ok := c.Get(service.OpsClientBusinessLimitedReasonKey); ok {
  828. businessLimitedReason, _ = v.(string)
  829. }
  830. })
  831. router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg)))
  832. router.GET("/t", func(c *gin.Context) {
  833. c.JSON(http.StatusOK, gin.H{"ok": true})
  834. })
  835. w := httptest.NewRecorder()
  836. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  837. req.RemoteAddr = "9.9.9.9:12345"
  838. req.Header.Set("x-api-key", apiKey.Key)
  839. req.Header.Set("X-Forwarded-For", "1.2.3.4")
  840. req.Header.Set("X-Real-IP", "1.2.3.4")
  841. req.Header.Set("CF-Connecting-IP", "1.2.3.4")
  842. router.ServeHTTP(w, req)
  843. require.Equal(t, http.StatusForbidden, w.Code)
  844. requireAPIKeyAuthError(t, w, "ACCESS_DENIED", "Access denied. Your IP is 9.9.9.9")
  845. require.True(t, markedBusinessLimited)
  846. require.Equal(t, service.OpsClientBusinessLimitedReasonIPRestriction, businessLimitedReason)
  847. }
  848. func TestAPIKeyAuthIPRestrictionIncludesClientIPForBlacklistDenial(t *testing.T) {
  849. gin.SetMode(gin.TestMode)
  850. user := &service.User{
  851. ID: 7,
  852. Role: service.RoleUser,
  853. Status: service.StatusActive,
  854. Balance: 10,
  855. Concurrency: 3,
  856. }
  857. apiKey := &service.APIKey{
  858. ID: 100,
  859. UserID: user.ID,
  860. Key: "test-key",
  861. Status: service.StatusActive,
  862. User: user,
  863. IPBlacklist: []string{"9.9.9.9"},
  864. }
  865. apiKeyRepo := &stubApiKeyRepo{
  866. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  867. if key != apiKey.Key {
  868. return nil, service.ErrAPIKeyNotFound
  869. }
  870. clone := *apiKey
  871. return &clone, nil
  872. },
  873. }
  874. cfg := &config.Config{RunMode: config.RunModeSimple}
  875. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  876. router := gin.New()
  877. require.NoError(t, router.SetTrustedProxies(nil))
  878. router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg)))
  879. router.GET("/t", func(c *gin.Context) {
  880. c.JSON(http.StatusOK, gin.H{"ok": true})
  881. })
  882. w := httptest.NewRecorder()
  883. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  884. req.RemoteAddr = "9.9.9.9:12345"
  885. req.Header.Set("x-api-key", apiKey.Key)
  886. router.ServeHTTP(w, req)
  887. require.Equal(t, http.StatusForbidden, w.Code)
  888. requireAPIKeyAuthError(t, w, "ACCESS_DENIED", "Access denied. Your IP is 9.9.9.9")
  889. }
  890. func TestAPIKeyAuthIPRestrictionUsesConfiguredTrustedProxy(t *testing.T) {
  891. gin.SetMode(gin.TestMode)
  892. user := &service.User{
  893. ID: 7,
  894. Role: service.RoleUser,
  895. Status: service.StatusActive,
  896. Balance: 10,
  897. Concurrency: 3,
  898. }
  899. apiKey := &service.APIKey{
  900. ID: 100,
  901. UserID: user.ID,
  902. Key: "test-key",
  903. Status: service.StatusActive,
  904. User: user,
  905. IPWhitelist: []string{"1.2.3.4"},
  906. }
  907. apiKeyRepo := &stubApiKeyRepo{
  908. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  909. if key != apiKey.Key {
  910. return nil, service.ErrAPIKeyNotFound
  911. }
  912. clone := *apiKey
  913. return &clone, nil
  914. },
  915. }
  916. cfg := &config.Config{RunMode: config.RunModeSimple}
  917. cfg.SetTrustForwardedIPForAPIKeyACL(false)
  918. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  919. router := gin.New()
  920. require.NoError(t, router.SetTrustedProxies([]string{"9.9.9.9"}))
  921. router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg)))
  922. router.GET("/t", func(c *gin.Context) {
  923. c.JSON(http.StatusOK, gin.H{"ok": true})
  924. })
  925. w := httptest.NewRecorder()
  926. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  927. req.RemoteAddr = "9.9.9.9:12345"
  928. req.Header.Set("x-api-key", apiKey.Key)
  929. req.Header.Set("X-Forwarded-For", "1.2.3.4")
  930. req.Header.Set("X-Real-IP", "1.2.3.4")
  931. req.Header.Set("CF-Connecting-IP", "1.2.3.4")
  932. router.ServeHTTP(w, req)
  933. require.Equal(t, http.StatusOK, w.Code)
  934. }
  935. func TestAPIKeyAuthIPRestrictionUsesForwardedClientIPInDenialWhenTrusted(t *testing.T) {
  936. gin.SetMode(gin.TestMode)
  937. user := &service.User{
  938. ID: 7,
  939. Role: service.RoleUser,
  940. Status: service.StatusActive,
  941. Balance: 10,
  942. Concurrency: 3,
  943. }
  944. apiKey := &service.APIKey{
  945. ID: 100,
  946. UserID: user.ID,
  947. Key: "test-key",
  948. Status: service.StatusActive,
  949. User: user,
  950. IPWhitelist: []string{"9.9.9.9"},
  951. }
  952. apiKeyRepo := &stubApiKeyRepo{
  953. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  954. if key != apiKey.Key {
  955. return nil, service.ErrAPIKeyNotFound
  956. }
  957. clone := *apiKey
  958. return &clone, nil
  959. },
  960. }
  961. cfg := &config.Config{RunMode: config.RunModeSimple}
  962. cfg.SetTrustForwardedIPForAPIKeyACL(false)
  963. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  964. router := gin.New()
  965. require.NoError(t, router.SetTrustedProxies([]string{"9.9.9.9"}))
  966. router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg)))
  967. router.GET("/t", func(c *gin.Context) {
  968. c.JSON(http.StatusOK, gin.H{"ok": true})
  969. })
  970. w := httptest.NewRecorder()
  971. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  972. req.RemoteAddr = "9.9.9.9:12345"
  973. req.Header.Set("x-api-key", apiKey.Key)
  974. req.Header.Set("X-Forwarded-For", "1.2.3.4")
  975. req.Header.Set("X-Real-IP", "1.2.3.4")
  976. req.Header.Set("CF-Connecting-IP", "1.2.3.4")
  977. router.ServeHTTP(w, req)
  978. require.Equal(t, http.StatusForbidden, w.Code)
  979. requireAPIKeyAuthError(t, w, "ACCESS_DENIED", "Access denied. Your IP is 1.2.3.4")
  980. }
  981. func TestAPIKeyAuthTouchesLastUsedOnSuccess(t *testing.T) {
  982. gin.SetMode(gin.TestMode)
  983. user := &service.User{
  984. ID: 7,
  985. Role: service.RoleUser,
  986. Status: service.StatusActive,
  987. Balance: 10,
  988. Concurrency: 3,
  989. }
  990. apiKey := &service.APIKey{
  991. ID: 100,
  992. UserID: user.ID,
  993. Key: "touch-ok",
  994. Status: service.StatusActive,
  995. User: user,
  996. }
  997. var touchedID int64
  998. var touchedAt time.Time
  999. apiKeyRepo := &stubApiKeyRepo{
  1000. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  1001. if key != apiKey.Key {
  1002. return nil, service.ErrAPIKeyNotFound
  1003. }
  1004. clone := *apiKey
  1005. return &clone, nil
  1006. },
  1007. updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
  1008. touchedID = id
  1009. touchedAt = usedAt
  1010. return nil
  1011. },
  1012. }
  1013. cfg := &config.Config{RunMode: config.RunModeSimple}
  1014. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  1015. router := newAuthTestRouter(apiKeyService, nil, cfg)
  1016. w := httptest.NewRecorder()
  1017. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  1018. req.Header.Set("x-api-key", apiKey.Key)
  1019. router.ServeHTTP(w, req)
  1020. require.Equal(t, http.StatusOK, w.Code)
  1021. require.Equal(t, apiKey.ID, touchedID)
  1022. require.False(t, touchedAt.IsZero(), "expected touch timestamp")
  1023. }
  1024. func TestAPIKeyAuthTouchLastUsedFailureDoesNotBlock(t *testing.T) {
  1025. gin.SetMode(gin.TestMode)
  1026. user := &service.User{
  1027. ID: 8,
  1028. Role: service.RoleUser,
  1029. Status: service.StatusActive,
  1030. Balance: 10,
  1031. Concurrency: 3,
  1032. }
  1033. apiKey := &service.APIKey{
  1034. ID: 101,
  1035. UserID: user.ID,
  1036. Key: "touch-fail",
  1037. Status: service.StatusActive,
  1038. User: user,
  1039. }
  1040. touchCalls := 0
  1041. apiKeyRepo := &stubApiKeyRepo{
  1042. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  1043. if key != apiKey.Key {
  1044. return nil, service.ErrAPIKeyNotFound
  1045. }
  1046. clone := *apiKey
  1047. return &clone, nil
  1048. },
  1049. updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
  1050. touchCalls++
  1051. return errors.New("db unavailable")
  1052. },
  1053. }
  1054. cfg := &config.Config{RunMode: config.RunModeSimple}
  1055. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  1056. router := newAuthTestRouter(apiKeyService, nil, cfg)
  1057. w := httptest.NewRecorder()
  1058. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  1059. req.Header.Set("x-api-key", apiKey.Key)
  1060. router.ServeHTTP(w, req)
  1061. require.Equal(t, http.StatusOK, w.Code, "touch failure should not block request")
  1062. require.Equal(t, 1, touchCalls)
  1063. }
  1064. func TestAPIKeyAuthTouchesLastUsedInStandardMode(t *testing.T) {
  1065. gin.SetMode(gin.TestMode)
  1066. user := &service.User{
  1067. ID: 9,
  1068. Role: service.RoleUser,
  1069. Status: service.StatusActive,
  1070. Balance: 10,
  1071. Concurrency: 3,
  1072. }
  1073. apiKey := &service.APIKey{
  1074. ID: 102,
  1075. UserID: user.ID,
  1076. Key: "touch-standard",
  1077. Status: service.StatusActive,
  1078. User: user,
  1079. }
  1080. touchCalls := 0
  1081. apiKeyRepo := &stubApiKeyRepo{
  1082. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  1083. if key != apiKey.Key {
  1084. return nil, service.ErrAPIKeyNotFound
  1085. }
  1086. clone := *apiKey
  1087. return &clone, nil
  1088. },
  1089. updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
  1090. touchCalls++
  1091. return nil
  1092. },
  1093. }
  1094. cfg := &config.Config{RunMode: config.RunModeStandard}
  1095. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  1096. router := newAuthTestRouter(apiKeyService, nil, cfg)
  1097. w := httptest.NewRecorder()
  1098. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  1099. req.Header.Set("x-api-key", apiKey.Key)
  1100. router.ServeHTTP(w, req)
  1101. require.Equal(t, http.StatusOK, w.Code)
  1102. require.Equal(t, 1, touchCalls)
  1103. }
  1104. func TestAPIKeyAuthBillingInfoSkipsBillingAndSideEffects(t *testing.T) {
  1105. gin.SetMode(gin.TestMode)
  1106. group := &service.Group{
  1107. ID: 42,
  1108. Name: "subscription",
  1109. Status: service.StatusActive,
  1110. Hydrated: true,
  1111. SubscriptionType: service.SubscriptionTypeSubscription,
  1112. }
  1113. user := &service.User{
  1114. ID: 7,
  1115. Role: service.RoleUser,
  1116. Status: service.StatusActive,
  1117. Balance: 0,
  1118. Concurrency: 3,
  1119. }
  1120. expiredAt := time.Now().Add(-time.Hour)
  1121. apiKey := &service.APIKey{
  1122. ID: 100,
  1123. UserID: user.ID,
  1124. Key: "billing-info-auth-only",
  1125. Status: service.StatusAPIKeyQuotaExhausted,
  1126. User: user,
  1127. GroupID: &group.ID,
  1128. Group: group,
  1129. Quota: 1,
  1130. QuotaUsed: 1,
  1131. ExpiresAt: &expiredAt,
  1132. }
  1133. touchCalls := 0
  1134. subscriptionCalls := 0
  1135. apiKeyRepo := &stubApiKeyRepo{
  1136. getByKey: func(context.Context, string) (*service.APIKey, error) {
  1137. clone := *apiKey
  1138. return &clone, nil
  1139. },
  1140. updateLastUsed: func(context.Context, int64, time.Time) error {
  1141. touchCalls++
  1142. return nil
  1143. },
  1144. }
  1145. subscriptionRepo := &stubUserSubscriptionRepo{
  1146. getActive: func(context.Context, int64, int64) (*service.UserSubscription, error) {
  1147. subscriptionCalls++
  1148. return nil, service.ErrSubscriptionNotFound
  1149. },
  1150. }
  1151. cfg := &config.Config{RunMode: config.RunModeStandard}
  1152. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  1153. subscriptionService := service.NewSubscriptionService(nil, subscriptionRepo, nil, nil, cfg)
  1154. t.Cleanup(subscriptionService.Stop)
  1155. router := newAuthTestRouter(apiKeyService, subscriptionService, cfg)
  1156. w := httptest.NewRecorder()
  1157. req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil)
  1158. req.Header.Set("x-api-key", apiKey.Key)
  1159. router.ServeHTTP(w, req)
  1160. require.Equal(t, http.StatusOK, w.Code)
  1161. require.Zero(t, subscriptionCalls)
  1162. require.Zero(t, touchCalls)
  1163. }
  1164. func TestAPIKeyAuthBillingInfoSkipsLastUsedInSimpleMode(t *testing.T) {
  1165. gin.SetMode(gin.TestMode)
  1166. user := &service.User{ID: 7, Role: service.RoleUser, Status: service.StatusActive}
  1167. apiKey := &service.APIKey{ID: 100, UserID: user.ID, Key: "billing-info-simple", Status: service.StatusActive, User: user}
  1168. touchCalls := 0
  1169. apiKeyRepo := &stubApiKeyRepo{
  1170. getByKey: func(context.Context, string) (*service.APIKey, error) {
  1171. clone := *apiKey
  1172. return &clone, nil
  1173. },
  1174. updateLastUsed: func(context.Context, int64, time.Time) error {
  1175. touchCalls++
  1176. return nil
  1177. },
  1178. }
  1179. cfg := &config.Config{RunMode: config.RunModeSimple}
  1180. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  1181. router := newAuthTestRouter(apiKeyService, nil, cfg)
  1182. w := httptest.NewRecorder()
  1183. req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil)
  1184. req.Header.Set("x-api-key", apiKey.Key)
  1185. router.ServeHTTP(w, req)
  1186. require.Equal(t, http.StatusOK, w.Code)
  1187. require.Zero(t, touchCalls)
  1188. }
  1189. func TestAPIKeyAuthUsageStillTouchesLastUsed(t *testing.T) {
  1190. gin.SetMode(gin.TestMode)
  1191. user := &service.User{ID: 7, Role: service.RoleUser, Status: service.StatusActive, Balance: 10}
  1192. apiKey := &service.APIKey{ID: 100, UserID: user.ID, Key: "usage-touch", Status: service.StatusActive, User: user}
  1193. touchCalls := 0
  1194. apiKeyRepo := &stubApiKeyRepo{
  1195. getByKey: func(context.Context, string) (*service.APIKey, error) {
  1196. clone := *apiKey
  1197. return &clone, nil
  1198. },
  1199. updateLastUsed: func(context.Context, int64, time.Time) error {
  1200. touchCalls++
  1201. return nil
  1202. },
  1203. }
  1204. cfg := &config.Config{RunMode: config.RunModeStandard}
  1205. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  1206. router := newAuthTestRouter(apiKeyService, nil, cfg)
  1207. w := httptest.NewRecorder()
  1208. req := httptest.NewRequest(http.MethodGet, "/v1/usage", nil)
  1209. req.Header.Set("x-api-key", apiKey.Key)
  1210. router.ServeHTTP(w, req)
  1211. require.Equal(t, http.StatusOK, w.Code)
  1212. require.Equal(t, 1, touchCalls)
  1213. }
  1214. func TestAPIKeyAuthAllowsBalanceBelowMinimumReserve(t *testing.T) {
  1215. gin.SetMode(gin.TestMode)
  1216. user := &service.User{
  1217. ID: 10,
  1218. Role: service.RoleUser,
  1219. Status: service.StatusActive,
  1220. Balance: 0.005,
  1221. Concurrency: 3,
  1222. }
  1223. apiKey := &service.APIKey{
  1224. ID: 103,
  1225. UserID: user.ID,
  1226. Key: "held-balance-low",
  1227. Status: service.StatusActive,
  1228. User: user,
  1229. }
  1230. apiKeyRepo := &stubApiKeyRepo{
  1231. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  1232. if key != apiKey.Key {
  1233. return nil, service.ErrAPIKeyNotFound
  1234. }
  1235. clone := *apiKey
  1236. userClone := *user
  1237. clone.User = &userClone
  1238. return &clone, nil
  1239. },
  1240. }
  1241. cfg := &config.Config{RunMode: config.RunModeStandard}
  1242. cfg.Billing.MinimumBalanceReserve = 0.01
  1243. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  1244. router := newAuthTestRouter(apiKeyService, nil, cfg)
  1245. w := httptest.NewRecorder()
  1246. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  1247. req.Header.Set("x-api-key", apiKey.Key)
  1248. router.ServeHTTP(w, req)
  1249. // 鉴权层保持历史语义:MinimumBalanceReserve 只用于 billing-cache 预检,
  1250. // 0 < balance < reserve 不得被鉴权中间件硬 403(存量部署静默行为变更)。
  1251. require.Equal(t, http.StatusOK, w.Code)
  1252. }
  1253. func TestAPIKeyAuthRejectsExhaustedBalance(t *testing.T) {
  1254. gin.SetMode(gin.TestMode)
  1255. user := &service.User{
  1256. ID: 10,
  1257. Role: service.RoleUser,
  1258. Status: service.StatusActive,
  1259. Balance: 0,
  1260. Concurrency: 3,
  1261. }
  1262. apiKey := &service.APIKey{
  1263. ID: 104,
  1264. UserID: user.ID,
  1265. Key: "held-balance-zero",
  1266. Status: service.StatusActive,
  1267. User: user,
  1268. }
  1269. apiKeyRepo := &stubApiKeyRepo{
  1270. getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  1271. if key != apiKey.Key {
  1272. return nil, service.ErrAPIKeyNotFound
  1273. }
  1274. clone := *apiKey
  1275. userClone := *user
  1276. clone.User = &userClone
  1277. return &clone, nil
  1278. },
  1279. }
  1280. cfg := &config.Config{RunMode: config.RunModeStandard}
  1281. apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
  1282. router := newAuthTestRouter(apiKeyService, nil, cfg)
  1283. w := httptest.NewRecorder()
  1284. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  1285. req.Header.Set("x-api-key", apiKey.Key)
  1286. router.ServeHTTP(w, req)
  1287. require.Equal(t, http.StatusForbidden, w.Code)
  1288. requireAPIKeyAuthError(t, w, "INSUFFICIENT_BALANCE", "Insufficient account balance")
  1289. }
  1290. func TestAPIKeyAuthOpenAIQuotaErrorFormat(t *testing.T) {
  1291. gin.SetMode(gin.TestMode)
  1292. user := &service.User{ID: 11, Role: service.RoleUser, Status: service.StatusActive, Balance: 10}
  1293. group := &service.Group{ID: 8, Platform: service.PlatformOpenAI, Status: service.StatusActive}
  1294. apiKey := &service.APIKey{
  1295. ID: 105, UserID: user.ID, Key: "openai-quota-exhausted", Status: service.StatusAPIKeyQuotaExhausted,
  1296. User: user, Group: group, GroupID: &group.ID,
  1297. }
  1298. apiKeyRepo := &stubApiKeyRepo{getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  1299. if key != apiKey.Key {
  1300. return nil, service.ErrAPIKeyNotFound
  1301. }
  1302. clone := *apiKey
  1303. userClone := *user
  1304. clone.User = &userClone
  1305. return &clone, nil
  1306. }}
  1307. cfg := &config.Config{RunMode: config.RunModeStandard}
  1308. router := newAuthTestRouter(service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg), nil, cfg)
  1309. w := httptest.NewRecorder()
  1310. req := httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  1311. req.Header.Set("x-api-key", apiKey.Key)
  1312. router.ServeHTTP(w, req)
  1313. require.Equal(t, http.StatusTooManyRequests, w.Code)
  1314. var response struct {
  1315. Error struct {
  1316. Message string `json:"message"`
  1317. Type string `json:"type"`
  1318. Param *string `json:"param"`
  1319. Code string `json:"code"`
  1320. } `json:"error"`
  1321. }
  1322. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &response))
  1323. require.Equal(t, "API key 额度已用完", response.Error.Message)
  1324. require.Equal(t, "insufficient_quota", response.Error.Type)
  1325. require.Nil(t, response.Error.Param)
  1326. require.Equal(t, "insufficient_quota", response.Error.Code)
  1327. }
  1328. func TestAPIKeyAuthQuotaErrorKeepsLegacyFormatOutsideResponses(t *testing.T) {
  1329. gin.SetMode(gin.TestMode)
  1330. user := &service.User{ID: 11, Role: service.RoleUser, Status: service.StatusActive, Balance: 10}
  1331. group := &service.Group{ID: 8, Platform: service.PlatformOpenAI, Status: service.StatusActive}
  1332. apiKey := &service.APIKey{
  1333. ID: 105, UserID: user.ID, Key: "openai-quota-exhausted", Status: service.StatusAPIKeyQuotaExhausted,
  1334. User: user, Group: group, GroupID: &group.ID,
  1335. }
  1336. apiKeyRepo := &stubApiKeyRepo{getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
  1337. if key != apiKey.Key {
  1338. return nil, service.ErrAPIKeyNotFound
  1339. }
  1340. clone := *apiKey
  1341. userClone := *user
  1342. clone.User = &userClone
  1343. return &clone, nil
  1344. }}
  1345. cfg := &config.Config{RunMode: config.RunModeStandard}
  1346. router := newAuthTestRouter(service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg), nil, cfg)
  1347. w := httptest.NewRecorder()
  1348. req := httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
  1349. req.Header.Set("x-api-key", apiKey.Key)
  1350. router.ServeHTTP(w, req)
  1351. require.Equal(t, http.StatusTooManyRequests, w.Code)
  1352. requireAPIKeyAuthError(t, w, "API_KEY_QUOTA_EXHAUSTED", "API key 额度已用完")
  1353. }
  1354. func newAuthTestRouter(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) *gin.Engine {
  1355. router := gin.New()
  1356. router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, subscriptionService, cfg)))
  1357. ok := func(c *gin.Context) {
  1358. c.JSON(http.StatusOK, gin.H{"ok": true})
  1359. }
  1360. router.GET("/t", ok)
  1361. router.POST("/v1/responses", ok)
  1362. router.POST("/v1/messages", ok)
  1363. router.GET("/v1/usage", ok)
  1364. router.GET("/v1/sub2api/billing", ok)
  1365. return router
  1366. }
  1367. func requireAPIKeyAuthError(t *testing.T, w *httptest.ResponseRecorder, code, message string) {
  1368. t.Helper()
  1369. var resp ErrorResponse
  1370. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  1371. require.Equal(t, code, resp.Code)
  1372. require.Equal(t, message, resp.Message)
  1373. }
  1374. type stubApiKeyRepo struct {
  1375. getByKey func(ctx context.Context, key string) (*service.APIKey, error)
  1376. updateLastUsed func(ctx context.Context, id int64, usedAt time.Time) error
  1377. }
  1378. func (r *stubApiKeyRepo) Create(ctx context.Context, key *service.APIKey) error {
  1379. return errors.New("not implemented")
  1380. }
  1381. func (r *stubApiKeyRepo) GetByID(ctx context.Context, id int64) (*service.APIKey, error) {
  1382. return nil, errors.New("not implemented")
  1383. }
  1384. func (r *stubApiKeyRepo) GetKeyAndOwnerID(ctx context.Context, id int64) (string, int64, error) {
  1385. return "", 0, errors.New("not implemented")
  1386. }
  1387. func (r *stubApiKeyRepo) GetByKey(ctx context.Context, key string) (*service.APIKey, error) {
  1388. if r.getByKey != nil {
  1389. return r.getByKey(ctx, key)
  1390. }
  1391. return nil, errors.New("not implemented")
  1392. }
  1393. func (r *stubApiKeyRepo) GetByKeyForAuth(ctx context.Context, key string) (*service.APIKey, error) {
  1394. return r.GetByKey(ctx, key)
  1395. }
  1396. func (r *stubApiKeyRepo) Update(ctx context.Context, key *service.APIKey, _ service.APIKeyUpdateFields) error {
  1397. return errors.New("not implemented")
  1398. }
  1399. func (r *stubApiKeyRepo) Delete(ctx context.Context, id int64) error {
  1400. return errors.New("not implemented")
  1401. }
  1402. func (r *stubApiKeyRepo) DeleteWithAudit(ctx context.Context, id int64) error {
  1403. return errors.New("not implemented")
  1404. }
  1405. func (r *stubApiKeyRepo) ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams, _ service.APIKeyListFilters) ([]service.APIKey, *pagination.PaginationResult, error) {
  1406. return nil, nil, errors.New("not implemented")
  1407. }
  1408. func (r *stubApiKeyRepo) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
  1409. return nil, errors.New("not implemented")
  1410. }
  1411. func (r *stubApiKeyRepo) CountByUserID(ctx context.Context, userID int64) (int64, error) {
  1412. return 0, errors.New("not implemented")
  1413. }
  1414. func (r *stubApiKeyRepo) ExistsByKey(ctx context.Context, key string) (bool, error) {
  1415. return false, errors.New("not implemented")
  1416. }
  1417. func (r *stubApiKeyRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.APIKey, *pagination.PaginationResult, error) {
  1418. return nil, nil, errors.New("not implemented")
  1419. }
  1420. func (r *stubApiKeyRepo) SearchAPIKeys(ctx context.Context, userID int64, keyword string, limit int) ([]service.APIKey, error) {
  1421. return nil, errors.New("not implemented")
  1422. }
  1423. func (r *stubApiKeyRepo) ClearGroupIDByGroupID(ctx context.Context, groupID int64) (int64, error) {
  1424. return 0, errors.New("not implemented")
  1425. }
  1426. func (r *stubApiKeyRepo) UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (int64, error) {
  1427. return 0, errors.New("not implemented")
  1428. }
  1429. func (r *stubApiKeyRepo) CountByGroupID(ctx context.Context, groupID int64) (int64, error) {
  1430. return 0, errors.New("not implemented")
  1431. }
  1432. func (r *stubApiKeyRepo) ListKeysByUserID(ctx context.Context, userID int64) ([]string, error) {
  1433. return nil, errors.New("not implemented")
  1434. }
  1435. func (r *stubApiKeyRepo) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) {
  1436. return nil, errors.New("not implemented")
  1437. }
  1438. func (r *stubApiKeyRepo) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error) {
  1439. return 0, errors.New("not implemented")
  1440. }
  1441. func (r *stubApiKeyRepo) UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error {
  1442. if r.updateLastUsed != nil {
  1443. return r.updateLastUsed(ctx, id, usedAt)
  1444. }
  1445. return nil
  1446. }
  1447. func (r *stubApiKeyRepo) IncrementRateLimitUsage(ctx context.Context, id int64, cost float64) error {
  1448. return nil
  1449. }
  1450. func (r *stubApiKeyRepo) ResetRateLimitWindows(ctx context.Context, id int64) error {
  1451. return nil
  1452. }
  1453. func (r *stubApiKeyRepo) GetRateLimitData(ctx context.Context, id int64) (*service.APIKeyRateLimitData, error) {
  1454. return nil, nil
  1455. }
  1456. type stubUserSubscriptionRepo struct {
  1457. getByID func(ctx context.Context, id int64) (*service.UserSubscription, error)
  1458. getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error)
  1459. updateStatus func(ctx context.Context, subscriptionID int64, status string) error
  1460. activateWindow func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error
  1461. resetDaily func(ctx context.Context, id int64, start time.Time) error
  1462. resetWeekly func(ctx context.Context, id int64, start time.Time) error
  1463. resetMonthly func(ctx context.Context, id int64, start time.Time) error
  1464. }
  1465. type fakeSettingRepo struct {
  1466. values map[string]string
  1467. }
  1468. func (r fakeSettingRepo) Get(ctx context.Context, key string) (*service.Setting, error) {
  1469. return nil, errors.New("not implemented")
  1470. }
  1471. func (r fakeSettingRepo) GetValue(ctx context.Context, key string) (string, error) {
  1472. if v, ok := r.values[key]; ok {
  1473. return v, nil
  1474. }
  1475. return "", service.ErrSettingNotFound
  1476. }
  1477. func (r fakeSettingRepo) Set(ctx context.Context, key, value string) error {
  1478. return errors.New("not implemented")
  1479. }
  1480. func (r fakeSettingRepo) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
  1481. return nil, errors.New("not implemented")
  1482. }
  1483. func (r fakeSettingRepo) SetMultiple(ctx context.Context, settings map[string]string) error {
  1484. return errors.New("not implemented")
  1485. }
  1486. func (r fakeSettingRepo) GetAll(ctx context.Context) (map[string]string, error) {
  1487. return nil, errors.New("not implemented")
  1488. }
  1489. func (r fakeSettingRepo) Delete(ctx context.Context, key string) error {
  1490. return errors.New("not implemented")
  1491. }
  1492. func (r *stubUserSubscriptionRepo) Create(ctx context.Context, sub *service.UserSubscription) error {
  1493. return errors.New("not implemented")
  1494. }
  1495. func (r *stubUserSubscriptionRepo) GetByID(ctx context.Context, id int64) (*service.UserSubscription, error) {
  1496. if r.getByID != nil {
  1497. return r.getByID(ctx, id)
  1498. }
  1499. return nil, errors.New("not implemented")
  1500. }
  1501. func (r *stubUserSubscriptionRepo) GetByIDForUpdate(ctx context.Context, id int64) (*service.UserSubscription, error) {
  1502. return r.GetByID(ctx, id)
  1503. }
  1504. func (r *stubUserSubscriptionRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.UserSubscription, error) {
  1505. return nil, errors.New("not implemented")
  1506. }
  1507. func (r *stubUserSubscriptionRepo) GetByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
  1508. return nil, errors.New("not implemented")
  1509. }
  1510. func (r *stubUserSubscriptionRepo) GetActiveByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
  1511. if r.getActive != nil {
  1512. return r.getActive(ctx, userID, groupID)
  1513. }
  1514. return nil, errors.New("not implemented")
  1515. }
  1516. func (r *stubUserSubscriptionRepo) Update(ctx context.Context, sub *service.UserSubscription) error {
  1517. return errors.New("not implemented")
  1518. }
  1519. func (r *stubUserSubscriptionRepo) Delete(ctx context.Context, id int64) error {
  1520. return errors.New("not implemented")
  1521. }
  1522. func (r *stubUserSubscriptionRepo) Restore(ctx context.Context, subscriptionID int64, restoredStatus string) (*service.UserSubscription, error) {
  1523. return nil, errors.New("not implemented")
  1524. }
  1525. func (r *stubUserSubscriptionRepo) ListByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
  1526. return nil, errors.New("not implemented")
  1527. }
  1528. func (r *stubUserSubscriptionRepo) ListActiveByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
  1529. return nil, errors.New("not implemented")
  1530. }
  1531. func (r *stubUserSubscriptionRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.UserSubscription, *pagination.PaginationResult, error) {
  1532. return nil, nil, errors.New("not implemented")
  1533. }
  1534. func (r *stubUserSubscriptionRepo) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, platform, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
  1535. return nil, nil, errors.New("not implemented")
  1536. }
  1537. func (r *stubUserSubscriptionRepo) ExistsByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
  1538. return false, errors.New("not implemented")
  1539. }
  1540. func (r *stubUserSubscriptionRepo) ExistsActiveByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
  1541. return false, errors.New("not implemented")
  1542. }
  1543. func (r *stubUserSubscriptionRepo) ExtendExpiry(ctx context.Context, subscriptionID int64, newExpiresAt time.Time) error {
  1544. return errors.New("not implemented")
  1545. }
  1546. func (r *stubUserSubscriptionRepo) UpdateStatus(ctx context.Context, subscriptionID int64, status string) error {
  1547. if r.updateStatus != nil {
  1548. return r.updateStatus(ctx, subscriptionID, status)
  1549. }
  1550. return errors.New("not implemented")
  1551. }
  1552. func (r *stubUserSubscriptionRepo) UpdateNotes(ctx context.Context, subscriptionID int64, notes string) error {
  1553. return errors.New("not implemented")
  1554. }
  1555. func (r *stubUserSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error {
  1556. if r.activateWindow != nil {
  1557. return r.activateWindow(ctx, id, dailyStart, periodicStart)
  1558. }
  1559. return errors.New("not implemented")
  1560. }
  1561. func (r *stubUserSubscriptionRepo) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time, time.Time) error {
  1562. return errors.New("not implemented")
  1563. }
  1564. func (r *stubUserSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, _ *time.Time, newWindowStart time.Time) error {
  1565. if r.resetDaily != nil {
  1566. return r.resetDaily(ctx, id, newWindowStart)
  1567. }
  1568. return errors.New("not implemented")
  1569. }
  1570. func (r *stubUserSubscriptionRepo) ResetWeeklyUsage(ctx context.Context, id int64, _ *time.Time, newWindowStart time.Time) error {
  1571. if r.resetWeekly != nil {
  1572. return r.resetWeekly(ctx, id, newWindowStart)
  1573. }
  1574. return errors.New("not implemented")
  1575. }
  1576. func (r *stubUserSubscriptionRepo) ResetMonthlyUsage(ctx context.Context, id int64, _ *time.Time, newWindowStart time.Time) error {
  1577. if r.resetMonthly != nil {
  1578. return r.resetMonthly(ctx, id, newWindowStart)
  1579. }
  1580. return errors.New("not implemented")
  1581. }
  1582. func (r *stubUserSubscriptionRepo) IncrementUsage(ctx context.Context, id int64, costUSD float64) error {
  1583. return errors.New("not implemented")
  1584. }
  1585. func (r *stubUserSubscriptionRepo) BatchUpdateExpiredStatus(ctx context.Context) (int64, error) {
  1586. return 0, errors.New("not implemented")
  1587. }