keystore: port copy_seed and retain_seed to Rust

diff --git a/src/keystore.c b/src/keystore.c
index ef6d802..bc39f26 100644
--- a/src/keystore.c
+++ b/src/keystore.c
@@ -29,14 +29,6 @@
 #include <rust/rust.h>
 #include <secp256k1_ecdsa_s2c.h>
 
-// Change this ONLY via keystore_unlock() or keystore_lock()
-static bool _is_unlocked_device = false;
-// Stores a random key after unlock which, after stretching, is used to encrypt the retained seed.
-static uint8_t _unstretched_retained_seed_encryption_key[32] = {0};
-// Must be defined if is_unlocked is true. ONLY ACCESS THIS WITH keystore_copy_seed().
-// Stores the encrypted seed after unlock.
-static uint8_t _retained_seed_encrypted[KEYSTORE_MAX_SEED_LENGTH + 64] = {0};
-static size_t _retained_seed_encrypted_len = 0;
 // A hash of the unencrypted retained seed, used for comparing seeds without knowing their
 // plaintext.
 static uint8_t _retained_seed_hash[32] = {0};
@@ -53,39 +45,6 @@ static bool _validate_seed_length(size_t seed_len)
     return seed_len == 16 || seed_len == 24 || seed_len == 32;
 }
 
