From 5f6abb6de5baefd0610760b7cf27b356d9741a8b Mon Sep 17 00:00:00 2001
From: Salvatore Ingala <6681844+bigspider@users.noreply.github.com>
Date: Fri, 6 Dec 2024 12:01:10 +0100
Subject: [PATCH] Zeroize private keys on drop

---
 app-sdk/Cargo.toml           |  1 +
 app-sdk/src/curve.rs         | 12 ++++++------
 apps/sadik/app/src/main.rs   |  2 +-
 vm/Cargo.toml                |  1 +
 vm/src/handlers/lib/ecall.rs | 12 +++++-------
 5 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/app-sdk/Cargo.toml b/app-sdk/Cargo.toml
index da65ae1..817c8a9 100644
--- a/app-sdk/Cargo.toml
+++ b/app-sdk/Cargo.toml
@@ -10,6 +10,7 @@ ctor = "0.2.8"
 embedded-alloc = "0.5.1"
 subtle = { version="2.6.1", default-features = false }
 hex-literal = "0.4.1"
+zeroize = "1.8.1"
 
 [target.'cfg(target_arch = "x86_64")'.dependencies]
 bip32 = "0.5.2"
diff --git a/app-sdk/src/curve.rs b/app-sdk/src/curve.rs
index a16187b..cc5fc9e 100644
--- a/app-sdk/src/curve.rs
+++ b/app-sdk/src/curve.rs
@@ -1,9 +1,9 @@
+use zeroize::Zeroizing;
+
 use common::ecall_constants::CurveKind;
 
 use crate::ecalls::{Ecall, EcallsInterface};
 
-// TODO: make sure private data zeroed out on Drop.
-
 /// A struct representing a Hierarchical Deterministic (HD) node composed of a private key, and a 32-byte chaincode.
 ///
 /// # Type Parameters
@@ -16,14 +16,14 @@ use crate::ecalls::{Ecall, EcallsInterface};
 /// * `privkey` - An array of bytes representing the private key, with a length defined by `SCALAR_LENGTH`.
 pub struct HDPrivNode<const SCALAR_LENGTH: usize> {
     pub chaincode: [u8; 32],
-    pub privkey: [u8; SCALAR_LENGTH],
+    pub privkey: Zeroizing<[u8; SCALAR_LENGTH]>,
 }
 
 impl<const SCALAR_LENGTH: usize> Default for HDPrivNode<SCALAR_LENGTH> {
     fn default() -> Self {
         Self {
             chaincode: [0u8; 32],
-            privkey: [0u8; SCALAR_LENGTH],
+            privkey: Zeroizing::new([0u8; SCALAR_LENGTH]),
         }
     }
 }
@@ -109,7 +109,7 @@ mod tests {
             hex!("eb473a0fa0af5031f14db9fe7c37bb8416a4ff01bb69dae9966dc83b5e5bf921")
         );
         assert_eq!(
-            node.privkey,
+            node.privkey[..],
             hex!("34ac5d784ebb4df4727bcddf6a6743f5d5d46d83dd74aa825866390c694f2938")
         );
 
@@ -120,7 +120,7 @@ mod tests {
             hex!("6da5f32f47232b3b9b2d6b59b802e2b313afa7cbda242f73da607139d8e04989")
         );
         assert_eq!(
-            node.privkey,
+            node.privkey[..],
             hex!("239841e64103fd024b01283e752a213fee1a8969f6825204ee3617a45c5e4a91")
         );
     }
diff --git a/apps/sadik/app/src/main.rs b/apps/sadik/app/src/main.rs
index 0bde1cb..bf16368 100644
--- a/apps/sadik/app/src/main.rs
+++ b/apps/sadik/app/src/main.rs
@@ -189,7 +189,7 @@ pub fn main(_: isize, _: *const *const u8) -> isize {
                 Curve::Secp256k1 => {
                     let node = sdk::curve::Secp256k1::derive_hd_node(&path).unwrap();
                     let mut result = node.chaincode.to_vec();
-                    result.extend_from_slice(&node.privkey);
+                    result.extend_from_slice(&node.privkey[..]);
                     result
                 }
             },
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
index 8cd277b..21769eb 100644
--- a/vm/Cargo.toml
+++ b/vm/Cargo.toml
@@ -14,6 +14,7 @@ hex = { version = "0.4.3", default-features = false, features = ["serde", "alloc
 numtoa = "0.2.4"
 postcard = { version = "1.0.8", features = ["alloc"] }
 ledger_secure_sdk_sys = "1.5.3"
+zeroize = "1.8.1"
 
 [profile.release]
 opt-level = 3
diff --git a/vm/src/handlers/lib/ecall.rs b/vm/src/handlers/lib/ecall.rs
index 49cec94..626c298 100644
--- a/vm/src/handlers/lib/ecall.rs
+++ b/vm/src/handlers/lib/ecall.rs
@@ -19,6 +19,8 @@ use crate::{AppSW, Instruction};
 
 use super::outsourced_mem::OutsourcedMemory;
 
+use zeroize::Zeroizing;
+
 // BIP32 supports up to 255, but we don't want that many, and it would be very slow anyway
 const MAX_BIP32_PATH: usize = 16;
 
@@ -704,7 +706,7 @@ impl<'a> CommEcallHandler<'a> {
         };
 
         // derive the key
-        let mut private_key_local: [u8; 32] = [0; 32];
+        let mut private_key_local = Zeroizing::new([0u8; 32]);
         let mut chain_code_local: [u8; 32] = [0; 32];
         unsafe {
             ledger_secure_sdk_sys::os_perso_derive_node_bip32(
@@ -718,12 +720,10 @@ impl<'a> CommEcallHandler<'a> {
 
         // copy private_key and chain_code to V-App memory
         cpu.get_segment(private_key.0)?
-            .write_buffer(private_key.0, &private_key_local)?;
+            .write_buffer(private_key.0, &private_key_local[..])?;
         cpu.get_segment(chain_code.0)?
             .write_buffer(chain_code.0, &chain_code_local)?;
 
-        // TODO: we should to make sure the private key is zeroed before returning
-
         Ok(())
     }
 
@@ -737,7 +737,7 @@ impl<'a> CommEcallHandler<'a> {
         }
 
         // derive the key
-        let mut private_key_local: [u8; 32] = [0; 32];
+        let mut private_key_local = Zeroizing::new([0u8; 32]);
         let mut chain_code_local: [u8; 32] = [0; 32];
 
         let mut pubkey: ledger_secure_sdk_sys::cx_ecfp_public_key_t = Default::default();
@@ -770,8 +770,6 @@ impl<'a> CommEcallHandler<'a> {
             if ret1 != CX_OK || ret2 != CX_OK {
                 return Err("Failed to generate key pair");
             }
-
-            // TODO: make sure that the private key is deleted
         }
 
         let mut sha_hasher = ledger_device_sdk::hash::sha2::Sha2_256::new();