ivila commented on issue #153:
URL:
https://github.com/apache/incubator-teaclave-trustzone-sdk/issues/153#issuecomment-2516754958
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`
--
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
