Direct data copy between normal-world shared memory and TA memory

litebox_common_optee/src/lib.rs

@@ -582,12 +582,25 @@ impl UteeParams {
 #[derive(Clone)]
 pub enum UteeParamOwned {
     None,
-    ValueInput { value_a: u64, value_b: u64 },
+    ValueInput {
+        value_a: u64,
+        value_b: u64,
+    },
     ValueOutput,
-    ValueInout { value_a: u64, value_b: u64 },
-    MemrefInput { data: Box<[u8]> },
-    MemrefOutput { buffer_size: usize },
-    MemrefInout { data: Box<[u8]>, buffer_size: usize },
+    ValueInout {
+        value_a: u64,
+        value_b: u64,
+    },
+    MemrefInput {
+        data: Option<Box<[u8]>>,
+    },
+    MemrefOutput {
+        buffer_size: usize,
+    },
+    MemrefInout {
+        data: Option<Box<[u8]>>,
+        buffer_size: usize,
+    },
 }
 
 impl UteeParamOwned {

litebox_runner_lvbs/src/lib.rs

@@ -43,7 +43,7 @@ use litebox_shim_optee::msg_handler::{
     decode_ta_request, handle_optee_msg_args, handle_optee_smc_args, update_optee_msg_args,
 };
 use litebox_shim_optee::session::{OpenSessionTarget, TaInstance, session_manager};
-use litebox_shim_optee::{NormalWorldConstPtr, NormalWorldMutPtr, UserConstPtr};
+use litebox_shim_optee::{NormalWorldConstPtr, NormalWorldMutPtr, TaMemrefAddresses, UserConstPtr};
 
 /// Seed the initial heap regions so the global allocator has enough memory
 /// for slab-backed allocations (the slab needs >= 2 MB backing pages).
@@ -645,13 +645,14 @@ fn open_session_single_instance(
     let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?;
 
     // Load TA context with parameters for OpenSession - pass actual session_id
-    instance
+    let memref_addresses = instance
         .loaded_program()
         .entrypoints
         .as_ref()
         .ok_or(OpteeSmcReturnCode::EBadCmd)?
-        .load_ta_context(
+        .load_ta_context_with_shm(
             params,
+            &ta_req_info.shm_info,
             runner_session_id,
             UteeEntryFunc::OpenSession as u32,
             None,
@@ -698,6 +699,7 @@ fn open_session_single_instance(
             None, // No session ID on failure
             Some(&ta_params),
             Some(ta_req_info),
+            Some(&memref_addresses),
         );
 
         // For single-instance TAs, only clean up on TARGET_DEAD (panic).
@@ -730,6 +732,7 @@ fn open_session_single_instance(
         Some(runner_session_id),
         Some(&ta_params),
         Some(ta_req_info),
+        Some(&memref_addresses),
     );
 
     // Write-back failure: OpenSession succeeded inside the TA, but we cannot
@@ -849,6 +852,7 @@ fn open_session_new_instance(
             None, // No session ID on failure
             None,
             Some(ta_req_info),
+            None,
         );
 
         // Safety: We are about to tear down this TA instance;
@@ -869,12 +873,13 @@ fn open_session_new_instance(
         unsafe { teardown_ta_page_table(&shim, task_pt_id) };
         OpteeSmcReturnCode::EBadCmd
     })?;
-    loaded_program
+    let memref_addresses = loaded_program
         .entrypoints
         .as_ref()
         .unwrap()
-        .load_ta_context(
+        .load_ta_context_with_shm(
             params,
+            &ta_req_info.shm_info,
             runner_session_id,
             UteeEntryFunc::OpenSession as u32,
             None,
@@ -931,6 +936,7 @@ fn open_session_new_instance(
             None, // No session ID on failure
             Some(&ta_params),
             Some(ta_req_info),
+            Some(&memref_addresses),
         );
 
         // Safety: We are about to tear down this TA instance;
@@ -951,6 +957,7 @@ fn open_session_new_instance(
         Some(runner_session_id),
         Some(&ta_params),
         Some(ta_req_info),
+        Some(&memref_addresses),
     )
     .inspect_err(|_| {
         // Safety: We are about to tear down this TA instance;
@@ -1041,9 +1048,10 @@ fn handle_invoke_command(
 
         // Set up the entry-point parameters for InvokeCommand.
         let entrypoints_ref = instance.loaded_program().entrypoints.as_ref().unwrap();
-        entrypoints_ref
-            .load_ta_context(
+        let memref_addresses = entrypoints_ref
+            .load_ta_context_with_shm(
                 params.as_slice(),
+                &ta_req_info.shm_info,
                 session_id,
                 UteeEntryFunc::InvokeCommand as u32,
                 Some(cmd_id),
@@ -1080,6 +1088,7 @@ fn handle_invoke_command(
             None,
             Some(&ta_params),
             Some(&ta_req_info),
+            Some(&memref_addresses),
         );
 
         // Per OP-TEE OS: if TA panics (TARGET_DEAD), the TA context is
@@ -1179,6 +1188,7 @@ fn handle_close_session(
             None,
             None,
             None,
+            None,
         );
 
         let removed_flags = session_manager().unregister_session(session_id);
@@ -1251,6 +1261,7 @@ fn write_msg_args_to_normal_world(
     session_id: Option<u32>,
     ta_params: Option<&UteeParams>,
     ta_req_info: Option<&litebox_shim_optee::msg_handler::TaRequestInfo<PAGE_SIZE>>,
+    memref_addresses: Option<&TaMemrefAddresses>,
 ) -> Result<(), OpteeSmcReturnCode> {
     // Ensure we're on a task page table, not the base page table.
     // Accessing TA userspace memory requires the TA's page table to be active.
@@ -1273,6 +1284,7 @@ fn write_msg_args_to_normal_world(
         session_id,
         ta_params,
         ta_req_info,
+        memref_addresses,
         msg_args,
     )?;
 

litebox_runner_optee_on_linux_userland/src/tests.rs

@@ -311,7 +311,7 @@ impl TaCommandParamsBase64 {
                 value_b: *value_b,
             },
             TaCommandParamsBase64::MemrefInput { data_base64 } => UteeParamOwned::MemrefInput {
-                data: Self::decode_base64(data_base64).into_boxed_slice(),
+                data: Some(Self::decode_base64(data_base64).into_boxed_slice()),
             },
             TaCommandParamsBase64::MemrefOutput { buffer_size } => UteeParamOwned::MemrefOutput {
                 buffer_size: usize::try_from(*buffer_size).unwrap(),
@@ -327,7 +327,7 @@ impl TaCommandParamsBase64 {
                     "Buffer size is smaller than input data size"
                 );
                 UteeParamOwned::MemrefInout {
-                    data: decoded_data.into_boxed_slice(),
+                    data: Some(decoded_data.into_boxed_slice()),
                     buffer_size,
                 }
             }

litebox_shim_optee/src/lib.rs

@@ -236,6 +236,7 @@ impl GlobalState {
 
 type UserMutPtr<T> = <Platform as litebox::platform::RawPointerProvider>::RawMutPointer<T>;
 pub type UserConstPtr<T> = <Platform as litebox::platform::RawPointerProvider>::RawConstPointer<T>;
+pub type TaMemrefAddresses = [Option<usize>; litebox_common_optee::UteeParams::TEE_NUM_PARAMS];
 
 type MutPtr<T> = <Platform as litebox::platform::RawPointerProvider>::RawMutPointer<T>;
 
@@ -341,11 +342,30 @@ impl OpteeShimEntrypoints {
         func_id: u32,
         cmd_id: Option<u32>,
     ) -> Result<(), loader::elf::ElfLoaderError> {
+        self.load_ta_context_with_shm(params, &[], session_id, func_id, cmd_id)
+            .map(|_| ())
+    }
+
+    /// Load the TA context with shared-memory sources for its input buffers.
+    pub fn load_ta_context_with_shm(
+        &self,
+        params: &[litebox_common_optee::UteeParamOwned],
+        shm_info: &[Option<msg_handler::ShmInfo<PAGE_SIZE>>],
+        session_id: u32,
+        func_id: u32,
+        cmd_id: Option<u32>,
+    ) -> Result<TaMemrefAddresses, loader::elf::ElfLoaderError> {
         let init_state = self
             .task
-            .load_ta_context(params, session_id, func_id, cmd_id)?;
+            .load_ta_context(params, shm_info, session_id, func_id, cmd_id)?;
+        let ThreadInitState::Ta {
+            memref_addresses, ..
+        } = init_state
+        else {
+            return Err(loader::elf::ElfLoaderError::InvalidStackAddr);
+        };
         self.task.thread.init_state.set(init_state);
-        Ok(())
+        Ok(memref_addresses)
     }
 }
 
@@ -658,6 +678,7 @@ impl Task {
                 func_id,
                 entry_point,
                 stack_top,
+                ..
             } => {
                 #[cfg(target_arch = "x86_64")]
                 {
@@ -792,6 +813,7 @@ impl Task {
     fn load_ta_context(
         &self,
         params: &[litebox_common_optee::UteeParamOwned],
+        shm_info: &[Option<msg_handler::ShmInfo<PAGE_SIZE>>],
         session_id: u32,
         func_id: u32,
         cmd_id: Option<u32>,
@@ -817,8 +839,8 @@ impl Task {
             crate::loader::ta_stack::allocate_stack(self, self.get_ta_stack_base_addr()).ok_or(
                 ElfLoaderError::MappingError(litebox::mm::linux::MappingError::OutOfMemory),
             )?;
-        ta_stack
-            .init(self.global.platform, params)
+        let memref_addresses = ta_stack
+            .init(self.global.platform, params, shm_info)
             .ok_or(ElfLoaderError::InvalidStackAddr)?;
 
         Ok(ThreadInitState::Ta {
@@ -828,6 +850,7 @@ impl Task {
             func_id: func_id as usize,
             entry_point: self.get_ta_entry_point(),
             stack_top: ta_stack.get_cur_stack_top(),
+            memref_addresses,
         })
     }
 
@@ -1447,6 +1470,7 @@ pub(crate) enum ThreadInitState {
         func_id: usize,
         entry_point: usize,
         stack_top: usize,
+        memref_addresses: TaMemrefAddresses,
     },
 }
 

litebox_shim_optee/src/loader/ta_stack.rs

@@ -10,7 +10,7 @@ use litebox::{
 use litebox_common_optee::{LdelfArg, TeeParamType, UteeParamOwned, UteeParams};
 use zerocopy::IntoBytes;
 
-use crate::{Platform, UserMutPtr};
+use crate::{Platform, TaMemrefAddresses, UserMutPtr, msg_handler::ShmInfo};
 
 #[inline]
 fn align_down(addr: usize, align: usize) -> usize {
@@ -185,34 +185,41 @@ impl TaStack {
         param_type: TeeParamType,
         bytes: Option<&[u8]>,
         len: usize,
-    ) -> Option<()> {
+    ) -> Option<usize> {
         if self.num_params >= UteeParams::TEE_NUM_PARAMS {
             return None;
         }
         match param_type {
             TeeParamType::MemrefInput | TeeParamType::MemrefInout => {
-                let bytes = bytes?;
+                let bytes = bytes.unwrap_or(&[]);
                 if len > bytes.len() {
                     self.pos = self.pos.checked_sub(len - bytes.len())?;
                 }
                 self.push_bytes(bytes)?;
-                self.params
-                    .set_values(self.num_params, self.get_cur_stack_top() as u64, len as u64)
-                    .ok()?;
-            }
-            TeeParamType::MemrefOutput => {
-                self.pos = self.pos.checked_sub(len)?;
-                self.params
-                    .set_values(self.num_params, self.get_cur_stack_top() as u64, len as u64)
-                    .ok()?;
-            }
-            _ => {
-                return None;
             }
+            TeeParamType::MemrefOutput => self.pos = self.pos.checked_sub(len)?,
+            _ => return None,
         }
+        let address = self.get_cur_stack_top();
+        self.params
+            .set_values(self.num_params, address as u64, len as u64)
+            .ok()?;
         self.params.set_type(self.num_params, param_type).ok()?;
         self.num_params += 1;
-        Some(())
+        Some(address)
+    }
+
+    fn push_param_memref_from_shm<const ALIGN: usize>(
+        &mut self,
+        param_type: TeeParamType,
+        shm_info: &ShmInfo<ALIGN>,
+        len: usize,
+    ) -> Option<usize> {
+        let address = self.push_param_memref(param_type, None, len)?;
+        shm_info
+            .copy_to_user(UserMutPtr::from_usize(address), len)
+            .ok()?;
+        Some(address)
     }
 
     /// Set `UteeParams` on the stack.
@@ -223,14 +230,20 @@ impl TaStack {
         Some(())
     }
 
-    pub(crate) fn init(&mut self, platform: &Platform, params: &[UteeParamOwned]) -> Option<()> {
+    pub(crate) fn init<const ALIGN: usize>(
+        &mut self,
+        platform: &Platform,
+        params: &[UteeParamOwned],
+        shm_info: &[Option<ShmInfo<ALIGN>>],
+    ) -> Option<TaMemrefAddresses> {
         if params.len() > UteeParams::TEE_NUM_PARAMS {
             return None;
         }
 
         self.scrub()?;
 
-        for param in params {
+        let mut memref_addresses = [None; UteeParams::TEE_NUM_PARAMS];
+        for (index, param) in params.iter().enumerate() {
             match param {
                 UteeParamOwned::ValueInput { value_a, value_b } => {
                     self.push_param_values(TeeParamType::ValueInput, Some((*value_a, *value_b)))?;
@@ -242,13 +255,37 @@ impl TaStack {
                     self.push_param_values(TeeParamType::ValueInout, Some((*value_a, *value_b)))?;
                 }
                 UteeParamOwned::MemrefInput { data } => {
-                    self.push_param_memref(TeeParamType::MemrefInput, Some(data), data.len())?;
+                    if let Some(shm_info) = shm_info.get(index).and_then(Option::as_ref) {
+                        let len = shm_info.len();
+                        self.push_param_memref_from_shm(TeeParamType::MemrefInput, shm_info, len)?;
+                    } else {
+                        let data = data.as_deref()?;
+                        self.push_param_memref(TeeParamType::MemrefInput, Some(data), data.len())?;
+                    }
                 }
                 UteeParamOwned::MemrefInout { data, buffer_size } => {
-                    self.push_param_memref(TeeParamType::MemrefInout, Some(data), *buffer_size)?;
+                    let address =
+                        if let Some(shm_info) = shm_info.get(index).and_then(Option::as_ref) {
+                            self.push_param_memref_from_shm(
+                                TeeParamType::MemrefInout,
+                                shm_info,
+                                *buffer_size,
+                            )?
+                        } else {
+                            self.push_param_memref(
+                                TeeParamType::MemrefInout,
+                                Some(data.as_deref()?),
+                                *buffer_size,
+                            )?
+                        };
+                    memref_addresses[index] = Some(address);
                 }
                 UteeParamOwned::MemrefOutput { buffer_size } => {
-                    self.push_param_memref(TeeParamType::MemrefOutput, None, *buffer_size)?;
+                    memref_addresses[index] = Some(self.push_param_memref(
+                        TeeParamType::MemrefOutput,
+                        None,
+                        *buffer_size,
+                    )?);
                 }
                 UteeParamOwned::None => self.push_param_none()?,
             }
@@ -269,7 +306,7 @@ impl TaStack {
             self.pos % Self::STACK_ALIGNMENT,
             core::mem::size_of::<usize>()
         );
-        Some(())
+        Some(memref_addresses)
     }
 
     pub(crate) fn init_with_ldelf_arg(&mut self, ldelf_arg: &LdelfArg) -> Option<()> {