fix(auth): spend the MFA enrollment code and throttle Login's MFA leg per account (#442)

internal/api/handler.go

@@ -220,6 +220,7 @@ func NewHandler(cfg HandlerConfig) *Handler {
 		"register",
 		"approve_cancel_public",
 		"mfa",
+		"login_mfa",
 	)
 
 	return h

internal/api/handler_auth.go

@@ -37,6 +37,10 @@ func (h *Handler) login(ctx context.Context, req *events.LambdaFunctionURLReques
 	}
 	loginReq.Password = decoded
 
+	if limitErr := h.checkLoginMFARateLimit(ctx, loginReq.Email, loginReq.MFACode); limitErr != nil {
+		return nil, limitErr
+	}
+
 	response, err := h.auth.Login(ctx, loginReq)
 	if err != nil {
 		// Map MFA sentinels to machine-readable error codes so the

internal/api/handler_auth_login_mfa_limit_test.go

@@ -0,0 +1,116 @@
+package api
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"testing"
+
+	"github.com/aws/aws-lambda-go/events"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func loginReqFrom(ip, email, code string) *events.LambdaFunctionURLRequest {
+	body := fmt.Sprintf(`{"email":%q,"password":%q,"mfa_code":%q}`, email, b64("pw"), code)
+	if code == "" {
+		body = fmt.Sprintf(`{"email":%q,"password":%q}`, email, b64("pw"))
+	}
+	req := &events.LambdaFunctionURLRequest{Body: body}
+	req.RequestContext.HTTP.SourceIP = ip
+	return req
+}
+
+func requireClientCode(t *testing.T, err error, code int, msgAndArgs ...any) {
+	t.Helper()
+	ce, ok := IsClientError(err)
+	require.True(t, ok, "expected a client error, got %v", err)
+	assert.Equal(t, code, ce.code, msgAndArgs...)
+}
+
+// Issue #442: Login's TOTP / recovery-code leg is throttled per account, so
+// rotating source IPs does not reset the code-guessing budget. Email variants
+// that Login resolves to the same address share one budget.
+func TestHandler_login_MFALegRateLimitedPerEmail(t *testing.T) {
+	limit := getDefaultRateLimits()["login_mfa"].MaxAttempts
+	require.Less(t, limit, 100)
+	emails := []string{"victim@example.com", "Victim@Example.com", `"Someone" <victim@example.com>`}
+
+	mockAuth := new(MockAuthService)
+	t.Cleanup(func() { mockAuth.AssertExpectations(t) })
+	mockAuth.On("Login", mock.Anything, mock.Anything).
+		Return((*LoginResponse)(nil), ErrInvalidMFACode_test()).Times(limit)
+	h := &Handler{auth: mockAuth, rateLimiter: NewInMemoryRateLimiter()}
+
+	for i := 0; i <= limit; i++ {
+		ip := fmt.Sprintf("198.51.100.%d", i+1) // fresh IP each time: the per-IP bucket never trips
+		_, err := h.login(context.Background(), loginReqFrom(ip, emails[i%len(emails)], "000000"))
+		if i < limit {
+			requireClientCode(t, err, 401, "attempt %d", i+1)
+		} else {
+			requireClientCode(t, err, 429, "attempt %d must be throttled", i+1)
+		}
+	}
+
+	// Another account keeps its own budget.
+	mockAuth.On("Login", mock.Anything, mock.Anything).
+		Return((*LoginResponse)(nil), ErrInvalidMFACode_test()).Once()
+	_, err := h.login(context.Background(), loginReqFrom("203.0.113.1", "other@example.com", "000000"))
+	requireClientCode(t, err, 401)
+}
+
+// The password-only first leg never touches the per-email bucket, so users
+// without MFA and the mfa_required round trip are unaffected.
+func TestHandler_login_NoMFACodeSkipsPerEmailBucket(t *testing.T) {
+	mockAuth := new(MockAuthService)
+	t.Cleanup(func() { mockAuth.AssertExpectations(t) })
+	mockAuth.On("Login", mock.Anything, mock.Anything).
+		Return((*LoginResponse)(nil), ErrMFARequired_test()).Once()
+	rl := new(MockRateLimiter)
+	t.Cleanup(func() { rl.AssertExpectations(t) })
+	rl.On("AllowWithIP", mock.Anything, "198.51.100.1", "login").Return(true, nil).Once()
+	h := &Handler{auth: mockAuth, rateLimiter: rl}
+
+	_, err := h.login(context.Background(), loginReqFrom("198.51.100.1", "u@example.com", ""))
+	requireClientCode(t, err, 401)
+	rl.AssertNotCalled(t, "AllowWithEmail", mock.Anything, mock.Anything, mock.Anything)
+}
+
+func TestHandler_login_MFALegRateLimiterErrorFailsClosed(t *testing.T) {
+	mockAuth := new(MockAuthService)
+	rl := new(MockRateLimiter)
+	t.Cleanup(func() { rl.AssertExpectations(t) })
+	rl.On("AllowWithIP", mock.Anything, mock.Anything, "login").Return(true, nil).Once()
+	rl.On("AllowWithEmail", mock.Anything, "u@example.com", "login_mfa").Return(false, errors.New("db down")).Once()
+	h := &Handler{auth: mockAuth, rateLimiter: rl}
+
+	_, err := h.login(context.Background(), loginReqFrom("198.51.100.1", "U@example.com", "000000"))
+	requireClientCode(t, err, 503)
+	mockAuth.AssertNotCalled(t, "Login", mock.Anything, mock.Anything)
+}
+
+// Issue #442: requireMFASession checks the session before the per-IP "mfa"
+// bucket, so requests without a valid session cannot spend the budget of
+// users behind the same IP.
+func TestHandler_mfaEndpoints_InvalidSessionDoesNotSpendIPBudget(t *testing.T) {
+	limit := getDefaultRateLimits()["mfa"].MaxAttempts
+	mockAuth := new(MockAuthService)
+	t.Cleanup(func() { mockAuth.AssertExpectations(t) })
+	mockAuth.On("ValidateSession", mock.Anything, "bad").Return((*Session)(nil), errors.New("invalid")).Times(limit + 1)
+	mockAuth.On("ValidateSession", mock.Anything, "tok").Return(&Session{UserID: "user-1"}, nil).Once()
+	mockAuth.On("MFAEnableAPI", mock.Anything, "user-1", "000000").Return([]string(nil), nil).Once()
+	h := &Handler{auth: mockAuth, rateLimiter: NewInMemoryRateLimiter()}
+
+	for i := 0; i <= limit; i++ {
+		req := authedReq("bad", `{"code":"000000"}`)
+		req.RequestContext.HTTP.SourceIP = "198.51.100.1"
+		_, err := h.mfaEnable(context.Background(), req)
+		requireClientCode(t, err, 401, "attempt %d", i+1)
+	}
+
+	req := authedReq("tok", `{"code":"000000"}`)
+	req.RequestContext.HTTP.SourceIP = "198.51.100.1"
+	_, err := h.mfaEnable(context.Background(), req)
+	require.NoError(t, err, "a valid session on the same IP must still have its budget")
+}

internal/api/middleware.go

@@ -535,16 +535,19 @@ func (h *Handler) checkRateLimitStrict(ctx context.Context, req *events.LambdaFu
 
 // requireMFASession is requireSession plus the "mfa" bucket applied per IP and
 // then per session user, both fail-closed, so a stolen session cannot reset its
-// guessing budget by rotating IPs (#94).
+// guessing budget by rotating IPs (#94). The session is checked first so
+// requests without a valid session are refused with 401 before they reach
+// either bucket and cannot spend the per-IP budget of users behind the same
+// NAT (#442).
 func (h *Handler) requireMFASession(ctx context.Context, req *events.LambdaFunctionURLRequest) (*Session, error) {
 	const endpoint = "mfa"
-	if err := h.checkRateLimitStrict(ctx, req, endpoint); err != nil {
-		return nil, err
-	}
 	session, err := h.requireSession(ctx, req)
 	if err != nil {
 		return nil, err
 	}
+	if limitErr := h.checkRateLimitStrict(ctx, req, endpoint); limitErr != nil {
+		return nil, limitErr
+	}
 	if h.rateLimiter == nil {
 		return session, nil
 	}

internal/api/rate_limiter.go

@@ -2,7 +2,12 @@
 package api
 
 import (
+	"context"
+	"net/mail"
+	"strings"
 	"time"
+
+	"github.com/LeanerCloud/cloud-commitments-go/pkg/logging"
 )
 
 // RateLimitConfig defines the rate limiting parameters for a specific endpoint/operation.
@@ -44,6 +49,10 @@ func getDefaultRateLimits() map[string]RateLimitConfig {
 		// recovery code. Applied per IP AND per session user (#94), so rotating
 		// source IPs does not reset one account's guessing budget.
 		"mfa": NewRateLimitConfig(10, 15*60), // 10 attempts / 15 minutes / IP and / user (#94)
+		// login_mfa is Login's TOTP / recovery-code leg, keyed per account email
+		// on top of the per-IP login bucket, so rotating source IPs does not
+		// reset one account's code-guessing budget (#442).
+		"login_mfa": NewRateLimitConfig(10, 15*60), // 10 attempts / 15 minutes / email (#442)
 	}
 }
 
@@ -61,3 +70,39 @@ func assertRateLimitKeysKnown(keys ...string) {
 		}
 	}
 }
