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