keystore: port bip39 unlock/retain to Rust

diff --git a/src/keystore.c b/src/keystore.c
index 1357f0d..ef6d802 100644
--- a/src/keystore.c
+++ b/src/keystore.c
@@ -41,16 +41,6 @@ static size_t _retained_seed_encrypted_len = 0;
 // plaintext.
 static uint8_t _retained_seed_hash[32] = {0};
 
-// Change this ONLY via keystore_unlock_bip39_finalize().
-static bool _is_unlocked_bip39 = false;
-// Stores a random key after bip39-unlock which, after stretching, is used to encrypt the retained
-// bip39 seed.
-static uint8_t _unstretched_retained_bip39_seed_encryption_key[32] = {0};
-// Must be defined if _is_unlocked is true. ONLY ACCESS THIS WITH keystore_copy_bip39_seed().
-// Stores the encrypted BIP-39 seed after bip39-unlock.
-static uint8_t _retained_bip39_seed_encrypted[64 + 64] = {0};
-static size_t _retained_bip39_seed_encrypted_len = 0;
-
 // Unlocking the keystore take longer than the 500ms watchdog we have setup. Reset the watchdog
 // counter to (~7s) to avoid incorrectly assuming we lost communication with the app.
 #define LONG_TIMEOUT (-70)
@@ -96,48 +86,6 @@ bool keystore_copy_seed(uint8_t* seed_out, size_t* length_out)
     return true;
 }
 
