admin_auth_test.go 8.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261
  1. //go:build unit
  2. package middleware
  3. import (
  4. "context"
  5. "net/http"
  6. "net/http/httptest"
  7. "testing"
  8. "time"
  9. "github.com/Wei-Shaw/sub2api/internal/config"
  10. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  11. "github.com/Wei-Shaw/sub2api/internal/service"
  12. "github.com/gin-gonic/gin"
  13. "github.com/stretchr/testify/require"
  14. )
  15. func TestAdminAuthJWTValidatesTokenVersion(t *testing.T) {
  16. gin.SetMode(gin.TestMode)
  17. cfg := &config.Config{JWT: config.JWTConfig{Secret: "test-secret", ExpireHour: 1}}
  18. authService := service.NewAuthService(nil, nil, nil, nil, cfg, nil, nil, nil, nil, nil, nil, nil, nil)
  19. admin := &service.User{
  20. ID: 1,
  21. Email: "admin@example.com",
  22. Role: service.RoleAdmin,
  23. Status: service.StatusActive,
  24. TokenVersion: 2,
  25. Concurrency: 1,
  26. }
  27. userRepo := &stubUserRepo{
  28. getByID: func(ctx context.Context, id int64) (*service.User, error) {
  29. if id != admin.ID {
  30. return nil, service.ErrUserNotFound
  31. }
  32. clone := *admin
  33. return &clone, nil
  34. },
  35. }
  36. userService := service.NewUserService(userRepo, nil, nil, nil)
  37. router := gin.New()
  38. router.Use(gin.HandlerFunc(NewAdminAuthMiddleware(authService, userService, nil, nil)))
  39. router.GET("/t", func(c *gin.Context) {
  40. c.JSON(http.StatusOK, gin.H{"ok": true})
  41. })
  42. t.Run("token_version_mismatch_rejected", func(t *testing.T) {
  43. token, err := authService.GenerateToken(context.Background(), &service.User{
  44. ID: admin.ID,
  45. Email: admin.Email,
  46. Role: admin.Role,
  47. TokenVersion: admin.TokenVersion - 1,
  48. })
  49. require.NoError(t, err)
  50. w := httptest.NewRecorder()
  51. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  52. req.Header.Set("Authorization", "Bearer "+token)
  53. router.ServeHTTP(w, req)
  54. require.Equal(t, http.StatusUnauthorized, w.Code)
  55. require.Contains(t, w.Body.String(), "TOKEN_REVOKED")
  56. })
  57. t.Run("token_version_match_allows", func(t *testing.T) {
  58. token, err := authService.GenerateToken(context.Background(), &service.User{
  59. ID: admin.ID,
  60. Email: admin.Email,
  61. Role: admin.Role,
  62. TokenVersion: admin.TokenVersion,
  63. })
  64. require.NoError(t, err)
  65. w := httptest.NewRecorder()
  66. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  67. req.Header.Set("Authorization", "Bearer "+token)
  68. router.ServeHTTP(w, req)
  69. require.Equal(t, http.StatusOK, w.Code)
  70. })
  71. t.Run("websocket_token_version_mismatch_rejected", func(t *testing.T) {
  72. token, err := authService.GenerateToken(context.Background(), &service.User{
  73. ID: admin.ID,
  74. Email: admin.Email,
  75. Role: admin.Role,
  76. TokenVersion: admin.TokenVersion - 1,
  77. })
  78. require.NoError(t, err)
  79. w := httptest.NewRecorder()
  80. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  81. req.Header.Set("Upgrade", "websocket")
  82. req.Header.Set("Connection", "Upgrade")
  83. req.Header.Set("Sec-WebSocket-Protocol", "sub2api-admin, jwt."+token)
  84. router.ServeHTTP(w, req)
  85. require.Equal(t, http.StatusUnauthorized, w.Code)
  86. require.Contains(t, w.Body.String(), "TOKEN_REVOKED")
  87. })
  88. t.Run("websocket_token_version_match_allows", func(t *testing.T) {
  89. token, err := authService.GenerateToken(context.Background(), &service.User{
  90. ID: admin.ID,
  91. Email: admin.Email,
  92. Role: admin.Role,
  93. TokenVersion: admin.TokenVersion,
  94. })
  95. require.NoError(t, err)
  96. w := httptest.NewRecorder()
  97. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  98. req.Header.Set("Upgrade", "websocket")
  99. req.Header.Set("Connection", "Upgrade")
  100. req.Header.Set("Sec-WebSocket-Protocol", "sub2api-admin, jwt."+token)
  101. router.ServeHTTP(w, req)
  102. require.Equal(t, http.StatusOK, w.Code)
  103. })
  104. }
  105. type stubUserRepo struct {
  106. getByID func(ctx context.Context, id int64) (*service.User, error)
  107. }
  108. func (s *stubUserRepo) Create(ctx context.Context, user *service.User) error {
  109. panic("unexpected Create call")
  110. }
  111. func (s *stubUserRepo) CreateWithEmailAliasGuard(ctx context.Context, user *service.User) error {
  112. panic("unexpected CreateWithEmailAliasGuard call")
  113. }
  114. func (s *stubUserRepo) GetByID(ctx context.Context, id int64) (*service.User, error) {
  115. if s.getByID == nil {
  116. panic("GetByID not stubbed")
  117. }
  118. return s.getByID(ctx, id)
  119. }
  120. func (s *stubUserRepo) GetByEmail(ctx context.Context, email string) (*service.User, error) {
  121. panic("unexpected GetByEmail call")
  122. }
  123. func (s *stubUserRepo) GetFirstAdmin(ctx context.Context) (*service.User, error) {
  124. panic("unexpected GetFirstAdmin call")
  125. }
  126. func (s *stubUserRepo) Update(ctx context.Context, user *service.User, fields service.UserUpdateFields) error {
  127. panic("unexpected Update call")
  128. }
  129. func (s *stubUserRepo) Delete(ctx context.Context, id int64) error {
  130. panic("unexpected Delete call")
  131. }
  132. func (s *stubUserRepo) GetUserAvatar(ctx context.Context, userID int64) (*service.UserAvatar, error) {
  133. return nil, nil
  134. }
  135. func (s *stubUserRepo) UpsertUserAvatar(ctx context.Context, userID int64, input service.UpsertUserAvatarInput) (*service.UserAvatar, error) {
  136. panic("unexpected UpsertUserAvatar call")
  137. }
  138. func (s *stubUserRepo) DeleteUserAvatar(ctx context.Context, userID int64) error {
  139. panic("unexpected DeleteUserAvatar call")
  140. }
  141. func (s *stubUserRepo) List(ctx context.Context, params pagination.PaginationParams) ([]service.User, *pagination.PaginationResult, error) {
  142. panic("unexpected List call")
  143. }
  144. func (s *stubUserRepo) ListWithFilters(ctx context.Context, params pagination.PaginationParams, filters service.UserListFilters) ([]service.User, *pagination.PaginationResult, error) {
  145. panic("unexpected ListWithFilters call")
  146. }
  147. func (s *stubUserRepo) GetLatestUsedAtByUserIDs(ctx context.Context, userIDs []int64) (map[int64]*time.Time, error) {
  148. panic("unexpected GetLatestUsedAtByUserIDs call")
  149. }
  150. func (s *stubUserRepo) GetLatestUsedAtByUserID(ctx context.Context, userID int64) (*time.Time, error) {
  151. panic("unexpected GetLatestUsedAtByUserID call")
  152. }
  153. func (s *stubUserRepo) UpdateUserLastActiveAt(ctx context.Context, userID int64, activeAt time.Time) error {
  154. panic("unexpected UpdateUserLastActiveAt call")
  155. }
  156. func (s *stubUserRepo) UpdateBalance(ctx context.Context, id int64, amount float64) error {
  157. panic("unexpected UpdateBalance call")
  158. }
  159. func (s *stubUserRepo) DeductBalance(ctx context.Context, id int64, amount float64) error {
  160. panic("unexpected DeductBalance call")
  161. }
  162. func (s *stubUserRepo) AdjustBalance(ctx context.Context, id int64, delta float64) (service.BalanceChange, error) {
  163. panic("unexpected AdjustBalance call")
  164. }
  165. func (s *stubUserRepo) SetBalance(ctx context.Context, id int64, value float64) (service.BalanceChange, error) {
  166. panic("unexpected SetBalance call")
  167. }
  168. func (s *stubUserRepo) UpdateConcurrency(ctx context.Context, id int64, amount int) error {
  169. panic("unexpected UpdateConcurrency call")
  170. }
  171. func (s *stubUserRepo) BatchSetConcurrency(context.Context, []int64, int) (int, error) { return 0, nil }
  172. func (s *stubUserRepo) BatchAddConcurrency(context.Context, []int64, int) (int, error) { return 0, nil }
  173. func (s *stubUserRepo) BatchUpdateLimits(context.Context, []int64, *int, *int) (int, error) {
  174. return 0, nil
  175. }
  176. func (s *stubUserRepo) ExistsByEmail(ctx context.Context, email string) (bool, error) {
  177. panic("unexpected ExistsByEmail call")
  178. }
  179. func (s *stubUserRepo) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
  180. panic("unexpected ExistsByEmailAlias call")
  181. }
  182. func (s *stubUserRepo) RemoveGroupFromAllowedGroups(ctx context.Context, groupID int64) (int64, error) {
  183. panic("unexpected RemoveGroupFromAllowedGroups call")
  184. }
  185. func (s *stubUserRepo) RemoveGroupFromUserAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
  186. panic("unexpected RemoveGroupFromUserAllowedGroups call")
  187. }
  188. func (s *stubUserRepo) AddGroupToAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
  189. panic("unexpected AddGroupToAllowedGroups call")
  190. }
  191. func (s *stubUserRepo) ListUserAuthIdentities(ctx context.Context, userID int64) ([]service.UserAuthIdentityRecord, error) {
  192. panic("unexpected ListUserAuthIdentities call")
  193. }
  194. func (s *stubUserRepo) UnbindUserAuthProvider(context.Context, int64, string) error {
  195. panic("unexpected UnbindUserAuthProvider call")
  196. }
  197. func (s *stubUserRepo) UpdateTotpSecret(ctx context.Context, userID int64, encryptedSecret *string) error {
  198. panic("unexpected UpdateTotpSecret call")
  199. }
  200. func (s *stubUserRepo) EnableTotp(ctx context.Context, userID int64) error {
  201. panic("unexpected EnableTotp call")
  202. }
  203. func (s *stubUserRepo) DisableTotp(ctx context.Context, userID int64) error {
  204. panic("unexpected DisableTotp call")
  205. }
  206. func (s *stubUserRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.User, error) {
  207. panic("unexpected GetByIDIncludeDeleted call")
  208. }