From ae7dab03bae850be2b3a7a801b471ed9e75e8648 Mon Sep 17 00:00:00 2001
From: Valentin David <valentin.david@canonical.com>
Date: Wed, 20 Nov 2024 17:21:22 +0100
Subject: [PATCH] crypt.go: allow adding names to legacy keyslots

When reprovisioning with a newer snapd with old disks, we need to be
able to convert old keyslots to new ones with names. Otherwise we
cannot remove after reprovisioning is done.
---
 crypt.go      | 57 +++++++++++++++++++++++++++++++
 crypt_test.go | 95 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 152 insertions(+)

diff --git a/crypt.go b/crypt.go
index dbe93974..3be6d76b 100755
--- a/crypt.go
+++ b/crypt.go
@@ -934,3 +934,60 @@ func RenameLUKS2ContainerKey(devicePath, oldName, newName string) error {
 
 	return nil
 }
+
+// KeyslotAlreadyHasAName may be returned bv
+// NameLegacyLUKS2ContainerKey when trying to create a token for a
+// keyslot that already used by a token.
+var KeyslotAlreadyHasAName = errors.New("keyslot already has a name")
+
+// NameLegacyLUKS2ContainerKey will add a token for a recovery key for
+// a specified keyslot. That keyslot must not be in use in any
+// existing token. If the keyslot does not exist, this function will
+// do nothing. This function is intended to be used to name keys on
+// an old container.
+func NameLegacyLUKS2ContainerKey(devicePath string, keyslot int, newName string) error {
+	view, err := newLUKSView(devicePath, luks2.LockModeBlocking)
+	if err != nil {
+		return xerrors.Errorf("cannot obtain LUKS header view: %w", err)
+	}
+
+	_, _, inUse := view.TokenByName(newName)
+	if inUse {
+		return errors.New("the new name is already in use")
+	}
+
+	for _, name := range view.TokenNames() {
+		token, _, inUse := view.TokenByName(name)
+		if inUse {
+			for _, usedKeyslot := range token.Keyslots() {
+				if usedKeyslot == keyslot {
+					return KeyslotAlreadyHasAName
+				}
+			}
+		}
+	}
+
+	keyslotExists := false
+	for _, usedSlot := range view.UsedKeyslots() {
+		if usedSlot == keyslot {
+			keyslotExists = true
+			break
+		}
+	}
+	if !keyslotExists {
+		return nil
+	}
+
+	token := &luksview.RecoveryToken{
+		TokenBase: luksview.TokenBase{
+			TokenName:    newName,
+			TokenKeyslot: keyslot,
+		},
+	}
+
+	if err := luks2ImportToken(devicePath, token, nil); err != nil {
+		return xerrors.Errorf("cannot import token: %w", err)
+	}
+
+	return nil
+}
diff --git a/crypt_test.go b/crypt_test.go
index 5a2bc5af..f97ccab5 100644
--- a/crypt_test.go
+++ b/crypt_test.go
@@ -4003,3 +4003,98 @@ func (s *cryptSuite) TestActivateVolumeWithLegacyPathsError(c *C) {
 
 	s.checkKeyDataKeysInKeyring(c, "", "/dev/some/path", unlockKey, primaryKey)
 }
+
+func (s *cryptSuite) TestNameLegacyLUKS2ContainerKey(c *C) {
+	firstKey := s.newPrimaryKey(c, 32)
+	secondKey := s.newPrimaryKey(c, 32)
+
+	m := &mockLUKS2Container{
+		tokens: map[int]luks2.Token{},
+		keyslots: map[int][]byte{
+			0: firstKey,
+			1: secondKey,
+		},
+	}
+
+	s.luks2.devices["/dev/foo1"] = m
+
+	// Nothing should happen
+	err := NameLegacyLUKS2ContainerKey("/dev/foo1", 2, "some-name")
+	c.Check(err, IsNil)
+	c.Check(m.tokens, HasLen, 0)
+
+	err = NameLegacyLUKS2ContainerKey("/dev/foo1", 0, "some-name")
+	c.Check(err, IsNil)
+
+	token, hasToken := m.tokens[0]
+	c.Assert(hasToken, Equals, true)
+	c.Check(token, DeepEquals, &luksview.RecoveryToken{
+		TokenBase: luksview.TokenBase{
+			TokenKeyslot: 0,
+			TokenName:    "some-name",
+		},
+	})
+
+	err = NameLegacyLUKS2ContainerKey("/dev/foo1", 1, "some-other-name")
+	c.Check(err, IsNil)
+
+	token, hasToken = m.tokens[1]
+	c.Assert(hasToken, Equals, true)
+	c.Check(token, DeepEquals, &luksview.RecoveryToken{
+		TokenBase: luksview.TokenBase{
+			TokenKeyslot: 1,
+			TokenName:    "some-other-name",
+		},
+	})
+}
+
+func (s *cryptSuite) TestNameLegacyLUKS2ContainerKeyNameExists(c *C) {
+	firstKey := s.newPrimaryKey(c, 32)
+	secondKey := s.newPrimaryKey(c, 32)
+
+	m := &mockLUKS2Container{
+		tokens: map[int]luks2.Token{
+			1: &luksview.KeyDataToken{
+				TokenBase: luksview.TokenBase{
+					TokenKeyslot: 1,
+					TokenName:    "already",
+				},
+			},
+		},
+		keyslots: map[int][]byte{
+			0: firstKey,
+			1: secondKey,
+		},
+	}
+
+	s.luks2.devices["/dev/foo1"] = m
+
+	err := NameLegacyLUKS2ContainerKey("/dev/foo1", 1, "some-name")
+	c.Check(err, ErrorMatches, `keyslot already has a name`)
+	c.Check(errors.Is(err, KeyslotAlreadyHasAName), Equals, true)
+}
+
+func (s *cryptSuite) TestNameLegacyLUKS2ContainerKeyNameAlreadyUsed(c *C) {
+	firstKey := s.newPrimaryKey(c, 32)
+	secondKey := s.newPrimaryKey(c, 32)
+
+	m := &mockLUKS2Container{
+		tokens: map[int]luks2.Token{
+			1: &luksview.KeyDataToken{
+				TokenBase: luksview.TokenBase{
+					TokenKeyslot: 1,
+					TokenName:    "already-used",
+				},
+			},
+		},
+		keyslots: map[int][]byte{
+			0: firstKey,
+			1: secondKey,
+		},
+	}
+
+	s.luks2.devices["/dev/foo1"] = m
+
+	err := NameLegacyLUKS2ContainerKey("/dev/foo1", 0, "already-used")
+	c.Check(err, ErrorMatches, `the new name is already in use`)
+}