From faf17be5515508291b330ba21245f693f7cb5ec4 Mon Sep 17 00:00:00 2001
From: Soon Siang <jymchng@hotmail.com>
Date: Sat, 25 Nov 2023 22:58:54 +0800
Subject: [PATCH] finally resolved the problem of not zeroizing when EC < MEC

---
 src/secret.rs       | 38 +++++++++++++++--------------------
 tests/extern_bin.rs | 48 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 22 deletions(-)

diff --git a/src/secret.rs b/src/secret.rs
index b160098..af8519b 100644
--- a/src/secret.rs
+++ b/src/secret.rs
@@ -1,6 +1,6 @@
 use core::{
     marker::PhantomData,
-    mem::ManuallyDrop,
+    mem::{forget, ManuallyDrop},
     ops::{Add, Deref, Drop},
 };
 
@@ -58,18 +58,14 @@ impl<
         for<'brand> ClosureType: FnOnce(ExposedSecret<'brand, &'brand T>) -> ReturnType,
     {
         let returned_value = scope(ExposedSecret(&self.0, PhantomData));
-        (
-            Secret(
-                // SAFETY: Since compile error prevents constructing a `Secret` with `EC` > `MEC`,
-                // `zeroize()` is only called when `Secret` is maximally exposed
-                // and it is not possible to call `expose_secret(...)`
-                // when `Secret` is maximally exposed to access **private** `self.0` field,
-                // therefore, this is safe.
-                ManuallyDrop::new(unsafe { ManuallyDrop::take(&mut self.0) }),
-                PhantomData,
-            ),
-            returned_value,
-        )
+        // SAFETY: Since compile error prevents constructing a `Secret` with `EC` > `MEC`,
+        // `zeroize()` is only called when `Secret` is maximally exposed
+        // and it is not possible to call `expose_secret(...)`
+        // when `Secret` is maximally exposed to access **private** `self.0` field,
+        // therefore, this is safe.
+        let inner = ManuallyDrop::new(unsafe { ManuallyDrop::take(&mut self.0) });
+        forget(self);
+        (Secret(inner, PhantomData), returned_value)
     }
 }
 
@@ -90,15 +86,13 @@ where
 {
     #[inline(always)]
     fn drop(&mut self) {
-        if EC::U64 == MEC::U64 {
-            // SAFETY: Since compile error prevents constructing a `Secret` with `EC` > `MEC`,
-            // `zeroize()` is only called when `Secret` is maximally exposed
-            // and it is not possible to call `expose_secret(...)`
-            // when `Secret` is maximally exposed to access **private** `self.0` field,
-            // therefore, this is safe.
-            let mut inner = unsafe { ManuallyDrop::take(&mut self.0) };
-            inner.zeroize();
-        }
+        // SAFETY: Since compile error prevents constructing a `Secret` with `EC` > `MEC`,
+        // `zeroize()` is only called when `Secret` is maximally exposed
+        // and it is not possible to call `expose_secret(...)`
+        // when `Secret` is maximally exposed to access **private** `self.0` field,
+        // therefore, this is safe.
+        let mut inner = unsafe { ManuallyDrop::take(&mut self.0) };
+        inner.zeroize();
     }
 }
 
diff --git a/tests/extern_bin.rs b/tests/extern_bin.rs
index c09ec61..b8c8c76 100644
--- a/tests/extern_bin.rs
+++ b/tests/extern_bin.rs
@@ -14,10 +14,58 @@ fn test_expose_secret_extern() {
     assert_eq!(returned_value.inner, "MySecret".to_owned());
 }
 
+#[test]
+fn test_expose_secret_with_wrapper() {
+    use typenum::U50;
+    use zeroize::Zeroize;
+
+    #[derive(Clone, Debug, PartialEq)]
+    struct SecretString(String);
+
+    impl Zeroize for SecretString {
+        fn zeroize(&mut self) {
+            self.0.zeroize();
+        }
+    }
+
+    let secret = SecretString("MySecret".to_owned());
+    let new_secret: Secret<_, U50, _> = Secret::new(secret);
+    let (new_secret, returned_value) = new_secret.expose_secret(|exposed_secret| {
+        let returned_value = UseSecret::new((*exposed_secret).to_owned());
+        returned_value
+    });
+    assert_eq!(returned_value.inner, SecretString("MySecret".to_owned()));
+
+    let (new_secret, returned_value) = new_secret.expose_secret(|exposed_secret| {
+        let returned_value = UseSecret::new((*exposed_secret).to_owned());
+        returned_value
+    });
+    assert_eq!(returned_value.inner, SecretString("MySecret".to_owned()));
+
+    let (new_secret, returned_value) = new_secret.expose_secret(|exposed_secret| {
+        let returned_value = UseSecret::new((*exposed_secret).to_owned());
+        returned_value
+    });
+    assert_eq!(returned_value.inner, SecretString("MySecret".to_owned()));
+
+    let (new_secret, returned_value) = new_secret.expose_secret(|exposed_secret| {
+        let returned_value = UseSecret::new((*exposed_secret).to_owned());
+        returned_value
+    });
+    assert_eq!(returned_value.inner, SecretString("MySecret".to_owned()));
+
+    let (_new_secret, returned_value) = new_secret.expose_secret(|exposed_secret| {
+        let returned_value = UseSecret::new((*exposed_secret).to_owned());
+        returned_value
+    });
+    assert_eq!(returned_value.inner, SecretString("MySecret".to_owned()));
+}
+
 // #[test]
 // fn test_expose_secret_in_a_loop() {
 //     use common::{self, UseSecret};
 //     use sosecrets_rs::traits::ExposeSecret;
+//     use typenum::U50;
 
 //     let secret = "MySecret".to_owned();
 //     let new_secret: Secret<_, U50, _> = Secret::new(secret);