chore(GODT-2551): Store and Recover Last User Agent from Vault

internal/bridge/bridge.go

@@ -237,6 +237,8 @@ func newBridge(
 		return nil, fmt.Errorf("failed to save last version indicator: %w", err)
 	}
 
+	identifier.SetClientString(vault.GetLastUserAgent())
+
 	imapServer, err := newIMAPServer(
 		gluonCacheDir,
 		gluonDataDir,

internal/bridge/bridge_test.go

@@ -171,6 +171,41 @@ func TestBridge_UserAgent(t *testing.T) {
 	})
 }
 
+func TestBridge_UserAgent_Persistence(t *testing.T) {
+	withEnv(t, func(ctx context.Context, s *server.Server, netCtl *proton.NetCtl, locator bridge.Locator, vaultKey []byte) {
+		withBridge(ctx, t, s.GetHostURL(), netCtl, locator, vaultKey, func(b *bridge.Bridge, mocks *bridge.Mocks) {
+			currentUserAgent := b.GetCurrentUserAgent()
+			require.Contains(t, currentUserAgent, vault.DefaultUserAgent)
+
+			imapClient, err := client.Dial(fmt.Sprintf("%v:%v", constants.Host, b.GetIMAPPort()))
+			require.NoError(t, err)
+			defer func() { _ = imapClient.Logout() }()
+
+			idClient := imapid.NewClient(imapClient)
+
+			// Set IMAP ID before Login to have the value capture in the Login API Call.
+			_, err = idClient.ID(imapid.ID{
+				imapid.FieldName:    "MyFancyClient",
+				imapid.FieldVersion: "0.1.2",
+			})
+
+			require.NoError(t, err)
+
+			// Login the user.
+			_, err = b.LoginFull(context.Background(), username, password, nil, nil)
+			require.NoError(t, err)
+
+			// Assert that the user agent then contains the platform.
+			require.Contains(t, b.GetCurrentUserAgent(), "MyFancyClient/0.1.2")
+		})
+
+		withBridge(ctx, t, s.GetHostURL(), netCtl, locator, vaultKey, func(bridge *bridge.Bridge, mocks *bridge.Mocks) {
+			currentUserAgent := bridge.GetCurrentUserAgent()
+			require.Contains(t, currentUserAgent, "MyFancyClient/0.1.2")
+		})
+	})
+}
+
 func TestBridge_UserAgentFromIMAPID(t *testing.T) {
 	withEnv(t, func(ctx context.Context, s *server.Server, netCtl *proton.NetCtl, locator bridge.Locator, vaultKey []byte) {
 		var (

internal/bridge/identifier.go

@@ -17,6 +17,8 @@
 
 package bridge
 
+import "github.com/sirupsen/logrus"
+
 func (bridge *Bridge) GetCurrentUserAgent() string {
 	return bridge.identifier.GetUserAgent()
 }
@@ -24,3 +26,17 @@ func (bridge *Bridge) GetCurrentUserAgent() string {
 func (bridge *Bridge) SetCurrentPlatform(platform string) {
 	bridge.identifier.SetPlatform(platform)
 }
+
+func (bridge *Bridge) setUserAgent(name, version string) {
+	currentUserAgent := bridge.identifier.GetClientString()
+
+	bridge.identifier.SetClient(name, version)
+
+	newUserAgent := bridge.identifier.GetClientString()
+
+	if currentUserAgent != newUserAgent {
+		if err := bridge.vault.SetLastUserAgent(newUserAgent); err != nil {
+			logrus.WithError(err).Error("Failed to write new user agent to vault")
+		}
+	}
+}

internal/bridge/imap.go

@@ -41,11 +41,6 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-const (
-	defaultClientName    = "UnknownClient"
-	defaultClientVersion = "0.0.1"
-)
-
 func (bridge *Bridge) serveIMAP() error {
 	port, err := func() (int, error) {
 		if bridge.imapServer == nil {
@@ -249,11 +244,6 @@ func (bridge *Bridge) handleIMAPEvent(event imapEvents.Event) {
 			}).Info("Received mailbox message count")
 		}
 
-	case imapEvents.SessionAdded:
-		if !bridge.identifier.HasClient() {
-			bridge.identifier.SetClient(defaultClientName, defaultClientVersion)
-		}
-
 	case imapEvents.IMAPID:
 		logrus.WithFields(logrus.Fields{
 			"sessionID": event.SessionID,
@@ -262,7 +252,7 @@ func (bridge *Bridge) handleIMAPEvent(event imapEvents.Event) {
 		}).Info("Received IMAP ID")
 
 		if event.IMAPID.Name != "" && event.IMAPID.Version != "" {
-			bridge.identifier.SetClient(event.IMAPID.Name, event.IMAPID.Version)
+			bridge.setUserAgent(event.IMAPID.Name, event.IMAPID.Version)
 		}
 
 	case imapEvents.LoginFailed:

internal/bridge/types.go

@@ -38,6 +38,8 @@ type Identifier interface {
 	HasClient() bool
 	SetClient(name, version string)
 	SetPlatform(platform string)
+	SetClientString(client string)
+	GetClientString() string
 }
 
 type ProxyController interface {

internal/bridge/user.go

@@ -550,7 +550,7 @@ func (bridge *Bridge) addUserWithVault(
 	// As such, if we find this ID in the context, we should use it to update our user agent.
 	client.AddPreRequestHook(func(_ *resty.Client, r *resty.Request) error {
 		if imapID, ok := imap.GetIMAPIDFromContext(r.Context()); ok {
-			bridge.identifier.SetClient(imapID.Name, imapID.Version)
+			bridge.setUserAgent(imapID.Name, imapID.Version)
 		}
 
 		return nil

internal/useragent/useragent.go

@@ -44,6 +44,20 @@ func (ua *UserAgent) SetClient(name, version string) {
 	ua.client = fmt.Sprintf("%v/%v", name, regexp.MustCompile(`(.*) \((.*)\)`).ReplaceAllString(version, "$1-$2"))
 }
 
+func (ua *UserAgent) SetClientString(client string) {
+	ua.lock.Lock()
+	defer ua.lock.Unlock()
+
+	ua.client = client
+}
+
+func (ua *UserAgent) GetClientString() string {
+	ua.lock.RLock()
+	defer ua.lock.RUnlock()
+
+	return ua.client
+}
+
 func (ua *UserAgent) HasClient() bool {
 	ua.lock.RLock()
 	defer ua.lock.RUnlock()

internal/vault/settings.go

@@ -225,3 +225,22 @@ func (vault *Vault) SetMaxSyncMemory(maxMemory uint64) error {
 		data.Settings.MaxSyncMemory = maxMemory
 	})
 }
+
+// GetLastUserAgent returns the last user agent recorded by bridge.
+func (vault *Vault) GetLastUserAgent() string {
+	v := vault.get().Settings.LastUserAgent
+
+	// Handle case where there may be no value.
+	if len(v) == 0 {
+		v = DefaultUserAgent
+	}
+
+	return v
+}
+
+// SetLastUserAgent store the last user agent recorded by bridge.
+func (vault *Vault) SetLastUserAgent(userAgent string) error {
+	return vault.mod(func(data *Data) {
+		data.Settings.LastUserAgent = userAgent
+	})
+}

internal/vault/settings_test.go

@@ -216,3 +216,11 @@ func TestVault_Settings_MaxSyncMemory(t *testing.T) {
 	// Check the default first start value.
 	require.Equal(t, vault.DefaultMaxSyncMemory, s.GetMaxSyncMemory())
 }
+
+func TestVault_Settings_LastUserAgent(t *testing.T) {
+	// create a new test vault.
+	s := newVault(t)
+
+	// Check the default first start value.
+	require.Equal(t, vault.DefaultUserAgent, s.GetLastUserAgent())
+}

internal/vault/types_settings.go

@@ -46,12 +46,15 @@ type Settings struct {
 
 	MaxSyncMemory uint64
 
+	LastUserAgent string
+
 	// **WARNING**: These entry can't be removed until they vault has proper migration support.
 	SyncWorkers int
 	SyncAttPool int
 }
 
 const DefaultMaxSyncMemory = 2 * 1024 * uint64(1024*1024)
+const DefaultUserAgent = "UnknownClient/0.0.1"
 
 func GetDefaultSyncWorkerCount() int {
 	const minSyncWorkers = 16
@@ -91,5 +94,7 @@ func newDefaultSettings(gluonDir string) Settings {
 		MaxSyncMemory: DefaultMaxSyncMemory,
 		SyncWorkers:   syncWorkers,
 		SyncAttPool:   syncWorkers,
+
+		LastUserAgent: DefaultUserAgent,
 	}
 }