Fix TOTP lockout state persistence errors

internal/auth/totp.go

@@ -24,12 +24,13 @@ import (
 )
 
 var (
-	ErrTOTPNotConfigured = errors.New("TOTP authentication is not configured")
-	ErrInvalidTOTP       = errors.New("invalid TOTP code")
-	ErrTOTPReplay        = errors.New("TOTP code was already used")
-	ErrTemporarilyLocked = errors.New("authentication temporarily locked")
-	ErrEnrollmentExpired = errors.New("TOTP enrollment expired")
-	ErrRecoveryExpired   = errors.New("recovery request expired")
+	ErrTOTPNotConfigured    = errors.New("TOTP authentication is not configured")
+	ErrInvalidTOTP          = errors.New("invalid TOTP code")
+	ErrTOTPStatePersistence = errors.New("TOTP state persistence failed")
+	ErrTOTPReplay           = errors.New("TOTP code was already used")
+	ErrTemporarilyLocked    = errors.New("authentication temporarily locked")
+	ErrEnrollmentExpired    = errors.New("TOTP enrollment expired")
+	ErrRecoveryExpired      = errors.New("recovery request expired")
 )
 
 const maxFailures = 5
@@ -249,7 +250,9 @@ func (service *TOTPService) Verify(accountID, code string) error {
 			factor.FailedCount = 0
 			factor.LockedUntil = now.Add(15 * time.Minute)
 		}
-		_ = service.put("factor", accountID, factor)
+		if err := service.put("factor", accountID, factor); err != nil {
+			return fmt.Errorf("%w: persist failed-attempt state: %w", ErrTOTPStatePersistence, err)
+		}
 		return ErrInvalidTOTP
 	}
 	if step <= factor.LastUsedStep {

internal/auth/totp_test.go

@@ -3,10 +3,13 @@ package auth
 import (
 	"encoding/base64"
 	"errors"
+	"fmt"
+	"path/filepath"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/dgraph-io/badger/v4"
 	"github.com/pquerna/otp/totp"
 	"github.com/wavefnd/wave-platform/internal/storage"
 )
@@ -55,6 +58,105 @@ func TestTOTPFactorEncryptsSecretAndRejectsReplay(t *testing.T) {
 	}
 }
 
