optee_utee/
tee_parameter.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
18use crate::{Error, ErrorKind, ParamType, Result};
19use core::ops::{Index, IndexMut};
20use optee_utee_sys as raw;
21
22#[derive(Copy, Clone, Debug)]
23pub enum ParamIndex {
24    Arg0,
25    Arg1,
26    Arg2,
27    Arg3,
28}
29
30impl ParamIndex {
31    fn to_usize(self) -> usize {
32        match self {
33            ParamIndex::Arg0 => 0,
34            ParamIndex::Arg1 => 1,
35            ParamIndex::Arg2 => 2,
36            ParamIndex::Arg3 => 3,
37        }
38    }
39}
40
41enum ParamContent<'a> {
42    None,
43    MemrefInput {
44        buffer: &'a [u8],
45    },
46    MemrefOutput {
47        buffer: &'a mut [u8],
48        written: usize,
49    },
50    MemrefInout {
51        buffer: &'a mut [u8],
52        written: usize,
53    },
54    ValueInput {
55        a: u32,
56        b: u32,
57    },
58    ValueOutput {
59        a: u32,
60        b: u32,
61    },
62    ValueInout {
63        a: u32,
64        b: u32,
65    },
66}
67
68pub struct Param<'a> {
69    content: ParamContent<'a>,
70}
71
72impl<'a> Param<'a> {
73    fn new() -> Self {
74        Self {
75            content: ParamContent::None,
76        }
77    }
78
79    /// Returns the written slice if available (for `MemrefOutput` or `MemrefInout`).
80    /// Returns `None` for other types. Developers can decide whether to treat absence as an error.
81    pub fn written_slice(&self) -> Option<&[u8]> {
82        match &self.content {
83            ParamContent::MemrefOutput { buffer, written } => Some(&buffer[..*written]),
84            ParamContent::MemrefInout { buffer, written } => Some(&buffer[..*written]),
85            _ => None,
86        }
87    }
88
89    /// Returns the output values if available (for `ValueOutput` or `ValueInout`).
90    /// Returns `None` for other types. Caller decides how to handle missing values.
91    pub fn output_value(&self) -> Option<(u32, u32)> {
92        match &self.content {
93            ParamContent::ValueOutput { a, b } => Some((*a, *b)),
94            ParamContent::ValueInout { a, b } => Some((*a, *b)),
95            _ => None,
96        }
97    }
98
99    fn get_type(&self) -> ParamType {
100        match &self.content {
101            ParamContent::None => ParamType::None,
102            ParamContent::MemrefInput { .. } => ParamType::MemrefInput,
103            ParamContent::MemrefOutput { .. } => ParamType::MemrefOutput,
104            ParamContent::MemrefInout { .. } => ParamType::MemrefInout,
105            ParamContent::ValueInput { .. } => ParamType::ValueInput,
106            ParamContent::ValueOutput { .. } => ParamType::ValueOutput,
107            ParamContent::ValueInout { .. } => ParamType::ValueInout,
108        }
109    }
110
111    fn get_raw_type(&self) -> u32 {
112        self.get_type() as u32
113    }
114
115    fn as_raw(&mut self) -> raw::TEE_Param {
116        match &mut self.content {
117            ParamContent::None => raw::TEE_Param {
118                memref: raw::Memref {
119                    buffer: core::ptr::null_mut(),
120                    size: 0,
121                },
122            },
123            ParamContent::MemrefInput { buffer } => raw::TEE_Param {
124                memref: raw::Memref {
125                    buffer: (*buffer).as_ptr() as *mut core::ffi::c_void,
126                    size: buffer.len(),
127                },
128            },
129            ParamContent::MemrefOutput { buffer, written: _ } => raw::TEE_Param {
130                memref: raw::Memref {
131                    buffer: (*buffer).as_mut_ptr() as *mut core::ffi::c_void,
132                    size: buffer.len(),
133                },
134            },
135            ParamContent::MemrefInout { buffer, written: _ } => raw::TEE_Param {
136                memref: raw::Memref {
137                    buffer: (*buffer).as_mut_ptr() as *mut core::ffi::c_void,
138                    size: buffer.len(),
139                },
140            },
141            ParamContent::ValueInput { a, b } => raw::TEE_Param {
142                value: raw::Value { a: *a, b: *b },
143            },
144            ParamContent::ValueInout { a, b } => raw::TEE_Param {
145                value: raw::Value { a: *a, b: *b },
146            },
147            ParamContent::ValueOutput { a, b } => raw::TEE_Param {
148                value: raw::Value { a: *a, b: *b },
149            },
150        }
151    }
152
153    fn update_size_from_raw(&mut self, raw_param: &raw::TEE_Param) -> Result<()> {
154        match &mut self.content {
155            ParamContent::MemrefOutput { buffer, written } => {
156                // SAFETY:
157                // The caller must ensure this param is of memref type and properly initialized.
158                // This is enforced by the variant match on `ParamContent::MemrefOutput`.
159                // Accessing `raw_param.memref.size` is safe under these assumptions.
160                let new_size = unsafe { raw_param.memref.size };
161                if new_size > (*buffer).len() {
162                    return Err(Error::new(ErrorKind::BadParameters));
163                }
164                *written = new_size;
165                Ok(())
166            }
167            ParamContent::MemrefInout { buffer, written } => {
168                // SAFETY:
169                // The caller must ensure this param is of memref type and properly initialized.
170                // This is enforced by the variant match on `ParamContent::MemrefOutput`.
171                // Accessing `raw_param.memref.size` is safe under these assumptions.
172                let new_size = unsafe { raw_param.memref.size };
173                if new_size > (*buffer).len() {
174                    return Err(Error::new(ErrorKind::BadParameters));
175                }
176                *written = new_size;
177                Ok(())
178            }
179            _ => Err(Error::new(ErrorKind::BadFormat)),
180        }
181    }
182
183    fn update_value_from_raw(&mut self, raw_param: &raw::TEE_Param) {
184        match &mut self.content {
185            ParamContent::ValueInout { a, b } => {
186                // SAFETY:
187                // The caller must ensure this param is of value type and properly initialized.
188                // This is guaranteed by matching against `ParamContent::ValueInout`.
189                // Accessing `raw_param.value.a` is safe under above assumption.
190                *a = unsafe { raw_param.value.a };
191                // SAFETY:
192                // Accessing `raw_param.value.b` is safe under above assumption.
193                *b = unsafe { raw_param.value.b };
194            }
195            ParamContent::ValueOutput { a, b } => {
196                // SAFETY:
197                // The caller must ensure this param is of value type and properly initialized.
198                // This is guaranteed by matching against `ParamContent::ValueInout`.
199                // Accessing `raw_param.value.a` is safe under above assumption.
200                *a = unsafe { raw_param.value.a };
201                // SAFETY:
202                // Accessing `raw_param.value.b` is safe under above assumption.
203                *b = unsafe { raw_param.value.b };
204            }
205            _ => {}
206        }
207    }
208}
209
210/// The TeeParams struct is used to manage the parameters for TEE commands.
211pub struct TeeParams<'a> {
212    params: [Param<'a>; 4],
213}
214
215impl<'a> Default for TeeParams<'a> {
216    fn default() -> Self {
217        Self::new()
218    }
219}
220
221impl<'a> TeeParams<'a> {
222    pub fn new() -> Self {
223        Self {
224            params: [Param::new(), Param::new(), Param::new(), Param::new()],
225        }
226    }
227
228    /// These functions allow for method-chaining to easily configure multiple parameters at once.
229    ///
230    /// The following methods can be chained:
231    /// - `with_memref_in`: Sets a memory reference input parameter.
232    /// - `with_memref_out`: Sets a memory reference output parameter.
233    /// - `with_memref_inout`: Sets a memory reference inout parameter.
234    /// - `with_value_in`: Sets a value input parameter.
235    /// - `with_value_out`: Sets a value output parameter.
236    /// - `with_value_inout`: Sets a value inout parameter.
237    ///
238    /// Example usage:
239    /// ``` rust,no_run
240    /// # use optee_utee::{TeeParams, ParamIndex};
241    /// # let input_buffer = [0_u8; 0];
242    /// # let mut output_buffer = [0_u8; 0];
243    ///
244    /// let params = TeeParams::new()
245    ///     .with_memref_in(ParamIndex::Arg0, &input_buffer)
246    ///     .with_memref_out(ParamIndex::Arg1, &mut output_buffer)
247    ///     .with_value_in(ParamIndex::Arg2, 42, 0)
248    ///     .with_value_out(ParamIndex::Arg3, 0, 0);
249    /// ```
250    pub fn with_memref_in(mut self, idx: ParamIndex, buffer: &'a [u8]) -> Self {
251        self[idx].content = ParamContent::MemrefInput { buffer };
252        self
253    }
254
255    pub fn with_memref_out(mut self, idx: ParamIndex, buffer: &'a mut [u8]) -> Self {
256        self[idx].content = ParamContent::MemrefOutput { buffer, written: 0 };
257        self
258    }
259
260    pub fn with_memref_inout(mut self, idx: ParamIndex, buffer: &'a mut [u8]) -> Self {
261        self[idx].content = ParamContent::MemrefInout { buffer, written: 0 };
262        self
263    }
264
265    pub fn with_value_in(mut self, idx: ParamIndex, a: u32, b: u32) -> Self {
266        self[idx].content = ParamContent::ValueInput { a, b };
267        self
268    }
269
270    pub fn with_value_out(mut self, idx: ParamIndex, a: u32, b: u32) -> Self {
271        self[idx].content = ParamContent::ValueOutput { a, b };
272        self
273    }
274
275    pub fn with_value_inout(mut self, idx: ParamIndex, a: u32, b: u32) -> Self {
276        self[idx].content = ParamContent::ValueInout { a, b };
277        self
278    }
279
280    /// These methods allow the user to set the content at a specific index.
281    ///
282    /// Example usage:
283    /// ``` rust,no_run
284    /// # use optee_utee::{TeeParams, ParamIndex};
285    /// # let input_buffer = [0_u8; 0];
286    /// # let mut output_buffer = [0_u8; 0];
287    ///
288    /// let mut params = TeeParams::new();
289    /// params.set_memref_in(ParamIndex::Arg0, &input_buffer);
290    /// params.set_memref_out(ParamIndex::Arg1, &mut output_buffer);
291    /// params.set_value_in(ParamIndex::Arg2, 42, 0);
292    /// params.set_value_out(ParamIndex::Arg3, 0, 0);
293    /// ```
294    pub fn set_memref_in(&mut self, idx: ParamIndex, buffer: &'a [u8]) -> &mut Self {
295        self[idx].content = ParamContent::MemrefInput { buffer };
296        self
297    }
298
299    pub fn set_memref_out(&mut self, idx: ParamIndex, buffer: &'a mut [u8]) -> &mut Self {
300        self[idx].content = ParamContent::MemrefOutput { buffer, written: 0 };
301        self
302    }
303
304    pub fn set_memref_inout(&mut self, idx: ParamIndex, buffer: &'a mut [u8]) -> &mut Self {
305        self[idx].content = ParamContent::MemrefInout { buffer, written: 0 };
306        self
307    }
308
309    pub fn set_value_in(&mut self, idx: ParamIndex, a: u32, b: u32) -> &mut Self {
310        self[idx].content = ParamContent::ValueInput { a, b };
311        self
312    }
313
314    pub fn set_value_out(&mut self, idx: ParamIndex, a: u32, b: u32) -> &mut Self {
315        self[idx].content = ParamContent::ValueOutput { a, b };
316        self
317    }
318
319    pub fn set_value_inout(&mut self, idx: ParamIndex, a: u32, b: u32) -> &mut Self {
320        self[idx].content = ParamContent::ValueInout { a, b };
321        self
322    }
323
324    pub(crate) fn raw_param_types(&self) -> u32 {
325        let mut param_types = 0;
326        for (i, param) in self.params.iter().enumerate() {
327            param_types |= param.get_raw_type() << (i * 4);
328        }
329        param_types
330    }
331
332    pub(crate) fn as_raw(&mut self) -> [raw::TEE_Param; 4] {
333        [
334            self.params[0].as_raw(),
335            self.params[1].as_raw(),
336            self.params[2].as_raw(),
337            self.params[3].as_raw(),
338        ]
339    }
340
341    /// Updates the parameters with results after each TEE call.
342    ///
343    /// This function updates the content of parameters for `MemrefInout`, `MemrefOutput`, `ValueInout`, and `ValueOutput`.
344    /// Parameters of other types are not modified.
345    pub(crate) fn update_from_raw(&mut self, raw_params: &[raw::TEE_Param; 4]) -> Result<()> {
346        // update the content for memref inout/out, and value inout/out
347        for (i, param) in self.params.iter_mut().enumerate() {
348            let raw_param = &raw_params[i];
349            match param.get_type() {
350                ParamType::MemrefOutput => {
351                    param.update_size_from_raw(raw_param)?;
352                }
353                ParamType::MemrefInout => {
354                    param.update_size_from_raw(raw_param)?;
355                }
356                ParamType::ValueOutput => {
357                    param.update_value_from_raw(raw_param);
358                }
359                ParamType::ValueInout => {
360                    param.update_value_from_raw(raw_param);
361                }
362                _ => {
363                    // No action needed for other types
364                }
365            }
366        }
367        Ok(())
368    }
369}
370
371// Index trait implementations for direct parameter access
372impl<'a> Index<ParamIndex> for TeeParams<'a> {
373    type Output = Param<'a>;
374
375    fn index(&self, index: ParamIndex) -> &Self::Output {
376        &self.params[index.to_usize()]
377    }
378}
379
380impl<'a> IndexMut<ParamIndex> for TeeParams<'a> {
381    fn index_mut(&mut self, index: ParamIndex) -> &mut Self::Output {
382        &mut self.params[index.to_usize()]
383    }
384}