Skip to main content

optee_teec_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 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/// Attribute to derive the injected init function from an existing function
27/// ``` no_run
28/// use optee_teec_macros::plugin_init;
29///
30/// #[plugin_init]
31/// fn plugin_init() -> optee_teec::Result<()> {
32///     Ok(())
33/// }
34/// ```
35#[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    // check the function signature
44    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
78// check if return_type of the function is `optee_teec::Result<()>`
79fn 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
89// check path match the expected values
90// it might still fail if developers re-export crate as other name.
91fn 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/// Attribute to derive the injected invoke function from an existing function
118/// ``` no_run
119/// use optee_teec_macros::plugin_invoke;
120///
121/// #[plugin_invoke]
122/// fn plugin_invoke(params: &mut optee_teec::PluginParameters) -> optee_teec::Result<()> {
123///     Ok(())
124/// }
125/// ```
126#[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    // 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_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}