From be9e03d917aa6368547f4affb110d5d77ed5642f Mon Sep 17 00:00:00 2001 From: Atanas Janeshliev Date: Thu, 26 Jun 2025 17:23:48 +0200 Subject: [PATCH] feat(BRIDGE-278): Rollout Feature Flag stickiness support; a new UUID 'sticky' value has been added to the vault" --- go.mod | 2 +- go.sum | 4 ++-- internal/bridge/bridge.go | 7 ++++++- internal/bridge/bridge_test.go | 25 +++++++++++++++++++++++++ internal/unleash/service.go | 11 +++++++++-- internal/vault/types_data.go | 9 +++++++-- internal/vault/vault.go | 21 +++++++++++++++++++++ 7 files changed, 71 insertions(+), 8 deletions(-) diff --git a/go.mod b/go.mod index ff6159c5..87852eb1 100644 --- a/go.mod +++ b/go.mod @@ -9,7 +9,7 @@ require ( github.com/Masterminds/semver/v3 v3.2.0 github.com/ProtonMail/gluon v0.17.1-0.20250627102828-b014b7cc8132 github.com/ProtonMail/go-autostart v0.0.0-20210130080809-00ed301c8e9a - github.com/ProtonMail/go-proton-api v0.4.1-0.20250627100628-e6d3af81dfdf + github.com/ProtonMail/go-proton-api v0.4.1-0.20250627135952-bf973947255c github.com/ProtonMail/gopenpgp/v2 v2.9.0-proton github.com/PuerkitoBio/goquery v1.8.1 github.com/abiosoft/ishell v2.0.0+incompatible diff --git a/go.sum b/go.sum index b40546bc..f8bb080e 100644 --- a/go.sum +++ b/go.sum @@ -45,8 +45,8 @@ github.com/ProtonMail/go-message v0.13.1-0.20240919135104-3bc88e6a9423 h1:p8nBDx github.com/ProtonMail/go-message v0.13.1-0.20240919135104-3bc88e6a9423/go.mod h1:NBAn21zgCJ/52WLDyed18YvYFm5tEoeDauubFqLokM4= github.com/ProtonMail/go-mime v0.0.0-20230322103455-7d82a3887f2f h1:tCbYj7/299ekTTXpdwKYF8eBlsYsDVoggDAuAjoK66k= github.com/ProtonMail/go-mime v0.0.0-20230322103455-7d82a3887f2f/go.mod h1:gcr0kNtGBqin9zDW9GOHcVntrwnjrK+qdJ06mWYBybw= -github.com/ProtonMail/go-proton-api v0.4.1-0.20250627100628-e6d3af81dfdf h1:lGyK9G1hyAzUBhkU1wmc6oYPnUTwrRBnVb6YA14Z/bk= -github.com/ProtonMail/go-proton-api v0.4.1-0.20250627100628-e6d3af81dfdf/go.mod h1:9t9+oQfH+6ssa7O2nLv34Uyjv8UmqTPGbVNcFToewck= +github.com/ProtonMail/go-proton-api v0.4.1-0.20250627135952-bf973947255c h1:FhfHIrGgehnTV/T2NkyVauKcJ3NzPq1uLcU/0eK661A= +github.com/ProtonMail/go-proton-api v0.4.1-0.20250627135952-bf973947255c/go.mod h1:9t9+oQfH+6ssa7O2nLv34Uyjv8UmqTPGbVNcFToewck= github.com/ProtonMail/go-smtp v0.0.0-20231109081432-2b3d50599865 h1:EP1gnxLL5Z7xBSymE9nSTM27nRYINuvssAtDmG0suD8= github.com/ProtonMail/go-smtp v0.0.0-20231109081432-2b3d50599865/go.mod h1:qm27SGYgoIPRot6ubfQ/GpiPy/g3PaZAVRxiO/sDUgQ= github.com/ProtonMail/go-srp v0.0.7 h1:Sos3Qk+th4tQR64vsxGIxYpN3rdnG9Wf9K4ZloC1JrI= diff --git a/internal/bridge/bridge.go b/internal/bridge/bridge.go index fbb677f1..1fecc60d 100644 --- a/internal/bridge/bridge.go +++ b/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() +} diff --git a/internal/bridge/bridge_test.go b/internal/bridge/bridge_test.go index a72a8b49..5fd9e78a 100644 --- a/internal/bridge/bridge_test.go +++ b/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. diff --git a/internal/unleash/service.go b/internal/unleash/service.go index 974f1a64..51b196a9 100644 --- a/internal/unleash/service.go +++ b/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) } diff --git a/internal/vault/types_data.go b/internal/vault/types_data.go index f444d288..e5c84e08 100644 --- a/internal/vault/types_data.go +++ b/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(), } } diff --git a/internal/vault/vault.go b/internal/vault/vault.go index 24b83ffe..43652ae4 100644 --- a/internal/vault/vault.go +++ b/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()