-bool keystore_copy_seed(uint8_t* seed_out, size_t* length_out)
-{
-    if (!_is_unlocked_device) {
-        return false;
-    }
-
-    uint8_t retained_seed_encryption_key[32] = {0};
-    UTIL_CLEANUP_32(retained_seed_encryption_key);
-    if (!rust_keystore_stretch_retained_seed_encryption_key(
-            rust_util_bytes(
-                _unstretched_retained_seed_encryption_key,
-                sizeof(_unstretched_retained_seed_encryption_key)),
-            "keystore_retained_seed_access_in",
-            "keystore_retained_seed_access_out",
-            rust_util_bytes_mut(
-                retained_seed_encryption_key, sizeof(retained_seed_encryption_key)))) {
-        return false;
-    }
-    size_t len = _retained_seed_encrypted_len - 48;
-    bool password_correct = cipher_aes_hmac_decrypt(
-        _retained_seed_encrypted,
-        _retained_seed_encrypted_len,
-        seed_out,
-        &len,
-        retained_seed_encryption_key);
-    if (!password_correct) {
-        // Should never happen.
-        return false;
-    }
-    *length_out = len;
-    return true;
-}
-
 /**
  * Retrieves the encrypted seed and attempts to decrypt it using the password.
  *
@@ -153,46 +112,15 @@ static keystore_error_t _hash_seed(const uint8_t* seed, size_t seed_len, uint8_t
 
 USE_RESULT static keystore_error_t _retain_seed(const uint8_t* seed, size_t seed_len)
 {
-#ifdef TESTING
-    const uint8_t test_unstretched_retained_seed_encryption_key[32] =
-        "\xfe\x09\x76\x01\x14\x52\xa7\x22\x12\xe4\xb8\xbd\x57\x2b\x5b\xe3\x01\x41\xa3\x56\xf1\x13"
-        "\x37\xd2\x9d\x35\xea\x8f\xf9\x97\xbe\xfc";
-    memcpy(
-        _unstretched_retained_seed_encryption_key,
-        test_unstretched_retained_seed_encryption_key,
-        32);
-#else
-    random_32_bytes(_unstretched_retained_seed_encryption_key);
-#endif
-    uint8_t retained_seed_encryption_key[32] = {0};
-    UTIL_CLEANUP_32(retained_seed_encryption_key);
-    bool stretched = rust_keystore_stretch_retained_seed_encryption_key(
-        rust_util_bytes(
-            _unstretched_retained_seed_encryption_key,
-            sizeof(_unstretched_retained_seed_encryption_key)),
-        "keystore_retained_seed_access_in",
-        "keystore_retained_seed_access_out",
-        rust_util_bytes_mut(retained_seed_encryption_key, sizeof(retained_seed_encryption_key)));
-    if (!stretched) {
+    if (!rust_keystore_retain_seed(rust_util_bytes(seed, seed_len))) {
         return KEYSTORE_ERR_STRETCH_RETAINED_SEED_KEY;
     }
-    size_t len = seed_len + 64;
-    if (!cipher_aes_hmac_encrypt(
-            seed, seed_len, _retained_seed_encrypted, &len, retained_seed_encryption_key)) {
-        return KEYSTORE_ERR_ENCRYPT;
-    }
-    _retained_seed_encrypted_len = len;
 
     return _hash_seed(seed, seed_len, _retained_seed_hash);
 }
 
 static void _delete_retained_seeds(void)
 {
-    util_zero(
-        _unstretched_retained_seed_encryption_key,
-        sizeof(_unstretched_retained_seed_encryption_key));
-    util_zero(_retained_seed_encrypted, sizeof(_retained_seed_encrypted));
-    _retained_seed_encrypted_len = 0;
     util_zero(_retained_seed_hash, sizeof(_retained_seed_hash));
 }
 
@@ -245,7 +173,6 @@ keystore_error_t keystore_encrypt_and_store_seed(
     if (retain_seed_result != KEYSTORE_OK) {
         return retain_seed_result;
     }
-    _is_unlocked_device = true;
 
     return KEYSTORE_OK;
 }
@@ -253,7 +180,7 @@ keystore_error_t keystore_encrypt_and_store_seed(
 // Checks if the retained seed matches the passed seed.
 static bool _check_retained_seed(const uint8_t* seed, size_t seed_length)
 {
-    if (!_is_unlocked_device) {
+    if (!rust_keystore_is_unlocked_device()) {
         return false;
     }
     uint8_t seed_hashed[32] = {0};
@@ -300,7 +227,7 @@ keystore_error_t keystore_unlock(
         return result;
     }
     if (result == KEYSTORE_OK) {
-        if (_is_unlocked_device) {
+        if (rust_keystore_is_unlocked_device()) {
             // Already unlocked. Fail if the seed changed under our feet (should never happen).
             if (!_check_retained_seed(seed, seed_len)) {
                 Abort("Seed has suddenly changed. This should never happen.");
@@ -310,7 +237,6 @@ keystore_error_t keystore_unlock(
             if (retain_seed_result != KEYSTORE_OK) {
                 return retain_seed_result;
             }
-            _is_unlocked_device = true;
         }
         bitbox02_smarteeprom_reset_unlock_attempts();
 
@@ -334,7 +260,7 @@ keystore_error_t keystore_unlock(
 
 bool keystore_unlock_bip39_check(const uint8_t* seed, size_t seed_length)
 {
-    if (!_is_unlocked_device) {
+    if (!rust_keystore_is_unlocked_device()) {
         return false;
     }
 
@@ -347,13 +273,12 @@ bool keystore_unlock_bip39_check(const uint8_t* seed, size_t seed_length)
 
 void keystore_lock(void)
 {
-    _is_unlocked_device = false;
     _delete_retained_seeds();
 }
 
 bool keystore_is_locked(void)
 {
-    bool unlocked = _is_unlocked_device && rust_keystore_is_unlocked_bip39();
+    bool unlocked = rust_keystore_is_unlocked_device() && rust_keystore_is_unlocked_bip39();
     return !unlocked;
 }
 
@@ -403,17 +328,10 @@ bool keystore_secp256k1_sign(
 #ifdef TESTING
 void keystore_mock_unlocked(const uint8_t* seed, size_t seed_len)
 {
-    _is_unlocked_device = seed != NULL;
     if (seed != NULL) {
         if (_retain_seed(seed, seed_len) != KEYSTORE_OK) {
             Abort("couldn't retain seed");
         }
     }
 }
-
-const uint8_t* keystore_test_get_retained_seed_encrypted(size_t* len_out)
-{
-    *len_out = _retained_seed_encrypted_len;
-    return _retained_seed_encrypted;
-}
 #endif
diff --git a/src/keystore.h b/src/keystore.h
index d8e45ed..f5de6dc 100644
--- a/src/keystore.h
+++ b/src/keystore.h
@@ -44,16 +44,6 @@ typedef enum {
     KEYSTORE_ERR_STRETCH_RETAINED_SEED_KEY,
 } keystore_error_t;
 
-/**
- * Copies the retained seed into the given buffer. The caller must
- * zero the seed with util_zero once it is no longer needed.
- * @param[out] seed_out The seed bytes copied from the retained seed.
- * The buffer should be KEYSTORE_MAX_SEED_LENGTH bytes long.
- * @param[out] length_out The seed length.
- * @return true if the seed was still retained.
- */
-USE_RESULT bool keystore_copy_seed(uint8_t* seed_out, size_t* length_out);
-
 /**
  * Restores a seed. This also unlocks the keystore with this seed.
  * @param[in] seed The seed that is to be restored.
@@ -174,8 +164,6 @@ USE_RESULT bool keystore_secp256k1_sign(
  * convenience to mock the keystore state (locked, seed) in tests.
  */
 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);
 #endif
 
 #endif
