macaroons: add SetRootKey method

This adds a SetRootKey method which can be used to set the default
macaroon root key to a specified value.

macaroons/service.go

@@ -68,6 +68,10 @@ type ExtendedRootKeyStore interface {
 	// GenerateNewRootKey calls the underlying root key store's
 	// GenerateNewRootKey and returns the result.
 	GenerateNewRootKey() error
+
+	// SetRootKey calls the underlying root key store's SetRootKey and
+	// returns the result.
+	SetRootKey(rootKey []byte) error
 }
 
 // Service encapsulates bakery.Bakery and adds a Close() method that zeroes the
@@ -300,6 +304,16 @@ func (svc *Service) GenerateNewRootKey() error {
 	return nil
 }
 
+// SetRootKey calls the underlying root key store's SetRootKey and returns the
+// result.
+func (svc *Service) SetRootKey(rootKey []byte) error {
+	if boltRKS, ok := svc.rks.(ExtendedRootKeyStore); ok {
+		return boltRKS.SetRootKey(rootKey)
+	}
+
+	return nil
+}
+
 // ChangePassword calls the underlying root key store's ChangePassword and
 // returns the result.
 func (svc *Service) ChangePassword(oldPw, newPw []byte) error {

macaroons/store.go

@@ -332,6 +332,32 @@ func (r *RootKeyStorage) GenerateNewRootKey() error {
 	}, func() {})
 }
 
+// SetRootKey sets the default macaroon root key, replacing the previous root
+// key if it existed.
+func (r *RootKeyStorage) SetRootKey(rootKey []byte) error {
+	if r.encKey == nil {
+		return ErrStoreLocked
+	}
+	if len(rootKey) != RootKeyLen {
+		return fmt.Errorf("root key must be %v bytes",
+			RootKeyLen)
+	}
+
+	encryptedKey, err := r.encKey.Encrypt(rootKey)
+	if err != nil {
+		return err
+	}
+
+	return kvdb.Update(r.Backend, func(tx kvdb.RwTx) error {
+		bucket := tx.ReadWriteBucket(rootKeyBucketName)
+		if bucket == nil {
+			return ErrRootKeyBucketNotFound
+		}
+
+		return bucket.Put(DefaultRootKeyID, encryptedKey)
+	}, func() {})
+}
+
 // Close closes the underlying database and zeroes the encryption key stored
 // in memory.
 func (r *RootKeyStorage) Close() error {

macaroons/store_test.go

@@ -2,6 +2,7 @@ package macaroons_test
 
 import (
 	"context"
+	"crypto/rand"
 	"io/ioutil"
 	"os"
 	"path"
@@ -169,6 +170,42 @@ func TestStoreGenerateNewRootKey(t *testing.T) {
 	require.NotEqual(t, oldRootKey, newRootKey)
 }
 
+// TestStoreSetRootKey tests that a root key can be set to a specified value.
+func TestStoreSetRootKey(t *testing.T) {
+	_, cleanup, store := newTestStore(t)
+	defer cleanup()
+
+	// Create a new random key
+	rootKey := make([]byte, 32)
+	_, err := rand.Read(rootKey)
+	require.NoError(t, err)
+
+	// The store must be unlocked to set the root key.
+	err = store.SetRootKey(rootKey)
+	require.Equal(t, macaroons.ErrStoreLocked, err)
+
+	// Unlock the store and read the current key.
+	pw := []byte("weks")
+	err = store.CreateUnlock(&pw)
+	require.NoError(t, err)
+	oldRootKey, _, err := store.RootKey(defaultRootKeyIDContext)
+	require.NoError(t, err)
+
+	// Ensure the new key is different from the old key.
+	require.NotEqual(t, oldRootKey, rootKey)
+
+	// Replace the root key with the new key.
+	err = store.SetRootKey(rootKey)
+	require.NoError(t, err)
+
+	// Finally, read the root key from the DB and compare it to the one
+	// we created earlier. This makes sure that the encryption/
+	// decryption of the key in the DB worked as expected too.
+	newRootKey, _, err := store.RootKey(defaultRootKeyIDContext)
+	require.NoError(t, err)
+	require.Equal(t, rootKey, newRootKey)
+}
+
 // TestStoreChangePassword tests that the password for the store can be changed
 // without changing the root key.
 func TestStoreChangePassword(t *testing.T) {