feat(BRIDGE-119): added support for Feature Flags

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)
+}

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 {

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 <https://www.gnu.org/licenses/>.
+
+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)
+		})
+	})
+}

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(

internal/service/types.go

@@ -19,4 +19,5 @@ package service
 
 type Locator interface {
 	ProvideSettingsPath() (string, error)
+	ProvideUnleashCachePath() (string, error)
 }

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 <https://www.gnu.org/licenses/>.
+
+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
+}