guarantee_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse_macro_input, ItemFn};
4
5/// Transforms an axum handler to automatically sign responses with TEE attestation.
6///
7/// The macro:
8/// 1. Extracts `Arc<EnclaveAttestor>` from axum Extension
9/// 2. Generates a request ID
10/// 3. Runs the original handler
11/// 4. Signs the response body
12/// 5. Attaches X-TEE-Attestation and X-TEE-Verified headers
13#[proc_macro_attribute]
14pub fn attest(_attr: TokenStream, item: TokenStream) -> TokenStream {
15    let input_fn = parse_macro_input!(item as ItemFn);
16    let fn_name = &input_fn.sig.ident;
17    let fn_vis = &input_fn.vis;
18    let fn_inputs = &input_fn.sig.inputs;
19    let fn_body = &input_fn.block;
20    let fn_attrs = &input_fn.attrs;
21
22    let expanded = quote! {
23        #(#fn_attrs)*
24        #fn_vis async fn #fn_name(
25            ::axum::extract::Extension(attestor): ::axum::extract::Extension<::std::sync::Arc<::guarantee::EnclaveAttestor>>,
26            #fn_inputs
27        ) -> impl ::axum::response::IntoResponse {
28            use ::axum::response::IntoResponse;
29            use ::axum::http::header::HeaderValue;
30
31            // Generate request ID
32            let request_id = ::uuid::Uuid::new_v4().to_string();
33
34            // Execute original handler
35            let inner_response = {
36                #fn_body
37            };
38
39            // Convert to axum response
40            let response = inner_response.into_response();
41            let (mut parts, body) = response.into_parts();
42
43            // Read body bytes — return 500 if body cannot be read
44            let body_bytes = match ::axum::body::to_bytes(body, usize::MAX).await {
45                Ok(bytes) => bytes,
46                Err(_) => {
47                    let error_response = ::axum::response::Response::builder()
48                        .status(::axum::http::StatusCode::INTERNAL_SERVER_ERROR)
49                        .header("content-type", "application/json")
50                        .body(::axum::body::Body::from(
51                            r#"{"error":{"code":"body_read_failed","message":"Failed to read response body for attestation"}}"#
52                        ))
53                        .expect("failed to build error response");
54                    return error_response.into_response();
55                }
56            };
57
58            // Sign the response
59            let header = attestor.sign_response(&body_bytes, &request_id);
60
61            // Insert attestation headers
62            if let Ok(val) = HeaderValue::from_str(&header.to_header_value()) {
63                parts.headers.insert("X-TEE-Attestation", val);
64            }
65            if let Ok(val) = HeaderValue::from_str("true") {
66                parts.headers.insert("X-TEE-Verified", val);
67            }
68            if let Ok(val) = HeaderValue::from_str(&request_id) {
69                parts.headers.insert("X-TEE-Request-Id", val);
70            }
71
72            ::axum::response::Response::from_parts(parts, ::axum::body::Body::from(body_bytes))
73        }
74    };
75
76    TokenStream::from(expanded)
77}
guarantee_macros/lib.rs

guarantee_macros/
lib.rs
