1extern crate proc_macro;
19
20use optee_teec_build::{DEFAULT_INIT_FN_NAME, DEFAULT_INVOKE_FN_NAME};
21use proc_macro::TokenStream;
22use quote::quote;
23use syn::spanned::Spanned;
24use syn::{FnArg, parse_macro_input};
25
26#[proc_macro_attribute]
36pub fn plugin_init(_args: TokenStream, input: TokenStream) -> TokenStream {
37 let f = parse_macro_input!(input as syn::ItemFn);
38 let f_vis = &f.vis;
39 let f_block = &f.block;
40 let f_sig = &f.sig;
41 let f_inputs = &f_sig.inputs;
42
43 let valid_signature = f_sig.constness.is_none()
45 && matches!(f_vis, syn::Visibility::Inherited)
46 && f_sig.abi.is_none()
47 && f_inputs.is_empty()
48 && f_sig.generics.where_clause.is_none()
49 && f_sig.variadic.is_none()
50 && check_return_type(&f);
51
52 if !valid_signature {
53 return syn::parse::Error::new(
54 f.span(),
55 "`#[plugin_init]` function must have signature `fn() -> optee_teec::Result<()>`",
56 )
57 .to_compile_error()
58 .into();
59 }
60
61 let bindgen_fn_name = quote::format_ident!("{}", DEFAULT_INIT_FN_NAME);
62 let origin_fn_name = &f_sig.ident;
63 quote!(
64 #f_vis #f_sig {
65 #f_block
66 }
67 const _: fn() -> optee_teec::Result<()> = #origin_fn_name;
68 unsafe extern "C" fn #bindgen_fn_name() -> optee_teec::raw::TEEC_Result {
69 match #origin_fn_name() {
70 Ok(()) => optee_teec::raw::TEEC_SUCCESS,
71 Err(err) => err.raw_code(),
72 }
73 }
74 )
75 .into()
76}
77
78fn check_return_type(item_fn: &syn::ItemFn) -> bool {
80 const EXPECTED: [&str; 2] = ["optee_teec", "Result"];
81 if let syn::ReturnType::Type(_, return_type) = item_fn.sig.output.to_owned()
82 && let syn::Type::Path(path) = return_type.as_ref()
83 {
84 return check_path_might_match(&path.path, &EXPECTED);
85 }
86 false
87}
88
89fn check_path_might_match(path: &syn::Path, exp: &[&str]) -> bool {
92 path.segments.len() <= exp.len()
93 && path
94 .segments
95 .iter()
96 .zip(&exp[exp.len() - path.segments.len()..])
97 .all(|(seg, exp)| seg.ident == *exp)
98}
99
100fn check_invoke_fn_params(item_fn: &syn::ItemFn) -> bool {
101 if item_fn.sig.inputs.len() != 1 {
102 return false;
103 }
104
105 let arg = item_fn.sig.inputs.first().expect("Infallible");
106 if let FnArg::Typed(typ) = arg
107 && let syn::Type::Reference(typ_ref) = typ.ty.as_ref()
108 && typ_ref.mutability.is_some()
109 && let syn::Type::Path(inner_type) = typ_ref.elem.as_ref()
110 {
111 const EXPECTED: [&str; 2] = ["optee_teec", "PluginParameters"];
112 return check_path_might_match(&inner_type.path, &EXPECTED);
113 }
114 false
115}
116
117#[proc_macro_attribute]
127pub fn plugin_invoke(_args: TokenStream, input: TokenStream) -> TokenStream {
128 let f = parse_macro_input!(input as syn::ItemFn);
129 let f_vis = &f.vis;
130 let f_block = &f.block;
131 let f_sig = &f.sig;
132 let f_inputs = &f_sig.inputs;
133
134 let valid_signature = f_sig.constness.is_none()
136 && matches!(f_vis, syn::Visibility::Inherited)
137 && f_sig.abi.is_none()
138 && f_inputs.len() == 1
139 && f_sig.generics.where_clause.is_none()
140 && f_sig.variadic.is_none()
141 && check_return_type(&f)
142 && check_invoke_fn_params(&f);
143
144 if !valid_signature {
145 return syn::parse::Error::new(
146 f.span(),
147 concat!(
148 "`#[plugin_invoke]` function must have signature",
149 " `fn(params: &mut PluginParameters) -> optee_teec::Result<()>`"
150 ),
151 )
152 .to_compile_error()
153 .into();
154 }
155
156 let bindgen_fn_name = quote::format_ident!("{}", DEFAULT_INVOKE_FN_NAME);
157 let origin_fn_name = &f_sig.ident;
158
159 quote!(
160 #f_vis #f_sig {
161 #f_block
162 }
163 const _: fn(_: &mut optee_teec::PluginParameters) -> optee_teec::Result<()> = #origin_fn_name;
164
165 unsafe extern "C" fn #bindgen_fn_name(
166 cmd: u32,
167 sub_cmd: u32,
168 data: *mut core::ffi::c_void,
169 in_len: optee_teec::raw::size_t,
170 out_len: *mut optee_teec::raw::size_t,
171 ) -> optee_teec::raw::TEEC_Result {
172 let mut parameter = match unsafe {
173 optee_teec::PluginParameters::from_raw(
174 cmd,
175 sub_cmd,
176 data,
177 in_len,
178 out_len,
179 )
180 } {
181 Ok(v) => v,
182 Err(err) => return err.raw_code(),
183 };
184
185 match #origin_fn_name(&mut parameter) {
186 Ok(()) => optee_teec::raw::TEEC_SUCCESS,
187 Err(err) => err.raw_code(),
188 }
189 }
190 )
191 .into()
192}