From 2f4044f808ef1e65cf7829226981cb3d5fde8a0c Mon Sep 17 00:00:00 2001
From: Tony Arcieri <tony@iqlusion.io>
Date: Sat, 23 Mar 2019 13:53:10 -0700
Subject: [PATCH] zeroize: Add ZeroizeOnDrop marker trait + custom derive

Adds support for deriving the `ZeroizeOnDrop` trait, which invokes
`zeroize()` on `self` inside of the `Drop` handler.
---
 zeroize/src/lib.rs        | 35 ++++++++++++++++++++++++++++++++++-
 zeroize_derive/src/lib.rs | 21 ++++++++++++++++++++-
 2 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/zeroize/src/lib.rs b/zeroize/src/lib.rs
index e024ef11..8ec84787 100644
--- a/zeroize/src/lib.rs
+++ b/zeroize/src/lib.rs
@@ -4,7 +4,6 @@
 //! ## Usage
 //!
 //! ```
-//! extern crate zeroize;
 //! use zeroize::Zeroize;
 //!
 //! fn main() {
@@ -32,6 +31,33 @@
 //! impl [Default], which implements [Zeroize] by overwriting a value with
 //! the default value.
 //!
+//! ## Custom Derive Support
+//!
+//! This crate has custom derive support for the `Zeroize` trait, which
+//! automatically calls `zeroize()` on all members of a struct or tuple struct:
+//!
+//! ```
+//! // Ensure you import the crate with `macro_use`:
+//! // #[macro_use]
+//! // extern crate zeroize;
+//!
+//! use zeroize::Zeroize;
+//!
+//! #[derive(Zeroize)]
+//! struct MyStruct([u8; 64]);
+//! ```
+//!
+//! Additionally, you can derive `ZeroizeOnDrop`, which will automatically
+//! derive a `Drop` handler that calls `zeroize()`:
+//!
+//! ```
+//! use zeroize::{Zeroize, ZeroizeOnDrop};
+//!
+//! // This struct will be zeroized on drop
+//! #[derive(Zeroize, ZeroizeOnDrop)]
+//! struct MyStruct([u8; 64]);
+//! ```
+//!
 //! ## About
 //!
 //! [Zeroing memory securely is hard] - compilers optimize for performance, and
@@ -235,6 +261,13 @@ pub trait Zeroize {
 /// Marker trait for types whose `Default` is the desired zeroization result
 pub trait DefaultIsZeroes: Copy + Default + Sized {}
 
+/// Marker trait intended for use with `zeroize_derive` which indicates that
+/// a type should have a drop handler which calls Zeroize.
+///
+/// Use `#[derive(ZeroizeOnDrop)]` to automatically impl this trait and an
+/// associated drop handler.
+pub trait ZeroizeOnDrop: Zeroize + Drop {}
+
 impl<Z> Zeroize for Z
 where
     Z: DefaultIsZeroes,
diff --git a/zeroize_derive/src/lib.rs b/zeroize_derive/src/lib.rs
index e76bde4e..db15b058 100644
--- a/zeroize_derive/src/lib.rs
+++ b/zeroize_derive/src/lib.rs
@@ -15,7 +15,7 @@ macro_rules! q {
 }
 
 #[proc_macro_derive(Zeroize)]
-pub fn derive(input: TokenStream) -> TokenStream {
+pub fn derive_zeroize(input: TokenStream) -> TokenStream {
     let ast: syn::DeriveInput = syn::parse(input).unwrap();
 
     let zeroizers = match ast.data {
@@ -38,6 +38,25 @@ pub fn derive(input: TokenStream) -> TokenStream {
     zeroize_impl.into()
 }
 
+#[proc_macro_derive(ZeroizeOnDrop)]
+pub fn derive_zeroize_on_drop(input: TokenStream) -> TokenStream {
+    let ast: syn::DeriveInput = syn::parse(input).unwrap();
+    let name = &ast.ident;
+    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
+
+    let zeroize_on_drop_impl = q! {
+        impl #impl_generics Drop for #name #ty_generics #where_clause {
+            fn drop(&mut self) {
+                self.zeroize();
+            }
+        }
+
+        impl #impl_generics ZeroizeOnDrop for #name #ty_generics #where_clause {}
+    };
+
+    zeroize_on_drop_impl.into()
+}
+
 fn derive_struct_zeroizers(fields: &syn::Fields) -> proc_macro2::TokenStream {
     let self_ident = syn::Ident::new("self", proc_macro2::Span::call_site());