-bool keystore_copy_bip39_seed(uint8_t* bip39_seed_out)
-{
-    if (!_is_unlocked_bip39) {
-        return false;
-    }
-
-    uint8_t retained_bip39_seed_encryption_key[32] = {0};
-    UTIL_CLEANUP_32(retained_bip39_seed_encryption_key);
-    if (!rust_keystore_stretch_retained_seed_encryption_key(
-            rust_util_bytes(
-                _unstretched_retained_bip39_seed_encryption_key,
-                sizeof(_unstretched_retained_bip39_seed_encryption_key)),
-            "keystore_retained_bip39_seed_access_in",
-            "keystore_retained_bip39_seed_access_out",
-            rust_util_bytes_mut(
-                retained_bip39_seed_encryption_key, sizeof(retained_bip39_seed_encryption_key)))) {
-        return false;
-    }
-    size_t len = _retained_bip39_seed_encrypted_len - 48;
-    bool password_correct = cipher_aes_hmac_decrypt(
-        _retained_bip39_seed_encrypted,
-        _retained_bip39_seed_encrypted_len,
-        bip39_seed_out,
-        &len,
-        retained_bip39_seed_encryption_key);
-    if (!password_correct) {
-        // Should never happen.
-        return false;
-    }
-    if (len != 64) {
-        // Should never happen.
-        return false;
-    }
-    // sanity check
-    uint8_t zero[64] = {0};
-    util_zero(zero, 64);
-    if (MEMEQ(bip39_seed_out, zero, 64)) {
-        return false;
-    }
-    return true;
-}
-
 /**
  * Retrieves the encrypted seed and attempts to decrypt it using the password.
  *
@@ -238,44 +186,6 @@ USE_RESULT static keystore_error_t _retain_seed(const uint8_t* seed, size_t seed
     return _hash_seed(seed, seed_len, _retained_seed_hash);
 }
 
-USE_RESULT static bool _retain_bip39_seed(const uint8_t* bip39_seed)
-{
-#ifdef TESTING
-    const uint8_t test_unstretched_retained_bip39_seed_encryption_key[32] =
-        "\x9b\x44\xc7\x04\x88\x93\xfa\xaf\x6e\x2d\x76\x25\xd1\x3d\x8f\x1c\xab\x07\x65\xfd\x61\xf1"
-        "\x59\xd9\x71\x3e\x08\x15\x5d\x06\x71\x7c";
-    memcpy(
-        _unstretched_retained_bip39_seed_encryption_key,
-        test_unstretched_retained_bip39_seed_encryption_key,
-        32);
-#else
-    random_32_bytes(_unstretched_retained_bip39_seed_encryption_key);
-#endif
-    uint8_t retained_bip39_seed_encryption_key[32] = {0};
-    UTIL_CLEANUP_32(retained_bip39_seed_encryption_key);
-    if (!rust_keystore_stretch_retained_seed_encryption_key(
-            rust_util_bytes(
-                _unstretched_retained_bip39_seed_encryption_key,
-                sizeof(_unstretched_retained_bip39_seed_encryption_key)),
-            "keystore_retained_bip39_seed_access_in",
-            "keystore_retained_bip39_seed_access_out",
-            rust_util_bytes_mut(
-                retained_bip39_seed_encryption_key, sizeof(retained_bip39_seed_encryption_key)))) {
-        return false;
-    }
-    size_t len = sizeof(_retained_bip39_seed_encrypted);
-    if (!cipher_aes_hmac_encrypt(
-            bip39_seed,
-            64,
-            _retained_bip39_seed_encrypted,
-            &len,
-            retained_bip39_seed_encryption_key)) {
-        return false;
-    }
-    _retained_bip39_seed_encrypted_len = len;
-    return true;
-}
-
 static void _delete_retained_seeds(void)
 {
     util_zero(
@@ -284,12 +194,6 @@ static void _delete_retained_seeds(void)
     util_zero(_retained_seed_encrypted, sizeof(_retained_seed_encrypted));
     _retained_seed_encrypted_len = 0;
     util_zero(_retained_seed_hash, sizeof(_retained_seed_hash));
-
-    util_zero(
-        _unstretched_retained_bip39_seed_encryption_key,
-        sizeof(_unstretched_retained_seed_encryption_key));
-    util_zero(_retained_bip39_seed_encrypted, sizeof(_retained_bip39_seed_encrypted));
-    _retained_bip39_seed_encrypted_len = 0;
 }
 
 keystore_error_t keystore_encrypt_and_store_seed(
@@ -441,25 +345,15 @@ bool keystore_unlock_bip39_check(const uint8_t* seed, size_t seed_length)
     return true;
 }
 
-bool keystore_unlock_bip39_finalize(const uint8_t* bip39_seed)
-{
-    if (!_retain_bip39_seed(bip39_seed)) {
-        return false;
-    }
-    _is_unlocked_bip39 = true;
-    return true;
-}
-
 void keystore_lock(void)
 {
     _is_unlocked_device = false;
-    _is_unlocked_bip39 = false;
     _delete_retained_seeds();
 }
 
 bool keystore_is_locked(void)
 {
-    bool unlocked = _is_unlocked_device && _is_unlocked_bip39;
+    bool unlocked = _is_unlocked_device && rust_keystore_is_unlocked_bip39();
     return !unlocked;
 }
 
@@ -515,7 +409,6 @@ void keystore_mock_unlocked(const uint8_t* seed, size_t seed_len)
             Abort("couldn't retain seed");
         }
     }
-    _is_unlocked_bip39 = false;
 }
 
 const uint8_t* keystore_test_get_retained_seed_encrypted(size_t* len_out)
@@ -523,10 +416,4 @@ const uint8_t* keystore_test_get_retained_seed_encrypted(size_t* len_out)
     *len_out = _retained_seed_encrypted_len;
     return _retained_seed_encrypted;
 }
-
-const uint8_t* keystore_test_get_retained_bip39_seed_encrypted(size_t* len_out)
-{
-    *len_out = _retained_bip39_seed_encrypted_len;
-    return _retained_bip39_seed_encrypted;
-}
 #endif
diff --git a/src/keystore.h b/src/keystore.h
index a00b14c..d8e45ed 100644
--- a/src/keystore.h
+++ b/src/keystore.h
@@ -54,15 +54,6 @@ typedef enum {
  */
 USE_RESULT bool keystore_copy_seed(uint8_t* seed_out, size_t* length_out);
 
