From 3d9cd398983032f012b8f54f0ab958f8df7f2609 Mon Sep 17 00:00:00 2001
From: daxpedda <daxpedda@gmail.com>
Date: Thu, 6 Jan 2022 20:19:08 +0100
Subject: [PATCH] Implement `ZeroizeOnDrop` for containers (#703)

---
 zeroize/src/lib.rs | 52 ++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 50 insertions(+), 2 deletions(-)

diff --git a/zeroize/src/lib.rs b/zeroize/src/lib.rs
index 05c29d9c..deba5f70 100644
--- a/zeroize/src/lib.rs
+++ b/zeroize/src/lib.rs
@@ -260,8 +260,7 @@ pub trait Zeroize {
 }
 
 /// Marker trait signifying that this type will [`zeroize`](Zeroize::zeroize) itself on [`Drop`].
-#[allow(drop_bounds)]
-pub trait ZeroizeOnDrop: Drop {}
+pub trait ZeroizeOnDrop {}
 
 /// Marker trait for types whose `Default` is the desired zeroization result
 pub trait DefaultIsZeroes: Copy + Default + Sized {}
@@ -324,6 +323,8 @@ where
         self.iter_mut().zeroize();
     }
 }
+/// Implement `ZeroizeOnDrop` on arrays of types that impl `ZeroizeOnDrop`
+impl<Z, const N: usize> ZeroizeOnDrop for [Z; N] where Z: ZeroizeOnDrop {}
 
 impl<'a, Z> Zeroize for IterMut<'a, Z>
 where
@@ -376,6 +377,8 @@ where
     }
 }
 
+impl<Z> ZeroizeOnDrop for Option<Z> where Z: ZeroizeOnDrop {}
+
 /// Impl `Zeroize` on slices of MaybeUninit types
 /// This impl can eventually be optimized using an memset intrinsic,
 /// such as `core::intrinsics::volatile_set_memory`.
@@ -453,6 +456,10 @@ where
     }
 }
 
+#[cfg(feature = "alloc")]
+#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
+impl<Z> ZeroizeOnDrop for Vec<Z> where Z: ZeroizeOnDrop {}
+
 #[cfg(feature = "alloc")]
 #[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
 impl<Z> Zeroize for Box<[Z]>
@@ -466,6 +473,10 @@ where
     }
 }
 
+#[cfg(feature = "alloc")]
+#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
+impl<Z> ZeroizeOnDrop for Box<[Z]> where Z: ZeroizeOnDrop {}
+
 #[cfg(feature = "alloc")]
 #[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
 impl Zeroize for String {
@@ -607,14 +618,20 @@ unsafe fn volatile_set<T: Copy + Sized>(dst: *mut T, src: T, count: usize) {
 impl<Z> Zeroize for PhantomData<Z> {
     fn zeroize(&mut self) {}
 }
+/// `PhantomData` is always zero sized so provide a ZeroizeOnDrop implementation.
+impl<Z> ZeroizeOnDrop for PhantomData<Z> {}
 /// `PhantomPinned` is zero sized so provide a Zeroize implementation.
 impl Zeroize for PhantomPinned {
     fn zeroize(&mut self) {}
 }
+/// `PhantomPinned` is zero sized so provide a ZeroizeOnDrop implementation.
+impl ZeroizeOnDrop for PhantomPinned {}
 /// `()` is zero sized so provide a Zeroize implementation.
 impl Zeroize for () {
     fn zeroize(&mut self) {}
 }
+/// `()` is zero sized so provide a ZeroizeOnDrop implementation.
+impl ZeroizeOnDrop for () {}
 
 /// Generic implementation of Zeroize for tuples up to 10 parameters.
 impl<A: Zeroize> Zeroize for (A,) {
@@ -622,6 +639,8 @@ impl<A: Zeroize> Zeroize for (A,) {
         self.0.zeroize();
     }
 }
+/// Generic implementation of ZeroizeOnDrop for tuples up to 10 parameters.
+impl<A: ZeroizeOnDrop> ZeroizeOnDrop for (A,) {}
 macro_rules! impl_zeroize_tuple {
     ( $( $type_name:ident )+ ) => {
         impl<$($type_name: Zeroize),+> Zeroize for ($($type_name),+) {
@@ -631,6 +650,8 @@ macro_rules! impl_zeroize_tuple {
                 $($type_name.zeroize());+
             }
         }
+
+        impl<$($type_name: ZeroizeOnDrop),+> ZeroizeOnDrop for ($($type_name),+) { }
     }
 }
 // Generic implementations for tuples up to 10 parameters.
@@ -653,6 +674,15 @@ mod tests {
     #[cfg(feature = "alloc")]
     use alloc::boxed::Box;
 
+    #[derive(Clone, Debug, PartialEq)]
+    struct ZeroizedOnDrop(u64);
+
+    impl Drop for ZeroizedOnDrop {
+        fn drop(&mut self) {
+            self.0.zeroize();
+        }
+    }
+
     #[test]
     fn non_zero() {
         macro_rules! non_zero_test {
@@ -688,6 +718,13 @@ mod tests {
         assert_eq!(arr.as_ref(), [0u8; 137].as_ref());
     }
 
+    #[test]
+    fn zeroize_on_drop_byte_arrays() {
+        let mut arr = [ZeroizedOnDrop(42); 1];
+        unsafe { core::ptr::drop_in_place(&mut arr) };
+        assert_eq!(arr.as_ref(), [ZeroizedOnDrop(0); 1].as_ref());
+    }
+
     #[test]
     fn zeroize_maybeuninit_byte_arrays() {
         let mut arr = [MaybeUninit::new(42u64); 64];
@@ -716,6 +753,17 @@ mod tests {
         assert_eq!(tup2, (0u8, 0u8));
     }
 
+    #[test]
+    fn zeroize_on_drop_check_tuple() {
+        let mut tup1 = (ZeroizedOnDrop(42),);
+        unsafe { core::ptr::drop_in_place(&mut tup1) };
+        assert_eq!(tup1, (ZeroizedOnDrop(0),));
+
+        let mut tup2 = (ZeroizedOnDrop(42), ZeroizedOnDrop(42));
+        unsafe { core::ptr::drop_in_place(&mut tup2) };
+        assert_eq!(tup2, (ZeroizedOnDrop(0), ZeroizedOnDrop(0)));
+    }
+
     #[cfg(feature = "alloc")]
     #[test]
     fn zeroize_vec() {