auth_rate_limit_test.go 1.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  1. package routes
  2. import (
  3. "net/http"
  4. "net/http/httptest"
  5. "strings"
  6. "testing"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/handler"
  9. servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
  10. "github.com/gin-gonic/gin"
  11. "github.com/redis/go-redis/v9"
  12. "github.com/stretchr/testify/require"
  13. )
  14. func newAuthRoutesTestRouter(redisClient *redis.Client) *gin.Engine {
  15. gin.SetMode(gin.TestMode)
  16. router := gin.New()
  17. v1 := router.Group("/api/v1")
  18. RegisterAuthRoutes(
  19. v1,
  20. &handler.Handlers{
  21. Auth: &handler.AuthHandler{},
  22. Setting: &handler.SettingHandler{},
  23. },
  24. servermiddleware.JWTAuthMiddleware(func(c *gin.Context) {
  25. c.Next()
  26. }),
  27. servermiddleware.AuditLogMiddleware(func(c *gin.Context) {
  28. c.Next()
  29. }),
  30. redisClient,
  31. nil,
  32. nil,
  33. )
  34. return router
  35. }
  36. func TestAuthRoutesRateLimitFailCloseWhenRedisUnavailable(t *testing.T) {
  37. rdb := redis.NewClient(&redis.Options{
  38. Addr: "127.0.0.1:1",
  39. DialTimeout: 50 * time.Millisecond,
  40. ReadTimeout: 50 * time.Millisecond,
  41. WriteTimeout: 50 * time.Millisecond,
  42. })
  43. t.Cleanup(func() {
  44. _ = rdb.Close()
  45. })
  46. router := newAuthRoutesTestRouter(rdb)
  47. paths := []string{
  48. "/api/v1/auth/register",
  49. "/api/v1/auth/login",
  50. "/api/v1/auth/login/2fa",
  51. "/api/v1/auth/send-verify-code",
  52. "/api/v1/auth/oauth/pending/send-verify-code",
  53. }
  54. for _, path := range paths {
  55. req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{}`))
  56. req.Header.Set("Content-Type", "application/json")
  57. req.RemoteAddr = "203.0.113.10:12345"
  58. w := httptest.NewRecorder()
  59. router.ServeHTTP(w, req)
  60. require.Equal(t, http.StatusTooManyRequests, w.Code, "path=%s", path)
  61. require.Contains(t, w.Body.String(), "rate limit exceeded", "path=%s", path)
  62. }
  63. }