ivila commented on issue #153:
URL:
https://github.com/apache/incubator-teaclave-trustzone-sdk/issues/153#issuecomment-2519256939
> After trying some different ways, I can now introduce a new way, which
keeps compatibility and simplicity, and can prevent developers from misuse
(while rest of my other ways don't).
>
> ## Solution
> First, we introduce an enum, call Buffer, it could has its own buffer or
borrow from others.
>
> ```rust
> pub enum Buffer<'a> {
> Own(Vec<u8>),
> Borrow(&'a mut [u8]),
> }
>
> impl<'a> Buffer<'a> {
> // from immutable slice, create a new buffer instead
> pub fn new_from_parameter(data: &[u8], buffer_size: usize) -> Self {
> assert!(buffer_size >= data.len());
> let mut buffer = vec![0_u8; buffer_size];
> buffer[..data.len()].copy_from_slice(&data);
> Self::Own(buffer)
> }
> // from mutable slice, use that directly
> pub fn new_from_share_buffer(data: &'a mut [u8]) -> Self {
> Self::Borrow(data)
> }
> // get buffer for the parameters of invoke function
> pub fn buffer(&mut self) -> &mut [u8] {
> match self {
> Self::Own(buf) => buf.as_mut_slice(),
> Self::Borrow(buf) => buf,
> }
> }
> // into the result of invoke function
> pub fn into_sized_vector(self, size: usize) -> Vec<u8> {
> match self {
> Self::Own(mut buf) => {
> buf.resize(size, 0);
> buf
> }
> Self::Borrow(buf) => {
> let mut buffer = vec![0_u8; size];
> buffer[..core::cmp::min(size,
buf.len())].copy_from_slice(&buf);
> buffer
> }
> }
> }
> }
> ```
>
> Then, we introduce a wire Trait, we use this to **make our invoke function
flexible**
>
> ```rust
> pub trait Wire<'a, OutputT> {
> // convert the function input to buffer
> fn into_buffer(self) -> Buffer<'a>;
> // construct the result
> fn make_result(buffer: Buffer<'a>, outlen: usize) -> OutputT;
> }
> ```
>
> With both of them, the invoke method can rewrite to:
>
> ```rust
> impl Plugin {
> pub fn invoke<'a, OutputT, InputT: Wire<'a, OutputT>>(
> &mut self,
> command_id: u32,
> subcommand_id: u32,
> data: InputT,
> ) -> Result<OutputT> {
> let raw_uuid: Uuid = self.uuid;
> let mut outlen: usize = 0;
> let mut data = data.into_buffer(); // transfer data into buffer
> let buffer = data.buffer();
> match unsafe {
> raw::tee_invoke_supp_plugin(
> raw_uuid.as_raw_ptr(),
> command_id as u32,
> subcommand_id as u32,
> buffer.as_mut_ptr(),
> buffer.len(),
> &mut outlen as *mut usize,
> )
> } {
> raw::TEE_SUCCESS => {
> assert!(outlen <= buffer.len());
> // construct result by the type of data
> Ok(InputT::make_result(data, outlen))
> }
> code => Err(Error::from_raw_error(code)),
> }
> }
> }
> ```
>
> Finally, we can imply Wire trait for the (InputT, OutputT) pair we want:
>
> ```rust
> // use generic type pair (AsRef<[u8]>, Vec<u8>) to keeps compatibility
> impl<'a, T: AsRef<[u8]>> Wire<'a, Vec<u8>> for T {
> fn into_buffer(self) -> Buffer<'a> {
> let data = self.as_ref().into();
> Buffer::new_from_parameter(data, data.len())
> }
> fn make_result(buffer: Buffer<'a>, outlen: usize) -> Vec<u8> {
> buffer.into_sized_vector(outlen)
> }
> }
>
> // use exact type pair to avoid confliction and keep developers from
misuse.
> impl<'a> Wire<'a, usize> for &'a mut [u8] {
> fn into_buffer(self) -> Buffer<'a> {
> Buffer::Borrow(self)
> }
> fn make_result(_: Buffer<'a>, outlen: usize) -> usize {
> outlen
> }
> }
> ```
>
> And now, we can call the `invoke` function as we want, here is some
examples that may or may not compile:
>
> ```rust
> let mut plugin = Plugin {
> uuid:
Uuid::parse_str("7dd54ee6-a705-4e4d-8b6b-aa5024dfcd10").unwrap(),
> };
> let immutable_buffer = [0_u8; 32];
> let mut mutable_buffer = [0_u8; 32];
> // =================== old feature, pass immutable slice
==================
> // response1 infer to Vec<u8> automatically, keeps compatibility.
> let response1 = plugin.invoke(1, 1, &immutable_buffer).unwrap();
> // this won't compile, as we didn't imply the Wire trait for the
type pair
> // and I think codes like this must be mistake from developers,
developers
> // must get the buffer back
> let response2: usize = plugin.invoke(1, 1,
&immutable_buffer).unwrap();
>
> // =================== new feature, pass mutable slice
====================
> // developers must specific what they want, keeping developers
from misuse
> // currently the valid return types are: Vec<u8> or usize
> let response3: Vec<u8> = plugin.invoke(1, 1, &mut
mutable_buffer).unwrap();
> let response4: Vec<u8> = plugin.invoke(1, 1,
mutable_buffer.as_mut_slice()).unwrap();
> // this won't compile, as we need developers to specific the input
type to &mut [u8]
> let response5: usize = plugin.invoke(1, 1, &mut
mutable_buffer).unwrap();
> // this is working like std::io::Read
> let response5: usize = plugin.invoke(1, 1,
mutable_buffer.as_mut_slice()).unwrap();
> ```
>
> ## Full Codes Demo
> Here's the full codes demo.
>
> 1. cargo.toml
>
> ```toml
> [package]
> name = "plugin"
> version = "0.1.0"
> edition = "2024"
>
> [dependencies]
> optee-utee = { git =
"https://github.com/apache/incubator-teaclave-trustzone-sdk.git";, branch =
"main", version = "0.2.0", features = ["no_panic_handler"] }
> optee-utee-sys = { git =
"https://github.com/apache/incubator-teaclave-trustzone-sdk.git";, branch =
"main", version = "0.2.0" }
>
> [dev-dependencies]
> rand = "0.8.5"
> once_cell = "1.20.2"
> ```
>
> 2. lib.rs
>
> ```rust
> extern crate alloc;
> use alloc::vec;
> use alloc::vec::Vec;
> use optee_utee::{Error, Result, Uuid};
> use optee_utee_sys as raw;
>
> pub enum Buffer<'a> {
> Own(Vec<u8>),
> Borrow(&'a mut [u8]),
> }
>
> pub trait Wire<'a, OutputT> {
> // convert the function input to buffer
> fn into_buffer(self) -> Buffer<'a>;
> // construct the result
> fn make_result(buffer: Buffer<'a>, outlen: usize) -> OutputT;
> }
>
> pub struct Plugin {
> uuid: Uuid,
> }
>
> impl Plugin {
> pub fn invoke<'a, OutputT, InputT: Wire<'a, OutputT>>(
> &mut self,
> command_id: u32,
> subcommand_id: u32,
> data: InputT,
> ) -> Result<OutputT> {
> let raw_uuid: Uuid = self.uuid;
> let mut outlen: usize = 0;
> let mut data = data.into_buffer();
> let buffer = data.buffer();
> match unsafe {
> raw::tee_invoke_supp_plugin(
> raw_uuid.as_raw_ptr(),
> command_id as u32,
> subcommand_id as u32,
> buffer.as_mut_ptr(),
> buffer.len(),
> &mut outlen as *mut usize,
> )
> } {
> raw::TEE_SUCCESS => {
> assert!(outlen <= buffer.len());
> // construct result by the type of data
> Ok(InputT::make_result(data, outlen))
> }
> code => Err(Error::from_raw_error(code)),
> }
> }
> }
>
> impl<'a> Buffer<'a> {
> // from immutable slice, create a new buffer instead
> pub fn new_from_parameter(data: &[u8], buffer_size: usize) -> Self {
> assert!(buffer_size >= data.len());
> let mut buffer = vec![0_u8; buffer_size];
> buffer[..data.len()].copy_from_slice(&data);
> Self::Own(buffer)
> }
> // from mutable slice, use that directly
> pub fn new_from_buffer(data: &'a mut [u8]) -> Self {
> Self::Borrow(data)
> }
> // get buffer for the parameters of invoke function
> pub fn buffer(&mut self) -> &mut [u8] {
> match self {
> Self::Own(buf) => buf.as_mut_slice(),
> Self::Borrow(buf) => buf,
> }
> }
> // into the result of invoke function
> pub fn into_sized_vector(self, size: usize) -> Vec<u8> {
> match self {
> Self::Own(mut buf) => {
> buf.resize(size, 0);
> buf
> }
> Self::Borrow(buf) => {
> let mut buffer = vec![0_u8; size];
> buffer[..core::cmp::min(size,
buf.len())].copy_from_slice(&buf);
> buffer
> }
> }
> }
> }
>
> // use generic type pair (AsRef<[u8]>, Vec<u8>) to keeps compatibility
> impl<'a, T: AsRef<[u8]>> Wire<'a, Vec<u8>> for T {
> fn into_buffer(self) -> Buffer<'a> {
> let data = self.as_ref().into();
> Buffer::new_from_parameter(data, data.len())
> }
> fn make_result(buffer: Buffer<'a>, outlen: usize) -> Vec<u8> {
> buffer.into_sized_vector(outlen)
> }
> }
>
> // use exact type pair to avoid confliction and keep us from misuse.
> impl<'a> Wire<'a, usize> for &'a mut [u8] {
> fn into_buffer(self) -> Buffer<'a> {
> Buffer::Borrow(self)
> }
> fn make_result(_: Buffer<'a>, outlen: usize) -> usize {
> outlen
> }
> }
>
> #[cfg(test)]
> pub mod tests {
> extern crate std;
> use super::*;
> use core::ffi::c_char;
> use once_cell::sync::Lazy;
> use optee_utee_sys::{TEE_Result, TEE_UUID};
> use rand::Rng;
> use rand::distributions::Alphanumeric;
> use std::collections::HashMap;
> use std::string::String;
> use std::sync::RwLock;
>
> static REE_RETURN_VALUES: RwLock<Lazy<HashMap<u32, String>>> =
> RwLock::new(Lazy::new(|| HashMap::new()));
>
> fn set_value_for_test(sub_cmd: u32, value: String) {
> let mut values = REE_RETURN_VALUES.write().unwrap();
> assert!(!values.contains_key(&sub_cmd));
> values.insert(sub_cmd, value);
> }
>
> fn get_test_value(sub_cmd: u32) -> String {
> let values = REE_RETURN_VALUES.read().unwrap();
> values.get(&sub_cmd).unwrap().to_owned()
> }
>
> fn generate_test_pair(length: usize) -> (u32, String) {
> let sub_cmd: u32 = rand::thread_rng().r#gen();
> let random_string: String = rand::thread_rng()
> .sample_iter(&Alphanumeric)
> .take(length)
> .map(char::from)
> .collect();
> set_value_for_test(sub_cmd, random_string.clone());
> (sub_cmd, random_string)
> }
>
> #[unsafe(no_mangle)]
> pub extern "C" fn tee_invoke_supp_plugin(
> _uuid: *const TEE_UUID,
> _cmd: u32,
> sub_cmd: u32,
> buf: *mut c_char,
> len: usize,
> outlen: *mut usize,
> ) -> TEE_Result {
> let inbuf = unsafe { core::slice::from_raw_parts_mut(buf, len) };
> std::println!(
> "*plugin*: receive value: {:?} length {:?}",
> inbuf,
> inbuf.len()
> );
> let test_value = get_test_value(sub_cmd);
> assert!(test_value.len() <= len);
> std::println!("*plugin*: write value '{}' to buffer", test_value);
>
> inbuf[0..test_value.len()].copy_from_slice(test_value.as_bytes());
> unsafe {
> *outlen = test_value.len();
> }
> return raw::TEE_SUCCESS;
> }
>
> #[test]
> fn test_invoke_with_immutable_array_and_vec_result() {
> let mut plugin = Plugin {
> uuid:
Uuid::parse_str("7dd54ee6-a705-4e4d-8b6b-aa5024dfcd10").unwrap(),
> };
> const BUFFER_LEN: usize = 32;
>
> // test calling with &[u8; N], and output less than that
> let parameter = [0_u8; BUFFER_LEN];
> let (sub_cmd, random) = generate_test_pair(BUFFER_LEN / 2);
> let response = plugin.invoke(1, sub_cmd, ¶meter).unwrap();
> std::println!("*TA*: response is {:?}", response);
> assert_eq!(response.as_slice(), random.as_bytes());
> // test calling with &[u8; N], and output equals to that
> let parameter = [0_u8; BUFFER_LEN];
> let (sub_cmd, random) = generate_test_pair(BUFFER_LEN);
> let response = plugin.invoke(1, sub_cmd, ¶meter).unwrap();
> std::println!("*TA*: response is {:?}", response);
> assert_eq!(response.as_slice(), random.as_bytes());
> // test calling with &[u8], and output less than that
> let parameter = [0_u8; BUFFER_LEN];
> let (sub_cmd, random) = generate_test_pair(BUFFER_LEN / 2);
> let response = plugin.invoke(1, sub_cmd,
parameter.as_slice()).unwrap();
> std::println!("*TA*: response is {:?}", response);
> assert_eq!(response.as_slice(), random.as_bytes());
> // test calling with &[u8], and output equals to that
> let parameter = [0_u8; BUFFER_LEN];
> let (sub_cmd, random) = generate_test_pair(BUFFER_LEN);
> let response = plugin.invoke(1, sub_cmd,
parameter.as_slice()).unwrap();
> std::println!("*TA*: response is {:?}", response);
> assert_eq!(response.as_slice(), random.as_bytes());
> // test calling with &Vec, and output less than that
> let parameter = vec![0_u8; BUFFER_LEN];
> let (sub_cmd, random) = generate_test_pair(BUFFER_LEN / 2);
> let response = plugin.invoke(1, sub_cmd, ¶meter).unwrap();
> std::println!("*TA*: response is {:?}", response);
> assert_eq!(response.as_slice(), random.as_bytes());
> // test calling with &Vec, and output equals to that
> let parameter = vec![0_u8; BUFFER_LEN];
> let (sub_cmd, random) = generate_test_pair(BUFFER_LEN);
> let response = plugin.invoke(1, sub_cmd, ¶meter).unwrap();
> std::println!("*TA*: response is {:?}", response);
> assert_eq!(response.as_slice(), random.as_bytes());
> }
> #[test]
> fn test_invoke_with_mutable_array_and_length_result() {
> let mut plugin = Plugin {
> uuid:
Uuid::parse_str("7dd54ee6-a705-4e4d-8b6b-aa5024dfcd10").unwrap(),
> };
> const BUFFER_LEN: usize = 32;
>
> // test calling with &mut [u8; N], and output less than that
> let mut buffer = [0_u8; 32];
> let (sub_cmd, random) = generate_test_pair(BUFFER_LEN / 2);
> let response: usize = plugin.invoke(1, sub_cmd,
buffer.as_mut_slice()).unwrap();
> std::println!("*TA*: response is {:?}", response);
> assert_eq!(response, random.len());
> assert_eq!(&buffer[..random.len()], random.as_bytes());
> assert!(buffer[random.len()..].into_iter().all(|&byte| byte == 0));
> // test calling with &mut [u8; N], and output equals to that
> let mut buffer = [0_u8; 32];
> let (sub_cmd, random) = generate_test_pair(BUFFER_LEN);
> let response: usize = plugin.invoke(1, sub_cmd,
buffer.as_mut_slice()).unwrap();
> std::println!("*TA*: response is {:?}", response);
> assert_eq!(response, random.len());
> assert_eq!(&buffer, random.as_bytes());
> // test calling with &mut [u8], and output less than that
> let mut buffer = [0_u8; 32];
> let (sub_cmd, random) = generate_test_pair(BUFFER_LEN / 2);
> let response: usize = plugin.invoke(1, sub_cmd,
buffer.as_mut_slice()).unwrap();
> std::println!("*TA*: response is {:?}", response);
> assert_eq!(response, random.len());
> assert_eq!(&buffer[..random.len()], random.as_bytes());
> assert!(buffer[random.len()..].into_iter().all(|&byte| byte == 0));
> // test calling with &mut [u8], and output equals to that
> let mut buffer = [0_u8; 32];
> let (sub_cmd, random) = generate_test_pair(BUFFER_LEN);
> let response: usize = plugin.invoke(1, sub_cmd,
buffer.as_mut_slice()).unwrap();
> std::println!("*TA*: response is {:?}", response);
> assert_eq!(response, random.len());
> assert_eq!(&buffer, random.as_bytes());
> }
> }
> ```
>
> You can test it by `SYS_BUILD_TYPE=unit_test cargo test -- --show-output`
@DemesneGH Got some time to have a look?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: dev-unsubscr...@teaclave.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: dev-unsubscr...@teaclave.apache.org
For additional commands, e-mail: dev-h...@teaclave.apache.org
