From 812d6aecbebd42e937654224f2ed51a089c9e232 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Thu, 10 Sep 2026 04:17:01 +0000 Subject: [PATCH 1/7] optee: Add Dynamic TA RPC protocol helpers Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu --- litebox_common_optee/src/lib.rs | 116 ++++++++++++++++++++++++++++++-- 1 file changed, 111 insertions(+), 5 deletions(-) diff --git a/litebox_common_optee/src/lib.rs b/litebox_common_optee/src/lib.rs index 976e5ab26c..4878cd5267 100644 --- a/litebox_common_optee/src/lib.rs +++ b/litebox_common_optee/src/lib.rs @@ -692,6 +692,15 @@ impl TeeUuid { Self::from_bytes(bytes) } + #[allow(clippy::missing_panics_doc)] + pub fn to_u64_array(self) -> [u64; 2] { + let bytes = self.to_bytes(); + [ + u64::from_le_bytes(bytes[0..8].try_into().unwrap()), + u64::from_le_bytes(bytes[8..16].try_into().unwrap()), + ] + } + /// Converts the UUID to a 16-byte array with little-endian encoding. pub fn to_le_bytes(self) -> [u8; 16] { let mut bytes = [0u8; 16]; @@ -701,6 +710,16 @@ impl TeeUuid { bytes[8..16].copy_from_slice(&self.clock_seq_and_node); bytes } + + /// Converts the UUID to a 16-byte array with big-endian encoding (RFC 4122 format). + pub fn to_bytes(self) -> [u8; 16] { + let mut bytes = [0u8; 16]; + bytes[0..4].copy_from_slice(&self.time_low.to_be_bytes()); + bytes[4..6].copy_from_slice(&self.time_mid.to_be_bytes()); + bytes[6..8].copy_from_slice(&self.time_hi_and_version.to_be_bytes()); + bytes[8..16].copy_from_slice(&self.clock_seq_and_node); + bytes + } } /// TA flags from `optee_os/lib/libutee/include/user_ta_header.h`. @@ -1556,6 +1575,7 @@ const OPTEE_MSG_ATTR_TYPE_TMEM_INOUT: u8 = 0xb; /// Meta-parameter marker of the attribute word. Set on the `OpenSession` /// TA-UUID and client-identity params. const OPTEE_MSG_ATTR_META: u64 = 1 << 8; +const OPTEE_MSG_ATTR_NONCONTIG: u64 = 1 << 9; #[non_exhaustive] #[derive(Debug, PartialEq, TryFromPrimitive)] @@ -1604,7 +1624,7 @@ impl OpteeMsgAttr { /// Returns `true` when the noncontig bit (bit 9) is set. pub fn noncontig(&self) -> bool { - self.0 & (1 << 9) != 0 + self.0 & OPTEE_MSG_ATTR_NONCONTIG != 0 } } @@ -2116,6 +2136,45 @@ impl OpteeRpcArgs { } } + /// Access a TMEM output parameter with exact direction and flag validation. + /// + /// The NONCONTIG flag is permitted because an SHM_ALLOC response may return + /// either contiguous memory or an OP-TEE page-list descriptor. + pub fn get_param_tmem_output( + &self, + index: usize, + ) -> Result { + if index >= self.num_params as usize { + return Err(OpteeSmcReturnCode::ENotAvail); + } + + let param = &self.params[index]; + if param.attr.attr_type() != OpteeMsgAttrType::TmemOutput as u8 + || param.attr.0 & !(u64::from(u8::MAX) | OPTEE_MSG_ATTR_NONCONTIG) != 0 + { + return Err(OpteeSmcReturnCode::EBadCmd); + } + OpteeMsgParamTmem::read_from_bytes(¶m.data).map_err(|_| OpteeSmcReturnCode::EBadCmd) + } + + /// Access an RMEM output parameter with exact direction and flag validation. + pub fn get_param_rmem_output( + &self, + index: usize, + ) -> Result { + if index >= self.num_params as usize { + return Err(OpteeSmcReturnCode::ENotAvail); + } + + let param = &self.params[index]; + if param.attr.attr_type() != OpteeMsgAttrType::RmemOutput as u8 + || param.attr.0 & !u64::from(u8::MAX) != 0 + { + return Err(OpteeSmcReturnCode::EBadCmd); + } + OpteeMsgParamRmem::read_from_bytes(¶m.data).map_err(|_| OpteeSmcReturnCode::EBadCmd) + } + /// Set a value parameter by index with bounds checking against `num_params`. pub fn set_param_value( &mut self, @@ -2130,6 +2189,34 @@ impl OpteeRpcArgs { } } + /// Set a parameter's attribute type by index with bounds checking against `num_params`. + pub fn set_param_attr_type( + &mut self, + index: usize, + attr_type: OpteeMsgAttrType, + ) -> Result<(), OpteeSmcReturnCode> { + if index >= self.num_params as usize { + Err(OpteeSmcReturnCode::ENotAvail) + } else { + self.params[index].attr = OpteeMsgAttr(attr_type as u64); + Ok(()) + } + } + + /// Set an rmem parameter by index with bounds checking against `num_params`. + pub fn set_param_rmem( + &mut self, + index: usize, + rmem: OpteeMsgParamRmem, + ) -> Result<(), OpteeSmcReturnCode> { + if index >= self.num_params as usize { + Err(OpteeSmcReturnCode::ENotAvail) + } else { + self.params[index].data.copy_from_slice(rmem.as_bytes()); + Ok(()) + } + } + /// Set a tmem parameter by index with bounds checking against `num_params`. pub fn set_param_tmem( &mut self, @@ -2143,10 +2230,6 @@ impl OpteeRpcArgs { Ok(()) } } - - // Note: RPC does not use rmem params. Rmem requires pre-registered shared memory - // references from the normal-world driver, which is a main-messaging-path concept. - // RPC uses tmem for buffer references since OP-TEE provides physical addresses directly. } /// Serialize the params portion as raw bytes into `buf`. @@ -2530,6 +2613,29 @@ mod tests { uuid.clock_seq_and_node, [0xaf, 0x63, 0x00, 0x02, 0xa5, 0xd5, 0xc5, 0x1b] ); + assert_eq!( + uuid.to_u64_array(), + [0xe311f8e7_e0b34f38, 0x1bc5d5a5_020063af] + ); + assert_eq!(TeeUuid::from_u64_array(uuid.to_u64_array()), uuid); + } + + #[test] + fn test_tee_uuid_to_bytes() { + let uuid = TeeUuid { + time_low: 0x384f_b3e0, + time_mid: 0xe7f8, + time_hi_and_version: 0x11e3, + clock_seq_and_node: [0xaf, 0x63, 0x00, 0x02, 0xa5, 0xd5, 0xc5, 0x1b], + }; + + assert_eq!( + uuid.to_bytes(), + [ + 0x38, 0x4f, 0xb3, 0xe0, 0xe7, 0xf8, 0x11, 0xe3, 0xaf, 0x63, 0x00, 0x02, 0xa5, 0xd5, + 0xc5, 0x1b, + ] + ); } #[test] From 06f009049a75260cf1983788792296287aa5dfed Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Thu, 10 Sep 2026 04:21:21 +0000 Subject: [PATCH 2/7] optee: Add trusted context tracking for Dynamic TA RPCs Track trusted continuation state across the multi-call Dynamic TA loading sequence. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu --- dev_tests/src/ratchet.rs | 2 +- litebox_common_optee/src/lib.rs | 20 ++ litebox_shim_optee/src/lib.rs | 1 + litebox_shim_optee/src/rpc_context.rs | 425 ++++++++++++++++++++++++++ 4 files changed, 447 insertions(+), 1 deletion(-) create mode 100644 litebox_shim_optee/src/rpc_context.rs diff --git a/dev_tests/src/ratchet.rs b/dev_tests/src/ratchet.rs index 41dc17bb9e..5c6c76f4cb 100644 --- a/dev_tests/src/ratchet.rs +++ b/dev_tests/src/ratchet.rs @@ -42,7 +42,7 @@ fn ratchet_globals() -> Result<()> { ("litebox_runner_lvbs/", 8), ("litebox_runner_snp/", 2), ("litebox_shim_linux/", 1), - ("litebox_shim_optee/", 6), + ("litebox_shim_optee/", 7), ], |file| { Ok(file diff --git a/litebox_common_optee/src/lib.rs b/litebox_common_optee/src/lib.rs index 4878cd5267..a08dedf53c 100644 --- a/litebox_common_optee/src/lib.rs +++ b/litebox_common_optee/src/lib.rs @@ -2312,6 +2312,18 @@ impl OpteeSmcArgs { } } + /// Set the context ID used to identify an RPC call in the preserved `args[3]` register. + pub fn set_rpc_context_id(&mut self, context_id: u32) { + self.args[3] = context_id as usize; + } + + /// Get the context ID used to identify an RPC call from the preserved `args[3]` register. + pub fn get_rpc_context_id(&self) -> Result { + self.args[3] + .try_into() + .map_err(|_| OpteeSmcReturnCode::EBadCmd) + } + /// Set the return code of an OP-TEE SMC call pub fn set_return_code(&mut self, code: OpteeSmcReturnCode) { self.args[0] = code as usize; @@ -2582,6 +2594,14 @@ pub const HUK_SUBKEY_MAX_LEN: usize = 32; mod tests { use super::*; + #[cfg(target_pointer_width = "64")] + #[test] + fn test_rpc_context_id_rejects_upper_bits() { + let mut args = OpteeSmcArgs::default(); + args.args[3] = (u32::MAX as usize) + 1; + assert_eq!(args.get_rpc_context_id(), Err(OpteeSmcReturnCode::EBadCmd)); + } + #[test] fn test_optee_msg_args_header_size_and_layout() { use core::mem::{offset_of, size_of}; diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index 4d96ccd24b..c345967f56 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -30,6 +30,7 @@ use litebox_common_optee::{ }; pub mod loader; +pub mod rpc_context; pub mod session; pub(crate) mod syscalls; diff --git a/litebox_shim_optee/src/rpc_context.rs b/litebox_shim_optee/src/rpc_context.rs new file mode 100644 index 0000000000..54c7fde83b --- /dev/null +++ b/litebox_shim_optee/src/rpc_context.rs @@ -0,0 +1,425 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! RPC context tracking for multi-call OP-TEE operations. +//! +//! # Dynamic TA loading +//! +//! OP-TEE loads a Dynamic TA from the normal world with a sequence of RPCs. +//! The reference flow is implemented by `rpc_load()` in +//! `optee_os/core/kernel/ree_fs_ta.c`; the RPC transport and shared-memory +//! allocation are implemented by `thread_rpc_cmd()` and +//! `thread_rpc_alloc_payload()` in +//! `optee_os/core/arch/arm/kernel/thread_optee_smc.c`. +//! +//! LiteBox follows the same high-level protocol across the VTL boundary: +//! +//! ```text +//! VTL1 (LiteBox OP-TEE shim) VTL0 (driver / supplicant) +//! | | +//! |-- LOAD_TA(UUID, empty output TMEM) ---->| +//! |<------- TA size in TMEM.size ------------| +//! | | +//! |-- SHM_ALLOC(application, size, align) -->| +//! |<-- TMEM { buf_ptr, size, shm_ref } ------| +//! | | +//! | Register the allocation by shm_ref | +//! | | +//! |-- LOAD_TA(UUID, output RMEM) ------------>| +//! |<------ TA binary written to RMEM ---------| +//! | | +//! | Read, validate, and copy the TA | +//! | | +//! |-- SHM_FREE(application, shm_ref) ------->| +//! |<-------------- completion ---------------| +//! ``` +//! +//! The first `LOAD_TA` discovers the required binary size. `SHM_ALLOC` then +//! returns a temporary-memory reference containing the physical buffer address, +//! allocated size, and an opaque shared-memory reference. LiteBox records that +//! allocation and sends the second `LOAD_TA` as an RMEM referring to the same +//! `shm_ref`; the normal-world driver resolves it before asking the supplicant +//! to fill the buffer with the TA binary. +//! +//! # Why explicit contexts are needed +//! +//! OP-TEE OS executes this sequence on a secure-world thread. `thread_rpc()` +//! suspends that thread while normal world handles an RPC, preserving the +//! `rpc_load()` call stack, local variables, RPC arguments, and memory-object +//! references. Normal world returns the thread ID in register `a3`, allowing +//! `OPTEE_SMC_CALL_RETURN_FROM_RPC` to resume the suspended continuation. The +//! Dynamic TA stage is therefore implicit in the saved thread execution state; +//! OP-TEE does not need a separate protocol-stage enum. +//! +//! LiteBox has no equivalent resumable OP-TEE thread and call stack. Instead, +//! [`RpcContextMap`] associates the context ID carried in `args[3]` with trusted +//! continuation state. [`RpcContext`] records which RPC response is expected +//! and carries only the state valid for that stage. Stage-checked transitions +//! prevent a response from being interpreted as a different step of the +//! protocol. +//! +//! LiteBox copies the loaded binary into trusted memory, then releases the VTL0 +//! allocation with a `SHM_FREE` RPC tracked by [`RpcContext::ShmFree`]. The +//! trusted cached binary is dropped after ldelf loads it into TA runtime memory. + +use alloc::boxed::Box; +use hashbrown::HashMap; +use litebox::utils::id_pool::IdPool; +use litebox_common_optee::{OpteeSmcReturnCode, TeeUuid}; +use once_cell::race::OnceBox; +use spin::mutex::SpinMutex; + +const MAX_RPC_CONTEXTS: u32 = 1024; + +/// An RPC context map operation failed. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RpcContextError { + Full, + NotFound, + UnexpectedStage, +} + +/// Action to take after an in-flight shared-memory free RPC returns. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum RpcCompletion { + OpenSession, + ReturnError(OpteeSmcReturnCode), +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct RpcCommon { + ta_uuid: TeeUuid, + // RPC continuation reuses args[3] for the context ID. Use these fields to + // preserve the original registered SHM reference and offset + // before overwriting it. + registered_shm_ref: u64, + regd_shm_offset: usize, +} + +/// Trusted continuation state for an RPC-backed Dynamic TA request. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum RpcContext { + LoadTaSize { + common: RpcCommon, + }, + ShmAlloc { + common: RpcCommon, + requested_size: u64, + }, + LoadTaBinary { + common: RpcCommon, + requested_size: u64, + shm_ref: u64, + }, + ShmFree { + common: RpcCommon, + shm_ref: u64, + completion: RpcCompletion, + }, +} + +impl RpcContext { + fn new(ta_uuid: TeeUuid, registered_shm_ref: u64, regd_shm_offset: usize) -> Self { + Self::LoadTaSize { + common: RpcCommon { + ta_uuid, + registered_shm_ref, + regd_shm_offset, + }, + } + } + + fn common(&self) -> RpcCommon { + match self { + Self::LoadTaSize { common } + | Self::ShmAlloc { common, .. } + | Self::LoadTaBinary { common, .. } + | Self::ShmFree { common, .. } => *common, + } + } + + pub fn ta_uuid(&self) -> TeeUuid { + self.common().ta_uuid + } + + pub fn registered_shm_ref(&self) -> u64 { + self.common().registered_shm_ref + } + + pub fn registered_shm_offset(&self) -> usize { + self.common().regd_shm_offset + } + + pub fn requested_size(&self) -> Option { + match self { + Self::ShmAlloc { requested_size, .. } | Self::LoadTaBinary { requested_size, .. } => { + Some(*requested_size) + } + Self::LoadTaSize { .. } | Self::ShmFree { .. } => None, + } + } + + pub fn shm_ref(&self) -> Option { + match self { + Self::LoadTaBinary { shm_ref, .. } | Self::ShmFree { shm_ref, .. } => Some(*shm_ref), + Self::LoadTaSize { .. } | Self::ShmAlloc { .. } => None, + } + } + + pub fn completion(&self) -> Option { + match self { + Self::ShmFree { completion, .. } => Some(*completion), + Self::LoadTaSize { .. } | Self::ShmAlloc { .. } | Self::LoadTaBinary { .. } => None, + } + } + + fn into_shm_alloc(self, requested_size: u64) -> Result { + match self { + Self::LoadTaSize { common } => Ok(Self::ShmAlloc { + common, + requested_size, + }), + _ => Err(RpcContextError::UnexpectedStage), + } + } + + fn into_load_ta_binary(self, shm_ref: u64) -> Result { + match self { + Self::ShmAlloc { + common, + requested_size, + } => Ok(Self::LoadTaBinary { + common, + requested_size, + shm_ref, + }), + _ => Err(RpcContextError::UnexpectedStage), + } + } + + fn into_shm_free( + self, + shm_ref: u64, + completion: RpcCompletion, + ) -> Result { + match self { + Self::ShmAlloc { common, .. } => Ok(Self::ShmFree { + common, + shm_ref, + completion, + }), + Self::LoadTaBinary { + common, + shm_ref: expected_shm_ref, + .. + } if shm_ref == expected_shm_ref => Ok(Self::ShmFree { + common, + shm_ref, + completion, + }), + _ => Err(RpcContextError::UnexpectedStage), + } + } +} + +struct RpcContexts { + ids: IdPool, + contexts: HashMap, +} + +impl RpcContexts { + fn new() -> Self { + Self { + ids: IdPool::with_capacity(MAX_RPC_CONTEXTS), + contexts: HashMap::new(), + } + } +} + +/// Maps RPC context IDs to trusted continuation state. +pub struct RpcContextMap { + inner: SpinMutex, +} + +impl RpcContextMap { + /// Create an empty RPC context map. + pub fn new() -> Self { + Self { + inner: SpinMutex::new(RpcContexts::new()), + } + } + + /// Allocate a context for the first `LOAD_TA` response. + pub fn allocate( + &self, + ta_uuid: TeeUuid, + registered_shm_ref: u64, + regd_shm_offset: usize, + ) -> Result { + let mut inner = self.inner.lock(); + let context_id = inner.ids.allocate().ok_or(RpcContextError::Full)?; + inner.contexts.insert( + context_id, + RpcContext::new(ta_uuid, registered_shm_ref, regd_shm_offset), + ); + Ok(context_id) + } + + /// Return a snapshot of a context. + pub fn get(&self, context_id: u32) -> Option { + self.inner.lock().contexts.get(&context_id).copied() + } + + pub fn transition_to_shm_alloc( + &self, + context_id: u32, + requested_size: u64, + ) -> Result<(), RpcContextError> { + self.update(context_id, |context| context.into_shm_alloc(requested_size)) + } + + pub fn transition_to_load_ta_binary( + &self, + context_id: u32, + shm_ref: u64, + ) -> Result<(), RpcContextError> { + self.update(context_id, |context| context.into_load_ta_binary(shm_ref)) + } + + pub fn transition_to_shm_free( + &self, + context_id: u32, + shm_ref: u64, + completion: RpcCompletion, + ) -> Result<(), RpcContextError> { + self.update(context_id, |context| { + context.into_shm_free(shm_ref, completion) + }) + } + + fn update( + &self, + context_id: u32, + transition: impl FnOnce(RpcContext) -> Result, + ) -> Result<(), RpcContextError> { + let mut inner = self.inner.lock(); + let context = inner + .contexts + .get(&context_id) + .copied() + .ok_or(RpcContextError::NotFound)?; + inner.contexts.insert(context_id, transition(context)?); + Ok(()) + } + + /// Remove a context and recycle its ID. + pub fn take(&self, context_id: u32) -> Option { + let mut inner = self.inner.lock(); + let context = inner.contexts.remove(&context_id)?; + inner.ids.recycle(context_id); + Some(context) + } +} + +impl Default for RpcContextMap { + fn default() -> Self { + Self::new() + } +} + +/// Return the global RPC context map. +pub fn rpc_context_map() -> &'static RpcContextMap { + static RPC_CONTEXT_MAP: OnceBox = OnceBox::new(); + RPC_CONTEXT_MAP.get_or_init(|| Box::new(RpcContextMap::new())) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn test_uuid(value: u32) -> TeeUuid { + TeeUuid { + time_low: value, + time_mid: 0, + time_hi_and_version: 0, + clock_seq_and_node: [0; 8], + } + } + + #[test] + fn tracks_dynamic_ta_rpc_lifecycle() { + let contexts = RpcContextMap::new(); + let context_id = contexts.allocate(test_uuid(1), 0x10, 0x100).unwrap(); + + let initial = contexts.get(context_id).unwrap(); + assert!(matches!(initial, RpcContext::LoadTaSize { .. })); + assert_eq!(initial.ta_uuid(), test_uuid(1)); + assert_eq!(initial.registered_shm_ref(), 0x10); + assert_eq!(initial.registered_shm_offset(), 0x100); + + contexts + .transition_to_shm_alloc(context_id, 0x4000) + .unwrap(); + assert!(matches!( + contexts.get(context_id), + Some(RpcContext::ShmAlloc { + requested_size: 0x4000, + .. + }) + )); + + contexts + .transition_to_load_ta_binary(context_id, 0x1234) + .unwrap(); + assert!(matches!( + contexts.get(context_id), + Some(RpcContext::LoadTaBinary { + requested_size: 0x4000, + shm_ref: 0x1234, + .. + }) + )); + + let completion = RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd); + contexts + .transition_to_shm_free(context_id, 0x1234, completion) + .unwrap(); + assert_eq!( + contexts.take(context_id), + Some(RpcContext::ShmFree { + common: initial.common(), + shm_ref: 0x1234, + completion, + }) + ); + } + + #[test] + fn rejects_incomplete_and_replayed_transitions() { + let contexts = RpcContextMap::new(); + let context_id = contexts.allocate(test_uuid(1), 1, 0).unwrap(); + + contexts.transition_to_shm_alloc(context_id, 1).unwrap(); + assert_eq!( + contexts.transition_to_shm_alloc(context_id, 2), + Err(RpcContextError::UnexpectedStage) + ); + contexts + .transition_to_load_ta_binary(context_id, 1) + .unwrap(); + assert_eq!( + contexts.transition_to_load_ta_binary(context_id, 2), + Err(RpcContextError::UnexpectedStage) + ); + assert_eq!( + contexts.transition_to_shm_free( + context_id, + 2, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ), + Err(RpcContextError::UnexpectedStage) + ); + assert!(contexts.take(context_id).is_some()); + assert_eq!(contexts.take(context_id), None); + } +} From bd9e004819485b07bd6c0fbd6ac84d1358811a77 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Thu, 10 Sep 2026 17:59:06 +0000 Subject: [PATCH 3/7] optee: Write RPC arguments through registered shared memory Replace direct physical-address writes with bounded writes through registered shared-memory bookkeeping. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu --- litebox_runner_lvbs/src/lib.rs | 36 ++------------------- litebox_shim_optee/src/msg_handler.rs | 46 +++++++++++++++++++++++++++ litebox_shim_optee/src/rpc_context.rs | 23 -------------- 3 files changed, 48 insertions(+), 57 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 442eba6ebc..75c709e760 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -15,8 +15,8 @@ use litebox::{ use litebox_common_linux::errno::Errno; use litebox_common_lvbs::{NUM_VTLCALL_PARAMS, VsmError, VsmFunction}; use litebox_common_optee::{ - OpteeMessageCommand, OpteeMsgArgs, OpteeRpcArgs, OpteeSmcArgs, OpteeSmcResult, - OpteeSmcReturnCode, TeeOrigin, TeeResult, UteeEntryFunc, UteeParams, optee_msg_args_total_size, + OpteeMessageCommand, OpteeMsgArgs, OpteeSmcArgs, OpteeSmcResult, OpteeSmcReturnCode, TeeOrigin, + TeeResult, UteeEntryFunc, UteeParams, optee_msg_args_total_size, }; use litebox_platform_lvbs::host::LvbsLinuxKernel as Platform; use litebox_platform_lvbs::mshv::vsm::{LvbsVtl0Gate, LvbsVtl0PrivilegedWriter, LvbsVtl1Gate}; @@ -1377,38 +1377,6 @@ fn write_non_ta_msg_args_to_normal_world( Ok(()) } -/// Write `OpteeRpcArgs` to the normal world. Its write address is determined by -/// `msg_args_phys_addr` and the size of `OpteeMsgArgs`. -/// -/// Unlike [`write_msg_args_to_normal_world`], this function does not access TA userspace -/// memory and can be called from the base page table context. It simply serializes the -/// rpc_args and writes it to the normal world physical address. -#[expect(dead_code)] -#[inline] -fn write_rpc_args_to_normal_world( - platform: &'static Platform, - msg_args: &OpteeMsgArgs, - msg_args_phys_addr: u64, - rpc_args: &OpteeRpcArgs, -) -> Result<(), OpteeSmcReturnCode> { - let msg_args_size = optee_msg_args_total_size(msg_args.num_params); - - let rpc_args_size = optee_msg_args_total_size(rpc_args.num_params); - let mut blob = vec![0u8; rpc_args_size]; - rpc_args.serialize(&mut blob)?; - - let rpc_pa: usize = >::trunc(msg_args_phys_addr) - .checked_add(msg_args_size) - .ok_or(OpteeSmcReturnCode::EBadAddr)?; // RPC args are placed right after the main msg_args blob - let ptr = NormalWorldMutPtr::::with_contiguous_pages( - platform, - rpc_pa, - rpc_args_size, - )?; - ptr.write_slice_at_offset(0, &blob)?; - Ok(()) -} - // use include_bytes! to include ldelf const LDELF_BINARY: &[u8] = &[0u8; 0]; const TA_BINARY: &[u8] = &[0u8; 0]; diff --git a/litebox_shim_optee/src/msg_handler.rs b/litebox_shim_optee/src/msg_handler.rs index c368099e41..6ab2810ffa 100644 --- a/litebox_shim_optee/src/msg_handler.rs +++ b/litebox_shim_optee/src/msg_handler.rs @@ -758,6 +758,32 @@ impl ShmInfo { Ok(()) } + /// Write `buffer` to the normal-world shared memory pages referenced by `self`, + /// starting at byte `offset` within the view. + fn write_at>( + &self, + platform: &Platform, + offset: usize, + buffer: &[u8], + ) -> Result<(), OpteeSmcReturnCode> { + if offset + .checked_add(buffer.len()) + .is_none_or(|end| end > self.len) + { + return Err(OpteeSmcReturnCode::EBadAddr); + } + if buffer.is_empty() { + return Ok(()); + } + let ptr = NormalWorldMutPtr::::new( + platform, + &self.page_addrs, + self.page_offset, + )?; + ptr.write_slice_at_offset(offset, buffer)?; + Ok(()) + } + /// Copy from this normal-world shared memory into TA userspace. pub(crate) fn copy_to_user( &self, @@ -935,6 +961,26 @@ impl ShmRefMap { } } +/// Serialize RPC arguments immediately after the main message in registered shared memory. +pub fn write_rpc_args_to_regd_shm>( + platform: &Platform, + shm_ref: u64, + msg_args_offset: usize, + msg_args_num_params: u32, + rpc_args: &OpteeRpcArgs, +) -> Result<(), OpteeSmcReturnCode> { + let shm_info = shm_ref_map() + .get(shm_ref) + .ok_or(OpteeSmcReturnCode::EBadAddr)?; + let rpc_args_offset = msg_args_offset + .checked_add(optee_msg_args_total_size(msg_args_num_params)) + .ok_or(OpteeSmcReturnCode::EBadAddr)?; + let rpc_args_size = optee_msg_args_total_size(rpc_args.num_params); + let mut blob = alloc::vec![0u8; rpc_args_size]; + rpc_args.serialize(&mut blob)?; + shm_info.write_at(platform, rpc_args_offset, &blob) +} + fn shm_ref_map() -> &'static ShmRefMap { static SHM_REF_MAP: OnceBox> = OnceBox::new(); SHM_REF_MAP.get_or_init(|| Box::new(ShmRefMap::new())) diff --git a/litebox_shim_optee/src/rpc_context.rs b/litebox_shim_optee/src/rpc_context.rs index 54c7fde83b..ca52b5a66d 100644 --- a/litebox_shim_optee/src/rpc_context.rs +++ b/litebox_shim_optee/src/rpc_context.rs @@ -150,29 +150,6 @@ impl RpcContext { self.common().regd_shm_offset } - pub fn requested_size(&self) -> Option { - match self { - Self::ShmAlloc { requested_size, .. } | Self::LoadTaBinary { requested_size, .. } => { - Some(*requested_size) - } - Self::LoadTaSize { .. } | Self::ShmFree { .. } => None, - } - } - - pub fn shm_ref(&self) -> Option { - match self { - Self::LoadTaBinary { shm_ref, .. } | Self::ShmFree { shm_ref, .. } => Some(*shm_ref), - Self::LoadTaSize { .. } | Self::ShmAlloc { .. } => None, - } - } - - pub fn completion(&self) -> Option { - match self { - Self::ShmFree { completion, .. } => Some(*completion), - Self::LoadTaSize { .. } | Self::ShmAlloc { .. } | Self::LoadTaBinary { .. } => None, - } - } - fn into_shm_alloc(self, requested_size: u64) -> Result { match self { Self::LoadTaSize { common } => Ok(Self::ShmAlloc { From 0927f0da73ba102ed1e7306b2ba0e59813e176d5 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Thu, 10 Sep 2026 04:30:56 +0000 Subject: [PATCH 4/7] optee: Initiate LOAD_TA if TA is not cached in VTL1 Initiate the first LOAD_TA if the TA is not present in secure-world cache. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu --- litebox_common_optee/src/lib.rs | 46 +++++++++ litebox_runner_lvbs/src/lib.rs | 124 +++++++++++++++++++---- litebox_shim_optee/src/lib.rs | 33 +++--- litebox_shim_optee/src/syscalls/ldelf.rs | 2 +- 4 files changed, 171 insertions(+), 34 deletions(-) diff --git a/litebox_common_optee/src/lib.rs b/litebox_common_optee/src/lib.rs index a08dedf53c..56f7278ce1 100644 --- a/litebox_common_optee/src/lib.rs +++ b/litebox_common_optee/src/lib.rs @@ -2230,6 +2230,52 @@ impl OpteeRpcArgs { Ok(()) } } + + /// Prepare a LOAD_TA RPC request to be sent to normal world. + pub fn prepare_load_ta_rpc( + &mut self, + ta_uuid: TeeUuid, + memref: Option, + ) -> Result<(), OpteeSmcReturnCode> { + self.cmd = OpteeRpcCommand::LoadTa; + // Match OP-TEE's get_rpc_arg(): default to failure in case normal world + // returns without updating the RPC result. + self.ret = TeeResult::GenericError; + self.num_params = 2; + + self.set_param_attr_type(0, OpteeMsgAttrType::ValueInput)?; + let uuid_bytes = ta_uuid.to_u64_array(); + self.set_param_value( + 0, + OpteeMsgParamValue { + a: uuid_bytes[0], + b: uuid_bytes[1], + c: 0, + }, + )?; + + match memref { + None => { + // First LOAD_TA call: normal world returns the TA size in `tmem.size`. + self.set_param_attr_type(1, OpteeMsgAttrType::TmemOutput)?; + self.set_param_tmem( + 1, + OpteeMsgParamTmem { + buf_ptr: 0, + size: 0, + shm_ref: 0, + }, + )?; + } + Some(rmem) => { + // Second LOAD_TA call: normal world populates VTL0-owned memory. + self.set_param_attr_type(1, OpteeMsgAttrType::RmemOutput)?; + self.set_param_rmem(1, rmem)?; + } + } + + Ok(()) + } } /// Serialize the params portion as raw bytes into `buf`. diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 75c709e760..1bb2bfde21 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -15,19 +15,22 @@ use litebox::{ use litebox_common_linux::errno::Errno; use litebox_common_lvbs::{NUM_VTLCALL_PARAMS, VsmError, VsmFunction}; use litebox_common_optee::{ - OpteeMessageCommand, OpteeMsgArgs, OpteeSmcArgs, OpteeSmcResult, OpteeSmcReturnCode, TeeOrigin, - TeeResult, UteeEntryFunc, UteeParams, optee_msg_args_total_size, + OpteeMessageCommand, OpteeMsgArgs, OpteeRpcArgs, OpteeSmcArgs, OpteeSmcFunction, + OpteeSmcResult, OpteeSmcReturnCode, TeeOrigin, TeeResult, UteeEntryFunc, UteeParams, + optee_msg_args_total_size, }; use litebox_platform_lvbs::host::LvbsLinuxKernel as Platform; use litebox_platform_lvbs::mshv::vsm::{LvbsVtl0Gate, LvbsVtl0PrivilegedWriter, LvbsVtl1Gate}; use litebox_platform_lvbs::{ - arch::{gdt, instrs::hlt_loop, interrupts, timer}, + arch::instrs::hlt_loop, mshv::vsm_intercept::raise_vtl0_gp_fault, serial_println, +}; +use litebox_platform_lvbs::{ + arch::{gdt, interrupts, timer}, debug_serial_println, host::{bootparam::get_vtl1_memory_info, per_cpu_variables}, mm::MemoryProvider, mshv::{ hvcall, - vsm_intercept::raise_vtl0_gp_fault, vtl_switch::{vtl_switch, vtl_switch_init}, vtl1_mem_layout::{ VSM_SK_PTE_PAGES_COUNT, VTL1_INIT_HEAP_SIZE, VTL1_INIT_HEAP_START_PAGE, @@ -37,13 +40,16 @@ use litebox_platform_lvbs::{ get_text_start_address, }, }, - serial_println, -}; -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, SessionManager, TaInstance}; use litebox_shim_optee::{NormalWorldConstPtr, NormalWorldMutPtr, TaMemrefAddresses, UserConstPtr}; +use litebox_shim_optee::{ + msg_handler::{ + decode_ta_request, handle_optee_msg_args, handle_optee_smc_args, update_optee_msg_args, + write_rpc_args_to_regd_shm, + }, + rpc_context::rpc_context_map, +}; /// The session registry shared by all shims in this runner. fn session_manager() -> &'static SessionManager { @@ -548,14 +554,16 @@ fn optee_smc_handler(platform: &'static Platform, smc_args_addr: usize) -> Optee }; if let OpteeSmcResult::CallWithArg { msg_args, - rpc_args: _, + mut rpc_args, msg_args_phys_addr, } = smc_result { let mut msg_args = *msg_args; debug_serial_println!("OP-TEE SMC with MsgArgs Command: {:?}", msg_args.cmd); let result = match msg_args.cmd { - OpenSession => handle_open_session(platform, &mut msg_args, msg_args_phys_addr), + OpenSession => { + handle_open_session(platform, &mut msg_args, &mut rpc_args, msg_args_phys_addr) + } InvokeCommand => handle_invoke_command(platform, &mut msg_args, msg_args_phys_addr), CloseSession => handle_close_session(platform, &mut msg_args, msg_args_phys_addr), _ => { @@ -577,7 +585,72 @@ fn optee_smc_handler(platform: &'static Platform, smc_args_addr: usize) -> Optee unsafe { switch_to_base_page_table(platform) }; if let Err(e) = result { - smc_args.set_return_code(e); + if e == OpteeSmcReturnCode::RpcCmd { + debug_serial_println!("OP-TEE SMC returning RPC command to normal world"); + + // Dynamic TA RPC continuation requires CallWithRegdArg because the request + // buffer must be identified by a registered SHM reference and offset. + // CallWithArg and CallWithRpcArg do not provide that registered-SHM identity. + if smc_args.func_id() != Ok(OpteeSmcFunction::CallWithRegdArg) { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + } + let Some(rpc_args_ref) = rpc_args.as_ref() else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + }; + + let (registered_shm_ref, regd_shm_offset) = + match smc_args.optee_regd_shm_ref_and_offset() { + Ok(location) => location, + Err(error) => { + smc_args.set_return_code(error); + return *smc_args; + } + }; + let ta_uuid = match decode_ta_request(platform, &msg_args) + .ok() + .and_then(|request| request.uuid) + { + Some(ta_uuid) => ta_uuid, + None => { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + } + }; + let context_id = match rpc_context_map().allocate( + ta_uuid, + registered_shm_ref, + regd_shm_offset, + ) { + Ok(context_id) => context_id, + Err(error) => { + debug_serial_println!( + "Failed to allocate RPC context for LOAD_TA request: {:?}", + error + ); + smc_args.set_return_code(OpteeSmcReturnCode::EThreadLimit); + return *smc_args; + } + }; + smc_args.set_rpc_context_id(context_id); + + if let Err(e) = write_rpc_args_to_regd_shm( + platform, + registered_shm_ref, + regd_shm_offset, + msg_args.num_params, + rpc_args_ref, + ) { + let _ = rpc_context_map().take(context_id); + smc_args.set_return_code(e); + } else { + smc_args.set_return_code(OpteeSmcReturnCode::RpcCmd); + } + } else { + debug_serial_println!("OP-TEE SMC returning error code: {:?}", e); + smc_args.set_return_code(e); + } } else { smc_args.set_return_code(OpteeSmcReturnCode::Ok); } @@ -599,6 +672,7 @@ fn optee_smc_handler(platform: &'static Platform, smc_args_addr: usize) -> Optee fn handle_open_session( platform: &'static Platform, msg_args: &mut OpteeMsgArgs, + rpc_args: &mut Option>, msg_args_phys_addr: u64, ) -> Result<(), OpteeSmcReturnCode> { let ta_req_info = @@ -625,6 +699,7 @@ fn handle_open_session( platform, msg_args, msg_args_phys_addr, + rpc_args, params, ta_uuid, client_identity, @@ -814,22 +889,37 @@ fn open_session_single_instance( /// /// If ldelf loading or OpenSession entry point fails, the page table is torn down. /// Per OP-TEE OS semantics: if OpenSession returns non-success, cleanup happens. +#[allow(clippy::too_many_arguments)] fn open_session_new_instance( platform: &'static Platform, msg_args: &mut OpteeMsgArgs, msg_args_phys_addr: u64, + rpc_args: &mut Option>, params: &[litebox_common_optee::UteeParamOwned], ta_uuid: litebox_common_optee::TeeUuid, client_identity: Option, ta_req_info: &litebox_shim_optee::msg_handler::TaRequestInfo, ) -> Result<(), OpteeSmcReturnCode> { let shim = litebox_shim_optee::OpteeShimBuilder::new(platform, session_manager()).build(); - if shim.get_ta_bin(&ta_uuid).is_none() { - msg_args.session = 0; - msg_args.ret = TeeResult::ItemNotFound; - msg_args.ret_origin = TeeOrigin::Tee; - write_non_ta_msg_args_to_normal_world(platform, msg_args, msg_args_phys_addr)?; - return Ok(()); + + if !shim.contains_ta_bin(&ta_uuid) { + debug_serial_println!( + "TA binary not found for uuid={:?}, requesting load from normal world", + ta_uuid + ); + + let Some(rpc) = rpc_args.as_deref_mut() else { + debug_serial_println!( + "RPC args not present in incoming request, cannot request LOAD_TA from normal world" + ); + msg_args.session = 0; + msg_args.ret = TeeResult::ItemNotFound; + msg_args.ret_origin = TeeOrigin::Tee; + write_non_ta_msg_args_to_normal_world(platform, msg_args, msg_args_phys_addr)?; + return Ok(()); + }; + rpc.prepare_load_ta_rpc(ta_uuid, None)?; + return Err(OpteeSmcReturnCode::RpcCmd); } // Token is declared before `task_pt_guard` so it drops AFTER it. diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index c345967f56..09db90f369 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -225,17 +225,14 @@ impl GlobalState { self.ta_uuid_map.insert(*ta_uuid, ta_bin.into()) } - /// Get the TA binary associated with the given TA UUID. + /// Get the cached TA binary associated with the given TA UUID. pub(crate) fn get_ta_bin(&self, ta_uuid: &TeeUuid) -> Option> { - if let Some(ta_bin) = self.ta_uuid_map.get(ta_uuid) { - Some(ta_bin) - } else { - let ta_bin = Self::rpc_get_ta_bin(ta_uuid)?; - if !self.store_ta_bin(ta_uuid, &ta_bin) { - return None; - } - Some(ta_bin) - } + self.ta_uuid_map.get(ta_uuid) + } + + /// Return whether a TA binary is cached for the given UUID. + pub(crate) fn contains_ta_bin(&self, ta_uuid: &TeeUuid) -> bool { + self.ta_uuid_map.contains(ta_uuid) } /// Get the TA flags associated with the given TA UUID. @@ -262,11 +259,6 @@ impl GlobalState { pub(crate) fn remove_ta_bin(&self, ta_uuid: &TeeUuid) { let _ = self.ta_uuid_map.remove(ta_uuid); } - - /// RPC to get the TA binary associated with the given TA UUID. Placeholder for now. - fn rpc_get_ta_bin(_ta_uuid: &TeeUuid) -> Option> { - None - } } type UserMutPtr = @@ -365,11 +357,16 @@ impl OpteeShim { self.0.store_ta_bin(ta_uuid, ta_bin) } - /// Get the TA binary associated with the given TA UUID. + /// Get the cached TA binary associated with the given TA UUID. pub fn get_ta_bin(&self, ta_uuid: &TeeUuid) -> Option> { self.0.get_ta_bin(ta_uuid) } + /// Return whether a TA binary is cached for the given UUID. + pub fn contains_ta_bin(&self, ta_uuid: &TeeUuid) -> bool { + self.0.contains_ta_bin(ta_uuid) + } + /// Release all user-space memory mappings owned by this shim instance. /// /// This must be called before switching to the base page table and deleting @@ -1447,6 +1444,10 @@ impl TaUuidMap { self.inner.read().get(uuid).map(|info| info.binary.clone()) } + pub(crate) fn contains(&self, uuid: &TeeUuid) -> bool { + self.inner.read().contains_key(uuid) + } + /// Get the TA flags for a given UUID. pub(crate) fn get_flags(&self, uuid: &TeeUuid) -> Option { self.inner.read().get(uuid).map(|info| info.flags) diff --git a/litebox_shim_optee/src/syscalls/ldelf.rs b/litebox_shim_optee/src/syscalls/ldelf.rs index e79ceba61b..f376465ae4 100644 --- a/litebox_shim_optee/src/syscalls/ldelf.rs +++ b/litebox_shim_optee/src/syscalls/ldelf.rs @@ -256,7 +256,7 @@ impl Task { "sys_open_bin" ); - if self.global.get_ta_bin(&ta_uuid).is_none() { + if !self.global.contains_ta_bin(&ta_uuid) { return Err(TeeResult::ItemNotFound); } let new_handle = self.ta_handle_map.insert(ta_uuid); From 28ed5bceba818b8166998693cd56234ef3d4355e Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Thu, 10 Sep 2026 04:44:02 +0000 Subject: [PATCH 5/7] optee: Discover TA size and allocate shared memory Read the TA Size from VTL0, send SHM_ALLOC to get VTL0 to allocate memory to store TA Binary. Receive the allocation in VTL1, register that memory into a shm object and send the final LOAD_TA request. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu --- litebox_common_optee/src/lib.rs | 99 ++++++ litebox_runner_lvbs/Cargo.toml | 6 + litebox_runner_lvbs/src/lib.rs | 414 ++++++++++++++++++++------ litebox_shim_optee/src/msg_handler.rs | 102 +++++-- 4 files changed, 494 insertions(+), 127 deletions(-) diff --git a/litebox_common_optee/src/lib.rs b/litebox_common_optee/src/lib.rs index 56f7278ce1..d94117eae2 100644 --- a/litebox_common_optee/src/lib.rs +++ b/litebox_common_optee/src/lib.rs @@ -1450,6 +1450,22 @@ const OPTEE_MSG_RPC_CMD_RPMB_PROBE_RESET: u32 = 22; const OPTEE_MSG_RPC_CMD_RPMB_PROBE_NEXT: u32 = 23; const OPTEE_MSG_RPC_CMD_RPMB_PROBE_FRAMES: u32 = 24; +/// Memory that can be shared with a non-secure user space application +const OPTEE_RPC_SHM_TYPE_APPL: u32 = 0; +/// Memory only shared with non-secure kernel +const OPTEE_RPC_SHM_TYPE_KERNEL: u32 = 1; +/// Memory shared with non-secure kernel and exported to a non-secure user +/// space application +const OPTEE_RPC_SHM_TYPE_GLOBAL: u32 = 2; + +/// OP-TEE RPC shared memory types +#[repr(u32)] +pub enum OpteeRpcShmType { + Appl = OPTEE_RPC_SHM_TYPE_APPL, + Kernel = OPTEE_RPC_SHM_TYPE_KERNEL, + Global = OPTEE_RPC_SHM_TYPE_GLOBAL, +} + /// RPC command IDs from `optee_os/core/include/optee_msg.h` /// /// These are the command IDs used in the `cmd` field of the RPC `optee_msg_arg`. @@ -2157,6 +2173,12 @@ impl OpteeRpcArgs { OpteeMsgParamTmem::read_from_bytes(¶m.data).map_err(|_| OpteeSmcReturnCode::EBadCmd) } + /// Return whether an exactly validated TMEM output parameter uses a page list. + fn is_param_tmem_output_noncontiguous(&self, index: usize) -> Result { + self.get_param_tmem_output(index)?; + Ok(self.params[index].attr.noncontig()) + } + /// Access an RMEM output parameter with exact direction and flag validation. pub fn get_param_rmem_output( &self, @@ -2231,6 +2253,71 @@ impl OpteeRpcArgs { } } + /// Validate a successful first LOAD_TA response and return the requested TA size. + pub fn load_ta_size_response(&self) -> Result { + if self.cmd != OpteeRpcCommand::LoadTa + || self.ret != TeeResult::Success + || self.num_params != 2 + { + return Err(OpteeSmcReturnCode::EBadCmd); + } + + let size = self.get_param_tmem_output(1)?.size; + if size == 0 { + return Err(OpteeSmcReturnCode::EBadCmd); + } + Ok(size) + } + + /// Validate a successful SHM_ALLOC response and return its memory reference. + pub fn shm_alloc_response( + &self, + requested_size: u64, + ) -> Result { + if self.cmd != OpteeRpcCommand::ShmAlloc + || self.ret != TeeResult::Success + || self.num_params != 1 + { + return Err(OpteeSmcReturnCode::EBadCmd); + } + + let tmem = self.get_param_tmem_output(0)?; + if !self.is_param_tmem_output_noncontiguous(0)? + || tmem.buf_ptr == 0 + || tmem.shm_ref == 0 + || tmem.size < requested_size + { + return Err(OpteeSmcReturnCode::EBadCmd); + } + Ok(tmem) + } + + /// Prepare a shared-memory allocation RPC request to be sent to normal world. + pub fn prepare_shm_alloc_rpc( + &mut self, + shm_type: OpteeRpcShmType, + size: u64, + alignment: u64, + ) -> Result<(), OpteeSmcReturnCode> { + self.cmd = OpteeRpcCommand::ShmAlloc; + // Match OP-TEE's get_rpc_arg(): default to failure in case normal world + // returns without updating the RPC result. + self.ret = TeeResult::GenericError; + self.num_params = 1; + + self.set_param_attr_type(0, OpteeMsgAttrType::ValueInput)?; + self.set_param_value( + 0, + OpteeMsgParamValue { + a: shm_type as u64, + b: size, + c: alignment, + }, + )?; + + Ok(()) + } + /// Prepare a LOAD_TA RPC request to be sent to normal world. pub fn prepare_load_ta_rpc( &mut self, @@ -2380,6 +2467,7 @@ impl OpteeSmcArgs { /// TODO: Add stuffs based on the OP-TEE driver that LVBS is using. const OPTEE_SMC_FUNCID_GET_OS_UUID: usize = 0x0; const OPTEE_SMC_FUNCID_GET_OS_REVISION: usize = 0x1; +const OPTEE_SMC_FUNCID_RETURN_FROM_RPC: usize = 0x3; const OPTEE_SMC_FUNCID_CALL_WITH_ARG: usize = 0x4; const OPTEE_SMC_FUNCID_EXCHANGE_CAPABILITIES: usize = 0x9; const OPTEE_SMC_FUNCID_DISABLE_SHM_CACHE: usize = 0xa; @@ -2394,6 +2482,7 @@ const OPTEE_SMC_FUNCID_CALLS_REVISION: usize = 0xff03; pub enum OpteeSmcFunction { GetOsUuid = OPTEE_SMC_FUNCID_GET_OS_UUID, GetOsRevision = OPTEE_SMC_FUNCID_GET_OS_REVISION, + ReturnFromRpc = OPTEE_SMC_FUNCID_RETURN_FROM_RPC, CallWithArg = OPTEE_SMC_FUNCID_CALL_WITH_ARG, ExchangeCapabilities = OPTEE_SMC_FUNCID_EXCHANGE_CAPABILITIES, DisableShmCache = OPTEE_SMC_FUNCID_DISABLE_SHM_CACHE, @@ -2443,6 +2532,11 @@ pub enum OpteeSmcResult<'a> { rpc_args: Option>, msg_args_phys_addr: u64, }, + ReturnFromRpc { + msg_args: Box, + rpc_args: Box, + msg_args_phys_addr: u64, + }, } impl From> for OpteeSmcArgs { @@ -2506,6 +2600,11 @@ impl From> for OpteeSmcArgs { "OpteeSmcResult::CallWithArg cannot be converted to OpteeSmcArgs directly. Handle the incorporated OpteeMsgArgs." ); } + OpteeSmcResult::ReturnFromRpc { .. } => { + panic!( + "OpteeSmcResult::ReturnFromRpc cannot be converted to OpteeSmcArgs directly. Handle the incorporated OpteeMsgArgs and OpteeRpcArgs." + ); + } } } } diff --git a/litebox_runner_lvbs/Cargo.toml b/litebox_runner_lvbs/Cargo.toml index 5ad183ec92..a99f275e6d 100644 --- a/litebox_runner_lvbs/Cargo.toml +++ b/litebox_runner_lvbs/Cargo.toml @@ -3,6 +3,12 @@ name = "litebox_runner_lvbs" version = "0.1.0" edition = "2024" +[[bin]] +name = "litebox_runner_lvbs" +path = "src/main.rs" +test = false +bench = false + [dependencies] arrayvec = { version = "0.7.6", default-features = false } litebox = { version = "0.1.0", path = "../litebox" } diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 1bb2bfde21..1a073cb121 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -15,9 +15,9 @@ use litebox::{ use litebox_common_linux::errno::Errno; use litebox_common_lvbs::{NUM_VTLCALL_PARAMS, VsmError, VsmFunction}; use litebox_common_optee::{ - OpteeMessageCommand, OpteeMsgArgs, OpteeRpcArgs, OpteeSmcArgs, OpteeSmcFunction, - OpteeSmcResult, OpteeSmcReturnCode, TeeOrigin, TeeResult, UteeEntryFunc, UteeParams, - optee_msg_args_total_size, + OpteeMessageCommand, OpteeMsgArgs, OpteeMsgParamRmem, OpteeRpcArgs, OpteeRpcShmType, + OpteeSmcArgs, OpteeSmcFunction, OpteeSmcResult, OpteeSmcReturnCode, TeeOrigin, TeeResult, + UteeEntryFunc, UteeParams, optee_msg_args_total_size, }; use litebox_platform_lvbs::host::LvbsLinuxKernel as Platform; use litebox_platform_lvbs::mshv::vsm::{LvbsVtl0Gate, LvbsVtl0PrivilegedWriter, LvbsVtl1Gate}; @@ -45,10 +45,11 @@ use litebox_shim_optee::session::{OpenSessionTarget, SessionManager, TaInstance} use litebox_shim_optee::{NormalWorldConstPtr, NormalWorldMutPtr, TaMemrefAddresses, UserConstPtr}; use litebox_shim_optee::{ msg_handler::{ - decode_ta_request, handle_optee_msg_args, handle_optee_smc_args, update_optee_msg_args, - write_rpc_args_to_regd_shm, + checked_memref_size, decode_ta_request, handle_optee_msg_args, handle_optee_smc_args, + read_optee_msg_args_from_regd_shm, register_rpc_shm, unregister_rpc_shm, + update_optee_msg_args, write_rpc_args_to_regd_shm, }, - rpc_context::rpc_context_map, + rpc_context::{RpcContext, rpc_context_map}, }; /// The session registry shared by all shims in this runner. @@ -548,115 +549,334 @@ fn optee_smc_handler(platform: &'static Platform, smc_args_addr: usize) -> Optee let Ok(mut smc_args) = smc_args_ptr.read_at_offset(0) else { return make_error_response(OpteeSmcReturnCode::EBadAddr); }; - let Ok(smc_result) = handle_optee_smc_args(platform, &mut smc_args) else { + let smc_result = if smc_args.func_id() == Ok(OpteeSmcFunction::ReturnFromRpc) { + let context_id = match smc_args.get_rpc_context_id() { + Ok(context_id) => context_id, + Err(error) => { + smc_args.set_return_code(error); + return *smc_args; + } + }; + let Some(context) = rpc_context_map().get(context_id) else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + }; + let result = read_optee_msg_args_from_regd_shm( + platform, + context.registered_shm_ref(), + context.registered_shm_offset(), + ) + .and_then(|(msg_args, rpc_args, msg_args_phys_addr)| { + Ok(OpteeSmcResult::ReturnFromRpc { + msg_args, + rpc_args: rpc_args.ok_or(OpteeSmcReturnCode::EBadAddr)?, + msg_args_phys_addr, + }) + }); + if result.is_err() { + discard_rpc_context(context_id); + } + result + } else { + handle_optee_smc_args(platform, &mut smc_args) + }; + let Ok(smc_result) = smc_result else { smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); return *smc_args; }; - if let OpteeSmcResult::CallWithArg { - msg_args, - mut rpc_args, - msg_args_phys_addr, - } = smc_result - { - let mut msg_args = *msg_args; - debug_serial_println!("OP-TEE SMC with MsgArgs Command: {:?}", msg_args.cmd); - let result = match msg_args.cmd { - OpenSession => { - handle_open_session(platform, &mut msg_args, &mut rpc_args, msg_args_phys_addr) - } - InvokeCommand => handle_invoke_command(platform, &mut msg_args, msg_args_phys_addr), - CloseSession => handle_close_session(platform, &mut msg_args, msg_args_phys_addr), - _ => { - let r = handle_optee_msg_args(platform, &msg_args); - if r.is_ok() { - msg_args.ret = TeeResult::Success; - } else { - msg_args.ret = TeeResult::BadParameters; + match smc_result { + OpteeSmcResult::CallWithArg { + msg_args, + mut rpc_args, + msg_args_phys_addr, + } => { + let mut msg_args = *msg_args; + debug_serial_println!("OP-TEE SMC with MsgArgs Command: {:?}", msg_args.cmd); + let result = match msg_args.cmd { + OpenSession => { + handle_open_session(platform, &mut msg_args, &mut rpc_args, msg_args_phys_addr) } - msg_args.ret_origin = TeeOrigin::Tee; - let _ = - write_non_ta_msg_args_to_normal_world(platform, &msg_args, msg_args_phys_addr); - r - } - }; + InvokeCommand => handle_invoke_command(platform, &mut msg_args, msg_args_phys_addr), + CloseSession => handle_close_session(platform, &mut msg_args, msg_args_phys_addr), + _ => { + let r = handle_optee_msg_args(platform, &msg_args); + if r.is_ok() { + msg_args.ret = TeeResult::Success; + } else { + msg_args.ret = TeeResult::BadParameters; + } + msg_args.ret_origin = TeeOrigin::Tee; + let _ = write_non_ta_msg_args_to_normal_world( + platform, + &msg_args, + msg_args_phys_addr, + ); + r + } + }; - // Always switch back to base page table before returning to VTL0 - // Safety: No user-space memory references are held after this point - unsafe { switch_to_base_page_table(platform) }; + // Always switch back to base page table before returning to VTL0 + // Safety: No user-space memory references are held after this point + unsafe { switch_to_base_page_table(platform) }; - if let Err(e) = result { - if e == OpteeSmcReturnCode::RpcCmd { - debug_serial_println!("OP-TEE SMC returning RPC command to normal world"); + if let Err(e) = result { + if e == OpteeSmcReturnCode::RpcCmd { + debug_serial_println!("OP-TEE SMC returning RPC command to normal world"); - // Dynamic TA RPC continuation requires CallWithRegdArg because the request - // buffer must be identified by a registered SHM reference and offset. - // CallWithArg and CallWithRpcArg do not provide that registered-SHM identity. - if smc_args.func_id() != Ok(OpteeSmcFunction::CallWithRegdArg) { - smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); - return *smc_args; - } - let Some(rpc_args_ref) = rpc_args.as_ref() else { - smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); - return *smc_args; - }; - - let (registered_shm_ref, regd_shm_offset) = - match smc_args.optee_regd_shm_ref_and_offset() { - Ok(location) => location, - Err(error) => { - smc_args.set_return_code(error); - return *smc_args; - } - }; - let ta_uuid = match decode_ta_request(platform, &msg_args) - .ok() - .and_then(|request| request.uuid) - { - Some(ta_uuid) => ta_uuid, - None => { + // Dynamic TA RPC continuation requires CallWithRegdArg because the request + // buffer must be identified by a registered SHM reference and offset. + // CallWithArg and CallWithRpcArg do not provide that registered-SHM identity. + if smc_args.func_id() != Ok(OpteeSmcFunction::CallWithRegdArg) { smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); return *smc_args; } - }; - let context_id = match rpc_context_map().allocate( - ta_uuid, - registered_shm_ref, - regd_shm_offset, - ) { - Ok(context_id) => context_id, - Err(error) => { - debug_serial_println!( - "Failed to allocate RPC context for LOAD_TA request: {:?}", - error - ); - smc_args.set_return_code(OpteeSmcReturnCode::EThreadLimit); + let Some(rpc_args_ref) = rpc_args.as_ref() else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); return *smc_args; + }; + let (registered_shm_ref, regd_shm_offset) = + match smc_args.optee_regd_shm_ref_and_offset() { + Ok(location) => location, + Err(error) => { + smc_args.set_return_code(error); + return *smc_args; + } + }; + let ta_uuid = match decode_ta_request(platform, &msg_args) + .ok() + .and_then(|request| request.uuid) + { + Some(ta_uuid) => ta_uuid, + None => { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + } + }; + let context_id = match rpc_context_map().allocate( + ta_uuid, + registered_shm_ref, + regd_shm_offset, + ) { + Ok(context_id) => context_id, + Err(error) => { + debug_serial_println!( + "Failed to allocate RPC context for LOAD_TA request: {:?}", + error + ); + smc_args.set_return_code(OpteeSmcReturnCode::EThreadLimit); + return *smc_args; + } + }; + smc_args.set_rpc_context_id(context_id); + if let Err(e) = write_rpc_args_to_regd_shm( + platform, + registered_shm_ref, + regd_shm_offset, + msg_args.num_params, + rpc_args_ref, + ) { + let _ = rpc_context_map().take(context_id); + smc_args.set_return_code(e); + } else { + smc_args.set_return_code(OpteeSmcReturnCode::RpcCmd); } - }; - smc_args.set_rpc_context_id(context_id); - - if let Err(e) = write_rpc_args_to_regd_shm( - platform, - registered_shm_ref, - regd_shm_offset, - msg_args.num_params, - rpc_args_ref, - ) { - let _ = rpc_context_map().take(context_id); - smc_args.set_return_code(e); } else { - smc_args.set_return_code(OpteeSmcReturnCode::RpcCmd); + debug_serial_println!("OP-TEE SMC returning error code: {:?}", e); + smc_args.set_return_code(e); } } else { - debug_serial_println!("OP-TEE SMC returning error code: {:?}", e); - smc_args.set_return_code(e); + smc_args.set_return_code(OpteeSmcReturnCode::Ok); } - } else { - smc_args.set_return_code(OpteeSmcReturnCode::Ok); + *smc_args + } + OpteeSmcResult::ReturnFromRpc { + msg_args, + rpc_args, + msg_args_phys_addr: _, + } => { + let msg_args = *msg_args; + let mut rpc_args = *rpc_args; + + let context_id = match smc_args.get_rpc_context_id() { + Ok(context_id) => context_id, + Err(error) => { + smc_args.set_return_code(error); + return *smc_args; + } + }; + let Some(context) = rpc_context_map().get(context_id) else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + }; + match context { + RpcContext::LoadTaSize { .. } => { + handle_return_from_load_ta_rpc( + platform, + context_id, + &mut smc_args, + &msg_args, + &mut rpc_args, + ); + } + RpcContext::ShmAlloc { .. } => { + handle_return_from_shm_alloc_rpc( + platform, + context_id, + &mut smc_args, + &msg_args, + &mut rpc_args, + ); + } + _ => { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + } + } + *smc_args } - *smc_args + _ => smc_result.into(), + } +} + +fn handle_return_from_load_ta_rpc( + platform: &Platform, + context_id: u32, + smc_args: &mut OpteeSmcArgs, + msg_args: &OpteeMsgArgs, + rpc_args: &mut OpteeRpcArgs, +) { + let ta_size = match rpc_args.load_ta_size_response() { + Ok(ta_size) => ta_size, + Err(error) => { + discard_rpc_context(context_id); + smc_args.set_return_code(error); + return; + } + }; + debug_serial_println!("First LOAD_TA request, TA size: {}", ta_size); + if checked_memref_size(ta_size).is_err() { + debug_serial_println!("Invalid TA size in first LOAD_TA request"); + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + + if rpc_args + .prepare_shm_alloc_rpc(OpteeRpcShmType::Appl, ta_size, 8) + .is_err() + { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + if rpc_context_map() + .transition_to_shm_alloc(context_id, ta_size) + .is_err() + { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + if !write_next_rpc(platform, smc_args, msg_args, rpc_args, context_id) { + discard_rpc_context(context_id); + } +} + +fn handle_return_from_shm_alloc_rpc( + platform: &Platform, + context_id: u32, + smc_args: &mut OpteeSmcArgs, + msg_args: &OpteeMsgArgs, + rpc_args: &mut OpteeRpcArgs, +) { + let Some(RpcContext::ShmAlloc { + common: _, + requested_size, + }) = rpc_context_map().get(context_id) + else { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + }; + let tmem = match rpc_args.shm_alloc_response(requested_size) { + Ok(tmem) => tmem, + Err(error) => { + discard_rpc_context(context_id); + smc_args.set_return_code(error); + return; + } + }; + if checked_memref_size(tmem.size).is_err() { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + if register_rpc_shm(platform, &tmem).is_err() { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + let Some(context) = rpc_context_map().get(context_id) else { + let _ = unregister_rpc_shm(tmem.shm_ref); + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + }; + let rmem = OpteeMsgParamRmem { + offs: 0, + size: requested_size, + shm_ref: tmem.shm_ref, + }; + if rpc_args + .prepare_load_ta_rpc(context.ta_uuid(), Some(rmem)) + .is_err() + || rpc_context_map() + .transition_to_load_ta_binary(context_id, tmem.shm_ref) + .is_err() + { + let _ = unregister_rpc_shm(tmem.shm_ref); + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + if !write_next_rpc(platform, smc_args, msg_args, rpc_args, context_id) { + discard_rpc_context(context_id); + } +} + +fn write_next_rpc( + platform: &Platform, + smc_args: &mut OpteeSmcArgs, + msg_args: &OpteeMsgArgs, + rpc_args: &OpteeRpcArgs, + context_id: u32, +) -> bool { + let result = rpc_context_map() + .get(context_id) + .ok_or(OpteeSmcReturnCode::EBadCmd) + .and_then(|context| { + write_rpc_args_to_regd_shm( + platform, + context.registered_shm_ref(), + context.registered_shm_offset(), + msg_args.num_params, + rpc_args, + ) + }); + if let Err(error) = result { + discard_rpc_context(context_id); + smc_args.set_return_code(error); + false } else { - smc_result.into() + smc_args.set_return_code(OpteeSmcReturnCode::RpcCmd); + true + } +} + +fn discard_rpc_context(context_id: u32) { + if let Some(context) = rpc_context_map().take(context_id) + && let Some(shm_ref) = context.shm_ref() + { + let _ = unregister_rpc_shm(shm_ref); } } diff --git a/litebox_shim_optee/src/msg_handler.rs b/litebox_shim_optee/src/msg_handler.rs index 6ab2810ffa..dbf3b4ef30 100644 --- a/litebox_shim_optee/src/msg_handler.rs +++ b/litebox_shim_optee/src/msg_handler.rs @@ -70,7 +70,7 @@ const MAX_SHM_MEMREF_SIZE: usize = 8 * 1024 * 1024; const MAX_SHM_REF_MAP_ENTRIES: usize = 1024; #[inline] -fn page_align_down(address: u64) -> u64 { +pub fn page_align_down(address: u64) -> u64 { address & !(PAGE_SIZE as u64 - 1) } @@ -80,7 +80,7 @@ fn page_align_up(len: u64) -> Option { } #[inline] -fn checked_memref_size(size: u64) -> Result { +pub fn checked_memref_size(size: u64) -> Result { if size > MAX_SHM_MEMREF_SIZE as u64 { return Err(OpteeSmcReturnCode::ENomem); } @@ -208,6 +208,72 @@ pub fn read_optee_msg_args_from_phys( parse_optee_msg_args(&blob, has_rpc_arg) } +/// Read main and RPC arguments from an explicitly identified registered SHM view. +#[allow(clippy::type_complexity)] +pub fn read_optee_msg_args_from_regd_shm< + Platform: litebox_common_linux::vmap::VmapManager, +>( + platform: &Platform, + shm_ref: u64, + offset: usize, +) -> Result<(Box, Option>, u64), OpteeSmcReturnCode> { + let shm_info = shm_ref_map() + .get(shm_ref) + .ok_or(OpteeSmcReturnCode::EBadAddr)?; + let main_max = optee_msg_args_total_size(OpteeMsgArgs::MAX_ARG_PARAM_COUNT.trunc()); + let copy_size = + main_max + optee_msg_args_total_size(OpteeRpcArgs::MAX_RPC_ARG_PARAM_COUNT.trunc()); + let mut blob = alloc::vec![0u8; copy_size]; + shm_info.read_at(platform, offset, &mut blob)?; + let (msg_args, rpc_args) = parse_optee_msg_args(&blob, true)?; + + let total_offset = shm_info + .page_offset + .checked_add(offset) + .ok_or(OpteeSmcReturnCode::EBadAddr)?; + let page_index = total_offset / PAGE_SIZE; + let offset_in_page = total_offset % PAGE_SIZE; + let msg_args_phys_addr = shm_info + .page_addrs + .get(page_index) + .ok_or(OpteeSmcReturnCode::EBadAddr)? + .as_usize() + .checked_add(offset_in_page) + .ok_or(OpteeSmcReturnCode::EBadAddr)? as u64; + + Ok((msg_args, rpc_args, msg_args_phys_addr)) +} + +/// Register a page-list-backed TMEM allocation returned by normal world. +pub fn register_rpc_shm>( + platform: &Platform, + tmem: &OpteeMsgParamTmem, +) -> Result<(), OpteeSmcReturnCode> { + checked_memref_size(tmem.size)?; + let pages_data_phys_addr = page_align_down(tmem.buf_ptr); + let page_offset = tmem + .buf_ptr + .checked_sub(pages_data_phys_addr) + .ok_or(OpteeSmcReturnCode::EBadAddr)?; + let total_size = page_offset + .checked_add(tmem.size) + .ok_or(OpteeSmcReturnCode::EBadAddr)?; + let aligned_size = page_align_up(total_size).ok_or(OpteeSmcReturnCode::EBadAddr)?; + shm_ref_map().register_shm( + platform, + pages_data_phys_addr, + page_offset, + tmem.size, + aligned_size, + tmem.shm_ref, + ) +} + +/// Remove a shared-memory mapping inserted for an RPC allocation. +pub fn unregister_rpc_shm(shm_ref: u64) -> bool { + shm_ref_map().remove(shm_ref).is_some() +} + /// This function handles `OpteeSmcArgs` passed from the normal world (VTL0) via an OP-TEE SMC call. /// It returns an `OpteeSmcResult` representing the result of the SMC call or `OpteeMsgArgs` it contains /// if the SMC call involves with an OP-TEE message which should be handled by @@ -244,41 +310,17 @@ pub fn handle_optee_smc_args<'a, Platform: crate::OpteeShimPlatform>( msg_args_phys_addr: msg_args_addr as u64, }) } + OpteeSmcFunction::ReturnFromRpc => Err(OpteeSmcReturnCode::EBadCmd), OpteeSmcFunction::CallWithRegdArg => { // `OpteeMsgArgs` is located at the offset specified in args[3] within the shared memory region pointed by args[1]:args[2]. let (shm_ref, offset) = smc.optee_regd_shm_ref_and_offset()?; - let shm_info = shm_ref_map() - .get(shm_ref) - .ok_or(OpteeSmcReturnCode::EBadAddr)?; - - // Compute copy size from known-good upper bounds — no untrusted data involved. - let main_max = optee_msg_args_total_size(OpteeMsgArgs::MAX_ARG_PARAM_COUNT.trunc()); - let copy_size = - main_max + optee_msg_args_total_size(OpteeRpcArgs::MAX_RPC_ARG_PARAM_COUNT.trunc()); - - let mut blob = alloc::vec![0u8; copy_size]; - shm_info.read_at(platform, offset, &mut blob)?; - let (msg_args, rpc_args) = parse_optee_msg_args(&blob, true)?; - - // Compute the physical address of `OpteeMsgArgs` - let total_offset = shm_info - .page_offset - .checked_add(offset) - .ok_or(OpteeSmcReturnCode::EBadAddr)?; - let page_index = total_offset / PAGE_SIZE; - let offset_in_page = total_offset % PAGE_SIZE; - if page_index >= shm_info.page_addrs.len() { - return Err(OpteeSmcReturnCode::EBadAddr); - } - let msg_args_addr = shm_info.page_addrs[page_index] - .as_usize() - .checked_add(offset_in_page) - .ok_or(OpteeSmcReturnCode::EBadAddr)?; + let (msg_args, rpc_args, msg_args_phys_addr) = + read_optee_msg_args_from_regd_shm(platform, shm_ref, offset)?; Ok(OpteeSmcResult::CallWithArg { msg_args, rpc_args, - msg_args_phys_addr: msg_args_addr as u64, + msg_args_phys_addr, }) } OpteeSmcFunction::ExchangeCapabilities => { From 3ae14cb43be6b2d09e8484f1fd156efb3820f8b9 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Thu, 10 Sep 2026 05:05:28 +0000 Subject: [PATCH 6/7] optee: Complete Dynamic TA loading and resume OpenSession Load TA binary and initiate OpenSession. Initiate SHM_FREE RPC to clean up the memory allocated in VTL0. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Praveen K Paladugu --- litebox_common_optee/src/lib.rs | 44 +++++ litebox_runner_lvbs/src/lib.rs | 243 +++++++++++++++++++++++--- litebox_shim_optee/src/lib.rs | 13 +- litebox_shim_optee/src/msg_handler.rs | 17 +- 4 files changed, 283 insertions(+), 34 deletions(-) diff --git a/litebox_common_optee/src/lib.rs b/litebox_common_optee/src/lib.rs index d94117eae2..93c8b774fb 100644 --- a/litebox_common_optee/src/lib.rs +++ b/litebox_common_optee/src/lib.rs @@ -2269,6 +2269,27 @@ impl OpteeRpcArgs { Ok(size) } + /// Validate a successful second LOAD_TA response and read the TA binary. + pub fn load_ta_binary_response( + &self, + expected_shm_ref: u64, + expected_size: u64, + read_binary: impl FnOnce(u64, u64) -> Result, OpteeSmcReturnCode>, + ) -> Result, OpteeSmcReturnCode> { + if self.cmd != OpteeRpcCommand::LoadTa + || self.ret != TeeResult::Success + || self.num_params != 2 + { + return Err(OpteeSmcReturnCode::EBadCmd); + } + + let rmem = self.get_param_rmem_output(1)?; + if rmem.shm_ref != expected_shm_ref || rmem.offs != 0 || rmem.size != expected_size { + return Err(OpteeSmcReturnCode::EBadCmd); + } + read_binary(rmem.shm_ref, rmem.size) + } + /// Validate a successful SHM_ALLOC response and return its memory reference. pub fn shm_alloc_response( &self, @@ -2318,6 +2339,29 @@ impl OpteeRpcArgs { Ok(()) } + /// Prepare a shared-memory free RPC request to be sent to normal world. + pub fn prepare_shm_free_rpc( + &mut self, + shm_type: OpteeRpcShmType, + shm_ref: u64, + ) -> Result<(), OpteeSmcReturnCode> { + self.cmd = OpteeRpcCommand::ShmFree; + // Match OP-TEE's get_rpc_arg(): default to failure in case normal world + // returns without updating the RPC result. + self.ret = TeeResult::GenericError; + self.num_params = 1; + + self.set_param_attr_type(0, OpteeMsgAttrType::ValueInput)?; + self.set_param_value( + 0, + OpteeMsgParamValue { + a: shm_type as u64, + b: shm_ref, + c: 0, + }, + ) + } + /// Prepare a LOAD_TA RPC request to be sent to normal world. pub fn prepare_load_ta_rpc( &mut self, diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 1a073cb121..957e0651f7 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -15,9 +15,9 @@ use litebox::{ use litebox_common_linux::errno::Errno; use litebox_common_lvbs::{NUM_VTLCALL_PARAMS, VsmError, VsmFunction}; use litebox_common_optee::{ - OpteeMessageCommand, OpteeMsgArgs, OpteeMsgParamRmem, OpteeRpcArgs, OpteeRpcShmType, - OpteeSmcArgs, OpteeSmcFunction, OpteeSmcResult, OpteeSmcReturnCode, TeeOrigin, TeeResult, - UteeEntryFunc, UteeParams, optee_msg_args_total_size, + OpteeMessageCommand, OpteeMsgArgs, OpteeMsgParamRmem, OpteeRpcArgs, OpteeRpcCommand, + OpteeRpcShmType, OpteeSmcArgs, OpteeSmcFunction, OpteeSmcResult, OpteeSmcReturnCode, TeeOrigin, + TeeResult, UteeEntryFunc, UteeParams, optee_msg_args_total_size, }; use litebox_platform_lvbs::host::LvbsLinuxKernel as Platform; use litebox_platform_lvbs::mshv::vsm::{LvbsVtl0Gate, LvbsVtl0PrivilegedWriter, LvbsVtl1Gate}; @@ -46,10 +46,10 @@ use litebox_shim_optee::{NormalWorldConstPtr, NormalWorldMutPtr, TaMemrefAddress use litebox_shim_optee::{ msg_handler::{ checked_memref_size, decode_ta_request, handle_optee_msg_args, handle_optee_smc_args, - read_optee_msg_args_from_regd_shm, register_rpc_shm, unregister_rpc_shm, + read_optee_msg_args_from_regd_shm, read_rpc_shm, register_rpc_shm, unregister_rpc_shm, update_optee_msg_args, write_rpc_args_to_regd_shm, }, - rpc_context::{RpcContext, rpc_context_map}, + rpc_context::{RpcCompletion, RpcContext, rpc_context_map}, }; /// The session registry shared by all shims in this runner. @@ -692,9 +692,9 @@ fn optee_smc_handler(platform: &'static Platform, smc_args_addr: usize) -> Optee OpteeSmcResult::ReturnFromRpc { msg_args, rpc_args, - msg_args_phys_addr: _, + msg_args_phys_addr, } => { - let msg_args = *msg_args; + let mut msg_args = *msg_args; let mut rpc_args = *rpc_args; let context_id = match smc_args.get_rpc_context_id() { @@ -727,8 +727,24 @@ fn optee_smc_handler(platform: &'static Platform, smc_args_addr: usize) -> Optee &mut rpc_args, ); } - _ => { - smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + RpcContext::LoadTaBinary { .. } => { + handle_return_from_load_ta_binary_rpc( + platform, + context_id, + &mut smc_args, + &msg_args, + &mut rpc_args, + ); + } + RpcContext::ShmFree { .. } => { + handle_return_from_shm_free_rpc( + platform, + context_id, + &mut smc_args, + &mut msg_args, + &rpc_args, + msg_args_phys_addr, + ); } } *smc_args @@ -781,6 +797,133 @@ fn handle_return_from_load_ta_rpc( } } +fn handle_return_from_load_ta_binary_rpc( + platform: &'static Platform, + context_id: u32, + smc_args: &mut OpteeSmcArgs, + msg_args: &OpteeMsgArgs, + rpc_args: &mut OpteeRpcArgs, +) { + let Some( + context @ RpcContext::LoadTaBinary { + requested_size, + shm_ref, + .. + }, + ) = rpc_context_map().get(context_id) + else { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + }; + let response = rpc_args.load_ta_binary_response(shm_ref, requested_size, |shm_ref, size| { + let ta_size = checked_memref_size(size)?; + let mut ta_binary = alloc::vec![0u8; ta_size]; + read_rpc_shm(platform, shm_ref, 0, &mut ta_binary)?; + Ok(ta_binary.into_boxed_slice()) + }); + + let completion = match response { + Ok(ta_binary) => { + let shim = + litebox_shim_optee::OpteeShimBuilder::new(platform, session_manager()).build(); + if shim.store_ta_bin(&context.ta_uuid(), &ta_binary) { + RpcCompletion::OpenSession + } else { + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd) + } + } + Err(error) => RpcCompletion::ReturnError(error), + }; + if !start_shm_free_rpc( + platform, smc_args, msg_args, rpc_args, context_id, shm_ref, completion, + ) && completion == RpcCompletion::OpenSession + { + litebox_shim_optee::OpteeShimBuilder::new(platform, session_manager()) + .build() + .remove_ta_bin(&context.ta_uuid()); + } +} + +#[allow(clippy::too_many_arguments)] +fn start_shm_free_rpc( + platform: &Platform, + smc_args: &mut OpteeSmcArgs, + msg_args: &OpteeMsgArgs, + rpc_args: &mut OpteeRpcArgs, + context_id: u32, + shm_ref: u64, + completion: RpcCompletion, +) -> bool { + let unregister_local_shm = matches!( + rpc_context_map().get(context_id), + Some(RpcContext::LoadTaBinary { .. }) + ); + if rpc_args + .prepare_shm_free_rpc(OpteeRpcShmType::Appl, shm_ref) + .is_err() + || rpc_context_map() + .transition_to_shm_free(context_id, shm_ref, completion) + .is_err() + { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return false; + } + if unregister_local_shm { + let _ = unregister_rpc_shm(shm_ref); + } + if !write_next_rpc(platform, smc_args, msg_args, rpc_args, context_id) { + discard_rpc_context(context_id); + return false; + } + true +} + +fn handle_return_from_shm_free_rpc( + platform: &'static Platform, + context_id: u32, + smc_args: &mut OpteeSmcArgs, + msg_args: &mut OpteeMsgArgs, + rpc_args: &OpteeRpcArgs, + msg_args_phys_addr: u64, +) { + let Some(context) = rpc_context_map().take(context_id) else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + }; + let RpcContext::ShmFree { completion, .. } = context else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + }; + match completion { + RpcCompletion::OpenSession => { + let shim = + litebox_shim_optee::OpteeShimBuilder::new(platform, session_manager()).build(); + if rpc_args.cmd != OpteeRpcCommand::ShmFree + || rpc_args.ret != TeeResult::Success + || rpc_args.num_params != 1 + { + shim.remove_ta_bin(&context.ta_uuid()); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + let mut no_rpc_args = None; + let result = + handle_open_session(platform, msg_args, &mut no_rpc_args, msg_args_phys_addr); + // Regardless of the result, remove the TA binary from VTL1 cache. + // If the TA is SINGLE_INSTANCE and has KEEP_ALIVE flag, the TA + // runtime will be cached in memory even after last session is + // closed. + shim.remove_ta_bin(&context.ta_uuid()); + smc_args.set_return_code(result.err().unwrap_or(OpteeSmcReturnCode::Ok)); + } + // ReturnError is only recorded before the TA binary is cached or when + // caching fails. Removing by UUID here could evict another load's entry. + RpcCompletion::ReturnError(error) => smc_args.set_return_code(error), + } +} + fn handle_return_from_shm_alloc_rpc( platform: &Platform, context_id: u32, @@ -806,40 +949,91 @@ fn handle_return_from_shm_alloc_rpc( } }; if checked_memref_size(tmem.size).is_err() { - discard_rpc_context(context_id); - smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + tmem.shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ); return; } if register_rpc_shm(platform, &tmem).is_err() { - discard_rpc_context(context_id); - smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + tmem.shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ); return; } - let Some(context) = rpc_context_map().get(context_id) else { + if rpc_context_map() + .transition_to_load_ta_binary(context_id, tmem.shm_ref) + .is_err() + { let _ = unregister_rpc_shm(tmem.shm_ref); - discard_rpc_context(context_id); - smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + tmem.shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ); return; - }; + } let rmem = OpteeMsgParamRmem { offs: 0, size: requested_size, shm_ref: tmem.shm_ref, }; + let Some(context) = rpc_context_map().get(context_id) else { + let _ = unregister_rpc_shm(tmem.shm_ref); + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + tmem.shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ); + return; + }; if rpc_args .prepare_load_ta_rpc(context.ta_uuid(), Some(rmem)) .is_err() - || rpc_context_map() - .transition_to_load_ta_binary(context_id, tmem.shm_ref) - .is_err() { let _ = unregister_rpc_shm(tmem.shm_ref); - discard_rpc_context(context_id); - smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + tmem.shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ); return; } if !write_next_rpc(platform, smc_args, msg_args, rpc_args, context_id) { - discard_rpc_context(context_id); + let _ = unregister_rpc_shm(tmem.shm_ref); + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + tmem.shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadAddr), + ); + return; } } @@ -863,7 +1057,6 @@ fn write_next_rpc( ) }); if let Err(error) = result { - discard_rpc_context(context_id); smc_args.set_return_code(error); false } else { @@ -874,7 +1067,7 @@ fn write_next_rpc( fn discard_rpc_context(context_id: u32) { if let Some(context) = rpc_context_map().take(context_id) - && let Some(shm_ref) = context.shm_ref() + && let RpcContext::LoadTaBinary { shm_ref, .. } = context { let _ = unregister_rpc_shm(shm_ref); } diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index 09db90f369..00e51a6b08 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -249,13 +249,7 @@ impl GlobalState { TimeProvider::now(self.platform).duration_since(&self.boot_instant) } - /// Remove the TA binary associated with the given TA UUID. - /// - /// Since a TA binary can be continuously loaded/used by multiple clients, we cache it - /// to avoid repeated RPCs and memory transfers. We remove it lazily if there is - /// a memory pressure. - /// - #[expect(dead_code)] + /// Remove a TA binary after it is no longer needed in the trusted cache. pub(crate) fn remove_ta_bin(&self, ta_uuid: &TeeUuid) { let _ = self.ta_uuid_map.remove(ta_uuid); } @@ -367,6 +361,11 @@ impl OpteeShim { self.0.contains_ta_bin(ta_uuid) } + /// Remove a TA binary from the trusted cache. + pub fn remove_ta_bin(&self, ta_uuid: &TeeUuid) { + self.0.remove_ta_bin(ta_uuid); + } + /// Release all user-space memory mappings owned by this shim instance. /// /// This must be called before switching to the base page table and deleting diff --git a/litebox_shim_optee/src/msg_handler.rs b/litebox_shim_optee/src/msg_handler.rs index dbf3b4ef30..7b9ec036c8 100644 --- a/litebox_shim_optee/src/msg_handler.rs +++ b/litebox_shim_optee/src/msg_handler.rs @@ -70,7 +70,7 @@ const MAX_SHM_MEMREF_SIZE: usize = 8 * 1024 * 1024; const MAX_SHM_REF_MAP_ENTRIES: usize = 1024; #[inline] -pub fn page_align_down(address: u64) -> u64 { +fn page_align_down(address: u64) -> u64 { address & !(PAGE_SIZE as u64 - 1) } @@ -84,7 +84,7 @@ pub fn checked_memref_size(size: u64) -> Result { if size > MAX_SHM_MEMREF_SIZE as u64 { return Err(OpteeSmcReturnCode::ENomem); } - Ok(size.trunc()) + usize::try_from(size).map_err(|_| OpteeSmcReturnCode::ENomem) } fn parse_optee_msg_args( @@ -274,6 +274,19 @@ pub fn unregister_rpc_shm(shm_ref: u64) -> bool { shm_ref_map().remove(shm_ref).is_some() } +/// Copy bytes from a registered RPC allocation into trusted memory. +pub fn read_rpc_shm>( + platform: &Platform, + shm_ref: u64, + offset: usize, + buffer: &mut [u8], +) -> Result<(), OpteeSmcReturnCode> { + shm_ref_map() + .get(shm_ref) + .ok_or(OpteeSmcReturnCode::EBadAddr)? + .read_at(platform, offset, buffer) +} + /// This function handles `OpteeSmcArgs` passed from the normal world (VTL0) via an OP-TEE SMC call. /// It returns an `OpteeSmcResult` representing the result of the SMC call or `OpteeMsgArgs` it contains /// if the SMC call involves with an OP-TEE message which should be handled by From 6f59e50c8a134602ca2545f18bdc5e7b535e6283 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Wed, 23 Sep 2026 15:53:13 +0000 Subject: [PATCH 7/7] lvbs: Fix clippy errors Signed-off-by: Praveen K Paladugu --- litebox_runner_lvbs/src/lib.rs | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 957e0651f7..91645a9293 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -642,15 +642,12 @@ fn optee_smc_handler(platform: &'static Platform, smc_args_addr: usize) -> Optee return *smc_args; } }; - let ta_uuid = match decode_ta_request(platform, &msg_args) + let Some(ta_uuid) = decode_ta_request(platform, &msg_args) .ok() .and_then(|request| request.uuid) - { - Some(ta_uuid) => ta_uuid, - None => { - smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); - return *smc_args; - } + else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; }; let context_id = match rpc_context_map().allocate( ta_uuid, @@ -1033,7 +1030,6 @@ fn handle_return_from_shm_alloc_rpc( tmem.shm_ref, RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadAddr), ); - return; } }