From 3bf4282037e36c3687aa20728a4757902d5baa1a Mon Sep 17 00:00:00 2001
From: Xavier Michelon <xavier@michelon.ch>
Date: Mon, 2 Oct 2023 15:15:30 +0200
Subject: [PATCH] feat(GODT-2940): allow 3 attempts for mailbox password.

---
 internal/bridge/user.go                   | 27 +++++++++++++----
 internal/frontend/grpc/service.go         | 37 ++++++++++++++++++-----
 internal/frontend/grpc/service_methods.go |  1 +
 3 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/internal/bridge/user.go b/internal/bridge/user.go
index 594ccb7e..d98f82d1 100644
--- a/internal/bridge/user.go
+++ b/internal/bridge/user.go
@@ -46,6 +46,8 @@ const (
 	Connected
 )
 
+var ErrFailedToUnlock = errors.New("failed to unlock user keys")
+
 type UserInfo struct {
 	// UserID is the user's API ID.
 	UserID string
@@ -157,11 +159,15 @@ func (bridge *Bridge) LoginUser(
 		func() (string, error) {
 			return bridge.loginUser(ctx, client, auth.UID, auth.RefreshToken, keyPass)
 		},
-		func() error {
-			return client.AuthDelete(ctx)
-		},
 	)
+
 	if err != nil {
+		// Failure to unlock will allow retries, so we do not delete auth.
+		if !errors.Is(err, ErrFailedToUnlock) {
+			if deleteErr := client.AuthDelete(ctx); deleteErr != nil {
+				logrus.WithError(deleteErr).Error("Failed to delete auth")
+			}
+		}
 		return "", fmt.Errorf("failed to login user: %w", err)
 	}
 
@@ -217,7 +223,16 @@ func (bridge *Bridge) LoginFull(
 		keyPass = password
 	}
 
-	return bridge.LoginUser(ctx, client, auth, keyPass)
+	userID, err := bridge.LoginUser(ctx, client, auth, keyPass)
+	if err != nil {
+		if deleteErr := client.AuthDelete(ctx); deleteErr != nil {
+			logrus.WithError(err).Error("Failed to delete auth")
+		}
+
+		return "", err
+	}
+
+	return userID, nil
 }
 
 // LogoutUser logs out the given user.
@@ -374,9 +389,9 @@ func (bridge *Bridge) loginUser(ctx context.Context, client *proton.Client, auth
 	}
 
 	if userKR, err := apiUser.Keys.Unlock(saltedKeyPass, nil); err != nil {
-		return "", fmt.Errorf("failed to unlock user keys: %w", err)
+		return "", fmt.Errorf("%w: %w", ErrFailedToUnlock, err)
 	} else if userKR.CountDecryptionEntities() == 0 {
-		return "", fmt.Errorf("failed to unlock user keys")
+		return "", ErrFailedToUnlock
 	}
 
 	if err := bridge.addUser(ctx, client, apiUser, authUID, authRef, saltedKeyPass, true); err != nil {
diff --git a/internal/frontend/grpc/service.go b/internal/frontend/grpc/service.go
index ef9f8f18..bb59784d 100644
--- a/internal/frontend/grpc/service.go
+++ b/internal/frontend/grpc/service.go
@@ -54,8 +54,9 @@ import (
 )
 
 const (
-	serverConfigFileName   = "grpcServerConfig.json"
-	serverTokenMetadataKey = "server-token"
+	serverConfigFileName        = "grpcServerConfig.json"
+	serverTokenMetadataKey      = "server-token"
+	twoPasswordsMaxAttemptCount = 3 // The number of attempts allowed for the mailbox password.
 )
 
 // Service is the RPC service struct.
@@ -82,9 +83,10 @@ type Service struct { // nolint:structcheck
 	target     updater.VersionInfo
 	targetLock safe.RWMutex
 
-	authClient *proton.Client
-	auth       proton.Auth
-	password   []byte
+	authClient              *proton.Client
+	auth                    proton.Auth
+	password                []byte
+	twoPasswordAttemptCount int
 
 	log                *logrus.Entry
 	initializing       sync.WaitGroup
@@ -408,7 +410,12 @@ func (s *Service) loginClean() {
 }
 
 func (s *Service) finishLogin() {
-	defer s.loginClean()
+	performCleanup := true
+	defer func() {
+		if performCleanup {
+			s.loginClean()
+		}
+	}()
 
 	wasSignedOut := s.bridge.HasUser(s.auth.UserID)
 
@@ -426,10 +433,24 @@ func (s *Service) finishLogin() {
 	eventCh, done := s.bridge.GetEvents(events.UserLoggedIn{})
 	defer done()
 
-	userID, err := s.bridge.LoginUser(context.Background(), s.authClient, s.auth, s.password)
+	ctx := context.Background()
+	userID, err := s.bridge.LoginUser(ctx, s.authClient, s.auth, s.password)
 	if err != nil {
 		s.log.WithError(err).Errorf("Finish login failed")
-		_ = s.SendEvent(NewLoginError(LoginErrorType_TWO_PASSWORDS_ABORT, err.Error()))
+		s.twoPasswordAttemptCount++
+		errType := LoginErrorType_TWO_PASSWORDS_ABORT
+		if errors.Is(err, bridge.ErrFailedToUnlock) {
+			if s.twoPasswordAttemptCount < twoPasswordsMaxAttemptCount {
+				performCleanup = false
+				errType = LoginErrorType_TWO_PASSWORDS_ERROR
+			} else {
+				if deleteErr := s.authClient.AuthDelete(ctx); deleteErr != nil {
+					s.log.WithError(deleteErr).Error("Failed to delete auth")
+				}
+			}
+		}
+
+		_ = s.SendEvent(NewLoginError(errType, err.Error()))
 		return
 	}
 
diff --git a/internal/frontend/grpc/service_methods.go b/internal/frontend/grpc/service_methods.go
index 009856c6..0c703e7b 100644
--- a/internal/frontend/grpc/service_methods.go
+++ b/internal/frontend/grpc/service_methods.go
@@ -384,6 +384,7 @@ func (s *Service) Login(_ context.Context, login *LoginRequest) (*emptypb.Empty,
 	go func() {
 		defer async.HandlePanic(s.panicHandler)
 
+		s.twoPasswordAttemptCount = 0
 		password, err := base64Decode(login.Password)
 		if err != nil {
 			s.log.WithError(err).Error("Cannot decode password")