auth_oauth_email_flow_test.go 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "testing"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/config"
  9. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  10. "github.com/stretchr/testify/require"
  11. )
  12. type redeemCodeRepoStub struct {
  13. codesByCode map[string]*RedeemCode
  14. useCalls []struct {
  15. id int64
  16. userID int64
  17. }
  18. updateCalls []*RedeemCode
  19. }
  20. func (s *redeemCodeRepoStub) Create(context.Context, *RedeemCode) error {
  21. panic("unexpected Create call")
  22. }
  23. func (s *redeemCodeRepoStub) CreateBatch(context.Context, []RedeemCode) error {
  24. panic("unexpected CreateBatch call")
  25. }
  26. func (s *redeemCodeRepoStub) GetByID(context.Context, int64) (*RedeemCode, error) {
  27. panic("unexpected GetByID call")
  28. }
  29. func (s *redeemCodeRepoStub) GetByCode(_ context.Context, code string) (*RedeemCode, error) {
  30. if s.codesByCode == nil {
  31. return nil, ErrRedeemCodeNotFound
  32. }
  33. redeemCode, ok := s.codesByCode[code]
  34. if !ok {
  35. return nil, ErrRedeemCodeNotFound
  36. }
  37. cloned := *redeemCode
  38. return &cloned, nil
  39. }
  40. func (s *redeemCodeRepoStub) Update(_ context.Context, code *RedeemCode) error {
  41. if code == nil {
  42. return nil
  43. }
  44. cloned := *code
  45. s.updateCalls = append(s.updateCalls, &cloned)
  46. if s.codesByCode == nil {
  47. s.codesByCode = make(map[string]*RedeemCode)
  48. }
  49. s.codesByCode[cloned.Code] = &cloned
  50. return nil
  51. }
  52. func (s *redeemCodeRepoStub) BatchUpdate(context.Context, []int64, RedeemCodeBatchUpdateFields) (int64, error) {
  53. panic("unexpected BatchUpdate call")
  54. }
  55. func (s *redeemCodeRepoStub) Delete(context.Context, int64) error {
  56. panic("unexpected Delete call")
  57. }
  58. func (s *redeemCodeRepoStub) Use(_ context.Context, id, userID int64) error {
  59. for code, redeemCode := range s.codesByCode {
  60. if redeemCode.ID != id {
  61. continue
  62. }
  63. now := time.Now().UTC()
  64. redeemCode.Status = StatusUsed
  65. redeemCode.UsedBy = &userID
  66. redeemCode.UsedAt = &now
  67. s.codesByCode[code] = redeemCode
  68. s.useCalls = append(s.useCalls, struct {
  69. id int64
  70. userID int64
  71. }{id: id, userID: userID})
  72. return nil
  73. }
  74. return ErrRedeemCodeNotFound
  75. }
  76. func (s *redeemCodeRepoStub) List(context.Context, pagination.PaginationParams) ([]RedeemCode, *pagination.PaginationResult, error) {
  77. panic("unexpected List call")
  78. }
  79. func (s *redeemCodeRepoStub) ListWithFilters(context.Context, pagination.PaginationParams, string, string, string) ([]RedeemCode, *pagination.PaginationResult, error) {
  80. panic("unexpected ListWithFilters call")
  81. }
  82. func (s *redeemCodeRepoStub) ListByUser(context.Context, int64, int) ([]RedeemCode, error) {
  83. panic("unexpected ListByUser call")
  84. }
  85. func (s *redeemCodeRepoStub) ListByUserPaginated(context.Context, int64, pagination.PaginationParams, string) ([]RedeemCode, *pagination.PaginationResult, error) {
  86. panic("unexpected ListByUserPaginated call")
  87. }
  88. func (s *redeemCodeRepoStub) SumPositiveBalanceByUser(context.Context, int64) (float64, error) {
  89. panic("unexpected SumPositiveBalanceByUser call")
  90. }
  91. func newOAuthEmailFlowAuthService(
  92. userRepo UserRepository,
  93. redeemRepo RedeemCodeRepository,
  94. refreshTokenCache RefreshTokenCache,
  95. settings map[string]string,
  96. emailCache EmailCache,
  97. quotaRepo UserPlatformQuotaRepository, // 新增
  98. ) *AuthService {
  99. cfg := &config.Config{
  100. JWT: config.JWTConfig{
  101. Secret: "test-secret",
  102. ExpireHour: 1,
  103. AccessTokenExpireMinutes: 60,
  104. RefreshTokenExpireDays: 7,
  105. },
  106. Default: config.DefaultConfig{
  107. UserBalance: 3.5,
  108. UserConcurrency: 2,
  109. },
  110. }
  111. settingService := NewSettingService(&settingRepoStub{values: settings}, cfg)
  112. emailService := NewEmailService(&settingRepoStub{values: settings}, emailCache)
  113. return NewAuthService(
  114. nil,
  115. userRepo,
  116. redeemRepo,
  117. refreshTokenCache,
  118. cfg,
  119. settingService,
  120. emailService,
  121. nil,
  122. nil,
  123. nil,
  124. nil,
  125. nil,
  126. quotaRepo, // 替换原来的 nil
  127. )
  128. }
  129. func TestRegisterOAuthEmailAccountRollsBackCreatedUserWhenTokenPairGenerationFails(t *testing.T) {
  130. userRepo := &userRepoStub{nextID: 42}
  131. redeemRepo := &redeemCodeRepoStub{
  132. codesByCode: map[string]*RedeemCode{
  133. "INVITE123": {
  134. ID: 7,
  135. Code: "INVITE123",
  136. Type: RedeemTypeInvitation,
  137. Status: StatusUnused,
  138. },
  139. },
  140. }
  141. emailCache := &emailCacheStub{
  142. data: &VerificationCodeData{
  143. Code: "246810",
  144. Attempts: 0,
  145. CreatedAt: time.Now().UTC(),
  146. ExpiresAt: time.Now().UTC().Add(15 * time.Minute),
  147. },
  148. }
  149. authService := newOAuthEmailFlowAuthService(
  150. userRepo,
  151. redeemRepo,
  152. nil,
  153. map[string]string{
  154. SettingKeyRegistrationEnabled: "true",
  155. SettingKeyInvitationCodeEnabled: "true",
  156. SettingKeyEmailVerifyEnabled: "true",
  157. },
  158. emailCache,
  159. nil,
  160. )
  161. tokenPair, user, err := authService.RegisterOAuthEmailAccount(
  162. context.Background(),
  163. "fresh@example.com",
  164. "secret-123",
  165. "246810",
  166. "INVITE123",
  167. "oidc",
  168. )
  169. require.Nil(t, tokenPair)
  170. require.Nil(t, user)
  171. require.Error(t, err)
  172. require.Contains(t, err.Error(), "generate token pair")
  173. require.Equal(t, []int64{42}, userRepo.deletedIDs)
  174. require.Len(t, userRepo.created, 1)
  175. require.Empty(t, redeemRepo.useCalls)
  176. require.Empty(t, redeemRepo.updateCalls)
  177. }
  178. func TestRegisterOAuthEmailAccount_NonWhitelistDomainLimit(t *testing.T) {
  179. userRepo := &userRepoStub{domainCounts: map[string]int{"custom.example": 1}}
  180. authService := newOAuthEmailFlowAuthService(
  181. userRepo,
  182. &redeemCodeRepoStub{},
  183. &refreshTokenCacheStub{},
  184. map[string]string{
  185. SettingKeyRegistrationEnabled: "true",
  186. SettingKeyRegistrationEmailSuffixWhitelist: `["@example.com"]`,
  187. SettingKeyRegistrationEmailDomainQuotaEnabled: "true",
  188. },
  189. &emailCacheStub{data: &VerificationCodeData{
  190. Code: "246810",
  191. CreatedAt: time.Now().UTC(),
  192. ExpiresAt: time.Now().UTC().Add(15 * time.Minute),
  193. }},
  194. nil,
  195. )
  196. _, _, err := authService.RegisterOAuthEmailAccount(
  197. context.Background(),
  198. "second@custom.example",
  199. "secret-123",
  200. "246810",
  201. "",
  202. "oidc",
  203. )
  204. require.ErrorIs(t, err, ErrEmailDomainRegistrationLimit)
  205. }
  206. func TestRegisterVerifiedOAuthEmailAccount_NonWhitelistDomainLimit(t *testing.T) {
  207. userRepo := &userRepoStub{domainCounts: map[string]int{"custom.example": 1}}
  208. authService := newOAuthEmailFlowAuthService(
  209. userRepo,
  210. nil,
  211. &refreshTokenCacheStub{},
  212. map[string]string{
  213. SettingKeyRegistrationEnabled: "true",
  214. SettingKeyRegistrationEmailSuffixWhitelist: `["@example.com"]`,
  215. SettingKeyRegistrationEmailDomainQuotaEnabled: "true",
  216. },
  217. &emailCacheStub{},
  218. nil,
  219. )
  220. _, _, err := authService.RegisterVerifiedOAuthEmailAccount(
  221. context.Background(),
  222. "second@custom.example",
  223. "secret-123",
  224. "",
  225. "oidc",
  226. )
  227. require.ErrorIs(t, err, ErrEmailDomainRegistrationLimit)
  228. }
  229. func TestSendPendingOAuthVerifyCode_NonWhitelistDomainLimit(t *testing.T) {
  230. userRepo := &userRepoStub{domainCounts: map[string]int{"custom.example": 1}}
  231. authService := newOAuthEmailFlowAuthService(
  232. userRepo,
  233. nil,
  234. nil,
  235. map[string]string{
  236. SettingKeyRegistrationEnabled: "true",
  237. SettingKeyRegistrationEmailSuffixWhitelist: `["@example.com"]`,
  238. SettingKeyRegistrationEmailDomainQuotaEnabled: "true",
  239. },
  240. &emailCacheStub{},
  241. nil,
  242. )
  243. _, err := authService.SendPendingOAuthVerifyCode(context.Background(), "second@custom.example")
  244. require.ErrorIs(t, err, ErrEmailDomainRegistrationLimit)
  245. }
  246. // 域名限量注册开关默认关闭:白名单外域名在 pending OAuth 发码阶段即被严格拒绝。
  247. func TestSendPendingOAuthVerifyCode_NonWhitelistDomainRejectedWhenQuotaDisabled(t *testing.T) {
  248. userRepo := &userRepoStub{domainCounts: map[string]int{"custom.example": 0}}
  249. authService := newOAuthEmailFlowAuthService(
  250. userRepo,
  251. nil,
  252. nil,
  253. map[string]string{
  254. SettingKeyRegistrationEnabled: "true",
  255. SettingKeyRegistrationEmailSuffixWhitelist: `["@example.com"]`,
  256. },
  257. &emailCacheStub{},
  258. nil,
  259. )
  260. _, err := authService.SendPendingOAuthVerifyCode(context.Background(), "first@custom.example")
  261. require.ErrorIs(t, err, ErrEmailSuffixNotAllowed)
  262. }
  263. func TestSendPendingOAuthVerifyCode_NilServiceReturnsUnavailable(t *testing.T) {
  264. var authService *AuthService
  265. _, err := authService.SendPendingOAuthVerifyCode(context.Background(), "fresh@example.com")
  266. require.ErrorIs(t, err, ErrServiceUnavailable)
  267. }
  268. func TestRegisterOAuthEmailAccountSetsNormalizedSignupSourceOnCreatedUser(t *testing.T) {
  269. userRepo := &userRepoStub{nextID: 42}
  270. emailCache := &emailCacheStub{
  271. data: &VerificationCodeData{
  272. Code: "246810",
  273. Attempts: 0,
  274. CreatedAt: time.Now().UTC(),
  275. ExpiresAt: time.Now().UTC().Add(15 * time.Minute),
  276. },
  277. }
  278. authService := newOAuthEmailFlowAuthService(
  279. userRepo,
  280. &redeemCodeRepoStub{},
  281. &refreshTokenCacheStub{},
  282. map[string]string{
  283. SettingKeyRegistrationEnabled: "true",
  284. SettingKeyEmailVerifyEnabled: "true",
  285. },
  286. emailCache,
  287. nil,
  288. )
  289. tokenPair, user, err := authService.RegisterOAuthEmailAccount(
  290. context.Background(),
  291. "fresh@example.com",
  292. "secret-123",
  293. "246810",
  294. "",
  295. " OIDC ",
  296. )
  297. require.NoError(t, err)
  298. require.NotNil(t, tokenPair)
  299. require.NotNil(t, user)
  300. require.Len(t, userRepo.created, 1)
  301. require.Equal(t, "oidc", userRepo.created[0].SignupSource)
  302. }
  303. func TestRegisterOAuthEmailAccountKeepsGitHubAndGoogleSignupSource(t *testing.T) {
  304. tests := []struct {
  305. name string
  306. email string
  307. signupSource string
  308. want string
  309. }{
  310. {
  311. name: "github",
  312. email: "github@example.com",
  313. signupSource: " GitHub ",
  314. want: "github",
  315. },
  316. {
  317. name: "google",
  318. email: "google@example.com",
  319. signupSource: " Google ",
  320. want: "google",
  321. },
  322. }
  323. for _, tt := range tests {
  324. t.Run(tt.name, func(t *testing.T) {
  325. userRepo := &userRepoStub{nextID: 43}
  326. emailCache := &emailCacheStub{
  327. data: &VerificationCodeData{
  328. Code: "246810",
  329. Attempts: 0,
  330. CreatedAt: time.Now().UTC(),
  331. ExpiresAt: time.Now().UTC().Add(15 * time.Minute),
  332. },
  333. }
  334. authService := newOAuthEmailFlowAuthService(
  335. userRepo,
  336. &redeemCodeRepoStub{},
  337. &refreshTokenCacheStub{},
  338. map[string]string{
  339. SettingKeyRegistrationEnabled: "true",
  340. SettingKeyEmailVerifyEnabled: "true",
  341. },
  342. emailCache,
  343. nil,
  344. )
  345. tokenPair, user, err := authService.RegisterOAuthEmailAccount(
  346. context.Background(),
  347. tt.email,
  348. "secret-123",
  349. "246810",
  350. "",
  351. tt.signupSource,
  352. )
  353. require.NoError(t, err)
  354. require.NotNil(t, tokenPair)
  355. require.NotNil(t, user)
  356. require.Len(t, userRepo.created, 1)
  357. require.Equal(t, tt.want, userRepo.created[0].SignupSource)
  358. })
  359. }
  360. }
  361. func TestRegisterOAuthEmailAccountFallsBackUnknownSignupSourceToEmail(t *testing.T) {
  362. userRepo := &userRepoStub{nextID: 43}
  363. emailCache := &emailCacheStub{
  364. data: &VerificationCodeData{
  365. Code: "246810",
  366. Attempts: 0,
  367. CreatedAt: time.Now().UTC(),
  368. ExpiresAt: time.Now().UTC().Add(15 * time.Minute),
  369. },
  370. }
  371. authService := newOAuthEmailFlowAuthService(
  372. userRepo,
  373. &redeemCodeRepoStub{},
  374. &refreshTokenCacheStub{},
  375. map[string]string{
  376. SettingKeyRegistrationEnabled: "true",
  377. SettingKeyEmailVerifyEnabled: "true",
  378. },
  379. emailCache,
  380. nil,
  381. )
  382. tokenPair, user, err := authService.RegisterOAuthEmailAccount(
  383. context.Background(),
  384. "fallback@example.com",
  385. "secret-123",
  386. "246810",
  387. "",
  388. "unknown-provider",
  389. )
  390. require.NoError(t, err)
  391. require.NotNil(t, tokenPair)
  392. require.NotNil(t, user)
  393. require.Len(t, userRepo.created, 1)
  394. require.Equal(t, "email", userRepo.created[0].SignupSource)
  395. }
  396. func TestRollbackOAuthEmailAccountCreationRestoresInvitationUsage(t *testing.T) {
  397. userRepo := &userRepoStub{}
  398. redeemRepo := &redeemCodeRepoStub{
  399. codesByCode: map[string]*RedeemCode{
  400. "INVITE123": {
  401. ID: 7,
  402. Code: "INVITE123",
  403. Type: RedeemTypeInvitation,
  404. Status: StatusUsed,
  405. UsedBy: func() *int64 {
  406. v := int64(42)
  407. return &v
  408. }(),
  409. UsedAt: func() *time.Time {
  410. v := time.Now().UTC()
  411. return &v
  412. }(),
  413. },
  414. },
  415. }
  416. authService := newOAuthEmailFlowAuthService(
  417. userRepo,
  418. redeemRepo,
  419. &refreshTokenCacheStub{},
  420. map[string]string{
  421. SettingKeyRegistrationEnabled: "true",
  422. SettingKeyInvitationCodeEnabled: "true",
  423. },
  424. &emailCacheStub{},
  425. nil,
  426. )
  427. err := authService.RollbackOAuthEmailAccountCreation(context.Background(), 42, "INVITE123")
  428. require.NoError(t, err)
  429. require.Equal(t, []int64{42}, userRepo.deletedIDs)
  430. require.Len(t, redeemRepo.updateCalls, 1)
  431. require.Equal(t, StatusUnused, redeemRepo.updateCalls[0].Status)
  432. require.Nil(t, redeemRepo.updateCalls[0].UsedBy)
  433. require.Nil(t, redeemRepo.updateCalls[0].UsedAt)
  434. }
  435. func TestRollbackOAuthEmailAccountCreationPropagatesDeleteError(t *testing.T) {
  436. userRepo := &userRepoStub{deleteErr: errors.New("delete failed")}
  437. authService := newOAuthEmailFlowAuthService(
  438. userRepo,
  439. &redeemCodeRepoStub{},
  440. &refreshTokenCacheStub{},
  441. map[string]string{
  442. SettingKeyRegistrationEnabled: "true",
  443. },
  444. &emailCacheStub{},
  445. nil,
  446. )
  447. err := authService.RollbackOAuthEmailAccountCreation(context.Background(), 42, "")
  448. require.Error(t, err)
  449. require.Contains(t, err.Error(), "delete created oauth user")
  450. }
  451. func TestFinalizeOAuthEmailAccount_SnapshotsPlatformQuotaDefaults(t *testing.T) {
  452. userRepo := &userRepoStub{nextID: 99}
  453. quotaRepo := &userPlatformQuotaRepoStub{}
  454. authService := newOAuthEmailFlowAuthService(
  455. userRepo,
  456. nil,
  457. &refreshTokenCacheStub{},
  458. map[string]string{
  459. SettingKeyRegistrationEnabled: "true",
  460. SettingKeyEmailVerifyEnabled: "true",
  461. SettingKeyDefaultPlatformQuotas: `{"anthropic": {"daily": 5.5}}`,
  462. },
  463. &emailCacheStub{},
  464. quotaRepo,
  465. )
  466. user := &User{
  467. ID: 99,
  468. Email: "newuser@example.com",
  469. Role: RoleUser,
  470. Status: StatusActive,
  471. SignupSource: "oidc",
  472. }
  473. err := authService.FinalizeOAuthEmailAccount(
  474. context.Background(),
  475. user,
  476. "",
  477. "oidc",
  478. "",
  479. )
  480. require.NoError(t, err)
  481. require.Len(t, quotaRepo.bulkInsertCalls, 1, "snapshotPlatformQuotaDefaults must call BulkInsertInitial once on successful OAuth signup")
  482. records := quotaRepo.bulkInsertCalls[0]
  483. var anthropicRecord *UserPlatformQuotaRecord
  484. for i := range records {
  485. if records[i].Platform == "anthropic" {
  486. anthropicRecord = &records[i]
  487. break
  488. }
  489. }
  490. require.NotNil(t, anthropicRecord, "expected anthropic platform record")
  491. require.Equal(t, int64(99), anthropicRecord.UserID)
  492. require.NotNil(t, anthropicRecord.DailyLimitUSD)
  493. require.InDelta(t, 5.5, *anthropicRecord.DailyLimitUSD, 0.0001)
  494. }