optee-utee: support setting capacity of shared buffer of LoadablePlugin

1) introduce a Command struct for flexible use
2) allow developers settings capacity of shared buffer when invoking
   command

Signed-off-by: ivila <[email protected]>
Reviewed-by: Yuan Zhuang <[email protected]>

optee-utee/Cargo.toml

@@ -30,7 +30,13 @@ optee-utee-macros = { version = "0.2.0", path = "macros" }
 bitflags = "1.0.4"
 uuid = { version = "0.8", default-features = false }
 hex = { version = "0.4", default-features = false, features = ["alloc"] }
-libc_alloc = "=1.0.5"
+libc_alloc = "1.0.5"
+
+[dev-dependencies]
+rand = "0.8.5"
+once_cell = "1.20.2"
+serde = { version = "1.0.215" }
+serde_json = { version = "1.0.133" }
 
 [features]
 no_panic_handler = []

optee-utee/src/extension.rs

@@ -26,32 +26,304 @@ pub struct LoadablePlugin {
     uuid: Uuid
 }
 
+pub struct LoadablePluginCommand<'a> {
+    plugin: &'a LoadablePlugin,
+    cmd_id: u32,
+    sub_cmd_id: u32,
+    buffer: Vec<u8>,
+}
+
 impl LoadablePlugin {
     pub fn new(uuid: &Uuid) -> Self {
         Self { uuid: uuid.to_owned() }
     }
-    pub fn invoke(&mut self, command_id: u32, subcommand_id: u32, data: &[u8]) -> Result<Vec<u8>> {
-        let raw_uuid: Uuid = self.uuid;
+    pub fn invoke(&self, command_id: u32, subcommand_id: u32, data: &[u8]) -> Result<Vec<u8>> {
+        self.new_command_with_capacity(command_id, data.len())
+            .sub_cmd(subcommand_id)
+            .chain_write_body(data)
+            .call()
+    }
+    pub fn invoke_with_capacity<'a>(
+        &'a self,
+        command_id: u32,
+        subcommand_id: u32,
+        capacity: usize,
+    ) -> LoadablePluginCommand<'a> {
+        self.new_command_with_capacity(command_id, capacity)
+            .sub_cmd(subcommand_id)
+    }
+    // for people who wants to set the buffer size for response manually
+    pub fn new_command_with_capacity<'a>(
+        &'a self,
+        command_id: u32,
+        capacity: usize,
+    ) -> LoadablePluginCommand<'a> {
+        LoadablePluginCommand::new_with_capacity(self, command_id, capacity)
+    }
+}
+
+impl<'a> LoadablePluginCommand<'a> {
+    // use this to set sub_cmd if needed
+    pub fn sub_cmd(mut self, sub_cmd_id: u32) -> Self {
+        self.sub_cmd_id = sub_cmd_id;
+        self
+    }
+    // use this to write request body if needed
+    pub fn write_body(&mut self, data: &[u8]) {
+        self.buffer.extend_from_slice(data);
+    }
+    // same with write_body, but chainable
+    pub fn chain_write_body(mut self, data: &[u8]) -> Self {
+        self.write_body(data);
+        self
+    }
+    // invoke the command, and get result from it
+    pub fn call(self) -> Result<Vec<u8>> {
         let mut outlen: usize = 0;
+        let mut buffer = self.buffer;
+        buffer.resize(buffer.capacity(), 0); // resize to capacity first
         match unsafe {
             raw::tee_invoke_supp_plugin(
-                raw_uuid.as_raw_ptr(),
-                command_id as u32,
-                subcommand_id as u32,
-                data.as_ptr() as _,
-                data.len(),
+                self.plugin.uuid.as_raw_ptr(),
+                self.cmd_id,
+                self.sub_cmd_id,
+                // convert the pointer manually, as in some platform c_char is i8
+                buffer.as_mut_slice().as_mut_ptr() as *mut _,
+                buffer.len(),
                 &mut outlen as *mut usize,
             )
         } {
             raw::TEE_SUCCESS => {
-                assert!(outlen <= (data.len()));
-                let mut outbuf = vec![0; outlen];
-                outbuf.copy_from_slice(&data[..outlen]);
-
-                Ok(outbuf)
-            },
+                assert!(outlen <= buffer.len());
+                buffer.resize(outlen, 0);
+                Ok(buffer)
+            }
             code => Err(Error::from_raw_error(code)),
         }
-
+    }
+}
+
+impl<'a> LoadablePluginCommand<'a> {
+    fn new_with_capacity(plugin: &'a LoadablePlugin, cmd_id: u32, capacity: usize) -> Self {
+        Self {
+            plugin,
+            cmd_id,
+            sub_cmd_id: 0,
+            buffer: Vec::with_capacity(capacity),
+        }
+    }
+}
+
+#[cfg(test)]
+pub mod test_loadable_plugin {
+    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::distributions::Alphanumeric;
+    use rand::Rng;
+    use std::collections::HashMap;
+    use std::sync::RwLock;
+
+    static REE_RETURN_VALUES: RwLock<Lazy<HashMap<(u32, u32), Vec<u8>>>> =
+        RwLock::new(Lazy::new(|| HashMap::new()));
+    static REE_EXPECTED_VALUES: RwLock<Lazy<HashMap<(u32, u32), Vec<u8>>>> =
+        RwLock::new(Lazy::new(|| HashMap::new()));
+
+    fn set_ree_return_value(cmd: u32, sub_cmd: u32, value: Vec<u8>) {
+        let mut values = REE_RETURN_VALUES.write().unwrap();
+        let key = (cmd, sub_cmd);
+        assert!(!values.contains_key(&key));
+        values.insert(key, value);
+    }
+
+    fn set_ree_expected_value(cmd: u32, sub_cmd: u32, value: Vec<u8>) {
+        let mut values = REE_EXPECTED_VALUES.write().unwrap();
+        let key = (cmd, sub_cmd);
+        assert!(!values.contains_key(&key));
+        values.insert(key, value);
+    }
+
+    fn get_ree_return_value(cmd: u32, sub_cmd: u32) -> Vec<u8> {
+        let values = REE_RETURN_VALUES.read().unwrap();
+        let key = (cmd, sub_cmd);
+        values.get(&key).unwrap().to_owned()
+    }
+
+    fn get_ree_expected_value(cmd: u32, sub_cmd: u32) -> Vec<u8> {
+        let values = REE_EXPECTED_VALUES.read().unwrap();
+        let key = (cmd, sub_cmd);
+        values.get(&key).unwrap().to_owned()
+    }
+
+    fn generate_random_bytes(len: usize) -> Vec<u8> {
+        rand::thread_rng()
+            .sample_iter(&Alphanumeric)
+            .take(len)
+            .collect()
+    }
+
+    fn generate_test_pairs(
+        request_size: usize,
+        response_size: usize,
+    ) -> (u32, u32, Vec<u8>, Vec<u8>) {
+        let cmd: u32 = rand::thread_rng().r#gen();
+        let sub_cmd: u32 = rand::thread_rng().r#gen();
+        let random_request: Vec<u8> = generate_random_bytes(request_size);
+        let random_response: Vec<u8> = generate_random_bytes(response_size);
+        (cmd, sub_cmd, random_request, random_response)
+    }
+
+    #[no_mangle]
+    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 {
+        // must convert buf to u8, for in some platform c_char was treated as i8
+        let inbuf = unsafe { core::slice::from_raw_parts_mut(buf as *mut u8, len) };
+        std::println!(
+            "*plugin*: receive value: {:?} length {:?}",
+            inbuf,
+            inbuf.len()
+        );
+        let expected_value = get_ree_expected_value(cmd, sub_cmd);
+        assert_eq!(inbuf, expected_value.as_slice());
+
+        let return_value = get_ree_return_value(cmd, sub_cmd);
+        assert!(return_value.len() <= len);
+        std::println!("*plugin*: write value '{:?}' to buffer", return_value);
+
+        inbuf[0..return_value.len()].copy_from_slice(&return_value);
+        unsafe {
+            *outlen = return_value.len();
+        }
+        return raw::TEE_SUCCESS;
+    }
+
+    #[test]
+    fn test_invoke() {
+        let plugin = LoadablePlugin {
+            uuid: Uuid::parse_str("7dd54ee6-a705-4e4d-8b6b-aa5024dfcd10").unwrap(),
+        };
+        const REQUEST_LEN: usize = 32;
+
+        // test calling with output size less than input
+        let (cmd, sub_cmd, request, exp_response) =
+            generate_test_pairs(REQUEST_LEN, REQUEST_LEN / 2);
+        set_ree_expected_value(cmd, sub_cmd, request.clone());
+        set_ree_return_value(cmd, sub_cmd, exp_response.clone());
+        let response = plugin.invoke(cmd, sub_cmd, &request).unwrap();
+        std::println!("*TA*: response is {:?}", response);
+        assert_eq!(response, exp_response);
+
+        // test calling with output size equals to input
+        let (cmd, sub_cmd, request, exp_response) = generate_test_pairs(REQUEST_LEN, REQUEST_LEN);
+        set_ree_expected_value(cmd, sub_cmd, request.clone());
+        set_ree_return_value(cmd, sub_cmd, exp_response.clone());
+        let response = plugin.invoke(cmd, sub_cmd, &request).unwrap();
+        std::println!("*TA*: response is {:?}", response);
+        assert_eq!(response, exp_response);
+    }
+
+    #[test]
+    fn test_invoke_with_capacity() {
+        let plugin = LoadablePlugin {
+            uuid: Uuid::parse_str("7dd54ee6-a705-4e4d-8b6b-aa5024dfcd10").unwrap(),
+        };
+        const RESPONSE_LEN: usize = 32;
+
+        // test calling with output size less than input
+        let (cmd, sub_cmd, request, exp_response) =
+            generate_test_pairs(2 * RESPONSE_LEN, RESPONSE_LEN);
+        set_ree_expected_value(cmd, sub_cmd, request.clone());
+        set_ree_return_value(cmd, sub_cmd, exp_response.clone());
+        let response = plugin
+            .invoke_with_capacity(cmd, sub_cmd, exp_response.len())
+            .chain_write_body(&request)
+            .call()
+            .unwrap();
+        std::println!("*TA*: response is {:?}", response);
+        assert_eq!(response, exp_response);
+
+        // test calling with output size equals to input
+        let (cmd, sub_cmd, request, exp_response) = generate_test_pairs(RESPONSE_LEN, RESPONSE_LEN);
+        set_ree_expected_value(cmd, sub_cmd, request.clone());
+        set_ree_return_value(cmd, sub_cmd, exp_response.clone());
+        let response = plugin
+            .invoke_with_capacity(cmd, sub_cmd, exp_response.len())
+            .chain_write_body(&request)
+            .call()
+            .unwrap();
+        std::println!("*TA*: response is {:?}", response);
+        assert_eq!(response, exp_response);
+
+        // test calling with output size greater than input
+        let (cmd, sub_cmd, mut request, exp_response) =
+            generate_test_pairs(RESPONSE_LEN / 2, RESPONSE_LEN);
+        request.resize(exp_response.len(), 0);
+        set_ree_expected_value(cmd, sub_cmd, request.clone());
+        set_ree_return_value(cmd, sub_cmd, exp_response.clone());
+        let response = plugin
+            .invoke_with_capacity(cmd, sub_cmd, exp_response.len())
+            .chain_write_body(&request)
+            .call()
+            .unwrap();
+        std::println!("*TA*: response is {:?}", response);
+        assert_eq!(response, exp_response);
+    }
+    #[test]
+    fn test_invoke_with_writer() {
+        let plugin = LoadablePlugin {
+            uuid: Uuid::parse_str("7dd54ee6-a705-4e4d-8b6b-aa5024dfcd10").unwrap(),
+        };
+        // impl a writer for Command
+        struct Wrapper<'a, 'b>(&'b mut LoadablePluginCommand<'a>);
+        impl<'a, 'b> std::io::Write for Wrapper<'a, 'b> {
+            fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+                self.0.write_body(buf);
+                Ok(buf.len())
+            }
+            fn flush(&mut self) -> std::io::Result<()> {
+                Ok(())
+            }
+        }
+        // serialize data into command directly
+        let test_data = serde_json::json!({
+            "code": 100,
+            "message": "error"
+        });
+        let mut exp_request = serde_json::to_vec(&test_data).unwrap();
+        let buffer_len = exp_request.len() * 2;
+        let (cmd, sub_cmd, _, exp_response) = generate_test_pairs(0, buffer_len);
+        let mut plugin_cmd = plugin.invoke_with_capacity(cmd, sub_cmd, buffer_len);
+        exp_request.resize(exp_response.len(), 0);
+        set_ree_expected_value(cmd, sub_cmd, exp_request);
+        set_ree_return_value(cmd, sub_cmd, exp_response.clone());
+        serde_json::to_writer(Wrapper(&mut plugin_cmd), &test_data).unwrap();
+        let response = plugin_cmd.call().unwrap();
+        std::println!("*TA*: response is {:?}", response);
+        assert_eq!(response, exp_response);
+    }
+    #[test]
+    fn test_invoke_with_no_data() {
+        let plugin = LoadablePlugin {
+            uuid: Uuid::parse_str("7dd54ee6-a705-4e4d-8b6b-aa5024dfcd10").unwrap(),
+        };
+        const OUTPUT_LEN: usize = 50;
+        let (cmd, sub_cmd, _, exp_response) = generate_test_pairs(0, OUTPUT_LEN);
+        let exp_request = vec![0_u8; OUTPUT_LEN];
+        set_ree_expected_value(cmd, sub_cmd, exp_request);
+        set_ree_return_value(cmd, sub_cmd, exp_response.clone());
+        let response = plugin
+            .invoke_with_capacity(cmd, sub_cmd, OUTPUT_LEN)
+            .call()
+            .unwrap();
+        std::println!("*TA*: response is {:?}", response);
+        assert_eq!(response, exp_response);
     }
 }

optee-utee/systest/Cargo.toml

@@ -26,7 +26,7 @@ edition = "2018"
 
 [dependencies]
 optee-utee-sys = { path = "../optee-utee-sys" }
-libc = "=0.2.59"
+libc = "0.2.59"
 
 [build-dependencies]
 ctest = "0.2"