From 9c46d88a8d0d1292845d8961fc7c8288dbde629d Mon Sep 17 00:00:00 2001
From: Chris Schinnerl <chris@sia.tech>
Date: Fri, 5 Apr 2024 16:07:48 +0200
Subject: [PATCH] implement Drop for secrets

---
 src/address.rs | 43 +++++++++++++++++++++++++++++--------------
 src/seed.rs    |  4 ++--
 2 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/src/address.rs b/src/address.rs
index 30fff2d..471f821 100644
--- a/src/address.rs
+++ b/src/address.rs
@@ -11,18 +11,20 @@ use ed25519_dalek::{SigningKey, Signer, VerifyingKey, Verifier};
 pub struct PublicKey([u8;32]);
 
 impl PublicKey {
-	pub fn as_array(&self) -> [u8;32] {
-		self.0
-	}
-
 	pub fn verify_hash(&self, hash: &[u8;32], signature: Signature) -> bool {
 		let pk = VerifyingKey::from_bytes(&self.0).unwrap();
 		pk.verify(hash, &signature.into()).is_ok()
 	}
 }
 
+impl AsRef<[u8]> for PublicKey {
+	fn as_ref(&self) -> &[u8] {
+		&self.0
+	}
+}
+
 /// An ed25519 private key that can be used to sign a hash
-#[derive(Debug, PartialEq, Clone, Copy)]
+#[derive(Debug, PartialEq, Clone)]
 pub struct PrivateKey([u8;64]);
 
 impl PrivateKey {
@@ -31,10 +33,6 @@ impl PrivateKey {
 		PrivateKey(sk.to_keypair_bytes())
 	}
 
-	pub fn as_array(&self) -> [u8;64] {
-		self.0
-	}
-
 	pub fn public_key(&self) -> PublicKey {
 		PublicKey(self.0[32..].try_into().unwrap())
 	}
@@ -45,12 +43,27 @@ impl PrivateKey {
 	}
 }
 
+impl AsRef<[u8]> for PrivateKey {
+	fn as_ref(&self) -> &[u8] {
+		&self.0
+	}
+}
+
 impl Into<UnlockKey> for PrivateKey {
 	fn into(self) -> UnlockKey {
 		UnlockKey::new(Algorithm::ED25519, self.public_key())
 	}
 }
 
+impl Drop for PrivateKey {
+	fn drop(&mut self) {
+		// Zero out the private key
+		for byte in self.0.iter_mut() {
+			*byte = 0;
+		}
+	}
+}
+
 /// An address that can be used to receive UTXOs
 #[derive(Debug, PartialEq, Clone, Copy)]
 pub struct Address([u8;32]);
@@ -60,10 +73,6 @@ impl Address {
 		Address(addr)
 	}
 
-	pub fn as_array(&self) -> [u8;32] {
-		self.0
-	}
-
 	pub fn parse_string(s: &str) -> Result<Self, HexParseError> {
 		let s = match s.split_once(":"){
 			Some((_prefix, suffix)) => suffix,
@@ -92,6 +101,12 @@ impl Address {
 	}
 }
 
+impl AsRef<[u8]> for Address {
+	fn as_ref(&self) -> &[u8] {
+		&self.0
+	}
+}
+
 impl fmt::Display for Address {
 	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 		let mut buf = [0u8;32+6];
@@ -193,7 +208,7 @@ impl SiaEncodable for UnlockKey {
 	fn encode(&self, buf: &mut Vec<u8>) {
 		self.algorithm.encode(buf);
 		buf.extend_from_slice(&(32 as u64).to_le_bytes());
-		buf.extend_from_slice(&self.public_key.as_array());
+		buf.extend_from_slice(self.public_key.as_ref());
 	}
 }
 
diff --git a/src/seed.rs b/src/seed.rs
index 8500d08..625b701 100644
--- a/src/seed.rs
+++ b/src/seed.rs
@@ -117,7 +117,7 @@ mod tests {
 		let seed = Seed::from_mnemonic(PHRASE).unwrap();
 		for (i, expected) in test_addresses {
 			let pk = seed.private_key(i);
-			assert_eq!(pk.as_array(), expected.as_slice(), "index {}", i);
+			assert_eq!(pk.as_ref(), expected, "index {}", i);
 		}
 	}
 	
@@ -143,7 +143,7 @@ mod tests {
 		let seed = Seed::from_mnemonic(PHRASE).unwrap();
 		for (i, expected) in test_addresses {
 			let pk = seed.private_key(i).public_key();
-			assert_eq!(pk.as_array(), expected.as_slice()[32..], "index {}", i);
+			assert_eq!(pk.as_ref(), expected[32..].as_ref(), "index {}", i);
 		}
 	}