Skip to main content

hyde_macros/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use quote::quote;
5use syn::{parse_macro_input, DeriveInput, Fields, Meta, Expr, Lit};
6
7/// Attribute macro that makes a struct protectable by a TEE.
8///
9/// Generates:
10/// - `fn protect(ctx, fields...) -> Result<Protected<Self>>` constructor
11/// - `Zeroize` on Drop (if `zeroize = true`, the default)
12/// - Derives `Serialize`, `Deserialize` on the struct
13///
14/// # Example
15/// ```ignore
16/// #[hyde::protect]
17/// struct DocumentKey {
18///     key_material: [u8; 32],
19/// }
20///
21/// // Generated API:
22/// let protected = DocumentKey::protect(&mut ctx, key_material)?;
23/// let doc_key: DocumentKey = protected.unprotect(&mut ctx)?;
24/// ```
25///
26/// # Attributes
27/// - `zeroize = false` — disable Zeroize on Drop (default: true)
28#[proc_macro_attribute]
29pub fn protect(attr: TokenStream, item: TokenStream) -> TokenStream {
30    let input = parse_macro_input!(item as DeriveInput);
31    let attrs = parse_protect_attrs(attr);
32
33    let name = &input.ident;
34    let vis = &input.vis;
35    let generics = &input.generics;
36
37    // Extract named fields
38    let fields = match &input.data {
39        syn::Data::Struct(data) => match &data.fields {
40            Fields::Named(named) => &named.named,
41            _ => {
42                return syn::Error::new_spanned(
43                    &input.ident,
44                    "#[hyde::protect] only supports structs with named fields",
45                )
46                .to_compile_error()
47                .into();
48            }
49        },
50        _ => {
51            return syn::Error::new_spanned(
52                &input.ident,
53                "#[hyde::protect] can only be applied to structs",
54            )
55            .to_compile_error()
56            .into();
57        }
58    };
59
60    // Build the `protect()` method parameters and struct construction
61    let field_params: Vec<_> = fields
62        .iter()
63        .map(|f| {
64            let fname = &f.ident;
65            let fty = &f.ty;
66            quote! { #fname: #fty }
67        })
68        .collect();
69
70    let field_names: Vec<_> = fields.iter().map(|f| &f.ident).collect();
71
72    // Generate Zeroize on Drop if enabled
73    let zeroize_impl = if attrs.zeroize {
74        let zeroize_fields: Vec<_> = fields
75            .iter()
76            .map(|f| {
77                let fname = &f.ident;
78                quote! { zeroize::Zeroize::zeroize(&mut self.#fname); }
79            })
80            .collect();
81
82        quote! {
83            impl #generics Drop for #name #generics {
84                fn drop(&mut self) {
85                    #(#zeroize_fields)*
86                }
87            }
88        }
89    } else {
90        quote! {}
91    };
92
93    // Keep existing attributes (except our #[hyde::protect])
94    let existing_attrs: Vec<_> = input
95        .attrs
96        .iter()
97        .filter(|a| !a.path().is_ident("protect"))
98        .collect();
99
100    let expanded = quote! {
101        #(#existing_attrs)*
102        #[derive(serde::Serialize, serde::Deserialize)]
103        #vis struct #name #generics {
104            #fields
105        }
106
107        impl #generics #name #generics {
108            /// Protect this value using a TEE backend.
109            ///
110            /// Serializes the struct, encrypts it via `HydeContext`, and returns
111            /// a `Protected<Self>` that can only be decrypted with the same TEE.
112            #vis fn protect(
113                ctx: &mut hyde_core::HydeContext,
114                #(#field_params),*
115            ) -> hyde_core::Result<hyde_core::Protected<Self>> {
116                let value = Self { #(#field_names),* };
117                hyde_core::Protected::new(ctx, &value)
118            }
119        }
120
121        #zeroize_impl
122    };
123
124    expanded.into()
125}
126
127struct ProtectAttrs {
128    zeroize: bool,
129}
130
131fn parse_protect_attrs(attr: TokenStream) -> ProtectAttrs {
132    let mut zeroize = true;
133
134    if !attr.is_empty() {
135        let parser = syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated;
136        if let Ok(metas) = syn::parse::Parser::parse(parser, attr) {
137            for meta in metas {
138                if let Meta::NameValue(nv) = meta {
139                    if nv.path.is_ident("zeroize") {
140                        if let Expr::Lit(expr_lit) = &nv.value {
141                            if let Lit::Bool(lit_bool) = &expr_lit.lit {
142                                zeroize = lit_bool.value;
143                            }
144                        }
145                    }
146                }
147            }
148        }
149    }
150
151    ProtectAttrs { zeroize }
152}