1#![crate_type = "proc-macro"]
4#![forbid(unsafe_code)]
5#![warn(rust_2018_idioms, trivial_casts, unused_qualifications)]
6
7use proc_macro2::{Ident, TokenStream};
8use quote::{format_ident, quote};
9use syn::{
10 parse::{Parse, ParseStream},
11 parse_quote,
12 punctuated::Punctuated,
13 token::Comma,
14 visit::Visit,
15 Attribute, Data, DeriveInput, Expr, ExprLit, Field, Fields, Lit, Meta, Result, Variant,
16 WherePredicate,
17};
18
19const ZEROIZE_ATTR: &str = "zeroize";
21
22#[proc_macro_derive(Zeroize, attributes(zeroize))]
34pub fn derive_zeroize(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
35 derive_zeroize_impl(syn::parse_macro_input!(input as DeriveInput)).into()
36}
37
38fn derive_zeroize_impl(input: DeriveInput) -> TokenStream {
39 let attributes = ZeroizeAttrs::parse(&input);
40
41 let mut generics = input.generics.clone();
42
43 let extra_bounds = match attributes.bound {
44 Some(bounds) => bounds.0,
45 None => attributes
46 .auto_params
47 .iter()
48 .map(|type_param| -> WherePredicate {
49 parse_quote! {#type_param: Zeroize}
50 })
51 .collect(),
52 };
53
54 generics.make_where_clause().predicates.extend(extra_bounds);
55
56 let ty_name = &input.ident;
57
58 let (impl_gen, type_gen, where_) = generics.split_for_impl();
59
60 let drop_impl = if attributes.drop {
61 quote! {
62 #[doc(hidden)]
63 impl #impl_gen Drop for #ty_name #type_gen #where_ {
64 fn drop(&mut self) {
65 self.zeroize()
66 }
67 }
68 }
69 } else {
70 quote! {}
71 };
72
73 let zeroizers = generate_fields(&input, quote! { zeroize });
74 let zeroize_impl = quote! {
75 impl #impl_gen ::zeroize::Zeroize for #ty_name #type_gen #where_ {
76 fn zeroize(&mut self) {
77 #zeroizers
78 }
79 }
80 };
81
82 quote! {
83 #zeroize_impl
84 #drop_impl
85 }
86}
87
88#[proc_macro_derive(ZeroizeOnDrop, attributes(zeroize))]
95pub fn derive_zeroize_on_drop(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
96 derive_zeroize_on_drop_impl(syn::parse_macro_input!(input as DeriveInput)).into()
97}
98
99fn derive_zeroize_on_drop_impl(input: DeriveInput) -> TokenStream {
100 let zeroizers = generate_fields(&input, quote! { zeroize_or_on_drop });
101
102 let (impl_gen, type_gen, where_) = input.generics.split_for_impl();
103 let name = input.ident.clone();
104
105 let drop_impl = quote! {
106 impl #impl_gen Drop for #name #type_gen #where_ {
107 fn drop(&mut self) {
108 use ::zeroize::__internal::AssertZeroize;
109 use ::zeroize::__internal::AssertZeroizeOnDrop;
110 #zeroizers
111 }
112 }
113 };
114 let zeroize_on_drop_impl = impl_zeroize_on_drop(&input);
115
116 quote! {
117 #drop_impl
118 #zeroize_on_drop_impl
119 }
120}
121
122#[derive(Default)]
124struct ZeroizeAttrs {
125 drop: bool,
127 bound: Option<Bounds>,
129 auto_params: Vec<Ident>,
131}
132
133struct Bounds(Punctuated<WherePredicate, Comma>);
135
136impl Parse for Bounds {
137 fn parse(input: ParseStream<'_>) -> Result<Self> {
138 Ok(Self(Punctuated::parse_terminated(input)?))
139 }
140}
141
142struct BoundAccumulator<'a> {
143 generics: &'a syn::Generics,
144 params: Vec<Ident>,
145}
146
147impl<'ast> Visit<'ast> for BoundAccumulator<'ast> {
148 fn visit_path(&mut self, path: &'ast syn::Path) {
149 if path.segments.len() != 1 {
150 return;
151 }
152
153 if let Some(segment) = path.segments.first() {
154 for param in &self.generics.params {
155 if let syn::GenericParam::Type(type_param) = param {
156 if type_param.ident == segment.ident && !self.params.contains(&segment.ident) {
157 self.params.push(type_param.ident.clone());
158 }
159 }
160 }
161 }
162 }
163}
164
165impl ZeroizeAttrs {
166 fn parse(input: &DeriveInput) -> Self {
168 let mut result = Self::default();
169 let mut bound_accumulator = BoundAccumulator {
170 generics: &input.generics,
171 params: Vec::new(),
172 };
173
174 for attr in &input.attrs {
175 result.parse_attr(attr, None, None);
176 }
177
178 match &input.data {
179 syn::Data::Enum(enum_) => {
180 for variant in &enum_.variants {
181 for attr in &variant.attrs {
182 result.parse_attr(attr, Some(variant), None);
183 }
184 for field in &variant.fields {
185 for attr in &field.attrs {
186 result.parse_attr(attr, Some(variant), Some(field));
187 }
188 if !attr_skip(&field.attrs) {
189 bound_accumulator.visit_type(&field.ty);
190 }
191 }
192 }
193 }
194 syn::Data::Struct(struct_) => {
195 for field in &struct_.fields {
196 for attr in &field.attrs {
197 result.parse_attr(attr, None, Some(field));
198 }
199 if !attr_skip(&field.attrs) {
200 bound_accumulator.visit_type(&field.ty);
201 }
202 }
203 }
204 syn::Data::Union(union_) => panic!("Unsupported untagged union {:?}", union_),
205 }
206
207 result.auto_params = bound_accumulator.params;
208
209 result
210 }
211
212 fn parse_attr(&mut self, attr: &Attribute, variant: Option<&Variant>, binding: Option<&Field>) {
214 let meta_list = match &attr.meta {
215 Meta::List(list) => list,
216 _ => return,
217 };
218
219 if !meta_list.path.is_ident(ZEROIZE_ATTR) {
221 return;
222 }
223
224 for meta in attr
225 .parse_args_with(Punctuated::<Meta, Comma>::parse_terminated)
226 .unwrap_or_else(|e| panic!("error parsing attribute: {:?} ({})", attr, e))
227 {
228 self.parse_meta(&meta, variant, binding);
229 }
230 }
231
232 fn parse_meta(&mut self, meta: &Meta, variant: Option<&Variant>, binding: Option<&Field>) {
234 if meta.path().is_ident("drop") {
235 assert!(!self.drop, "duplicate #[zeroize] drop flags");
236
237 match (variant, binding) {
238 (_variant, Some(_binding)) => {
239 let item_kind = match variant {
241 Some(_) => "enum",
242 None => "struct",
243 };
244 panic!(
245 concat!(
246 "The #[zeroize(drop)] attribute is not allowed on {} fields. ",
247 "Use it on the containing {} instead.",
248 ),
249 item_kind, item_kind,
250 )
251 }
252 (Some(_variant), None) => panic!(concat!(
253 "The #[zeroize(drop)] attribute is not allowed on enum variants. ",
254 "Use it on the containing enum instead.",
255 )),
256 (None, None) => (),
257 };
258
259 self.drop = true;
260 } else if meta.path().is_ident("bound") {
261 assert!(self.bound.is_none(), "duplicate #[zeroize] bound flags");
262
263 match (variant, binding) {
264 (_variant, Some(_binding)) => {
265 let item_kind = match variant {
267 Some(_) => "enum",
268 None => "struct",
269 };
270 panic!(
271 concat!(
272 "The #[zeroize(bound)] attribute is not allowed on {} fields. ",
273 "Use it on the containing {} instead.",
274 ),
275 item_kind, item_kind,
276 )
277 }
278 (Some(_variant), None) => panic!(concat!(
279 "The #[zeroize(bound)] attribute is not allowed on enum variants. ",
280 "Use it on the containing enum instead.",
281 )),
282 (None, None) => {
283 if let Meta::NameValue(meta_name_value) = meta {
284 if let Expr::Lit(ExprLit {
285 lit: Lit::Str(lit), ..
286 }) = &meta_name_value.value
287 {
288 if lit.value().is_empty() {
289 self.bound = Some(Bounds(Punctuated::new()));
290 } else {
291 self.bound = Some(lit.parse().unwrap_or_else(|e| {
292 panic!("error parsing bounds: {:?} ({})", lit, e)
293 }));
294 }
295
296 return;
297 }
298 }
299
300 panic!(concat!(
301 "The #[zeroize(bound)] attribute expects a name-value syntax with a string literal value.",
302 "E.g. #[zeroize(bound = \"T: MyTrait\")]."
303 ))
304 }
305 }
306 } else if meta.path().is_ident("skip") {
307 if variant.is_none() && binding.is_none() {
308 panic!(concat!(
309 "The #[zeroize(skip)] attribute is not allowed on a `struct` or `enum`. ",
310 "Use it on a field or variant instead.",
311 ))
312 }
313 } else {
314 panic!("unknown #[zeroize] attribute type: {:?}", meta.path());
315 }
316 }
317}
318
319fn field_ident(n: usize, field: &Field) -> Ident {
320 if let Some(ref name) = field.ident {
321 name.clone()
322 } else {
323 format_ident!("__zeroize_field_{}", n)
324 }
325}
326
327fn generate_fields(input: &DeriveInput, method: TokenStream) -> TokenStream {
328 let input_id = &input.ident;
329 let fields: Vec<_> = match input.data {
330 Data::Enum(ref enum_) => enum_
331 .variants
332 .iter()
333 .filter_map(|variant| {
334 if attr_skip(&variant.attrs) {
335 if variant.fields.iter().any(|field| attr_skip(&field.attrs)) {
336 panic!("duplicate #[zeroize] skip flags")
337 }
338 None
339 } else {
340 let variant_id = &variant.ident;
341 Some((quote! { #input_id :: #variant_id }, &variant.fields))
342 }
343 })
344 .collect(),
345 Data::Struct(ref struct_) => vec![(quote! { #input_id }, &struct_.fields)],
346 Data::Union(ref union_) => panic!("Cannot generate fields for untagged union {:?}", union_),
347 };
348
349 let arms = fields.into_iter().map(|(name, fields)| {
350 let method_field = fields.iter().enumerate().filter_map(|(n, field)| {
351 if attr_skip(&field.attrs) {
352 None
353 } else {
354 let name = field_ident(n, field);
355 Some(quote! { #name.#method() })
356 }
357 });
358
359 let field_bindings = fields
360 .iter()
361 .enumerate()
362 .map(|(n, field)| field_ident(n, field));
363
364 let binding = match fields {
365 Fields::Named(_) => quote! {
366 #name { #(#field_bindings),* }
367 },
368 Fields::Unnamed(_) => quote! {
369 #name ( #(#field_bindings),* )
370 },
371 Fields::Unit => quote! {
372 #name
373 },
374 };
375
376 quote! {
377 #[allow(unused_variables)]
378 #binding => {
379 #(#method_field);*
380 }
381 }
382 });
383
384 quote! {
385 match self {
386 #(#arms),*
387 _ => {}
388 }
389 }
390}
391
392fn attr_skip(attrs: &[Attribute]) -> bool {
393 let mut result = false;
394 for attr in attrs.iter().map(|attr| &attr.meta) {
395 if let Meta::List(list) = attr {
396 if list.path.is_ident(ZEROIZE_ATTR) {
397 for meta in list
398 .parse_args_with(Punctuated::<Meta, Comma>::parse_terminated)
399 .unwrap_or_else(|e| panic!("error parsing attribute: {:?} ({})", list, e))
400 {
401 if let Meta::Path(path) = meta {
402 if path.is_ident("skip") {
403 assert!(!result, "duplicate #[zeroize] skip flags");
404 result = true;
405 }
406 }
407 }
408 }
409 }
410 }
411 result
412}
413
414fn impl_zeroize_on_drop(input: &DeriveInput) -> TokenStream {
415 let name = input.ident.clone();
416 let (impl_gen, type_gen, where_) = input.generics.split_for_impl();
417 quote! {
418 #[doc(hidden)]
419 impl #impl_gen ::zeroize::ZeroizeOnDrop for #name #type_gen #where_ {}
420 }
421}
422
423#[cfg(test)]
424mod tests {
425 use super::*;
426
427 #[track_caller]
428 fn test_derive(
429 f: impl Fn(DeriveInput) -> TokenStream,
430 input: TokenStream,
431 expected_output: TokenStream,
432 ) {
433 let output = f(syn::parse2(input).unwrap());
434 assert_eq!(format!("{output}"), format!("{expected_output}"));
435 }
436
437 #[track_caller]
438 fn parse_zeroize_test(unparsed: &str) -> TokenStream {
439 derive_zeroize_impl(syn::parse_str(unparsed).expect("Failed to parse test input"))
440 }
441
442 #[test]
443 fn zeroize_without_drop() {
444 test_derive(
445 derive_zeroize_impl,
446 quote! {
447 struct Z {
448 a: String,
449 b: Vec<u8>,
450 c: [u8; 3],
451 }
452 },
453 quote! {
454 impl ::zeroize::Zeroize for Z {
455 fn zeroize(&mut self) {
456 match self {
457 #[allow(unused_variables)]
458 Z { a, b, c } => {
459 a.zeroize();
460 b.zeroize();
461 c.zeroize()
462 }
463 _ => {}
464 }
465 }
466 }
467 },
468 )
469 }
470
471 #[test]
472 fn zeroize_with_drop() {
473 test_derive(
474 derive_zeroize_impl,
475 quote! {
476 #[zeroize(drop)]
477 struct Z {
478 a: String,
479 b: Vec<u8>,
480 c: [u8; 3],
481 }
482 },
483 quote! {
484 impl ::zeroize::Zeroize for Z {
485 fn zeroize(&mut self) {
486 match self {
487 #[allow(unused_variables)]
488 Z { a, b, c } => {
489 a.zeroize();
490 b.zeroize();
491 c.zeroize()
492 }
493 _ => {}
494 }
495 }
496 }
497 #[doc(hidden)]
498 impl Drop for Z {
499 fn drop(&mut self) {
500 self.zeroize()
501 }
502 }
503 },
504 )
505 }
506
507 #[test]
508 fn zeroize_with_skip() {
509 test_derive(
510 derive_zeroize_impl,
511 quote! {
512 struct Z {
513 a: String,
514 b: Vec<u8>,
515 #[zeroize(skip)]
516 c: [u8; 3],
517 }
518 },
519 quote! {
520 impl ::zeroize::Zeroize for Z {
521 fn zeroize(&mut self) {
522 match self {
523 #[allow(unused_variables)]
524 Z { a, b, c } => {
525 a.zeroize();
526 b.zeroize()
527 }
528 _ => {}
529 }
530 }
531 }
532 },
533 )
534 }
535
536 #[test]
537 fn zeroize_with_bound() {
538 test_derive(
539 derive_zeroize_impl,
540 quote! {
541 #[zeroize(bound = "T: MyTrait")]
542 struct Z<T>(T);
543 },
544 quote! {
545 impl<T> ::zeroize::Zeroize for Z<T> where T: MyTrait {
546 fn zeroize(&mut self) {
547 match self {
548 #[allow(unused_variables)]
549 Z(__zeroize_field_0) => {
550 __zeroize_field_0.zeroize()
551 }
552 _ => {}
553 }
554 }
555 }
556 },
557 )
558 }
559
560 #[test]
561 fn zeroize_only_drop() {
562 test_derive(
563 derive_zeroize_on_drop_impl,
564 quote! {
565 struct Z {
566 a: String,
567 b: Vec<u8>,
568 c: [u8; 3],
569 }
570 },
571 quote! {
572 impl Drop for Z {
573 fn drop(&mut self) {
574 use ::zeroize::__internal::AssertZeroize;
575 use ::zeroize::__internal::AssertZeroizeOnDrop;
576 match self {
577 #[allow(unused_variables)]
578 Z { a, b, c } => {
579 a.zeroize_or_on_drop();
580 b.zeroize_or_on_drop();
581 c.zeroize_or_on_drop()
582 }
583 _ => {}
584 }
585 }
586 }
587 #[doc(hidden)]
588 impl ::zeroize::ZeroizeOnDrop for Z {}
589 },
590 )
591 }
592
593 #[test]
594 fn zeroize_on_struct() {
595 parse_zeroize_test(stringify!(
596 #[zeroize(drop)]
597 struct Z {
598 a: String,
599 b: Vec<u8>,
600 c: [u8; 3],
601 }
602 ));
603 }
604
605 #[test]
606 fn zeroize_on_enum() {
607 parse_zeroize_test(stringify!(
608 #[zeroize(drop)]
609 enum Z {
610 Variant1 { a: String, b: Vec<u8>, c: [u8; 3] },
611 }
612 ));
613 }
614
615 #[test]
616 #[should_panic(expected = "#[zeroize(drop)] attribute is not allowed on struct fields")]
617 fn zeroize_on_struct_field() {
618 parse_zeroize_test(stringify!(
619 struct Z {
620 #[zeroize(drop)]
621 a: String,
622 b: Vec<u8>,
623 c: [u8; 3],
624 }
625 ));
626 }
627
628 #[test]
629 #[should_panic(expected = "#[zeroize(drop)] attribute is not allowed on struct fields")]
630 fn zeroize_on_tuple_struct_field() {
631 parse_zeroize_test(stringify!(
632 struct Z(#[zeroize(drop)] String);
633 ));
634 }
635
636 #[test]
637 #[should_panic(expected = "#[zeroize(drop)] attribute is not allowed on struct fields")]
638 fn zeroize_on_second_field() {
639 parse_zeroize_test(stringify!(
640 struct Z {
641 a: String,
642 #[zeroize(drop)]
643 b: Vec<u8>,
644 c: [u8; 3],
645 }
646 ));
647 }
648
649 #[test]
650 #[should_panic(expected = "#[zeroize(drop)] attribute is not allowed on enum fields")]
651 fn zeroize_on_tuple_enum_variant_field() {
652 parse_zeroize_test(stringify!(
653 enum Z {
654 Variant(#[zeroize(drop)] String),
655 }
656 ));
657 }
658
659 #[test]
660 #[should_panic(expected = "#[zeroize(drop)] attribute is not allowed on enum fields")]
661 fn zeroize_on_enum_variant_field() {
662 parse_zeroize_test(stringify!(
663 enum Z {
664 Variant {
665 #[zeroize(drop)]
666 a: String,
667 b: Vec<u8>,
668 c: [u8; 3],
669 },
670 }
671 ));
672 }
673
674 #[test]
675 #[should_panic(expected = "#[zeroize(drop)] attribute is not allowed on enum fields")]
676 fn zeroize_on_enum_second_variant_field() {
677 parse_zeroize_test(stringify!(
678 enum Z {
679 Variant1 {
680 a: String,
681 b: Vec<u8>,
682 c: [u8; 3],
683 },
684 Variant2 {
685 #[zeroize(drop)]
686 a: String,
687 b: Vec<u8>,
688 c: [u8; 3],
689 },
690 }
691 ));
692 }
693
694 #[test]
695 #[should_panic(expected = "#[zeroize(drop)] attribute is not allowed on enum variants")]
696 fn zeroize_on_enum_variant() {
697 parse_zeroize_test(stringify!(
698 enum Z {
699 #[zeroize(drop)]
700 Variant,
701 }
702 ));
703 }
704
705 #[test]
706 #[should_panic(expected = "#[zeroize(drop)] attribute is not allowed on enum variants")]
707 fn zeroize_on_enum_second_variant() {
708 parse_zeroize_test(stringify!(
709 enum Z {
710 Variant1,
711 #[zeroize(drop)]
712 Variant2,
713 }
714 ));
715 }
716
717 #[test]
718 #[should_panic(
719 expected = "The #[zeroize(skip)] attribute is not allowed on a `struct` or `enum`. Use it on a field or variant instead."
720 )]
721 fn zeroize_skip_on_struct() {
722 parse_zeroize_test(stringify!(
723 #[zeroize(skip)]
724 struct Z {
725 a: String,
726 b: Vec<u8>,
727 c: [u8; 3],
728 }
729 ));
730 }
731
732 #[test]
733 #[should_panic(
734 expected = "The #[zeroize(skip)] attribute is not allowed on a `struct` or `enum`. Use it on a field or variant instead."
735 )]
736 fn zeroize_skip_on_enum() {
737 parse_zeroize_test(stringify!(
738 #[zeroize(skip)]
739 enum Z {
740 Variant1,
741 Variant2,
742 }
743 ));
744 }
745
746 #[test]
747 #[should_panic(expected = "duplicate #[zeroize] skip flags")]
748 fn zeroize_duplicate_skip() {
749 parse_zeroize_test(stringify!(
750 struct Z {
751 a: String,
752 #[zeroize(skip)]
753 #[zeroize(skip)]
754 b: Vec<u8>,
755 c: [u8; 3],
756 }
757 ));
758 }
759
760 #[test]
761 #[should_panic(expected = "duplicate #[zeroize] skip flags")]
762 fn zeroize_duplicate_skip_list() {
763 parse_zeroize_test(stringify!(
764 struct Z {
765 a: String,
766 #[zeroize(skip, skip)]
767 b: Vec<u8>,
768 c: [u8; 3],
769 }
770 ));
771 }
772
773 #[test]
774 #[should_panic(expected = "duplicate #[zeroize] skip flags")]
775 fn zeroize_duplicate_skip_enum() {
776 parse_zeroize_test(stringify!(
777 enum Z {
778 #[zeroize(skip)]
779 Variant {
780 a: String,
781 #[zeroize(skip)]
782 b: Vec<u8>,
783 c: [u8; 3],
784 },
785 }
786 ));
787 }
788
789 #[test]
790 #[should_panic(expected = "duplicate #[zeroize] bound flags")]
791 fn zeroize_duplicate_bound() {
792 parse_zeroize_test(stringify!(
793 #[zeroize(bound = "T: MyTrait")]
794 #[zeroize(bound = "")]
795 struct Z<T>(T);
796 ));
797 }
798
799 #[test]
800 #[should_panic(expected = "duplicate #[zeroize] bound flags")]
801 fn zeroize_duplicate_bound_list() {
802 parse_zeroize_test(stringify!(
803 #[zeroize(bound = "T: MyTrait", bound = "")]
804 struct Z<T>(T);
805 ));
806 }
807
808 #[test]
809 #[should_panic(
810 expected = "The #[zeroize(bound)] attribute is not allowed on struct fields. Use it on the containing struct instead."
811 )]
812 fn zeroize_bound_struct() {
813 parse_zeroize_test(stringify!(
814 struct Z<T> {
815 #[zeroize(bound = "T: MyTrait")]
816 a: T,
817 }
818 ));
819 }
820
821 #[test]
822 #[should_panic(
823 expected = "The #[zeroize(bound)] attribute is not allowed on enum variants. Use it on the containing enum instead."
824 )]
825 fn zeroize_bound_enum() {
826 parse_zeroize_test(stringify!(
827 enum Z<T> {
828 #[zeroize(bound = "T: MyTrait")]
829 A(T),
830 }
831 ));
832 }
833
834 #[test]
835 #[should_panic(
836 expected = "The #[zeroize(bound)] attribute is not allowed on enum fields. Use it on the containing enum instead."
837 )]
838 fn zeroize_bound_enum_variant_field() {
839 parse_zeroize_test(stringify!(
840 enum Z<T> {
841 A {
842 #[zeroize(bound = "T: MyTrait")]
843 a: T,
844 },
845 }
846 ));
847 }
848
849 #[test]
850 #[should_panic(
851 expected = "The #[zeroize(bound)] attribute expects a name-value syntax with a string literal value.E.g. #[zeroize(bound = \"T: MyTrait\")]."
852 )]
853 fn zeroize_bound_no_value() {
854 parse_zeroize_test(stringify!(
855 #[zeroize(bound)]
856 struct Z<T>(T);
857 ));
858 }
859
860 #[test]
861 #[should_panic(expected = "error parsing bounds: LitStr { token: \"T\" } (expected `:`)")]
862 fn zeroize_bound_no_where_predicate() {
863 parse_zeroize_test(stringify!(
864 #[zeroize(bound = "T")]
865 struct Z<T>(T);
866 ));
867 }
868}