From 22375547005da19df7a5524541442657a40600a7 Mon Sep 17 00:00:00 2001
From: punwai <pun@mystenlabs.com>
Date: Wed, 10 Aug 2022 09:40:46 -0700
Subject: [PATCH] [crypto] zeroize bls12381 secrets (#733)

* [crypto] implement zeroize on bls12381 secrets
---
 crypto/src/bls12381.rs             | 32 +++++++++++++++++++++++-
 crypto/src/tests/bls12381_tests.rs | 40 +++++++++++++++++++++++++++++-
 2 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/crypto/src/bls12381.rs b/crypto/src/bls12381.rs
index 41ded6994..c6f0c959e 100644
--- a/crypto/src/bls12381.rs
+++ b/crypto/src/bls12381.rs
@@ -13,6 +13,7 @@ use blst::min_sig as blst;
 
 use once_cell::sync::OnceCell;
 use rand::{rngs::OsRng, RngCore};
+use zeroize::Zeroize;
 
 use crate::{
     pubkey_bytes::PublicKeyBytes,
@@ -403,7 +404,7 @@ impl KeyPair for BLS12381KeyPair {
     }
 
     fn private(self) -> Self::PrivKey {
-        self.secret
+        BLS12381PrivateKey::from_bytes(self.secret.as_ref()).unwrap()
     }
 
     fn generate<R: rand::CryptoRng + rand::RngCore>(rng: &mut R) -> Self {
@@ -595,3 +596,32 @@ impl From<&BLS12381PublicKey> for BLS12381PublicKeyBytes {
         BLS12381PublicKeyBytes::from_bytes(pk.as_ref()).unwrap()
     }
 }
+
+impl zeroize::Zeroize for BLS12381PrivateKey {
+    fn zeroize(&mut self) {
+        self.bytes.take().zeroize();
+        self.privkey.zeroize();
+    }
+}
+
+impl zeroize::ZeroizeOnDrop for BLS12381PrivateKey {}
+
+impl Drop for BLS12381PrivateKey {
+    fn drop(&mut self) {
+        self.zeroize();
+    }
+}
+
+impl zeroize::Zeroize for BLS12381KeyPair {
+    fn zeroize(&mut self) {
+        self.secret.zeroize()
+    }
+}
+
+impl zeroize::ZeroizeOnDrop for BLS12381KeyPair {}
+
+impl Drop for BLS12381KeyPair {
+    fn drop(&mut self) {
+        self.zeroize();
+    }
+}
diff --git a/crypto/src/tests/bls12381_tests.rs b/crypto/src/tests/bls12381_tests.rs
index 4d2d67d15..af2994300 100644
--- a/crypto/src/tests/bls12381_tests.rs
+++ b/crypto/src/tests/bls12381_tests.rs
@@ -7,7 +7,9 @@ use crate::{
         BLS12381PublicKeyBytes, BLS12381Signature,
     },
     hkdf::hkdf_generate_from_ikm,
-    traits::{AggregateAuthenticator, EncodeDecodeBase64, KeyPair, ToFromBytes, VerifyingKey},
+    traits::{
+        AggregateAuthenticator, EncodeDecodeBase64, KeyPair, SigningKey, ToFromBytes, VerifyingKey,
+    },
 };
 use rand::{rngs::StdRng, SeedableRng as _};
 use sha3::Sha3_256;
@@ -400,3 +402,39 @@ async fn signature_service() {
     // Verify the signature we received.
     assert!(pk.verify(digest.as_ref(), &signature).is_ok());
 }
+
+// Checks if the private keys zeroed out
+#[test]
+fn test_sk_zeroization_on_drop() {
+    let ptr: *const u8;
+    let bytes_ptr: *const u8;
+
+    let mut sk_bytes = Vec::new();
+
+    {
+        let mut rng = StdRng::from_seed([9; 32]);
+        let kp = BLS12381KeyPair::generate(&mut rng);
+        let sk = kp.private();
+        sk_bytes.extend_from_slice(sk.as_ref());
+
+        ptr = std::ptr::addr_of!(sk.privkey) as *const u8;
+        bytes_ptr = &sk.as_ref()[0] as *const u8;
+
+        let sk_memory: &[u8] =
+            unsafe { ::std::slice::from_raw_parts(bytes_ptr, BLS12381PrivateKey::LENGTH) };
+        // Assert that this is equal to sk_bytes before deletion
+        assert_eq!(sk_memory, &sk_bytes[..]);
+    }
+
+    // Check that self.privkey is zeroized
+    unsafe {
+        for i in 0..BLS12381PrivateKey::LENGTH {
+            assert!(*ptr.add(i) == 0);
+        }
+    }
+
+    // Check that self.bytes is zeroized
+    let sk_memory: &[u8] =
+        unsafe { ::std::slice::from_raw_parts(bytes_ptr, BLS12381PrivateKey::LENGTH) };
+    assert_ne!(sk_memory, &sk_bytes[..]);
+}