diff --git a/src/rust/bitbox02-rust/src/keystore.rs b/src/rust/bitbox02-rust/src/keystore.rs
index a81efcb..836bd3d 100644
--- a/src/rust/bitbox02-rust/src/keystore.rs
+++ b/src/rust/bitbox02-rust/src/keystore.rs
@@ -101,6 +101,8 @@ impl RetainedEncryptedBuffer {
     }
 }
 
+// Stores the encrypted seed after unlock.
+static RETAINED_SEED: SyncCell<Option<RetainedEncryptedBuffer>> = SyncCell::new(None);
 // Stores the encrypted BIP-39 seed after bip39-unlock.
 static RETAINED_BIP39_SEED: SyncCell<Option<RetainedEncryptedBuffer>> = SyncCell::new(None);
 
@@ -110,6 +112,7 @@ static ROOT_FINGERPRINT: SyncCell<Option<[u8; 4]>> = SyncCell::new(None);
 pub fn lock() {
     keystore::_lock();
     ROOT_FINGERPRINT.write(None);
+    RETAINED_SEED.write(None);
     RETAINED_BIP39_SEED.write(None);
 }
 
@@ -135,6 +138,28 @@ fn verify_seed(encryption_key: &[u8], expected_seed: &[u8]) -> bool {
     decrypted.as_slice() == expected_seed
 }
 
+fn retain_seed(seed: &[u8]) -> Result<(), Error> {
+    // TODO: temporary inline mocked value for testing until we are able to pass through the HAL
+    // instance properly.
+    #[cfg(feature = "testing")]
+    let mut random = {
+        let mut r = crate::hal::testing::TestingRandom::new();
+        r.mock_next(hex_lit::hex!(
+            "fe0976011452a72212e4b8bd572b5be30141a356f11337d29d35ea8ff997befc"
+        ));
+        r
+    };
+    #[cfg(not(feature = "testing"))]
+    let mut random = crate::hal::BitBox02Random;
+
+    RETAINED_SEED.write(Some(RetainedEncryptedBuffer::from_buffer(
+        &mut random,
+        seed,
+        "keystore_retained_seed_access",
+    )?));
+    Ok(())
+}
+
 pub fn unlock(password: &str) -> Result<zeroize::Zeroizing<Vec<u8>>, Error> {
     keystore::_unlock(password)
 }
