optee: share TA binaries through a global UUID map

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

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,13 +783,14 @@ fn open_session_new_instance(
     client_identity: Option<litebox_common_optee::TeeIdentity>,
     ta_req_info: &litebox_shim_optee::msg_handler::TaRequestInfo<PAGE_SIZE>,
 ) -> Result<(), OpteeSmcReturnCode> {
-    let Some(ta_bin) = find_ta_binary(ta_uuid) else {
+    let shim = litebox_shim_optee::OpteeShimBuilder::new().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(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
@@ -801,16 +807,12 @@ 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))
-            .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;
 
@@ -1337,24 +1339,24 @@ 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];
 
-/// 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;
+/// 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) {
     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);
-        }
+        assert!(register_embedded_ta(shim, ta_binary));
     }
-    None
 }
 
 #[panic_handler]

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,14 +109,18 @@ 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];
 
         if func_id == UteeEntryFunc::OpenSession {
             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");
                 })

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<LoadedProgram> = 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");
                 })

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"] }

litebox_shim_optee/src/lib.rs

@@ -150,7 +150,7 @@ impl OpteeShimBuilder {
             boot_instant: TimeProvider::now(self.platform),
             pm: PageManager::new(&self.litebox),
             _litebox: self.litebox,
-            ta_uuid_map: TaUuidMap::new(),
+            ta_uuid_map: ta_uuid_map(),
             pta_busy: spin::mutex::SpinMutex::new(HashSet::new()),
         });
         OpteeShim(global)
@@ -170,7 +170,7 @@ struct GlobalState {
     /// The LiteBox instance used throughout the shim.
     _litebox: litebox::LiteBox<Platform>,
     /// The TA UUID to binary map for TA loading.
-    ta_uuid_map: TaUuidMap,
+    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.
@@ -191,7 +191,7 @@ impl GlobalState {
     }
 
     /// Get the TA binary associated with the given TA UUID.
-    pub(crate) fn get_ta_bin(&self, ta_uuid: &TeeUuid) -> Option<alloc::boxed::Box<[u8]>> {
+    pub(crate) fn get_ta_bin(&self, ta_uuid: &TeeUuid) -> Option<Arc<[u8]>> {
         if let Some(ta_bin) = self.ta_uuid_map.get(ta_uuid) {
             Some(ta_bin)
         } else {
@@ -223,15 +223,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 _ = 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<alloc::boxed::Box<[u8]>> {
+    fn rpc_get_ta_bin(_ta_uuid: &TeeUuid) -> Option<Arc<[u8]>> {
         None
     }
 }
@@ -257,7 +255,6 @@ impl OpteeShim {
         &self,
         ldelf_bin: &[u8],
         ta_uuid: TeeUuid,
-        ta_bin: Option<&[u8]>,
     ) -> Result<LoadedProgram, loader::elf::ElfLoaderError> {
         let entrypoints = crate::OpteeShimEntrypoints {
             _not_send: core::marker::PhantomData,
@@ -277,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() {
@@ -310,6 +302,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<Arc<[u8]>> {
+        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
@@ -1315,24 +1320,24 @@ 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,
 }
 
 /// Data structure to maintain a mapping from TA UUIDs to their binary data and flags.
 pub(crate) struct TaUuidMap {
-    inner: spin::mutex::SpinMutex<HashMap<TeeUuid, TaInfo>>,
+    inner: spin::rwlock::RwLock<HashMap<TeeUuid, TaInfo>>,
 }
 
 impl TaUuidMap {
     pub(crate) fn new() -> Self {
         Self {
-            inner: spin::mutex::SpinMutex::new(HashMap::new()),
+            inner: spin::rwlock::RwLock::new(HashMap::new()),
         }
     }
 
-    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;
@@ -1343,8 +1348,7 @@ impl TaUuidMap {
             return false;
         }
 
-        let mut inner = self.inner.lock();
-        inner.insert(
+        let _replaced = self.inner.write().insert(
             uuid,
             TaInfo {
                 binary: ta_bin,
@@ -1354,21 +1358,27 @@ impl TaUuidMap {
         true
     }
 
-    pub(crate) fn get(&self, uuid: &TeeUuid) -> Option<alloc::boxed::Box<[u8]>> {
-        self.inner.lock().get(uuid).map(|info| info.binary.clone())
+    pub(crate) fn get(&self, uuid: &TeeUuid) -> Option<Arc<[u8]>> {
+        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<TaFlags> {
-        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<alloc::boxed::Box<[u8]>> {
-        self.inner.lock().remove(uuid).map(|info| info.binary)
+    pub(crate) fn remove(&self, uuid: &TeeUuid) -> Option<Arc<[u8]>> {
+        self.inner.write().remove(uuid).map(|info| info.binary)
     }
 }
 
+/// Get the global TA UUID map.
+fn ta_uuid_map() -> &'static TaUuidMap {
+    static TA_UUID_MAP: once_cell::race::OnceBox<TaUuidMap> = 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
 /// a single-instance multi-session TA. The active session id is carried
 /// per entry (see [`Task::current_session_id`]).