feat(BRIDGE-278): Rollout Feature Flag stickiness support; a new UUID 'sticky' value has been added to the vault"

internal/bridge/bridge.go

@@ -57,6 +57,7 @@ import (
 	"github.com/bradenaw/juniper/xslices"
 	"github.com/elastic/go-sysinfo/types"
 	"github.com/go-resty/resty/v2"
+	uuid "github.com/google/uuid"
 	"github.com/sirupsen/logrus"
 )
 
@@ -271,7 +272,7 @@ func newBridge(
 		return nil, fmt.Errorf("failed to create focus service: %w", err)
 	}
 
-	unleashService := unleash.NewBridgeService(ctx, api, locator, panicHandler)
+	unleashService := unleash.NewBridgeService(ctx, api, locator, panicHandler, vault.GetFeatureFlagStickyKey())
 
 	observabilityService := observability.NewService(ctx, panicHandler)
 
@@ -785,3 +786,7 @@ func (bridge *Bridge) SetHostVersionGetterTest(fn func(host types.Host) string)
 func (bridge *Bridge) SetRolloutPercentageTest(rollout float64) error {
 	return bridge.vault.SetUpdateRollout(rollout)
 }
+
+func (bridge *Bridge) GetFeatureFlagStickyKey() uuid.UUID {
+	return bridge.vault.GetFeatureFlagStickyKey()
+}

internal/bridge/bridge_test.go

@@ -56,6 +56,7 @@ import (
 	imapid "github.com/emersion/go-imap-id"
 	"github.com/emersion/go-sasl"
 	"github.com/emersion/go-smtp"
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
 )
@@ -778,6 +779,30 @@ func TestBridge_ChangeCacheDirectory(t *testing.T) {
 	})
 }
 
+func TestBridge_FeatureFlagStickyKey_Persistence(t *testing.T) {
+	var uuidOne uuid.UUID
+	var uuidTwo uuid.UUID
+
+	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, _ *bridge.Mocks) {
+			uuidOne = b.GetFeatureFlagStickyKey()
+		})
+		withBridge(ctx, t, s.GetHostURL(), netCtl, locator, vaultKey, func(b *bridge.Bridge, _ *bridge.Mocks) {
+			require.Equal(t, uuidOne, b.GetFeatureFlagStickyKey())
+		})
+	})
+
+	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, _ *bridge.Mocks) {
+			uuidTwo = b.GetFeatureFlagStickyKey()
+			require.NotEqual(t, uuidOne, uuidTwo)
+		})
+		withBridge(ctx, t, s.GetHostURL(), netCtl, locator, vaultKey, func(b *bridge.Bridge, _ *bridge.Mocks) {
+			require.Equal(t, uuidTwo, b.GetFeatureFlagStickyKey())
+		})
+	})
+}
+
 func TestBridge_ChangeAddressOrder(t *testing.T) {
 	withEnv(t, func(ctx context.Context, s *server.Server, netCtl *proton.NetCtl, locator bridge.Locator, vaultKey []byte) {
 		// Create a user.

internal/unleash/service.go

@@ -28,6 +28,7 @@ import (
 	"github.com/ProtonMail/gluon/async"
 	"github.com/ProtonMail/go-proton-api"
 	"github.com/ProtonMail/proton-bridge/v3/internal/service"
+	"github.com/google/uuid"
 	"github.com/sirupsen/logrus"
 )
 
@@ -83,7 +84,13 @@ type Service struct {
 	getFeaturesFn func(ctx context.Context) (proton.FeatureFlagResult, error)
 }
 
-func NewBridgeService(ctx context.Context, api *proton.Manager, locator service.Locator, panicHandler async.PanicHandler) *Service {
+func NewBridgeService(
+	ctx context.Context,
+	api *proton.Manager,
+	locator service.Locator,
+	panicHandler async.PanicHandler,
+	featureFlagUUID uuid.UUID,
+) *Service {
 	log := logrus.WithField("service", "unleash")
 	cacheDir, err := locator.ProvideUnleashCachePath()
 	if err != nil {
@@ -92,7 +99,7 @@ func NewBridgeService(ctx context.Context, api *proton.Manager, locator service.
 	cachePath := filepath.Clean(filepath.Join(cacheDir, filename))
 
 	return newService(ctx, func(ctx context.Context) (proton.FeatureFlagResult, error) {
-		return api.GetFeatures(ctx)
+		return api.GetFeatures(ctx, featureFlagUUID)
 	}, log, cachePath, panicHandler)
 }
 

internal/vault/types_data.go

@@ -17,17 +17,22 @@
 
 package vault
 
+import "github.com/google/uuid"
+
 type Data struct {
 	Settings Settings
 	Users    []UserData
 	Cookies  []byte
 	Certs    Certs
 	Migrated bool
+	// FeatureFlagStickyKey a utility value for ensuring rollout feature flags "stick" to a particular Bridge client.
+	FeatureFlagStickyKey uuid.UUID
 }
 
 func newDefaultData(gluonDir string) Data {
 	return Data{
-		Settings: newDefaultSettings(gluonDir),
-		Certs:    newDefaultCerts(),
+		Settings:             newDefaultSettings(gluonDir),
+		Certs:                newDefaultCerts(),
+		FeatureFlagStickyKey: uuid.New(),
 	}
 }

internal/vault/vault.go

@@ -32,6 +32,7 @@ import (
 	"github.com/ProtonMail/gluon/async"
 	"github.com/bradenaw/juniper/parallel"
 	"github.com/bradenaw/juniper/xslices"
+	"github.com/google/uuid"
 	"github.com/sirupsen/logrus"
 )
 
@@ -76,6 +77,10 @@ func New(vaultDir, gluonCacheDir string, key []byte, panicHandler async.PanicHan
 		return nil, corrupt, err
 	}
 
+	if err := vault.setFeatureFlagStickyKeyIfEmpty(); err != nil {
+		return vault, corrupt, err
+	}
+
 	vault.panicHandler = panicHandler
 
 	return vault, corrupt, nil
@@ -91,6 +96,22 @@ func (vault *Vault) GetUserIDs() []string {
 	})
 }
 
+// GetFeatureFlagStickyKey - the sticky key is  a utility value for ensuring rollout feature flags "stick" to a particular Bridge client.
+func (vault *Vault) GetFeatureFlagStickyKey() uuid.UUID {
+	return vault.getSafe().FeatureFlagStickyKey
+}
+
+// setFeatureFlagStickyKeyIfEmpty - checks if the sticky key is nil in the vault and if so generates a new one.
+func (vault *Vault) setFeatureFlagStickyKeyIfEmpty() error {
+	if vault.getSafe().FeatureFlagStickyKey != uuid.Nil {
+		return nil
+	}
+
+	return vault.modSafe(func(data *Data) {
+		data.FeatureFlagStickyKey = uuid.New()
+	})
+}
+
 func (vault *Vault) getUsers() ([]*User, error) {
 	vault.lock.Lock()
 	defer vault.lock.Unlock()