diff --git a/dev_tests/src/ratchet.rs b/dev_tests/src/ratchet.rs index 41dc17bb9..5c6c76f4c 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 976e5ab26..801ea384b 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,54 @@ 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. + pub 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 +2214,34 @@ impl OpteeRpcArgs { } } + /// Set a parameter's attribute type by index with bounds checking against `num_params`. + pub fn set_param_attr_type( + &mut self, + index: usize, + attr_type: OpteeMsgAttrType, + ) -> Result<(), OpteeSmcReturnCode> { + if index >= self.num_params as usize { + Err(OpteeSmcReturnCode::ENotAvail) + } else { + self.params[index].attr = OpteeMsgAttr(attr_type as u64); + Ok(()) + } + } + + /// Set an rmem parameter by index with bounds checking against `num_params`. + pub fn set_param_rmem( + &mut self, + index: usize, + rmem: OpteeMsgParamRmem, + ) -> Result<(), OpteeSmcReturnCode> { + if index >= self.num_params as usize { + Err(OpteeSmcReturnCode::ENotAvail) + } else { + self.params[index].data.copy_from_slice(rmem.as_bytes()); + Ok(()) + } + } + /// Set a tmem parameter by index with bounds checking against `num_params`. pub fn set_param_tmem( &mut self, @@ -2143,10 +2255,102 @@ impl OpteeRpcArgs { Ok(()) } } +} + +/// Prepare a shared-memory allocation RPC request to be sent to normal world. +pub fn prepare_shm_alloc_rpc( + rpc_msg_args: &mut OpteeRpcArgs, + shm_type: OpteeRpcShmType, + size: u64, + alignment: u64, +) -> Result<(), OpteeSmcReturnCode> { + rpc_msg_args.cmd = OpteeRpcCommand::ShmAlloc; + // Match OP-TEE's get_rpc_arg(): default to failure in case normal world + // returns without updating the RPC result. + rpc_msg_args.ret = TeeResult::GenericError; + rpc_msg_args.num_params = 1; + + rpc_msg_args + .set_param_attr_type(0, OpteeMsgAttrType::ValueInput) + .map_err(|_| OpteeSmcReturnCode::EBadCmd)?; + + rpc_msg_args.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. +/// This will free the memory allocated by SHM_ALLOC request in normal world. +pub fn prepare_shm_free_rpc( + rpc_msg_args: &mut OpteeRpcArgs, + shm_type: OpteeRpcShmType, + shm_ref: u64, +) -> Result<(), OpteeSmcReturnCode> { + rpc_msg_args.cmd = OpteeRpcCommand::ShmFree; + rpc_msg_args.ret = TeeResult::GenericError; + rpc_msg_args.num_params = 1; + rpc_msg_args.set_param_attr_type(0, OpteeMsgAttrType::ValueInput)?; + rpc_msg_args.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( + rpc_msg_args: &mut OpteeRpcArgs, + ta_uuid: TeeUuid, + memref: Option, +) -> Result<(), OpteeSmcReturnCode> { + rpc_msg_args.cmd = OpteeRpcCommand::LoadTa; + // Match OP-TEE's get_rpc_arg(): default to failure in case normal world + // returns without updating the RPC result. + rpc_msg_args.ret = TeeResult::GenericError; + rpc_msg_args.num_params = 2; + + rpc_msg_args.set_param_attr_type(0, OpteeMsgAttrType::ValueInput)?; + let uuid_bytes = ta_uuid.to_u64_array(); + rpc_msg_args.set_param_value( + 0, + OpteeMsgParamValue { + a: uuid_bytes[0], + b: uuid_bytes[1], + c: 0, + }, + )?; + + if memref.is_none() { + // First call of LOAD_TA protocol: normal world returns the TA size in `tmem.size`. + rpc_msg_args.set_param_attr_type(1, OpteeMsgAttrType::TmemOutput)?; + rpc_msg_args.set_param_tmem( + 1, + OpteeMsgParamTmem { + buf_ptr: 0, + size: 0, + shm_ref: 0, + }, + )?; + } else { + // Second LOAD_TA: LiteBox passes VTL0-owned memory back to VTL0 + // so normal world can populate it with the TA binary. + rpc_msg_args.set_param_attr_type(1, OpteeMsgAttrType::RmemOutput)?; + if let Some(rmem) = memref { + rpc_msg_args.set_param_rmem(1, rmem)?; + } + } - // 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. + Ok(()) } /// Serialize the params portion as raw bytes into `buf`. @@ -2229,6 +2433,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 +2455,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 +2470,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 +2520,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 +2588,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 +2727,23 @@ pub const HUK_SUBKEY_MAX_LEN: usize = 32; mod tests { use super::*; + #[test] + fn test_rpc_context_id_roundtrip() { + for context_id in [0, 1, u32::MAX] { + let mut args = OpteeSmcArgs::default(); + args.set_rpc_context_id(context_id); + assert_eq!(args.get_rpc_context_id(), Ok(context_id)); + } + } + + #[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 +2775,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] @@ -2629,6 +2897,137 @@ mod tests { assert_eq!(header_out.num_params, 2); } + #[test] + fn test_optee_rpc_args_attr_and_rmem_setters() { + let header = OpteeMsgArgsHeader { + cmd: OpteeRpcCommand::LoadTa as u32, + func: 0, + session: 0, + cancel_id: 0, + pad: 0, + ret: 0, + ret_origin: 0, + num_params: 1, + }; + let raw_params = [0u8; size_of::()]; + let mut rpc_args = OpteeRpcArgs::from_header_and_raw_params(&header, &raw_params) + .expect("should parse RPC args"); + + rpc_args.params[0].attr = OpteeMsgAttr::META_VALUE_INPUT; + rpc_args + .set_param_attr_type(0, OpteeMsgAttrType::RmemOutput) + .expect("attribute index should be available"); + assert_eq!( + rpc_args.params[0].attr.attr_type(), + OpteeMsgAttrType::RmemOutput as u8 + ); + assert!(!rpc_args.params[0].attr.meta()); + assert!(!rpc_args.params[0].attr.noncontig()); + + let rmem = OpteeMsgParamRmem { + offs: 0x0102_0304_0506_0708, + size: 0x1112_1314_1516_1718, + shm_ref: 0x2122_2324_2526_2728, + }; + rpc_args + .set_param_rmem(0, rmem) + .expect("rmem index should be available"); + assert_eq!(&rpc_args.params[0].data[0..8], &rmem.offs.to_le_bytes()); + assert_eq!(&rpc_args.params[0].data[8..16], &rmem.size.to_le_bytes()); + assert_eq!( + &rpc_args.params[0].data[16..24], + &rmem.shm_ref.to_le_bytes() + ); + + assert_eq!( + rpc_args.set_param_attr_type(1, OpteeMsgAttrType::RmemOutput), + Err(OpteeSmcReturnCode::ENotAvail) + ); + assert_eq!( + rpc_args.set_param_rmem(1, rmem), + Err(OpteeSmcReturnCode::ENotAvail) + ); + } + + #[test] + fn test_optee_rpc_args_exact_output_getters() { + let header = OpteeMsgArgsHeader { + cmd: OpteeRpcCommand::LoadTa as u32, + func: 0, + session: 0, + cancel_id: 0, + pad: 0, + ret: 0, + ret_origin: 0, + num_params: 1, + }; + let raw_params = [0u8; size_of::()]; + let mut rpc_args = OpteeRpcArgs::from_header_and_raw_params(&header, &raw_params) + .expect("should parse RPC args"); + + rpc_args.params[0].attr = OpteeMsgAttr(OpteeMsgAttrType::TmemOutput as u64); + assert!(rpc_args.get_param_tmem_output(0).is_ok()); + assert_eq!(rpc_args.is_param_tmem_output_noncontiguous(0), Ok(false)); + rpc_args.params[0].attr = + OpteeMsgAttr(OpteeMsgAttrType::TmemOutput as u64 | OPTEE_MSG_ATTR_NONCONTIG); + assert!(rpc_args.get_param_tmem_output(0).is_ok()); + assert_eq!(rpc_args.is_param_tmem_output_noncontiguous(0), Ok(true)); + + rpc_args.params[0].attr = OpteeMsgAttr(OpteeMsgAttrType::TmemInout as u64); + assert!(matches!( + rpc_args.get_param_tmem_output(0), + Err(OpteeSmcReturnCode::EBadCmd) + )); + rpc_args.params[0].attr = + OpteeMsgAttr(OpteeMsgAttrType::TmemOutput as u64 | OPTEE_MSG_ATTR_META); + assert!(matches!( + rpc_args.get_param_tmem_output(0), + Err(OpteeSmcReturnCode::EBadCmd) + )); + + rpc_args.params[0].attr = OpteeMsgAttr(OpteeMsgAttrType::RmemOutput as u64); + assert!(rpc_args.get_param_rmem_output(0).is_ok()); + rpc_args.params[0].attr = + OpteeMsgAttr(OpteeMsgAttrType::RmemOutput as u64 | OPTEE_MSG_ATTR_NONCONTIG); + assert!(matches!( + rpc_args.get_param_rmem_output(0), + Err(OpteeSmcReturnCode::EBadCmd) + )); + assert!(matches!( + rpc_args.get_param_rmem_output(1), + Err(OpteeSmcReturnCode::ENotAvail) + )); + } + #[test] + fn test_prepare_shm_free_rpc() { + let header = OpteeMsgArgsHeader { + cmd: OpteeRpcCommand::LoadTa as u32, + func: 0, + session: 0, + cancel_id: 0, + pad: 0, + ret: TeeResult::Success as u32, + ret_origin: 0, + num_params: 1, + }; + let raw_params = [0u8; size_of::()]; + let mut rpc_args = OpteeRpcArgs::from_header_and_raw_params(&header, &raw_params) + .expect("should parse RPC args"); + + prepare_shm_free_rpc(&mut rpc_args, OpteeRpcShmType::Appl, 0x1234) + .expect("should prepare SHM_FREE"); + + assert_eq!(rpc_args.cmd, OpteeRpcCommand::ShmFree); + assert_eq!(rpc_args.ret, TeeResult::GenericError); + assert_eq!(rpc_args.num_params, 1); + let value = rpc_args + .get_param_value(0) + .expect("SHM_FREE value parameter should be readable"); + assert_eq!(value.a, OpteeRpcShmType::Appl as u64); + assert_eq!(value.b, 0x1234); + assert_eq!(value.c, 0); + } + #[test] fn test_rpc_args_rejects_main_cmd() { // Pick a cmd value that lies in the gap between Plugin (12) and I2C Transfer (21), diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 442eba6eb..30be9e708 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -15,8 +15,10 @@ 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, prepare_load_ta_rpc, + prepare_shm_alloc_rpc, prepare_shm_free_rpc, }; use litebox_platform_lvbs::host::LvbsLinuxKernel as Platform; use litebox_platform_lvbs::mshv::vsm::{LvbsVtl0Gate, LvbsVtl0PrivilegedWriter, LvbsVtl1Gate}; @@ -39,11 +41,16 @@ use litebox_platform_lvbs::{ }, 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, RpcStage, rpc_context_map}, +}; /// The session registry shared by all shims in this runner. fn session_manager() -> &'static SessionManager { @@ -542,48 +549,641 @@ 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 is_return_from_rpc = smc_args.func_id() == Ok(OpteeSmcFunction::ReturnFromRpc); + let smc_result = if is_return_from_rpc { + 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(registered_shm_ref) = rpc_context_map().get_registered_shm_ref(context_id) else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + }; + let Some(regd_shm_offset) = rpc_context_map().get_regd_shm_offset(context_id) else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + }; + read_optee_msg_args_from_regd_shm(platform, registered_shm_ref, regd_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, + }) + }, + ) + } else { + handle_optee_smc_args(platform, &mut smc_args) + }; + let Ok(smc_result) = smc_result else { + if is_return_from_rpc && let Ok(context_id) = smc_args.get_rpc_context_id() { + discard_rpc_context(context_id); + } 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"); + + // CallWithArg is not supported for Dynamic TA RPC continuation. + // This applies to how shm_ref is interpreted and used during RPC calls. + 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; + }; + // RPC continuation reuses args[3] for the context ID. Preserve the + // request identity in trusted state before overwriting it. + let (registered_shm_ref, regd_shm_offset) = + match smc_args.optee_regd_shm_ref_and_offset() { + Ok(location) => location, + Err(error) => { + smc_args.set_return_code(error); + return *smc_args; + } + }; + let ta_uuid = match decode_ta_request(platform, &msg_args) + .ok() + .and_then(|request| request.uuid) + { + Some(ta_uuid) => ta_uuid, + None => { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + } + }; + let context_id = match rpc_context_map().allocate( + RpcStage::LoadTaSize, + 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(curr_stage) = rpc_context_map().get_curr_stage(context_id) else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return *smc_args; + }; + match curr_stage { + RpcStage::LoadTaSize => { + handle_return_from_load_ta_rpc( + platform, + &mut smc_args, + &msg_args, + &mut rpc_args, + ); + } + RpcStage::ShmAlloc => { + handle_return_from_shm_alloc_rpc( + platform, + &mut smc_args, + &msg_args, + &mut rpc_args, + ); + } + RpcStage::LoadTaBinary => { + handle_return_from_load_ta_binary_rpc( + platform, + &mut smc_args, + &msg_args, + &mut rpc_args, + ); + } + RpcStage::ShmFree => { + handle_return_from_shm_free_rpc( + platform, + &mut smc_args, + &mut msg_args, + &rpc_args, + msg_args_phys_addr, + ); + } + } + *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( + platform: &Platform, + smc_args: &mut OpteeSmcArgs, + msg_args: &OpteeMsgArgs, + rpc_args: &mut OpteeRpcArgs, +) { + let context_id = match smc_args.get_rpc_context_id() { + Ok(context_id) => context_id, + Err(error) => { + smc_args.set_return_code(error); + return; + } + }; + if rpc_args.cmd != OpteeRpcCommand::LoadTa + || rpc_args.ret != TeeResult::Success + || rpc_args.num_params != 2 + { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + let tmem = match rpc_args.get_param_tmem_output(1) { + Ok(tmem) => tmem, + Err(error) => { + discard_rpc_context(context_id); + smc_args.set_return_code(error); + return; + } + }; + let ta_size = tmem.size; + debug_serial_println!("First LOAD_TA request, TA size: {}", ta_size); + if ta_size == 0 || checked_memref_size(ta_size).is_err() { + debug_serial_println!("Invalid TA size in first LOAD_TA request"); + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } - if let Err(e) = result { - smc_args.set_return_code(e); - } else { - smc_args.set_return_code(OpteeSmcReturnCode::Ok); + if rpc_context_map() + .set_requested_size(context_id, RpcStage::LoadTaSize, ta_size) + .is_err() + { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + if prepare_shm_alloc_rpc(rpc_args, 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(context_id, RpcStage::LoadTaSize, RpcStage::ShmAlloc) + .is_err() + { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + if !write_next_rpc(platform, smc_args, msg_args, rpc_args, context_id) { + discard_rpc_context(context_id); + } +} + +fn handle_return_from_load_ta_binary_rpc( + platform: &'static Platform, + smc_args: &mut OpteeSmcArgs, + msg_args: &OpteeMsgArgs, + rpc_args: &mut OpteeRpcArgs, +) { + let context_id = match smc_args.get_rpc_context_id() { + Ok(context_id) => context_id, + Err(error) => { + smc_args.set_return_code(error); + return; + } + }; + let Some(shm_ref) = rpc_context_map().get_shm_ref(context_id) else { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + }; + let response = (|| { + if rpc_args.cmd != OpteeRpcCommand::LoadTa + || rpc_args.ret != TeeResult::Success + || rpc_args.num_params != 2 + { + return Err(OpteeSmcReturnCode::EBadCmd); + } + let rmem = rpc_args.get_param_rmem_output(1)?; + let requested_size = rpc_context_map() + .get_requested_size(context_id) + .ok_or(OpteeSmcReturnCode::EBadCmd)?; + if rmem.shm_ref != shm_ref || rmem.offs != 0 || rmem.size != requested_size { + return Err(OpteeSmcReturnCode::EBadCmd); + } + let ta_size = checked_memref_size(rmem.size)?; + let mut ta_binary = alloc::vec![0u8; ta_size]; + read_rpc_shm(platform, shm_ref, 0, &mut ta_binary)?; + Ok(ta_binary) + })(); + + let completion = match response { + Ok(ta_binary) => { + let Some(ta_uuid) = rpc_context_map().get_ta_uuid(context_id) else { + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + RpcStage::LoadTaBinary, + shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ); + return; + }; + let shim = + litebox_shim_optee::OpteeShimBuilder::new(platform, session_manager()).build(); + if !shim.store_ta_bin(&ta_uuid, &ta_binary) { + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd) + } else { + RpcCompletion::OpenSession + } + } + Err(error) => RpcCompletion::ReturnError(error), + }; + let ta_uuid = rpc_context_map().get_ta_uuid(context_id); + if !start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + RpcStage::LoadTaBinary, + shm_ref, + completion, + ) && completion == RpcCompletion::OpenSession + && let Some(ta_uuid) = ta_uuid + { + litebox_shim_optee::OpteeShimBuilder::new(platform, session_manager()) + .build() + .remove_ta_bin(&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, + expected_stage: RpcStage, + shm_ref: u64, + completion: RpcCompletion, +) -> bool { + if rpc_context_map().is_local_shm_registered(context_id) == Some(true) { + let _ = unregister_rpc_shm(shm_ref); + } + if prepare_shm_free_rpc(rpc_args, OpteeRpcShmType::Appl, shm_ref).is_err() + || rpc_context_map() + .set_completion(context_id, expected_stage, completion) + .is_err() + || rpc_context_map() + .transition(context_id, expected_stage, RpcStage::ShmFree) + .is_err() + { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return false; + } + 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, + smc_args: &mut OpteeSmcArgs, + msg_args: &mut OpteeMsgArgs, + rpc_args: &OpteeRpcArgs, + msg_args_phys_addr: u64, +) { + let context_id = match smc_args.get_rpc_context_id() { + Ok(context_id) => context_id, + Err(error) => { + smc_args.set_return_code(error); + return; + } + }; + let Some(context) = rpc_context_map().take(context_id) else { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + }; + if context.stage() != RpcStage::ShmFree { + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + match context.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, + smc_args: &mut OpteeSmcArgs, + msg_args: &OpteeMsgArgs, + rpc_args: &mut OpteeRpcArgs, +) { + let context_id = match smc_args.get_rpc_context_id() { + Ok(context_id) => context_id, + Err(error) => { + smc_args.set_return_code(error); + return; + } + }; + if rpc_args.cmd != OpteeRpcCommand::ShmAlloc + || rpc_args.ret != TeeResult::Success + || rpc_args.num_params != 1 + { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + let tmem = match rpc_args.get_param_tmem_output(0) { + Ok(tmem) => tmem, + Err(_) => { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + }; + if tmem.shm_ref == 0 { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + if rpc_context_map() + .set_shm_ref(context_id, RpcStage::ShmAlloc, tmem.shm_ref) + .is_err() + { + discard_rpc_context(context_id); + smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); + return; + } + match rpc_args.is_param_tmem_output_noncontiguous(0) { + Ok(true) => {} + Ok(false) | Err(_) => { + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + RpcStage::ShmAlloc, + tmem.shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::ENotAvail), + ); + return; + } + } + let Some(requested_size) = rpc_context_map().get_requested_size(context_id) else { + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + RpcStage::ShmAlloc, + tmem.shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ); + return; + }; + if tmem.buf_ptr == 0 || tmem.size < requested_size || checked_memref_size(tmem.size).is_err() { + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + RpcStage::ShmAlloc, + 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, + RpcStage::ShmAlloc, + tmem.shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ); + return; + } + if rpc_context_map() + .set_local_shm_registered(context_id, RpcStage::ShmAlloc) + .is_err() + { + let _ = unregister_rpc_shm(tmem.shm_ref); + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + RpcStage::ShmAlloc, + tmem.shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ); + return; + } + let Some(ta_uuid) = rpc_context_map().get_ta_uuid(context_id) else { + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + RpcStage::ShmAlloc, + tmem.shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ); + return; + }; + let rmem = OpteeMsgParamRmem { + offs: 0, + size: requested_size, + shm_ref: tmem.shm_ref, + }; + if prepare_load_ta_rpc(rpc_args, ta_uuid, Some(rmem)).is_err() { + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + RpcStage::ShmAlloc, + tmem.shm_ref, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ); + return; + } + if rpc_context_map() + .transition(context_id, RpcStage::ShmAlloc, RpcStage::LoadTaBinary) + .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) { + start_shm_free_rpc( + platform, + smc_args, + msg_args, + rpc_args, + context_id, + RpcStage::LoadTaBinary, + 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 location = rpc_context_map() + .get_registered_shm_ref(context_id) + .zip(rpc_context_map().get_regd_shm_offset(context_id)); + let result = location.ok_or(OpteeSmcReturnCode::EBadCmd).and_then( + |(registered_shm_ref, regd_shm_offset)| { + write_rpc_args_to_regd_shm( + platform, + registered_shm_ref, + regd_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) + && context.local_shm_registered() + && let Some(shm_ref) = context.shm_ref() + { + let _ = unregister_rpc_shm(shm_ref); } } @@ -599,6 +1199,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 +1226,7 @@ fn handle_open_session( platform, msg_args, msg_args_phys_addr, + rpc_args, params, ta_uuid, client_identity, @@ -818,18 +1420,34 @@ 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(()); + }; + // LOAD_TA is a two-call protocol. By passing `None`, we indicate that + // this is the first call, and the normal world should return TA binary size. + prepare_load_ta_rpc(rpc, ta_uuid, None)?; + return Err(OpteeSmcReturnCode::RpcCmd); } // Token is declared before `task_pt_guard` so it drops AFTER it. @@ -1377,40 +1995,9 @@ 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 LDELF_BINARY: &[u8] = + include_bytes!("../../litebox_runner_optee_on_linux_userland/tests/ldelf.elf"); const TA_BINARY: &[u8] = &[0u8; 0]; const TA_BINARIES: &[&[u8]] = &[TA_BINARY]; diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index 4d96ccd24..00e51a6b0 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 c368099e4..2feb960a2 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,38 @@ 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, rpc_args_size) = + rpc_args_range(msg_args_offset, msg_args_num_params, 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 rpc_args_range( + msg_args_offset: usize, + msg_args_num_params: u32, + rpc_args_num_params: u32, +) -> Result<(usize, usize), OpteeSmcReturnCode> { + let rpc_args_offset = msg_args_offset + .checked_add(optee_msg_args_total_size(msg_args_num_params)) + .ok_or(OpteeSmcReturnCode::EBadAddr)?; + Ok(( + rpc_args_offset, + optee_msg_args_total_size(rpc_args_num_params), + )) +} + fn shm_ref_map() -> &'static ShmRefMap { static SHM_REF_MAP: OnceBox> = OnceBox::new(); SHM_REF_MAP.get_or_init(|| Box::new(ShmRefMap::new())) @@ -986,7 +1099,7 @@ fn get_shm_info_from_optee_msg_param_tmem( /// /// `rmem.offs` must be an offset within the shared memory region registered with `rmem.shm_ref` before /// and `rmem.offs + rmem.size` must not exceed the size of the registered shared memory region. -fn get_shm_info_from_optee_msg_param_rmem( +pub fn get_shm_info_from_optee_msg_param_rmem( rmem: OpteeMsgParamRmem, ) -> Result, OpteeSmcReturnCode> { let Some(shm_info) = shm_ref_map().get(rmem.shm_ref) else { @@ -1022,3 +1135,28 @@ fn get_shm_info_from_optee_msg_param_rmem( rmem.size.trunc(), ) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rpc_args_range_can_cross_page_boundary() { + let main_size = optee_msg_args_total_size(2); + let rpc_size = optee_msg_args_total_size(2); + let msg_args_offset = PAGE_SIZE - main_size - rpc_size / 2; + + let (offset, size) = rpc_args_range(msg_args_offset, 2, 2).unwrap(); + + assert!(offset < PAGE_SIZE); + assert!(offset + size > PAGE_SIZE); + } + + #[test] + fn rpc_args_range_rejects_offset_overflow() { + assert_eq!( + rpc_args_range(usize::MAX, 0, 0), + Err(OpteeSmcReturnCode::EBadAddr) + ); + } +} diff --git a/litebox_shim_optee/src/rpc_context.rs b/litebox_shim_optee/src/rpc_context.rs new file mode 100644 index 000000000..3b1fc4435 --- /dev/null +++ b/litebox_shim_optee/src/rpc_context.rs @@ -0,0 +1,596 @@ +// 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. [`RpcStage`] records which RPC response is expected, +//! while [`RpcContext`] retains the registered-memory offset, allocation +//! reference, and action to perform after cleanup. Stage-checked transitions +//! prevent a response from being interpreted as a different step of the +//! protocol. +//! +//! Upstream OP-TEE retains the allocated memory object after a successful +//! `rpc_load()` and releases it when the TA store handle closes (or immediately +//! on an error). LiteBox instead copies or caches the loaded binary in trusted +//! memory and tracks the subsequent `SHM_FREE` round trip explicitly with +//! [`RpcStage::ShmFree`]. + +use alloc::boxed::Box; +use core::sync::atomic::{AtomicU32, Ordering}; +use hashbrown::HashMap; +use litebox_common_optee::{OpteeSmcReturnCode, TeeUuid}; +use once_cell::race::OnceBox; +use spin::mutex::SpinMutex; + +/// Maximum number of RPC contexts that may be active at once. +pub const MAX_RPC_CONTEXTS: usize = 1024; + +/// Progress of an RPC-backed Dynamic TA request. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RpcStage { + LoadTaSize, + ShmAlloc, + LoadTaBinary, + ShmFree, +} + +/// An RPC context map operation failed. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RpcContextError { + Full, + AlreadyExists, + 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), +} + +/// Trusted state for an RPC-backed Dynamic TA request. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct RpcContext { + stage: RpcStage, + ta_uuid: TeeUuid, + registered_shm_ref: u64, + regd_shm_offset: usize, + requested_size: Option, + /// Opaque normal-world allocation reference used for the `SHM_FREE` RPC. + shm_ref: Option, + /// Whether this context owns an entry for the allocation in the local SHM map. + local_shm_registered: bool, + completion: RpcCompletion, +} + +impl RpcContext { + fn new( + stage: RpcStage, + ta_uuid: TeeUuid, + registered_shm_ref: u64, + regd_shm_offset: usize, + ) -> Self { + Self { + stage, + ta_uuid, + registered_shm_ref, + regd_shm_offset, + requested_size: None, + shm_ref: None, + local_shm_registered: false, + completion: RpcCompletion::OpenSession, + } + } + + /// Return the current RPC stage. + pub fn stage(&self) -> RpcStage { + self.stage + } + + /// Return the TA UUID captured when the RPC sequence started. + pub fn ta_uuid(&self) -> TeeUuid { + self.ta_uuid + } + + /// Return the registered shared-memory reference containing the RPC arguments. + pub fn registered_shm_ref(&self) -> u64 { + self.registered_shm_ref + } + + /// Return the offset of the message arguments in registered shared memory. + pub fn registered_shm_offset(&self) -> usize { + self.regd_shm_offset + } + + /// Return the shared-memory reference associated with this context. + pub fn shm_ref(&self) -> Option { + self.shm_ref + } + + /// Return whether this context inserted its allocation into the local SHM map. + pub fn local_shm_registered(&self) -> bool { + self.local_shm_registered + } + + /// Return the action to perform when this RPC sequence finishes. + pub fn completion(&self) -> RpcCompletion { + self.completion + } +} + +/// Maps RPC context IDs to trusted continuation state. +pub struct RpcContextMap { + next_id: AtomicU32, + inner: SpinMutex>, + max_contexts: usize, +} + +impl RpcContextMap { + /// Create an empty RPC context map. + pub fn new() -> Self { + Self::with_limits(0, MAX_RPC_CONTEXTS) + } + + fn with_limits(next_id: u32, max_contexts: usize) -> Self { + Self { + next_id: AtomicU32::new(next_id), + inner: SpinMutex::new(HashMap::new()), + max_contexts, + } + } + + /// Allocate a unique context ID and associate it with trusted continuation state. + pub fn allocate( + &self, + stage: RpcStage, + ta_uuid: TeeUuid, + registered_shm_ref: u64, + regd_shm_offset: usize, + ) -> Result { + let mut contexts = self.inner.lock(); + if contexts.len() >= self.max_contexts { + return Err(RpcContextError::Full); + } + + // With N active entries, at least one of N + 1 consecutive IDs is free. + for _ in 0..=contexts.len() { + let context_id = self.next_id.fetch_add(1, Ordering::Relaxed); + if let hashbrown::hash_map::Entry::Vacant(entry) = contexts.entry(context_id) { + entry.insert(RpcContext::new( + stage, + ta_uuid, + registered_shm_ref, + regd_shm_offset, + )); + return Ok(context_id); + } + } + + Err(RpcContextError::Full) + } + + /// Insert a context with a caller-provided ID. + #[cfg(test)] + fn insert( + &self, + context_id: u32, + stage: RpcStage, + ta_uuid: TeeUuid, + registered_shm_ref: u64, + regd_shm_offset: usize, + ) -> Result<(), RpcContextError> { + let mut contexts = self.inner.lock(); + if contexts.contains_key(&context_id) { + return Err(RpcContextError::AlreadyExists); + } + if contexts.len() >= self.max_contexts { + return Err(RpcContextError::Full); + } + contexts.insert( + context_id, + RpcContext::new(stage, ta_uuid, registered_shm_ref, regd_shm_offset), + ); + Ok(()) + } + + /// Get the current stage for `context_id`. + pub fn get_curr_stage(&self, context_id: u32) -> Option { + self.inner.lock().get(&context_id).map(RpcContext::stage) + } + + /// Get the registered-SHM offset captured before `args[3]` became the context ID. + pub fn get_regd_shm_offset(&self, context_id: u32) -> Option { + self.inner + .lock() + .get(&context_id) + .map(|context| context.regd_shm_offset) + } + + /// Get the TA UUID captured when `context_id` was allocated. + pub fn get_ta_uuid(&self, context_id: u32) -> Option { + self.inner.lock().get(&context_id).map(RpcContext::ta_uuid) + } + + /// Get the registered shared-memory reference containing the RPC arguments. + pub fn get_registered_shm_ref(&self, context_id: u32) -> Option { + self.inner + .lock() + .get(&context_id) + .map(RpcContext::registered_shm_ref) + } + + /// Record the requested allocation size associated with `context_id`. + pub fn set_requested_size( + &self, + context_id: u32, + expected_stage: RpcStage, + requested_size: u64, + ) -> Result<(), RpcContextError> { + let mut contexts = self.inner.lock(); + let context = contexts + .get_mut(&context_id) + .ok_or(RpcContextError::NotFound)?; + if context.stage != expected_stage { + return Err(RpcContextError::UnexpectedStage); + } + context.requested_size = Some(requested_size); + Ok(()) + } + + /// Get the trusted allocation size requested for `context_id`. + pub fn get_requested_size(&self, context_id: u32) -> Option { + self.inner + .lock() + .get(&context_id) + .and_then(|context| context.requested_size) + } + + /// Record the shared-memory allocation associated with `context_id`. + pub fn set_shm_ref( + &self, + context_id: u32, + expected_stage: RpcStage, + shm_ref: u64, + ) -> Result<(), RpcContextError> { + let mut contexts = self.inner.lock(); + let context = contexts + .get_mut(&context_id) + .ok_or(RpcContextError::NotFound)?; + if context.stage != expected_stage { + return Err(RpcContextError::UnexpectedStage); + } + context.shm_ref = Some(shm_ref); + Ok(()) + } + + /// Get the trusted shared-memory reference for `context_id`. + pub fn get_shm_ref(&self, context_id: u32) -> Option { + self.inner + .lock() + .get(&context_id) + .and_then(RpcContext::shm_ref) + } + + /// Record that this context inserted its allocation into the local SHM map. + pub fn set_local_shm_registered( + &self, + context_id: u32, + expected_stage: RpcStage, + ) -> Result<(), RpcContextError> { + let mut contexts = self.inner.lock(); + let context = contexts + .get_mut(&context_id) + .ok_or(RpcContextError::NotFound)?; + if context.stage != expected_stage { + return Err(RpcContextError::UnexpectedStage); + } + context.local_shm_registered = true; + Ok(()) + } + + /// Return whether `context_id` owns an entry in the local SHM map. + pub fn is_local_shm_registered(&self, context_id: u32) -> Option { + self.inner + .lock() + .get(&context_id) + .map(RpcContext::local_shm_registered) + } + + /// Set the action to perform after the current RPC sequence is cleaned up. + pub fn set_completion( + &self, + context_id: u32, + expected_stage: RpcStage, + completion: RpcCompletion, + ) -> Result<(), RpcContextError> { + let mut contexts = self.inner.lock(); + let context = contexts + .get_mut(&context_id) + .ok_or(RpcContextError::NotFound)?; + if context.stage != expected_stage { + return Err(RpcContextError::UnexpectedStage); + } + context.completion = completion; + Ok(()) + } + + /// Atomically advance the stage of a context from `expected` to `next` if it matches the current stage. + /// TODO: Check if the new stage is a valid transition from the current stage. + pub fn transition( + &self, + context_id: u32, + expected: RpcStage, + next: RpcStage, + ) -> Result<(), RpcContextError> { + let mut contexts = self.inner.lock(); + let context = contexts + .get_mut(&context_id) + .ok_or(RpcContextError::NotFound)?; + if context.stage != expected { + return Err(RpcContextError::UnexpectedStage); + } + context.stage = next; + Ok(()) + } + + /// Remove and return a context. + pub fn take(&self, context_id: u32) -> Option { + self.inner.lock().remove(&context_id) + } + + /// Return the number of active RPC contexts. + pub fn len(&self) -> usize { + self.inner.lock().len() + } + + /// Return whether there are no active RPC contexts. + pub fn is_empty(&self) -> bool { + self.inner.lock().is_empty() + } +} + +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 allocates_unique_context_ids() { + let contexts = RpcContextMap::new(); + let first = contexts + .allocate(RpcStage::LoadTaSize, test_uuid(1), 0x10, 0x100) + .unwrap(); + let second = contexts + .allocate(RpcStage::LoadTaSize, test_uuid(2), 0x20, 0x200) + .unwrap(); + + assert_ne!(first, second); + assert_eq!(contexts.len(), 2); + assert_eq!(contexts.get_regd_shm_offset(first), Some(0x100)); + assert_eq!(contexts.get_regd_shm_offset(second), Some(0x200)); + assert_eq!(contexts.get_registered_shm_ref(first), Some(0x10)); + assert_eq!(contexts.get_registered_shm_ref(second), Some(0x20)); + assert_eq!(contexts.get_ta_uuid(first), Some(test_uuid(1))); + assert_eq!(contexts.get_ta_uuid(second), Some(test_uuid(2))); + } + + #[test] + fn transitions_a_stable_context_id() { + let contexts = RpcContextMap::new(); + let context_id = contexts + .allocate(RpcStage::LoadTaSize, test_uuid(1), 1, 0) + .unwrap(); + + assert_eq!( + contexts.transition(context_id, RpcStage::LoadTaSize, RpcStage::ShmAlloc), + Ok(()) + ); + assert_eq!( + contexts.get_curr_stage(context_id), + Some(RpcStage::ShmAlloc) + ); + assert_eq!( + contexts.transition(context_id, RpcStage::LoadTaSize, RpcStage::LoadTaBinary), + Err(RpcContextError::UnexpectedStage) + ); + } + + #[test] + fn taking_a_context_rejects_replay() { + let contexts = RpcContextMap::new(); + let context_id = contexts + .allocate(RpcStage::ShmFree, test_uuid(1), 1, 0) + .unwrap(); + + assert_eq!( + contexts.take(context_id).map(|context| context.stage()), + Some(RpcStage::ShmFree) + ); + assert_eq!(contexts.take(context_id), None); + assert!(contexts.is_empty()); + } + + #[test] + fn tracks_shm_ref_in_trusted_context() { + let contexts = RpcContextMap::new(); + let context_id = contexts + .allocate(RpcStage::ShmAlloc, test_uuid(1), 1, 0) + .unwrap(); + + assert_eq!( + contexts.set_shm_ref(context_id, RpcStage::ShmAlloc, 0x1234), + Ok(()) + ); + assert_eq!(contexts.get_shm_ref(context_id), Some(0x1234)); + assert_eq!( + contexts.set_shm_ref(context_id, RpcStage::LoadTaBinary, 0x5678), + Err(RpcContextError::UnexpectedStage) + ); + assert_eq!(contexts.get_shm_ref(context_id), Some(0x1234)); + } + + #[test] + fn tracks_requested_size_in_trusted_context() { + let contexts = RpcContextMap::new(); + let context_id = contexts + .allocate(RpcStage::LoadTaSize, test_uuid(1), 1, 0) + .unwrap(); + + assert_eq!( + contexts.set_requested_size(context_id, RpcStage::LoadTaSize, 0x4000), + Ok(()) + ); + assert_eq!(contexts.get_requested_size(context_id), Some(0x4000)); + assert_eq!( + contexts.set_requested_size(context_id, RpcStage::ShmAlloc, 0x8000), + Err(RpcContextError::UnexpectedStage) + ); + assert_eq!(contexts.get_requested_size(context_id), Some(0x4000)); + } + + #[test] + fn tracks_local_shm_registration_ownership() { + let contexts = RpcContextMap::new(); + let context_id = contexts + .allocate(RpcStage::ShmAlloc, test_uuid(1), 1, 0) + .unwrap(); + + assert_eq!(contexts.is_local_shm_registered(context_id), Some(false)); + assert_eq!( + contexts.set_local_shm_registered(context_id, RpcStage::ShmAlloc), + Ok(()) + ); + assert_eq!(contexts.is_local_shm_registered(context_id), Some(true)); + assert_eq!( + contexts.set_local_shm_registered(context_id, RpcStage::LoadTaBinary), + Err(RpcContextError::UnexpectedStage) + ); + } + + #[test] + fn tracks_post_free_completion() { + let contexts = RpcContextMap::new(); + let context_id = contexts + .allocate(RpcStage::LoadTaBinary, test_uuid(1), 1, 0) + .unwrap(); + + contexts + .set_completion( + context_id, + RpcStage::LoadTaBinary, + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd), + ) + .unwrap(); + assert_eq!( + contexts.take(context_id).unwrap().completion(), + RpcCompletion::ReturnError(OpteeSmcReturnCode::EBadCmd) + ); + } + + #[test] + fn enforces_capacity_and_duplicate_ids() { + let contexts = RpcContextMap::with_limits(0, 1); + assert_eq!( + contexts.insert(7, RpcStage::LoadTaSize, test_uuid(1), 1, 0), + Ok(()) + ); + assert_eq!( + contexts.insert(7, RpcStage::ShmAlloc, test_uuid(1), 1, 0), + Err(RpcContextError::AlreadyExists) + ); + assert_eq!( + contexts.allocate(RpcStage::LoadTaSize, test_uuid(1), 1, 0), + Err(RpcContextError::Full) + ); + } + + #[test] + fn allocation_wraps_and_skips_active_ids() { + let contexts = RpcContextMap::with_limits(u32::MAX, 3); + contexts + .insert(u32::MAX, RpcStage::LoadTaSize, test_uuid(1), 1, 0) + .unwrap(); + + let context_id = contexts + .allocate(RpcStage::ShmAlloc, test_uuid(2), 2, 0) + .unwrap(); + assert_eq!(context_id, 0); + assert_eq!( + contexts.get_curr_stage(u32::MAX), + Some(RpcStage::LoadTaSize) + ); + assert_eq!(contexts.get_curr_stage(0), Some(RpcStage::ShmAlloc)); + } +} diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 1e1cb1b40..ab5d09ec2 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -1073,6 +1073,29 @@ mod tests { assert!(manager.single_instance_cache.get(&uuid).is_some()); } + #[test] + fn keep_alive_instance_is_reused_after_last_session_closes() { + let manager = SessionManager::new(); + let uuid = make_uuid(0xA5); + let flags = single_instance_flags() | TaFlags::INSTANCE_KEEP_ALIVE; + + register_for_test(&manager, 107, flags, 12, uuid); + let instance = manager.single_instance_cache.get(&uuid).unwrap(); + + assert_eq!(manager.unregister_session(107), Some(flags)); + assert_eq!(manager.count_sessions_for_instance(&instance), 0); + + manager + .with_ta(&uuid, |target| { + let OpenSessionTarget::Sibling(reused) = target else { + panic!("keep-alive instance was not reused"); + }; + assert_eq!(reused.task_page_table_id(), 12); + Ok(()) + }) + .unwrap(); + } + /// `mark_sessions_dead_for_instance` retires the cached single-instance /// TA: Live entries become Dead, stop counting for /// `count_sessions_for_instance`, `with_session` thereafter sees `None`, diff --git a/litebox_shim_optee/src/syscalls/ldelf.rs b/litebox_shim_optee/src/syscalls/ldelf.rs index e79ceba61..f376465ae 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);