+func TestVerifySurfacesFailedAttemptPersistenceErrors(t *testing.T) {
+	for _, failedCount := range []int{0, maxFailures - 1} {
+		t.Run(fmt.Sprintf("failed-count-%d", failedCount), func(t *testing.T) {
+			database, err := storage.Open(t.TempDir())
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer database.Close()
+
+			service, err := NewTOTPService(database, base64.StdEncoding.EncodeToString([]byte("0123456789abcdef0123456789abcdef")), "Wave Test")
+			if err != nil {
+				t.Fatal(err)
+			}
+			if err := service.PutFactor("account-1", "JBSWY3DPEHPK3PXPJBSWY3DPEHPK3PXP", "recover@example.com", false); err != nil {
+				t.Fatal(err)
+			}
+			factor, err := service.Factor("account-1")
+			if err != nil {
+				t.Fatal(err)
+			}
+			factor.FailedCount = failedCount
+			if err := service.put("factor", "account-1", factor); err != nil {
+				t.Fatal(err)
+			}
+
+			var closeErr error
+			service.now = func() time.Time {
+				closeErr = database.Close()
+				return time.Unix(1_800_000_000, 0).UTC()
+			}
+
+			err = service.Verify("account-1", "not-a-code")
+			if closeErr != nil {
+				t.Fatal(closeErr)
+			}
+			if err == nil || errors.Is(err, ErrInvalidTOTP) || !strings.Contains(err.Error(), "TOTP state persistence failed") {
+				t.Fatalf("Verify() = %v; want a persistence error, not an ordinary invalid-code error", err)
+			}
+			writable, err := badger.Open(badger.DefaultOptions(filepath.Join(database.Root, "badger")).WithLogger(nil))
+			if err != nil {
+				t.Fatal(err)
+			}
+			database.DB = writable
+			stored, err := service.Factor("account-1")
+			if err != nil {
+				t.Fatal(err)
+			}
+			if stored.FailedCount != failedCount {
+				t.Fatalf("failed count = %d, want unchanged %d after failed write", stored.FailedCount, failedCount)
+			}
+		})
+	}
+}
+
+func TestVerifyPersistsFailedAttemptsAndLockoutThreshold(t *testing.T) {
+	for _, failedCount := range []int{0, maxFailures - 1} {
+		t.Run(fmt.Sprintf("failed-count-%d", failedCount), func(t *testing.T) {
+			database, err := storage.Open(t.TempDir())
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer database.Close()
+
+			service, err := NewTOTPService(database, base64.StdEncoding.EncodeToString([]byte("0123456789abcdef0123456789abcdef")), "Wave Test")
+			if err != nil {
+				t.Fatal(err)
+			}
+			if err := service.PutFactor("account-1", "JBSWY3DPEHPK3PXPJBSWY3DPEHPK3PXP", "recover@example.com", false); err != nil {
+				t.Fatal(err)
+			}
+			factor, err := service.Factor("account-1")
+			if err != nil {
+				t.Fatal(err)
+			}
+			factor.FailedCount = failedCount
+			if err := service.put("factor", "account-1", factor); err != nil {
+				t.Fatal(err)
+			}
+
+			now := time.Unix(1_800_000_000, 0).UTC()
+			service.now = func() time.Time { return now }
+			if err := service.Verify("account-1", "not-a-code"); !errors.Is(err, ErrInvalidTOTP) {
+				t.Fatalf("Verify() = %v; want ErrInvalidTOTP after a successful write", err)
+			}
+			stored, err := service.Factor("account-1")
+			if err != nil {
+				t.Fatal(err)
+			}
+			if failedCount < maxFailures-1 {
+				if stored.FailedCount != failedCount+1 || !stored.LockedUntil.IsZero() {
+					t.Fatalf("stored factor = %#v; want incremented count without lockout", stored)
+				}
+			} else if stored.FailedCount != 0 || !stored.LockedUntil.Equal(now.Add(15*time.Minute)) {
+				t.Fatalf("stored factor = %#v; want reset count and a 15-minute lockout", stored)
+			}
+		})
+	}
+}
+
 func TestNormalizeRecoveryEmailRejectsDisplayName(t *testing.T) {
 	if _, err := NormalizeRecoveryEmail("User <user@example.com>"); err == nil {
 		t.Fatal("display-name address should be rejected")

internal/identity/service.go

@@ -39,6 +39,13 @@ var (
 	ErrAddressChoiceRequired = errors.New("a different Wave mail address is required for this duplicate display name")
 )
 
+func totpVerificationError(err error) error {
+	if errors.Is(err, auth.ErrTOTPStatePersistence) {
+		return err
+	}
+	return ErrInvalidCredentials
+}
+
 const ManagementMailboxAccountID = "role/platform-management"
 const PatchMailboxAccountID = "service/patchs"
 
@@ -422,15 +429,15 @@ func (service *Service) AuthenticateTOTP(identifier, code, userAgent string) (ac
 		return account.Account{}, "", session.Session{}, ErrAccountInactive
 	}
 	if err := service.totp.Verify(item.ID, code); err != nil {
-		return account.Account{}, "", session.Session{}, ErrInvalidCredentials
+		return account.Account{}, "", session.Session{}, totpVerificationError(err)
 	}
 	token, currentSession, err := service.sessions.Create(item.ID, userAgent)
 	return item, token, currentSession, err
 }
 
 func (service *Service) BeginTOTPRotation(accountID, currentCode string) (auth.EnrollmentResult, error) {
 	if err := service.totp.Verify(accountID, currentCode); err != nil {
-		return auth.EnrollmentResult{}, ErrInvalidCredentials
+		return auth.EnrollmentResult{}, totpVerificationError(err)
 	}
 	item, err := service.accounts.Account(accountID)
 	if err != nil {
@@ -500,7 +507,7 @@ func (service *Service) SecurityStatus(accountID string) (auth.TOTPFactor, error
 
 func (service *Service) ChangeRecoveryEmail(accountID, currentCode, email string) error {
 	if err := service.totp.Verify(accountID, currentCode); err != nil {
-		return ErrInvalidCredentials
+		return totpVerificationError(err)
 	}
 	return service.sendRecoveryVerification(accountID, email)
 }
@@ -541,7 +548,7 @@ func (service *Service) UpdateProfile(accountID, displayName, bio, timeZone stri
 
 func (service *Service) ChangeWaveAddress(accountID, currentCode, localPart string) (account.Account, error) {
 	if err := service.totp.Verify(accountID, strings.TrimSpace(currentCode)); err != nil {
-		return account.Account{}, ErrInvalidCredentials
+		return account.Account{}, totpVerificationError(err)
 	}
 	localPart, err := account.MailLocalPart(localPart)
 	if err != nil {

internal/identity/service_test.go

@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"github.com/pquerna/otp/totp"
+	"github.com/wavefnd/wave-platform/internal/auth"
 	maildomain "github.com/wavefnd/wave-platform/internal/mail"
 	"github.com/wavefnd/wave-platform/internal/permission"
 	"github.com/wavefnd/wave-platform/internal/storage"
@@ -34,6 +35,16 @@ func testRegistration(displayName string) Registration {
 	return Registration{DisplayName: displayName, TOTPSecret: testTOTPSecret, RecoveryEmail: "recovery@example.net"}
 }
 
+func TestTOTPVerificationErrorPreservesStatePersistenceFailure(t *testing.T) {
+	persistenceErr := errors.Join(auth.ErrTOTPStatePersistence, errors.New("injected write failure"))
+	if got := totpVerificationError(persistenceErr); !errors.Is(got, auth.ErrTOTPStatePersistence) || errors.Is(got, ErrInvalidCredentials) {
+		t.Fatalf("totpVerificationError() = %v; want the persistence failure without invalid credentials", got)
+	}
+	if got := totpVerificationError(auth.ErrInvalidTOTP); !errors.Is(got, ErrInvalidCredentials) {
+		t.Fatalf("invalid-code error = %v; want invalid credentials", got)
+	}
+}
+
 func TestRegisterCreatesAddressAuthenticatorMailboxAndSession(t *testing.T) {
 	database, err := storage.Open(t.TempDir())
 	if err != nil {

internal/web/handler/auth.go

@@ -185,7 +185,7 @@ func (handler AuthHandler) Login(writer http.ResponseWriter, request *http.Reque
 	}
 	item, token, currentSession, err := handler.Service.AuthenticateTOTP(input.Identifier, input.Code, request.UserAgent())
 	if err != nil {
-		writeAPIError(writer, http.StatusUnauthorized, "invalid-credentials", "The username, email address, or authenticator code is incorrect.")
+		handler.writeLoginError(writer, err)
 		return
 	}
 	handler.setCookie(writer, token, currentSession)
@@ -307,7 +307,15 @@ func (handler AuthHandler) ChangeRecoveryEmail(writer http.ResponseWriter, reque
 		return
 	}
 	var input RecoveryEmailRequest
-	if xmlcodec.Decode(request.Body, &input) != nil || handler.Service.ChangeRecoveryEmail(actor.ID, input.Code, input.Email) != nil {
+	if xmlcodec.Decode(request.Body, &input) != nil {
+		writeAPIError(writer, http.StatusUnauthorized, "invalid-code", "The authenticator code or email address is invalid.")
+		return
+	}
+	if err := handler.Service.ChangeRecoveryEmail(actor.ID, input.Code, input.Email); err != nil {
+		if errors.Is(err, auth.ErrTOTPStatePersistence) {
+			handler.writeAuthError(writer, err)
+			return
+		}
 		writeAPIError(writer, http.StatusUnauthorized, "invalid-code", "The authenticator code or email address is invalid.")
 		return
 	}
@@ -368,6 +376,8 @@ func (handler AuthHandler) writeAuthError(writer http.ResponseWriter, err error)
 		writeAPIError(writer, http.StatusConflict, "account-conflict", "That Wave address is already in use.")
 	case errors.Is(err, auth.ErrTOTPNotConfigured):
 		writeAPIError(writer, http.StatusServiceUnavailable, "totp-not-configured", "TOTP authentication is not configured on this server.")
+	case errors.Is(err, auth.ErrTOTPStatePersistence):
+		writeAPIError(writer, http.StatusInternalServerError, "internal-error", "The authentication request could not be completed.")
 	case errors.Is(err, identity.ErrAddressChoiceRequired):
 		writeAPIError(writer, http.StatusConflict, "address-choice-required", err.Error())
 	case errors.Is(err, identity.ErrInvalidRegistration):
@@ -377,6 +387,14 @@ func (handler AuthHandler) writeAuthError(writer http.ResponseWriter, err error)
 	}
 }
 
+func (handler AuthHandler) writeLoginError(writer http.ResponseWriter, err error) {
+	if errors.Is(err, auth.ErrTOTPStatePersistence) {
+		handler.writeAuthError(writer, err)
+		return
+	}
+	writeAPIError(writer, http.StatusUnauthorized, "invalid-credentials", "The username, email address, or authenticator code is incorrect.")
+}
+
 func (handler AuthHandler) allowMutation(writer http.ResponseWriter, request *http.Request, _ string) bool {
 	if sameOrigin(request) {
 		return true

internal/web/handler/auth_test.go

@@ -3,13 +3,15 @@ package handler
 import (
 	"encoding/base64"
 	"encoding/xml"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/pquerna/otp/totp"
+	"github.com/wavefnd/wave-platform/internal/auth"
 	"github.com/wavefnd/wave-platform/internal/identity"
 	"github.com/wavefnd/wave-platform/internal/storage"
 )
@@ -79,6 +81,20 @@ func TestTOTPRegisterCurrentAndLogout(t *testing.T) {
 	}
 }
 
+func TestWriteAuthErrorReportsTOTPStatePersistenceFailureAsServerError(t *testing.T) {
+	response := httptest.NewRecorder()
+	handler := AuthHandler{}
+	handler.writeLoginError(response, fmt.Errorf("%w: injected write error", auth.ErrTOTPStatePersistence))
+	if response.Code != http.StatusInternalServerError {
+		t.Fatalf("TOTP persistence status = %d, want %d; body=%s", response.Code, http.StatusInternalServerError, response.Body.String())
+	}
+	invalid := httptest.NewRecorder()
+	handler.writeLoginError(invalid, identity.ErrInvalidCredentials)
+	if invalid.Code != http.StatusUnauthorized {
+		t.Fatalf("invalid-code status = %d, want %d; body=%s", invalid.Code, http.StatusUnauthorized, invalid.Body.String())
+	}
+}
+
 func TestSessionTokenWithoutCookie(t *testing.T) {
 	request := httptest.NewRequest(http.MethodGet, "http://wave.test/", nil)
 	if token, ok := sessionToken(request); ok || token != "" {