From c3637a9ed21545ef31385bd99b07da82e7d480b4 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Thu, 6 Aug 2026 19:40:46 +0000 Subject: [PATCH 1/9] optee: share TA binaries through a global UUID map Move the TA UUID-to-binary map out of individual shim instances so TA binaries can be registered once and reused across instances. This global map is required to support Dyanamically Loading TAs and using them across sessions. Protect the map with a spin-based RwLock and expose shim methods for storing and retrieving TA binaries. Register the LVBS runner's embedded TA during BSP initialization and resolve new TA instances through the shared map. Signed-off-by: Praveen K Paladugu --- litebox_runner_lvbs/src/lib.rs | 53 +++++++++++++++++++++++----------- litebox_shim_optee/Cargo.toml | 2 +- litebox_shim_optee/src/lib.rs | 45 ++++++++++++++++++++--------- 3 files changed, 68 insertions(+), 32 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 8e5350bdc..3f0b54139 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -224,6 +224,11 @@ pub fn init(is_bsp: bool) -> Option<&'static Platform> { // Per-CPU; safe to call on BSP and APs. timer::init(); + if is_bsp { + let shim = litebox_shim_optee::OpteeShimBuilder::new().build(); + register_embedded_tas(&shim); + } + ret } @@ -778,7 +783,8 @@ fn open_session_new_instance( client_identity: Option, ta_req_info: &litebox_shim_optee::msg_handler::TaRequestInfo, ) -> Result<(), OpteeSmcReturnCode> { - let Some(ta_bin) = find_ta_binary(ta_uuid) else { + let shim = litebox_shim_optee::OpteeShimBuilder::new().build(); + let Some(ta_bin) = find_ta_binary(&shim, &ta_uuid) else { msg_args.session = 0; msg_args.ret = TeeResult::ItemNotFound; msg_args.ret_origin = TeeOrigin::Tee; @@ -801,9 +807,8 @@ fn open_session_new_instance( })?; // Load ldelf and TA - Box immediately to keep at fixed heap address - let shim = litebox_shim_optee::OpteeShimBuilder::new().build(); let loaded_program = Box::new( - shim.load_ldelf(LDELF_BINARY, ta_uuid, Some(ta_bin)) + shim.load_ldelf(LDELF_BINARY, ta_uuid, Some(&ta_bin)) .map_err(|_| { // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. @@ -1337,24 +1342,38 @@ fn write_rpc_args_to_normal_world( Ok(()) } -// use include_bytes! to include ldelf and (KMPP) TA binaries +// use include_bytes! to include ldelf const LDELF_BINARY: &[u8] = &[0u8; 0]; -const TA_BINARY: &[u8] = &[0u8; 0]; -const TA_BINARIES: &[&[u8]] = &[TA_BINARY]; + +/// Register a TA binary embedded in the runner image. +fn register_embedded_ta(shim: &litebox_shim_optee::OpteeShim, ta_binary: &'static [u8]) -> bool { + let Some(ta_head) = litebox_common_optee::parse_ta_head(ta_binary) else { + return false; + }; + shim.store_ta_bin(&ta_head.uuid, ta_binary) +} + +/// Register all TA binaries embedded in the runner image. +fn register_embedded_tas(shim: &litebox_shim_optee::OpteeShim) { + static REGISTERED: spin::Once<()> = spin::Once::new(); + REGISTERED.call_once(|| { + assert!(register_embedded_ta( + shim, + include_bytes!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../litebox_runner_optee_on_linux_userland/tests/hello-ta.elf" + )) + )); + }); +} /// Look up TA binary by UUID. /// TODO: Handle PTA UUIDs -fn find_ta_binary(ta_uuid: litebox_common_optee::TeeUuid) -> Option<&'static [u8]> { - use litebox_common_optee::parse_ta_head; - - for ta_binary in TA_BINARIES { - if let Some(ta_head) = parse_ta_head(ta_binary) - && ta_head.uuid == ta_uuid - { - return Some(ta_binary); - } - } - None +fn find_ta_binary( + shim: &litebox_shim_optee::OpteeShim, + ta_uuid: &litebox_common_optee::TeeUuid, +) -> Option> { + shim.get_ta_bin(ta_uuid) } #[panic_handler] diff --git a/litebox_shim_optee/Cargo.toml b/litebox_shim_optee/Cargo.toml index 057a59ef7..459c3d729 100644 --- a/litebox_shim_optee/Cargo.toml +++ b/litebox_shim_optee/Cargo.toml @@ -19,7 +19,7 @@ num_enum = { version = "0.7.3", default-features = false } rangemap = { version = "1.5.1", features = ["const_fn"] } once_cell = { version = "1.20.2", default-features = false, features = ["alloc", "race"] } sha2 = { version = "0.10", default-features = false } -spin = { version = "0.10.0", default-features = false, features = ["spin_mutex", "once"] } +spin = { version = "0.10.0", default-features = false, features = ["spin_mutex", "rwlock", "once"] } thiserror = { version = "2.0.6", default-features = false } zerocopy = { version = "0.8", default-features = false, features = ["derive"] } zeroize = { version = "1.8", default-features = false, features = ["alloc"] } diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index f4bd367c2..c19c68493 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -150,7 +150,6 @@ impl OpteeShimBuilder { boot_instant: TimeProvider::now(self.platform), pm: PageManager::new(&self.litebox), _litebox: self.litebox, - ta_uuid_map: TaUuidMap::new(), pta_busy: spin::mutex::SpinMutex::new(HashSet::new()), }); OpteeShim(global) @@ -169,8 +168,6 @@ struct GlobalState { pm: litebox::mm::PageManager, /// The LiteBox instance used throughout the shim. _litebox: litebox::LiteBox, - /// The TA UUID to binary map for TA loading. - ta_uuid_map: TaUuidMap, /// Tracks which non-concurrent PTAs (i.e., PTAs w/o `TaFlags::CONCURRENT`) /// are currently busy. A busy PTA is *rejected* with `TeeResult::Busy` /// rather than queued. @@ -187,12 +184,12 @@ impl GlobalState { /// Returns `true` if the binary was successfully stored, `false` if the binary's /// UUID (from `.ta_head` section) doesn't match the provided UUID or parsing failed. pub(crate) fn store_ta_bin(&self, ta_uuid: &TeeUuid, ta_bin: &[u8]) -> bool { - self.ta_uuid_map.insert(*ta_uuid, ta_bin.into()) + ta_uuid_map().insert(*ta_uuid, ta_bin.into()) } /// Get the 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) { + if let Some(ta_bin) = ta_uuid_map().get(ta_uuid) { Some(ta_bin) } else { let ta_bin = Self::rpc_get_ta_bin(ta_uuid)?; @@ -205,7 +202,7 @@ impl GlobalState { /// Get the TA flags associated with the given TA UUID. pub(crate) fn get_ta_flags(&self, ta_uuid: &TeeUuid) -> TaFlags { - self.ta_uuid_map.get_flags(ta_uuid).unwrap_or_default() + ta_uuid_map().get_flags(ta_uuid).unwrap_or_default() } /// Monotonic time elapsed since this instance was created, used as GP @@ -227,7 +224,7 @@ impl GlobalState { /// this TA binary #[expect(dead_code)] pub(crate) fn remove_ta_bin(&self, ta_uuid: &TeeUuid) { - let _ = self.ta_uuid_map.remove(ta_uuid); + let _ = ta_uuid_map().remove(ta_uuid); } /// RPC to get the TA binary associated with the given TA UUID. Placeholder for now. @@ -310,6 +307,19 @@ impl OpteeShim { &self.0.pm } + /// Store a TA binary associated with the given TA UUID. + /// + /// Returns `true` if the binary was successfully stored, `false` if the binary's + /// UUID (from `.ta_head` section) doesn't match the provided UUID or parsing failed. + pub fn store_ta_bin(&self, ta_uuid: &TeeUuid, ta_bin: &[u8]) -> bool { + self.0.store_ta_bin(ta_uuid, ta_bin) + } + + /// Get the 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) + } + /// Release all user-space memory mappings owned by this shim instance. /// /// This must be called before switching to the base page table and deleting @@ -1322,13 +1332,13 @@ struct TaInfo { /// Data structure to maintain a mapping from TA UUIDs to their binary data and flags. pub(crate) struct TaUuidMap { - inner: spin::mutex::SpinMutex>, + inner: spin::rwlock::RwLock>, } impl TaUuidMap { pub(crate) fn new() -> Self { Self { - inner: spin::mutex::SpinMutex::new(HashMap::new()), + inner: spin::rwlock::RwLock::new(HashMap::new()), } } @@ -1343,8 +1353,7 @@ impl TaUuidMap { return false; } - let mut inner = self.inner.lock(); - inner.insert( + let _replaced = self.inner.write().insert( uuid, TaInfo { binary: ta_bin, @@ -1355,20 +1364,28 @@ impl TaUuidMap { } pub(crate) fn get(&self, uuid: &TeeUuid) -> Option> { - self.inner.lock().get(uuid).map(|info| info.binary.clone()) + self.inner.read().get(uuid).map(|info| info.binary.clone()) } /// Get the TA flags for a given UUID. pub(crate) fn get_flags(&self, uuid: &TeeUuid) -> Option { - self.inner.lock().get(uuid).map(|info| info.flags) + self.inner.read().get(uuid).map(|info| info.flags) } // Lazy removal of TA binaries when they are no longer needed. pub(crate) fn remove(&self, uuid: &TeeUuid) -> Option> { - self.inner.lock().remove(uuid).map(|info| info.binary) + self.inner.write().remove(uuid).map(|info| info.binary) } } +/// Get the global TA UUID map. +fn ta_uuid_map() -> Arc { + static TA_UUID_MAP: once_cell::race::OnceBox> = once_cell::race::OnceBox::new(); + TA_UUID_MAP + .get_or_init(|| alloc::boxed::Box::new(Arc::new(TaUuidMap::new()))) + .clone() +} + /// Per-instance TA state which can be shared between sessions if it is /// a single-instance multi-session TA. The active session id is carried /// per entry (see [`Task::current_session_id`]). From 29eaa9465da21f5d9ddb916c5dc40bf3ab8967af Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Tue, 11 Aug 2026 19:40:39 +0000 Subject: [PATCH 2/9] Drop the redundant call_once while registering TAs Restore TA_BINARIES, to keep supporting mulitple embedded binaries Signed-off-by: Praveen K Paladugu --- litebox_runner_lvbs/src/lib.rs | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 3f0b54139..f68bcca44 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -1344,6 +1344,8 @@ fn write_rpc_args_to_normal_world( // use include_bytes! to include ldelf const LDELF_BINARY: &[u8] = &[0u8; 0]; +const TA_BINARY: &[u8] = &[0u8; 0]; +const TA_BINARIES: &[&[u8]] = &[TA_BINARY]; /// Register a TA binary embedded in the runner image. fn register_embedded_ta(shim: &litebox_shim_optee::OpteeShim, ta_binary: &'static [u8]) -> bool { @@ -1355,16 +1357,9 @@ fn register_embedded_ta(shim: &litebox_shim_optee::OpteeShim, ta_binary: &'stati /// Register all TA binaries embedded in the runner image. fn register_embedded_tas(shim: &litebox_shim_optee::OpteeShim) { - static REGISTERED: spin::Once<()> = spin::Once::new(); - REGISTERED.call_once(|| { - assert!(register_embedded_ta( - shim, - include_bytes!(concat!( - env!("CARGO_MANIFEST_DIR"), - "/../litebox_runner_optee_on_linux_userland/tests/hello-ta.elf" - )) - )); - }); + for ta_binary in TA_BINARIES { + assert!(register_embedded_ta(shim, ta_binary)); + } } /// Look up TA binary by UUID. From 1539154f75e2b234e2638f464e3711102fedfae9 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Tue, 11 Aug 2026 19:43:34 +0000 Subject: [PATCH 3/9] drop Arc from global TaUuidMap Singleton TaUuidMap is already sychronized with an interior RwLock. An Arc wrapping adds un-necessary reference counting overhead. Signed-off-by: Praveen K Paladugu --- litebox_shim_optee/src/lib.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index c19c68493..cdfe88539 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -1379,11 +1379,9 @@ impl TaUuidMap { } /// Get the global TA UUID map. -fn ta_uuid_map() -> Arc { - static TA_UUID_MAP: once_cell::race::OnceBox> = once_cell::race::OnceBox::new(); - TA_UUID_MAP - .get_or_init(|| alloc::boxed::Box::new(Arc::new(TaUuidMap::new()))) - .clone() +fn ta_uuid_map() -> &'static TaUuidMap { + static TA_UUID_MAP: once_cell::race::OnceBox = once_cell::race::OnceBox::new(); + TA_UUID_MAP.get_or_init(|| alloc::boxed::Box::new(TaUuidMap::new())) } /// Per-instance TA state which can be shared between sessions if it is From c25825fe29841f13f3670e49f53501f271e51b54 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Tue, 11 Aug 2026 20:21:09 +0000 Subject: [PATCH 4/9] optee: store shared TA binaries in Arc This drops the need for deep copy while getting a binary. Signed-off-by: Praveen K Paladugu --- litebox_runner_lvbs/src/lib.rs | 2 +- litebox_shim_optee/src/lib.rs | 16 +++++++--------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index f68bcca44..6fce6ea31 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -1367,7 +1367,7 @@ fn register_embedded_tas(shim: &litebox_shim_optee::OpteeShim) { fn find_ta_binary( shim: &litebox_shim_optee::OpteeShim, ta_uuid: &litebox_common_optee::TeeUuid, -) -> Option> { +) -> Option> { shim.get_ta_bin(ta_uuid) } diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index cdfe88539..826ec9ab0 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -188,7 +188,7 @@ impl GlobalState { } /// Get the TA binary associated with the given TA UUID. - pub(crate) fn get_ta_bin(&self, ta_uuid: &TeeUuid) -> Option> { + pub(crate) fn get_ta_bin(&self, ta_uuid: &TeeUuid) -> Option> { if let Some(ta_bin) = ta_uuid_map().get(ta_uuid) { Some(ta_bin) } else { @@ -220,15 +220,13 @@ impl GlobalState { /// to avoid repeated RPCs and memory transfers. We remove it lazily if there is /// a memory pressure. /// - /// TODO: Use something like `Arc` to to ensure no active ldelf/TA holds a handle to - /// this TA binary #[expect(dead_code)] pub(crate) fn remove_ta_bin(&self, ta_uuid: &TeeUuid) { let _ = 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> { + fn rpc_get_ta_bin(_ta_uuid: &TeeUuid) -> Option> { None } } @@ -316,7 +314,7 @@ impl OpteeShim { } /// Get the TA binary associated with the given TA UUID. - pub fn get_ta_bin(&self, ta_uuid: &TeeUuid) -> Option> { + pub fn get_ta_bin(&self, ta_uuid: &TeeUuid) -> Option> { self.0.get_ta_bin(ta_uuid) } @@ -1325,7 +1323,7 @@ impl TaHandleMap { /// Entry in the TA UUID map containing binary data and parsed flags. struct TaInfo { /// The raw TA binary - binary: alloc::boxed::Box<[u8]>, + binary: Arc<[u8]>, /// Parsed TA flags from .ta_head section flags: TaFlags, } @@ -1342,7 +1340,7 @@ impl TaUuidMap { } } - pub(crate) fn insert(&self, uuid: TeeUuid, ta_bin: alloc::boxed::Box<[u8]>) -> bool { + pub(crate) fn insert(&self, uuid: TeeUuid, ta_bin: Arc<[u8]>) -> bool { // Parse TA head from the binary's .ta_head section let Some(ta_head) = litebox_common_optee::parse_ta_head(&ta_bin) else { return false; @@ -1363,7 +1361,7 @@ impl TaUuidMap { true } - pub(crate) fn get(&self, uuid: &TeeUuid) -> Option> { + pub(crate) fn get(&self, uuid: &TeeUuid) -> Option> { self.inner.read().get(uuid).map(|info| info.binary.clone()) } @@ -1373,7 +1371,7 @@ impl TaUuidMap { } // Lazy removal of TA binaries when they are no longer needed. - pub(crate) fn remove(&self, uuid: &TeeUuid) -> Option> { + pub(crate) fn remove(&self, uuid: &TeeUuid) -> Option> { self.inner.write().remove(uuid).map(|info| info.binary) } } From 389fbc7281056840f99deb5b0595848ca99b791c Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Wed, 12 Aug 2026 19:24:55 +0000 Subject: [PATCH 5/9] optee: re-introduce ta_uuid_map in Shim Object Now references to ta_uuid_map in shim object will be scoped to the lifetime of the shim object. Signed-off-by: Praveen K Paladugu --- litebox_shim_optee/src/lib.rs | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index 826ec9ab0..8232c8f14 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -150,6 +150,7 @@ impl OpteeShimBuilder { boot_instant: TimeProvider::now(self.platform), pm: PageManager::new(&self.litebox), _litebox: self.litebox, + ta_uuid_map: ta_uuid_map(), pta_busy: spin::mutex::SpinMutex::new(HashSet::new()), }); OpteeShim(global) @@ -168,6 +169,8 @@ struct GlobalState { pm: litebox::mm::PageManager, /// The LiteBox instance used throughout the shim. _litebox: litebox::LiteBox, + /// The TA UUID to binary map for TA loading. + ta_uuid_map: &'static TaUuidMap, /// Tracks which non-concurrent PTAs (i.e., PTAs w/o `TaFlags::CONCURRENT`) /// are currently busy. A busy PTA is *rejected* with `TeeResult::Busy` /// rather than queued. @@ -184,12 +187,12 @@ impl GlobalState { /// Returns `true` if the binary was successfully stored, `false` if the binary's /// UUID (from `.ta_head` section) doesn't match the provided UUID or parsing failed. pub(crate) fn store_ta_bin(&self, ta_uuid: &TeeUuid, ta_bin: &[u8]) -> bool { - ta_uuid_map().insert(*ta_uuid, ta_bin.into()) + self.ta_uuid_map.insert(*ta_uuid, ta_bin.into()) } /// Get the TA binary associated with the given TA UUID. pub(crate) fn get_ta_bin(&self, ta_uuid: &TeeUuid) -> Option> { - if let Some(ta_bin) = ta_uuid_map().get(ta_uuid) { + 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)?; @@ -202,7 +205,7 @@ impl GlobalState { /// Get the TA flags associated with the given TA UUID. pub(crate) fn get_ta_flags(&self, ta_uuid: &TeeUuid) -> TaFlags { - ta_uuid_map().get_flags(ta_uuid).unwrap_or_default() + self.ta_uuid_map.get_flags(ta_uuid).unwrap_or_default() } /// Monotonic time elapsed since this instance was created, used as GP @@ -222,7 +225,7 @@ impl GlobalState { /// #[expect(dead_code)] pub(crate) fn remove_ta_bin(&self, ta_uuid: &TeeUuid) { - let _ = ta_uuid_map().remove(ta_uuid); + let _ = self.ta_uuid_map.remove(ta_uuid); } /// RPC to get the TA binary associated with the given TA UUID. Placeholder for now. From 488c7957e589e5d40767d5d3483e6f5358b7c0a7 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Wed, 12 Aug 2026 19:51:57 +0000 Subject: [PATCH 6/9] ci: ratchet fix in litebox_shim_optee Signed-off-by: Praveen K Paladugu --- dev_tests/src/ratchet.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev_tests/src/ratchet.rs b/dev_tests/src/ratchet.rs index 54a30a964..b1da9344e 100644 --- a/dev_tests/src/ratchet.rs +++ b/dev_tests/src/ratchet.rs @@ -43,7 +43,7 @@ fn ratchet_globals() -> Result<()> { ("litebox_runner_lvbs/", 6), ("litebox_runner_snp/", 2), ("litebox_shim_linux/", 1), - ("litebox_shim_optee/", 5), + ("litebox_shim_optee/", 6), ], |file| { Ok(file From 9b9c5956daf6b000ad944731db49c71356446858 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Thu, 13 Aug 2026 18:18:14 +0000 Subject: [PATCH 7/9] optee: Don't store TA during ldelf ldelf stores TA in the UUID Map while loading it. Drop this redundant step. Signed-off-by: Praveen K Paladugu --- litebox_runner_lvbs/src/lib.rs | 19 ++++++++----------- litebox_shim_optee/src/lib.rs | 6 ------ 2 files changed, 8 insertions(+), 17 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 6fce6ea31..22b847d96 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -784,13 +784,13 @@ fn open_session_new_instance( ta_req_info: &litebox_shim_optee::msg_handler::TaRequestInfo, ) -> Result<(), OpteeSmcReturnCode> { let shim = litebox_shim_optee::OpteeShimBuilder::new().build(); - let Some(ta_bin) = find_ta_binary(&shim, &ta_uuid) else { + if shim.get_ta_bin(&ta_uuid).is_none() { msg_args.session = 0; msg_args.ret = TeeResult::ItemNotFound; msg_args.ret_origin = TeeOrigin::Tee; write_non_ta_msg_args_to_normal_world(msg_args, msg_args_phys_addr)?; return Ok(()); - }; + } // Token is declared before `task_pt_guard` so it drops AFTER it. // Marker only releases once CR3 is back to base. See @@ -807,15 +807,12 @@ fn open_session_new_instance( })?; // Load ldelf and TA - Box immediately to keep at fixed heap address - let loaded_program = Box::new( - shim.load_ldelf(LDELF_BINARY, ta_uuid, Some(&ta_bin)) - .map_err(|_| { - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. - unsafe { teardown_ta_page_table(&shim, task_pt_id) }; - OpteeSmcReturnCode::ENomem - })?, - ); + let loaded_program = Box::new(shim.load_ldelf(LDELF_BINARY, ta_uuid).map_err(|_| { + // Safety: We are about to tear down this TA instance; + // no references to user-space memory will be held afterwards. + unsafe { teardown_ta_page_table(&shim, task_pt_id) }; + OpteeSmcReturnCode::ENomem + })?); let ta_flags = loaded_program.ta_flags; diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index 8232c8f14..cd316fed8 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -255,7 +255,6 @@ impl OpteeShim { &self, ldelf_bin: &[u8], ta_uuid: TeeUuid, - ta_bin: Option<&[u8]>, ) -> Result { let entrypoints = crate::OpteeShimEntrypoints { _not_send: core::marker::PhantomData, @@ -275,11 +274,6 @@ impl OpteeShim { tls_base_addr: Cell::new(0), }, }; - if let Some(ta_bin) = ta_bin - && !entrypoints.task.global.store_ta_bin(&ta_uuid, ta_bin) - { - return Err(loader::elf::ElfLoaderError::InvalidUuid); - } let elf_loader = loader::elf::ElfLoader::new(&entrypoints.task, ldelf_bin, true)?; entrypoints.task.load_ldelf(elf_loader, ta_uuid)?; let params_address = if entrypoints.task.get_ta_stack_base_addr().is_some() { From efb08543b6b2376d9fa51341232f84bf6fc38451 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Thu, 13 Aug 2026 18:34:52 +0000 Subject: [PATCH 8/9] optee: Drop ta_bin from ldelf in Linux Userland Drop ta_bin arg in ldelf casls from Linux Userland. Signed-off-by: Praveen K Paladugu --- litebox_runner_optee_on_linux_userland/src/lib.rs | 8 ++++++-- litebox_runner_optee_on_linux_userland/src/tests.rs | 7 ++++--- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/litebox_runner_optee_on_linux_userland/src/lib.rs b/litebox_runner_optee_on_linux_userland/src/lib.rs index e2f0ea7d0..69875c13e 100644 --- a/litebox_runner_optee_on_linux_userland/src/lib.rs +++ b/litebox_runner_optee_on_linux_userland/src/lib.rs @@ -3,7 +3,7 @@ use anyhow::{Context as _, Result}; use clap::Parser; -use litebox_common_optee::{TeeUuid, UteeEntryFunc, UteeParamOwned}; +use litebox_common_optee::{UteeEntryFunc, UteeParamOwned}; use litebox_platform_multiplex::Platform; use litebox_shim_optee::session::session_manager; use std::path::PathBuf; @@ -109,6 +109,10 @@ fn run_ta_with_default_commands( ldelf_bin: &[u8], ta_bin: &[u8], ) { + let ta_uuid = litebox_common_optee::parse_ta_head(ta_bin) + .expect("Failed to parse TA header from ta_bin") + .uuid; + assert!(shim.store_ta_bin(&ta_uuid, ta_bin)); for func_id in [UteeEntryFunc::OpenSession, UteeEntryFunc::CloseSession] { let params = [const { UteeParamOwned::None }; UteeParamOwned::TEE_NUM_PARAMS]; @@ -116,7 +120,7 @@ fn run_ta_with_default_commands( let session_token = session_manager().try_acquire_open_session_token().unwrap(); let session_id = session_token.session_id().unwrap(); let loaded_program = shim - .load_ldelf(ldelf_bin, TeeUuid::default(), Some(ta_bin)) + .load_ldelf(ldelf_bin, ta_uuid) .map_err(|_| { panic!("Failed to load ldelf"); }) diff --git a/litebox_runner_optee_on_linux_userland/src/tests.rs b/litebox_runner_optee_on_linux_userland/src/tests.rs index 645055431..b930c8efd 100644 --- a/litebox_runner_optee_on_linux_userland/src/tests.rs +++ b/litebox_runner_optee_on_linux_userland/src/tests.rs @@ -27,6 +27,9 @@ pub fn run_ta_with_test_commands( let json_str = std::fs::read_to_string(json_path).unwrap(); serde_json::from_str(&json_str).unwrap() }; + let ta_head = + litebox_common_optee::parse_ta_head(ta_bin).expect("Failed to parse TA header from ta_bin"); + assert!(shim.store_ta_bin(&ta_head.uuid, ta_bin)); let mut ta_info: Option = None; // The active session id for the TA. Set at OpenSession and reused for the // subsequent InvokeCommand entries on the same persistent session. @@ -52,8 +55,6 @@ pub fn run_ta_with_test_commands( continue; } if func_id == UteeEntryFunc::OpenSession { - let ta_head = litebox_common_optee::parse_ta_head(ta_bin) - .expect("Failed to parse TA header from ta_bin"); let mut session_token = session_manager().try_acquire_open_session_token().unwrap(); let open_session_id = session_token.session_id().unwrap(); session_id = Some(open_session_id); @@ -67,7 +68,7 @@ pub fn run_ta_with_test_commands( ); session_manager().set_session_client_identity(open_session_id, Some(client_identity)); let loaded = shim - .load_ldelf(ldelf_bin, ta_head.uuid, Some(ta_bin)) + .load_ldelf(ldelf_bin, ta_head.uuid) .map_err(|_| { panic!("Failed to load TA"); }) From d25d45fa5344aa717c03dfe40e17cfc02e805a73 Mon Sep 17 00:00:00 2001 From: Praveen K Paladugu Date: Thu, 13 Aug 2026 18:36:45 +0000 Subject: [PATCH 9/9] cleanup: Drop unused find_ta_binary method Signed-off-by: Praveen K Paladugu --- litebox_runner_lvbs/src/lib.rs | 9 --------- 1 file changed, 9 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 22b847d96..9aa1644be 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -1359,15 +1359,6 @@ fn register_embedded_tas(shim: &litebox_shim_optee::OpteeShim) { } } -/// Look up TA binary by UUID. -/// TODO: Handle PTA UUIDs -fn find_ta_binary( - shim: &litebox_shim_optee::OpteeShim, - ta_uuid: &litebox_common_optee::TeeUuid, -) -> Option> { - shim.get_ta_bin(ta_uuid) -} - #[panic_handler] fn panic(info: &PanicInfo) -> ! { serial_println!("{}", info);