+
+// checkLoginMFARateLimit throttles Login's TOTP / recovery-code leg per account
+// (#442): the "login" bucket is per IP only, so a caller holding the password
+// could rotate source IPs to keep guessing codes. It runs only when the request
+// carries an MFA code, before any password or code is checked, and fails closed
+// like checkRateLimitStrict.
+//
+// The key is the email exactly as Login resolves it (display name stripped by
+// mail.ParseAddress), lower-cased so case variants share one budget. An email
+// Login would reject as malformed is not counted; Login fails before any code
+// is checked. A caller without the password can spend a victim's budget, but
+// wrong passwords already lock that account for AccountLockoutDuration, so this
+// adds no new denial of service. The 429 does not depend on whether the
+// account exists.
+func (h *Handler) checkLoginMFARateLimit(ctx context.Context, email, mfaCode string) error {
+	if h.rateLimiter == nil || mfaCode == "" {
+		return nil
+	}
+	parsed, err := mail.ParseAddress(email)
+	if err != nil {
+		return nil
+	}
+	const endpoint = "login_mfa"
+	key := strings.ToLower(parsed.Address)
+	allowed, err := h.rateLimiter.AllowWithEmail(ctx, key, endpoint)
+	if err != nil {
+		logging.Errorf("ALERT: rate limiter error on credential endpoint %s for %s; request denied (02-M1): %v",
+			endpoint, redactEmail(key), err)
+		return NewClientError(503, "service temporarily unavailable, please try again later")
+	}
+	if !allowed {
+		logging.Warnf("Rate limit exceeded for %s for: %s", endpoint, redactEmail(key))
+		return NewClientError(429, "too many requests, please try again later")
+	}
+	return nil
+}

