1 // Licensed to the Apache Software Foundation (ASF) under one
2 // or more contributor license agreements.  See the NOTICE file
3 // distributed with this work for additional information
4 // regarding copyright ownership.  The ASF licenses this file
5 // to you under the Apache License, Version 2.0 (the
6 // "License"); you may not use this file except in compliance
7 // with the License.  You may obtain a copy of the License at
8 //
9 //   http://www.apache.org/licenses/LICENSE-2.0
10 //
11 // Unless required by applicable law or agreed to in writing,
12 // software distributed under the License is distributed on an
13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14 // KIND, either express or implied.  See the License for the
15 // specific language governing permissions and limitations
16 // under the License.
17 
18 extern crate alloc;
19 extern crate proc_macro;
20 
21 #[cfg(not(feature = "std"))]
22 use alloc::vec::Vec;
23 use proc_macro::TokenStream;
24 use quote::quote;
25 use syn::parse_macro_input;
26 use syn::spanned::Spanned;
27 
28 /// Attribute to declare the entry point of creating TA.
29 ///
30 /// # Examples
31 ///
32 /// ``` no_run
33 /// #[ta_crate]
34 /// fn ta_crate() -> Result<()> { }
35 /// ```
36 #[proc_macro_attribute]
ta_create(_args: TokenStream, input: TokenStream) -> TokenStream37 pub fn ta_create(_args: TokenStream, input: TokenStream) -> TokenStream {
38     let f = parse_macro_input!(input as syn::ItemFn);
39     let ident = &f.ident;
40 
41     // check the function signature
42     let valid_signature = f.constness.is_none()
43         && match f.vis {
44             syn::Visibility::Inherited => true,
45             _ => false,
46         }
47         && f.abi.is_none()
48         && f.decl.inputs.is_empty()
49         && f.decl.generics.where_clause.is_none()
50         && f.decl.variadic.is_none();
51 
52     if !valid_signature {
53         return syn::parse::Error::new(
54             f.span(),
55             "`#[ta_crate]` function must have signature `fn() -> optee_utee::Result<()>`",
56         )
57         .to_compile_error()
58         .into();
59     }
60 
61     quote!(
62         #[no_mangle]
63         pub extern "C" fn TA_CreateEntryPoint() -> optee_utee_sys::TEE_Result {
64             match #ident() {
65                 Ok(_) => optee_utee_sys::TEE_SUCCESS,
66                 Err(e) => e.raw_code()
67             }
68         }
69 
70         #f
71     )
72     .into()
73 }
74 
75 /// Attribute to declare the entry point of destroying TA.
76 ///
77 /// # Examples
78 ///
79 /// ``` no_run
80 /// #[ta_destroy]
81 /// fn ta_destroy() { }
82 /// ```
83 #[proc_macro_attribute]
ta_destroy(_args: TokenStream, input: TokenStream) -> TokenStream84 pub fn ta_destroy(_args: TokenStream, input: TokenStream) -> TokenStream {
85     let f = parse_macro_input!(input as syn::ItemFn);
86     let ident = &f.ident;
87 
88     // check the function signature
89     let valid_signature = f.constness.is_none()
90         && match f.vis {
91             syn::Visibility::Inherited => true,
92             _ => false,
93         }
94         && f.abi.is_none()
95         && f.decl.inputs.is_empty()
96         && f.decl.generics.where_clause.is_none()
97         && f.decl.variadic.is_none()
98         && match f.decl.output {
99             syn::ReturnType::Default => true,
100             _ => false,
101         };
102 
103     if !valid_signature {
104         return syn::parse::Error::new(
105             f.span(),
106             "`#[ta_destroy]` function must have signature `fn()`",
107         )
108         .to_compile_error()
109         .into();
110     }
111 
112     quote!(
113         #[no_mangle]
114         pub extern "C" fn TA_DestroyEntryPoint() {
115             #ident();
116         }
117 
118         #f
119     )
120     .into()
121 }
122 
123 /// Attribute to declare the entry point of opening a session. Pointer to
124 /// session context pointer (*mut *mut T) can be defined as an optional
125 /// parameter.
126 ///
127 /// # Examples
128 ///
129 /// ``` no_run
130 /// #[ta_open_session]
131 /// fn open_session(params: &mut Parameters) -> Result<()> { }
132 ///
133 /// // T is the sess_ctx struct and is required to implement default trait
134 /// #[ta_open_session]
135 /// fn open_session(params: &mut Parameters, sess_ctx: &mut T) -> Result<()> { }
136 /// ```
137 #[proc_macro_attribute]
ta_open_session(_args: TokenStream, input: TokenStream) -> TokenStream138 pub fn ta_open_session(_args: TokenStream, input: TokenStream) -> TokenStream {
139     let f = parse_macro_input!(input as syn::ItemFn);
140     let ident = &f.ident;
141 
142     // check the function signature
143     let valid_signature = f.constness.is_none()
144         && match f.vis {
145             syn::Visibility::Inherited => true,
146             _ => false,
147         }
148         && f.abi.is_none()
149         && (f.decl.inputs.len() == 1 || f.decl.inputs.len() == 2)
150         && f.decl.generics.where_clause.is_none()
151         && f.decl.variadic.is_none();
152 
153     if !valid_signature {
154         return syn::parse::Error::new(
155             f.span(),
156             "`#[ta_open_session]` function must have signature `fn(&mut Parameters) -> Result<()>` or `fn(&mut Parameters, &mut T) -> Result<()>`",
157         )
158         .to_compile_error()
159         .into();
160     }
161 
162     match f.decl.inputs.len() {
163         1 => quote!(
164             #[no_mangle]
165             pub extern "C" fn TA_OpenSessionEntryPoint(
166                 param_types: u32,
167                 params: &mut [optee_utee_sys::TEE_Param; 4],
168                 sess_ctx: *mut *mut c_void,
169             ) -> optee_utee_sys::TEE_Result {
170                 let mut parameters = Parameters::from_raw(params, param_types);
171                 match #ident(&mut parameters) {
172                     Ok(_) => optee_utee_sys::TEE_SUCCESS,
173                     Err(e) => e.raw_code()
174                 }
175             }
176 
177             #f
178         )
179         .into(),
180 
181         2 => {
182             let input_types: Vec<_> = f
183                 .decl
184                 .inputs
185                 .iter()
186                 .map(|arg| match arg {
187                     &syn::FnArg::Captured(ref val) => &val.ty,
188                     _ => unreachable!(),
189                 })
190                 .collect();
191             let ctx_type = match input_types.last().unwrap() {
192                 &syn::Type::Reference(ref r) => &r.elem,
193                 _ => unreachable!(),
194             };
195 
196             quote!(
197                 #[no_mangle]
198                 pub extern "C" fn TA_OpenSessionEntryPoint(
199                     param_types: u32,
200                     params: &mut [optee_utee_sys::TEE_Param; 4],
201                     sess_ctx: *mut *mut c_void,
202                 ) -> optee_utee_sys::TEE_Result {
203                     let mut parameters = Parameters::from_raw(params, param_types);
204                     let mut ctx: #ctx_type = Default::default();
205                     match #ident(&mut parameters, &mut ctx) {
206                         Ok(_) =>
207                         {
208                             unsafe { *sess_ctx = Box::into_raw(Box::new(ctx)) as _; }
209                             optee_utee_sys::TEE_SUCCESS
210                         }
211                         Err(e) => e.raw_code()
212                     }
213                 }
214 
215                 #f
216             )
217             .into()
218         }
219         _ => unreachable!(),
220     }
221 }
222 
223 /// Attribute to declare the entry point of closing a session. Session context
224 /// raw pointer (`*mut T`) can be defined as an optional parameter.
225 ///
226 /// # Examples
227 ///
228 /// ``` no_run
229 /// #[ta_close_session]
230 /// fn close_session(sess_ctx: &mut T) { }
231 ///
232 /// #[ta_close_session]
233 /// fn close_session() { }
234 /// ```
235 #[proc_macro_attribute]
ta_close_session(_args: TokenStream, input: TokenStream) -> TokenStream236 pub fn ta_close_session(_args: TokenStream, input: TokenStream) -> TokenStream {
237     let f = parse_macro_input!(input as syn::ItemFn);
238     let ident = &f.ident;
239 
240     // check the function signature
241     let valid_signature = f.constness.is_none()
242         && match f.vis {
243             syn::Visibility::Inherited => true,
244             _ => false,
245         }
246         && f.abi.is_none()
247         && (f.decl.inputs.len() == 0 || f.decl.inputs.len() == 1)
248         && f.decl.generics.where_clause.is_none()
249         && f.decl.variadic.is_none()
250         && match f.decl.output {
251             syn::ReturnType::Default => true,
252             _ => false,
253         };
254 
255     if !valid_signature {
256         return syn::parse::Error::new(
257             f.span(),
258             "`#[ta_close_session]` function must have signature `fn(&mut T)` or `fn()`",
259         )
260         .to_compile_error()
261         .into();
262     }
263 
264     match f.decl.inputs.len() {
265         0 => quote!(
266             #[no_mangle]
267             pub extern "C" fn TA_CloseSessionEntryPoint(sess_ctx: *mut c_void) {
268                 #ident();
269             }
270 
271             #f
272         )
273         .into(),
274         1 => {
275             let input_types: Vec<_> = f
276                 .decl
277                 .inputs
278                 .iter()
279                 .map(|arg| match arg {
280                     &syn::FnArg::Captured(ref val) => &val.ty,
281                     _ => unreachable!(),
282                 })
283                 .collect();
284             let t = match input_types.first().unwrap() {
285                 &syn::Type::Reference(ref r) => &r.elem,
286                 _ => unreachable!(),
287             };
288 
289             quote!(
290                 #[no_mangle]
291                 pub extern "C" fn TA_CloseSessionEntryPoint(sess_ctx: *mut c_void) {
292                     if sess_ctx.is_null() {
293                         panic!("sess_ctx is null");
294                     }
295                     let mut b = unsafe {Box::from_raw(sess_ctx as *mut #t)};
296                     #ident(&mut b);
297                     drop(b);
298                 }
299 
300                 #f
301             )
302             .into()
303         }
304         _ => unreachable!(),
305     }
306 }
307 
308 /// Attribute to declare the entry point of invoking commands. Session context
309 /// reference (`&mut T`) can be defined as an optional parameter.
310 ///
311 /// # Examples
312 ///
313 /// ``` no_run
314 /// #[ta_invoke_command]
315 /// fn invoke_command(sess_ctx: &mut T, cmd_id: u32, params: &mut Parameters) -> Result<()> { }
316 ///
317 /// #[ta_invoke_command]
318 /// fn invoke_command(cmd_id: u32, params: &mut Parameters) -> Result<()> { }
319 /// ```
320 #[proc_macro_attribute]
ta_invoke_command(_args: TokenStream, input: TokenStream) -> TokenStream321 pub fn ta_invoke_command(_args: TokenStream, input: TokenStream) -> TokenStream {
322     let f = parse_macro_input!(input as syn::ItemFn);
323     let ident = &f.ident;
324 
325     // check the function signature
326     let valid_signature = f.constness.is_none()
327         && match f.vis {
328             syn::Visibility::Inherited => true,
329             _ => false,
330         }
331         && f.abi.is_none()
332         && (f.decl.inputs.len() == 2 || f.decl.inputs.len() == 3)
333         && f.decl.generics.where_clause.is_none()
334         && f.decl.variadic.is_none();
335 
336     if !valid_signature {
337         return syn::parse::Error::new(
338             f.span(),
339             "`#[ta_invoke_command]` function must have signature `fn(&mut T, u32, &mut Parameters) -> Result<()>` or `fn(u32, &mut Parameters) -> Result<()>`",
340         )
341         .to_compile_error()
342         .into();
343     }
344 
345     match f.decl.inputs.len() {
346         2 => quote!(
347             #[no_mangle]
348             pub extern "C" fn TA_InvokeCommandEntryPoint(
349                 sess_ctx: *mut c_void,
350                 cmd_id: u32,
351                 param_types: u32,
352                 params: &mut [optee_utee_sys::TEE_Param; 4],
353             ) -> optee_utee_sys::TEE_Result {
354                 let mut parameters = Parameters::from_raw(params, param_types);
355                 match #ident(cmd_id, &mut parameters) {
356                     Ok(_) => {
357                         optee_utee_sys::TEE_SUCCESS
358                     },
359                     Err(e) => e.raw_code()
360                 }
361             }
362 
363             #f
364         )
365         .into(),
366         3 => {
367             let input_types: Vec<_> = f
368                 .decl
369                 .inputs
370                 .iter()
371                 .map(|arg| match arg {
372                     &syn::FnArg::Captured(ref val) => &val.ty,
373                     _ => unreachable!(),
374                 })
375                 .collect();
376             let t = match input_types.first().unwrap() {
377                 &syn::Type::Reference(ref r) => &r.elem,
378                 _ => unreachable!(),
379             };
380 
381             quote!(
382                 #[no_mangle]
383                 pub extern "C" fn TA_InvokeCommandEntryPoint(
384                     sess_ctx: *mut c_void,
385                     cmd_id: u32,
386                     param_types: u32,
387                     params: &mut [optee_utee_sys::TEE_Param; 4],
388                 ) -> optee_utee_sys::TEE_Result {
389                     if sess_ctx.is_null() {
390                         return optee_utee_sys::TEE_ERROR_SECURITY;
391                     }
392                     let mut parameters = Parameters::from_raw(params, param_types);
393                     let mut b = unsafe {Box::from_raw(sess_ctx as *mut #t)};
394                     match #ident(&mut b, cmd_id, &mut parameters) {
395                         Ok(_) => {
396                             core::mem::forget(b);
397                             optee_utee_sys::TEE_SUCCESS
398                         },
399                         Err(e) => e.raw_code()
400                     }
401                 }
402 
403                 #f
404             )
405             .into()
406         }
407         _ => unreachable!(),
408     }
409 }
410