GODT-2022: Fix change between address modes

internal/bridge/imap.go

@@ -113,10 +113,16 @@ func (bridge *Bridge) addIMAPUser(ctx context.Context, user *user.User) error {
 
 // removeIMAPUser disconnects the given user from gluon, optionally also removing its files.
 func (bridge *Bridge) removeIMAPUser(ctx context.Context, user *user.User, withFiles bool) error {
-	for _, gluonID := range user.GetGluonIDs() {
+	for addrID, gluonID := range user.GetGluonIDs() {
 		if err := bridge.imapServer.RemoveUser(ctx, gluonID, withFiles); err != nil {
 			return fmt.Errorf("failed to remove IMAP user: %w", err)
 		}
+
+		if withFiles {
+			if err := user.RemoveGluonID(addrID, gluonID); err != nil {
+				return fmt.Errorf("failed to remove IMAP user ID: %w", err)
+			}
+		}
 	}
 
 	return nil

internal/bridge/user.go

@@ -258,10 +258,14 @@ func (bridge *Bridge) SetAddressMode(ctx context.Context, userID string, mode va
 			return fmt.Errorf("address mode is already %q", mode)
 		}
 
-		for _, gluonID := range user.GetGluonIDs() {
+		for addrID, gluonID := range user.GetGluonIDs() {
 			if err := bridge.imapServer.RemoveUser(ctx, gluonID, true); err != nil {
 				return fmt.Errorf("failed to remove user from IMAP server: %w", err)
 			}
+
+			if err := user.RemoveGluonID(addrID, gluonID); err != nil {
+				return fmt.Errorf("failed to remove gluon ID from user: %w", err)
+			}
 		}
 
 		if err := user.SetAddressMode(ctx, mode); err != nil {

internal/bridge/user_events.go

@@ -93,6 +93,10 @@ func (bridge *Bridge) handleUserAddressDeleted(ctx context.Context, user *user.U
 		if err := bridge.imapServer.RemoveUser(ctx, gluonID, true); err != nil {
 			return fmt.Errorf("failed to remove user from IMAP server: %w", err)
 		}
+
+		if err := user.RemoveGluonID(event.AddressID, gluonID); err != nil {
+			return fmt.Errorf("failed to remove gluon ID for address: %w", err)
+		}
 	}
 
 	return nil

internal/bridge/user_test.go

@@ -462,15 +462,26 @@ func TestBridge_AddressMode(t *testing.T) {
 			// The user is in combined mode by default.
 			require.Equal(t, vault.CombinedMode, info.AddressMode)
 
-			// Put the user in split mode.
-			require.NoError(t, bridge.SetAddressMode(ctx, userID, vault.SplitMode))
+			// Repeatedly switch between address modes.
+			for i := 1; i <= 10; i++ {
+				var target vault.AddressMode
 
-			// Get the user's info.
-			info, err = bridge.GetUserInfo(userID)
-			require.NoError(t, err)
+				if i%2 == 0 {
+					target = vault.CombinedMode
+				} else {
+					target = vault.SplitMode
+				}
 
-			// The user is in split mode.
-			require.Equal(t, vault.SplitMode, info.AddressMode)
+				// Put the user in the target mode.
+				require.NoError(t, bridge.SetAddressMode(ctx, userID, target))
+
+				// Get the user's info.
+				info, err = bridge.GetUserInfo(userID)
+				require.NoError(t, err)
+
+				// The user is in the target mode.
+				require.Equal(t, target, info.AddressMode)
+			}
 		})
 	})
 }

internal/frontend/cli/accounts.go

@@ -265,7 +265,16 @@ func (f *frontendCLI) changeMode(c *ishell.Context) {
 		targetMode = vault.CombinedMode
 	}
 
-	if !f.yesNoQuestion("Are you sure you want to change the mode for account " + bold(user.Username) + " to " + bold(targetMode)) {
+	var targetModeName string
+
+	switch targetMode {
+	case vault.CombinedMode:
+		targetModeName = "combined"
+	case vault.SplitMode:
+		targetModeName = "split"
+	}
+
+	if !f.yesNoQuestion("Are you sure you want to change the mode for account " + bold(user.Username) + " to " + bold(targetModeName)) {
 		return
 	}
 
@@ -273,7 +282,7 @@ func (f *frontendCLI) changeMode(c *ishell.Context) {
 		f.printAndLogError("Cannot switch address mode:", err)
 	}
 
-	f.Printf("Address mode for account %s changed to %s\n", user.Username, targetMode)
+	f.Printf("Address mode for account %s changed to %s\n", user.Username, targetModeName)
 }
 
 func (f *frontendCLI) configureAppleMail(c *ishell.Context) {

internal/frontend/grpc/service.go

@@ -259,6 +259,9 @@ func (s *Service) watchEvents() { //nolint:funlen
 		case events.UserDeleted:
 			_ = s.SendEvent(NewUserChangedEvent(event.UserID))
 
+		case events.AddressModeChanged:
+			_ = s.SendEvent(NewUserChangedEvent(event.UserID))
+
 		case events.UserDeauth:
 			if user, err := s.bridge.GetUserInfo(event.UserID); err != nil {
 				s.log.WithError(err).Error("Failed to get user info")

internal/frontend/grpc/service_user.go

@@ -21,7 +21,6 @@ import (
 	"context"
 
 	"github.com/ProtonMail/proton-bridge/v2/internal/vault"
-	"github.com/sirupsen/logrus"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/emptypb"
@@ -76,15 +75,21 @@ func (s *Service) SetUserSplitMode(ctx context.Context, splitMode *UserSplitMode
 
 		var targetMode vault.AddressMode
 
-		if splitMode.Active && user.AddressMode == vault.CombinedMode {
+		if splitMode.Active {
 			targetMode = vault.SplitMode
-		} else if !splitMode.Active && user.AddressMode == vault.SplitMode {
+		} else {
 			targetMode = vault.CombinedMode
 		}
 
-		if err := s.bridge.SetAddressMode(context.Background(), user.UserID, targetMode); err != nil {
-			logrus.WithError(err).Error("Failed to set address mode")
+		if user.AddressMode == targetMode {
+			return
 		}
+
+		if err := s.bridge.SetAddressMode(context.Background(), user.UserID, targetMode); err != nil {
+			s.log.WithError(err).Error("Failed to set address mode")
+		}
+
+		s.log.WithField("userID", user.UserID).WithField("mode", targetMode).Info("Address mode changed")
 	}()
 
 	return &emptypb.Empty{}, nil
@@ -101,7 +106,7 @@ func (s *Service) LogoutUser(ctx context.Context, userID *wrapperspb.StringValue
 		defer s.panicHandler.HandlePanic()
 
 		if err := s.bridge.LogoutUser(context.Background(), userID.Value); err != nil {
-			logrus.WithError(err).Error("Failed to log user out")
+			s.log.WithError(err).Error("Failed to log user out")
 		}
 	}()
 

internal/user/sync.go

@@ -39,7 +39,7 @@ import (
 
 const (
 	maxUpdateSize = 1 << 25
-	maxBatchSize  = 1 << 8
+	maxBatchSize  = 1 << 6
 )
 
 // doSync begins syncing the users data.

internal/user/user.go

@@ -344,6 +344,11 @@ func (user *User) SetGluonID(addrID, gluonID string) error {
 	return user.vault.SetGluonID(addrID, gluonID)
 }
 
+// RemoveGluonID removes the gluon ID for the given address.
+func (user *User) RemoveGluonID(addrID, gluonID string) error {
+	return user.vault.RemoveGluonID(addrID, gluonID)
+}
+
 // GluonKey returns the user's gluon key from the vault.
 func (user *User) GluonKey() []byte {
 	return user.vault.GluonKey()

internal/vault/user.go

@@ -18,6 +18,8 @@
 package vault
 
 import (
+	"fmt"
+
 	"github.com/ProtonMail/gluon/imap"
 )
 
@@ -49,6 +51,22 @@ func (user *User) SetGluonID(addrID, gluonID string) error {
 	})
 }
 
+func (user *User) RemoveGluonID(addrID, gluonID string) error {
+	var err error
+
+	if modErr := user.vault.modUser(user.userID, func(data *UserData) {
+		if data.GluonIDs[addrID] != gluonID {
+			err = fmt.Errorf("gluon ID mismatch: %s != %s", data.GluonIDs[addrID], gluonID)
+		} else {
+			delete(data.GluonIDs, addrID)
+		}
+	}); modErr != nil {
+		return modErr
+	}
+
+	return err
+}
+
 func (user *User) GetUIDValidity(addrID string) (imap.UID, bool) {
 	validity, ok := user.vault.getUser(user.userID).UIDValidity[addrID]
 	if !ok {