fix(auth): lock the account on wrong MFA codes at enable, disable and regenerate

internal/auth/service_mfa.go

@@ -7,12 +7,14 @@ import (
 	"crypto/sha1" // #nosec G505 -- SHA-1 retained for broad authenticator app compatibility and existing otpauth provisioning; RFC 6238 permits SHA-256/SHA-512 but most authenticator apps default to SHA-1
 	"crypto/subtle"
 	"encoding/base32"
+	"errors"
 	"fmt"
 	"net/url"
 	"slices"
 	"strings"
 	"time"
 
+	"github.com/LeanerCloud/cloud-commitments-go/pkg/logging"
 	"github.com/LeanerCloud/cloud-commitments-platform/internal/config"
 	"golang.org/x/crypto/bcrypt"
 )
@@ -390,10 +392,16 @@ func (s *Service) MFAEnable(ctx context.Context, userID, code string) ([]string,
 	if err != nil || user == nil {
 		return nil, fmt.Errorf("%w", ErrMFAAuthFailed)
 	}
+	if lockErr := s.rejectLockedMFAUser(user); lockErr != nil {
+		return nil, lockErr
+	}
 	if user.MFAEnabled {
 		return nil, fmt.Errorf("%w", ErrMFAAlreadyEnabled)
 	}
 	err = validatePendingMFAEnrollment(user, code)
+	if errors.Is(err, ErrMFAInvalidCode) {
+		return nil, s.rejectMFACode(ctx, user, code)
+	}
 	if err != nil {
 		return nil, err
 	}
@@ -415,6 +423,40 @@ func (s *Service) disableMFAAlreadyOff(ctx context.Context, user *User) error {
 	return s.store.ClearPendingMFASecret(ctx, user.ID)
 }
 
+// rejectLockedMFAUser refuses MFA lifecycle calls while the account is locked
+// by failed attempts, using the same opaque error as a missing user so a
+// stolen session learns nothing about the lock (same stance as Login).
+func (s *Service) rejectLockedMFAUser(user *User) error {
+	if user.LockedUntil != nil && time.Now().Before(*user.LockedUntil) {
+		logging.Warnf("MFA operation refused for locked account (locked for %v more)", time.Until(*user.LockedUntil).Round(time.Minute))
+		return fmt.Errorf("%w", ErrMFAAuthFailed)
+	}
+	return nil
+}
+
+// rejectMFACode counts a wrong or replayed code against the account lockout,
+// the same counter Login uses, and returns the invalid-code error. An empty
+// code is a missing field, not a guess, and (as in Login) is not counted.
+func (s *Service) rejectMFACode(ctx context.Context, user *User, code string) error {
+	if code != "" {
+		s.recordFailedLogin(ctx, user)
+	}
+	return fmt.Errorf("%w", ErrMFAInvalidCode)
+}
+
+// claimTOTPOrReject spends an unused TOTP code for the user, counting a wrong
+// or replayed one against the lockout.
+func (s *Service) claimTOTPOrReject(ctx context.Context, user *User, code string) error {
+	valid, err := s.verifyAndClaimTOTP(ctx, user.ID, user.MFASecret, code)
+	if err != nil {
+		return fmt.Errorf("failed to record TOTP use: %w", err)
+	}
+	if !valid {
+		return s.rejectMFACode(ctx, user, code)
+	}
+	return nil
+}
+
 // verifyMFAProof accepts an unspent TOTP code, or else consumes a recovery code
 // in memory; the caller persists the consumed slice through its own CAS write.
 func (s *Service) verifyMFAProof(ctx context.Context, user *User, codeOrRecovery string) error {
@@ -423,7 +465,7 @@ func (s *Service) verifyMFAProof(ctx context.Context, user *User, codeOrRecovery
 		return fmt.Errorf("failed to record TOTP use: %w", err)
 	}
 	if !matched && !s.consumeRecoveryCode(user, codeOrRecovery) {
-		return fmt.Errorf("%w", ErrMFAInvalidCode)
+		return s.rejectMFACode(ctx, user, codeOrRecovery)
 	}
 	return nil
 }
@@ -445,6 +487,9 @@ func (s *Service) MFADisable(ctx context.Context, userID, password, codeOrRecove
 	if err != nil || user == nil {
 		return fmt.Errorf("%w", ErrMFAAuthFailed)
 	}
+	if lockErr := s.rejectLockedMFAUser(user); lockErr != nil {
+		return lockErr
+	}
 	if !s.verifyPassword(password, user.PasswordHash) {
 		return fmt.Errorf("%w", ErrMFAInvalidPassword)
 	}
@@ -482,15 +527,14 @@ func (s *Service) MFARegenerateRecoveryCodes(ctx context.Context, userID, code s
 	if err != nil || user == nil {
 		return nil, fmt.Errorf("%w", ErrMFAAuthFailed)
 	}
+	if lockErr := s.rejectLockedMFAUser(user); lockErr != nil {
+		return nil, lockErr
+	}
 	if !user.MFAEnabled || user.MFASecret == "" {
 		return nil, fmt.Errorf("%w", ErrMFANotEnabled)
 	}
-	valid, err := s.verifyAndClaimTOTP(ctx, user.ID, user.MFASecret, code)
-	if err != nil {
-		return nil, fmt.Errorf("failed to record TOTP use: %w", err)
-	}
-	if !valid {
-		return nil, fmt.Errorf("%w", ErrMFAInvalidCode)
+	if claimErr := s.claimTOTPOrReject(ctx, user, code); claimErr != nil {
+		return nil, claimErr
 	}
 
 	plaintext, hashes, err := s.generateAndHashRecoveryCodes()

internal/auth/service_mfa_lockout_db_test.go

@@ -0,0 +1,162 @@
+//go:build integration
+
+package auth
+
+import (
+	"strconv"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// Issue #442: wrong codes on the MFA lifecycle endpoints share Login's durable
+// failed-attempt counter and lockout, so a stolen session cannot guess codes
+// without limit.
+
+// wrongMFACode returns a six-digit code outside the accepted TOTP window.
+func wrongMFACode(t *testing.T, secret string) string {
+	t.Helper()
+	for n := 0; n < 1000; n++ {
+		code := strconv.Itoa(100000 + n)
+		if _, ok := matchTOTP(secret, code); !ok {
+			return code
+		}
+	}
+	t.Fatal("no wrong code found")
+	return ""
+}
+
+func (f *credentialRaceFixture) currentTOTP() string {
+	return generateTOTP(f.stored().MFASecret, time.Now().Unix()/30)
+}
+
+// mfaLockoutOps are the endpoints under test. Each takes a code and runs the
+// operation against an account prepared by prepare.
+var mfaLockoutOps = []struct {
+	name    string
+	enrolls bool // true: starts with MFA on; false: starts with a pending enrollment
+	call    func(f *credentialRaceFixture, code string) error
+}{
+	{"disable", true, func(f *credentialRaceFixture, code string) error {
+		return f.svc.MFADisable(f.t.Context(), f.user.ID, credentialRacePassword, code)
+	}},
+	{"regenerate", true, func(f *credentialRaceFixture, code string) error {
+		_, err := f.svc.MFARegenerateRecoveryCodes(f.t.Context(), f.user.ID, code)
+		return err
+	}},
+	{"enable", false, func(f *credentialRaceFixture, code string) error {
+		_, err := f.svc.MFAEnable(f.t.Context(), f.user.ID, code)
+		return err
+	}},
+}
+
+// prepare returns the secret the endpoint checks codes against.
+func prepareMFALockoutAccount(f *credentialRaceFixture, enrolled bool) string {
+	if enrolled {
+		f.enrollMFA()
+		return f.stored().MFASecret
+	}
+	setup, err := f.svc.MFASetup(f.t.Context(), f.user.ID, credentialRacePassword)
+	require.NoError(f.t, err)
+	return setup.Secret
+}
+
+func TestIntegration_MFALifecycleWrongCodesLockAccount(t *testing.T) {
+	store := NewPostgresStore(setupAuthTestDB(t))
+
+	// Baseline: the same number of wrong codes at Login locks the account.
+	t.Run("login-baseline", func(t *testing.T) {
+		f := newCredentialRaceFixture(t, store, "lockout-login@example.com")
+		f.enrollMFA()
+		secret := f.stored().MFASecret
+		for i := 0; i < MaxFailedLoginAttempts; i++ {
+			_, err := f.svc.Login(t.Context(), LoginRequest{Email: f.user.Email, Password: credentialRacePassword, MFACode: wrongMFACode(t, secret)})
+			require.ErrorIs(t, err, ErrInvalidMFACode)
+		}
+		require.NotNil(t, f.stored().LockedUntil)
+	})
+
+	for _, op := range mfaLockoutOps {
+		t.Run(op.name+"-locks-after-max-failures", func(t *testing.T) {
+			f := newCredentialRaceFixture(t, store, "lockout-"+op.name+"@example.com")
+			secret := prepareMFALockoutAccount(f, op.enrolls)
+			before := f.stored()
+
+			for i := 0; i < MaxFailedLoginAttempts; i++ {
+				require.ErrorIs(t, op.call(f, wrongMFACode(t, secret)), ErrMFAInvalidCode, "attempt %d", i+1)
+			}
+			locked := f.stored()
+			require.NotNil(t, locked.LockedUntil, "%d wrong codes must lock the account", MaxFailedLoginAttempts)
+			assert.Equal(t, MaxFailedLoginAttempts, locked.FailedLoginAttempts)
+
+			// A correct code is now refused, with no state change.
+			require.ErrorIs(t, op.call(f, generateTOTP(secret, time.Now().Unix()/30)), ErrMFAAuthFailed)
+			after := f.stored()
+			assert.Equal(t, locked.FailedLoginAttempts, after.FailedLoginAttempts)
+			assert.Equal(t, before.MFAEnabled, after.MFAEnabled)
+			assert.Equal(t, before.MFARecoveryCodes, after.MFARecoveryCodes)
+
+			// The lock is the account's, so Login refuses it too.
+			if op.enrolls {
+				_, err := f.svc.Login(t.Context(), LoginRequest{Email: f.user.Email, Password: credentialRacePassword, MFACode: generateTOTP(secret, time.Now().Unix()/30+1)})
+				require.Error(t, err)
+			}
+		})
+	}
+
+	t.Run("below-threshold-correct-code-succeeds-and-keeps-count", func(t *testing.T) {
+		f := newCredentialRaceFixture(t, store, "lockout-below@example.com")
+		f.enrollMFA()
+		secret := f.stored().MFASecret
+		for i := 0; i < MaxFailedLoginAttempts-1; i++ {
+			require.ErrorIs(t, mfaLockoutOps[1].call(f, wrongMFACode(t, secret)), ErrMFAInvalidCode)
+		}
+		require.NoError(t, mfaLockoutOps[1].call(f, generateTOTP(secret, time.Now().Unix()/30)))
+		stored := f.stored()
+		assert.Nil(t, stored.LockedUntil)
+		assert.Equal(t, MaxFailedLoginAttempts-1, stored.FailedLoginAttempts,
+			"like Login, only a completed login resets the counter")
+
+		// One more failure reaches the threshold: the correct code did not buy a fresh budget.
+		require.ErrorIs(t, mfaLockoutOps[1].call(f, wrongMFACode(t, secret)), ErrMFAInvalidCode)
+		require.NotNil(t, f.stored().LockedUntil)
+	})
+
+	t.Run("replayed-code-counts-as-failure", func(t *testing.T) {
+		f := newCredentialRaceFixture(t, store, "lockout-replay@example.com")
+		f.enrollMFA()
+		code := f.currentTOTP()
+		require.NoError(t, mfaLockoutOps[1].call(f, code))
+		require.Equal(t, 0, f.stored().FailedLoginAttempts)
+		require.ErrorIs(t, mfaLockoutOps[1].call(f, code), ErrMFAInvalidCode)
+		assert.Equal(t, 1, f.stored().FailedLoginAttempts)
+	})
+
+	t.Run("empty-code-is-not-counted", func(t *testing.T) {
+		f := newCredentialRaceFixture(t, store, "lockout-empty@example.com")
+		f.enrollMFA()
+		require.Error(t, mfaLockoutOps[1].call(f, ""))
+		assert.Equal(t, 0, f.stored().FailedLoginAttempts)
+	})
+
+	t.Run("concurrent-failures-count-atomically", func(t *testing.T) {
+		f := newCredentialRaceFixture(t, store, "lockout-concurrent@example.com")
+		f.enrollMFA()
+		secret := f.stored().MFASecret
+		var wg sync.WaitGroup
+		for i := 0; i < MaxFailedLoginAttempts; i++ {
+			wg.Add(1)
+			go func(code string) {
+				defer wg.Done()
+				assert.ErrorIs(t, mfaLockoutOps[0].call(f, code), ErrMFAInvalidCode)
+			}(wrongMFACode(t, secret))
+		}
+		wg.Wait()
+		stored := f.stored()
+		assert.Equal(t, MaxFailedLoginAttempts, stored.FailedLoginAttempts, "no failure may be lost")
+		assert.NotNil(t, stored.LockedUntil)
+	})
+}

internal/auth/service_mfa_test.go

@@ -302,6 +302,7 @@ func TestMFAEnable_WrongCode(t *testing.T) {
 	user.MFAPendingSecret = secret
 	user.MFAPendingSecretExpiresAt = &expiresAt
 	mockStore.On("GetUserByID", ctx, user.ID).Return(user, nil)
+	mockStore.On("RecordFailedLogin", ctx, user.ID).Return(nil).Once()
 
 	_, err := service.MFAEnable(ctx, user.ID, "000000")
 	require.Error(t, err)
@@ -455,6 +456,7 @@ func TestMFARegenerateRecoveryCodes_WrongTOTP(t *testing.T) {
 	user.MFAEnabled = true
 	user.MFASecret = "JBSWY3DPEHPK3PXP"
 	mockStore.On("GetUserByID", ctx, user.ID).Return(user, nil)
+	mockStore.On("RecordFailedLogin", ctx, user.ID).Return(nil).Once()
 
 	_, err := service.MFARegenerateRecoveryCodes(ctx, user.ID, "000000")
 	require.Error(t, err)
@@ -595,6 +597,7 @@ func TestMFAEnable_WrongCode_ReturnsSentinel(t *testing.T) {
 	user.MFAPendingSecret = secret
 	user.MFAPendingSecretExpiresAt = &expiresAt
 	mockStore.On("GetUserByID", ctx, user.ID).Return(user, nil)
+	mockStore.On("RecordFailedLogin", ctx, user.ID).Return(nil).Once()
 
 	_, err := service.MFAEnable(ctx, user.ID, "000000")
 	require.Error(t, err)
@@ -647,6 +650,7 @@ func TestMFADisable_WrongCode_ReturnsSentinel(t *testing.T) {
 	user.MFASecret = "JBSWY3DPEHPK3PXP"
 	user.MFARecoveryCodes = []string{"$2a$04$hashedstub"} // won't match any real code
 	mockStore.On("GetUserByID", ctx, user.ID).Return(user, nil)
+	mockStore.On("RecordFailedLogin", ctx, user.ID).Return(nil).Once()
 
 	err := service.MFADisable(ctx, user.ID, "SecurePass@123", "000000")
 	require.Error(t, err)
@@ -680,6 +684,7 @@ func TestMFARegenerateRecoveryCodes_WrongCode_ReturnsSentinel(t *testing.T) {
 	user.MFAEnabled = true
 	user.MFASecret = "JBSWY3DPEHPK3PXP"
 	mockStore.On("GetUserByID", ctx, user.ID).Return(user, nil)
+	mockStore.On("RecordFailedLogin", ctx, user.ID).Return(nil).Once()
 
 	_, err := service.MFARegenerateRecoveryCodes(ctx, user.ID, "000000")
 	require.Error(t, err)
@@ -711,10 +716,12 @@ func TestMFADisable_ReplayedTOTPIsRejected(t *testing.T) {
 	user.MFASecret = secret
 	mockStore.On("GetUserByID", ctx, user.ID).Return(user, nil)
 	mockStore.On("ClaimTOTPCounter", ctx, user.ID, mock.AnythingOfType("int64")).Return(false, nil).Once()
+	mockStore.On("RecordFailedLogin", ctx, user.ID).Return(nil).Once()
 
 	err := service.MFADisable(ctx, user.ID, "SecurePass@123", totpFor(secret))
 	require.ErrorIs(t, err, ErrMFAInvalidCode)
 	mockStore.AssertNotCalled(t, "DisableMFA", mock.Anything, mock.Anything, mock.Anything, mock.Anything)
+	mockStore.AssertExpectations(t)
 }
 
 func TestMFARegenerateRecoveryCodes_ReplayedTOTPIsRejected(t *testing.T) {
@@ -728,6 +735,7 @@ func TestMFARegenerateRecoveryCodes_ReplayedTOTPIsRejected(t *testing.T) {
 	user.MFASecret = secret
 	mockStore.On("GetUserByID", ctx, user.ID).Return(user, nil)
 	mockStore.On("ClaimTOTPCounter", ctx, user.ID, mock.AnythingOfType("int64")).Return(false, nil).Once()
+	mockStore.On("RecordFailedLogin", ctx, user.ID).Return(nil).Once()
 
 	_, err := service.MFARegenerateRecoveryCodes(ctx, user.ID, totpFor(secret))
 	require.ErrorIs(t, err, ErrMFAInvalidCode)
@@ -767,3 +775,49 @@ func TestMFADisable_ClaimFailureFailsClosed(t *testing.T) {
 	require.ErrorContains(t, err, "db down")
 	mockStore.AssertNotCalled(t, "DisableMFA", mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 }
+
+// Issue #442: a locked account is refused before any code is checked, with
+// the opaque error and without spending a TOTP counter or writing state.
+func TestMFALifecycle_LockedAccountRefusedBeforeCodeCheck(t *testing.T) {
+	secret := "JBSWY3DPEHPK3PXP"
+	calls := map[string]func(s *Service, user *User) error{
+		"enable": func(s *Service, user *User) error {
+			_, err := s.MFAEnable(context.Background(), user.ID, totpFor(secret))
+			return err
+		},
+		"disable": func(s *Service, user *User) error {
+			return s.MFADisable(context.Background(), user.ID, "SecurePass@123", totpFor(secret))
+		},
+		"regenerate": func(s *Service, user *User) error {
+			_, err := s.MFARegenerateRecoveryCodes(context.Background(), user.ID, totpFor(secret))
+			return err
+		},
+	}
+	for name, call := range calls {
+		t.Run(name, func(t *testing.T) {
+			mockStore := new(MockStore)
+			service := createTestService(mockStore, new(MockEmailSender))
+			until := time.Now().Add(time.Minute)
+			expires := time.Now().Add(mfaPendingExpiry)
+			user := createTestUser(t, "SecurePass@123")
+			user.MFAEnabled, user.MFASecret, user.LockedUntil = true, secret, &until
+			user.MFAPendingSecret, user.MFAPendingSecretExpiresAt = secret, &expires
+			mockStore.On("GetUserByID", mock.Anything, user.ID).Return(user, nil)
+			// 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("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()
+			mockStore.On("RecordFailedLogin", mock.Anything, user.ID).Return(nil).Maybe()
+
+			require.ErrorIs(t, call(service, user), ErrMFAAuthFailed)
+			for _, method := range writes {
+				for _, c := range mockStore.Calls {
+					assert.NotEqual(t, method, c.Method, "locked account must not reach %s", method)
+				}
+			}
+		})
+	}
+}