From 51f0082c53db6c55668958b511d7c746695cd215 Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Wed, 17 May 2023 11:51:56 +0200
Subject: [PATCH] macaroons: Let GenerateNewRootKey be applied to all root keys

With this commit, GenerateNewRootKey will regenerate the Default root
key and will then also check if any other root keys exist and regenerate
those as well.
---
 macaroons/store.go      | 26 +++++++++++++++++++++++++-
 macaroons/store_test.go | 12 ++++--------
 2 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/macaroons/store.go b/macaroons/store.go
index 09da98885..7b40dff30 100644
--- a/macaroons/store.go
+++ b/macaroons/store.go
@@ -325,10 +325,34 @@ func (r *RootKeyStorage) GenerateNewRootKey() error {
 		if bucket == nil {
 			return ErrRootKeyBucketNotFound
 		}
+
+		// The default root key should be created even if it does not
+		// yet exist, so we do this separately from the rest of the
+		// root keys.
 		_, err := generateAndStoreNewRootKey(
 			bucket, DefaultRootKeyID, r.encKey,
 		)
-		return err
+		if err != nil {
+			return err
+		}
+
+		// Now iterate over all the other root keys that may exist
+		// and re-generate each of them.
+		return bucket.ForEach(func(k, v []byte) error {
+			if bytes.Equal(k, encryptionKeyID) {
+				return nil
+			}
+
+			if bytes.Equal(k, DefaultRootKeyID) {
+				return nil
+			}
+
+			_, err := generateAndStoreNewRootKey(
+				bucket, k, r.encKey,
+			)
+
+			return err
+		})
 	}, func() {})
 }
 
diff --git a/macaroons/store_test.go b/macaroons/store_test.go
index 83692b627..175635836 100644
--- a/macaroons/store_test.go
+++ b/macaroons/store_test.go
@@ -135,10 +135,8 @@ func TestStore(t *testing.T) {
 	require.Equal(t, rootID, id)
 }
 
-// TestStoreGenerateNewRootKey tests that a root key can be replaced with a new
-// one in the store without changing the password. Also demonstrate that at the
-// moment, only the default root key will be replaced. This is a bug that will
-// be fixed in an upcoming commit.
+// TestStoreGenerateNewRootKey tests that root keys can be replaced with new
+// ones in the store without changing the password.
 func TestStoreGenerateNewRootKey(t *testing.T) {
 	_, store := newTestStore(t)
 
@@ -159,22 +157,20 @@ func TestStoreGenerateNewRootKey(t *testing.T) {
 	oldRootKey2, _, err := store.RootKey(nonDefaultRootKeyIDContext)
 	require.NoError(t, err)
 
-	// Attempt to replace the root keys with new random keys.
+	// Replace the root keys with new random keys.
 	err = store.GenerateNewRootKey()
 	require.NoError(t, err)
 
 	// Finally, read both root keys from the DB and compare them to the ones
 	// we got returned earlier. This makes sure that the encryption/
 	// decryption of the key in the DB worked as expected too.
-	// Currently, this is only successful for the default root key and not
-	// for non-default key. This will be fixed in an upcoming commit.
 	newRootKey1, _, err := store.RootKey(defaultRootKeyIDContext)
 	require.NoError(t, err)
 	require.NotEqual(t, oldRootKey1, newRootKey1)
 
 	newRootKey2, _, err := store.RootKey(nonDefaultRootKeyIDContext)
 	require.NoError(t, err)
-	require.Equal(t, oldRootKey2, newRootKey2)
+	require.NotEqual(t, oldRootKey2, newRootKey2)
 }
 
 // TestStoreSetRootKey tests that a root key can be set to a specified value.