From c4768ad8cf3363bd4acf4640d0738e18ad2b727c Mon Sep 17 00:00:00 2001
From: Caleb Doxsey <cdoxsey@pomerium.com>
Date: Fri, 26 Apr 2024 15:32:11 -0600
Subject: [PATCH] add runtime flag

---
 config/runtime_flags.go                    |  3 +++
 databroker/cache.go                        | 15 +++++++++++++++
 internal/identity/legacymanager/manager.go |  2 +-
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/config/runtime_flags.go b/config/runtime_flags.go
index c0c6627ae..957e048f6 100644
--- a/config/runtime_flags.go
+++ b/config/runtime_flags.go
@@ -8,6 +8,9 @@ var (
 
 	// RuntimeFlagMatchAnyIncomingPort enables ignoring the incoming port when matching routes
 	RuntimeFlagMatchAnyIncomingPort = runtimeFlag("match_any_incoming_port", true)
+
+	// RuntimeFlagLegacyIdentityManager enables the legacy identity manager
+	RuntimeFlagLegacyIdentityManager = runtimeFlag("legacy_identity_manager", false)
 )
 
 // RuntimeFlag is a runtime flag that can flip on/off certain features
diff --git a/databroker/cache.go b/databroker/cache.go
index 79ab014d7..f3559e835 100644
--- a/databroker/cache.go
+++ b/databroker/cache.go
@@ -18,6 +18,7 @@ import (
 	"github.com/pomerium/pomerium/internal/atomicutil"
 	"github.com/pomerium/pomerium/internal/events"
 	"github.com/pomerium/pomerium/internal/identity"
+	"github.com/pomerium/pomerium/internal/identity/legacymanager"
 	"github.com/pomerium/pomerium/internal/identity/manager"
 	"github.com/pomerium/pomerium/internal/log"
 	"github.com/pomerium/pomerium/internal/telemetry"
@@ -34,6 +35,7 @@ import (
 type DataBroker struct {
 	dataBrokerServer *dataBrokerServer
 	manager          *manager.Manager
+	legacyManager    *legacymanager.Manager
 	eventsMgr        *events.Manager
 
 	localListener       net.Listener
@@ -158,6 +160,12 @@ func (c *DataBroker) update(ctx context.Context, cfg *config.Config) error {
 	options := []manager.Option{
 		manager.WithDataBrokerClient(dataBrokerClient),
 		manager.WithEventManager(c.eventsMgr),
+		manager.WithEnabled(!cfg.Options.IsRuntimeFlagSet(config.RuntimeFlagLegacyIdentityManager)),
+	}
+	legacyOptions := []legacymanager.Option{
+		legacymanager.WithDataBrokerClient(dataBrokerClient),
+		legacymanager.WithEventManager(c.eventsMgr),
+		legacymanager.WithEnabled(cfg.Options.IsRuntimeFlagSet(config.RuntimeFlagLegacyIdentityManager)),
 	}
 
 	if cfg.Options.SupportsUserRefresh() {
@@ -166,6 +174,7 @@ func (c *DataBroker) update(ctx context.Context, cfg *config.Config) error {
 			log.Error(ctx).Err(err).Msg("databroker: failed to create authenticator")
 		} else {
 			options = append(options, manager.WithAuthenticator(authenticator))
+			legacyOptions = append(legacyOptions, legacymanager.WithAuthenticator(authenticator))
 		}
 	} else {
 		log.Info(ctx).Msg("databroker: disabling refresh of user sessions")
@@ -177,6 +186,12 @@ func (c *DataBroker) update(ctx context.Context, cfg *config.Config) error {
 		c.manager.UpdateConfig(options...)
 	}
 
+	if c.legacyManager == nil {
+		c.legacyManager = legacymanager.New(legacyOptions...)
+	} else {
+		c.legacyManager.UpdateConfig(legacyOptions...)
+	}
+
 	return nil
 }
 
diff --git a/internal/identity/legacymanager/manager.go b/internal/identity/legacymanager/manager.go
index 296f414ff..25965125f 100644
--- a/internal/identity/legacymanager/manager.go
+++ b/internal/identity/legacymanager/manager.go
@@ -88,7 +88,7 @@ func (mgr *Manager) UpdateConfig(options ...Option) {
 
 // RunEnabled runs the manager. This method blocks until an error occurs or the given context is canceled.
 func (mgr *Manager) RunEnabled(ctx context.Context) error {
-	leaser := databroker.NewLeaser("identity_manager", time.Second*30, mgr)
+	leaser := databroker.NewLeaser("legacy_identity_manager", time.Second*30, mgr)
 	return leaser.Run(ctx)
 }