step_up_test.go 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166
  1. package middleware
  2. import (
  3. "context"
  4. "errors"
  5. "net/http"
  6. "net/http/httptest"
  7. "testing"
  8. "github.com/Wei-Shaw/sub2api/internal/service"
  9. "github.com/gin-gonic/gin"
  10. "github.com/stretchr/testify/require"
  11. )
  12. type stubStepUpGrantChecker struct {
  13. granted bool
  14. err error
  15. }
  16. func (s stubStepUpGrantChecker) HasStepUpGrant(ctx context.Context, userID int64, sessionKey string) (bool, error) {
  17. return s.granted, s.err
  18. }
  19. type stubStepUpUserReader struct {
  20. user *service.User
  21. err error
  22. }
  23. func (s stubStepUpUserReader) GetByID(ctx context.Context, id int64) (*service.User, error) {
  24. return s.user, s.err
  25. }
  26. type stubStepUpSettingReader struct {
  27. enabled bool
  28. }
  29. func (s stubStepUpSettingReader) IsStepUpEnabled(ctx context.Context) bool {
  30. return s.enabled
  31. }
  32. // stepUpEnabled 功能开关开启的设置桩,供既有门控分支测试使用。
  33. var stepUpEnabled = stubStepUpSettingReader{enabled: true}
  34. func newStepUpTestContext(t *testing.T) (*gin.Context, *httptest.ResponseRecorder) {
  35. t.Helper()
  36. gin.SetMode(gin.TestMode)
  37. rec := httptest.NewRecorder()
  38. c, _ := gin.CreateTestContext(rec)
  39. c.Request = httptest.NewRequest(http.MethodPost, "/sensitive", nil)
  40. return c, rec
  41. }
  42. func TestEnforceStepUpRejectsAdminAPIKey(t *testing.T) {
  43. c, rec := newStepUpTestContext(t)
  44. c.Set("auth_method", service.AuditAuthMethodAdminAPIKey)
  45. ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{TotpEnabled: true}}, stepUpEnabled)
  46. require.False(t, ok)
  47. require.True(t, c.IsAborted())
  48. require.Equal(t, http.StatusForbidden, rec.Code)
  49. require.Contains(t, rec.Body.String(), "STEP_UP_ADMIN_API_KEY_FORBIDDEN")
  50. }
  51. func TestEnforceStepUpRequiresAuthSubject(t *testing.T) {
  52. c, rec := newStepUpTestContext(t)
  53. ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{TotpEnabled: true}}, stepUpEnabled)
  54. require.False(t, ok)
  55. require.Equal(t, http.StatusUnauthorized, rec.Code)
  56. }
  57. func TestEnforceStepUpRequiresTotpEnabled(t *testing.T) {
  58. c, rec := newStepUpTestContext(t)
  59. c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
  60. ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: false}}, stepUpEnabled)
  61. require.False(t, ok)
  62. require.Equal(t, http.StatusForbidden, rec.Code)
  63. require.Contains(t, rec.Body.String(), "STEP_UP_TOTP_NOT_ENABLED")
  64. }
  65. func TestEnforceStepUpFailsClosedOnGrantError(t *testing.T) {
  66. c, rec := newStepUpTestContext(t)
  67. c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
  68. ok := enforceStepUp(c, stubStepUpGrantChecker{err: errors.New("redis down")}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, stepUpEnabled)
  69. require.False(t, ok)
  70. require.Equal(t, http.StatusServiceUnavailable, rec.Code)
  71. require.Contains(t, rec.Body.String(), "STEP_UP_UNAVAILABLE")
  72. }
  73. func TestEnforceStepUpRequiresGrant(t *testing.T) {
  74. c, rec := newStepUpTestContext(t)
  75. c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
  76. ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, stepUpEnabled)
  77. require.False(t, ok)
  78. require.Equal(t, http.StatusForbidden, rec.Code)
  79. require.Contains(t, rec.Body.String(), "STEP_UP_REQUIRED")
  80. }
  81. func TestEnforceStepUpPassesWithGrant(t *testing.T) {
  82. c, _ := newStepUpTestContext(t)
  83. c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
  84. ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, stepUpEnabled)
  85. require.True(t, ok)
  86. require.False(t, c.IsAborted())
  87. }
  88. // 功能开关关闭时:不论 TOTP/grant/凭证类型,一律放行(恢复门控引入前行为)。
  89. func TestEnforceStepUpDisabledSkipsAllChecks(t *testing.T) {
  90. disabled := stubStepUpSettingReader{enabled: false}
  91. t.Run("no totp, no grant", func(t *testing.T) {
  92. c, _ := newStepUpTestContext(t)
  93. c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
  94. ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: false}}, disabled)
  95. require.True(t, ok)
  96. require.False(t, c.IsAborted())
  97. })
  98. t.Run("admin api key", func(t *testing.T) {
  99. c, _ := newStepUpTestContext(t)
  100. c.Set("auth_method", service.AuditAuthMethodAdminAPIKey)
  101. ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: nil, err: errors.New("should not be called")}, disabled)
  102. require.True(t, ok)
  103. require.False(t, c.IsAborted())
  104. })
  105. }
  106. // settings 为 nil 时保持门控(fail-closed),避免装配缺陷静默关闭安全控制。
  107. func TestEnforceStepUpNilSettingsFailsClosed(t *testing.T) {
  108. c, rec := newStepUpTestContext(t)
  109. c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
  110. ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, nil)
  111. require.False(t, ok)
  112. require.Equal(t, http.StatusForbidden, rec.Code)
  113. require.Contains(t, rec.Body.String(), "STEP_UP_REQUIRED")
  114. }
  115. // EnforceStepUp 收到 nil *service.SettingService 时不得因 typed-nil 装箱绕过门控:
  116. // 未认证请求仍应被拦截(401),而不是当作"开关关闭"放行。
  117. func TestEnforceStepUpTypedNilSettingServiceFailsClosed(t *testing.T) {
  118. require.Nil(t, stepUpSettingsOrNil(nil))
  119. c, rec := newStepUpTestContext(t)
  120. ok := EnforceStepUp(c, nil, nil, nil)
  121. require.False(t, ok)
  122. require.Equal(t, http.StatusUnauthorized, rec.Code)
  123. }