internal/auth/interfaces.go

@@ -19,7 +19,7 @@ type StoreInterface interface {
 	ConsumePasswordResetToken(ctx context.Context, userID, readResetToken string) error
 	ConsumeMFARecoveryCode(ctx context.Context, userID string, readCodes, remaining []string) error
 	SetPendingMFASecret(ctx context.Context, userID, secret string, expiresAt time.Time) error
-	EnableMFA(ctx context.Context, userID, pendingSecret string, recoveryHashes []string) error
+	EnableMFA(ctx context.Context, userID, pendingSecret string, totpCounter int64, recoveryHashes []string) error
 	ReplaceMFARecoveryCodes(ctx context.Context, userID, readSecret string, readCodes, newHashes []string) error
 	DisableMFA(ctx context.Context, userID, readSecret string, readCodes []string) error
 	ClearPendingMFASecret(ctx context.Context, userID string) error

internal/auth/service_credentials_db_test.go

@@ -421,7 +421,9 @@ func TestIntegration_ResetConfirmRejectsStaleRead(t *testing.T) {
 func (f *credentialRaceFixture) enrollMFA() []string {
 	setup, err := f.svc.MFASetup(f.t.Context(), f.user.ID, credentialRacePassword)
 	require.NoError(f.t, err)
-	codes, err := f.svc.MFAEnable(f.t.Context(), f.user.ID, generateTOTP(setup.Secret, time.Now().Unix()/30))
+	// Enroll with the previous step's code: MFAEnable spends its step (#442), and
+	// callers then sign in with the current one.
+	codes, err := f.svc.MFAEnable(f.t.Context(), f.user.ID, generateTOTP(setup.Secret, time.Now().Unix()/30-1))
 	require.NoError(f.t, err)
 	return codes
 }

internal/auth/service_mfa.go

@@ -356,19 +356,21 @@ func (s *Service) generateAndHashRecoveryCodes() (plaintext, hashes []string, er
 }
 
 // validatePendingMFAEnrollment checks that a pending enrollment
-// exists, hasn't expired, and the supplied code matches. It writes
-// nothing: the next setup overwrites an expired pending secret.
-func validatePendingMFAEnrollment(user *User, code string) error {
+// exists, hasn't expired, and the supplied code matches, returning the
+// matched TOTP counter. It writes nothing: the next setup overwrites an
+// expired pending secret.
+func validatePendingMFAEnrollment(user *User, code string) (int64, error) {
 	if user.MFAPendingSecret == "" || user.MFAPendingSecretExpiresAt == nil {
-		return fmt.Errorf("%w", ErrMFANoEnrollmentInProgress)
+		return 0, fmt.Errorf("%w", ErrMFANoEnrollmentInProgress)
 	}
 	if time.Now().After(*user.MFAPendingSecretExpiresAt) {
-		return fmt.Errorf("%w", ErrMFAEnrollmentExpired)
+		return 0, fmt.Errorf("%w", ErrMFAEnrollmentExpired)
 	}
-	if !verifyTOTP(user.MFAPendingSecret, code) {
-		return fmt.Errorf("%w", ErrMFAInvalidCode)
+	counter, ok := matchTOTP(user.MFAPendingSecret, code)
+	if !ok {
+		return 0, fmt.Errorf("%w", ErrMFAInvalidCode)
 	}
-	return nil
+	return counter, nil
 }
 
 // MFAEnable finalizes an MFA enrollment. Validates the supplied TOTP
@@ -398,7 +400,7 @@ func (s *Service) MFAEnable(ctx context.Context, userID, code string) ([]string,
 	if user.MFAEnabled {
 		return nil, fmt.Errorf("%w", ErrMFAAlreadyEnabled)
 	}
-	err = validatePendingMFAEnrollment(user, code)
+	counter, err := validatePendingMFAEnrollment(user, code)
 	if errors.Is(err, ErrMFAInvalidCode) {
 		return nil, s.rejectMFACode(ctx, user, code)
 	}
@@ -410,7 +412,10 @@ func (s *Service) MFAEnable(ctx context.Context, userID, code string) ([]string,
 	if err != nil {
 		return nil, err
 	}
-	if err := s.store.EnableMFA(ctx, user.ID, user.MFAPendingSecret, hashes); err != nil {
+	// EnableMFA records the enrollment code's counter in the same write, so the
+	// code typed to turn MFA on cannot be replayed at Login inside its window
+	// (#442).
+	if err := s.store.EnableMFA(ctx, user.ID, user.MFAPendingSecret, counter, hashes); err != nil {
 		return nil, err
 	}
 	return plaintext, nil

internal/auth/service_mfa_enroll_counter_test.go

@@ -0,0 +1,42 @@
+package auth
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+// currentTOTPStep matches the counter of a code generated for the current step,
+// allowing the clock to cross one step boundary while the test runs.
+func currentTOTPStep() any {
+	return mock.MatchedBy(func(counter int64) bool {
+		now := time.Now().Unix() / 30
+		return counter == now || counter == now-1
+	})
+}
+
+// Issue #442: MFAEnable hands the enrollment code's counter to EnableMFA, so the
+// code that turned MFA on cannot be spent again at Login.
+func TestMFAEnable_RecordsEnrollmentCodeCounter(t *testing.T) {
+	ctx := context.Background()
+	mockStore := new(MockStore)
+	service := createTestService(mockStore, new(MockEmailSender))
+
+	secret := "JBSWY3DPEHPK3PXP"
+	expiresAt := time.Now().Add(mfaPendingExpiry)
+	user := createTestUser(t, "SecurePass@123")
+	user.MFAPendingSecret = secret
+	user.MFAPendingSecretExpiresAt = &expiresAt
+	step := time.Now().Unix()/30 - 1
+
+	mockStore.On("GetUserByID", ctx, user.ID).Return(user, nil)
+	mockStore.On("EnableMFA", ctx, user.ID, secret, step, mock.AnythingOfType("[]string")).Return(nil).Once()
+
+	_, err := service.MFAEnable(ctx, user.ID, generateTOTP(secret, step))
+	require.NoError(t, err)
+	mockStore.AssertExpectations(t)
+	mockStore.AssertNotCalled(t, "ClaimTOTPCounter", mock.Anything, mock.Anything, mock.Anything)
+}

internal/auth/service_mfa_test.go

@@ -250,8 +250,8 @@ func TestMFAEnable_HappyPath(t *testing.T) {
 	user.MFAPendingSecretExpiresAt = &expiresAt
 
 	mockStore.On("GetUserByID", ctx, user.ID).Return(user, nil)
-	mockStore.On("EnableMFA", ctx, user.ID, secret, mock.AnythingOfType("[]string")).Run(func(args mock.Arguments) {
-		assert.Len(t, args.Get(3).([]string), recoveryCodeCount)
+	mockStore.On("EnableMFA", ctx, user.ID, secret, currentTOTPStep(), mock.AnythingOfType("[]string")).Run(func(args mock.Arguments) {
+		assert.Len(t, args.Get(4).([]string), recoveryCodeCount)
 	}).Return(nil).Once()
 
 	codes, err := service.MFAEnable(ctx, user.ID, totpFor(secret))
@@ -806,7 +806,7 @@ func TestMFALifecycle_LockedAccountRefusedBeforeCodeCheck(t *testing.T) {
 			// Permissive stubs so an unfixed service fails by assertion, not by panic.
 			writes := []string{"ClaimTOTPCounter", "EnableMFA", "DisableMFA", "ReplaceMFARecoveryCodes", "ClearPendingMFASecret", "RecordFailedLogin"}
 			mockStore.On("ClaimTOTPCounter", mock.Anything, user.ID, mock.Anything).Return(true, nil).Maybe()
-			mockStore.On("EnableMFA", mock.Anything, user.ID, mock.Anything, mock.Anything).Return(nil).Maybe()
+			mockStore.On("EnableMFA", mock.Anything, user.ID, mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
 			mockStore.On("DisableMFA", mock.Anything, user.ID, mock.Anything, mock.Anything).Return(nil).Maybe()
 			mockStore.On("ReplaceMFARecoveryCodes", mock.Anything, user.ID, mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
 			mockStore.On("ClearPendingMFASecret", mock.Anything, user.ID).Return(nil).Maybe()