Support skipping variants with ZeroizeOnDrop (#109)

Author: daxpedda

CHANGELOG.md

@@ -11,6 +11,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Use the `Copy` implementation for `Clone` and the `Ord` implementation for
   `PartialOrd` when custom bounds are present.
 
+### Fixed
+- Support skipping only some variants with `ZeroizeOnDrop`.
+- Only generate single call to `zeroize()` with `ZeroizeOnDrop` on multiple
+  fields or variants.
+
 ## [1.2.7] - 2023-12-14
 
 ### Fixed

src/test/zeroize.rs

@@ -230,3 +230,46 @@ fn fqs() -> Result<()> {
 		},
 	)
 }
+
+#[test]
+fn enum_skip() -> Result<()> {
+	test_derive(
+		quote! {
+			#[derive_where(ZeroizeOnDrop)]
+			enum Test<T> {
+				A(std::marker::PhantomData<T>),
+				#[derive_where(skip_inner)]
+				B(std::marker::PhantomData<T>),
+			}
+		},
+		#[cfg(not(feature = "zeroize-on-drop"))]
+		quote! {
+			#[automatically_derived]
+			impl<T> ::core::ops::Drop for Test<T> {
+				fn drop(&mut self) {
+					::zeroize::Zeroize::zeroize(self);
+				}
+			}
+		},
+		#[cfg(feature = "zeroize-on-drop")]
+		quote! {
+			#[automatically_derived]
+			impl <T> ::core::ops::Drop for Test<T> {
+				fn drop(&mut self) {
+					use ::zeroize::__internal::AssertZeroize;
+					use ::zeroize::__internal::AssertZeroizeOnDrop;
+
+					match self {
+						Test::A(ref mut __field_0) => {
+							__field_0.zeroize_or_on_drop();
+						}
+						Test::B(ref mut __field_0) => { }
+					}
+				}
+			}
+
+			#[automatically_derived]
+			impl<T> ::zeroize::ZeroizeOnDrop for Test<T> { }
+		},
+	)
+}

src/trait_/zeroize_on_drop.rs

@@ -7,9 +7,9 @@ use syn::{
 	Token,
 };
 
-use crate::{
-	util, Data, DeriveTrait, DeriveWhere, Error, Item, SimpleType, SplitGenerics, TraitImpl,
-};
+use crate::{util, DeriveTrait, DeriveWhere, Error, Item, SplitGenerics, TraitImpl};
+#[cfg(feature = "zeroize-on-drop")]
+use crate::{Data, SimpleType};
 
 /// Dummy-struct implement [`Trait`](crate::Trait) for [`ZeroizeOnDrop`](https://docs.rs/zeroize/latest/zeroize/trait.ZeroizeOnDrop.html) .
 pub struct ZeroizeOnDrop;
@@ -97,80 +97,82 @@ impl TraitImpl for ZeroizeOnDrop {
 			Item::Item(data) if data.is_empty(**trait_) => quote! {
 				fn drop(&mut self) { }
 			},
+			#[cfg(feature = "zeroize-on-drop")]
 			_ => {
-				#[cfg(feature = "zeroize-on-drop")]
-				{
-					let crate_ = trait_.crate_();
-					let internal = util::path_segment("__internal");
+				let crate_ = trait_.crate_();
+				let internal = util::path_segment("__internal");
 
-					let mut assert_zeroize = crate_.clone();
-					assert_zeroize
-						.segments
-						.extend([internal.clone(), util::path_segment("AssertZeroize")]);
+				let mut assert_zeroize = crate_.clone();
+				assert_zeroize
+					.segments
+					.extend([internal.clone(), util::path_segment("AssertZeroize")]);
 
-					let mut assert_zeroize_on_drop = crate_;
-					assert_zeroize_on_drop
-						.segments
-						.extend([internal, util::path_segment("AssertZeroizeOnDrop")]);
+				let mut assert_zeroize_on_drop = crate_;
+				assert_zeroize_on_drop
+					.segments
+					.extend([internal, util::path_segment("AssertZeroizeOnDrop")]);
 
-					quote! {
-						fn drop(&mut self) {
-							use #assert_zeroize;
-							use #assert_zeroize_on_drop;
+				quote! {
+					fn drop(&mut self) {
+						use #assert_zeroize;
+						use #assert_zeroize_on_drop;
 
-							match self {
-								#body
-							}
+						match self {
+							#body
 						}
 					}
 				}
-				#[cfg(not(feature = "zeroize-on-drop"))]
+			}
+			#[cfg(not(feature = "zeroize-on-drop"))]
+			_ => {
+				// Use unused variables.
+				let _ = body;
+
+				let path = util::path_from_root_and_strs(trait_.crate_(), &["Zeroize"]);
+
 				quote! {
 					fn drop(&mut self) {
-						#body
+						#path::zeroize(self);
 					}
 				}
 			}
 		}
 	}
 
+	#[cfg(feature = "zeroize-on-drop")]
 	fn build_body(
 		&self,
 		_derive_where: &DeriveWhere,
 		trait_: &DeriveTrait,
 		data: &Data,
 	) -> TokenStream {
-		if data.is_empty(**trait_) {
-			TokenStream::new()
-		} else {
-			match data.simple_type() {
-				SimpleType::Struct(fields) | SimpleType::Tuple(fields) => {
-					#[cfg(feature = "zeroize-on-drop")]
-					{
-						let self_pattern = fields.self_pattern_mut();
-						let self_ident = data.iter_self_ident(**trait_);
-
-						quote! {
-							#self_pattern => {
-								#(#self_ident.zeroize_or_on_drop();)*
-							}
+		match data.simple_type() {
+			SimpleType::Struct(fields) | SimpleType::Tuple(fields) => {
+				#[cfg(feature = "zeroize-on-drop")]
+				{
+					let self_pattern = fields.self_pattern_mut();
+					let self_ident = data.iter_self_ident(**trait_);
+
+					quote! {
+						#self_pattern => {
+							#(#self_ident.zeroize_or_on_drop();)*
 						}
 					}
-					#[cfg(not(feature = "zeroize-on-drop"))]
-					{
-						// Use unused variables.
-						let _ = fields;
+				}
+				#[cfg(not(feature = "zeroize-on-drop"))]
+				{
+					// Use unused variables.
+					let _ = fields;
 
-						let path = util::path_from_root_and_strs(trait_.crate_(), &["Zeroize"]);
+					let path = util::path_from_root_and_strs(trait_.crate_(), &["Zeroize"]);
 
-						quote! {
-							#path::zeroize(self);
-						}
+					quote! {
+						#path::zeroize(self);
 					}
 				}
-				SimpleType::Unit(_) => TokenStream::new(),
-				SimpleType::Union => unreachable!("unexpected trait for union"),
 			}
+			SimpleType::Unit(_) => TokenStream::new(),
+			SimpleType::Union => unreachable!("unexpected trait for union"),
 		}
 	}
 }