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 976e5ab26c..e571195dbd 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`. @@ -1431,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`. @@ -1556,6 +1591,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 +1640,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 +2152,35 @@ impl OpteeRpcArgs { } } + /// Access a TMEM output parameter with exact direction and flag validation. + 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]; + let tmem_output = OpteeMsgAttrType::TmemOutput as u64; + // The first LOAD_TA response uses contiguous TMEM to report the binary + // size. SHM_ALLOC may return NONCONTIG TMEM pointing to a page list. + if param.attr.0 != tmem_output && param.attr.0 != tmem_output | OPTEE_MSG_ATTR_NONCONTIG { + 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. + 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.0 != OpteeMsgAttrType::RmemOutput as u64 { + 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 +2195,34 @@ impl OpteeRpcArgs { } } + /// Set a parameter's attribute type by index with bounds checking against `num_params`. + 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`. + 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, @@ -2144,9 +2237,158 @@ impl OpteeRpcArgs { } } - // 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. + /// 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 second LOAD_TA response. + pub fn load_ta_binary_response( + &self, + expected_shm_ref: u64, + expected_size: u64, + ) -> Result { + 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); + } + Ok(rmem) + } + + /// 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.params[0].attr.noncontig() + || 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, + ) -> 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: 8, + }, + )?; + + 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, + 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`. @@ -2229,6 +2471,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; @@ -2239,6 +2493,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; @@ -2253,6 +2508,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, @@ -2302,6 +2558,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 { @@ -2365,6 +2626,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." + ); + } } } } @@ -2499,6 +2765,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}; @@ -2530,6 +2804,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] diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 405b1e4711..35057fad89 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -5,7 +5,7 @@ extern crate alloc; -use alloc::{boxed::Box, vec}; +use alloc::{boxed::Box, sync::Arc, vec}; use core::{ops::Neg, panic::PanicInfo}; use litebox::{ platform::RawConstPointer, @@ -14,19 +14,22 @@ use litebox::{ use litebox_common_linux::{errno::Errno, vmem::PAGE_SIZE}; 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, 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}; 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, @@ -36,13 +39,17 @@ 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::{ + checked_memref_size, decode_ta_request, handle_optee_msg_args, handle_optee_smc_args, + 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::{RpcCompletion, RpcContext, rpc_context_map}, +}; /// The session registry shared by all shims in this runner. fn session_manager() -> &'static SessionManager { @@ -225,7 +232,7 @@ pub fn init(is_bsp: bool) -> &'static Platform { per_cpu_variables::allocate_xsave_area(); if let Err(e) = hvcall::init(is_bsp) { - panic!("Err: {:?}", e); + panic!("Err: {e:?}"); } gdt::init(); interrupts::init_idt(); @@ -447,6 +454,17 @@ impl Drop for TaskPageTableGuard { } } +struct DynamicTaBinaryGuard { + shim: litebox_shim_optee::OpteeShim, + ta_uuid: litebox_common_optee::TeeUuid, +} + +impl Drop for DynamicTaBinaryGuard { + fn drop(&mut self) { + self.shim.remove_ta_bin(&self.ta_uuid); + } +} + /// Switches to base and unregisters the task table. /// /// All user-memory accesses on this core must be complete. @@ -511,48 +529,402 @@ 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, rpc_context) = 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().take(context_id) else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + }; + let common = match &context { + RpcContext::LoadTaSize { common } + | RpcContext::ShmAlloc { common, .. } + | RpcContext::LoadTaBinary { common, .. } + | RpcContext::ShmFree { common, .. } => common, + }; + let result = match read_optee_msg_args_from_regd_shm( + platform, + common.registered_shm_ref, + common.regd_shm_offset, + ) { + Ok((msg_args, Some(rpc_args), msg_args_phys_addr)) => { + Ok(OpteeSmcResult::ReturnFromRpc { + msg_args, + rpc_args, + msg_args_phys_addr, + }) + } + Ok((_msg_args, None, _msg_args_phys_addr)) => Err(OpteeSmcReturnCode::EBadAddr), + Err(error) => Err(error), + }; + if result.is_err() { + if let RpcContext::LoadTaBinary { shm_ref, .. } = &context { + let _ = unregister_rpc_shm(*shm_ref); + } + rpc_context_map().release(context_id); + } + (result, Some((context_id, context))) + } else { + (handle_optee_smc_args(platform, &mut smc_args), None) + }; + let Ok(smc_result) = smc_result else { smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); return *smc_args; }; - if let OpteeSmcResult::CallWithArg { - msg_args, - 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), - 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; + 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, + None, + ), + 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) }; + + 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 Some(ta_uuid) = decode_ta_request(platform, &msg_args) + .ok() + .and_then(|request| request.uuid) + else { + 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); + rpc_context_map().release(context_id); + smc_args.set_return_code(e); + } else { + smc_args.set_return_code(OpteeSmcReturnCode::RpcCmd); + } } else { - msg_args.ret = TeeResult::BadParameters; + debug_serial_println!("OP-TEE SMC returning error code: {:?}", e); + smc_args.set_return_code(e); } - msg_args.ret_origin = TeeOrigin::Tee; - let _ = - write_non_ta_msg_args_to_normal_world(platform, &msg_args, msg_args_phys_addr); - r + } else { + smc_args.set_return_code(OpteeSmcReturnCode::Ok); } - }; + *smc_args + } + OpteeSmcResult::ReturnFromRpc { + msg_args, + rpc_args, + msg_args_phys_addr, + } => { + let mut msg_args = *msg_args; + let mut rpc_args = *rpc_args; + + let Some((context_id, context)) = rpc_context else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + }; + let next_context = match context { + RpcContext::LoadTaSize { common } => { + handle_return_from_load_ta_rpc(common, &mut rpc_args) + } + RpcContext::ShmAlloc { + common, + requested_size, + } => handle_return_from_shm_alloc_rpc( + platform, + common, + requested_size, + &mut rpc_args, + ), + RpcContext::LoadTaBinary { + common, + requested_size, + shm_ref, + } => handle_return_from_load_ta_binary_rpc( + platform, + common, + requested_size, + shm_ref, + &mut rpc_args, + ), + RpcContext::ShmFree { + common, + shm_ref, + completion, + } => { + let return_code = handle_return_from_shm_free_rpc( + platform, + &mut msg_args, + &rpc_args, + msg_args_phys_addr, + &completion, + ); + if return_code == OpteeSmcReturnCode::EThreadLimit { + if rpc_context_map() + .insert( + context_id, + RpcContext::ShmFree { + common, + shm_ref, + completion, + }, + ) + .is_err() + { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + } else { + smc_args.set_return_code(return_code); + } + } else { + rpc_context_map().release(context_id); + smc_args.set_return_code(return_code); + } + return *smc_args; + } + }; + match next_context { + Err(error) => { + rpc_context_map().release(context_id); + smc_args.set_return_code(error); + } + Ok(context) => { + let (RpcContext::ShmAlloc { common, .. } + | RpcContext::LoadTaBinary { common, .. } + | RpcContext::ShmFree { common, .. }) = &context + else { + rpc_context_map().release(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + }; + let cleanup_shm_ref = match &context { + RpcContext::LoadTaBinary { shm_ref, .. } => Some(*shm_ref), + _ => None, + }; + let result = match write_rpc_args_to_regd_shm( + platform, + common.registered_shm_ref, + common.regd_shm_offset, + msg_args.num_params, + &rpc_args, + ) { + Ok(()) => Ok(()), + Err(error) => Err(error), + }; + if let Err(error) = result { + if let Some(shm_ref) = cleanup_shm_ref { + let _ = unregister_rpc_shm(shm_ref); + } + rpc_context_map().release(context_id); + smc_args.set_return_code(error); + } else if rpc_context_map().insert(context_id, context).is_err() { + if let Some(shm_ref) = cleanup_shm_ref { + let _ = unregister_rpc_shm(shm_ref); + } + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + } else { + smc_args.set_return_code(OpteeSmcReturnCode::RpcCmd); + } + } + } + *smc_args + } + _ => smc_result.into(), + } +} - // 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) }; +fn handle_return_from_load_ta_rpc( + common: litebox_shim_optee::rpc_context::RpcCommon, + rpc_args: &mut OpteeRpcArgs, +) -> Result { + let ta_size = rpc_args.load_ta_size_response()?; + 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"); + return Err(OpteeSmcReturnCode::EBadCmd); + } + + rpc_args.prepare_shm_alloc_rpc(OpteeRpcShmType::Appl, ta_size)?; + Ok(RpcContext::ShmAlloc { + common, + requested_size: ta_size, + }) +} - if let Err(e) = result { - smc_args.set_return_code(e); +fn handle_return_from_load_ta_binary_rpc( + platform: &Platform, + common: litebox_shim_optee::rpc_context::RpcCommon, + requested_size: u64, + shm_ref: u64, + rpc_args: &mut OpteeRpcArgs, +) -> Result { + let completion = match rpc_args.load_ta_binary_response(shm_ref, requested_size) { + Ok(rmem) => match checked_memref_size(rmem.size) { + Ok(ta_size) => { + let mut ta_binary = alloc::vec![0u8; ta_size]; + match read_rpc_shm(platform, rmem.shm_ref, 0, &mut ta_binary) { + Ok(()) => RpcCompletion::OpenSession { + ta_binary: Arc::from(ta_binary), + }, + Err(error) => RpcCompletion::ReturnError(error), + } + } + Err(error) => RpcCompletion::ReturnError(error), + }, + Err(error) => RpcCompletion::ReturnError(error), + }; + let _ = unregister_rpc_shm(shm_ref); + rpc_args.prepare_shm_free_rpc(OpteeRpcShmType::Appl, shm_ref)?; + Ok(RpcContext::ShmFree { + common, + shm_ref, + completion, + }) +} + +fn handle_return_from_shm_free_rpc( + platform: &'static Platform, + msg_args: &mut OpteeMsgArgs, + rpc_args: &OpteeRpcArgs, + msg_args_phys_addr: u64, + completion: &RpcCompletion, +) -> OpteeSmcReturnCode { + if rpc_args.cmd != OpteeRpcCommand::ShmFree + || rpc_args.ret != TeeResult::Success + || rpc_args.num_params != 1 + { + return OpteeSmcReturnCode::EBadCmd; + } + + match completion { + RpcCompletion::OpenSession { ta_binary } => { + let mut no_rpc_args = None; + handle_open_session( + platform, + msg_args, + &mut no_rpc_args, + msg_args_phys_addr, + Some(ta_binary.clone()), + ) + .err() + .unwrap_or(OpteeSmcReturnCode::Ok) + } + RpcCompletion::ReturnError(error) => *error, + } +} + +fn handle_return_from_shm_alloc_rpc( + platform: &Platform, + common: litebox_shim_optee::rpc_context::RpcCommon, + requested_size: u64, + rpc_args: &mut OpteeRpcArgs, +) -> Result { + let tmem = rpc_args.shm_alloc_response(requested_size)?; + let mut registered = false; + let load_result = + if checked_memref_size(tmem.size).is_err() || register_rpc_shm(platform, &tmem).is_err() { + Err(OpteeSmcReturnCode::EBadCmd) } else { - smc_args.set_return_code(OpteeSmcReturnCode::Ok); + registered = true; + rpc_args.prepare_load_ta_rpc( + common.ta_uuid, + Some(OpteeMsgParamRmem { + offs: 0, + size: requested_size, + shm_ref: tmem.shm_ref, + }), + ) + }; + match load_result { + Ok(()) => Ok(RpcContext::LoadTaBinary { + common, + requested_size, + shm_ref: tmem.shm_ref, + }), + Err(error) => { + if registered { + let _ = unregister_rpc_shm(tmem.shm_ref); + } + rpc_args.prepare_shm_free_rpc(OpteeRpcShmType::Appl, tmem.shm_ref)?; + Ok(RpcContext::ShmFree { + common, + shm_ref: tmem.shm_ref, + completion: RpcCompletion::ReturnError(error), + }) } - *smc_args - } else { - smc_result.into() } } @@ -568,7 +940,9 @@ 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, + ta_binary: Option>, ) -> Result<(), OpteeSmcReturnCode> { let ta_req_info = decode_ta_request(platform, msg_args).map_err(|_| OpteeSmcReturnCode::EBadCmd)?; @@ -594,10 +968,12 @@ fn handle_open_session( platform, msg_args, msg_args_phys_addr, + rpc_args, params, ta_uuid, client_identity, &ta_req_info, + ta_binary, ), OpenSessionTarget::Busy => { // Single-instance TA without MULTI_SESSION already has a live @@ -775,24 +1151,40 @@ 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, + ta_binary: Option>, ) -> Result<(), OpteeSmcReturnCode> { let shim = litebox_shim_optee::OpteeShimBuilder::new(platform, session_manager()) .with_ta_signing_cert(TA_SIGNING_CERT_DER) .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 ta_binary.is_none() && !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. @@ -825,6 +1217,17 @@ fn open_session_new_instance( OpteeSmcReturnCode::ENotAvail })?; + if let Some(ta_binary) = ta_binary + && !shim.store_ta_bin(&ta_uuid, &ta_binary) + { + teardown_ta_page_table(platform, task_pt_id); + return Err(OpteeSmcReturnCode::EBadCmd); + } + let _ta_binary_guard = DynamicTaBinaryGuard { + shim: shim.clone(), + ta_uuid, + }; + // Load ldelf and TA - Box immediately to keep at fixed heap address let loaded_program = Box::new(shim.load_ldelf(LDELF_BINARY, ta_uuid).map_err(|_| { teardown_ta_page_table(platform, task_pt_id); @@ -1329,38 +1732,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]; @@ -1375,7 +1746,7 @@ fn register_embedded_ta( let Some(ta_head) = litebox_common_optee::parse_ta_head(ta_binary) else { return false; }; - shim.store_ta_bin(&ta_head.uuid, ta_binary) + shim.store_embedded_ta_bin(&ta_head.uuid, ta_binary) } /// Register all TA binaries embedded in the runner image. diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index 6e90e136f5..dd8e0639b9 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -30,6 +30,7 @@ use litebox_common_optee::{ use sha2::{Digest, Sha256}; pub mod loader; +pub mod rpc_context; pub mod session; pub(crate) mod syscalls; @@ -229,7 +230,8 @@ impl OpteeShimBuilder { boot_instant: TimeProvider::now(self.platform), mm: MemoryManager::new(self.platform), _litebox: self.litebox, - ta_uuid_map: ta_uuid_map(), + embedded_ta_uuid_map: embedded_ta_uuid_map(), + dynamic_ta_uuid_map: TaUuidMap::new(), ta_signing_cert: self.ta_signing_cert, pta_busy: spin::mutex::SpinMutex::new(HashSet::new()), page_table_keepalive: None, @@ -252,8 +254,10 @@ struct GlobalState { mm: MemoryManager, /// The LiteBox instance used throughout the shim. _litebox: litebox::LiteBox, - /// The TA UUID to binary map for TA loading. - ta_uuid_map: &'static TaUuidMap, + /// TA binaries embedded in the runner image and shared across shim instances. + embedded_ta_uuid_map: &'static TaUuidMap, + /// Dynamically supplied TA binaries scoped to this shim instance. + dynamic_ta_uuid_map: TaUuidMap, ta_signing_cert: &'static [u8], /// Tracks which non-concurrent PTAs (i.e., PTAs w/o `TaFlags::CONCURRENT`) /// are currently busy. A busy PTA is *rejected* with `TeeResult::Busy` @@ -274,20 +278,19 @@ impl GlobalState { /// Returns `true` if the binary was successfully stored, `false` if the binary's /// UUID (from `.ta_head` section) doesn't match the provided UUID or parsing failed. pub(crate) fn store_ta_bin(&self, ta_uuid: &TeeUuid, ta_bin: &[u8]) -> bool { - self.ta_uuid_map.insert(*ta_uuid, ta_bin.into()) + self.dynamic_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.dynamic_ta_uuid_map + .get(ta_uuid) + .or_else(|| self.embedded_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.dynamic_ta_uuid_map.contains(ta_uuid) || self.embedded_ta_uuid_map.contains(ta_uuid) } /// Monotonic time elapsed since this instance was created, used as GP @@ -299,20 +302,9 @@ 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); - } - - /// 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 + let _ = self.dynamic_ta_uuid_map.remove(ta_uuid); } } @@ -363,8 +355,9 @@ impl OpteeShim { ) -> Result, loader::elf::ElfLoaderError> { let (ta_binary, ta_flags) = self .0 - .ta_uuid_map + .dynamic_ta_uuid_map .get_with_flags(&ta_uuid) + .or_else(|| self.0.embedded_ta_uuid_map.get_with_flags(&ta_uuid)) .ok_or(loader::elf::ElfLoaderError::OpenError(Errno::ENOENT))?; let ta_digest = Sha256::digest(&ta_binary).into(); let entrypoints = crate::OpteeShimEntrypoints { @@ -436,7 +429,12 @@ impl OpteeShim { self.0.store_ta_bin(ta_uuid, ta_bin) } - /// Get the TA binary associated with the given TA UUID. + /// Store a TA binary embedded in the runner image. + pub fn store_embedded_ta_bin(&self, ta_uuid: &TeeUuid, ta_bin: &'static [u8]) -> bool { + self.0.embedded_ta_uuid_map.insert(*ta_uuid, ta_bin.into()) + } + + /// 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) } @@ -456,6 +454,16 @@ impl OpteeShim { let _ = self.memory_manager().release_memory(release); } } + + /// 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) + } + + /// Remove a TA binary from the trusted cache. + pub fn remove_ta_bin(&self, ta_uuid: &TeeUuid) { + self.0.remove_ta_bin(ta_uuid); + } } impl OpteeShimEntrypoints { @@ -1518,6 +1526,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) + } + fn get_with_flags(&self, uuid: &TeeUuid) -> Option<(Arc<[u8]>, TaFlags)> { self.inner .read() @@ -1531,10 +1543,11 @@ impl TaUuidMap { } } -/// Get the global TA UUID map. -fn ta_uuid_map() -> &'static TaUuidMap { - static TA_UUID_MAP: once_cell::race::OnceBox = once_cell::race::OnceBox::new(); - TA_UUID_MAP.get_or_init(|| alloc::boxed::Box::new(TaUuidMap::new())) +/// Get the TA binaries embedded in the runner image. +fn embedded_ta_uuid_map() -> &'static TaUuidMap { + static EMBEDDED_TA_UUID_MAP: once_cell::race::OnceBox = + once_cell::race::OnceBox::new(); + EMBEDDED_TA_UUID_MAP.get_or_init(|| alloc::boxed::Box::new(TaUuidMap::new())) } /// Per-instance TA state which can be shared between sessions if it is diff --git a/litebox_shim_optee/src/msg_handler.rs b/litebox_shim_optee/src/msg_handler.rs index f06bd28fe4..3e7a34f4bf 100644 --- a/litebox_shim_optee/src/msg_handler.rs +++ b/litebox_shim_optee/src/msg_handler.rs @@ -79,11 +79,11 @@ 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); } - Ok(size.trunc()) + usize::try_from(size).map_err(|_| OpteeSmcReturnCode::ENomem) } fn parse_optee_msg_args( @@ -207,6 +207,85 @@ 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() +} + +/// 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 @@ -243,41 +322,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 => { @@ -934,6 +989,41 @@ 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)?; + if rpc_args_offset + .checked_add(blob.len()) + .is_none_or(|end| end > shm_info.len) + { + return Err(OpteeSmcReturnCode::EBadAddr); + } + if blob.is_empty() { + return Ok(()); + } + let ptr = NormalWorldMutPtr::::new( + platform, + &shm_info.page_addrs, + shm_info.page_offset, + )?; + ptr.write_slice_at_offset(rpc_args_offset, &blob)?; + Ok(()) +} + 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 new file mode 100644 index 0000000000..b3241cffba --- /dev/null +++ b/litebox_shim_optee/src/rpc_context.rs @@ -0,0 +1,237 @@ +// 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, sync::Arc}; +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, + UnexpectedStage, +} + +/// Action to take after an in-flight shared-memory free RPC returns. +#[derive(Clone, Debug, PartialEq)] +pub enum RpcCompletion { + OpenSession { ta_binary: Arc<[u8]> }, + ReturnError(OpteeSmcReturnCode), +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct RpcCommon { + pub 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. + pub registered_shm_ref: u64, + pub regd_shm_offset: usize, +} + +/// Continuation state for an RPC-backed Dynamic TA request. +#[derive(Clone, 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, + }, + } + } +} + +struct RpcContexts { + ids: IdPool, + contexts: HashMap, +} + +/// Maps RPC context IDs to continuation state. +pub struct RpcContextMap { + inner: SpinMutex, +} + +impl RpcContextMap { + pub fn new() -> Self { + Self::with_capacity(MAX_RPC_CONTEXTS) + } + + fn with_capacity(capacity: u32) -> Self { + Self { + inner: SpinMutex::new(RpcContexts { + ids: IdPool::with_capacity(capacity), + contexts: HashMap::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) + } + + /// Remove a context from the map while keeping its ID reserved. + pub fn take(&self, context_id: u32) -> Option { + self.inner.lock().contexts.remove(&context_id) + } + + /// Reinsert a context before returning for another RPC or retryable result. + pub fn insert(&self, context_id: u32, context: RpcContext) -> Result<(), RpcContextError> { + let mut inner = self.inner.lock(); + if inner.contexts.contains_key(&context_id) { + return Err(RpcContextError::UnexpectedStage); + } + inner.contexts.insert(context_id, context); + Ok(()) + } + + /// Recycle a context ID after its RPC sequence finishes. + pub fn release(&self, context_id: u32) { + let mut inner = self.inner.lock(); + debug_assert!(!inner.contexts.contains_key(&context_id)); + inner.ids.recycle(context_id); + } +} + +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 taken_context_keeps_its_id_reserved() { + let contexts = RpcContextMap::with_capacity(1); + let first_id = contexts.allocate(test_uuid(1), 1, 0).unwrap(); + let first = contexts.take(first_id).unwrap(); + + assert_eq!( + contexts.allocate(test_uuid(2), 2, 0), + Err(RpcContextError::Full) + ); + + contexts.insert(first_id, first.clone()).unwrap(); + + assert_eq!(contexts.take(first_id), Some(first)); + contexts.release(first_id); + assert_eq!(contexts.allocate(test_uuid(3), 3, 0).unwrap(), first_id); + } +} diff --git a/litebox_shim_optee/src/syscalls/ldelf.rs b/litebox_shim_optee/src/syscalls/ldelf.rs index c3da561983..dc58dd2063 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);