optee_utee_macros/
lib.rs

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

optee_utee_macros/
lib.rs
