From c0dabe28f85036b4d8c50cc58ad8086927a8c434 Mon Sep 17 00:00:00 2001
From: Casey Waldren <cwaldren@launchdarkly.com>
Date: Wed, 5 Jun 2024 11:22:20 -0700
Subject: [PATCH] temp working trying to get key roration working

---
 relay/filedata_actions.go       | 46 ++++++++++++++++++++++++++++++++-
 relay/filedata_actions_test.go  | 28 ++++++++++++++++++++
 relay/filedata_testdata_test.go | 22 ++++++++++++++++
 3 files changed, 95 insertions(+), 1 deletion(-)

diff --git a/relay/filedata_actions.go b/relay/filedata_actions.go
index 8e600321..9ebcca20 100644
--- a/relay/filedata_actions.go
+++ b/relay/filedata_actions.go
@@ -49,11 +49,23 @@ func (a *relayFileDataActions) AddEnvironment(ae filedata.ArchiveEnvironment) {
 		return config
 	}
 	envConfig := envfactory.NewEnvConfigFactoryForOfflineMode(a.r.config.OfflineMode).MakeEnvironmentConfig(ae.Params)
-	_, _, err := a.r.addEnvironment(ae.Params.Identifiers, envConfig, transformConfig)
+	env, _, err := a.r.addEnvironment(ae.Params.Identifiers, envConfig, transformConfig)
 	if err != nil {
 		a.r.loggers.Errorf(logMsgAutoConfEnvInitError, ae.Params.Identifiers.GetDisplayName(), err)
 		return
 	}
+
+	// Note: the following lines (until end marker) are
+	// copied from autoconfig_actions.go.
+	if ae.Params.ExpiringSDKKey != "" {
+		if foundEnvWithOldKey, _ := a.r.getEnvironment(ae.Params.ExpiringSDKKey); foundEnvWithOldKey == nil {
+			env.AddCredential(ae.Params.ExpiringSDKKey)
+			env.DeprecateCredential(ae.Params.ExpiringSDKKey)
+			a.r.addedEnvironmentCredential(env, ae.Params.ExpiringSDKKey)
+		}
+	}
+	// End.
+
 	select {
 	case updates := <-updatesCh:
 		if a.envUpdates == nil {
@@ -83,6 +95,38 @@ func (a *relayFileDataActions) UpdateEnvironment(ae filedata.ArchiveEnvironment)
 	env.SetTTL(ae.Params.TTL)
 	env.SetSecureMode(ae.Params.SecureMode)
 
+	// Note: the following lines (until end marker) are
+	// copied from autoconfig_actions.go.
+	var oldSDKKey config.SDKKey
+	var oldMobileKey config.MobileKey
+	for _, c := range env.GetCredentials() {
+		switch c := c.(type) {
+		case config.SDKKey:
+			oldSDKKey = c
+		case config.MobileKey:
+			oldMobileKey = c
+		}
+	}
+
+	if ae.Params.SDKKey != oldSDKKey {
+		env.AddCredential(ae.Params.SDKKey)
+		a.r.addedEnvironmentCredential(env, ae.Params.SDKKey)
+		if ae.Params.ExpiringSDKKey == oldSDKKey {
+			env.DeprecateCredential(oldSDKKey)
+		} else {
+			a.r.removingEnvironmentCredential(oldSDKKey)
+			env.RemoveCredential(oldSDKKey)
+		}
+	}
+
+	if ae.Params.MobileKey != oldMobileKey {
+		env.AddCredential(ae.Params.MobileKey)
+		a.r.addedEnvironmentCredential(env, ae.Params.MobileKey)
+		a.r.removingEnvironmentCredential(oldMobileKey)
+		env.RemoveCredential(oldMobileKey)
+	}
+	// End.
+
 	// SDKData will be non-nil only if the flag/segment data for the environment has actually changed.
 	if ae.SDKData != nil {
 		updates.Init(ae.SDKData)
diff --git a/relay/filedata_actions_test.go b/relay/filedata_actions_test.go
index 221239ab..cdd66f63 100644
--- a/relay/filedata_actions_test.go
+++ b/relay/filedata_actions_test.go
@@ -205,3 +205,31 @@ func TestOfflineModeEventsAreAcceptedAndDiscardedIfSendEventsIsTrue(t *testing.T
 		})
 	})
 }
+
+func TestOfflineModeDeprecateSDKKeyIsAccepted(t *testing.T) {
+	offlineModeTest(t, config.Config{}, func(p offlineModeTestParams) {
+		p.updateHandler.AddEnvironment(testFileDataExpiringSDKKey)
+
+		client1 := p.awaitClient()
+		assert.Equal(t, testFileDataExpiringSDKKey.Params.SDKKey, client1.Key)
+
+		env := p.awaitEnvironment(testFileDataEnv1.Params.EnvID)
+
+		var sdkKey config.SDKKey
+		for _, c := range env.GetCredentials() {
+			if c, ok := c.(config.SDKKey); ok {
+				sdkKey = c
+				break
+			}
+		}
+		var expiringSDKKey config.SDKKey
+		for _, c := range env.GetDeprecatedCredentials() {
+			if c, ok := c.(config.SDKKey); ok {
+				expiringSDKKey = c
+				break
+			}
+		}
+		assert.Equal(t, testFileDataExpiringSDKKey.Params.SDKKey, sdkKey)
+		assert.Equal(t, testFileDataExpiringSDKKey.Params.ExpiringSDKKey, expiringSDKKey)
+	})
+}
diff --git a/relay/filedata_testdata_test.go b/relay/filedata_testdata_test.go
index 3f6b0c79..76e15b52 100644
--- a/relay/filedata_testdata_test.go
+++ b/relay/filedata_testdata_test.go
@@ -58,3 +58,25 @@ var testFileDataEnv2 = filedata.ArchiveEnvironment{
 		},
 	},
 }
+
+var testFileDataExpiringSDKKey = filedata.ArchiveEnvironment{
+	Params: envfactory.EnvironmentParams{
+		EnvID:          config.EnvironmentID("env1"),
+		SDKKey:         config.SDKKey("sdkkey2"),
+		ExpiringSDKKey: config.SDKKey("sdkkey1"),
+		Identifiers: relayenv.EnvIdentifiers{
+			ProjName: "Project",
+			ProjKey:  "project",
+			EnvName:  "Env1",
+			EnvKey:   "env1",
+		},
+	},
+	SDKData: []ldstoretypes.Collection{
+		{
+			Kind: ldstoreimpl.Features(),
+			Items: []ldstoretypes.KeyedItemDescriptor{
+				{Key: testFileDataFlag1.Key, Item: sharedtest.FlagDesc(testFileDataFlag1)},
+			},
+		},
+	},
+}