From bd145a52524b08d7cb30324fc8a474b335d5edfd Mon Sep 17 00:00:00 2001 From: Angelina Vu Date: Fri, 6 Mar 2026 00:53:16 +0000 Subject: [PATCH 1/2] Support OP-TEE dynamic shared memory --- litebox_common_optee/src/lib.rs | 16 ++++++++++++++- litebox_runner_lvbs/src/lib.rs | 5 +---- litebox_shim_optee/src/msg_handler.rs | 29 +++++++++++++++++++++++++-- 3 files changed, 43 insertions(+), 7 deletions(-) diff --git a/litebox_common_optee/src/lib.rs b/litebox_common_optee/src/lib.rs index 66a24be35..234080651 100644 --- a/litebox_common_optee/src/lib.rs +++ b/litebox_common_optee/src/lib.rs @@ -2025,7 +2025,7 @@ impl From<&OpteeSmcArgsPage> for OpteeSmcArgs { /// OP-TEE SMC call arguments. #[derive(Clone, Copy, Default, FromBytes)] pub struct OpteeSmcArgs { - args: [usize; Self::NUM_OPTEE_SMC_ARGS], + pub args: [usize; Self::NUM_OPTEE_SMC_ARGS], } impl OpteeSmcArgs { @@ -2050,6 +2050,19 @@ impl OpteeSmcArgs { } } + /// Get the shared memory reference and offset for the physical address of `OpteeMsgArgs`. + pub fn optee_regd_shm_ref_and_offset(&self) -> Result<(u64, usize), OpteeSmcReturnCode> { + // args[1]:args[2] contains the shared memory reference (pointer) + // and args[3] contains the offset within that shared memory. + if self.args[1] & 0xffff_ffff_0000_0000 == 0 && self.args[2] & 0xffff_ffff_0000_0000 == 0 { + let shm_ref = (self.args[1] << 32) | self.args[2]; + let offset = self.args[3]; + Ok((shm_ref as u64, offset)) + } else { + Err(OpteeSmcReturnCode::EBadAddr) + } + } + /// 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; @@ -2121,6 +2134,7 @@ pub enum OpteeSmcResult<'a> { CallWithArg { msg_args: Box, rpc_args: Option>, + msg_args_phys_addr: u64, }, } diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index de8c3abd4..ee598e4be 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -417,10 +417,6 @@ fn optee_smc_handler(smc_args_addr: usize) -> OpteeSmcArgs { let Ok(mut smc_args) = (unsafe { smc_args_ptr.read_at_offset(0) }) else { return make_error_response(OpteeSmcReturnCode::EBadAddr); }; - let Ok(msg_args_phys_addr) = smc_args.optee_msg_args_phys_addr() else { - smc_args.set_return_code(OpteeSmcReturnCode::EBadAddr); - return *smc_args; - }; let Ok(smc_result) = handle_optee_smc_args(&mut smc_args) else { smc_args.set_return_code(OpteeSmcReturnCode::EBadCmd); return *smc_args; @@ -428,6 +424,7 @@ fn optee_smc_handler(smc_args_addr: usize) -> OpteeSmcArgs { if let OpteeSmcResult::CallWithArg { msg_args, rpc_args: _, + msg_args_phys_addr, } = smc_result { let mut msg_args = *msg_args; diff --git a/litebox_shim_optee/src/msg_handler.rs b/litebox_shim_optee/src/msg_handler.rs index f0c70acf7..5ca6243aa 100644 --- a/litebox_shim_optee/src/msg_handler.rs +++ b/litebox_shim_optee/src/msg_handler.rs @@ -185,13 +185,38 @@ pub fn handle_optee_smc_args( Ok(OpteeSmcResult::CallWithArg { msg_args, rpc_args: None, + msg_args_phys_addr: msg_args_addr as u64, }) } - OpteeSmcFunction::CallWithRpcArg | OpteeSmcFunction::CallWithRegdArg => { + OpteeSmcFunction::CallWithRpcArg => { let msg_args_addr = smc.optee_msg_args_phys_addr()?; let msg_args_addr: usize = msg_args_addr.truncate(); let (msg_args, rpc_args) = read_optee_msg_args_from_phys(msg_args_addr, true)?; - Ok(OpteeSmcResult::CallWithArg { msg_args, rpc_args }) + Ok(OpteeSmcResult::CallWithArg { + msg_args, + rpc_args, + msg_args_phys_addr: msg_args_addr as u64, + }) + } + 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)?; + let page_index = (shm_info.page_offset + offset) / PAGE_SIZE; + let offset_in_page = (shm_info.page_offset + 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() + offset_in_page; + + let (msg_args, rpc_args) = read_optee_msg_args_from_phys(msg_args_addr, true)?; + Ok(OpteeSmcResult::CallWithArg { + msg_args, + rpc_args, + msg_args_phys_addr: msg_args_addr as u64, + }) } OpteeSmcFunction::ExchangeCapabilities => { // TODO: update the below when we support more features From 57e77fad480299c489e9e3f5e78e18c0b60a6b68 Mon Sep 17 00:00:00 2001 From: Angelina Vu Date: Mon, 16 Mar 2026 18:57:56 +0000 Subject: [PATCH 2/2] Refactor code to use read_data_from_shm --- litebox_common_optee/src/lib.rs | 2 +- litebox_shim_optee/src/msg_handler.rs | 107 ++++++++++++++++---------- 2 files changed, 69 insertions(+), 40 deletions(-) diff --git a/litebox_common_optee/src/lib.rs b/litebox_common_optee/src/lib.rs index 234080651..c4384fcf0 100644 --- a/litebox_common_optee/src/lib.rs +++ b/litebox_common_optee/src/lib.rs @@ -2025,7 +2025,7 @@ impl From<&OpteeSmcArgsPage> for OpteeSmcArgs { /// OP-TEE SMC call arguments. #[derive(Clone, Copy, Default, FromBytes)] pub struct OpteeSmcArgs { - pub args: [usize; Self::NUM_OPTEE_SMC_ARGS], + args: [usize; Self::NUM_OPTEE_SMC_ARGS], } impl OpteeSmcArgs { diff --git a/litebox_shim_optee/src/msg_handler.rs b/litebox_shim_optee/src/msg_handler.rs index 5ca6243aa..fb45d9edb 100644 --- a/litebox_shim_optee/src/msg_handler.rs +++ b/litebox_shim_optee/src/msg_handler.rs @@ -64,6 +64,47 @@ fn page_align_up(len: u64) -> u64 { len.next_multiple_of(PAGE_SIZE as u64) } +fn parse_optee_msg_args( + blob: &[u8], + has_rpc_arg: bool, +) -> Result<(Box, Option>), OpteeSmcReturnCode> { + // Parse main header from the private buffer. + let main_header = OpteeMsgArgsHeader::read_from_prefix(blob) + .map_err(|_| OpteeSmcReturnCode::EBadAddr)? + .0; + + // Validate num_params from the snapshot. + if main_header.num_params as usize > OpteeMsgArgs::MAX_ARG_PARAM_COUNT { + return Err(OpteeSmcReturnCode::EBadCmd); + } + + let main_size = optee_msg_args_total_size(main_header.num_params); + let main_params = &blob[size_of::()..main_size]; + let main_args = OpteeMsgArgs::from_header_and_raw_params(&main_header, main_params)?; + + // Parse RPC args if present. + // The Linux kernel driver places the RPC arg at offset main_size (based on the actual + // num_params, not MAX_ARG_PARAM_COUNT). Since we copied main_max bytes which is >= main_size, + // and main_size is computed from our own validated snapshot, this is safe. + let rpc_args = if has_rpc_arg { + let rpc_blob = &blob[main_size..]; + let rpc_header = OpteeMsgArgsHeader::read_from_prefix(rpc_blob) + .map_err(|_| OpteeSmcReturnCode::EBadAddr)? + .0; + // Re-validate RPC num_params from the snapshot against our negotiated limit. + if rpc_header.num_params as usize > OpteeRpcArgs::MAX_RPC_ARG_PARAM_COUNT { + return Err(OpteeSmcReturnCode::EBadCmd); + } + let rpc_params = &rpc_blob[size_of::()..]; + let rpc = OpteeRpcArgs::from_header_and_raw_params(&rpc_header, rpc_params)?; + Some(Box::new(rpc)) + } else { + None + }; + + Ok((Box::new(main_args), rpc_args)) +} + /// Read `OpteeMsgArgs` (and optional `OpteeRpcArgs`) from a VTL0 physical address. /// /// Copies the maximum possible size in one shot into a private VTL1 buffer, then @@ -126,41 +167,7 @@ pub fn read_optee_msg_args_from_phys( unsafe { blob_ptr.read_slice_at_offset(0, &mut blob) } .map_err(|_| OpteeSmcReturnCode::EBadAddr)?; - // Parse main header from the private buffer. - let main_header = OpteeMsgArgsHeader::read_from_prefix(&blob) - .map_err(|_| OpteeSmcReturnCode::EBadAddr)? - .0; - - // Validate num_params from the snapshot. - if main_header.num_params as usize > OpteeMsgArgs::MAX_ARG_PARAM_COUNT { - return Err(OpteeSmcReturnCode::EBadCmd); - } - - let main_size = optee_msg_args_total_size(main_header.num_params); - let main_params = &blob[size_of::()..main_size]; - let main_args = OpteeMsgArgs::from_header_and_raw_params(&main_header, main_params)?; - - // Parse RPC args if present. - // The Linux kernel driver places the RPC arg at offset main_size (based on the actual - // num_params, not MAX_ARG_PARAM_COUNT). Since we copied main_max bytes which is >= main_size, - // and main_size is computed from our own validated snapshot, this is safe. - let rpc_args = if has_rpc_arg { - let rpc_blob = &blob[main_size..]; - let rpc_header = OpteeMsgArgsHeader::read_from_prefix(rpc_blob) - .map_err(|_| OpteeSmcReturnCode::EBadAddr)? - .0; - // Re-validate RPC num_params from the snapshot against our negotiated limit. - if rpc_header.num_params as usize > OpteeRpcArgs::MAX_RPC_ARG_PARAM_COUNT { - return Err(OpteeSmcReturnCode::EBadCmd); - } - let rpc_params = &rpc_blob[size_of::()..]; - let rpc = OpteeRpcArgs::from_header_and_raw_params(&rpc_header, rpc_params)?; - Some(Box::new(rpc)) - } else { - None - }; - - Ok((Box::new(main_args), rpc_args)) + parse_optee_msg_args(&blob, has_rpc_arg) } /// This function handles `OpteeSmcArgs` passed from the normal world (VTL0) via an OP-TEE SMC call. @@ -204,14 +211,28 @@ pub fn handle_optee_smc_args( let shm_info = shm_ref_map() .get(shm_ref) .ok_or(OpteeSmcReturnCode::EBadAddr)?; - let page_index = (shm_info.page_offset + offset) / PAGE_SIZE; - let offset_in_page = (shm_info.page_offset + offset) % PAGE_SIZE; + + // 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.truncate()); + let copy_size = main_max + + optee_msg_args_total_size(OpteeRpcArgs::MAX_RPC_ARG_PARAM_COUNT.truncate()); + + let mut blob = alloc::vec![0u8; copy_size]; + read_data_from_shm_with_offset(&shm_info, 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() + offset_in_page; - let (msg_args, rpc_args) = read_optee_msg_args_from_phys(msg_args_addr, true)?; Ok(OpteeSmcResult::CallWithArg { msg_args, rpc_args, @@ -773,11 +794,19 @@ fn get_shm_info_from_optee_msg_param_rmem( fn read_data_from_shm( shm_info: &ShmInfo, buffer: &mut [u8], +) -> Result<(), OpteeSmcReturnCode> { + read_data_from_shm_with_offset(shm_info, 0, buffer) +} + +fn read_data_from_shm_with_offset( + shm_info: &ShmInfo, + offset: usize, + buffer: &mut [u8], ) -> Result<(), OpteeSmcReturnCode> { let mut ptr: NormalWorldConstPtr = shm_info.clone().try_into()?; // SAFETY: The data is copied into a buffer owned by LiteBox to avoid TOCTOU issues. unsafe { - ptr.read_slice_at_offset(0, buffer)?; + ptr.read_slice_at_offset(offset, buffer)?; } Ok(()) }