@@ -174,7 +199,7 @@ pub async fn unlock_bip39(
 
 /// Returns a copy of the retained seed. Errors if the keystore is locked.
 pub fn copy_seed() -> Result<zeroize::Zeroizing<Vec<u8>>, ()> {
-    keystore::_copy_seed()
+    RETAINED_SEED.read().ok_or(())?.decrypt().map_err(|_| ())
 }
 
 /// Returns a copy of the retained bip39 seed. Errors if the keystore is locked.
@@ -362,40 +387,19 @@ pub extern "C" fn rust_keystore_lock() {
     lock()
 }
 
+#[unsafe(no_mangle)]
+pub extern "C" fn rust_keystore_is_unlocked_device() -> bool {
+    RETAINED_SEED.read().is_some()
+}
+
 #[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.
-/// `purpose_in` and `purpose_out` must be null-terminated C strings.
 #[unsafe(no_mangle)]
-pub unsafe extern "C" fn rust_keystore_stretch_retained_seed_encryption_key(
-    encryption_key: util::bytes::Bytes,
-    purpose_in: *const core::ffi::c_char,
-    purpose_out: *const core::ffi::c_char,
-    mut out: util::bytes::BytesMut,
-) -> bool {
-    let encryption_key: [u8; 32] = match encryption_key.as_ref().try_into() {
-        Ok(key) => key,
-        Err(_) => return false,
-    };
-    let purpose_in = unsafe { bitbox02::util::str_from_null_terminated_ptr(purpose_in) };
-    let purpose_out = unsafe { bitbox02::util::str_from_null_terminated_ptr(purpose_out) };
-    let (purpose_in, purpose_out) = match (purpose_in, purpose_out) {
-        (Ok(purpose_in), Ok(purpose_out)) => (purpose_in, purpose_out),
-        _ => return false,
-    };
-
-    match stretch_retained_seed_encryption_key(&encryption_key, purpose_in, purpose_out) {
-        Ok(stretched) => {
-            out.as_mut().copy_from_slice(stretched.as_slice());
-            true
-        }
-        Err(_) => false,
-    }
+pub extern "C" fn rust_keystore_retain_seed(seed: util::bytes::Bytes) -> bool {
+    retain_seed(seed.as_ref()).is_ok()
 }
 
 #[unsafe(no_mangle)]
@@ -791,12 +795,14 @@ mod tests {
 
         // Also check that the retained seed was encrypted with the expected encryption key.
         let decrypted = {
-            let retained_seed_encrypted: &[u8] = keystore::test_get_retained_seed_encrypted();
             let expected_retained_seed_secret =
                 hex::decode("b156be416530c6fc00018844161774a3546a53ac6dd4a0462608838e216008f7")
                     .unwrap();
-            bitbox_aes::decrypt_with_hmac(&expected_retained_seed_secret, retained_seed_encrypted)
-                .unwrap()
+            bitbox_aes::decrypt_with_hmac(
+                &expected_retained_seed_secret,
+                RETAINED_SEED.read().unwrap().encrypted_seed.as_slice(),
+            )
+            .unwrap()
         };
         assert_eq!(decrypted.as_slice(), seed.as_slice());
 
@@ -1085,37 +1091,6 @@ mod tests {
         assert_eq!(stretched.as_slice(), expected.as_slice());
     }
 
-    #[test]
-    fn test_rust_keystore_stretch_retained_seed_encryption_key_success() {
-        mock_memory();
-        let salt_root =
-            hex::decode("0000000000000000111111111111111122222222222222223333333333333333")
-                .unwrap();
-        bitbox02::memory::set_salt_root(salt_root.as_slice().try_into().unwrap()).unwrap();
-
-        let encryption_key_vec =
-            hex::decode("00112233445566778899aabbccddeeff112233445566778899aabbccddeeff00")
-                .unwrap();
-
-        let mut out = [0u8; 32];
-        let purpose_in = c"keystore_retained_seed_access_in";
-        let purpose_out = c"keystore_retained_seed_access_out";
-
-        let success = unsafe {
-            rust_keystore_stretch_retained_seed_encryption_key(
-                util::bytes::rust_util_bytes(encryption_key_vec.as_ptr(), encryption_key_vec.len()),
-                purpose_in.as_ptr(),
-                purpose_out.as_ptr(),
-                util::bytes::rust_util_bytes_mut(out.as_mut_ptr(), out.len()),
-            )
-        };
-        assert!(success);
-        let expected =
-            hex::decode("b6b20683810aee16b5603ae95d14eaae5ae2c8d9df9b66e1b67c698e627bb208")
-                .unwrap();
-        assert_eq!(out, expected.as_slice());
-    }
-
     #[test]
     fn test_stretch_retained_seed_encryption_key_salt_error() {
         mock_memory();
diff --git a/src/rust/bitbox02-sys/build.rs b/src/rust/bitbox02-sys/build.rs
index 42c3837..5ae05d3 100644
--- a/src/rust/bitbox02-sys/build.rs
+++ b/src/rust/bitbox02-sys/build.rs
@@ -73,7 +73,6 @@ const ALLOWLIST_FNS: &[&str] = &[
     "empty_create",
     "unlock_animation_create",
     "keystore_bip39_mnemonic_to_seed",
-    "keystore_copy_seed",
     "keystore_encrypt_and_store_seed",
     "keystore_get_bip39_word",
     "keystore_is_locked",
@@ -83,7 +82,6 @@ const ALLOWLIST_FNS: &[&str] = &[
     "keystore_secp256k1_sign",
     "keystore_unlock",
     "keystore_unlock_bip39_check",
-    "keystore_test_get_retained_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 31b73f5..b3225ba 100644
--- a/src/rust/bitbox02/src/keystore.rs
+++ b/src/rust/bitbox02/src/keystore.rs
@@ -108,27 +108,6 @@ pub fn unlock_bip39_check(seed: &[u8]) -> Result<(), Error> {
     }
 }
 
-#[cfg(feature = "testing")]
-pub fn test_get_retained_seed_encrypted() -> &'static [u8] {
-    unsafe {
-        let mut len = 0usize;
-        let ptr = bitbox02_sys::keystore_test_get_retained_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;
-    match unsafe { bitbox02_sys::keystore_copy_seed(seed.as_mut_ptr(), &mut seed_len) } {
-        true => {
-            seed.truncate(seed_len);
-            Ok(seed)
-        }
-        false => Err(()),
-    }
-}
-
 pub struct SignResult {
     pub signature: [u8; 64],
     pub recid: u8,