diff --git a/dev_tests/src/ratchet.rs b/dev_tests/src/ratchet.rs index 54a30a964e..b1da9344e5 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 diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 8e5350bdc5..9aa1644be5 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,13 +783,14 @@ 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(); + 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] diff --git a/litebox_runner_optee_on_linux_userland/src/lib.rs b/litebox_runner_optee_on_linux_userland/src/lib.rs index e2f0ea7d0b..69875c13e1 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 645055431e..b930c8efdc 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"); }) diff --git a/litebox_shim_optee/Cargo.toml b/litebox_shim_optee/Cargo.toml index 057a59ef7a..459c3d7299 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 f4bd367c28..cd316fed80 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/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, /// 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> { + 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 { @@ -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> { + fn rpc_get_ta_bin(_ta_uuid: &TeeUuid) -> Option> { None } } @@ -257,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, @@ -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> { + 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>, + 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()), } } - 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> { - self.inner.lock().get(uuid).map(|info| info.binary.clone()) + pub(crate) fn get(&self, uuid: &TeeUuid) -> Option> { + 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) + pub(crate) fn remove(&self, uuid: &TeeUuid) -> Option> { + 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 = 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`]).