From e290cd308bca365cef06f763044f12e9ef360392 Mon Sep 17 00:00:00 2001 From: Atanas Janeshliev Date: Wed, 7 Aug 2024 17:04:54 +0200 Subject: [PATCH] feat(BRIDGE-119): added support for Feature Flags --- go.mod | 2 +- go.sum | 6 + internal/bridge/bridge.go | 16 +++ internal/bridge/types.go | 1 + internal/bridge/unleash_test.go | 90 ++++++++++++ internal/locations/locations.go | 12 ++ internal/service/types.go | 1 + internal/unleash/service.go | 234 ++++++++++++++++++++++++++++++++ tests/ctx_bridge_test.go | 5 + tests/ctx_test.go | 3 + tests/environment_test.go | 8 ++ 11 files changed, 377 insertions(+), 1 deletion(-) create mode 100644 internal/bridge/unleash_test.go create mode 100644 internal/unleash/service.go diff --git a/go.mod b/go.mod index ea623bdf..05281c68 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.20240514133734-79cdd0fec41c github.com/ProtonMail/go-autostart v0.0.0-20210130080809-00ed301c8e9a - github.com/ProtonMail/go-proton-api v0.4.1-0.20240612082117-0f92424eed80 + github.com/ProtonMail/go-proton-api v0.4.1-0.20240821081056-dd607af0f917 github.com/ProtonMail/gopenpgp/v2 v2.7.4-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 0d1a5fc3..c4dd64ba 100644 --- a/go.sum +++ b/go.sum @@ -40,6 +40,12 @@ github.com/ProtonMail/go-mime v0.0.0-20230322103455-7d82a3887f2f h1:tCbYj7/299ek 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.20240612082117-0f92424eed80 h1:cP4+6RFn9vVgYnoDwxBU4EtIAZA+eM4rzOaSZNqZ1xg= github.com/ProtonMail/go-proton-api v0.4.1-0.20240612082117-0f92424eed80/go.mod h1:3A0cpdo0BIenIPjTG6u8EbzJ8uuJy7rVvM/NaynjCKA= +github.com/ProtonMail/go-proton-api v0.4.1-0.20240808145610-88df257767f6 h1:nERxOYS4ndSgWEr834YYkb1j0bZK/dJAmhoyYB1MtNY= +github.com/ProtonMail/go-proton-api v0.4.1-0.20240808145610-88df257767f6/go.mod h1:3A0cpdo0BIenIPjTG6u8EbzJ8uuJy7rVvM/NaynjCKA= +github.com/ProtonMail/go-proton-api v0.4.1-0.20240819131705-149e50199c5b h1:zifGh4LS5HwQIaVCccSe5/oJGTOjFeVObMRl3QJoJ3k= +github.com/ProtonMail/go-proton-api v0.4.1-0.20240819131705-149e50199c5b/go.mod h1:3A0cpdo0BIenIPjTG6u8EbzJ8uuJy7rVvM/NaynjCKA= +github.com/ProtonMail/go-proton-api v0.4.1-0.20240821081056-dd607af0f917 h1:Ma6PfXFDuw7rYYq28FXNW6ubhYquRUmBuLyZrjJWHUE= +github.com/ProtonMail/go-proton-api v0.4.1-0.20240821081056-dd607af0f917/go.mod h1:3A0cpdo0BIenIPjTG6u8EbzJ8uuJy7rVvM/NaynjCKA= 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 3fd6c250..e18270c3 100644 --- a/internal/bridge/bridge.go +++ b/internal/bridge/bridge.go @@ -47,6 +47,7 @@ import ( "github.com/ProtonMail/proton-bridge/v3/internal/services/imapsmtpserver" "github.com/ProtonMail/proton-bridge/v3/internal/services/syncservice" "github.com/ProtonMail/proton-bridge/v3/internal/telemetry" + "github.com/ProtonMail/proton-bridge/v3/internal/unleash" "github.com/ProtonMail/proton-bridge/v3/internal/user" "github.com/ProtonMail/proton-bridge/v3/internal/vault" "github.com/ProtonMail/proton-bridge/v3/pkg/keychain" @@ -136,6 +137,8 @@ type Bridge struct { serverManager *imapsmtpserver.Service syncService *syncservice.Service + // unleashService is responsible for polling the feature flags and caching + unleashService *unleash.Service } var logPkg = logrus.WithField("pkg", "bridge") //nolint:gochecknoglobals @@ -253,6 +256,8 @@ func newBridge( return nil, fmt.Errorf("failed to create focus service: %w", err) } + unleashService := unleash.NewBridgeService(ctx, api, locator, panicHandler) + bridge := &Bridge{ vault: vault, @@ -293,6 +298,8 @@ func newBridge( tasks: tasks, syncService: syncservice.NewService(reporter, panicHandler), + + unleashService: unleashService, } bridge.serverManager = imapsmtpserver.NewService(context.Background(), @@ -320,6 +327,8 @@ func newBridge( bridge.syncService.Run() + bridge.unleashService.Run() + return bridge, nil } @@ -470,6 +479,9 @@ func (bridge *Bridge) Close(ctx context.Context) { // Close the focus service. bridge.focusService.Close() + // Close the unleash service. + bridge.unleashService.Close() + // Close the watchers. bridge.watchersLock.Lock() defer bridge.watchersLock.Unlock() @@ -674,3 +686,7 @@ func GetUpdatedCachePath(gluonDBPath, gluonCachePath string) string { return strings.Replace(gluonCachePath, "/Users/"+cacheUsername+"/", "/Users/"+dbUsername+"/", 1) } + +func (bridge *Bridge) GetFeatureFlagValue(key string) bool { + return bridge.unleashService.GetFlagValue(key) +} diff --git a/internal/bridge/types.go b/internal/bridge/types.go index ed232803..d2a8b8eb 100644 --- a/internal/bridge/types.go +++ b/internal/bridge/types.go @@ -33,6 +33,7 @@ type Locator interface { GetDependencyLicensesLink() string Clear(...string) error ProvideIMAPSyncConfigPath() (string, error) + ProvideUnleashCachePath() (string, error) } type ProxyController interface { diff --git a/internal/bridge/unleash_test.go b/internal/bridge/unleash_test.go new file mode 100644 index 00000000..8bb9afa3 --- /dev/null +++ b/internal/bridge/unleash_test.go @@ -0,0 +1,90 @@ +// Copyright (c) 2024 Proton AG +// +// This file is part of Proton Mail Bridge. +// +// Proton Mail Bridge is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// Proton Mail Bridge is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with Proton Mail Bridge. If not, see . + +package bridge_test + +import ( + "context" + "testing" + "time" + + "github.com/ProtonMail/go-proton-api" + "github.com/ProtonMail/go-proton-api/server" + "github.com/ProtonMail/proton-bridge/v3/internal/bridge" + "github.com/ProtonMail/proton-bridge/v3/internal/unleash" + "github.com/stretchr/testify/require" +) + +func Test_UnleashService(t *testing.T) { + withEnv(t, func(ctx context.Context, s *server.Server, netCtl *proton.NetCtl, locator bridge.Locator, storeKey []byte) { + unleash.ModifyPollPeriodAndJitter(500*time.Millisecond, 0) + + withBridge(ctx, t, s.GetHostURL(), netCtl, locator, storeKey, func(b *bridge.Bridge, _ *bridge.Mocks) { + // Initial startup assumes there is no cached feature flags. + require.Equal(t, b.GetFeatureFlagValue("test-1"), false) + require.Equal(t, b.GetFeatureFlagValue("test-2"), false) + require.Equal(t, b.GetFeatureFlagValue("test-3"), false) + + s.PushFeatureFlag("test-1") + s.PushFeatureFlag("test-2") + + // Wait for poll. + time.Sleep(time.Millisecond * 700) + require.Equal(t, b.GetFeatureFlagValue("test-1"), true) + require.Equal(t, b.GetFeatureFlagValue("test-2"), true) + require.Equal(t, b.GetFeatureFlagValue("test-3"), false) + + s.PushFeatureFlag("test-3") + time.Sleep(time.Millisecond * 700) // Wait for poll again + require.Equal(t, b.GetFeatureFlagValue("test-1"), true) + require.Equal(t, b.GetFeatureFlagValue("test-2"), true) + require.Equal(t, b.GetFeatureFlagValue("test-3"), true) + }) + + // Wait for Bridge to close. + time.Sleep(time.Millisecond * 500) + + // Second instance should have a feature flag cache file available. Therefore, all of the flags should evaluate to true on startup. + withBridge(ctx, t, s.GetHostURL(), netCtl, locator, storeKey, func(b *bridge.Bridge, _ *bridge.Mocks) { + require.Equal(t, b.GetFeatureFlagValue("test-1"), true) + require.Equal(t, b.GetFeatureFlagValue("test-2"), true) + require.Equal(t, b.GetFeatureFlagValue("test-3"), true) + + s.DeleteFeatureFlags() + + require.Equal(t, b.GetFeatureFlagValue("test-1"), true) + require.Equal(t, b.GetFeatureFlagValue("test-2"), true) + require.Equal(t, b.GetFeatureFlagValue("test-3"), true) + + time.Sleep(time.Millisecond * 700) + + require.Equal(t, b.GetFeatureFlagValue("test-1"), false) + require.Equal(t, b.GetFeatureFlagValue("test-2"), false) + require.Equal(t, b.GetFeatureFlagValue("test-3"), false) + + s.PushFeatureFlag("test-3") + require.Equal(t, b.GetFeatureFlagValue("test-1"), false) + require.Equal(t, b.GetFeatureFlagValue("test-2"), false) + require.Equal(t, b.GetFeatureFlagValue("test-3"), false) + + time.Sleep(time.Millisecond * 700) + require.Equal(t, b.GetFeatureFlagValue("test-1"), false) + require.Equal(t, b.GetFeatureFlagValue("test-2"), false) + require.Equal(t, b.GetFeatureFlagValue("test-3"), true) + }) + }) +} diff --git a/internal/locations/locations.go b/internal/locations/locations.go index eeaa245d..1c9367cf 100644 --- a/internal/locations/locations.go +++ b/internal/locations/locations.go @@ -206,6 +206,16 @@ func (l *Locations) ProvideIMAPSyncConfigPath() (string, error) { return l.getIMAPSyncConfigPath(), nil } +// ProvideUnleashCachePath returns a location for the unleash cache data (e.g. ~/.cache/protonmail/bridge-v3). +// It creates it if it doesn't already exist. +func (l *Locations) ProvideUnleashCachePath() (string, error) { + if err := os.MkdirAll(l.getUnleashCachePath(), 0o700); err != nil { + return "", err + } + + return l.getUnleashCachePath(), nil +} + func (l *Locations) getGluonCachePath() string { return filepath.Join(l.userData, "gluon") } @@ -242,6 +252,8 @@ func (l *Locations) getStatsPath() string { return filepath.Join(l.userData, "stats") } +func (l *Locations) getUnleashCachePath() string { return filepath.Join(l.userCache, "unleash_cache") } + // Clear removes everything except the lock and update files. func (l *Locations) Clear(except ...string) error { return files.Remove( diff --git a/internal/service/types.go b/internal/service/types.go index 898e6498..4e895f7e 100644 --- a/internal/service/types.go +++ b/internal/service/types.go @@ -19,4 +19,5 @@ package service type Locator interface { ProvideSettingsPath() (string, error) + ProvideUnleashCachePath() (string, error) } diff --git a/internal/unleash/service.go b/internal/unleash/service.go new file mode 100644 index 00000000..f856b3e7 --- /dev/null +++ b/internal/unleash/service.go @@ -0,0 +1,234 @@ +// Copyright (c) 2024 Proton AG +// +// This file is part of Proton Mail Bridge. +// +// Proton Mail Bridge is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// Proton Mail Bridge is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with Proton Mail Bridge. If not, see . + +package unleash + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "sync" + "time" + + "github.com/ProtonMail/gluon/async" + "github.com/ProtonMail/go-proton-api" + "github.com/ProtonMail/proton-bridge/v3/internal/service" + "github.com/sirupsen/logrus" +) + +var pollPeriod = 10 * time.Minute //nolint:gochecknoglobals +var pollJitter = 2 * time.Minute //nolint:gochecknoglobals + +const filename = "unleash_flags" + +type requestFeaturesFn func(ctx context.Context) (proton.FeatureFlagResult, error) + +type Service struct { + panicHandler async.PanicHandler + timer *proton.Ticker + + ctx context.Context + cancel context.CancelFunc + + log *logrus.Entry + + ffStore map[string]bool + ffStoreLock sync.Mutex + + cacheFilepath string + cacheFileLock sync.Mutex + + channel chan map[string]bool + + getFeaturesFn func(ctx context.Context) (proton.FeatureFlagResult, error) +} + +func NewBridgeService(ctx context.Context, api *proton.Manager, locator service.Locator, panicHandler async.PanicHandler) *Service { + log := logrus.WithField("service", "unleash") + cacheDir, err := locator.ProvideUnleashCachePath() + if err != nil { + log.Warn("Could not find or create unleash cache directory") + } + cachePath := filepath.Clean(filepath.Join(cacheDir, filename)) + + return newService(ctx, func(ctx context.Context) (proton.FeatureFlagResult, error) { + return api.GetFeatures(ctx) + }, log, cachePath, panicHandler) +} + +func newService(ctx context.Context, fn requestFeaturesFn, log *logrus.Entry, cachePath string, panicHandler async.PanicHandler) *Service { + ctx, cancel := context.WithCancel(ctx) + + unleashService := &Service{ + panicHandler: panicHandler, + timer: proton.NewTicker(pollPeriod, pollJitter, panicHandler), + + ctx: ctx, + cancel: cancel, + + log: log, + + ffStore: make(map[string]bool), + cacheFilepath: cachePath, + + channel: make(chan map[string]bool), + + getFeaturesFn: fn, + } + + unleashService.readCacheFile() + return unleashService +} + +func readResponseData(data proton.FeatureFlagResult) map[string]bool { + featureData := make(map[string]bool) + for _, el := range data.Toggles { + featureData[el.Name] = el.Enabled + } + + return featureData +} + +func (s *Service) readCacheFile() { + defer s.cacheFileLock.Unlock() + s.cacheFileLock.Lock() + + file, err := os.Open(s.cacheFilepath) + if err != nil { + s.log.WithError(err).Error("Unable to open cache file") + return + } + + defer func(file *os.File) { + err := file.Close() + if err != nil { + s.log.WithError(err).Error("Unable to close cache file after read") + } + }(file) + + s.ffStoreLock.Lock() + defer s.ffStoreLock.Unlock() + if err = json.NewDecoder(file).Decode(&s.ffStore); err != nil { + s.log.WithError(err).Error("Unable to decode cache file") + } +} + +func (s *Service) writeCacheFile() { + defer s.cacheFileLock.Unlock() + s.cacheFileLock.Lock() + + file, err := os.Create(s.cacheFilepath) + if err != nil { + s.log.WithError(err).Error("Unable to create cache file") + return + } + + defer func(file *os.File) { + err := file.Close() + if err != nil { + s.log.WithError(err).Error("Unable to close cache file after write") + } + }(file) + + s.ffStoreLock.Lock() + defer s.ffStoreLock.Unlock() + if err = json.NewEncoder(file).Encode(s.ffStore); err != nil { + s.log.WithError(err).Error("Unable to encode data to cache file") + } +} + +func (s *Service) Run() { + s.log.Info("Starting service") + + go func() { + s.runFlagPoll() + }() + + go func() { + s.runReceiver() + }() +} + +func (s *Service) runFlagPoll() { + defer async.HandlePanic(s.panicHandler) + defer s.timer.Stop() + s.log.Info("Starting poll service") + + data, err := s.getFeaturesFn(s.ctx) + if err != nil { + s.log.WithError(err).Error("Failed to get flags from server") + } else { + s.channel <- readResponseData(data) + } + + for { + select { + case <-s.ctx.Done(): + return + case <-s.timer.C: + s.log.Info("Polling flag service") + data, err := s.getFeaturesFn(s.ctx) + if err != nil { + s.log.WithError(err).Error("Failed to get feature flags from server") + continue + } + s.channel <- readResponseData(data) + } + } +} + +func (s *Service) runReceiver() { + defer async.HandlePanic(s.panicHandler) + s.log.Info("Starting receiver service") + + for { + select { + case <-s.ctx.Done(): + return + case res := <-s.channel: + s.ffStoreLock.Lock() + s.ffStore = res + s.ffStoreLock.Unlock() + s.writeCacheFile() + } + } +} + +func (s *Service) GetFlagValue(key string) bool { + defer s.ffStoreLock.Unlock() + s.ffStoreLock.Lock() + + val, ok := s.ffStore[key] + if !ok { + return false + } + + return val +} + +func (s *Service) Close() { + s.log.Info("Closing service") + s.cancel() + close(s.channel) +} + +// ModifyPollPeriodAndJitter is only used for testing. +func ModifyPollPeriodAndJitter(pollInterval, jitterInterval time.Duration) { + pollPeriod = pollInterval + pollJitter = jitterInterval +} diff --git a/tests/ctx_bridge_test.go b/tests/ctx_bridge_test.go index ad17a00b..e171fbd7 100644 --- a/tests/ctx_bridge_test.go +++ b/tests/ctx_bridge_test.go @@ -149,6 +149,11 @@ func (t *testCtx) initBridge() (<-chan events.Event, error) { } rt := t.netCtl.NewRoundTripper(&tls.Config{InsecureSkipVerify: true}) + + // We store the round tripper in the testing context so we can cancel the connection + // when we're turning it down/up + t.rt = &rt + if isBlack() { // GODT-1602 make sure we don't time out test server t, ok := rt.(*http.Transport) diff --git a/tests/ctx_test.go b/tests/ctx_test.go index 07e03bd0..aab58206 100644 --- a/tests/ctx_test.go +++ b/tests/ctx_test.go @@ -21,6 +21,7 @@ import ( "context" "fmt" "net" + "net/http" "net/smtp" "net/url" "regexp" @@ -164,6 +165,8 @@ type testCtx struct { imapServerStarted bool smtpServerStarted bool + + rt *http.RoundTripper } type imapClient struct { diff --git a/tests/environment_test.go b/tests/environment_test.go index 17184a5b..1b887494 100644 --- a/tests/environment_test.go +++ b/tests/environment_test.go @@ -59,11 +59,19 @@ func (s *scenario) itFailsWithError(wantErr string) error { func (s *scenario) internetIsTurnedOff() error { s.t.netCtl.SetCanDial(false) + t, ok := (*s.t.rt).(*http.Transport) + if ok { + t.CloseIdleConnections() + } return nil } func (s *scenario) internetIsTurnedOn() error { s.t.netCtl.SetCanDial(true) + t, ok := (*s.t.rt).(*http.Transport) + if ok { + t.CloseIdleConnections() + } return nil }