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..93c8b774fb 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,51 @@ 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) + } + + /// 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, + 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 +2211,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, @@ -2144,9 +2253,160 @@ 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 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, + 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 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 +2489,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 +2511,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 +2526,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 +2576,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 +2644,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 +2783,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 +2822,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/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 442eba6ebc..91645a9293 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, 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, @@ -37,13 +40,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 { @@ -542,48 +549,523 @@ 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, - 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) + } + 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); + 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 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, + ); + } + 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 + } + _ => 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; + } - // 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 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); + } +} - if let Err(e) = result { - smc_args.set_return_code(e); - } else { - smc_args.set_return_code(OpteeSmcReturnCode::Ok); +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)); } - *smc_args + // 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, + 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() { + 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() { + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + tmem.shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ); + return; + } + if rpc_context_map() + .transition_to_load_ta_binary(context_id, tmem.shm_ref) + .is_err() + { + 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; + } + 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() + { + 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 !write_next_rpc(platform, smc_args, msg_args, rpc_args, 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), + ); + } +} + +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 { + 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 RpcContext::LoadTaBinary { shm_ref, .. } = context + { + let _ = unregister_rpc_shm(shm_ref); } } @@ -599,6 +1081,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 +1108,7 @@ fn handle_open_session( platform, msg_args, msg_args_phys_addr, + rpc_args, params, ta_uuid, client_identity, @@ -814,22 +1298,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. @@ -1377,38 +1876,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/lib.rs b/litebox_shim_optee/src/lib.rs index 4d96ccd24b..00e51a6b08 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; @@ -224,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. @@ -251,21 +249,10 @@ 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 - } } type UserMutPtr = @@ -364,11 +351,21 @@ 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) + } + + /// 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 @@ -1446,6 +1443,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/msg_handler.rs b/litebox_shim_optee/src/msg_handler.rs index c368099e41..7b9ec036c8 100644 --- a/litebox_shim_optee/src/msg_handler.rs +++ b/litebox_shim_optee/src/msg_handler.rs @@ -80,11 +80,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( @@ -208,6 +208,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 @@ -244,41 +323,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 => { @@ -758,6 +813,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 +1016,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 new file mode 100644 index 0000000000..ca52b5a66d --- /dev/null +++ b/litebox_shim_optee/src/rpc_context.rs @@ -0,0 +1,402 @@ +// 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 + } + + 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); + } +} 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);