From 6fdded20bc5dc6265fe0239ac773b0f6ffc37800 Mon Sep 17 00:00:00 2001
From: Kris Nuttycombe <kris@nutty.land>
Date: Wed, 3 Nov 2021 12:03:37 -0600
Subject: [PATCH] Restore legacy HDSeed encryption tests.

---
 src/gtest/test_keystore.cpp | 63 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 62 insertions(+), 1 deletion(-)

diff --git a/src/gtest/test_keystore.cpp b/src/gtest/test_keystore.cpp
index 8c8dff1d3..e3473806f 100644
--- a/src/gtest/test_keystore.cpp
+++ b/src/gtest/test_keystore.cpp
@@ -312,7 +312,7 @@ public:
     bool Unlock(const CKeyingMaterial& vMasterKeyIn) { return CCryptoKeyStore::Unlock(vMasterKeyIn); }
 };
 
-TEST(KeystoreTests, StoreAndRetrieveHDSeedInEncryptedStore) {
+TEST(KeystoreTests, StoreAndRetrieveMnemonicSeedInEncryptedStore) {
     TestCCryptoKeyStore keyStore;
     CKeyingMaterial vMasterKey(32, 0);
     GetRandBytes(vMasterKey.data(), 32);
@@ -378,6 +378,67 @@ TEST(KeystoreTests, StoreAndRetrieveHDSeedInEncryptedStore) {
     EXPECT_EQ(seed3, seedOut.value());
 }
 
+TEST(KeystoreTests, StoreAndRetrieveLegacyHDSeedInEncryptedStore) {
+    TestCCryptoKeyStore keyStore;
+    CKeyingMaterial vMasterKey(32, 0);
+    GetRandBytes(vMasterKey.data(), 32);
+
+    // 1) Test adding a seed to an unencrypted key store, then encrypting it
+    // We use a mnemonic seed, then disregard the mnemonic itself.
+    auto seed = MnemonicSeed::Random(SLIP44_TESTNET_TYPE);
+    auto seedOut = keyStore.GetLegacyHDSeed();
+    EXPECT_FALSE(seedOut.has_value());
+
+    ASSERT_TRUE(keyStore.SetLegacyHDSeed(seed));
+    seedOut = keyStore.GetLegacyHDSeed();
+    ASSERT_TRUE(seedOut.has_value());
+
+    ASSERT_TRUE(keyStore.EncryptKeys(vMasterKey));
+    seedOut = keyStore.GetLegacyHDSeed();
+    EXPECT_FALSE(seedOut.has_value());
+
+    // Unlocking with a random key should fail
+    CKeyingMaterial vRandomKey(32, 0);
+    GetRandBytes(vRandomKey.data(), 32);
+    EXPECT_FALSE(keyStore.Unlock(vRandomKey));
+
+    // Unlocking with a slightly-modified vMasterKey should fail
+    CKeyingMaterial vModifiedKey(vMasterKey);
+    vModifiedKey[0] += 1;
+    EXPECT_FALSE(keyStore.Unlock(vModifiedKey));
+
+    // Unlocking with vMasterKey should succeed
+    ASSERT_TRUE(keyStore.Unlock(vMasterKey));
+    seedOut = keyStore.GetLegacyHDSeed();
+    ASSERT_TRUE(seedOut.has_value());
+    EXPECT_EQ(seed, seedOut.value());
+
+    // 2) Test replacing the seed in an already-encrypted key store fails
+    auto seed2 = MnemonicSeed::Random(SLIP44_TESTNET_TYPE);
+    EXPECT_FALSE(keyStore.SetLegacyHDSeed(seed2));
+    seedOut = keyStore.GetLegacyHDSeed();
+    ASSERT_TRUE(seedOut.has_value());
+    EXPECT_EQ(seed, seedOut.value());
+
+    // 3) Test adding a new seed to an already-encrypted key store
+    TestCCryptoKeyStore keyStore2;
+
+    // Add a Sprout address so the wallet has something to test when decrypting
+    ASSERT_TRUE(keyStore2.AddSproutSpendingKey(libzcash::SproutSpendingKey::random()));
+
+    ASSERT_TRUE(keyStore2.EncryptKeys(vMasterKey));
+    ASSERT_TRUE(keyStore2.Unlock(vMasterKey));
+
+    seedOut = keyStore2.GetLegacyHDSeed();
+    EXPECT_FALSE(seedOut.has_value());
+
+    auto seed3 = MnemonicSeed::Random(SLIP44_TESTNET_TYPE);
+    ASSERT_TRUE(keyStore2.SetLegacyHDSeed(seed3));
+    seedOut = keyStore2.GetLegacyHDSeed();
+    ASSERT_TRUE(seedOut.has_value());
+    EXPECT_EQ(seed3, seedOut.value());
+}
+
 TEST(KeystoreTests, StoreAndRetrieveSpendingKeyInEncryptedStore) {
     TestCCryptoKeyStore keyStore;
     uint256 r {GetRandHash()};