From 5648de0aacba9cd5dfc5163c507c5d094e877486 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Tue, 25 Aug 2026 15:23:59 +0000 Subject: [PATCH 1/5] optee: Introduce a method to check if a TA is cached This check eliminates the need to clone a TA binary while checking if a TA exists in cache. Signed-off-by: Praveen K Paladugu --- litebox_shim_optee/src/lib.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index cd316fed8..b83e2f157 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -203,6 +203,11 @@ impl GlobalState { } } + /// 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. pub(crate) fn get_ta_flags(&self, ta_uuid: &TeeUuid) -> TaFlags { self.ta_uuid_map.get_flags(ta_uuid).unwrap_or_default() @@ -315,6 +320,11 @@ impl OpteeShim { 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 @@ -1362,6 +1372,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) From 8e8416881b80f687b4b1f7ae6fb0c4e1b41a791c Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Tue, 25 Aug 2026 17:02:12 +0000 Subject: [PATCH 2/5] optee: Add UUID serialization helpers Signed-off-by: Praveen K Paladugu --- litebox_common_optee/src/lib.rs | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/litebox_common_optee/src/lib.rs b/litebox_common_optee/src/lib.rs index dbad3229f..377df0f45 100644 --- a/litebox_common_optee/src/lib.rs +++ b/litebox_common_optee/src/lib.rs @@ -679,6 +679,14 @@ impl TeeUuid { Self::from_bytes(bytes) } + 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]; @@ -688,6 +696,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`. @@ -2517,6 +2535,11 @@ 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] From 9bdfc1de526d58ec2a67eae5c8b74b3c1463109e Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Tue, 25 Aug 2026 17:02:54 +0000 Subject: [PATCH 3/5] optee: Add RPC parameter setters Signed-off-by: Praveen K Paladugu --- litebox_common_optee/src/lib.rs | 84 +++++++++++++++++++++++++++++++-- 1 file changed, 80 insertions(+), 4 deletions(-) diff --git a/litebox_common_optee/src/lib.rs b/litebox_common_optee/src/lib.rs index 377df0f45..08081974b 100644 --- a/litebox_common_optee/src/lib.rs +++ b/litebox_common_optee/src/lib.rs @@ -2135,6 +2135,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, @@ -2148,10 +2176,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`. @@ -2639,6 +2663,58 @@ mod tests { assert_eq!(header_out.num_params, 2); } + #[test] + fn test_optee_rpc_args_attr_and_rmem_setters() { + let header = OpteeMsgArgsHeader { + cmd: OpteeRpcCommand::LoadTa as u32, + func: 0, + session: 0, + cancel_id: 0, + pad: 0, + ret: 0, + ret_origin: 0, + num_params: 1, + }; + let raw_params = [0u8; size_of::()]; + let mut rpc_args = OpteeRpcArgs::from_header_and_raw_params(&header, &raw_params) + .expect("should parse RPC args"); + + rpc_args.params[0].attr = OpteeMsgAttr::META_VALUE_INPUT; + rpc_args + .set_param_attr_type(0, OpteeMsgAttrType::RmemOutput) + .expect("attribute index should be available"); + assert_eq!( + rpc_args.params[0].attr.attr_type(), + OpteeMsgAttrType::RmemOutput as u8 + ); + assert!(!rpc_args.params[0].attr.meta()); + assert!(!rpc_args.params[0].attr.noncontig()); + + let rmem = OpteeMsgParamRmem { + offs: 0x0102_0304_0506_0708, + size: 0x1112_1314_1516_1718, + shm_ref: 0x2122_2324_2526_2728, + }; + rpc_args + .set_param_rmem(0, rmem) + .expect("rmem index should be available"); + assert_eq!(&rpc_args.params[0].data[0..8], &rmem.offs.to_le_bytes()); + assert_eq!(&rpc_args.params[0].data[8..16], &rmem.size.to_le_bytes()); + assert_eq!( + &rpc_args.params[0].data[16..24], + &rmem.shm_ref.to_le_bytes() + ); + + assert_eq!( + rpc_args.set_param_attr_type(1, OpteeMsgAttrType::RmemOutput), + Err(OpteeSmcReturnCode::ENotAvail) + ); + assert_eq!( + rpc_args.set_param_rmem(1, rmem), + Err(OpteeSmcReturnCode::ENotAvail) + ); + } + #[test] fn test_rpc_args_rejects_main_cmd() { // Pick a cmd value that lies in the gap between Plugin (12) and I2C Transfer (21), From 929b5b9d32767f252a9288aa652fea37f949d009 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Thu, 13 Aug 2026 19:39:10 +0000 Subject: [PATCH 4/5] optee: Initiate LOAD_TA if TA not found If the TA binary with the input uuid does not exist in VTL1 send a LOAD_TA RPC to VTL0. The first LOAD_TA should be sent with memref size of 0 to get the size of the TA binary from VTL0. Signed-off-by: Praveen K Paladugu --- litebox_common_optee/src/lib.rs | 46 +++++++++++++++++++ litebox_runner_lvbs/src/lib.rs | 57 +++++++++++++++++++----- litebox_shim_optee/src/syscalls/ldelf.rs | 2 +- 3 files changed, 93 insertions(+), 12 deletions(-) diff --git a/litebox_common_optee/src/lib.rs b/litebox_common_optee/src/lib.rs index 08081974b..bb96b9469 100644 --- a/litebox_common_optee/src/lib.rs +++ b/litebox_common_optee/src/lib.rs @@ -2178,6 +2178,52 @@ impl OpteeRpcArgs { } } +/// Prepare a LOAD_TA RPC request to be sent to normal world. +/// +/// When `memref` is `None`, the request asks normal world to return the TA size. +/// When it is `Some`, the request provides registered memory for the TA binary. +pub fn prepare_load_ta_rpc( + rpc_msg_args: &mut OpteeRpcArgs, + ta_uuid: TeeUuid, + memref_size: u64, + memref: Option, +) -> Result<(), OpteeSmcReturnCode> { + rpc_msg_args.cmd = OpteeRpcCommand::LoadTa; + rpc_msg_args.num_params = 2; + + rpc_msg_args.set_param_attr_type(0, OpteeMsgAttrType::ValueInput)?; + let uuid_bytes = ta_uuid.to_u64_array(); + rpc_msg_args.set_param_value( + 0, + OpteeMsgParamValue { + a: uuid_bytes[0], + b: uuid_bytes[1], + c: 0, + }, + )?; + + if memref.is_none() { + // First call of LOAD_TA protocol: normal world returns the TA size in memref_size. + rpc_msg_args.set_param_attr_type(1, OpteeMsgAttrType::TmemOutput)?; + rpc_msg_args.set_param_tmem( + 1, + OpteeMsgParamTmem { + buf_ptr: 0, + size: memref_size, + shm_ref: 0, + }, + )?; + } else { + // Second call of LOAD_TA protocol: secure world provides a memref for the TA binary. + rpc_msg_args.set_param_attr_type(1, OpteeMsgAttrType::RmemOutput)?; + if let Some(rmem) = memref { + rpc_msg_args.set_param_rmem(1, rmem)?; + } + } + + Ok(()) +} + /// Serialize the params portion as raw bytes into `buf`. #[inline] fn write_optee_msg_params_to_buf(params: &[OpteeMsgParam], buf: &mut [u8]) { diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 7a0635551..33f20806d 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -17,6 +17,7 @@ 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, + prepare_load_ta_rpc, }; use litebox_platform_lvbs::mshv::vsm::{LvbsVtl0Gate, LvbsVtl0PrivilegedWriter, LvbsVtl1Gate}; use litebox_platform_lvbs::{ @@ -520,14 +521,14 @@ fn optee_smc_handler(smc_args_addr: usize) -> OpteeSmcArgs { }; 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(&mut msg_args, msg_args_phys_addr), + OpenSession => handle_open_session(&mut msg_args, &mut rpc_args, msg_args_phys_addr), InvokeCommand => handle_invoke_command(&mut msg_args, msg_args_phys_addr), CloseSession => handle_close_session(&mut msg_args, msg_args_phys_addr), _ => { @@ -548,7 +549,23 @@ fn optee_smc_handler(smc_args_addr: usize) -> OpteeSmcArgs { unsafe { switch_to_base_page_table() }; 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"); + let Some(rpc_args_ref) = rpc_args.as_ref() else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + }; + if let Err(e) = + write_rpc_args_to_normal_world(&msg_args, msg_args_phys_addr, rpc_args_ref) + { + 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); } @@ -569,18 +586,44 @@ fn optee_smc_handler(smc_args_addr: usize) -> OpteeSmcArgs { /// instance cleanup for TARGET_DEAD on single-instance TAs). fn handle_open_session( msg_args: &mut OpteeMsgArgs, + rpc_args: &mut Option>, msg_args_phys_addr: u64, ) -> Result<(), OpteeSmcReturnCode> { let ta_req_info = decode_ta_request(msg_args).map_err(|_| OpteeSmcReturnCode::EBadCmd)?; if ta_req_info.entry_func != UteeEntryFunc::OpenSession { return Err(OpteeSmcReturnCode::EBadCmd); } + let shim: litebox_shim_optee::OpteeShim = litebox_shim_optee::OpteeShimBuilder::new().build(); let ta_uuid = ta_req_info.uuid.ok_or(OpteeSmcReturnCode::EBadCmd)?; + 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(msg_args, msg_args_phys_addr)?; + return Ok(()); + }; + // LOAD_TA is a two-call protocol. The first call uses a zero-sized + // memref (a NULL buffer) so normal world returns the TA size. + prepare_load_ta_rpc(rpc, ta_uuid, 0, None)?; + return Err(OpteeSmcReturnCode::RpcCmd); + } + let client_identity = ta_req_info.client_identity; let params = &ta_req_info.params; session_manager().with_ta(&ta_uuid, |target| match target { + // A sibling points to an already-loaded single-instance TA, so no + // binary cache lookup or LOAD_TA request is needed on this path. OpenSessionTarget::Sibling(instance) => open_session_single_instance( msg_args, msg_args_phys_addr, @@ -784,13 +827,6 @@ fn open_session_new_instance( ta_req_info: &litebox_shim_optee::msg_handler::TaRequestInfo, ) -> Result<(), OpteeSmcReturnCode> { let shim = litebox_shim_optee::OpteeShimBuilder::new().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(msg_args, msg_args_phys_addr)?; - return Ok(()); - } // Token is declared before `task_pt_guard` so it drops AFTER it. // Marker only releases once CR3 is back to base. See @@ -1318,7 +1354,6 @@ fn write_non_ta_msg_args_to_normal_world( /// 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( msg_args: &OpteeMsgArgs, diff --git a/litebox_shim_optee/src/syscalls/ldelf.rs b/litebox_shim_optee/src/syscalls/ldelf.rs index 5f46d67c9..d5e74c27a 100644 --- a/litebox_shim_optee/src/syscalls/ldelf.rs +++ b/litebox_shim_optee/src/syscalls/ldelf.rs @@ -252,7 +252,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 5576449c95a23b1211868f7faf1e9d4ebadf7e0a Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Thu, 13 Aug 2026 19:56:21 +0000 Subject: [PATCH 5/5] optee: Drop rpc_get_ta_bin In order to send an RPC request, rpc_args and smc_args should be accessible. Access to these objects is not possible from optee shim, so drop the rpc_get_ta_bin placeholder. Signed-off-by: Praveen K Paladugu --- litebox_shim_optee/src/lib.rs | 19 +++---------------- 1 file changed, 3 insertions(+), 16 deletions(-) diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index b83e2f157..ecc937336 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -190,17 +190,9 @@ 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. @@ -232,11 +224,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 = ::RawMutPointer; @@ -315,7 +302,7 @@ 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) }