-/**
- * Copies the retained bip39 seed into the given buffer. The caller must
- * zero the seed once it is no longer needed.
- * @param[out] bip39_seed_out The seed bytes copied from the retained bip39 seed.
- * The buffer must be 64 bytes long.
- * @return true if the bip39 seed is available.
- */
-USE_RESULT bool keystore_copy_bip39_seed(uint8_t* bip32_seed_out);
-
 /**
  * Restores a seed. This also unlocks the keystore with this seed.
  * @param[in] seed The seed that is to be restored.
@@ -109,12 +100,6 @@ USE_RESULT keystore_error_t keystore_unlock(
  */
 USE_RESULT bool keystore_unlock_bip39_check(const uint8_t* seed, size_t seed_length);
 
-/**
- * Retains the given bip39 seed and marks the keystore as unlocked.
- * @param[in] bip39_seed 64 byte bip39 seed.
- */
-USE_RESULT bool keystore_unlock_bip39_finalize(const uint8_t* bip39_seed);
-
 /**
  * Locks the keystore (resets to state before `keystore_unlock()`).
  */
@@ -191,7 +176,6 @@ USE_RESULT bool keystore_secp256k1_sign(
 void keystore_mock_unlocked(const uint8_t* seed, size_t seed_len);
 
 const uint8_t* keystore_test_get_retained_seed_encrypted(size_t* len_out);
-const uint8_t* keystore_test_get_retained_bip39_seed_encrypted(size_t* len_out);
 #endif
 
 #endif
diff --git a/src/rust/bitbox02-rust/src/keystore.rs b/src/rust/bitbox02-rust/src/keystore.rs
index 2ba610d..a81efcb 100644
--- a/src/rust/bitbox02-rust/src/keystore.rs
+++ b/src/rust/bitbox02-rust/src/keystore.rs
@@ -33,12 +33,84 @@ use bitcoin::hashes::{Hash, HashEngine, Hmac, HmacEngine, sha256, sha512};
 /// Length of a compressed secp256k1 pubkey.
 const EC_PUBLIC_KEY_LEN: usize = 33;
 
+/// aes256cbc-hmac cipher adds 16 bytes IV, 16 bytes padding, 32 bytes hmac.
+const ENCRYPTION_OVERHEAD: usize = 64;
+
+#[derive(Copy, Clone)]
+struct ReadOnlyBuffer {
+    // 64 is the biggest retained buffer (bip39 seed) we will store, and 64 is added for the
+    // aes256cbc-hmac overhead.
+    data: [u8; 64 + ENCRYPTION_OVERHEAD],
+    len: usize,
+}
+
+impl ReadOnlyBuffer {
+    fn from_slice(data: &[u8]) -> Self {
+        let mut result = ReadOnlyBuffer {
+            data: [0; 64 + ENCRYPTION_OVERHEAD],
+            len: data.len(),
+        };
+        result.data[..data.len()].copy_from_slice(data);
+        result
+    }
+
+    fn as_slice(&self) -> &[u8] {
+        &self.data[..self.len]
+    }
+}
+
+/// Helper struct for retaining the seed and bip39 seed.
+#[derive(Copy, Clone)]
+struct RetainedEncryptedBuffer {
+    // Stores a random key which, after stretching, is used to encrypt the retained (bip39) seed.
+    unstretched_encryption_key: [u8; 32],
+    // Stores the encrypted (bip39) seed using aes256cbc.
+    encrypted_seed: ReadOnlyBuffer,
+    purpose: &'static str,
+}
+
+impl RetainedEncryptedBuffer {
+    fn from_buffer(
+        random: &mut impl crate::hal::Random,
+        data: &[u8],
+        purpose: &'static str,
+    ) -> Result<Self, Error> {
+        let rand: [u8; 32] = random.random_32_bytes().as_slice().try_into().unwrap();
+        let encryption_key = stretch_retained_seed_encryption_key(
+            &rand,
+            &format!("{}_in", purpose),
+            &format!("{}_out", purpose),
+        )?;
+        let iv: [u8; 16] = random.random_32_bytes()[..16].try_into().unwrap();
+        let encrypted = bitbox_aes::encrypt_with_hmac(&iv, &encryption_key, data);
+        Ok(RetainedEncryptedBuffer {
+            unstretched_encryption_key: rand,
+            encrypted_seed: ReadOnlyBuffer::from_slice(&encrypted),
+            purpose,
+        })
+    }
+
+    fn decrypt(&self) -> Result<zeroize::Zeroizing<Vec<u8>>, Error> {
+        let encryption_key = stretch_retained_seed_encryption_key(
+            &self.unstretched_encryption_key,
+            &format!("{}_in", self.purpose),
+            &format!("{}_out", self.purpose),
+        )?;
+        bitbox_aes::decrypt_with_hmac(&encryption_key, self.encrypted_seed.as_slice())
+            .map_err(|_| Error::Decrypt)
+    }
+}
+
+// Stores the encrypted BIP-39 seed after bip39-unlock.
+static RETAINED_BIP39_SEED: SyncCell<Option<RetainedEncryptedBuffer>> = SyncCell::new(None);
+
 static ROOT_FINGERPRINT: SyncCell<Option<[u8; 4]>> = SyncCell::new(None);
 
 /// Locks the keystore (resets to state before `unlock()`).
 pub fn lock() {
     keystore::_lock();
-    ROOT_FINGERPRINT.write(None)
+    ROOT_FINGERPRINT.write(None);
+    RETAINED_BIP39_SEED.write(None);
 }
 
 /// Returns false if the keystore is unlocked (unlock() followed by unlock_bip39()), true otherwise.
@@ -72,6 +144,7 @@ pub fn unlock(password: &str) -> Result<zeroize::Zeroizing<Vec<u8>>, Error> {
 /// `mnemonic_passphrase` is the bip39 passphrase used in the derivation. Use the empty string if no
 /// passphrase is needed or provided.
 pub async fn unlock_bip39(
+    random: &mut impl crate::hal::Random,
     seed: &[u8],
     mnemonic_passphrase: &str,
     yield_now: impl AsyncFn(),
@@ -88,7 +161,11 @@ pub async fn unlock_bip39(
         return Err(Error::Memory);
     }
 
-    keystore::unlock_bip39_finalize(bip39_seed.as_slice().try_into().unwrap())?;
+    RETAINED_BIP39_SEED.write(Some(RetainedEncryptedBuffer::from_buffer(
+        random,
+        bip39_seed.as_slice(),
+        "keystore_retained_bip39_seed_access",
+    )?));
 
     // Store root fingerprint.
     ROOT_FINGERPRINT.write(Some(root_fingerprint));
@@ -102,7 +179,11 @@ pub fn copy_seed() -> Result<zeroize::Zeroizing<Vec<u8>>, ()> {
 
 /// Returns a copy of the retained bip39 seed. Errors if the keystore is locked.
 pub fn copy_bip39_seed() -> Result<zeroize::Zeroizing<Vec<u8>>, ()> {
-    keystore::_copy_bip39_seed()
+    RETAINED_BIP39_SEED
+        .read()
+        .ok_or(())?
+        .decrypt()
+        .map_err(|_| ())
 }
 
 /// Restores a seed. This also unlocks the keystore with this seed.
@@ -281,6 +362,11 @@ pub extern "C" fn rust_keystore_lock() {
     lock()
 }
 
+#[unsafe(no_mangle)]
+pub extern "C" fn rust_keystore_is_unlocked_bip39() -> bool {
+    RETAINED_BIP39_SEED.read().is_some()
+}
+
 /// # Safety
 ///
 /// `encryption_key` must refer to a 32-byte buffer and `out` must have space for 32 bytes.
@@ -512,7 +598,13 @@ pub mod testing {
     pub fn mock_unlocked_using_mnemonic(mnemonic: &str, passphrase: &str) {
         let seed = crate::bip39::mnemonic_to_seed(mnemonic).unwrap();
         bitbox02::keystore::mock_unlocked(&seed);
-        util::bb02_async::block_on(super::unlock_bip39(&seed, passphrase, async || {})).unwrap();
+        util::bb02_async::block_on(super::unlock_bip39(
+            &mut crate::hal::testing::TestingRandom::new(),
+            &seed,
+            passphrase,
+            async || {},
+        ))
+        .unwrap();
     }
 
     pub const TEST_MNEMONIC: &str = "purity concert above invest pigeon category peace tuition hazard vivid latin since legal speak nation session onion library travel spell region blast estate stay";
@@ -651,6 +743,7 @@ mod tests {
 
     #[test]
     fn test_lock() {
+        let mut random = crate::hal::testing::TestingRandom::new();
         lock();
         assert!(is_locked());
 
@@ -658,7 +751,7 @@ mod tests {
             .unwrap();
         assert!(encrypt_and_store_seed(&seed, "password").is_ok());
         assert!(is_locked()); // still locked, it is only unlocked after unlock_bip39.
-        assert!(block_on(unlock_bip39(&seed, "foo", async || {})).is_ok());
+        assert!(block_on(unlock_bip39(&mut random, &seed, "foo", async || {})).is_ok());
         assert!(!is_locked());
         lock();
         assert!(is_locked());
@@ -749,6 +842,7 @@ mod tests {
         // Incorrect seed passed
         assert!(
             block_on(unlock_bip39(
+                &mut crate::hal::testing::TestingRandom::new(),
                 b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
                 "foo",
                 async || {}
@@ -756,8 +850,14 @@ mod tests {
             .is_err()
         );
         // Correct seed passed.
+        let mut random = crate::hal::testing::TestingRandom::new();
+        // Mock random value used for creating the unstretched bip39 seed encryption key.
+        random.mock_next(hex!(
+            "9b44c7048893faaf6e2d7625d13d8f1cab0765fd61f159d9713e08155d06717c"
+        ));
+
         bitbox02::securechip::fake_event_counter_reset();
-        assert!(block_on(unlock_bip39(&seed, "foo", async || {})).is_ok());
+        assert!(block_on(unlock_bip39(&mut random, &seed, "foo", async || {})).is_ok());
         assert_eq!(bitbox02::securechip::fake_event_counter(), 1);
         assert_eq!(root_fingerprint(), Ok(vec![0xf1, 0xbc, 0x3c, 0x46]),);
 
@@ -770,14 +870,16 @@ mod tests {
 
         // Check that the retained bip39 seed was encrypted with the expected encryption key.
         let decrypted = {
-            let retained_bip39_seed_encrypted: &[u8] =
-                keystore::test_get_retained_bip39_seed_encrypted();
             let expected_retained_bip39_seed_secret =
                 hex::decode("856d9a8c1ea42a69ae76324244ace674397ff1360a4ba4c85ffbd42cee8a7f29")
                     .unwrap();
             bitbox_aes::decrypt_with_hmac(
                 &expected_retained_bip39_seed_secret,
-                retained_bip39_seed_encrypted,
+                RETAINED_BIP39_SEED
+                    .read()
+                    .unwrap()
+                    .encrypted_seed
+                    .as_slice(),
             )
             .unwrap()
         };
@@ -1153,7 +1255,15 @@ mod tests {
             lock();
             let seed = &seed[..test.seed_len];
 
-            assert!(block_on(unlock_bip39(seed, test.mnemonic_passphrase, async || {})).is_err());
+            assert!(
+                block_on(unlock_bip39(
+                    &mut crate::hal::testing::TestingRandom::new(),
+                    seed,
+                    test.mnemonic_passphrase,
+                    async || {}
+                ))
+                .is_err()
+            );
 
             bitbox02::securechip::fake_event_counter_reset();
             assert!(encrypt_and_store_seed(seed, "foo").is_ok());
@@ -1162,7 +1272,15 @@ mod tests {
             assert!(is_locked());
 
             bitbox02::securechip::fake_event_counter_reset();
-            assert!(block_on(unlock_bip39(seed, test.mnemonic_passphrase, async || {})).is_ok());
+            assert!(
+                block_on(unlock_bip39(
+                    &mut crate::hal::testing::TestingRandom::new(),
+                    seed,
+                    test.mnemonic_passphrase,
+                    async || {}
+                ))
+                .is_ok()
+            );
             assert_eq!(bitbox02::securechip::fake_event_counter(), 1);
 
             assert!(!is_locked());
diff --git a/src/rust/bitbox02-rust/src/workflow/unlock.rs b/src/rust/bitbox02-rust/src/workflow/unlock.rs
index 8a11144..9cc5f4c 100644
--- a/src/rust/bitbox02-rust/src/workflow/unlock.rs
+++ b/src/rust/bitbox02-rust/src/workflow/unlock.rs
@@ -137,6 +137,7 @@ pub async fn unlock_bip39(hal: &mut impl crate::hal::Hal, seed: &[u8]) {
     let ((), result) = futures_lite::future::zip(
         super::unlock_animation::animate(),
         crate::keystore::unlock_bip39(
+            hal.random(),
             seed,
             &mnemonic_passphrase,
             // for the simulator, we don't yield at all, otherwise unlock becomes very slow in the
diff --git a/src/rust/bitbox02-sys/build.rs b/src/rust/bitbox02-sys/build.rs
index 8851031..42c3837 100644
--- a/src/rust/bitbox02-sys/build.rs
+++ b/src/rust/bitbox02-sys/build.rs
@@ -74,7 +74,6 @@ const ALLOWLIST_FNS: &[&str] = &[
     "unlock_animation_create",
     "keystore_bip39_mnemonic_to_seed",
     "keystore_copy_seed",
-    "keystore_copy_bip39_seed",
     "keystore_encrypt_and_store_seed",
     "keystore_get_bip39_word",
     "keystore_is_locked",
@@ -84,9 +83,7 @@ const ALLOWLIST_FNS: &[&str] = &[
     "keystore_secp256k1_sign",
     "keystore_unlock",
     "keystore_unlock_bip39_check",
-    "keystore_unlock_bip39_finalize",
     "keystore_test_get_retained_seed_encrypted",
-    "keystore_test_get_retained_bip39_seed_encrypted",
     "label_create",
     "gmtime",
     "memory_set_salt_root",
diff --git a/src/rust/bitbox02/src/keystore.rs b/src/rust/bitbox02/src/keystore.rs
index 8b5b9ff..31b73f5 100644
--- a/src/rust/bitbox02/src/keystore.rs
+++ b/src/rust/bitbox02/src/keystore.rs
@@ -14,7 +14,6 @@
 
 extern crate alloc;
 
-use alloc::vec;
 use alloc::vec::Vec;
 
 use bitcoin::secp256k1::{All, Secp256k1};
@@ -109,14 +108,6 @@ pub fn unlock_bip39_check(seed: &[u8]) -> Result<(), Error> {
     }
 }
 
-pub fn unlock_bip39_finalize(bip39_seed: &[u8; 64]) -> Result<(), Error> {
-    if unsafe { bitbox02_sys::keystore_unlock_bip39_finalize(bip39_seed.as_ptr()) } {
-        Ok(())
-    } else {
-        Err(Error::CannotUnlockBIP39)
-    }
-}
-
 #[cfg(feature = "testing")]
 pub fn test_get_retained_seed_encrypted() -> &'static [u8] {
     unsafe {
@@ -126,15 +117,6 @@ pub fn test_get_retained_seed_encrypted() -> &'static [u8] {
     }
 }
 
-#[cfg(feature = "testing")]
-pub fn test_get_retained_bip39_seed_encrypted() -> &'static [u8] {
-    unsafe {
-        let mut len = 0usize;
-        let ptr = bitbox02_sys::keystore_test_get_retained_bip39_seed_encrypted(&mut len);
-        core::slice::from_raw_parts(ptr, len)
-    }
-}
-
 pub fn _copy_seed() -> Result<zeroize::Zeroizing<Vec<u8>>, ()> {
     let mut seed = zeroize::Zeroizing::new([0u8; MAX_SEED_LENGTH].to_vec());
     let mut seed_len: usize = 0;
@@ -147,14 +129,6 @@ pub fn _copy_seed() -> Result<zeroize::Zeroizing<Vec<u8>>, ()> {
     }
 }
 
-pub fn _copy_bip39_seed() -> Result<zeroize::Zeroizing<Vec<u8>>, ()> {
-    let mut bip39_seed = zeroize::Zeroizing::new(vec![0u8; 64]);
-    match unsafe { bitbox02_sys::keystore_copy_bip39_seed(bip39_seed.as_mut_ptr()) } {
-        true => Ok(bip39_seed),
-        false => Err(()),
-    }
-}
-
 pub struct SignResult {
     pub signature: [u8; 64],
     pub recid: u8,