From 3ae4bca27b05d164a601bafc34dbec11cd9f0f29 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Fri, 22 May 2026 03:05:08 +0000 Subject: [PATCH 01/28] improve single-instance TA handling --- litebox_runner_lvbs/src/lib.rs | 378 +++++++++++++++++------------- litebox_shim_optee/src/lib.rs | 4 +- litebox_shim_optee/src/session.rs | 257 ++++++++++++-------- 3 files changed, 384 insertions(+), 255 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 0ae319e59d..b1bc16fe5f 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -44,11 +44,11 @@ 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::{ - CreationReservation, SessionIdGuard, SessionManager, TaInstance, allocate_session_id, + CreationReservation, SessionIdGuard, SessionManager, SessionTarget, TaInstance, + allocate_session_id, }; use litebox_shim_optee::{NormalWorldConstPtr, NormalWorldMutPtr, UserConstPtr}; use once_cell::race::OnceBox; -use spin::mutex::SpinMutex; /// Seed the initial heap regions so the global allocator has enough memory /// for slab-backed allocations (the slab needs >= 2 MB backing pages). @@ -358,6 +358,35 @@ unsafe fn delete_task_page_table(task_pt_id: usize) -> Result<(), OpteeSmcReturn } } +/// Guard that restores the base page table when leaving a TA page table scope. +struct TaskPageTableGuard { + active: bool, +} + +impl TaskPageTableGuard { + fn enter(task_pt_id: usize) -> Result { + unsafe { switch_to_task_page_table(task_pt_id)? }; + Ok(Self { active: true }) + } + + fn leave(mut self) { + unsafe { switch_to_base_page_table() }; + self.active = false; + } + + fn deactivate(&mut self) { + self.active = false; + } +} + +impl Drop for TaskPageTableGuard { + fn drop(&mut self) { + if self.active { + unsafe { switch_to_base_page_table() }; + } + } +} + /// Tears down a TA's memory mappings and page table. /// /// This performs the following steps in order: @@ -380,6 +409,15 @@ unsafe fn teardown_ta_page_table(shim: &litebox_shim_optee::OpteeShim, task_pt_i } } +unsafe fn teardown_active_ta_page_table( + guard: &mut TaskPageTableGuard, + shim: &litebox_shim_optee::OpteeShim, + task_pt_id: usize, +) { + guard.deactivate(); + unsafe { teardown_ta_page_table(shim, task_pt_id) }; +} + /// Handler for OP-TEE SMC calls. /// /// This function processes SMC calls from the normal world (VTL0) and dispatches them @@ -504,13 +542,17 @@ fn handle_open_session( .get_known_flags(&ta_uuid) .is_none_or(|f| f.is_single_instance()); - // Resolve or create the TA instance. - // For single-instance TAs, `with_creation_slot` re-checks the cache - // under its lock and serializes instance creation per UUID. - // If a cache hit returns a zombie (an instance torn down by a - // concurrent close/panic), evict the dead entry and ask the Linux driver - // to retry so it can create a fresh TA instance. - match session_manager().with_creation_slot(&ta_uuid, is_single_instance, || { + let single_instance_lock = + is_single_instance.then(|| session_manager().single_instance_lock(ta_uuid)); + let _single_instance_guard = if let Some(lock) = single_instance_lock.as_ref() { + Some(lock.try_lock().ok_or(OpteeSmcReturnCode::EThreadLimit)?) + } else { + None + }; + + // Resolve or create the TA instance. For single-instance TAs, the UUID + // lock above serializes creation, reuse, and teardown for this TA. + let result = match session_manager().with_creation_slot(&ta_uuid, is_single_instance, || { open_session_new_instance( msg_args, msg_args_phys_addr, @@ -519,68 +561,48 @@ fn handle_open_session( client_identity, &ta_req_info, ) - })? { - CreationReservation::ExistingSingleInstance(existing) => { - match open_session_single_instance( - msg_args, - msg_args_phys_addr, - existing.clone(), - params, - ta_uuid, - &ta_req_info, - )? { - OpenSessionOutcome::Handled => Ok(()), - OpenSessionOutcome::InstanceDestroyed => { - // Evict the zombie. Analog of OP-TEE's `maybe_release_ta_ctx` - // removing the dead ctx from `tee_ctxes`. - let _ = session_manager().remove_single_instance_if_same(&ta_uuid, &existing); - Err(OpteeSmcReturnCode::EThreadLimit) - } - } - } - CreationReservation::SlotReserved => Ok(()), + }) { + Ok(CreationReservation::ExistingSingleInstance(existing)) => open_session_single_instance( + msg_args, + msg_args_phys_addr, + existing.clone(), + params, + ta_uuid, + &ta_req_info, + ), + Ok(CreationReservation::SlotReserved) => Ok(()), + Err(e) => Err(e), + }; + + // If we conservatively held the per-UUID lock for an unknown TA that + // turned out to be multi-instance, evict the lock entry we created. + // Subsequent OpenSessions for this UUID skip the lock entirely (their + // `get_known_flags` will now return the multi-instance flags), so the + // entry would otherwise leak. + if single_instance_lock.is_some() + && let Some(actual_flags) = session_manager().get_known_flags(&ta_uuid) + && !actual_flags.is_single_instance() + { + session_manager().evict_single_instance_lock_if_unused(&ta_uuid); } -} -/// Outcome of [`open_session_single_instance`]. -enum OpenSessionOutcome { - /// Session was successfully opened, TA returned a non-fatal error, or TA panicked - /// and the instance was destroyed inline. No extra cleanup effort is needed. - Handled, - /// The cached `TaInstance` is `closed` and must not be entered. - InstanceDestroyed, + result } /// Open a new session on an existing single-instance TA. /// -/// Returns `Err(OpteeSmcReturnCode::EThreadLimit)` if the TA instance is currently in use. -/// The Linux driver will wait and retry automatically. -/// Returns [`OpenSessionOutcome::InstanceDestroyed`] if the cached TA is closed. -/// /// If the TA's OpenSession entry point returns an error, the session is not registered. -/// On TARGET_DEAD the cached instance is destroyed unconditionally; any sibling sessions -/// become orphans that fail-fast on next access via the `instance.closed` check. +/// On TARGET_DEAD the cached instance is destroyed unconditionally. /// For cleanup semantics, see OP-TEE OS `tee_ta_open_session()` in `tee_ta_manager.c`. -#[allow(clippy::type_complexity)] fn open_session_single_instance( msg_args: &mut OpteeMsgArgs, msg_args_phys_addr: u64, - instance_arc: Arc>, + instance_arc: Arc, params: &[litebox_common_optee::UteeParamOwned], ta_uuid: litebox_common_optee::TeeUuid, ta_req_info: &litebox_shim_optee::msg_handler::TaRequestInfo, -) -> Result { - // Use try_lock to avoid spinning - return EThreadLimit if TA is in use - // The Linux driver will handle this by waiting and retrying - let mut instance = instance_arc - .try_lock() - .ok_or(OpteeSmcReturnCode::EThreadLimit)?; - - // `closed == true` means the instance is terminal and must not be entered. - if instance.closed { - return Ok(OpenSessionOutcome::InstanceDestroyed); - } - let task_pt_id = instance.task_page_table_id; +) -> Result<(), OpteeSmcReturnCode> { + let task_pt_id = instance_arc.task_page_table_id; // Allocate session ID BEFORE calling load_ta_context so TA gets correct ID. // Use SessionIdGuard to ensure the ID is recycled on any error path @@ -597,13 +619,12 @@ fn open_session_single_instance( runner_session_id ); - let ta_flags = instance.loaded_program.ta_flags; + let ta_flags = instance_arc.loaded_program.ta_flags; - // Switch to the existing TA's page table - unsafe { switch_to_task_page_table(task_pt_id)? }; + let mut task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; // Load TA context with parameters for OpenSession - pass actual session_id - instance + instance_arc .loaded_program .entrypoints .as_ref() @@ -620,13 +641,13 @@ fn open_session_single_instance( let mut ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::reenter_thread_ref( - instance.loaded_program.entrypoints.as_ref().unwrap(), + instance_arc.loaded_program.entrypoints.as_ref().unwrap(), &mut ctx, ); } // Read TA output parameters from the stack buffer - let params_address = instance + let params_address = instance_arc .loaded_program .params_address .ok_or(OpteeSmcReturnCode::EBadAddr)?; @@ -647,8 +668,9 @@ fn open_session_single_instance( ); // Write error response BEFORE switching page tables (accesses user memory). - // Keep the instance lock held until this completes so another core cannot - // tear down the active page table while this core is copying TA outputs. + // The per-UUID lock held by the caller of `handle_open_session` keeps + // another core from tearing down the active page table while this core + // is copying TA outputs. let write_result = write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, @@ -664,21 +686,27 @@ fn open_session_single_instance( if return_code == TeeResult::TargetDead { debug_serial_println!("Single-instance TA panicked during OpenSession, cleaning up"); + // Mark sibling sessions dead BEFORE evicting the per-UUID lock + // (inside `remove_single_instance_if_same`). Otherwise a racing + // handler could allocate a fresh per-UUID lock and walk past a + // still-Live session entry. + session_manager() + .sessions() + .mark_sessions_dead_for_instance(&instance_arc); let _ = session_manager().remove_single_instance_if_same(&ta_uuid, &instance_arc); - instance.closed = true; - // 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(&instance.shim, task_pt_id) }; + unsafe { + teardown_active_ta_page_table(&mut task_pt_guard, &instance_arc.shim, task_pt_id) + }; // TODO: Per OP-TEE OS semantics, if the TA has INSTANCE_KEEP_ALIVE but not // INSTANCE_KEEP_CRASHED, we should respawn the TA here instead of just // cleaning it up. Currently we always clean up on panic. } - drop(instance); write_result?; - return Ok(OpenSessionOutcome::Handled); + return Ok(()); } // Treat write-back failure as OpenSession failure: do not publish the session. @@ -708,15 +736,14 @@ fn open_session_single_instance( == 0 { let _ = session_manager().remove_single_instance_if_same(&ta_uuid, &instance_arc); - instance.closed = true; - // 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(&instance.shim, task_pt_id) }; + unsafe { + teardown_active_ta_page_table(&mut task_pt_guard, &instance_arc.shim, task_pt_id) + }; } else { let _ = session_id_guard.disarm(); } - drop(instance); return Err(e); } @@ -724,14 +751,12 @@ fn open_session_single_instance( session_manager().register_session(runner_session_id, instance_arc.clone(), ta_uuid, ta_flags); session_id_guard.disarm(); - drop(instance); - debug_serial_println!( "OpenSession complete on single-instance TA: session_id={}", runner_session_id ); - Ok(OpenSessionOutcome::Handled) + Ok(()) } /// Create a new TA instance for a session. @@ -756,23 +781,20 @@ fn open_session_new_instance( debug_serial_println!("Created task page table ID: {}", task_pt_id); - unsafe { - switch_to_task_page_table(task_pt_id).inspect_err(|_| { - // Safety: switch_to_task_page_table failed, so task page table is not active. - let _ = delete_task_page_table(task_pt_id); - })?; - } + let mut task_pt_guard = TaskPageTableGuard::enter(task_pt_id).inspect_err(|_| { + // Safety: switch_to_task_page_table failed, so task page table is not active. + let _ = unsafe { delete_task_page_table(task_pt_id) }; + })?; // Allocate session ID before loading - return EBusy to normal world if exhausted. // Use SessionIdGuard to ensure the ID is recycled on any error path // (before it is registered with the session manager). - let session_id_guard = SessionIdGuard::new(allocate_session_id().ok_or_else(|| { - // Safety: We're switching to base page table; no user-space refs held. - unsafe { switch_to_base_page_table() }; - // Safety: We've switched to the base page table above. + let Some(session_id) = allocate_session_id() else { + task_pt_guard.leave(); let _ = unsafe { delete_task_page_table(task_pt_id) }; - OpteeSmcReturnCode::EBusy - })?); + return Err(OpteeSmcReturnCode::EBusy); + }; + let session_id_guard = SessionIdGuard::new(session_id); // Safe to unwrap: guard was just created with Some(id). let runner_session_id = session_id_guard.id().unwrap(); @@ -789,7 +811,7 @@ fn open_session_new_instance( .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) }; + unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; OpteeSmcReturnCode::ENomem })?, ); @@ -833,7 +855,7 @@ fn open_session_new_instance( // 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) }; + unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; write_result?; return Ok(()); @@ -843,7 +865,7 @@ fn open_session_new_instance( loaded_program.entrypoints.as_ref().ok_or_else(|| { // 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) }; + unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; OpteeSmcReturnCode::EBadCmd })?; loaded_program @@ -859,7 +881,7 @@ fn open_session_new_instance( .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) }; + unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; OpteeSmcReturnCode::EBadCmd })?; @@ -876,7 +898,7 @@ fn open_session_new_instance( let params_address = loaded_program.params_address.ok_or_else(|| { // 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) }; + unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; OpteeSmcReturnCode::EBadAddr })?; let ta_params = UserConstPtr::::from_usize(params_address) @@ -884,7 +906,7 @@ fn open_session_new_instance( .ok_or_else(|| { // 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) }; + unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; OpteeSmcReturnCode::EBadAddr })?; @@ -912,7 +934,7 @@ fn open_session_new_instance( // 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) }; + unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; write_result?; return Ok(()); @@ -934,16 +956,15 @@ fn open_session_new_instance( .inspect_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) }; + unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; })?; // Success: create TA instance - loaded_program is already boxed, no move happens - let instance = Arc::new(SpinMutex::new(TaInstance { + let instance = Arc::new(TaInstance { shim, loaded_program, task_page_table_id: task_pt_id, - closed: false, - })); + }); // Success: register session and disarm the guard (ownership transfers to session map) session_manager().register_session(runner_session_id, instance.clone(), ta_uuid, ta_flags); @@ -985,29 +1006,48 @@ fn handle_invoke_command( let session_entry = session_manager() .get_session_entry(session_id) .ok_or(OpteeSmcReturnCode::EBadCmd)?; - // Use try_lock to avoid spinning - return EThreadLimit if TA is in use - // The Linux driver will handle this by waiting and retrying - let Some(mut instance) = session_entry.instance.try_lock() else { - return Err(OpteeSmcReturnCode::EThreadLimit); + // Reserve this session id against concurrent SMC entry by another core. + let active_guard = session_manager() + .try_activate_session(session_id) + .ok_or(OpteeSmcReturnCode::EThreadLimit)?; + // For single-instance TAs, also take the per-UUID lock so sibling sessions + // on the same TA serialize against us. + let single_instance_lock = session_entry + .ta_flags + .is_single_instance() + .then(|| session_manager().single_instance_lock(session_entry.ta_uuid)); + let _single_instance_guard = if let Some(lock) = single_instance_lock.as_ref() { + Some(lock.try_lock().ok_or(OpteeSmcReturnCode::EThreadLimit)?) + } else { + None }; - // `closed == true` means the TA instance is terminal and must not be entered. - // The session is orphaned. Report TARGET_DEAD to the client. - if instance.closed { - drop(instance); + // Re-read after acquiring both serialization primitives to pick up any + // concurrent transition to `Dead` or removal that happened while we were + // waiting on the locks above. + let session_entry = session_manager() + .get_session_entry(session_id) + .ok_or(OpteeSmcReturnCode::EBadCmd)?; + let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { session_manager().unregister_session(session_id); + // Release the active-session slot before the recycled id can race a + // new OpenSession; subsequent SMCs for `session_id` see the new + // session (or `None`), neither of which depends on our guard. + drop(active_guard); + // We may have just resurrected the per-UUID lock entry via + // `single_instance_lock()`. If the cached TA is already gone, drop it. + session_manager().evict_single_instance_lock_if_unused(&session_entry.ta_uuid); msg_args.ret = TeeResult::TargetDead; msg_args.ret_origin = TeeOrigin::Tee; write_non_ta_msg_args_to_normal_world(msg_args, msg_args_phys_addr)?; debug_serial_println!( - "InvokeCommand: session_id={} on closed TA instance", + "InvokeCommand: session_id={} on dead TA session", session_id ); return Ok(()); - } - let task_pt_id = instance.task_page_table_id; + }; + let task_pt_id = instance_arc.task_page_table_id; - // Switch to the TA instance's page table - unsafe { switch_to_task_page_table(task_pt_id)? }; + let mut task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; debug_serial_println!( "InvokeCommand: session_id={}, task_pt_id={}, cmd_id={}", @@ -1017,7 +1057,7 @@ fn handle_invoke_command( ); // Load TA context with parameters and cmd_id - pass actual session_id - let entrypoints_ref = instance.loaded_program.entrypoints.as_ref().unwrap(); + let entrypoints_ref = instance_arc.loaded_program.entrypoints.as_ref().unwrap(); entrypoints_ref .load_ta_context( params.as_slice(), @@ -1031,13 +1071,13 @@ fn handle_invoke_command( let mut ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::reenter_thread_ref( - instance.loaded_program.entrypoints.as_ref().unwrap(), + instance_arc.loaded_program.entrypoints.as_ref().unwrap(), &mut ctx, ); } // params_address is constant - stack buffer is reused across invocations - let params_address = instance + let params_address = instance_arc .loaded_program .params_address .ok_or(OpteeSmcReturnCode::EBadAddr)?; @@ -1049,8 +1089,9 @@ fn handle_invoke_command( let return_code = TeeResult::try_from(return_code).unwrap_or(TeeResult::GenericError); // Write response BEFORE switching page tables (accesses user memory). - // Keep the instance lock held until this completes so another core cannot - // tear down the active page table while this core is copying TA outputs. + // The active-session guard and (for single-instance) per-UUID lock prevent + // another core from tearing down the active page table while this core is + // copying TA outputs. let write_result = write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, @@ -1072,26 +1113,28 @@ fn handle_invoke_command( let ta_uuid = session_entry.ta_uuid; let ta_flags = session_entry.ta_flags; - // Remove this session from the map. Sibling sessions on the same - // single-instance TA will be cleaned up lazily on their next - // invoke/close via the `instance.closed` check. - session_manager().unregister_session(session_id); - - // Clear single-instance cache so new OpenSessions for this UUID - // create a fresh instance instead of hitting the zombie one. if ta_flags.is_single_instance() { - let _ = - session_manager().remove_single_instance_if_same(&ta_uuid, &session_entry.instance); + // Mark siblings dead BEFORE evicting the per-UUID lock (inside + // `remove_single_instance_if_same`). Otherwise a racing handler + // could allocate a fresh per-UUID lock and walk past a still-Live + // session entry. + session_manager() + .sessions() + .mark_sessions_dead_for_instance(&instance_arc); + let _ = session_manager().remove_single_instance_if_same(&ta_uuid, &instance_arc); } - instance.closed = true; + session_manager().unregister_session(session_id); + // Release the active-session slot before the recycled id can race a + // new OpenSession; the session is gone from the map so any concurrent + // SMC for this id sees `None` regardless. + drop(active_guard); // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. - // The lock is held, so no other core can enter the TA. - unsafe { teardown_ta_page_table(&instance.shim, task_pt_id) }; - - drop(instance); + unsafe { + teardown_active_ta_page_table(&mut task_pt_guard, &instance_arc.shim, task_pt_id) + }; debug_serial_println!( "InvokeCommand: cleaned up dead TA instance, task_pt_id={}", @@ -1127,33 +1170,49 @@ fn handle_close_session( let session_entry = session_manager() .get_session_entry(session_id) .ok_or(OpteeSmcReturnCode::EBadCmd)?; - // Use try_lock to avoid spinning - return EThreadLimit if TA is in use - // The Linux driver will handle this by waiting and retrying - let Some(mut instance) = session_entry.instance.try_lock() else { - return Err(OpteeSmcReturnCode::EThreadLimit); + // Reserve this session id against concurrent SMC entry by another core. + let active_guard = session_manager() + .try_activate_session(session_id) + .ok_or(OpteeSmcReturnCode::EThreadLimit)?; + // For single-instance TAs, also take the per-UUID lock so sibling sessions + // on the same TA serialize against us. + let single_instance_lock = session_entry + .ta_flags + .is_single_instance() + .then(|| session_manager().single_instance_lock(session_entry.ta_uuid)); + let _single_instance_guard = if let Some(lock) = single_instance_lock.as_ref() { + Some(lock.try_lock().ok_or(OpteeSmcReturnCode::EThreadLimit)?) + } else { + None }; - // `closed == true` means the TA instance is terminal and must not be entered. - // From the client's perspective the session no longer exists, so - // CloseSession is trivially successful. - if instance.closed { - drop(instance); + // Re-read after acquiring both serialization primitives. + let session_entry = session_manager() + .get_session_entry(session_id) + .ok_or(OpteeSmcReturnCode::EBadCmd)?; + let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { session_manager().unregister_session(session_id); + // Release the active-session slot before the recycled id can race a + // new OpenSession; subsequent SMCs for `session_id` see the new + // session (or `None`), neither of which depends on our guard. + drop(active_guard); + // We may have just resurrected the per-UUID lock entry via + // `single_instance_lock()`. If the cached TA is already gone, drop it. + session_manager().evict_single_instance_lock_if_unused(&session_entry.ta_uuid); msg_args.ret = TeeResult::Success; msg_args.ret_origin = TeeOrigin::Tee; write_non_ta_msg_args_to_normal_world(msg_args, msg_args_phys_addr)?; debug_serial_println!( - "CloseSession complete: session_id={}, TA instance closed", + "CloseSession complete: session_id={}, dead TA session", session_id ); return Ok(()); - } - let task_pt_id = instance.task_page_table_id; + }; + let task_pt_id = instance_arc.task_page_table_id; - // Switch to the TA instance's page table - unsafe { switch_to_task_page_table(task_pt_id)? }; + let mut task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; // Load TA context for CloseSession (no params, no cmd_id) - pass actual session_id - instance + instance_arc .loaded_program .entrypoints .as_ref() @@ -1170,7 +1229,7 @@ fn handle_close_session( let mut ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::reenter_thread_ref( - instance.loaded_program.entrypoints.as_ref().unwrap(), + instance_arc.loaded_program.entrypoints.as_ref().unwrap(), &mut ctx, ); } @@ -1185,11 +1244,12 @@ fn handle_close_session( None, ); - // Clone the instance Arc before dropping the lock for later cleanup check - let instance_arc = session_entry.instance.clone(); - // Remove the session entry from the map let removed_entry = session_manager().unregister_session(session_id); + // Release the active-session slot before the recycled id can race a + // new OpenSession; the session is gone from the map so any concurrent + // SMC for this id sees `None` regardless. + drop(active_guard); // Check if this was the last session using the TA instance by counting // remaining sessions that reference this instance. @@ -1203,7 +1263,6 @@ fn handle_close_session( // If this is a single-instance TA with keep_alive flag, don't remove it from memory. // Note: keep_alive is only meaningful for single-instance TAs. if entry.ta_flags.is_single_instance() && entry.ta_flags.is_keep_alive() { - drop(instance); debug_serial_println!( "CloseSession complete: session_id={}, TA kept alive (INSTANCE_KEEP_ALIVE flag)", session_id @@ -1211,20 +1270,18 @@ fn handle_close_session( return write_result; } - // Clear single-instance cache if this was a single-instance TA + // Clear single-instance cache (and evict its per-UUID lock) if this + // was a single-instance TA. No sibling sessions remain on this + // instance, so we don't need to mark anything `Dead` first. if entry.ta_flags.is_single_instance() { let _ = session_manager().remove_single_instance_if_same(&entry.ta_uuid, &instance_arc); } - instance.closed = true; - // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. - // The lock is held, so no other core can enter the TA. - unsafe { teardown_ta_page_table(&instance.shim, task_pt_id) }; - - // Drop the instance to release shim/loaded_program resources - drop(instance); + unsafe { + teardown_active_ta_page_table(&mut task_pt_guard, &instance_arc.shim, task_pt_id) + }; debug_serial_println!( "CloseSession complete: deleted task_pt_id={} (last session)", @@ -1232,7 +1289,6 @@ fn handle_close_session( ); } } else { - drop(instance); debug_serial_println!( "CloseSession complete: session_id={}, other sessions remaining on TA", session_id diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index d4f411f0f4..8394dda8a2 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -39,8 +39,8 @@ pub mod ptr; // Re-export session management types for convenience pub use session::{ - CreationReservation, MAX_TA_INSTANCES, SessionEntry, SessionManager, SessionMap, - SingleInstanceCache, TaInstance, allocate_session_id, + ActiveSessionGuard, CreationReservation, MAX_TA_INSTANCES, SessionEntry, SessionManager, + SessionMap, SessionTarget, SingleInstanceCache, TaInstance, allocate_session_id, }; const MAX_KERNEL_BUF_SIZE: usize = 0x80_000; diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 5d21b404a8..7af0ac9785 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -10,29 +10,29 @@ //! //! ## Concurrency Model //! -//! Single-instance TAs (with `TA_FLAG_SINGLE_INSTANCE | TA_FLAG_MULTI_SESSION`) share -//! one TA instance across multiple sessions. When multiple CPUs try to invoke commands -//! on the same TA instance concurrently, we use `try_lock()` and return -//! `OPTEE_SMC_RETURN_ETHREAD_LIMIT` at the SMC level if the lock is held. +//! TA execution is serialized externally; [`TaInstance`] is shared as a plain +//! `Arc` without an inner mutex. The exclusivity invariant lives in +//! [`SessionManager`]: //! -//! ### Difference from OP-TEE OS +//! - **Single-instance TAs** (with `TA_FLAG_SINGLE_INSTANCE | TA_FLAG_MULTI_SESSION`) +//! share one [`TaInstance`] across all sessions. Open/Invoke/Close serialize on +//! a per-UUID `SpinMutex` returned by [`SessionManager::single_instance_lock`], +//! acquired non-blockingly at handler entry. //! -//! OP-TEE OS uses RPC-based waiting: when a TA is busy, it returns to normal world -//! via `mutex_lock()` issuing an RPC, allowing the Linux kernel to schedule other -//! work while waiting. This is efficient but fundamentally insecure because normal -//! world is untrusted. +//! - **Multi-instance TAs** have one [`TaInstance`] per session. Invoke/Close +//! serialize on a per-`session_id` entry in [`SessionManager::active_sessions`], +//! also acquired non-blockingly. Different sessions run in parallel on their +//! own instances. //! -//! ### LiteBox Behavior +//! Contention on either primitive returns `OPTEE_SMC_RETURN_ETHREAD_LIMIT`; the +//! Linux OP-TEE driver waits on its completion queue and retries the SMC. //! -//! We return `OPTEE_SMC_RETURN_ETHREAD_LIMIT` at the SMC level instead of RPC-waiting. -//! The Linux OP-TEE driver handles this by: -//! 1. Adding the caller to a wait queue (`optee_cq_wait_for_completion`) -//! 2. Sleeping until another call completes (`optee_cq_wait_final` wakes waiters) -//! 3. Automatically retrying the SMC -//! -//! This provides transparent retry behavior for client applications while keeping -//! the waiting logic in normal world (where scheduling is appropriate), without -//! requiring RPCs that would give untrusted code control over secure world execution. +//! On panic teardown or last-session close, sibling sessions of a single-instance +//! TA are flipped to [`SessionTarget::Dead`] *before* the per-UUID lock entry is +//! evicted (see [`SessionManager::remove_single_instance_if_same`]). A racing +//! handler that subsequently allocates a fresh per-UUID lock will therefore +//! observe `Dead` on its re-read of the session entry and short-circuit through +//! the dead-target path. //! //! Reference: //! @@ -115,26 +115,31 @@ pub struct TaInstance { /// Boxed to keep it at a fixed heap address - the Task inside must not be moved /// after initialization because it contains internal state that may not survive moves. pub loaded_program: alloc::boxed::Box, - /// The task page table ID associated with this TA instance. Valid only - /// while `closed == false`. + /// The task page table ID associated with this TA instance. pub task_page_table_id: usize, - /// Set when the TA is committed to teardown (panic or last session closed). Any lock - /// holders should check `closed` before touching `task_page_table_id` and bail if true. - /// - /// The per-instance lock must be held when setting `closed = true` and across - /// the subsequent `teardown_ta_page_table`. - pub closed: bool, } -// SAFETY: TaInstance is protected by SpinMutex and try_lock (`SessionEntry`) +// SAFETY: TaInstance is shared as `Arc`, but only one core is ever +// inside the TA at a time: single-instance TAs serialize on the per-UUID lock, +// multi-instance TAs on the per-`session_id` entry in `active_sessions`. See +// the module-level "Concurrency Model" doc. unsafe impl Send for TaInstance {} unsafe impl Sync for TaInstance {} +/// The target associated with a normal-world session ID. +#[derive(Clone)] +pub enum SessionTarget { + /// The session still targets a live TA instance. + Live(Arc), + /// The TA died, but normal world may still issue Invoke/Close for this ID. + Dead, +} + /// Per-session entry in the session map. #[derive(Clone)] pub struct SessionEntry { - /// The TA instance (may be shared with other sessions for single-instance TAs). - pub instance: Arc>, + /// The TA target (may be shared with other sessions for single-instance TAs). + pub target: SessionTarget, /// The TA UUID (needed for cleanup of single-instance TAs). pub ta_uuid: TeeUuid, /// TA flags parsed from the `.ta_head` section. @@ -157,11 +162,11 @@ impl SessionMap { } /// Get a session's TA instance by session ID. - pub fn get(&self, session_id: u32) -> Option>> { - self.inner - .lock() - .get(&session_id) - .map(|e| e.instance.clone()) + pub fn get(&self, session_id: u32) -> Option> { + match self.inner.lock().get(&session_id).map(|e| &e.target) { + Some(SessionTarget::Live(instance)) => Some(instance.clone()), + Some(SessionTarget::Dead) | None => None, + } } /// Get full session entry by session ID. @@ -173,14 +178,14 @@ impl SessionMap { pub fn insert( &self, session_id: u32, - instance: Arc>, + instance: Arc, ta_uuid: TeeUuid, ta_flags: TaFlags, ) { self.inner.lock().insert( session_id, SessionEntry { - instance, + target: SessionTarget::Live(instance), ta_uuid, ta_flags, }, @@ -203,13 +208,26 @@ impl SessionMap { } /// Count sessions for a specific TA instance (by Arc pointer equality). - pub fn count_sessions_for_instance(&self, instance: &Arc>) -> usize { + pub fn count_sessions_for_instance(&self, instance: &Arc) -> usize { self.inner .lock() .values() - .filter(|e| Arc::ptr_eq(&e.instance, instance)) + .filter(|e| match &e.target { + SessionTarget::Live(current) => Arc::ptr_eq(current, instance), + SessionTarget::Dead => false, + }) .count() } + + /// Mark all sessions pointing at `instance` as dead. + pub fn mark_sessions_dead_for_instance(&self, instance: &Arc) { + for entry in self.inner.lock().values_mut() { + if matches!(&entry.target, SessionTarget::Live(current) if Arc::ptr_eq(current, instance)) + { + entry.target = SessionTarget::Dead; + } + } + } } impl Default for SessionMap { @@ -223,7 +241,7 @@ impl Default for SessionMap { /// Single-instance TAs (with `TA_FLAG_SINGLE_INSTANCE`) share a single TA instance /// across all sessions. This cache stores instances by UUID for fast reuse lookup. pub struct SingleInstanceCache { - inner: SpinMutex>>>, + inner: SpinMutex>>, } impl SingleInstanceCache { @@ -235,17 +253,17 @@ impl SingleInstanceCache { } /// Get a cached single-instance TA by UUID. - pub fn get(&self, uuid: &TeeUuid) -> Option>> { + pub fn get(&self, uuid: &TeeUuid) -> Option> { self.inner.lock().get(uuid).cloned() } /// Cache a single-instance TA by UUID. - pub fn insert(&self, uuid: TeeUuid, instance: Arc>) { + pub fn insert(&self, uuid: TeeUuid, instance: Arc) { self.inner.lock().insert(uuid, instance); } /// Remove a cached single-instance TA only if it is the expected instance. - fn remove_if_same(&self, uuid: &TeeUuid, expected: &Arc>) -> bool { + fn remove_if_same(&self, uuid: &TeeUuid, expected: &Arc) -> bool { let mut guard = self.inner.lock(); match guard.get(uuid) { Some(current) if Arc::ptr_eq(current, expected) => { @@ -329,27 +347,35 @@ impl Drop for SessionIdGuard { } } +/// RAII guard returned by [`SessionManager::try_activate_session`] that removes +/// the session id from the active set on drop. Provides per-session-id +/// serialization for Invoke/Close handlers without requiring a mutex on the +/// `TaInstance` itself. +pub struct ActiveSessionGuard<'a> { + manager: &'a SessionManager, + session_id: u32, +} + +impl Drop for ActiveSessionGuard<'_> { + fn drop(&mut self) { + self.manager.active_sessions.lock().remove(&self.session_id); + } +} + /// Result of [`SessionManager::with_creation_slot`]. pub enum CreationReservation { /// An existing single-instance TA was found (another core cached it /// between our initial lookup and the reservation). Reuse this instance. - ExistingSingleInstance(Arc>), + ExistingSingleInstance(Arc), /// The creation closure ran successfully inside the reserved slot. SlotReserved, } /// State for coordinating concurrent instance creation. /// -/// Guarded by a single lock to provide atomic capacity checks and -/// duplicate-UUID prevention. +/// Guarded by a single lock to provide atomic capacity checks. struct CreationState { - /// UUIDs of single-instance TAs currently being loaded. Prevents multiple cores - /// from simultaneously creating a new instance for the same single-instance - /// TA UUID (which would violate the single-instance invariant). - /// Multi-instance TAs are not tracked here. They can be created concurrently. - pending_uuids: HashSet, - /// Number of instances currently being created (not yet registered). This - /// covers both single-instance and multi-instance TAs. + /// Number of instances currently being created (not yet registered). /// Added to [`SessionManager::instance_count`] for accurate capacity checks. pending_count: usize, } @@ -369,6 +395,11 @@ pub struct SessionManager { creation_state: SpinMutex, /// Cached TA flags by UUID, populated on first successful session registration. known_flags: SpinMutex>, + /// Per-UUID serialization locks for single-instance TA handling. + single_instance_locks: SpinMutex>>>, + /// Session ids currently being handled (Invoke/Close). Guards a session + /// against concurrent SMC entry by another core that targets the same id. + active_sessions: SpinMutex>, } impl SessionManager { @@ -377,11 +408,10 @@ impl SessionManager { Self { sessions: SessionMap::new(), single_instance_cache: SingleInstanceCache::new(), - creation_state: SpinMutex::new(CreationState { - pending_uuids: HashSet::new(), - pending_count: 0, - }), + creation_state: SpinMutex::new(CreationState { pending_count: 0 }), known_flags: SpinMutex::new(HashMap::new()), + single_instance_locks: SpinMutex::new(HashMap::new()), + active_sessions: SpinMutex::new(HashSet::new()), } } @@ -396,12 +426,12 @@ impl SessionManager { } /// Cache a single-instance TA. - pub fn cache_single_instance(&self, uuid: TeeUuid, instance: Arc>) { + pub fn cache_single_instance(&self, uuid: TeeUuid, instance: Arc) { self.single_instance_cache.insert(uuid, instance); } /// Get a session by ID. - pub fn get_session(&self, session_id: u32) -> Option>> { + pub fn get_session(&self, session_id: u32) -> Option> { self.sessions.get(session_id) } @@ -418,11 +448,34 @@ impl SessionManager { self.known_flags.lock().get(uuid).copied() } + /// Get the serialization lock for a single-instance TA UUID. + pub fn single_instance_lock(&self, uuid: TeeUuid) -> Arc> { + self.single_instance_locks + .lock() + .entry(uuid) + .or_insert_with(|| Arc::new(SpinMutex::new(()))) + .clone() + } + + /// Try to mark `session_id` as actively handled. Returns a guard that + /// removes the marker on drop, or `None` if another core is already + /// inside an Invoke/Close handler for this session id. + pub fn try_activate_session(&self, session_id: u32) -> Option> { + if self.active_sessions.lock().insert(session_id) { + Some(ActiveSessionGuard { + manager: self, + session_id, + }) + } else { + None + } + } + /// Register a new session. pub fn register_session( &self, session_id: u32, - instance: Arc>, + instance: Arc, ta_uuid: TeeUuid, ta_flags: TaFlags, ) { @@ -441,13 +494,51 @@ impl SessionManager { } /// Remove a single-instance TA from the cache only if the currently - /// cached `Arc` is the same as `expected`. + /// cached `Arc` is the same as `expected`, and evict the per-UUID + /// serialization lock so its memory is reclaimed. + /// + /// Callers tearing down on TA panic must have already called + /// [`SessionMap::mark_sessions_dead_for_instance`] before invoking this, + /// so any handler that subsequently allocates a fresh per-UUID lock will + /// observe `Dead` on its re-read of the session entry. Callers on the + /// last-session-close path may skip the mark step — by that point there + /// are no sibling sessions to fence out. pub fn remove_single_instance_if_same( &self, uuid: &TeeUuid, - expected: &Arc>, + expected: &Arc, ) -> bool { - self.single_instance_cache.remove_if_same(uuid, expected) + let removed = self.single_instance_cache.remove_if_same(uuid, expected); + if removed { + self.single_instance_locks.lock().remove(uuid); + } + removed + } + + /// Evict the per-UUID serialization lock entry if no single-instance TA is + /// currently cached under that UUID. + /// + /// Two cleanup scenarios use this: + /// + /// 1. **Late stale Invoke/Close on a `Dead` session.** The cached instance + /// was torn down earlier (which evicted the lock entry), but our call + /// to [`single_instance_lock`] resurrected the entry via + /// `or_insert_with`. Evicting it on the way out keeps the lock map + /// bounded. + /// + /// 2. **First OpenSession of a previously-unknown TA that turns out to be + /// multi-instance.** [`handle_open_session`] conservatively assumes + /// single-instance when flags are unknown and grabs the per-UUID lock. + /// Once the TA loads and we learn it's actually multi-instance, the + /// entry serves no future purpose and would otherwise leak. + /// + /// The cache emptiness check is the safety guard: if the cache still has + /// an instance for this UUID, the entry is in legitimate use and we leave + /// it alone. + pub fn evict_single_instance_lock_if_unused(&self, uuid: &TeeUuid) { + if self.single_instance_cache.get(uuid).is_none() { + self.single_instance_locks.lock().remove(uuid); + } } /// Get the total count of unique TA instances (for limit checking). @@ -478,15 +569,16 @@ impl SessionManager { /// Atomically reserve a creation slot and run `f` to create a new TA instance. /// - /// Behavior depends on whether the TA is: - /// - /// - **Single-instance**: Re-checks the single-instance cache under the lock to - /// close TOCTOU windows, and prevents duplicate concurrent creation of - /// the same UUID via `pending_uuids`. + /// For single-instance TAs, the caller must already hold the per-UUID + /// serialization lock (see [`SessionManager::single_instance_lock`]), + /// which guarantees no other core is creating an instance for this UUID + /// concurrently. This function re-checks the single-instance cache under + /// the creation lock so a freshly-cached instance from a prior holder of + /// the per-UUID lock is observed before starting a new load. /// - /// - **Multi-instance**: Each session gets its own independent TA instance, - /// matching OP-TEE OS behavior. Multiple cores may create instances of - /// the same UUID concurrently. + /// For multi-instance TAs, each session gets its own independent + /// `TaInstance`. Multiple cores may create instances of the same UUID + /// concurrently. pub fn with_creation_slot( &self, uuid: &TeeUuid, @@ -499,21 +591,8 @@ impl SessionManager { { let mut state = self.creation_state.lock(); - if is_single_instance { - // Check the single-instance cache under the creation lock. A - // hit means another core finished creating the instance for - // this UUID; reuse it instead of starting a new load. - if let Some(existing) = self.single_instance_cache.get(uuid) { - return Ok(CreationReservation::ExistingSingleInstance(existing)); - } - - // Another core is currently in the middle of creating an instance - // for this single-instance UUID. The instance isn't cached yet, - // so we cannot reuse it. Return EThreadLimit to have the - // normal-world driver wait and retry. - if state.pending_uuids.contains(uuid) { - return Err(OpteeSmcReturnCode::EThreadLimit); - } + if is_single_instance && let Some(existing) = self.single_instance_cache.get(uuid) { + return Ok(CreationReservation::ExistingSingleInstance(existing)); } // Capacity check including in-flight creations. @@ -522,9 +601,6 @@ impl SessionManager { return Err(OpteeSmcReturnCode::ENomem); } - if is_single_instance { - state.pending_uuids.insert(*uuid); - } state.pending_count += 1; } @@ -532,9 +608,6 @@ impl SessionManager { { let mut state = self.creation_state.lock(); - if is_single_instance { - state.pending_uuids.remove(uuid); - } state.pending_count = state.pending_count.saturating_sub(1); } From 69c7145a694a94f614fffdf8a027d37e2b89064c Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Fri, 22 May 2026 21:22:42 +0000 Subject: [PATCH 02/28] revise --- litebox_runner_lvbs/src/lib.rs | 164 ++++++++++++++++----------------- 1 file changed, 81 insertions(+), 83 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index b1bc16fe5f..775a774bc4 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -44,8 +44,8 @@ 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::{ - CreationReservation, SessionIdGuard, SessionManager, SessionTarget, TaInstance, - allocate_session_id, + ActiveSessionGuard, CreationReservation, SessionEntry, SessionIdGuard, SessionManager, + SessionTarget, TaInstance, allocate_session_id, }; use litebox_shim_optee::{NormalWorldConstPtr, NormalWorldMutPtr, UserConstPtr}; use once_cell::race::OnceBox; @@ -359,31 +359,23 @@ unsafe fn delete_task_page_table(task_pt_id: usize) -> Result<(), OpteeSmcReturn } /// Guard that restores the base page table when leaving a TA page table scope. -struct TaskPageTableGuard { - active: bool, -} +/// +/// `switch_to_base_page_table` is an idempotent CR3 write, so it is fine if +/// teardown paths (which switch to base internally before deleting the task +/// page table) run before this guard's `Drop` — the redundant write at drop +/// time is benign. +struct TaskPageTableGuard; impl TaskPageTableGuard { fn enter(task_pt_id: usize) -> Result { unsafe { switch_to_task_page_table(task_pt_id)? }; - Ok(Self { active: true }) - } - - fn leave(mut self) { - unsafe { switch_to_base_page_table() }; - self.active = false; - } - - fn deactivate(&mut self) { - self.active = false; + Ok(Self) } } impl Drop for TaskPageTableGuard { fn drop(&mut self) { - if self.active { - unsafe { switch_to_base_page_table() }; - } + unsafe { switch_to_base_page_table() }; } } @@ -409,15 +401,6 @@ unsafe fn teardown_ta_page_table(shim: &litebox_shim_optee::OpteeShim, task_pt_i } } -unsafe fn teardown_active_ta_page_table( - guard: &mut TaskPageTableGuard, - shim: &litebox_shim_optee::OpteeShim, - task_pt_id: usize, -) { - guard.deactivate(); - unsafe { teardown_ta_page_table(shim, task_pt_id) }; -} - /// Handler for OP-TEE SMC calls. /// /// This function processes SMC calls from the normal world (VTL0) and dispatches them @@ -574,15 +557,13 @@ fn handle_open_session( Err(e) => Err(e), }; - // If we conservatively held the per-UUID lock for an unknown TA that - // turned out to be multi-instance, evict the lock entry we created. - // Subsequent OpenSessions for this UUID skip the lock entirely (their - // `get_known_flags` will now return the multi-instance flags), so the - // entry would otherwise leak. - if single_instance_lock.is_some() - && let Some(actual_flags) = session_manager().get_known_flags(&ta_uuid) - && !actual_flags.is_single_instance() - { + // If we conservatively held the per-UUID lock but no single-instance TA + // ended up cached under this UUID, evict the lock entry we just inserted. + // This covers (a) the TA turned out to be multi-instance and (b) load + // failed entirely. The helper's cache-empty check is the safety guard: + // if a single-instance TA is now cached for this UUID, the entry is in + // legitimate use and `evict_single_instance_lock_if_unused` is a no-op. + if single_instance_lock.is_some() { session_manager().evict_single_instance_lock_if_unused(&ta_uuid); } @@ -621,7 +602,7 @@ fn open_session_single_instance( let ta_flags = instance_arc.loaded_program.ta_flags; - let mut task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; + let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; // Load TA context with parameters for OpenSession - pass actual session_id instance_arc @@ -697,7 +678,7 @@ fn open_session_single_instance( // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. unsafe { - teardown_active_ta_page_table(&mut task_pt_guard, &instance_arc.shim, task_pt_id) + teardown_ta_page_table(&instance_arc.shim, task_pt_id); }; // TODO: Per OP-TEE OS semantics, if the TA has INSTANCE_KEEP_ALIVE but not @@ -739,7 +720,7 @@ fn open_session_single_instance( // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. unsafe { - teardown_active_ta_page_table(&mut task_pt_guard, &instance_arc.shim, task_pt_id) + teardown_ta_page_table(&instance_arc.shim, task_pt_id); }; } else { let _ = session_id_guard.disarm(); @@ -781,7 +762,7 @@ fn open_session_new_instance( debug_serial_println!("Created task page table ID: {}", task_pt_id); - let mut task_pt_guard = TaskPageTableGuard::enter(task_pt_id).inspect_err(|_| { + let task_pt_guard = TaskPageTableGuard::enter(task_pt_id).inspect_err(|_| { // Safety: switch_to_task_page_table failed, so task page table is not active. let _ = unsafe { delete_task_page_table(task_pt_id) }; })?; @@ -790,7 +771,7 @@ fn open_session_new_instance( // Use SessionIdGuard to ensure the ID is recycled on any error path // (before it is registered with the session manager). let Some(session_id) = allocate_session_id() else { - task_pt_guard.leave(); + drop(task_pt_guard); let _ = unsafe { delete_task_page_table(task_pt_id) }; return Err(OpteeSmcReturnCode::EBusy); }; @@ -811,7 +792,7 @@ fn open_session_new_instance( .map_err(|_| { // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. - unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; + unsafe { teardown_ta_page_table(&shim, task_pt_id) }; OpteeSmcReturnCode::ENomem })?, ); @@ -855,7 +836,7 @@ fn open_session_new_instance( // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. - unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; + unsafe { teardown_ta_page_table(&shim, task_pt_id) }; write_result?; return Ok(()); @@ -865,7 +846,7 @@ fn open_session_new_instance( loaded_program.entrypoints.as_ref().ok_or_else(|| { // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. - unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; + unsafe { teardown_ta_page_table(&shim, task_pt_id) }; OpteeSmcReturnCode::EBadCmd })?; loaded_program @@ -881,7 +862,7 @@ fn open_session_new_instance( .map_err(|_| { // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. - unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; + unsafe { teardown_ta_page_table(&shim, task_pt_id) }; OpteeSmcReturnCode::EBadCmd })?; @@ -898,7 +879,7 @@ fn open_session_new_instance( let params_address = loaded_program.params_address.ok_or_else(|| { // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. - unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; + unsafe { teardown_ta_page_table(&shim, task_pt_id) }; OpteeSmcReturnCode::EBadAddr })?; let ta_params = UserConstPtr::::from_usize(params_address) @@ -906,7 +887,7 @@ fn open_session_new_instance( .ok_or_else(|| { // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. - unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; + unsafe { teardown_ta_page_table(&shim, task_pt_id) }; OpteeSmcReturnCode::EBadAddr })?; @@ -934,7 +915,7 @@ fn open_session_new_instance( // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. - unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; + unsafe { teardown_ta_page_table(&shim, task_pt_id) }; write_result?; return Ok(()); @@ -956,7 +937,7 @@ fn open_session_new_instance( .inspect_err(|_| { // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. - unsafe { teardown_active_ta_page_table(&mut task_pt_guard, &shim, task_pt_id) }; + unsafe { teardown_ta_page_table(&shim, task_pt_id) }; })?; // Success: create TA instance - loaded_program is already boxed, no move happens @@ -984,6 +965,37 @@ fn open_session_new_instance( Ok(()) } +/// Tear down a `Dead` session entry observed at Invoke/Close handler entry. +/// +/// Consumes `active_guard` so it is dropped immediately after `unregister_session` +/// recycles the id — the ordering matters: subsequent SMCs for `session_id` then +/// see the new session (or `None`), neither of which depends on our guard. +/// +/// The caller's per-UUID lock guard (if any) stays in its own scope, so the +/// lock remains held across `evict_single_instance_lock_if_unused`. +fn finalize_dead_session( + session_id: u32, + session_entry: &SessionEntry, + active_guard: ActiveSessionGuard<'_>, + msg_args: &mut OpteeMsgArgs, + msg_args_phys_addr: u64, + return_code: TeeResult, + log_prefix: &str, +) -> Result<(), OpteeSmcReturnCode> { + session_manager().unregister_session(session_id); + drop(active_guard); + session_manager().evict_single_instance_lock_if_unused(&session_entry.ta_uuid); + msg_args.ret = return_code; + msg_args.ret_origin = TeeOrigin::Tee; + write_non_ta_msg_args_to_normal_world(msg_args, msg_args_phys_addr)?; + debug_serial_println!( + "{}: session_id={} on dead TA session", + log_prefix, + session_id + ); + Ok(()) +} + /// Handle InvokeCommand. /// /// Looks up the session by ID, switches to its page table, and runs the command. @@ -1028,26 +1040,19 @@ fn handle_invoke_command( .get_session_entry(session_id) .ok_or(OpteeSmcReturnCode::EBadCmd)?; let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { - session_manager().unregister_session(session_id); - // Release the active-session slot before the recycled id can race a - // new OpenSession; subsequent SMCs for `session_id` see the new - // session (or `None`), neither of which depends on our guard. - drop(active_guard); - // We may have just resurrected the per-UUID lock entry via - // `single_instance_lock()`. If the cached TA is already gone, drop it. - session_manager().evict_single_instance_lock_if_unused(&session_entry.ta_uuid); - msg_args.ret = TeeResult::TargetDead; - msg_args.ret_origin = TeeOrigin::Tee; - write_non_ta_msg_args_to_normal_world(msg_args, msg_args_phys_addr)?; - debug_serial_println!( - "InvokeCommand: session_id={} on dead TA session", - session_id + return finalize_dead_session( + session_id, + &session_entry, + active_guard, + msg_args, + msg_args_phys_addr, + TeeResult::TargetDead, + "InvokeCommand", ); - return Ok(()); }; let task_pt_id = instance_arc.task_page_table_id; - let mut task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; + let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; debug_serial_println!( "InvokeCommand: session_id={}, task_pt_id={}, cmd_id={}", @@ -1133,7 +1138,7 @@ fn handle_invoke_command( // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. unsafe { - teardown_active_ta_page_table(&mut task_pt_guard, &instance_arc.shim, task_pt_id) + teardown_ta_page_table(&instance_arc.shim, task_pt_id); }; debug_serial_println!( @@ -1190,26 +1195,19 @@ fn handle_close_session( .get_session_entry(session_id) .ok_or(OpteeSmcReturnCode::EBadCmd)?; let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { - session_manager().unregister_session(session_id); - // Release the active-session slot before the recycled id can race a - // new OpenSession; subsequent SMCs for `session_id` see the new - // session (or `None`), neither of which depends on our guard. - drop(active_guard); - // We may have just resurrected the per-UUID lock entry via - // `single_instance_lock()`. If the cached TA is already gone, drop it. - session_manager().evict_single_instance_lock_if_unused(&session_entry.ta_uuid); - msg_args.ret = TeeResult::Success; - msg_args.ret_origin = TeeOrigin::Tee; - write_non_ta_msg_args_to_normal_world(msg_args, msg_args_phys_addr)?; - debug_serial_println!( - "CloseSession complete: session_id={}, dead TA session", - session_id + return finalize_dead_session( + session_id, + &session_entry, + active_guard, + msg_args, + msg_args_phys_addr, + TeeResult::Success, + "CloseSession", ); - return Ok(()); }; let task_pt_id = instance_arc.task_page_table_id; - let mut task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; + let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; // Load TA context for CloseSession (no params, no cmd_id) - pass actual session_id instance_arc @@ -1280,7 +1278,7 @@ fn handle_close_session( // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. unsafe { - teardown_active_ta_page_table(&mut task_pt_guard, &instance_arc.shim, task_pt_id) + teardown_ta_page_table(&instance_arc.shim, task_pt_id); }; debug_serial_println!( From 5fe6add708ff23d974ea6a98340caebf157ae0b5 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Fri, 22 May 2026 21:35:43 +0000 Subject: [PATCH 03/28] fix --- litebox_shim_optee/src/session.rs | 30 +++++++++++++++++++++++------- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 7af0ac9785..499e0b3712 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -20,12 +20,28 @@ //! acquired non-blockingly at handler entry. //! //! - **Multi-instance TAs** have one [`TaInstance`] per session. Invoke/Close -//! serialize on a per-`session_id` entry in [`SessionManager::active_sessions`], -//! also acquired non-blockingly. Different sessions run in parallel on their -//! own instances. +//! serialize on a per-`session_id` entry acquired via +//! [`SessionManager::try_activate_session`], also acquired non-blockingly. +//! Different sessions run in parallel on their own instances. //! -//! Contention on either primitive returns `OPTEE_SMC_RETURN_ETHREAD_LIMIT`; the -//! Linux OP-TEE driver waits on its completion queue and retries the SMC. +//! ### Difference from OP-TEE OS +//! +//! OP-TEE OS uses RPC-based waiting: when a TA is busy, it returns to normal world +//! via `mutex_lock()` issuing an RPC, allowing the Linux kernel to schedule other +//! work while waiting. This is efficient but fundamentally insecure because normal +//! world is untrusted. +//! +//! ### LiteBox Behavior +//! +//! We return `OPTEE_SMC_RETURN_ETHREAD_LIMIT` at the SMC level instead of RPC-waiting. +//! The Linux OP-TEE driver handles this by: +//! 1. Adding the caller to a wait queue (`optee_cq_wait_for_completion`) +//! 2. Sleeping until another call completes (`optee_cq_wait_final` wakes waiters) +//! 3. Automatically retrying the SMC +//! +//! This provides transparent retry behavior for client applications while keeping +//! the waiting logic in normal world (where scheduling is appropriate), without +//! requiring RPCs that would give untrusted code control over secure world execution. //! //! On panic teardown or last-session close, sibling sessions of a single-instance //! TA are flipped to [`SessionTarget::Dead`] *before* the per-UUID lock entry is @@ -522,12 +538,12 @@ impl SessionManager { /// /// 1. **Late stale Invoke/Close on a `Dead` session.** The cached instance /// was torn down earlier (which evicted the lock entry), but our call - /// to [`single_instance_lock`] resurrected the entry via + /// to [`Self::single_instance_lock`] resurrected the entry via /// `or_insert_with`. Evicting it on the way out keeps the lock map /// bounded. /// /// 2. **First OpenSession of a previously-unknown TA that turns out to be - /// multi-instance.** [`handle_open_session`] conservatively assumes + /// multi-instance.** The OpenSession handler conservatively assumes /// single-instance when flags are unknown and grabs the per-UUID lock. /// Once the TA loads and we learn it's actually multi-instance, the /// entry serves no future purpose and would otherwise leak. From c87e93b895d86c76f4aa655465eecedb598fcba4 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Fri, 22 May 2026 21:54:01 +0000 Subject: [PATCH 04/28] fix --- litebox_shim_optee/src/session.rs | 37 +++++++------------------------ 1 file changed, 8 insertions(+), 29 deletions(-) diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 499e0b3712..4622be7ffa 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -510,8 +510,7 @@ impl SessionManager { } /// Remove a single-instance TA from the cache only if the currently - /// cached `Arc` is the same as `expected`, and evict the per-UUID - /// serialization lock so its memory is reclaimed. + /// cached `Arc` is the same as `expected`. /// /// Callers tearing down on TA panic must have already called /// [`SessionMap::mark_sessions_dead_for_instance`] before invoking this, @@ -524,37 +523,17 @@ impl SessionManager { uuid: &TeeUuid, expected: &Arc, ) -> bool { - let removed = self.single_instance_cache.remove_if_same(uuid, expected); - if removed { - self.single_instance_locks.lock().remove(uuid); - } - removed + self.single_instance_cache.remove_if_same(uuid, expected) } - /// Evict the per-UUID serialization lock entry if no single-instance TA is - /// currently cached under that UUID. - /// - /// Two cleanup scenarios use this: - /// - /// 1. **Late stale Invoke/Close on a `Dead` session.** The cached instance - /// was torn down earlier (which evicted the lock entry), but our call - /// to [`Self::single_instance_lock`] resurrected the entry via - /// `or_insert_with`. Evicting it on the way out keeps the lock map - /// bounded. + /// Keep per-UUID locks for the lifetime of the manager. /// - /// 2. **First OpenSession of a previously-unknown TA that turns out to be - /// multi-instance.** The OpenSession handler conservatively assumes - /// single-instance when flags are unknown and grabs the per-UUID lock. - /// Once the TA loads and we learn it's actually multi-instance, the - /// entry serves no future purpose and would otherwise leak. - /// - /// The cache emptiness check is the safety guard: if the cache still has - /// an instance for this UUID, the entry is in legitimate use and we leave - /// it alone. + /// Replacing a lock while another core still holds the old one would let a + /// racing Open/Invoke/Close path bypass serialization for the same UUID. + /// This is intentionally a no-op, matching `known_flags` as per-UUID + /// lifecycle metadata retained for observed TAs. pub fn evict_single_instance_lock_if_unused(&self, uuid: &TeeUuid) { - if self.single_instance_cache.get(uuid).is_none() { - self.single_instance_locks.lock().remove(uuid); - } + let _ = uuid; } /// Get the total count of unique TA instances (for limit checking). From 1cdc9e31c7a5f834a59caf11d97ce969b7aa3f39 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Sat, 23 May 2026 00:27:15 +0000 Subject: [PATCH 05/28] simplify --- litebox_runner_lvbs/src/lib.rs | 25 +++---------------------- litebox_shim_optee/src/session.rs | 31 ++++++++++++++----------------- 2 files changed, 17 insertions(+), 39 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 775a774bc4..afcb46c235 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -44,8 +44,8 @@ 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::{ - ActiveSessionGuard, CreationReservation, SessionEntry, SessionIdGuard, SessionManager, - SessionTarget, TaInstance, allocate_session_id, + ActiveSessionGuard, CreationReservation, SessionIdGuard, SessionManager, SessionTarget, + TaInstance, allocate_session_id, }; use litebox_shim_optee::{NormalWorldConstPtr, NormalWorldMutPtr, UserConstPtr}; use once_cell::race::OnceBox; @@ -535,7 +535,7 @@ fn handle_open_session( // Resolve or create the TA instance. For single-instance TAs, the UUID // lock above serializes creation, reuse, and teardown for this TA. - let result = match session_manager().with_creation_slot(&ta_uuid, is_single_instance, || { + match session_manager().with_creation_slot(&ta_uuid, is_single_instance, || { open_session_new_instance( msg_args, msg_args_phys_addr, @@ -555,19 +555,7 @@ fn handle_open_session( ), Ok(CreationReservation::SlotReserved) => Ok(()), Err(e) => Err(e), - }; - - // If we conservatively held the per-UUID lock but no single-instance TA - // ended up cached under this UUID, evict the lock entry we just inserted. - // This covers (a) the TA turned out to be multi-instance and (b) load - // failed entirely. The helper's cache-empty check is the safety guard: - // if a single-instance TA is now cached for this UUID, the entry is in - // legitimate use and `evict_single_instance_lock_if_unused` is a no-op. - if single_instance_lock.is_some() { - session_manager().evict_single_instance_lock_if_unused(&ta_uuid); } - - result } /// Open a new session on an existing single-instance TA. @@ -970,12 +958,8 @@ fn open_session_new_instance( /// Consumes `active_guard` so it is dropped immediately after `unregister_session` /// recycles the id — the ordering matters: subsequent SMCs for `session_id` then /// see the new session (or `None`), neither of which depends on our guard. -/// -/// The caller's per-UUID lock guard (if any) stays in its own scope, so the -/// lock remains held across `evict_single_instance_lock_if_unused`. fn finalize_dead_session( session_id: u32, - session_entry: &SessionEntry, active_guard: ActiveSessionGuard<'_>, msg_args: &mut OpteeMsgArgs, msg_args_phys_addr: u64, @@ -984,7 +968,6 @@ fn finalize_dead_session( ) -> Result<(), OpteeSmcReturnCode> { session_manager().unregister_session(session_id); drop(active_guard); - session_manager().evict_single_instance_lock_if_unused(&session_entry.ta_uuid); msg_args.ret = return_code; msg_args.ret_origin = TeeOrigin::Tee; write_non_ta_msg_args_to_normal_world(msg_args, msg_args_phys_addr)?; @@ -1042,7 +1025,6 @@ fn handle_invoke_command( let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { return finalize_dead_session( session_id, - &session_entry, active_guard, msg_args, msg_args_phys_addr, @@ -1197,7 +1179,6 @@ fn handle_close_session( let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { return finalize_dead_session( session_id, - &session_entry, active_guard, msg_args, msg_args_phys_addr, diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 4622be7ffa..45b52decee 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -526,16 +526,6 @@ impl SessionManager { self.single_instance_cache.remove_if_same(uuid, expected) } - /// Keep per-UUID locks for the lifetime of the manager. - /// - /// Replacing a lock while another core still holds the old one would let a - /// racing Open/Invoke/Close path bypass serialization for the same UUID. - /// This is intentionally a no-op, matching `known_flags` as per-UUID - /// lifecycle metadata retained for observed TAs. - pub fn evict_single_instance_lock_if_unused(&self, uuid: &TeeUuid) { - let _ = uuid; - } - /// Get the total count of unique TA instances (for limit checking). /// /// This counts: @@ -564,16 +554,23 @@ impl SessionManager { /// Atomically reserve a creation slot and run `f` to create a new TA instance. /// - /// For single-instance TAs, the caller must already hold the per-UUID - /// serialization lock (see [`SessionManager::single_instance_lock`]), - /// which guarantees no other core is creating an instance for this UUID - /// concurrently. This function re-checks the single-instance cache under - /// the creation lock so a freshly-cached instance from a prior holder of - /// the per-UUID lock is observed before starting a new load. + /// # Caller contract + /// + /// For single-instance TAs, **the caller MUST already hold the per-UUID + /// serialization lock** returned by [`SessionManager::single_instance_lock`] + /// for `uuid`. That lock is the sole guarantee that no other core is + /// creating, reusing, or tearing down an instance for this UUID + /// concurrently — this function does NOT re-establish that exclusion + /// itself. Violating this contract can produce duplicate single-instance + /// TAs cached under the same UUID and break the single-instance invariant. + /// + /// Under that lock, this function re-checks the single-instance cache so a + /// freshly-cached instance from a prior holder of the per-UUID lock is + /// observed before starting a new load. /// /// For multi-instance TAs, each session gets its own independent /// `TaInstance`. Multiple cores may create instances of the same UUID - /// concurrently. + /// concurrently and no per-UUID lock is required. pub fn with_creation_slot( &self, uuid: &TeeUuid, From f3d84c86223cedda820e337ef6cf7a89951008ae Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Wed, 27 May 2026 00:45:00 +0000 Subject: [PATCH 06/28] refactoring --- litebox_runner_lvbs/src/lib.rs | 137 ++++----------- litebox_shim_optee/src/lib.rs | 8 +- litebox_shim_optee/src/session.rs | 270 +++++++++++++++++++++--------- 3 files changed, 235 insertions(+), 180 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index afcb46c235..dc01b10cd9 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -44,8 +44,7 @@ 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::{ - ActiveSessionGuard, CreationReservation, SessionIdGuard, SessionManager, SessionTarget, - TaInstance, allocate_session_id, + SessionIdGuard, SessionManager, SessionTarget, SessionToken, TaInstance, allocate_session_id, }; use litebox_shim_optee::{NormalWorldConstPtr, NormalWorldMutPtr, UserConstPtr}; use once_cell::race::OnceBox; @@ -518,44 +517,24 @@ fn handle_open_session( let client_identity = ta_req_info.client_identity; let params = &ta_req_info.params; - // Look up cached TA flags to determine single vs multi-instance. - // For the first-ever load of a UUID (no cached flags), conservatively - // assume single-instance to preserve all safety invariants. - let is_single_instance = session_manager() - .get_known_flags(&ta_uuid) - .is_none_or(|f| f.is_single_instance()); - - let single_instance_lock = - is_single_instance.then(|| session_manager().single_instance_lock(ta_uuid)); - let _single_instance_guard = if let Some(lock) = single_instance_lock.as_ref() { - Some(lock.try_lock().ok_or(OpteeSmcReturnCode::EThreadLimit)?) - } else { - None - }; - - // Resolve or create the TA instance. For single-instance TAs, the UUID - // lock above serializes creation, reuse, and teardown for this TA. - match session_manager().with_creation_slot(&ta_uuid, is_single_instance, || { - open_session_new_instance( + session_manager().with_creation_slot(&ta_uuid, |existing| match existing { + Some(instance) => open_session_single_instance( msg_args, msg_args_phys_addr, + instance, params, ta_uuid, - client_identity, &ta_req_info, - ) - }) { - Ok(CreationReservation::ExistingSingleInstance(existing)) => open_session_single_instance( + ), + None => open_session_new_instance( msg_args, msg_args_phys_addr, - existing.clone(), params, ta_uuid, + client_identity, &ta_req_info, ), - Ok(CreationReservation::SlotReserved) => Ok(()), - Err(e) => Err(e), - } + }) } /// Open a new session on an existing single-instance TA. @@ -637,9 +616,9 @@ fn open_session_single_instance( ); // Write error response BEFORE switching page tables (accesses user memory). - // The per-UUID lock held by the caller of `handle_open_session` keeps - // another core from tearing down the active page table while this core - // is copying TA outputs. + // The session token held by `with_creation_slot` keeps another core + // from tearing down the active page table while this core is copying + // TA outputs. let write_result = write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, @@ -655,10 +634,10 @@ fn open_session_single_instance( if return_code == TeeResult::TargetDead { debug_serial_println!("Single-instance TA panicked during OpenSession, cleaning up"); - // Mark sibling sessions dead BEFORE evicting the per-UUID lock - // (inside `remove_single_instance_if_same`). Otherwise a racing - // handler could allocate a fresh per-UUID lock and walk past a - // still-Live session entry. + // Mark sibling sessions dead BEFORE evicting the cached single + // instance. Otherwise a racing handler that subsequently + // acquires its own session token for the UUID could walk past + // a still-Live session entry. session_manager() .sessions() .mark_sessions_dead_for_instance(&instance_arc); @@ -955,19 +934,16 @@ fn open_session_new_instance( /// Tear down a `Dead` session entry observed at Invoke/Close handler entry. /// -/// Consumes `active_guard` so it is dropped immediately after `unregister_session` -/// recycles the id — the ordering matters: subsequent SMCs for `session_id` then -/// see the new session (or `None`), neither of which depends on our guard. +/// Consumes `token`; serialization is released at the end of this function. fn finalize_dead_session( session_id: u32, - active_guard: ActiveSessionGuard<'_>, + _token: SessionToken<'_>, msg_args: &mut OpteeMsgArgs, msg_args_phys_addr: u64, return_code: TeeResult, log_prefix: &str, ) -> Result<(), OpteeSmcReturnCode> { session_manager().unregister_session(session_id); - drop(active_guard); msg_args.ret = return_code; msg_args.ret_origin = TeeOrigin::Tee; write_non_ta_msg_args_to_normal_world(msg_args, msg_args_phys_addr)?; @@ -997,35 +973,16 @@ fn handle_invoke_command( let params = &ta_req_info.params; let session_id = ta_req_info.session; - // Get the session entry from the session map (need full entry for potential cleanup) - let session_entry = session_manager() - .get_session_entry(session_id) - .ok_or(OpteeSmcReturnCode::EBadCmd)?; - // Reserve this session id against concurrent SMC entry by another core. - let active_guard = session_manager() - .try_activate_session(session_id) - .ok_or(OpteeSmcReturnCode::EThreadLimit)?; - // For single-instance TAs, also take the per-UUID lock so sibling sessions - // on the same TA serialize against us. - let single_instance_lock = session_entry - .ta_flags - .is_single_instance() - .then(|| session_manager().single_instance_lock(session_entry.ta_uuid)); - let _single_instance_guard = if let Some(lock) = single_instance_lock.as_ref() { - Some(lock.try_lock().ok_or(OpteeSmcReturnCode::EThreadLimit)?) - } else { - None - }; - // Re-read after acquiring both serialization primitives to pick up any - // concurrent transition to `Dead` or removal that happened while we were - // waiting on the locks above. + let token = session_manager().try_acquire_for_session(session_id)?; + // Re-read under the token to pick up any concurrent transition to + // `Dead` or removal that happened during acquisition. let session_entry = session_manager() .get_session_entry(session_id) .ok_or(OpteeSmcReturnCode::EBadCmd)?; let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { return finalize_dead_session( session_id, - active_guard, + token, msg_args, msg_args_phys_addr, TeeResult::TargetDead, @@ -1076,9 +1033,8 @@ fn handle_invoke_command( let return_code = TeeResult::try_from(return_code).unwrap_or(TeeResult::GenericError); // Write response BEFORE switching page tables (accesses user memory). - // The active-session guard and (for single-instance) per-UUID lock prevent - // another core from tearing down the active page table while this core is - // copying TA outputs. + // The session token prevents another core from tearing down the active + // page table while this core is copying TA outputs. let write_result = write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, @@ -1101,9 +1057,9 @@ fn handle_invoke_command( let ta_flags = session_entry.ta_flags; if ta_flags.is_single_instance() { - // Mark siblings dead BEFORE evicting the per-UUID lock (inside - // `remove_single_instance_if_same`). Otherwise a racing handler - // could allocate a fresh per-UUID lock and walk past a still-Live + // Mark siblings dead BEFORE evicting the cached single instance. + // Otherwise a racing handler that subsequently acquires its own + // session token for the UUID could walk past a still-Live // session entry. session_manager() .sessions() @@ -1112,10 +1068,6 @@ fn handle_invoke_command( } session_manager().unregister_session(session_id); - // Release the active-session slot before the recycled id can race a - // new OpenSession; the session is gone from the map so any concurrent - // SMC for this id sees `None` regardless. - drop(active_guard); // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. @@ -1153,33 +1105,15 @@ fn handle_close_session( debug_serial_println!("CloseSession: session_id={}", session_id); - // Get the session entry from the session map - let session_entry = session_manager() - .get_session_entry(session_id) - .ok_or(OpteeSmcReturnCode::EBadCmd)?; - // Reserve this session id against concurrent SMC entry by another core. - let active_guard = session_manager() - .try_activate_session(session_id) - .ok_or(OpteeSmcReturnCode::EThreadLimit)?; - // For single-instance TAs, also take the per-UUID lock so sibling sessions - // on the same TA serialize against us. - let single_instance_lock = session_entry - .ta_flags - .is_single_instance() - .then(|| session_manager().single_instance_lock(session_entry.ta_uuid)); - let _single_instance_guard = if let Some(lock) = single_instance_lock.as_ref() { - Some(lock.try_lock().ok_or(OpteeSmcReturnCode::EThreadLimit)?) - } else { - None - }; - // Re-read after acquiring both serialization primitives. + let token = session_manager().try_acquire_for_session(session_id)?; + // Re-read under the token. let session_entry = session_manager() .get_session_entry(session_id) .ok_or(OpteeSmcReturnCode::EBadCmd)?; let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { return finalize_dead_session( session_id, - active_guard, + token, msg_args, msg_args_phys_addr, TeeResult::Success, @@ -1223,12 +1157,9 @@ fn handle_close_session( None, ); - // Remove the session entry from the map + // Remove the session entry from the map. The session token drops at + // the end of this function. let removed_entry = session_manager().unregister_session(session_id); - // Release the active-session slot before the recycled id can race a - // new OpenSession; the session is gone from the map so any concurrent - // SMC for this id sees `None` regardless. - drop(active_guard); // Check if this was the last session using the TA instance by counting // remaining sessions that reference this instance. @@ -1249,9 +1180,9 @@ fn handle_close_session( return write_result; } - // Clear single-instance cache (and evict its per-UUID lock) if this - // was a single-instance TA. No sibling sessions remain on this - // instance, so we don't need to mark anything `Dead` first. + // Clear the cached single instance if this was a single-instance TA. + // No sibling sessions remain, so we don't need to mark anything + // `Dead` first. if entry.ta_flags.is_single_instance() { let _ = session_manager().remove_single_instance_if_same(&entry.ta_uuid, &instance_arc); diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index 8394dda8a2..55f4591b2b 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -39,8 +39,8 @@ pub mod ptr; // Re-export session management types for convenience pub use session::{ - ActiveSessionGuard, CreationReservation, MAX_TA_INSTANCES, SessionEntry, SessionManager, - SessionMap, SessionTarget, SingleInstanceCache, TaInstance, allocate_session_id, + MAX_TA_INSTANCES, SessionEntry, SessionManager, SessionMap, SessionTarget, SessionToken, + SingleInstanceCache, TaInstance, allocate_session_id, }; const MAX_KERNEL_BUF_SIZE: usize = 0x80_000; @@ -1449,6 +1449,10 @@ impl SessionIdPool { } /// Recycle a session ID for reuse. Fallback IDs are not recycled. + /// + /// "Recycled" only marks the bit free; [`IdPool`](litebox::utils::id_pool::IdPool) + /// is hint+wrap, so the ID is not handed out again until every higher ID + /// has been allocated first. pub fn recycle(session_id: u32) { if session_id == 0 || session_id > Self::MAX_RECYCLABLE_SESSION_ID { return; diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 45b52decee..db37643997 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -12,17 +12,22 @@ //! //! TA execution is serialized externally; [`TaInstance`] is shared as a plain //! `Arc` without an inner mutex. The exclusivity invariant lives in -//! [`SessionManager`]: +//! [`SessionManager`] and is acquired through a single RAII [`SessionToken`] +//! that bundles whichever locks the current operation requires: //! //! - **Single-instance TAs** (with `TA_FLAG_SINGLE_INSTANCE | TA_FLAG_MULTI_SESSION`) -//! share one [`TaInstance`] across all sessions. Open/Invoke/Close serialize on -//! a per-UUID `SpinMutex` returned by [`SessionManager::single_instance_lock`], -//! acquired non-blockingly at handler entry. +//! share one [`TaInstance`] across all sessions. The token internally holds a +//! per-UUID `SpinMutex` so Open/Invoke/Close serialize on the same UUID. //! -//! - **Multi-instance TAs** have one [`TaInstance`] per session. Invoke/Close -//! serialize on a per-`session_id` entry acquired via -//! [`SessionManager::try_activate_session`], also acquired non-blockingly. -//! Different sessions run in parallel on their own instances. +//! - **Multi-instance TAs** have one [`TaInstance`] per session. The token +//! internally holds a per-`session_id` marker so Invoke/Close cannot +//! re-enter the same session concurrently, while different sessions run +//! in parallel on their own instances. +//! +//! OpenSession callers go through [`SessionManager::with_creation_slot`], +//! which manages the token internally. Invoke/Close callers acquire a token +//! via [`SessionManager::try_acquire_for_session`]. Both are non-blocking +//! and return `EThreadLimit` on contention. //! //! ### Difference from OP-TEE OS //! @@ -44,11 +49,11 @@ //! requiring RPCs that would give untrusted code control over secure world execution. //! //! On panic teardown or last-session close, sibling sessions of a single-instance -//! TA are flipped to [`SessionTarget::Dead`] *before* the per-UUID lock entry is +//! TA are flipped to [`SessionTarget::Dead`] *before* the cached instance is //! evicted (see [`SessionManager::remove_single_instance_if_same`]). A racing -//! handler that subsequently allocates a fresh per-UUID lock will therefore -//! observe `Dead` on its re-read of the session entry and short-circuit through -//! the dead-target path. +//! handler that subsequently acquires its own [`SessionToken`] for the UUID +//! will therefore observe `Dead` on its re-read of the session entry and +//! short-circuit through the dead-target path. //! //! Reference: //! @@ -363,30 +368,51 @@ impl Drop for SessionIdGuard { } } -/// RAII guard returned by [`SessionManager::try_activate_session`] that removes -/// the session id from the active set on drop. Provides per-session-id -/// serialization for Invoke/Close handlers without requiring a mutex on the -/// `TaInstance` itself. -pub struct ActiveSessionGuard<'a> { +/// RAII token bundling the serialization primitives required to safely +/// execute an OP-TEE TA operation. +/// +/// Acquired non-blockingly via [`SessionManager::try_acquire_for_session`] +/// for Invoke/Close on an existing session. OpenSession uses the same +/// token type internally through [`SessionManager::with_creation_slot`], +/// which manages acquisition and release on the caller's behalf. +/// +/// Holds whichever combination of locks is required for the operation: +/// +/// - **Single-instance TAs**: a per-UUID `SpinMutex` that serializes all +/// sessions on the same TA. +/// - **Existing-session operations** (Invoke/Close): a per-session-id marker +/// that prevents concurrent SMC entry by another core for the same id. +/// +/// For multi-instance OpenSession, the token holds nothing (each session +/// gets its own private instance, so no exclusion is required). +/// +/// On drop, the per-UUID lock is released first, then the per-session-id +/// marker. +pub struct SessionToken<'a> { manager: &'a SessionManager, - session_id: u32, + /// Held `Arc` of the per-UUID `SpinMutex`. The guard returned by + /// `try_lock()` was [`core::mem::forget`]-ed at acquisition time; this + /// type's `Drop` calls `force_unlock` to release the mutex. The `Arc` + /// keeps the mutex alive across acquisition and release. + uuid_lock: Option>>, + /// Session id reserved in [`SessionManager::active_sessions`]. + active_session_id: Option, } -impl Drop for ActiveSessionGuard<'_> { +impl Drop for SessionToken<'_> { fn drop(&mut self) { - self.manager.active_sessions.lock().remove(&self.session_id); + if let Some(lock) = self.uuid_lock.take() { + // SAFETY: This token holds the per-UUID lock because the + // acquisition path called `try_lock()` and forgot the resulting + // guard. No other holder exists, so `force_unlock` is sound. + unsafe { lock.force_unlock() }; + } + if let Some(id) = self.active_session_id.take() { + self.manager.active_sessions.lock().remove(&id); + } } } -/// Result of [`SessionManager::with_creation_slot`]. -pub enum CreationReservation { - /// An existing single-instance TA was found (another core cached it - /// between our initial lookup and the reservation). Reuse this instance. - ExistingSingleInstance(Arc), - /// The creation closure ran successfully inside the reserved slot. - SlotReserved, -} - /// State for coordinating concurrent instance creation. /// /// Guarded by a single lock to provide atomic capacity checks. @@ -464,8 +490,18 @@ impl SessionManager { self.known_flags.lock().get(uuid).copied() } - /// Get the serialization lock for a single-instance TA UUID. - pub fn single_instance_lock(&self, uuid: TeeUuid) -> Arc> { + /// Whether `uuid` should be treated as single-instance for serialization. + /// + /// Returns the cached `is_single_instance()` if known, or `true` for the + /// first-ever load (we have not yet observed the TA's flags) to preserve + /// safety invariants conservatively. + fn assume_single_instance(&self, uuid: &TeeUuid) -> bool { + self.get_known_flags(uuid) + .is_none_or(|f| f.is_single_instance()) + } + + /// Get or create the per-UUID serialization mutex `Arc`. + fn uuid_lock_arc(&self, uuid: TeeUuid) -> Arc> { self.single_instance_locks .lock() .entry(uuid) @@ -473,18 +509,108 @@ impl SessionManager { .clone() } - /// Try to mark `session_id` as actively handled. Returns a guard that - /// removes the marker on drop, or `None` if another core is already - /// inside an Invoke/Close handler for this session id. - pub fn try_activate_session(&self, session_id: u32) -> Option> { - if self.active_sessions.lock().insert(session_id) { - Some(ActiveSessionGuard { - manager: self, - session_id, - }) + /// Try to take the per-UUID serialization mutex non-blockingly. On + /// success returns the `Arc` whose forgotten guard is owned by the + /// caller — release via `force_unlock` on the returned `Arc`. + fn try_acquire_uuid_lock(&self, uuid: TeeUuid) -> Option>> { + let lock = self.uuid_lock_arc(uuid); + let guard = lock.try_lock()?; + // The lock now belongs to the SessionToken about to wrap us. Forget + // the guard so its `Drop` does not unlock; the token's `Drop` calls + // `force_unlock` via the retained `Arc`. + core::mem::forget(guard); + Some(lock) + } + + /// Acquire a [`SessionToken`] for an OpenSession request. + /// + /// For single-instance TAs (including first-ever load of an unknown + /// UUID) this takes the per-UUID `SpinMutex` non-blockingly. For + /// already-known multi-instance TAs the returned token holds no locks — + /// each session creates its own private instance, so no exclusion is + /// required. + /// + /// Returns `Err(EThreadLimit)` if another core is currently inside an + /// operation on the same single-instance UUID. + fn try_acquire_for_open(&self, uuid: TeeUuid) -> Result, OpteeSmcReturnCode> { + let uuid_lock = if self.assume_single_instance(&uuid) { + Some( + self.try_acquire_uuid_lock(uuid) + .ok_or(OpteeSmcReturnCode::EThreadLimit)?, + ) } else { None + }; + Ok(SessionToken { + manager: self, + uuid_lock, + active_session_id: None, + }) + } + + /// Acquire a [`SessionToken`] for an Invoke/Close on an existing session. + /// + /// Always reserves the per-session-id slot in `active_sessions`. For + /// single-instance TAs additionally takes the per-UUID `SpinMutex` so + /// sibling sessions on the same TA serialize against this operation. + /// + /// Returns `Err(EBadCmd)` if `session_id` is not registered, or + /// `Err(EThreadLimit)` if another core is inside the same session or + /// holds the per-UUID lock for the same single-instance TA. On failure + /// any partial acquisition is released via the token's `Drop`. + /// + /// Defense in depth: after the per-session-id marker is held, the entry + /// is re-read and its `(uuid, flags)` validated against the pre-marker + /// snapshot used to decide whether to take the per-UUID lock. If they + /// diverge (the id was recycled and reused under a different TA between + /// our first read and the marker insert), `Err(EThreadLimit)` is + /// returned so the Linux driver retries — a fresh acquisition will see + /// the new entry from the start. This guards against the read/marker + /// TOCTOU without relying on + /// [`IdPool`](litebox::utils::id_pool::IdPool)'s hint+wrap to quarantine + /// recycled ids. + pub fn try_acquire_for_session( + &self, + session_id: u32, + ) -> Result, OpteeSmcReturnCode> { + let entry = self + .sessions + .get_entry(session_id) + .ok_or(OpteeSmcReturnCode::EBadCmd)?; + let snapshot_uuid = entry.ta_uuid; + let snapshot_single = entry.ta_flags.is_single_instance(); + drop(entry); + + if !self.active_sessions.lock().insert(session_id) { + return Err(OpteeSmcReturnCode::EThreadLimit); } + let mut token = SessionToken { + manager: self, + uuid_lock: None, + active_session_id: Some(session_id), + }; + + // Validate the snapshot under the marker. If the entry has changed + // identity (or vanished), our snapshot is stale; bail so the caller + // retries with a fresh view. Token's `Drop` releases the marker. + let entry_now = self + .sessions + .get_entry(session_id) + .ok_or(OpteeSmcReturnCode::EBadCmd)?; + if entry_now.ta_uuid != snapshot_uuid + || entry_now.ta_flags.is_single_instance() != snapshot_single + { + return Err(OpteeSmcReturnCode::EThreadLimit); + } + + if snapshot_single { + // On failure, dropping `token` releases the marker we just took. + token.uuid_lock = Some( + self.try_acquire_uuid_lock(snapshot_uuid) + .ok_or(OpteeSmcReturnCode::EThreadLimit)?, + ); + } + Ok(token) } /// Register a new session. @@ -514,10 +640,10 @@ impl SessionManager { /// /// Callers tearing down on TA panic must have already called /// [`SessionMap::mark_sessions_dead_for_instance`] before invoking this, - /// so any handler that subsequently allocates a fresh per-UUID lock will - /// observe `Dead` on its re-read of the session entry. Callers on the - /// last-session-close path may skip the mark step — by that point there - /// are no sibling sessions to fence out. + /// so any handler that subsequently acquires its own [`SessionToken`] + /// for the UUID will observe `Dead` on its re-read of the session + /// entry. Callers on the last-session-close path may skip the mark + /// step — by that point there are no sibling sessions to fence out. pub fn remove_single_instance_if_same( &self, uuid: &TeeUuid, @@ -552,58 +678,52 @@ impl SessionManager { self.instance_count() >= MAX_TA_INSTANCES } - /// Atomically reserve a creation slot and run `f` to create a new TA instance. - /// - /// # Caller contract + /// Drive an OpenSession to completion under the right serialization. /// - /// For single-instance TAs, **the caller MUST already hold the per-UUID - /// serialization lock** returned by [`SessionManager::single_instance_lock`] - /// for `uuid`. That lock is the sole guarantee that no other core is - /// creating, reusing, or tearing down an instance for this UUID - /// concurrently — this function does NOT re-establish that exclusion - /// itself. Violating this contract can produce duplicate single-instance - /// TAs cached under the same UUID and break the single-instance invariant. + /// Internally acquires the per-UUID `SpinMutex` for single-instance TAs + /// (or no lock for known multi-instance TAs), then either: /// - /// Under that lock, this function re-checks the single-instance cache so a - /// freshly-cached instance from a prior holder of the per-UUID lock is - /// observed before starting a new load. + /// - Calls `f(Some(existing))` if a cached single-instance TA is found + /// for `uuid`. The lock is held throughout the call so the existing + /// instance cannot be torn down or replaced beneath `f`. + /// - Reserves a creation slot (atomic capacity check including in-flight + /// creations) and calls `f(None)` to load and register a new instance. + /// The slot is released when `f` returns, regardless of outcome. /// - /// For multi-instance TAs, each session gets its own independent - /// `TaInstance`. Multiple cores may create instances of the same UUID - /// concurrently and no per-UUID lock is required. - pub fn with_creation_slot( - &self, - uuid: &TeeUuid, - is_single_instance: bool, - f: F, - ) -> Result + /// The per-UUID lock is released when this function returns; `f` runs + /// under it. For multi-instance TAs each session gets its own + /// independent `TaInstance`, so no per-UUID exclusion is required. + pub fn with_creation_slot(&self, uuid: &TeeUuid, f: F) -> Result<(), OpteeSmcReturnCode> where - F: FnOnce() -> Result<(), OpteeSmcReturnCode>, + F: FnOnce(Option>) -> Result<(), OpteeSmcReturnCode>, { - { - let mut state = self.creation_state.lock(); + let token = self.try_acquire_for_open(*uuid)?; + let is_single_instance = token.uuid_lock.is_some(); - if is_single_instance && let Some(existing) = self.single_instance_cache.get(uuid) { - return Ok(CreationReservation::ExistingSingleInstance(existing)); - } + // For single-instance TAs the per-UUID lock above keeps our UUID's + // cache entry stable. For multi-instance we don't consult the cache. + if is_single_instance && let Some(existing) = self.single_instance_cache.get(uuid) { + return f(Some(existing)); + } + { + let mut state = self.creation_state.lock(); // Capacity check including in-flight creations. let total = self.instance_count() + state.pending_count; if total >= MAX_TA_INSTANCES { return Err(OpteeSmcReturnCode::ENomem); } - state.pending_count += 1; } - let result = f(); + let result = f(None); { let mut state = self.creation_state.lock(); state.pending_count = state.pending_count.saturating_sub(1); } - result.map(|()| CreationReservation::SlotReserved) + result } } From ded0f7f5b57a8db923d3947fd8c16ded6a8ef3a4 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Wed, 27 May 2026 04:02:55 +0000 Subject: [PATCH 07/28] refactoring further --- litebox_runner_lvbs/src/lib.rs | 339 +++++++++++++++--------------- litebox_shim_optee/src/lib.rs | 4 +- litebox_shim_optee/src/session.rs | 139 +++++++----- 3 files changed, 249 insertions(+), 233 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index dc01b10cd9..b986aacda6 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -44,7 +44,7 @@ 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::{ - SessionIdGuard, SessionManager, SessionTarget, SessionToken, TaInstance, allocate_session_id, + SessionIdGuard, SessionManager, SessionTarget, TaInstance, allocate_session_id, }; use litebox_shim_optee::{NormalWorldConstPtr, NormalWorldMutPtr, UserConstPtr}; use once_cell::race::OnceBox; @@ -517,7 +517,7 @@ fn handle_open_session( let client_identity = ta_req_info.client_identity; let params = &ta_req_info.params; - session_manager().with_creation_slot(&ta_uuid, |existing| match existing { + session_manager().with_ta(&ta_uuid, |existing| match existing { Some(instance) => open_session_single_instance( msg_args, msg_args_phys_addr, @@ -616,7 +616,7 @@ fn open_session_single_instance( ); // Write error response BEFORE switching page tables (accesses user memory). - // The session token held by `with_creation_slot` keeps another core + // The session token held by `with_ta` keeps another core // from tearing down the active page table while this core is copying // TA outputs. let write_result = write_msg_args_to_normal_world( @@ -709,7 +709,7 @@ fn open_session_single_instance( /// Create a new TA instance for a session. /// -/// The caller must invoke this inside [`SessionManager::with_creation_slot`] +/// The caller must invoke this inside [`SessionManager::with_ta`] /// to ensure a creation slot is held during execution and released afterward. /// /// If ldelf loading or OpenSession entry point fails, the page table is torn down. @@ -934,10 +934,10 @@ fn open_session_new_instance( /// Tear down a `Dead` session entry observed at Invoke/Close handler entry. /// -/// Consumes `token`; serialization is released at the end of this function. +/// Runs inside `with_session`'s closure, so the session token is alive for +/// the duration of this call and released when the closure returns. fn finalize_dead_session( session_id: u32, - _token: SessionToken<'_>, msg_args: &mut OpteeMsgArgs, msg_args_phys_addr: u64, return_code: TeeResult, @@ -973,119 +973,113 @@ fn handle_invoke_command( let params = &ta_req_info.params; let session_id = ta_req_info.session; - let token = session_manager().try_acquire_for_session(session_id)?; - // Re-read under the token to pick up any concurrent transition to - // `Dead` or removal that happened during acquisition. - let session_entry = session_manager() - .get_session_entry(session_id) - .ok_or(OpteeSmcReturnCode::EBadCmd)?; - let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { - return finalize_dead_session( - session_id, - token, - msg_args, - msg_args_phys_addr, - TeeResult::TargetDead, - "InvokeCommand", - ); - }; - let task_pt_id = instance_arc.task_page_table_id; - - let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; - - debug_serial_println!( - "InvokeCommand: session_id={}, task_pt_id={}, cmd_id={}", - session_id, - task_pt_id, - cmd_id - ); + session_manager().with_session(session_id, |session_entry| { + let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { + return finalize_dead_session( + session_id, + msg_args, + msg_args_phys_addr, + TeeResult::TargetDead, + "InvokeCommand", + ); + }; + let task_pt_id = instance_arc.task_page_table_id; - // Load TA context with parameters and cmd_id - pass actual session_id - let entrypoints_ref = instance_arc.loaded_program.entrypoints.as_ref().unwrap(); - entrypoints_ref - .load_ta_context( - params.as_slice(), - Some(session_id), - UteeEntryFunc::InvokeCommand as u32, - Some(cmd_id), - ) - .map_err(|_| OpteeSmcReturnCode::EBadCmd)?; + let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; - // Run the TA entry function using reference-based reenter to avoid moving the shim - let mut ctx = litebox_common_linux::PtRegs::default(); - unsafe { - litebox_platform_lvbs::reenter_thread_ref( - instance_arc.loaded_program.entrypoints.as_ref().unwrap(), - &mut ctx, + debug_serial_println!( + "InvokeCommand: session_id={}, task_pt_id={}, cmd_id={}", + session_id, + task_pt_id, + cmd_id ); - } - - // params_address is constant - stack buffer is reused across invocations - let params_address = instance_arc - .loaded_program - .params_address - .ok_or(OpteeSmcReturnCode::EBadAddr)?; - let ta_params = UserConstPtr::::from_usize(params_address) - .read_at_offset(0) - .ok_or(OpteeSmcReturnCode::EBadAddr)?; - - let return_code: u32 = ctx.rax.trunc(); - let return_code = TeeResult::try_from(return_code).unwrap_or(TeeResult::GenericError); - // Write response BEFORE switching page tables (accesses user memory). - // The session token prevents another core from tearing down the active - // page table while this core is copying TA outputs. - let write_result = write_msg_args_to_normal_world( - msg_args, - msg_args_phys_addr, - return_code, - None, - Some(&ta_params), - Some(&ta_req_info), - ); + // Load TA context with parameters and cmd_id - pass actual session_id + let entrypoints_ref = instance_arc.loaded_program.entrypoints.as_ref().unwrap(); + entrypoints_ref + .load_ta_context( + params.as_slice(), + Some(session_id), + UteeEntryFunc::InvokeCommand as u32, + Some(cmd_id), + ) + .map_err(|_| OpteeSmcReturnCode::EBadCmd)?; + + // Run the TA entry function using reference-based reenter to avoid moving the shim + let mut ctx = litebox_common_linux::PtRegs::default(); + unsafe { + litebox_platform_lvbs::reenter_thread_ref( + instance_arc.loaded_program.entrypoints.as_ref().unwrap(), + &mut ctx, + ); + } - // Per OP-TEE OS: if TA panics (TARGET_DEAD), the TA context is - // unrecoverable; all sessions on the same single-instance TA are - // implicitly dead (Ref: tee_ta_invoke_command() in tee_ta_manager.c). - if return_code == TeeResult::TargetDead { - debug_serial_println!( - "InvokeCommand: TA panicked (TARGET_DEAD), session_id={}", - session_id + // params_address is constant - stack buffer is reused across invocations + let params_address = instance_arc + .loaded_program + .params_address + .ok_or(OpteeSmcReturnCode::EBadAddr)?; + let ta_params = UserConstPtr::::from_usize(params_address) + .read_at_offset(0) + .ok_or(OpteeSmcReturnCode::EBadAddr)?; + + let return_code: u32 = ctx.rax.trunc(); + let return_code = TeeResult::try_from(return_code).unwrap_or(TeeResult::GenericError); + + // Write response BEFORE switching page tables (accesses user memory). + // The session token prevents another core from tearing down the active + // page table while this core is copying TA outputs. + let write_result = write_msg_args_to_normal_world( + msg_args, + msg_args_phys_addr, + return_code, + None, + Some(&ta_params), + Some(&ta_req_info), ); - let ta_uuid = session_entry.ta_uuid; - let ta_flags = session_entry.ta_flags; + // Per OP-TEE OS: if TA panics (TARGET_DEAD), the TA context is + // unrecoverable; all sessions on the same single-instance TA are + // implicitly dead (Ref: tee_ta_invoke_command() in tee_ta_manager.c). + if return_code == TeeResult::TargetDead { + debug_serial_println!( + "InvokeCommand: TA panicked (TARGET_DEAD), session_id={}", + session_id + ); - if ta_flags.is_single_instance() { - // Mark siblings dead BEFORE evicting the cached single instance. - // Otherwise a racing handler that subsequently acquires its own - // session token for the UUID could walk past a still-Live - // session entry. - session_manager() - .sessions() - .mark_sessions_dead_for_instance(&instance_arc); - let _ = session_manager().remove_single_instance_if_same(&ta_uuid, &instance_arc); - } + let ta_uuid = session_entry.ta_uuid; + let ta_flags = session_entry.ta_flags; + + if ta_flags.is_single_instance() { + // Mark siblings dead BEFORE evicting the cached single instance. + // Otherwise a racing handler entering with_session/with_ta + // for the UUID could walk past a still-Live session entry. + session_manager() + .sessions() + .mark_sessions_dead_for_instance(&instance_arc); + let _ = session_manager().remove_single_instance_if_same(&ta_uuid, &instance_arc); + } - session_manager().unregister_session(session_id); + session_manager().unregister_session(session_id); - // 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(&instance_arc.shim, task_pt_id); - }; + // 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(&instance_arc.shim, task_pt_id); + }; - debug_serial_println!( - "InvokeCommand: cleaned up dead TA instance, task_pt_id={}", - task_pt_id - ); + debug_serial_println!( + "InvokeCommand: cleaned up dead TA instance, task_pt_id={}", + task_pt_id + ); - // TODO: Per OP-TEE OS semantics, if the TA has INSTANCE_KEEP_ALIVE but not - // INSTANCE_KEEP_CRASHED, we should respawn the TA here instead of just - // cleaning it up. Currently we always clean up on panic. - } + // TODO: Per OP-TEE OS semantics, if the TA has INSTANCE_KEEP_ALIVE but not + // INSTANCE_KEEP_CRASHED, we should respawn the TA here instead of just + // cleaning it up. Currently we always clean up on panic. + } - write_result + write_result + }) } /// Handle CloseSession command. @@ -1105,71 +1099,67 @@ fn handle_close_session( debug_serial_println!("CloseSession: session_id={}", session_id); - let token = session_manager().try_acquire_for_session(session_id)?; - // Re-read under the token. - let session_entry = session_manager() - .get_session_entry(session_id) - .ok_or(OpteeSmcReturnCode::EBadCmd)?; - let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { - return finalize_dead_session( - session_id, - token, + session_manager().with_session(session_id, |session_entry| { + let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { + return finalize_dead_session( + session_id, + msg_args, + msg_args_phys_addr, + TeeResult::Success, + "CloseSession", + ); + }; + let task_pt_id = instance_arc.task_page_table_id; + + let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; + + // Load TA context for CloseSession (no params, no cmd_id) - pass actual session_id + instance_arc + .loaded_program + .entrypoints + .as_ref() + .unwrap() + .load_ta_context( + &[], + Some(session_id), + UteeEntryFunc::CloseSession as u32, + None, + ) + .map_err(|_| OpteeSmcReturnCode::EBadCmd)?; + + // Run the TA entry function (TA_CloseSessionEntryPoint) + let mut ctx = litebox_common_linux::PtRegs::default(); + unsafe { + litebox_platform_lvbs::reenter_thread_ref( + instance_arc.loaded_program.entrypoints.as_ref().unwrap(), + &mut ctx, + ); + } + + // CloseSession always succeeds (TA_CloseSessionEntryPoint returns void) + let write_result = write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, TeeResult::Success, - "CloseSession", - ); - }; - let task_pt_id = instance_arc.task_page_table_id; - - let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; - - // Load TA context for CloseSession (no params, no cmd_id) - pass actual session_id - instance_arc - .loaded_program - .entrypoints - .as_ref() - .unwrap() - .load_ta_context( - &[], - Some(session_id), - UteeEntryFunc::CloseSession as u32, None, - ) - .map_err(|_| OpteeSmcReturnCode::EBadCmd)?; - - // Run the TA entry function (TA_CloseSessionEntryPoint) - let mut ctx = litebox_common_linux::PtRegs::default(); - unsafe { - litebox_platform_lvbs::reenter_thread_ref( - instance_arc.loaded_program.entrypoints.as_ref().unwrap(), - &mut ctx, + None, + None, ); - } - - // CloseSession always succeeds (TA_CloseSessionEntryPoint returns void) - let write_result = write_msg_args_to_normal_world( - msg_args, - msg_args_phys_addr, - TeeResult::Success, - None, - None, - None, - ); - // Remove the session entry from the map. The session token drops at - // the end of this function. - let removed_entry = session_manager().unregister_session(session_id); + // Remove the session entry from the map. The session token drops + // when the enclosing `with_session` closure returns. + let removed_entry = session_manager().unregister_session(session_id); - // Check if this was the last session using the TA instance by counting - // remaining sessions that reference this instance. - let remaining_sessions = session_manager() - .sessions() - .count_sessions_for_instance(&instance_arc); + // Check if this was the last session using the TA instance by counting + // remaining sessions that reference this instance. + let remaining_sessions = session_manager() + .sessions() + .count_sessions_for_instance(&instance_arc); - // If this was the last session using the TA instance, clean up (unless keep_alive is set) - if remaining_sessions == 0 { - if let Some(entry) = removed_entry { + // If this was the last session using the TA instance, clean up (unless keep_alive is set) + if remaining_sessions == 0 + && let Some(entry) = removed_entry + { // If this is a single-instance TA with keep_alive flag, don't remove it from memory. // Note: keep_alive is only meaningful for single-instance TAs. if entry.ta_flags.is_single_instance() && entry.ta_flags.is_keep_alive() { @@ -1180,12 +1170,11 @@ fn handle_close_session( return write_result; } - // Clear the cached single instance if this was a single-instance TA. - // No sibling sessions remain, so we don't need to mark anything - // `Dead` first. + // If this was a single-instance TA, clear the cached instance. This is safe because + // we confirm no sibling sessions remain. We don't need to mark anything `Dead` first. if entry.ta_flags.is_single_instance() { - let _ = - session_manager().remove_single_instance_if_same(&entry.ta_uuid, &instance_arc); + let _ = session_manager() + .remove_single_instance_if_same(&entry.ta_uuid, &instance_arc); } // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. @@ -1197,15 +1186,15 @@ fn handle_close_session( "CloseSession complete: deleted task_pt_id={} (last session)", task_pt_id ); + } else { + debug_serial_println!( + "CloseSession complete: session_id={}, other sessions remaining on TA", + session_id + ); } - } else { - debug_serial_println!( - "CloseSession complete: session_id={}, other sessions remaining on TA", - session_id - ); - } - write_result + write_result + }) } /// Update msg_args with return values and write back to normal world memory. diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index 55f4591b2b..f4adceffba 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -39,8 +39,8 @@ pub mod ptr; // Re-export session management types for convenience pub use session::{ - MAX_TA_INSTANCES, SessionEntry, SessionManager, SessionMap, SessionTarget, SessionToken, - SingleInstanceCache, TaInstance, allocate_session_id, + MAX_TA_INSTANCES, SessionEntry, SessionManager, SessionMap, SessionTarget, SingleInstanceCache, + TaInstance, allocate_session_id, }; const MAX_KERNEL_BUF_SIZE: usize = 0x80_000; diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index db37643997..2f1adfe270 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -12,7 +12,7 @@ //! //! TA execution is serialized externally; [`TaInstance`] is shared as a plain //! `Arc` without an inner mutex. The exclusivity invariant lives in -//! [`SessionManager`] and is acquired through a single RAII [`SessionToken`] +//! [`SessionManager`] and is acquired through an internal RAII `SessionToken` //! that bundles whichever locks the current operation requires: //! //! - **Single-instance TAs** (with `TA_FLAG_SINGLE_INSTANCE | TA_FLAG_MULTI_SESSION`) @@ -24,10 +24,10 @@ //! re-enter the same session concurrently, while different sessions run //! in parallel on their own instances. //! -//! OpenSession callers go through [`SessionManager::with_creation_slot`], -//! which manages the token internally. Invoke/Close callers acquire a token -//! via [`SessionManager::try_acquire_for_session`]. Both are non-blocking -//! and return `EThreadLimit` on contention. +//! Both [`SessionManager::with_ta`] (OpenSession) and +//! [`SessionManager::with_session`] (Invoke/Close) acquire the token +//! non-blockingly, run the caller's closure under it, and release on return. +//! On contention they return `EThreadLimit`. //! //! ### Difference from OP-TEE OS //! @@ -51,9 +51,10 @@ //! On panic teardown or last-session close, sibling sessions of a single-instance //! TA are flipped to [`SessionTarget::Dead`] *before* the cached instance is //! evicted (see [`SessionManager::remove_single_instance_if_same`]). A racing -//! handler that subsequently acquires its own [`SessionToken`] for the UUID -//! will therefore observe `Dead` on its re-read of the session entry and -//! short-circuit through the dead-target path. +//! handler that subsequently enters [`SessionManager::with_ta`] or +//! [`SessionManager::with_session`] for the UUID will therefore observe `Dead` +//! on its re-read of the session entry and short-circuit through the +//! dead-target path. //! //! Reference: //! @@ -140,10 +141,11 @@ pub struct TaInstance { pub task_page_table_id: usize, } -// SAFETY: TaInstance is shared as `Arc`, but only one core is ever -// inside the TA at a time: single-instance TAs serialize on the per-UUID lock, -// multi-instance TAs on the per-`session_id` entry in `active_sessions`. See -// the module-level "Concurrency Model" doc. +// SAFETY: `TaInstance`'s interior (`shim`, `loaded_program`) is not +// auto-`Send`/`Sync`, but every access goes through a `SessionToken` that +// serializes execution on the per-UUID lock (single-instance TAs) or the +// per-`session_id` marker (multi-instance TAs), so at most one core is +// ever inside a given instance. See the module-level "Concurrency Model". unsafe impl Send for TaInstance {} unsafe impl Sync for TaInstance {} @@ -371,12 +373,9 @@ impl Drop for SessionIdGuard { /// RAII token bundling the serialization primitives required to safely /// execute an OP-TEE TA operation. /// -/// Acquired non-blockingly via [`SessionManager::try_acquire_for_session`] -/// for Invoke/Close on an existing session. OpenSession uses the same -/// token type internally through [`SessionManager::with_creation_slot`], -/// which manages acquisition and release on the caller's behalf. -/// -/// Holds whichever combination of locks is required for the operation: +/// Held only inside [`SessionManager::with_ta`] (OpenSession) and +/// [`SessionManager::with_session`] (Invoke/Close); never exposed to +/// external callers. Bundles whichever combination of locks is required: /// /// - **Single-instance TAs**: a per-UUID `SpinMutex` that serializes all /// sessions on the same TA. @@ -388,7 +387,7 @@ impl Drop for SessionIdGuard { /// /// On drop, the per-UUID lock is released first, then the per-session-id /// marker. -pub struct SessionToken<'a> { +pub(crate) struct SessionToken<'a> { manager: &'a SessionManager, /// Held `Arc` of the per-UUID `SpinMutex`. The guard returned by /// `try_lock()` was [`core::mem::forget`]-ed at acquisition time; this @@ -413,15 +412,6 @@ impl Drop for SessionToken<'_> { } } -/// State for coordinating concurrent instance creation. -/// -/// Guarded by a single lock to provide atomic capacity checks. -struct CreationState { - /// Number of instances currently being created (not yet registered). - /// Added to [`SessionManager::instance_count`] for accurate capacity checks. - pending_count: usize, -} - /// Session manager that coordinates session and instance lifecycle. /// /// This provides a unified interface for: @@ -433,8 +423,11 @@ pub struct SessionManager { sessions: SessionMap, /// Cache of single-instance TAs by UUID. single_instance_cache: SingleInstanceCache, - /// Coordination state for concurrent instance creation. - creation_state: SpinMutex, + /// Number of instances currently being created (not yet registered). + /// Added to [`SessionManager::instance_count`] for the capacity check + /// in [`SessionManager::with_ta`] so two concurrent loads cannot both + /// pass the limit before either registers. + pending_count: SpinMutex, /// Cached TA flags by UUID, populated on first successful session registration. known_flags: SpinMutex>, /// Per-UUID serialization locks for single-instance TA handling. @@ -450,7 +443,7 @@ impl SessionManager { Self { sessions: SessionMap::new(), single_instance_cache: SingleInstanceCache::new(), - creation_state: SpinMutex::new(CreationState { pending_count: 0 }), + pending_count: SpinMutex::new(0), known_flags: SpinMutex::new(HashMap::new()), single_instance_locks: SpinMutex::new(HashMap::new()), active_sessions: SpinMutex::new(HashSet::new()), @@ -522,7 +515,7 @@ impl SessionManager { Some(lock) } - /// Acquire a [`SessionToken`] for an OpenSession request. + /// Acquire a `SessionToken` for an OpenSession request. /// /// For single-instance TAs (including first-ever load of an unknown /// UUID) this takes the per-UUID `SpinMutex` non-blockingly. For @@ -548,7 +541,9 @@ impl SessionManager { }) } - /// Acquire a [`SessionToken`] for an Invoke/Close on an existing session. + /// Acquire a token + validated entry for an Invoke/Close on an existing + /// session. Returns the entry that survived the post-marker re-read so + /// callers don't need to look it up again. /// /// Always reserves the per-session-id slot in `active_sessions`. For /// single-instance TAs additionally takes the per-UUID `SpinMutex` so @@ -565,21 +560,17 @@ impl SessionManager { /// diverge (the id was recycled and reused under a different TA between /// our first read and the marker insert), `Err(EThreadLimit)` is /// returned so the Linux driver retries — a fresh acquisition will see - /// the new entry from the start. This guards against the read/marker - /// TOCTOU without relying on - /// [`IdPool`](litebox::utils::id_pool::IdPool)'s hint+wrap to quarantine - /// recycled ids. - pub fn try_acquire_for_session( + /// the new entry from the start. + fn try_acquire_for_session( &self, session_id: u32, - ) -> Result, OpteeSmcReturnCode> { + ) -> Result<(SessionToken<'_>, SessionEntry), OpteeSmcReturnCode> { let entry = self .sessions .get_entry(session_id) .ok_or(OpteeSmcReturnCode::EBadCmd)?; let snapshot_uuid = entry.ta_uuid; let snapshot_single = entry.ta_flags.is_single_instance(); - drop(entry); if !self.active_sessions.lock().insert(session_id) { return Err(OpteeSmcReturnCode::EThreadLimit); @@ -603,14 +594,41 @@ impl SessionManager { return Err(OpteeSmcReturnCode::EThreadLimit); } - if snapshot_single { + // Only take the per-UUID lock for `Live` single-instance sessions. + // A `Dead` entry needs no sibling serialization — its instance is + // already gone, and contending with live siblings (or a freshly + // created instance for the same UUID) just to call + // `finalize_dead_session` would needlessly delay them. + if snapshot_single && matches!(entry_now.target, SessionTarget::Live(_)) { // On failure, dropping `token` releases the marker we just took. token.uuid_lock = Some( self.try_acquire_uuid_lock(snapshot_uuid) .ok_or(OpteeSmcReturnCode::EThreadLimit)?, ); } - Ok(token) + Ok((token, entry_now)) + } + + /// Drive an Invoke/Close to completion under the right serialization. + /// + /// Internally acquires the per-session-id marker (and, for single- + /// instance TAs, the per-UUID lock), passes the validated `SessionEntry` + /// to `f`, and releases the locks when `f` returns. `f` runs entirely + /// under the token: state mutations it performs on the session manager + /// (e.g. `unregister_session`, `mark_sessions_dead_for_instance`, + /// `remove_single_instance_if_same`) are serialized against other + /// cores' Invoke/Close on the same session and (for single-instance) + /// the same UUID. + /// + /// Returns `Err(EBadCmd)` if `session_id` is not registered, or + /// `Err(EThreadLimit)` on lock contention; the Linux OP-TEE driver + /// retries `EThreadLimit` transparently. + pub fn with_session(&self, session_id: u32, f: F) -> Result<(), OpteeSmcReturnCode> + where + F: FnOnce(SessionEntry) -> Result<(), OpteeSmcReturnCode>, + { + let (_token, entry) = self.try_acquire_for_session(session_id)?; + f(entry) } /// Register a new session. @@ -640,10 +658,12 @@ impl SessionManager { /// /// Callers tearing down on TA panic must have already called /// [`SessionMap::mark_sessions_dead_for_instance`] before invoking this, - /// so any handler that subsequently acquires its own [`SessionToken`] - /// for the UUID will observe `Dead` on its re-read of the session - /// entry. Callers on the last-session-close path may skip the mark - /// step — by that point there are no sibling sessions to fence out. + /// so any handler that subsequently enters + /// [`SessionManager::with_ta`] or + /// [`SessionManager::with_session`] for the UUID will observe `Dead` + /// on its re-read of the session entry. Callers on the last-session- + /// close path may skip the mark step — by that point there are no + /// sibling sessions to fence out. pub fn remove_single_instance_if_same( &self, uuid: &TeeUuid, @@ -686,14 +706,22 @@ impl SessionManager { /// - Calls `f(Some(existing))` if a cached single-instance TA is found /// for `uuid`. The lock is held throughout the call so the existing /// instance cannot be torn down or replaced beneath `f`. - /// - Reserves a creation slot (atomic capacity check including in-flight - /// creations) and calls `f(None)` to load and register a new instance. - /// The slot is released when `f` returns, regardless of outcome. + /// - Reserves a creation slot (atomic capacity check against + /// `instance_count() + pending_count`) and calls `f(None)` to load + /// and register a new instance. The slot is released when `f` + /// returns, regardless of outcome. /// /// The per-UUID lock is released when this function returns; `f` runs /// under it. For multi-instance TAs each session gets its own /// independent `TaInstance`, so no per-UUID exclusion is required. - pub fn with_creation_slot(&self, uuid: &TeeUuid, f: F) -> Result<(), OpteeSmcReturnCode> + /// + /// `pending_count` exists only for capacity accounting (so two + /// multi-instance loads can't both pass the limit check before either + /// registers). Duplicate-prevention for single-instance TAs is provided + /// by the per-UUID lock above — it serializes the cache check and any + /// new load for the same UUID, so two concurrent loads cannot both miss + /// the cache and create rival instances. + pub fn with_ta(&self, uuid: &TeeUuid, f: F) -> Result<(), OpteeSmcReturnCode> where F: FnOnce(Option>) -> Result<(), OpteeSmcReturnCode>, { @@ -707,20 +735,19 @@ impl SessionManager { } { - let mut state = self.creation_state.lock(); + let mut pending = self.pending_count.lock(); // Capacity check including in-flight creations. - let total = self.instance_count() + state.pending_count; - if total >= MAX_TA_INSTANCES { + if self.instance_count() + *pending >= MAX_TA_INSTANCES { return Err(OpteeSmcReturnCode::ENomem); } - state.pending_count += 1; + *pending += 1; } let result = f(None); { - let mut state = self.creation_state.lock(); - state.pending_count = state.pending_count.saturating_sub(1); + let mut pending = self.pending_count.lock(); + *pending = pending.saturating_sub(1); } result From 5ee3892f146ac36bbd7b08786d02c27cb1775f40 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Wed, 27 May 2026 05:24:21 +0000 Subject: [PATCH 08/28] improve safety --- litebox_runner_lvbs/src/lib.rs | 112 ++++++------ litebox_shim_optee/src/lib.rs | 4 +- litebox_shim_optee/src/session.rs | 273 ++++++++++++++++++++---------- 3 files changed, 239 insertions(+), 150 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index b986aacda6..7827ef0765 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -6,7 +6,6 @@ extern crate alloc; use alloc::boxed::Box; -use alloc::sync::Arc; use alloc::vec; use core::{ops::Neg, panic::PanicInfo}; use litebox::{ @@ -44,7 +43,7 @@ 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::{ - SessionIdGuard, SessionManager, SessionTarget, TaInstance, allocate_session_id, + InstanceRef, SessionIdGuard, SessionManager, TaInstance, TargetView, allocate_session_id, }; use litebox_shim_optee::{NormalWorldConstPtr, NormalWorldMutPtr, UserConstPtr}; use once_cell::race::OnceBox; @@ -545,12 +544,12 @@ fn handle_open_session( fn open_session_single_instance( msg_args: &mut OpteeMsgArgs, msg_args_phys_addr: u64, - instance_arc: Arc, + instance: InstanceRef<'_>, params: &[litebox_common_optee::UteeParamOwned], ta_uuid: litebox_common_optee::TeeUuid, ta_req_info: &litebox_shim_optee::msg_handler::TaRequestInfo, ) -> Result<(), OpteeSmcReturnCode> { - let task_pt_id = instance_arc.task_page_table_id; + let task_pt_id = instance.task_page_table_id(); // Allocate session ID BEFORE calling load_ta_context so TA gets correct ID. // Use SessionIdGuard to ensure the ID is recycled on any error path @@ -567,13 +566,13 @@ fn open_session_single_instance( runner_session_id ); - let ta_flags = instance_arc.loaded_program.ta_flags; + let ta_flags = instance.loaded_program().ta_flags; let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; // Load TA context with parameters for OpenSession - pass actual session_id - instance_arc - .loaded_program + instance + .loaded_program() .entrypoints .as_ref() .ok_or(OpteeSmcReturnCode::EBadCmd)? @@ -589,14 +588,14 @@ fn open_session_single_instance( let mut ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::reenter_thread_ref( - instance_arc.loaded_program.entrypoints.as_ref().unwrap(), + instance.loaded_program().entrypoints.as_ref().unwrap(), &mut ctx, ); } // Read TA output parameters from the stack buffer - let params_address = instance_arc - .loaded_program + let params_address = instance + .loaded_program() .params_address .ok_or(OpteeSmcReturnCode::EBadAddr)?; let ta_params = UserConstPtr::::from_usize(params_address) @@ -638,14 +637,12 @@ fn open_session_single_instance( // instance. Otherwise a racing handler that subsequently // acquires its own session token for the UUID could walk past // a still-Live session entry. - session_manager() - .sessions() - .mark_sessions_dead_for_instance(&instance_arc); - let _ = session_manager().remove_single_instance_if_same(&ta_uuid, &instance_arc); + session_manager().mark_sessions_dead_for_instance(instance); + let _ = session_manager().remove_single_instance_if_same(&ta_uuid, instance); // 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(&instance_arc.shim, task_pt_id); + teardown_ta_page_table(instance.shim(), task_pt_id); }; // TODO: Per OP-TEE OS semantics, if the TA has INSTANCE_KEEP_ALIVE but not @@ -677,17 +674,13 @@ fn open_session_single_instance( // so we forget the id (disarm the guard) to prevent a future OpenSession // from reusing it and colliding with the orphaned TA-side bookkeeping. if let Err(e) = write_result { - if !ta_flags.is_keep_alive() - && session_manager() - .sessions() - .count_sessions_for_instance(&instance_arc) - == 0 + if !ta_flags.is_keep_alive() && session_manager().count_sessions_for_instance(instance) == 0 { - let _ = session_manager().remove_single_instance_if_same(&ta_uuid, &instance_arc); + let _ = session_manager().remove_single_instance_if_same(&ta_uuid, instance); // 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(&instance_arc.shim, task_pt_id); + teardown_ta_page_table(instance.shim(), task_pt_id); }; } else { let _ = session_id_guard.disarm(); @@ -695,8 +688,8 @@ fn open_session_single_instance( return Err(e); } - // Success: register session and disarm the guard (ownership transfers to session map) - session_manager().register_session(runner_session_id, instance_arc.clone(), ta_uuid, ta_flags); + // Success: register a sibling session pointing at the existing instance. + session_manager().register_sibling_session(runner_session_id, instance, ta_uuid, ta_flags); session_id_guard.disarm(); debug_serial_println!( @@ -907,22 +900,19 @@ fn open_session_new_instance( unsafe { teardown_ta_page_table(&shim, task_pt_id) }; })?; - // Success: create TA instance - loaded_program is already boxed, no move happens - let instance = Arc::new(TaInstance { - shim, - loaded_program, - task_page_table_id: task_pt_id, - }); - - // Success: register session and disarm the guard (ownership transfers to session map) - session_manager().register_session(runner_session_id, instance.clone(), ta_uuid, ta_flags); + // Success: hand the instance to the session manager. `register_new_session` + // wraps it in an `Arc` owned solely by the manager (no clone returned), + // and — for single-instance TAs — also inserts it into the cache. The + // runner never holds an `Arc`, so it cannot violate the + // `unsafe impl Send/Sync for TaInstance` invariant. + session_manager().register_new_session( + runner_session_id, + TaInstance::new(shim, loaded_program, task_pt_id), + ta_uuid, + ta_flags, + ); session_id_guard.disarm(); - // Cache single-instance TAs only after the opening session owns the instance. - if ta_flags.is_single_instance() { - session_manager().cache_single_instance(ta_uuid, instance.clone()); - } - debug_serial_println!( "OpenSession complete: session_id={}, single_instance={}", runner_session_id, @@ -974,7 +964,7 @@ fn handle_invoke_command( let session_id = ta_req_info.session; session_manager().with_session(session_id, |session_entry| { - let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { + let TargetView::Live(instance) = session_entry.target else { return finalize_dead_session( session_id, msg_args, @@ -983,7 +973,7 @@ fn handle_invoke_command( "InvokeCommand", ); }; - let task_pt_id = instance_arc.task_page_table_id; + let task_pt_id = instance.task_page_table_id(); let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; @@ -995,7 +985,7 @@ fn handle_invoke_command( ); // Load TA context with parameters and cmd_id - pass actual session_id - let entrypoints_ref = instance_arc.loaded_program.entrypoints.as_ref().unwrap(); + let entrypoints_ref = instance.loaded_program().entrypoints.as_ref().unwrap(); entrypoints_ref .load_ta_context( params.as_slice(), @@ -1009,14 +999,14 @@ fn handle_invoke_command( let mut ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::reenter_thread_ref( - instance_arc.loaded_program.entrypoints.as_ref().unwrap(), + instance.loaded_program().entrypoints.as_ref().unwrap(), &mut ctx, ); } // params_address is constant - stack buffer is reused across invocations - let params_address = instance_arc - .loaded_program + let params_address = instance + .loaded_program() .params_address .ok_or(OpteeSmcReturnCode::EBadAddr)?; let ta_params = UserConstPtr::::from_usize(params_address) @@ -1054,10 +1044,8 @@ fn handle_invoke_command( // Mark siblings dead BEFORE evicting the cached single instance. // Otherwise a racing handler entering with_session/with_ta // for the UUID could walk past a still-Live session entry. - session_manager() - .sessions() - .mark_sessions_dead_for_instance(&instance_arc); - let _ = session_manager().remove_single_instance_if_same(&ta_uuid, &instance_arc); + session_manager().mark_sessions_dead_for_instance(instance); + let _ = session_manager().remove_single_instance_if_same(&ta_uuid, instance); } session_manager().unregister_session(session_id); @@ -1065,7 +1053,7 @@ fn handle_invoke_command( // 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(&instance_arc.shim, task_pt_id); + teardown_ta_page_table(instance.shim(), task_pt_id); }; debug_serial_println!( @@ -1100,7 +1088,8 @@ fn handle_close_session( debug_serial_println!("CloseSession: session_id={}", session_id); session_manager().with_session(session_id, |session_entry| { - let SessionTarget::Live(instance_arc) = session_entry.target.clone() else { + let ta_uuid = session_entry.ta_uuid; + let TargetView::Live(instance) = session_entry.target else { return finalize_dead_session( session_id, msg_args, @@ -1109,13 +1098,13 @@ fn handle_close_session( "CloseSession", ); }; - let task_pt_id = instance_arc.task_page_table_id; + let task_pt_id = instance.task_page_table_id(); let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; // Load TA context for CloseSession (no params, no cmd_id) - pass actual session_id - instance_arc - .loaded_program + instance + .loaded_program() .entrypoints .as_ref() .unwrap() @@ -1131,7 +1120,7 @@ fn handle_close_session( let mut ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::reenter_thread_ref( - instance_arc.loaded_program.entrypoints.as_ref().unwrap(), + instance.loaded_program().entrypoints.as_ref().unwrap(), &mut ctx, ); } @@ -1148,21 +1137,20 @@ fn handle_close_session( // Remove the session entry from the map. The session token drops // when the enclosing `with_session` closure returns. - let removed_entry = session_manager().unregister_session(session_id); + let removed_flags = session_manager().unregister_session(session_id); // Check if this was the last session using the TA instance by counting // remaining sessions that reference this instance. let remaining_sessions = session_manager() - .sessions() - .count_sessions_for_instance(&instance_arc); + .count_sessions_for_instance(instance); // If this was the last session using the TA instance, clean up (unless keep_alive is set) if remaining_sessions == 0 - && let Some(entry) = removed_entry + && let Some(flags) = removed_flags { // If this is a single-instance TA with keep_alive flag, don't remove it from memory. // Note: keep_alive is only meaningful for single-instance TAs. - if entry.ta_flags.is_single_instance() && entry.ta_flags.is_keep_alive() { + if flags.is_single_instance() && flags.is_keep_alive() { debug_serial_println!( "CloseSession complete: session_id={}, TA kept alive (INSTANCE_KEEP_ALIVE flag)", session_id @@ -1172,14 +1160,14 @@ fn handle_close_session( // If this was a single-instance TA, clear the cached instance. This is safe because // we confirm no sibling sessions remain. We don't need to mark anything `Dead` first. - if entry.ta_flags.is_single_instance() { + if flags.is_single_instance() { let _ = session_manager() - .remove_single_instance_if_same(&entry.ta_uuid, &instance_arc); + .remove_single_instance_if_same(&ta_uuid, instance); } // 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(&instance_arc.shim, task_pt_id); + teardown_ta_page_table(instance.shim(), task_pt_id); }; debug_serial_println!( diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index f4adceffba..fe7a0cd32e 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -39,8 +39,8 @@ pub mod ptr; // Re-export session management types for convenience pub use session::{ - MAX_TA_INSTANCES, SessionEntry, SessionManager, SessionMap, SessionTarget, SingleInstanceCache, - TaInstance, allocate_session_id, + InstanceRef, MAX_TA_INSTANCES, SessionManager, SessionView, TaInstance, TargetView, + allocate_session_id, }; const MAX_KERNEL_BUF_SIZE: usize = 0x80_000; diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 2f1adfe270..bb97d0e4af 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -49,7 +49,7 @@ //! requiring RPCs that would give untrusted code control over secure world execution. //! //! On panic teardown or last-session close, sibling sessions of a single-instance -//! TA are flipped to [`SessionTarget::Dead`] *before* the cached instance is +//! TA are flipped to `Dead` *before* the cached instance is //! evicted (see [`SessionManager::remove_single_instance_if_same`]). A racing //! handler that subsequently enters [`SessionManager::with_ta`] or //! [`SessionManager::with_session`] for the UUID will therefore observe `Dead` @@ -130,15 +130,51 @@ pub const MAX_TA_INSTANCES: usize = 16; /// the last session closes (or with `TA_FLAG_INSTANCE_KEEP_ALIVE`, until explicit destroy). /// /// Each instance has its own task page table that provides memory isolation from other TAs. +/// +/// Fields are private; external callers obtain instances through +/// [`SessionManager::with_ta`] / [`SessionManager::with_session`] closures +/// (which run under a `SessionToken` that serializes access) and reach the +/// internals via the accessor methods on this type. pub struct TaInstance { /// The shim must be kept alive to keep the loaded program's memory mappings valid. - pub shim: OpteeShim, + shim: OpteeShim, /// The loaded TA program state including entrypoints. /// Boxed to keep it at a fixed heap address - the Task inside must not be moved /// after initialization because it contains internal state that may not survive moves. - pub loaded_program: alloc::boxed::Box, + loaded_program: alloc::boxed::Box, /// The task page table ID associated with this TA instance. - pub task_page_table_id: usize, + task_page_table_id: usize, +} + +impl TaInstance { + /// Construct a new instance from its three constituent parts. + pub fn new( + shim: OpteeShim, + loaded_program: alloc::boxed::Box, + task_page_table_id: usize, + ) -> Self { + Self { + shim, + loaded_program, + task_page_table_id, + } + } + + /// Task page table ID associated with this instance. + pub fn task_page_table_id(&self) -> usize { + self.task_page_table_id + } + + /// Reference to the underlying shim (needed for releasing user mappings + /// during teardown). + pub fn shim(&self) -> &OpteeShim { + &self.shim + } + + /// Reference to the loaded program (entry points, parameter address, TA flags). + pub fn loaded_program(&self) -> &LoadedProgram { + &self.loaded_program + } } // SAFETY: `TaInstance`'s interior (`shim`, `loaded_program`) is not @@ -150,17 +186,77 @@ unsafe impl Send for TaInstance {} unsafe impl Sync for TaInstance {} /// The target associated with a normal-world session ID. +/// +/// This is the in-map representation, kept private to the module. External +/// callers see [`TargetView`] inside a [`SessionView`] delivered by the +/// session-token-bound closure APIs ([`SessionManager::with_ta`], +/// [`SessionManager::with_session`]). #[derive(Clone)] -pub enum SessionTarget { +pub(crate) enum SessionTarget { /// The session still targets a live TA instance. Live(Arc), /// The TA died, but normal world may still issue Invoke/Close for this ID. Dead, } -/// Per-session entry in the session map. +/// Borrowed view of a [`TaInstance`], valid only inside the +/// [`SessionManager::with_ta`] / [`SessionManager::with_session`] closure +/// that received it. +/// +/// The lifetime ties the view to the `SessionToken` held internally by +/// the closure; the closure cannot smuggle this borrow out (HRTB on the +/// closure forces it to be valid for any lifetime the manager picks), +/// cannot clone it into an owned `Arc`, and cannot extract +/// the wrapped `Arc`. All mutation paths that need instance identity +/// (sibling marking, count, cache eviction, sibling-session registration) +/// take an `InstanceRef<'_>`. `Copy` is fine — it just produces another +/// borrow with the same lifetime, not an escape. +#[derive(Clone, Copy)] +pub struct InstanceRef<'a> { + arc: &'a Arc, +} + +impl<'a> InstanceRef<'a> { + fn new(arc: &'a Arc) -> Self { + Self { arc } + } + + /// Task page table ID associated with this instance. + pub fn task_page_table_id(&self) -> usize { + self.arc.task_page_table_id() + } + + /// Reference to the underlying shim. + pub fn shim(&self) -> &OpteeShim { + self.arc.shim() + } + + /// Reference to the loaded program (entry points, parameter address, TA flags). + pub fn loaded_program(&self) -> &LoadedProgram { + self.arc.loaded_program() + } +} + +/// Closure-bound view of a session's target. Mirrors the internal +/// `SessionTarget` enum but exposes [`InstanceRef`] instead of +/// `Arc`. +pub enum TargetView<'a> { + Live(InstanceRef<'a>), + Dead, +} + +/// Closure-bound view of a session entry. Delivered by +/// [`SessionManager::with_session`]. +pub struct SessionView<'a> { + pub ta_uuid: TeeUuid, + pub ta_flags: TaFlags, + pub target: TargetView<'a>, +} + +/// Per-session entry in the session map. Module-private; the closure-bound +/// public view is [`SessionView`]. #[derive(Clone)] -pub struct SessionEntry { +pub(crate) struct SessionEntry { /// The TA target (may be shared with other sessions for single-instance TAs). pub target: SessionTarget, /// The TA UUID (needed for cleanup of single-instance TAs). @@ -172,33 +268,25 @@ pub struct SessionEntry { /// Session map for tracking active sessions. /// /// Maps runner-allocated session IDs to session entries. -pub struct SessionMap { +pub(crate) struct SessionMap { inner: SpinMutex>, } impl SessionMap { /// Create a new empty session map. - pub fn new() -> Self { + pub(crate) fn new() -> Self { Self { inner: SpinMutex::new(HashMap::new()), } } - /// Get a session's TA instance by session ID. - pub fn get(&self, session_id: u32) -> Option> { - match self.inner.lock().get(&session_id).map(|e| &e.target) { - Some(SessionTarget::Live(instance)) => Some(instance.clone()), - Some(SessionTarget::Dead) | None => None, - } - } - /// Get full session entry by session ID. - pub fn get_entry(&self, session_id: u32) -> Option { + pub(crate) fn get_entry(&self, session_id: u32) -> Option { self.inner.lock().get(&session_id).cloned() } /// Insert a session into the map. - pub fn insert( + pub(crate) fn insert( &self, session_id: u32, instance: Arc, @@ -216,22 +304,12 @@ impl SessionMap { } /// Remove a session from the map. - pub fn remove(&self, session_id: u32) -> Option { + pub(crate) fn remove(&self, session_id: u32) -> Option { self.inner.lock().remove(&session_id) } - /// Get the number of active sessions. - pub fn len(&self) -> usize { - self.inner.lock().len() - } - - /// Check if the session map is empty. - pub fn is_empty(&self) -> bool { - self.inner.lock().is_empty() - } - /// Count sessions for a specific TA instance (by Arc pointer equality). - pub fn count_sessions_for_instance(&self, instance: &Arc) -> usize { + pub(crate) fn count_sessions_for_instance(&self, instance: &Arc) -> usize { self.inner .lock() .values() @@ -243,7 +321,7 @@ impl SessionMap { } /// Mark all sessions pointing at `instance` as dead. - pub fn mark_sessions_dead_for_instance(&self, instance: &Arc) { + pub(crate) fn mark_sessions_dead_for_instance(&self, instance: &Arc) { for entry in self.inner.lock().values_mut() { if matches!(&entry.target, SessionTarget::Live(current) if Arc::ptr_eq(current, instance)) { @@ -263,25 +341,25 @@ impl Default for SessionMap { /// /// Single-instance TAs (with `TA_FLAG_SINGLE_INSTANCE`) share a single TA instance /// across all sessions. This cache stores instances by UUID for fast reuse lookup. -pub struct SingleInstanceCache { +pub(crate) struct SingleInstanceCache { inner: SpinMutex>>, } impl SingleInstanceCache { /// Create a new empty cache. - pub fn new() -> Self { + pub(crate) fn new() -> Self { Self { inner: SpinMutex::new(HashMap::new()), } } /// Get a cached single-instance TA by UUID. - pub fn get(&self, uuid: &TeeUuid) -> Option> { + pub(crate) fn get(&self, uuid: &TeeUuid) -> Option> { self.inner.lock().get(uuid).cloned() } /// Cache a single-instance TA by UUID. - pub fn insert(&self, uuid: TeeUuid, instance: Arc) { + pub(crate) fn insert(&self, uuid: TeeUuid, instance: Arc) { self.inner.lock().insert(uuid, instance); } @@ -298,14 +376,9 @@ impl SingleInstanceCache { } /// Get the number of cached single-instance TAs. - pub fn len(&self) -> usize { + pub(crate) fn len(&self) -> usize { self.inner.lock().len() } - - /// Check if the cache is empty. - pub fn is_empty(&self) -> bool { - self.inner.lock().is_empty() - } } impl Default for SingleInstanceCache { @@ -325,14 +398,15 @@ pub fn allocate_session_id() -> Option { /// Recycle a session ID for potential future reuse. /// /// Delegates to `SessionIdPool::recycle`. -pub fn recycle_session_id(session_id: u32) { +fn recycle_session_id(session_id: u32) { SessionIdPool::recycle(session_id); } /// RAII guard that recycles a session ID on drop unless disarmed. /// /// Session IDs are allocated before the TA is invoked and only registered on -/// success via [`SessionManager::register_session`]. This guard ensures it is +/// success via [`SessionManager::register_new_session`] or +/// [`SessionManager::register_sibling_session`]. This guard ensures it is /// recycled on all error paths before this registration. pub struct SessionIdGuard { session_id: Option, @@ -450,29 +524,19 @@ impl SessionManager { } } - /// Get the session map. - pub fn sessions(&self) -> &SessionMap { - &self.sessions - } - - /// Get the single-instance cache. - pub fn single_instance_cache(&self) -> &SingleInstanceCache { - &self.single_instance_cache - } - - /// Cache a single-instance TA. - pub fn cache_single_instance(&self, uuid: TeeUuid, instance: Arc) { - self.single_instance_cache.insert(uuid, instance); - } - - /// Get a session by ID. - pub fn get_session(&self, session_id: u32) -> Option> { - self.sessions.get(session_id) + /// Mark every session currently pointing at `instance` as `Dead`. + /// + /// Must be called before evicting `instance` from the single-instance + /// cache (see [`SessionManager::remove_single_instance_if_same`]) so + /// a racing handler re-reads `Dead` on its sibling session entry. + pub fn mark_sessions_dead_for_instance(&self, instance: InstanceRef<'_>) { + self.sessions.mark_sessions_dead_for_instance(instance.arc); } - /// Get full session entry by ID. - pub fn get_session_entry(&self, session_id: u32) -> Option { - self.sessions.get_entry(session_id) + /// Count sessions currently pointing at `instance`. Used by the + /// last-close path to detect whether teardown is appropriate. + pub fn count_sessions_for_instance(&self, instance: InstanceRef<'_>) -> usize { + self.sessions.count_sessions_for_instance(instance.arc) } /// Look up previously observed TA flags for a UUID. @@ -625,51 +689,88 @@ impl SessionManager { /// retries `EThreadLimit` transparently. pub fn with_session(&self, session_id: u32, f: F) -> Result<(), OpteeSmcReturnCode> where - F: FnOnce(SessionEntry) -> Result<(), OpteeSmcReturnCode>, + F: for<'a> FnOnce(SessionView<'a>) -> Result<(), OpteeSmcReturnCode>, { let (_token, entry) = self.try_acquire_for_session(session_id)?; - f(entry) + let view = SessionView { + ta_uuid: entry.ta_uuid, + ta_flags: entry.ta_flags, + target: match &entry.target { + SessionTarget::Live(arc) => TargetView::Live(InstanceRef::new(arc)), + SessionTarget::Dead => TargetView::Dead, + }, + }; + f(view) } - /// Register a new session. - pub fn register_session( + /// Register a session for a freshly-loaded TA instance. + /// + /// Takes `TaInstance` by value, wraps it in an `Arc` owned solely by the + /// manager (no clone returned to the caller), and — for single-instance + /// TAs — caches that `Arc` under `ta_uuid`. The caller therefore cannot + /// retain an `Arc` past this call, which is what makes the + /// `unsafe impl Send/Sync for TaInstance` invariant enforceable. + pub fn register_new_session( &self, session_id: u32, - instance: Arc, + instance: TaInstance, + ta_uuid: TeeUuid, + ta_flags: TaFlags, + ) { + let arc = Arc::new(instance); + self.known_flags.lock().entry(ta_uuid).or_insert(ta_flags); + self.sessions + .insert(session_id, arc.clone(), ta_uuid, ta_flags); + if ta_flags.is_single_instance() { + self.single_instance_cache.insert(ta_uuid, arc); + } + } + + /// Register a session that re-uses an existing single-instance TA. + /// + /// `instance` is the borrow handed to the [`SessionManager::with_ta`] + /// closure on the cache-hit branch. + pub fn register_sibling_session( + &self, + session_id: u32, + instance: InstanceRef<'_>, ta_uuid: TeeUuid, ta_flags: TaFlags, ) { self.known_flags.lock().entry(ta_uuid).or_insert(ta_flags); self.sessions - .insert(session_id, instance, ta_uuid, ta_flags); + .insert(session_id, instance.arc.clone(), ta_uuid, ta_flags); } - /// Unregister a session, recycle its session ID, and return the entry. - pub fn unregister_session(&self, session_id: u32) -> Option { + /// Unregister a session and recycle its session ID. Returns whether + /// the session was registered and what flags it had (the latter for + /// callers that need to dispatch on `is_single_instance` / + /// `is_keep_alive` after removal). + pub fn unregister_session(&self, session_id: u32) -> Option { let entry = self.sessions.remove(session_id); if entry.is_some() { recycle_session_id(session_id); } - entry + entry.map(|e| e.ta_flags) } /// Remove a single-instance TA from the cache only if the currently - /// cached `Arc` is the same as `expected`. + /// cached `Arc` is the same as `instance`. /// /// Callers tearing down on TA panic must have already called - /// [`SessionMap::mark_sessions_dead_for_instance`] before invoking this, - /// so any handler that subsequently enters - /// [`SessionManager::with_ta`] or - /// [`SessionManager::with_session`] for the UUID will observe `Dead` - /// on its re-read of the session entry. Callers on the last-session- - /// close path may skip the mark step — by that point there are no - /// sibling sessions to fence out. + /// [`SessionManager::mark_sessions_dead_for_instance`] before invoking + /// this, so any handler that subsequently enters + /// [`SessionManager::with_ta`] or [`SessionManager::with_session`] for + /// the UUID will observe `Dead` on its re-read of the session entry. + /// Callers on the last-session-close path may skip the mark step — by + /// that point there are no sibling sessions to fence out. pub fn remove_single_instance_if_same( &self, uuid: &TeeUuid, - expected: &Arc, + instance: InstanceRef<'_>, ) -> bool { - self.single_instance_cache.remove_if_same(uuid, expected) + self.single_instance_cache + .remove_if_same(uuid, instance.arc) } /// Get the total count of unique TA instances (for limit checking). @@ -723,7 +824,7 @@ impl SessionManager { /// the cache and create rival instances. pub fn with_ta(&self, uuid: &TeeUuid, f: F) -> Result<(), OpteeSmcReturnCode> where - F: FnOnce(Option>) -> Result<(), OpteeSmcReturnCode>, + F: for<'a> FnOnce(Option>) -> Result<(), OpteeSmcReturnCode>, { let token = self.try_acquire_for_open(*uuid)?; let is_single_instance = token.uuid_lock.is_some(); @@ -731,7 +832,7 @@ impl SessionManager { // For single-instance TAs the per-UUID lock above keeps our UUID's // cache entry stable. For multi-instance we don't consult the cache. if is_single_instance && let Some(existing) = self.single_instance_cache.get(uuid) { - return f(Some(existing)); + return f(Some(InstanceRef::new(&existing))); } { From a2ef2d05ec23b4c0e4f253ad3b24c5fc9ce30c1f Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Wed, 27 May 2026 14:40:28 +0000 Subject: [PATCH 09/28] refine interface and doc --- litebox_runner_lvbs/src/lib.rs | 34 ++-- litebox_shim_optee/src/lib.rs | 5 +- litebox_shim_optee/src/session.rs | 255 +++++++++++++----------------- 3 files changed, 130 insertions(+), 164 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 7827ef0765..e02530f58f 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -43,7 +43,7 @@ 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::{ - InstanceRef, SessionIdGuard, SessionManager, TaInstance, TargetView, allocate_session_id, + SessionIdGuard, SessionManager, TaInstance, allocate_session_id, }; use litebox_shim_optee::{NormalWorldConstPtr, NormalWorldMutPtr, UserConstPtr}; use once_cell::race::OnceBox; @@ -544,7 +544,7 @@ fn handle_open_session( fn open_session_single_instance( msg_args: &mut OpteeMsgArgs, msg_args_phys_addr: u64, - instance: InstanceRef<'_>, + instance: &TaInstance, params: &[litebox_common_optee::UteeParamOwned], ta_uuid: litebox_common_optee::TeeUuid, ta_req_info: &litebox_shim_optee::msg_handler::TaRequestInfo, @@ -689,7 +689,7 @@ fn open_session_single_instance( } // Success: register a sibling session pointing at the existing instance. - session_manager().register_sibling_session(runner_session_id, instance, ta_uuid, ta_flags); + session_manager().register_sibling_session(runner_session_id, instance, ta_uuid, ta_flags)?; session_id_guard.disarm(); debug_serial_println!( @@ -900,14 +900,16 @@ fn open_session_new_instance( unsafe { teardown_ta_page_table(&shim, task_pt_id) }; })?; - // Success: hand the instance to the session manager. `register_new_session` - // wraps it in an `Arc` owned solely by the manager (no clone returned), - // and — for single-instance TAs — also inserts it into the cache. The - // runner never holds an `Arc`, so it cannot violate the - // `unsafe impl Send/Sync for TaInstance` invariant. + // Success: hand the three parts to the session manager. It wraps them + // in an `Arc` owned solely by itself (no clone returned), + // and — for single-instance TAs — also caches that `Arc`. The runner + // never holds a `TaInstance` or `Arc`, which is what makes + // the internal `unsafe impl Send/Sync for TaInstance` enforceable. session_manager().register_new_session( runner_session_id, - TaInstance::new(shim, loaded_program, task_pt_id), + shim, + loaded_program, + task_pt_id, ta_uuid, ta_flags, ); @@ -963,8 +965,8 @@ fn handle_invoke_command( let params = &ta_req_info.params; let session_id = ta_req_info.session; - session_manager().with_session(session_id, |session_entry| { - let TargetView::Live(instance) = session_entry.target else { + session_manager().with_session(session_id, |session| { + let Some(instance) = session.live() else { return finalize_dead_session( session_id, msg_args, @@ -1037,8 +1039,8 @@ fn handle_invoke_command( session_id ); - let ta_uuid = session_entry.ta_uuid; - let ta_flags = session_entry.ta_flags; + let ta_uuid = session.ta_uuid; + let ta_flags = session.ta_flags; if ta_flags.is_single_instance() { // Mark siblings dead BEFORE evicting the cached single instance. @@ -1087,9 +1089,9 @@ fn handle_close_session( debug_serial_println!("CloseSession: session_id={}", session_id); - session_manager().with_session(session_id, |session_entry| { - let ta_uuid = session_entry.ta_uuid; - let TargetView::Live(instance) = session_entry.target else { + session_manager().with_session(session_id, |session| { + let ta_uuid = session.ta_uuid; + let Some(instance) = session.live() else { return finalize_dead_session( session_id, msg_args, diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index fe7a0cd32e..45acf64dcd 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -38,10 +38,7 @@ pub mod msg_handler; pub mod ptr; // Re-export session management types for convenience -pub use session::{ - InstanceRef, MAX_TA_INSTANCES, SessionManager, SessionView, TaInstance, TargetView, - allocate_session_id, -}; +pub use session::{MAX_TA_INSTANCES, SessionManager, TaInstance, allocate_session_id}; const MAX_KERNEL_BUF_SIZE: usize = 0x80_000; diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index bb97d0e4af..5fa87fa978 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -10,10 +10,10 @@ //! //! ## Concurrency Model //! -//! TA execution is serialized externally; [`TaInstance`] is shared as a plain -//! `Arc` without an inner mutex. The exclusivity invariant lives in -//! [`SessionManager`] and is acquired through an internal RAII `SessionToken` -//! that bundles whichever locks the current operation requires: +//! TA execution is serialized externally; [`TaInstance`] is shared without an +//! inner mutex. The exclusivity invariant lives in [`SessionManager`] and is +//! acquired through an internal RAII `SessionToken` that bundles whichever +//! locks the current operation requires: //! //! - **Single-instance TAs** (with `TA_FLAG_SINGLE_INSTANCE | TA_FLAG_MULTI_SESSION`) //! share one [`TaInstance`] across all sessions. The token internally holds a @@ -123,18 +123,19 @@ use spin::mutex::SpinMutex; /// Maximum number of concurrent TA instances to avoid out of memory situations. pub const MAX_TA_INSTANCES: usize = 16; -/// A loaded TA instance that can be shared across multiple sessions. +/// A loaded TA instance. /// -/// For single-instance TAs (with `TA_FLAG_SINGLE_INSTANCE`), one TA instance -/// is shared across all sessions. The TA is loaded once and stays in memory until -/// the last session closes (or with `TA_FLAG_INSTANCE_KEEP_ALIVE`, until explicit destroy). +/// Fields are private; external callers never construct one (the three +/// constituent parts are passed to [`SessionManager::register_new_session`], +/// which builds the instance internally). Closures running under +/// [`SessionManager::with_ta`] / [`SessionManager::with_session`] observe +/// the instance through `&TaInstance`, with the borrow lifetime pinned to +/// the session token via HRTB on the closure type. /// -/// Each instance has its own task page table that provides memory isolation from other TAs. -/// -/// Fields are private; external callers obtain instances through -/// [`SessionManager::with_ta`] / [`SessionManager::with_session`] closures -/// (which run under a `SessionToken` that serializes access) and reach the -/// internals via the accessor methods on this type. +/// For single-instance TAs one instance is shared across all sessions; the +/// TA stays in memory until the last session closes (or, with +/// `TA_FLAG_INSTANCE_KEEP_ALIVE`, until explicit destroy). Each instance +/// has its own task page table that provides memory isolation from other TAs. pub struct TaInstance { /// The shim must be kept alive to keep the loaded program's memory mappings valid. shim: OpteeShim, @@ -143,35 +144,22 @@ pub struct TaInstance { /// after initialization because it contains internal state that may not survive moves. loaded_program: alloc::boxed::Box, /// The task page table ID associated with this TA instance. + /// + /// Also serves as the instance's identity for sibling-tracking + /// operations: page table ids are minted by `create_task_page_table()` + /// and not reused until the owning instance is fully torn down. task_page_table_id: usize, } impl TaInstance { - /// Construct a new instance from its three constituent parts. - pub fn new( - shim: OpteeShim, - loaded_program: alloc::boxed::Box, - task_page_table_id: usize, - ) -> Self { - Self { - shim, - loaded_program, - task_page_table_id, - } - } - - /// Task page table ID associated with this instance. pub fn task_page_table_id(&self) -> usize { self.task_page_table_id } - /// Reference to the underlying shim (needed for releasing user mappings - /// during teardown). pub fn shim(&self) -> &OpteeShim { &self.shim } - /// Reference to the loaded program (entry points, parameter address, TA flags). pub fn loaded_program(&self) -> &LoadedProgram { &self.loaded_program } @@ -188,73 +176,39 @@ unsafe impl Sync for TaInstance {} /// The target associated with a normal-world session ID. /// /// This is the in-map representation, kept private to the module. External -/// callers see [`TargetView`] inside a [`SessionView`] delivered by the -/// session-token-bound closure APIs ([`SessionManager::with_ta`], +/// callers see liveness via [`Session::live`] on the [`Session`] delivered +/// by the session-token-bound closure APIs ([`SessionManager::with_ta`], /// [`SessionManager::with_session`]). #[derive(Clone)] pub(crate) enum SessionTarget { - /// The session still targets a live TA instance. Live(Arc), /// The TA died, but normal world may still issue Invoke/Close for this ID. Dead, } -/// Borrowed view of a [`TaInstance`], valid only inside the -/// [`SessionManager::with_ta`] / [`SessionManager::with_session`] closure -/// that received it. +/// Closure-bound snapshot of a session, delivered by +/// [`SessionManager::with_session`]. /// -/// The lifetime ties the view to the `SessionToken` held internally by -/// the closure; the closure cannot smuggle this borrow out (HRTB on the -/// closure forces it to be valid for any lifetime the manager picks), -/// cannot clone it into an owned `Arc`, and cannot extract -/// the wrapped `Arc`. All mutation paths that need instance identity -/// (sibling marking, count, cache eviction, sibling-session registration) -/// take an `InstanceRef<'_>`. `Copy` is fine — it just produces another -/// borrow with the same lifetime, not an escape. -#[derive(Clone, Copy)] -pub struct InstanceRef<'a> { - arc: &'a Arc, +/// Holds the session's `ta_uuid` / `ta_flags` and, for live sessions, a +/// borrow of the [`TaInstance`]. The borrow's lifetime is pinned to the +/// `SessionToken` held internally by `with_session` via HRTB on the +/// closure type, so the closure cannot smuggle it out. +pub struct Session<'a> { + pub ta_uuid: TeeUuid, + pub ta_flags: TaFlags, + instance: Option<&'a TaInstance>, } -impl<'a> InstanceRef<'a> { - fn new(arc: &'a Arc) -> Self { - Self { arc } - } - - /// Task page table ID associated with this instance. - pub fn task_page_table_id(&self) -> usize { - self.arc.task_page_table_id() - } - - /// Reference to the underlying shim. - pub fn shim(&self) -> &OpteeShim { - self.arc.shim() - } - - /// Reference to the loaded program (entry points, parameter address, TA flags). - pub fn loaded_program(&self) -> &LoadedProgram { - self.arc.loaded_program() +impl<'a> Session<'a> { + /// Returns `Some(instance)` if the session's target is live, or `None` + /// if the TA has died and the caller should run dead-session cleanup. + pub fn live(&self) -> Option<&'a TaInstance> { + self.instance } } -/// Closure-bound view of a session's target. Mirrors the internal -/// `SessionTarget` enum but exposes [`InstanceRef`] instead of -/// `Arc`. -pub enum TargetView<'a> { - Live(InstanceRef<'a>), - Dead, -} - -/// Closure-bound view of a session entry. Delivered by -/// [`SessionManager::with_session`]. -pub struct SessionView<'a> { - pub ta_uuid: TeeUuid, - pub ta_flags: TaFlags, - pub target: TargetView<'a>, -} - /// Per-session entry in the session map. Module-private; the closure-bound -/// public view is [`SessionView`]. +/// public view is [`Session`]. #[derive(Clone)] pub(crate) struct SessionEntry { /// The TA target (may be shared with other sessions for single-instance TAs). @@ -273,19 +227,16 @@ pub(crate) struct SessionMap { } impl SessionMap { - /// Create a new empty session map. pub(crate) fn new() -> Self { Self { inner: SpinMutex::new(HashMap::new()), } } - /// Get full session entry by session ID. pub(crate) fn get_entry(&self, session_id: u32) -> Option { self.inner.lock().get(&session_id).cloned() } - /// Insert a session into the map. pub(crate) fn insert( &self, session_id: u32, @@ -303,27 +254,27 @@ impl SessionMap { ); } - /// Remove a session from the map. pub(crate) fn remove(&self, session_id: u32) -> Option { self.inner.lock().remove(&session_id) } - /// Count sessions for a specific TA instance (by Arc pointer equality). - pub(crate) fn count_sessions_for_instance(&self, instance: &Arc) -> usize { + /// Count live sessions whose instance has the given page table id. + pub(crate) fn count_sessions_for_pt(&self, task_page_table_id: usize) -> usize { self.inner .lock() .values() .filter(|e| match &e.target { - SessionTarget::Live(current) => Arc::ptr_eq(current, instance), + SessionTarget::Live(arc) => arc.task_page_table_id == task_page_table_id, SessionTarget::Dead => false, }) .count() } - /// Mark all sessions pointing at `instance` as dead. - pub(crate) fn mark_sessions_dead_for_instance(&self, instance: &Arc) { + /// Mark all live sessions whose instance has the given page table id + /// as `Dead`. + pub(crate) fn mark_sessions_dead_for_pt(&self, task_page_table_id: usize) { for entry in self.inner.lock().values_mut() { - if matches!(&entry.target, SessionTarget::Live(current) if Arc::ptr_eq(current, instance)) + if matches!(&entry.target, SessionTarget::Live(arc) if arc.task_page_table_id == task_page_table_id) { entry.target = SessionTarget::Dead; } @@ -346,28 +297,27 @@ pub(crate) struct SingleInstanceCache { } impl SingleInstanceCache { - /// Create a new empty cache. pub(crate) fn new() -> Self { Self { inner: SpinMutex::new(HashMap::new()), } } - /// Get a cached single-instance TA by UUID. pub(crate) fn get(&self, uuid: &TeeUuid) -> Option> { self.inner.lock().get(uuid).cloned() } - /// Cache a single-instance TA by UUID. pub(crate) fn insert(&self, uuid: TeeUuid, instance: Arc) { self.inner.lock().insert(uuid, instance); } - /// Remove a cached single-instance TA only if it is the expected instance. - fn remove_if_same(&self, uuid: &TeeUuid, expected: &Arc) -> bool { + /// Evict only if the cached instance matches `task_page_table_id`. + /// Distinguishes the live instance from a freshly-created one with the + /// same UUID when the caller wants to remove a specific one. + fn remove_if_pt(&self, uuid: &TeeUuid, task_page_table_id: usize) -> bool { let mut guard = self.inner.lock(); match guard.get(uuid) { - Some(current) if Arc::ptr_eq(current, expected) => { + Some(current) if current.task_page_table_id == task_page_table_id => { guard.remove(uuid); true } @@ -375,7 +325,6 @@ impl SingleInstanceCache { } } - /// Get the number of cached single-instance TAs. pub(crate) fn len(&self) -> usize { self.inner.lock().len() } @@ -387,17 +336,11 @@ impl Default for SingleInstanceCache { } } -/// Allocate a new unique session ID. -/// -/// Delegates to `SessionIdPool::allocate` for unified session ID management. /// Returns `None` if all session IDs are exhausted. pub fn allocate_session_id() -> Option { SessionIdPool::allocate() } -/// Recycle a session ID for potential future reuse. -/// -/// Delegates to `SessionIdPool::recycle`. fn recycle_session_id(session_id: u32) { SessionIdPool::recycle(session_id); } @@ -413,14 +356,13 @@ pub struct SessionIdGuard { } impl SessionIdGuard { - /// Create a new guard that will recycle `session_id` on drop. pub fn new(session_id: u32) -> Self { Self { session_id: Some(session_id), } } - /// Return the guarded session ID, or `None` if already disarmed. + /// Returns `None` if already disarmed. pub fn id(&self) -> Option { self.session_id } @@ -488,10 +430,11 @@ impl Drop for SessionToken<'_> { /// Session manager that coordinates session and instance lifecycle. /// -/// This provides a unified interface for: -/// - Opening sessions (with single-instance TA reuse) -/// - Looking up sessions -/// - Closing sessions (with proper cleanup) +/// The public entry points are the closure-bound [`SessionManager::with_ta`] +/// (OpenSession) and [`SessionManager::with_session`] (Invoke/Close), which +/// run the caller's closure under an internal `SessionToken`. State +/// mutations the closure performs on the manager (registration, +/// sibling-marking, cache eviction) are serialized by that token. pub struct SessionManager { /// Active sessions mapped by session ID. sessions: SessionMap, @@ -512,7 +455,6 @@ pub struct SessionManager { } impl SessionManager { - /// Create a new session manager. pub fn new() -> Self { Self { sessions: SessionMap::new(), @@ -529,14 +471,17 @@ impl SessionManager { /// Must be called before evicting `instance` from the single-instance /// cache (see [`SessionManager::remove_single_instance_if_same`]) so /// a racing handler re-reads `Dead` on its sibling session entry. - pub fn mark_sessions_dead_for_instance(&self, instance: InstanceRef<'_>) { - self.sessions.mark_sessions_dead_for_instance(instance.arc); + pub fn mark_sessions_dead_for_instance(&self, instance: &TaInstance) { + self.sessions + .mark_sessions_dead_for_pt(instance.task_page_table_id); } - /// Count sessions currently pointing at `instance`. Used by the - /// last-close path to detect whether teardown is appropriate. - pub fn count_sessions_for_instance(&self, instance: InstanceRef<'_>) -> usize { - self.sessions.count_sessions_for_instance(instance.arc) + /// Count live sessions currently pointing at `instance` (`Dead` entries + /// are skipped). Used by the last-close path to detect whether teardown + /// is appropriate. + pub fn count_sessions_for_instance(&self, instance: &TaInstance) -> usize { + self.sessions + .count_sessions_for_pt(instance.task_page_table_id) } /// Look up previously observed TA flags for a UUID. @@ -676,7 +621,7 @@ impl SessionManager { /// Drive an Invoke/Close to completion under the right serialization. /// /// Internally acquires the per-session-id marker (and, for single- - /// instance TAs, the per-UUID lock), passes the validated `SessionEntry` + /// instance TAs, the per-UUID lock), passes a validated [`Session`] /// to `f`, and releases the locks when `f` returns. `f` runs entirely /// under the token: state mutations it performs on the session manager /// (e.g. `unregister_session`, `mark_sessions_dead_for_instance`, @@ -689,35 +634,43 @@ impl SessionManager { /// retries `EThreadLimit` transparently. pub fn with_session(&self, session_id: u32, f: F) -> Result<(), OpteeSmcReturnCode> where - F: for<'a> FnOnce(SessionView<'a>) -> Result<(), OpteeSmcReturnCode>, + F: for<'a> FnOnce(Session<'a>) -> Result<(), OpteeSmcReturnCode>, { let (_token, entry) = self.try_acquire_for_session(session_id)?; - let view = SessionView { + let session = Session { ta_uuid: entry.ta_uuid, ta_flags: entry.ta_flags, - target: match &entry.target { - SessionTarget::Live(arc) => TargetView::Live(InstanceRef::new(arc)), - SessionTarget::Dead => TargetView::Dead, + instance: match &entry.target { + SessionTarget::Live(arc) => Some(&**arc), + SessionTarget::Dead => None, }, }; - f(view) + f(session) } - /// Register a session for a freshly-loaded TA instance. + /// Register a session for a freshly-loaded TA. The three parts (`shim`, + /// `loaded_program`, `task_page_table_id`) are taken by value and stored + /// inside the manager (no handle returned to the caller); for + /// single-instance TAs the instance is also cached under `ta_uuid` for + /// later reuse. /// - /// Takes `TaInstance` by value, wraps it in an `Arc` owned solely by the - /// manager (no clone returned to the caller), and — for single-instance - /// TAs — caches that `Arc` under `ta_uuid`. The caller therefore cannot - /// retain an `Arc` past this call, which is what makes the - /// `unsafe impl Send/Sync for TaInstance` invariant enforceable. + /// The caller never retains a `TaInstance`, which is what makes the + /// internal `unsafe impl Send/Sync for TaInstance` invariant enforceable + /// against the public API. pub fn register_new_session( &self, session_id: u32, - instance: TaInstance, + shim: OpteeShim, + loaded_program: alloc::boxed::Box, + task_page_table_id: usize, ta_uuid: TeeUuid, ta_flags: TaFlags, ) { - let arc = Arc::new(instance); + let arc = Arc::new(TaInstance { + shim, + loaded_program, + task_page_table_id, + }); self.known_flags.lock().entry(ta_uuid).or_insert(ta_flags); self.sessions .insert(session_id, arc.clone(), ta_uuid, ta_flags); @@ -729,17 +682,30 @@ impl SessionManager { /// Register a session that re-uses an existing single-instance TA. /// /// `instance` is the borrow handed to the [`SessionManager::with_ta`] - /// closure on the cache-hit branch. + /// closure on the cache-hit branch. The cached instance for `ta_uuid` + /// is matched against `task_page_table_id`. Under correct usage the + /// caller holds the per-UUID lock (via `with_ta`'s token) for the + /// duration, so the cache entry is stable and the lookup succeeds. + /// + /// Returns `Err(EBadCmd)` if no matching cached instance is found. + /// This is an internal-consistency check rather than a recoverable + /// runtime condition; in kernel code we surface it as an error rather + /// than panicking. pub fn register_sibling_session( &self, session_id: u32, - instance: InstanceRef<'_>, + instance: &TaInstance, ta_uuid: TeeUuid, ta_flags: TaFlags, - ) { + ) -> Result<(), OpteeSmcReturnCode> { + let arc = self + .single_instance_cache + .get(&ta_uuid) + .filter(|cached| cached.task_page_table_id == instance.task_page_table_id) + .ok_or(OpteeSmcReturnCode::EBadCmd)?; self.known_flags.lock().entry(ta_uuid).or_insert(ta_flags); - self.sessions - .insert(session_id, instance.arc.clone(), ta_uuid, ta_flags); + self.sessions.insert(session_id, arc, ta_uuid, ta_flags); + Ok(()) } /// Unregister a session and recycle its session ID. Returns whether @@ -755,7 +721,8 @@ impl SessionManager { } /// Remove a single-instance TA from the cache only if the currently - /// cached `Arc` is the same as `instance`. + /// cached instance is the same as `instance` (matched by + /// `task_page_table_id`). /// /// Callers tearing down on TA panic must have already called /// [`SessionManager::mark_sessions_dead_for_instance`] before invoking @@ -767,10 +734,10 @@ impl SessionManager { pub fn remove_single_instance_if_same( &self, uuid: &TeeUuid, - instance: InstanceRef<'_>, + instance: &TaInstance, ) -> bool { self.single_instance_cache - .remove_if_same(uuid, instance.arc) + .remove_if_pt(uuid, instance.task_page_table_id) } /// Get the total count of unique TA instances (for limit checking). @@ -824,7 +791,7 @@ impl SessionManager { /// the cache and create rival instances. pub fn with_ta(&self, uuid: &TeeUuid, f: F) -> Result<(), OpteeSmcReturnCode> where - F: for<'a> FnOnce(Option>) -> Result<(), OpteeSmcReturnCode>, + F: for<'a> FnOnce(Option<&'a TaInstance>) -> Result<(), OpteeSmcReturnCode>, { let token = self.try_acquire_for_open(*uuid)?; let is_single_instance = token.uuid_lock.is_some(); @@ -832,7 +799,7 @@ impl SessionManager { // For single-instance TAs the per-UUID lock above keeps our UUID's // cache entry stable. For multi-instance we don't consult the cache. if is_single_instance && let Some(existing) = self.single_instance_cache.get(uuid) { - return f(Some(InstanceRef::new(&existing))); + return f(Some(&existing)); } { From 76d0abe9d20e94044b02a6d73f5d18b48515bea3 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Wed, 27 May 2026 15:28:23 +0000 Subject: [PATCH 10/28] revise --- litebox_runner_lvbs/src/lib.rs | 45 ++++++----------- litebox_shim_optee/src/session.rs | 84 ++++++++++++++----------------- 2 files changed, 52 insertions(+), 77 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index e02530f58f..2c94779c74 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -633,14 +633,10 @@ fn open_session_single_instance( if return_code == TeeResult::TargetDead { debug_serial_println!("Single-instance TA panicked during OpenSession, cleaning up"); - // Mark sibling sessions dead BEFORE evicting the cached single - // instance. Otherwise a racing handler that subsequently - // acquires its own session token for the UUID could walk past - // a still-Live session entry. + // Mark-then-evict ordering: see SessionManager::remove_single_instance_if_same. session_manager().mark_sessions_dead_for_instance(instance); let _ = session_manager().remove_single_instance_if_same(&ta_uuid, instance); - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. + // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(instance.shim(), task_pt_id); }; @@ -677,8 +673,7 @@ fn open_session_single_instance( if !ta_flags.is_keep_alive() && session_manager().count_sessions_for_instance(instance) == 0 { let _ = session_manager().remove_single_instance_if_same(&ta_uuid, instance); - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. + // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(instance.shim(), task_pt_id); }; @@ -750,8 +745,7 @@ fn open_session_new_instance( runner_session_id, ) .map_err(|_| { - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. + // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(&shim, task_pt_id) }; OpteeSmcReturnCode::ENomem })?, @@ -794,8 +788,7 @@ fn open_session_new_instance( Some(ta_req_info), ); - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. + // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(&shim, task_pt_id) }; write_result?; @@ -804,8 +797,7 @@ fn open_session_new_instance( // Load TA context with parameters for OpenSession - pass actual session_id loaded_program.entrypoints.as_ref().ok_or_else(|| { - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. + // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(&shim, task_pt_id) }; OpteeSmcReturnCode::EBadCmd })?; @@ -820,8 +812,7 @@ fn open_session_new_instance( None, ) .map_err(|_| { - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. + // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(&shim, task_pt_id) }; OpteeSmcReturnCode::EBadCmd })?; @@ -837,16 +828,14 @@ fn open_session_new_instance( // Read TA output parameters from the stack buffer let params_address = loaded_program.params_address.ok_or_else(|| { - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. + // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(&shim, task_pt_id) }; OpteeSmcReturnCode::EBadAddr })?; let ta_params = UserConstPtr::::from_usize(params_address) .read_at_offset(0) .ok_or_else(|| { - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. + // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(&shim, task_pt_id) }; OpteeSmcReturnCode::EBadAddr })?; @@ -873,8 +862,7 @@ fn open_session_new_instance( Some(ta_req_info), ); - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. + // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(&shim, task_pt_id) }; write_result?; @@ -895,8 +883,7 @@ fn open_session_new_instance( Some(ta_req_info), ) .inspect_err(|_| { - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. + // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(&shim, task_pt_id) }; })?; @@ -1043,17 +1030,14 @@ fn handle_invoke_command( let ta_flags = session.ta_flags; if ta_flags.is_single_instance() { - // Mark siblings dead BEFORE evicting the cached single instance. - // Otherwise a racing handler entering with_session/with_ta - // for the UUID could walk past a still-Live session entry. + // Mark-then-evict ordering: see SessionManager::remove_single_instance_if_same. session_manager().mark_sessions_dead_for_instance(instance); let _ = session_manager().remove_single_instance_if_same(&ta_uuid, instance); } session_manager().unregister_session(session_id); - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. + // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(instance.shim(), task_pt_id); }; @@ -1166,8 +1150,7 @@ fn handle_close_session( let _ = session_manager() .remove_single_instance_if_same(&ta_uuid, instance); } - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. + // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(instance.shim(), task_pt_id); }; diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 5fa87fa978..55e5e87e5e 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -48,13 +48,9 @@ //! the waiting logic in normal world (where scheduling is appropriate), without //! requiring RPCs that would give untrusted code control over secure world execution. //! -//! On panic teardown or last-session close, sibling sessions of a single-instance -//! TA are flipped to `Dead` *before* the cached instance is -//! evicted (see [`SessionManager::remove_single_instance_if_same`]). A racing -//! handler that subsequently enters [`SessionManager::with_ta`] or -//! [`SessionManager::with_session`] for the UUID will therefore observe `Dead` -//! on its re-read of the session entry and short-circuit through the -//! dead-target path. +//! Cleanup paths flip sibling sessions to `Dead` before evicting the +//! cached instance; see [`SessionManager::remove_single_instance_if_same`] +//! for the ordering rationale. //! //! Reference: //! @@ -187,12 +183,8 @@ pub(crate) enum SessionTarget { } /// Closure-bound snapshot of a session, delivered by -/// [`SessionManager::with_session`]. -/// -/// Holds the session's `ta_uuid` / `ta_flags` and, for live sessions, a -/// borrow of the [`TaInstance`]. The borrow's lifetime is pinned to the -/// `SessionToken` held internally by `with_session` via HRTB on the -/// closure type, so the closure cannot smuggle it out. +/// [`SessionManager::with_session`]. Holds `ta_uuid` / `ta_flags` and, +/// for live sessions, a borrow of the [`TaInstance`]. pub struct Session<'a> { pub ta_uuid: TeeUuid, pub ta_flags: TaFlags, @@ -466,11 +458,9 @@ impl SessionManager { } } - /// Mark every session currently pointing at `instance` as `Dead`. - /// - /// Must be called before evicting `instance` from the single-instance - /// cache (see [`SessionManager::remove_single_instance_if_same`]) so - /// a racing handler re-reads `Dead` on its sibling session entry. + /// Mark every session currently pointing at `instance` as `Dead`. Must + /// be paired with [`SessionManager::remove_single_instance_if_same`] + /// in the documented order — see that function for the rationale. pub fn mark_sessions_dead_for_instance(&self, instance: &TaInstance) { self.sessions .mark_sessions_dead_for_pt(instance.task_page_table_id); @@ -551,7 +541,7 @@ impl SessionManager { } /// Acquire a token + validated entry for an Invoke/Close on an existing - /// session. Returns the entry that survived the post-marker re-read so + /// session. Returns the entry that survived the post-lock re-read so /// callers don't need to look it up again. /// /// Always reserves the per-session-id slot in `active_sessions`. For @@ -563,13 +553,22 @@ impl SessionManager { /// holds the per-UUID lock for the same single-instance TA. On failure /// any partial acquisition is released via the token's `Drop`. /// - /// Defense in depth: after the per-session-id marker is held, the entry - /// is re-read and its `(uuid, flags)` validated against the pre-marker - /// snapshot used to decide whether to take the per-UUID lock. If they - /// diverge (the id was recycled and reused under a different TA between - /// our first read and the marker insert), `Err(EThreadLimit)` is - /// returned so the Linux driver retries — a fresh acquisition will see - /// the new entry from the start. + /// # Ordering + /// + /// The per-UUID lock is acquired *before* the final session-map + /// re-read. This excludes concurrent `mark_sessions_dead_for_instance` + /// and cache eviction (both of which require the UUID lock), so the + /// `Live` / `Dead` state observed in the re-read remains authoritative + /// for the lifetime of the returned token. Reading the entry before + /// taking the UUID lock would let a sibling complete the entire + /// mark-dead / evict / teardown sequence between our read and our + /// lock acquisition, leaving us holding a stale `Live` entry pointing + /// at a torn-down page table. + /// + /// Defense in depth: the entry's `(uuid, flags)` are validated against + /// the pre-marker snapshot. If they diverge (the id was recycled and + /// reused under a different TA between our first read and the marker + /// insert), we return `EThreadLimit` so the Linux driver retries. fn try_acquire_for_session( &self, session_id: u32, @@ -590,9 +589,18 @@ impl SessionManager { active_session_id: Some(session_id), }; - // Validate the snapshot under the marker. If the entry has changed - // identity (or vanished), our snapshot is stale; bail so the caller - // retries with a fresh view. Token's `Drop` releases the marker. + // Take the per-UUID lock BEFORE the final re-read for single- + // instance TAs. This blocks any concurrent mark-dead / cache + // eviction so the re-read result is stable. On failure, the + // token's `Drop` releases the marker we already took. + if snapshot_single { + token.uuid_lock = Some( + self.try_acquire_uuid_lock(snapshot_uuid) + .ok_or(OpteeSmcReturnCode::EThreadLimit)?, + ); + } + + // Re-read under both locks and validate against the snapshot. let entry_now = self .sessions .get_entry(session_id) @@ -603,18 +611,6 @@ impl SessionManager { return Err(OpteeSmcReturnCode::EThreadLimit); } - // Only take the per-UUID lock for `Live` single-instance sessions. - // A `Dead` entry needs no sibling serialization — its instance is - // already gone, and contending with live siblings (or a freshly - // created instance for the same UUID) just to call - // `finalize_dead_session` would needlessly delay them. - if snapshot_single && matches!(entry_now.target, SessionTarget::Live(_)) { - // On failure, dropping `token` releases the marker we just took. - token.uuid_lock = Some( - self.try_acquire_uuid_lock(snapshot_uuid) - .ok_or(OpteeSmcReturnCode::EThreadLimit)?, - ); - } Ok((token, entry_now)) } @@ -731,11 +727,7 @@ impl SessionManager { /// the UUID will observe `Dead` on its re-read of the session entry. /// Callers on the last-session-close path may skip the mark step — by /// that point there are no sibling sessions to fence out. - pub fn remove_single_instance_if_same( - &self, - uuid: &TeeUuid, - instance: &TaInstance, - ) -> bool { + pub fn remove_single_instance_if_same(&self, uuid: &TeeUuid, instance: &TaInstance) -> bool { self.single_instance_cache .remove_if_pt(uuid, instance.task_page_table_id) } From 423159b17ea157cfa86b754d7f299609c43c3728 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Wed, 27 May 2026 17:31:26 +0000 Subject: [PATCH 11/28] simplification --- litebox_runner_lvbs/src/lib.rs | 32 +++--- litebox_shim_optee/src/session.rs | 170 ++++++++++++++---------------- 2 files changed, 93 insertions(+), 109 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 2c94779c74..5e8bf40e35 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -522,7 +522,6 @@ fn handle_open_session( msg_args_phys_addr, instance, params, - ta_uuid, &ta_req_info, ), None => open_session_new_instance( @@ -546,10 +545,10 @@ fn open_session_single_instance( msg_args_phys_addr: u64, instance: &TaInstance, params: &[litebox_common_optee::UteeParamOwned], - ta_uuid: litebox_common_optee::TeeUuid, ta_req_info: &litebox_shim_optee::msg_handler::TaRequestInfo, ) -> Result<(), OpteeSmcReturnCode> { let task_pt_id = instance.task_page_table_id(); + let ta_uuid = instance.uuid(); // Allocate session ID BEFORE calling load_ta_context so TA gets correct ID. // Use SessionIdGuard to ensure the ID is recycled on any error path @@ -633,9 +632,9 @@ fn open_session_single_instance( if return_code == TeeResult::TargetDead { debug_serial_println!("Single-instance TA panicked during OpenSession, cleaning up"); - // Mark-then-evict ordering: see SessionManager::remove_single_instance_if_same. + // Mark-then-evict ordering: see SessionManager::evict_cached_instance. session_manager().mark_sessions_dead_for_instance(instance); - let _ = session_manager().remove_single_instance_if_same(&ta_uuid, instance); + let _ = session_manager().evict_cached_instance(instance); // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(instance.shim(), task_pt_id); @@ -672,7 +671,7 @@ fn open_session_single_instance( if let Err(e) = write_result { if !ta_flags.is_keep_alive() && session_manager().count_sessions_for_instance(instance) == 0 { - let _ = session_manager().remove_single_instance_if_same(&ta_uuid, instance); + let _ = session_manager().evict_cached_instance(instance); // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(instance.shim(), task_pt_id); @@ -684,7 +683,7 @@ fn open_session_single_instance( } // Success: register a sibling session pointing at the existing instance. - session_manager().register_sibling_session(runner_session_id, instance, ta_uuid, ta_flags)?; + session_manager().register_sibling_session(runner_session_id, instance)?; session_id_guard.disarm(); debug_serial_println!( @@ -898,7 +897,6 @@ fn open_session_new_instance( loaded_program, task_pt_id, ta_uuid, - ta_flags, ); session_id_guard.disarm(); @@ -952,8 +950,8 @@ fn handle_invoke_command( let params = &ta_req_info.params; let session_id = ta_req_info.session; - session_manager().with_session(session_id, |session| { - let Some(instance) = session.live() else { + session_manager().with_session(session_id, |instance| { + let Some(instance) = instance else { return finalize_dead_session( session_id, msg_args, @@ -1026,13 +1024,10 @@ fn handle_invoke_command( session_id ); - let ta_uuid = session.ta_uuid; - let ta_flags = session.ta_flags; - - if ta_flags.is_single_instance() { - // Mark-then-evict ordering: see SessionManager::remove_single_instance_if_same. + if instance.loaded_program().ta_flags.is_single_instance() { + // Mark-then-evict ordering: see SessionManager::evict_cached_instance. session_manager().mark_sessions_dead_for_instance(instance); - let _ = session_manager().remove_single_instance_if_same(&ta_uuid, instance); + let _ = session_manager().evict_cached_instance(instance); } session_manager().unregister_session(session_id); @@ -1073,9 +1068,8 @@ fn handle_close_session( debug_serial_println!("CloseSession: session_id={}", session_id); - session_manager().with_session(session_id, |session| { - let ta_uuid = session.ta_uuid; - let Some(instance) = session.live() else { + session_manager().with_session(session_id, |instance| { + let Some(instance) = instance else { return finalize_dead_session( session_id, msg_args, @@ -1148,7 +1142,7 @@ fn handle_close_session( // we confirm no sibling sessions remain. We don't need to mark anything `Dead` first. if flags.is_single_instance() { let _ = session_manager() - .remove_single_instance_if_same(&ta_uuid, instance); + .evict_cached_instance(instance); } // SAFETY: no references to user-space memory will be held after this call. unsafe { diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 55e5e87e5e..495732c184 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -49,7 +49,7 @@ //! requiring RPCs that would give untrusted code control over secure world execution. //! //! Cleanup paths flip sibling sessions to `Dead` before evicting the -//! cached instance; see [`SessionManager::remove_single_instance_if_same`] +//! cached instance; see [`SessionManager::evict_cached_instance`] //! for the ordering rationale. //! //! Reference: @@ -145,6 +145,7 @@ pub struct TaInstance { /// operations: page table ids are minted by `create_task_page_table()` /// and not reused until the owning instance is fully torn down. task_page_table_id: usize, + ta_uuid: TeeUuid, } impl TaInstance { @@ -159,6 +160,10 @@ impl TaInstance { pub fn loaded_program(&self) -> &LoadedProgram { &self.loaded_program } + + pub fn uuid(&self) -> TeeUuid { + self.ta_uuid + } } // SAFETY: `TaInstance`'s interior (`shim`, `loaded_program`) is not @@ -169,46 +174,37 @@ impl TaInstance { unsafe impl Send for TaInstance {} unsafe impl Sync for TaInstance {} -/// The target associated with a normal-world session ID. +/// Per-session entry in the session map. +/// +/// Module-private. External callers see liveness via the +/// `Option<&TaInstance>` delivered to closures by +/// [`SessionManager::with_ta`] / [`SessionManager::with_session`]: +/// `Some` for live, `None` for dead. /// -/// This is the in-map representation, kept private to the module. External -/// callers see liveness via [`Session::live`] on the [`Session`] delivered -/// by the session-token-bound closure APIs ([`SessionManager::with_ta`], -/// [`SessionManager::with_session`]). +/// Live entries carry the `Arc` (uuid and flags reachable via +/// it). Dead entries retain only the historical `(ta_uuid, ta_flags)` — +/// enough to drive cleanup paths and the `try_acquire_for_session` +/// snapshot check. #[derive(Clone)] -pub(crate) enum SessionTarget { +pub(crate) enum SessionEntry { Live(Arc), - /// The TA died, but normal world may still issue Invoke/Close for this ID. - Dead, + Dead { ta_uuid: TeeUuid, ta_flags: TaFlags }, } -/// Closure-bound snapshot of a session, delivered by -/// [`SessionManager::with_session`]. Holds `ta_uuid` / `ta_flags` and, -/// for live sessions, a borrow of the [`TaInstance`]. -pub struct Session<'a> { - pub ta_uuid: TeeUuid, - pub ta_flags: TaFlags, - instance: Option<&'a TaInstance>, -} - -impl<'a> Session<'a> { - /// Returns `Some(instance)` if the session's target is live, or `None` - /// if the TA has died and the caller should run dead-session cleanup. - pub fn live(&self) -> Option<&'a TaInstance> { - self.instance +impl SessionEntry { + fn ta_uuid(&self) -> TeeUuid { + match self { + SessionEntry::Live(arc) => arc.ta_uuid, + SessionEntry::Dead { ta_uuid, .. } => *ta_uuid, + } } -} -/// Per-session entry in the session map. Module-private; the closure-bound -/// public view is [`Session`]. -#[derive(Clone)] -pub(crate) struct SessionEntry { - /// The TA target (may be shared with other sessions for single-instance TAs). - pub target: SessionTarget, - /// The TA UUID (needed for cleanup of single-instance TAs). - pub ta_uuid: TeeUuid, - /// TA flags parsed from the `.ta_head` section. - pub ta_flags: TaFlags, + fn ta_flags(&self) -> TaFlags { + match self { + SessionEntry::Live(arc) => arc.loaded_program.ta_flags, + SessionEntry::Dead { ta_flags, .. } => *ta_flags, + } + } } /// Session map for tracking active sessions. @@ -229,21 +225,10 @@ impl SessionMap { self.inner.lock().get(&session_id).cloned() } - pub(crate) fn insert( - &self, - session_id: u32, - instance: Arc, - ta_uuid: TeeUuid, - ta_flags: TaFlags, - ) { - self.inner.lock().insert( - session_id, - SessionEntry { - target: SessionTarget::Live(instance), - ta_uuid, - ta_flags, - }, - ); + pub(crate) fn insert_live(&self, session_id: u32, instance: Arc) { + self.inner + .lock() + .insert(session_id, SessionEntry::Live(instance)); } pub(crate) fn remove(&self, session_id: u32) -> Option { @@ -255,20 +240,26 @@ impl SessionMap { self.inner .lock() .values() - .filter(|e| match &e.target { - SessionTarget::Live(arc) => arc.task_page_table_id == task_page_table_id, - SessionTarget::Dead => false, + .filter(|e| match e { + SessionEntry::Live(arc) => arc.task_page_table_id == task_page_table_id, + SessionEntry::Dead { .. } => false, }) .count() } /// Mark all live sessions whose instance has the given page table id - /// as `Dead`. + /// as `Dead`, capturing the instance's uuid and flags on the way out + /// so cleanup paths still have them. pub(crate) fn mark_sessions_dead_for_pt(&self, task_page_table_id: usize) { for entry in self.inner.lock().values_mut() { - if matches!(&entry.target, SessionTarget::Live(arc) if arc.task_page_table_id == task_page_table_id) - { - entry.target = SessionTarget::Dead; + let dead = match entry { + SessionEntry::Live(arc) if arc.task_page_table_id == task_page_table_id => { + Some((arc.ta_uuid, arc.loaded_program.ta_flags)) + } + _ => None, + }; + if let Some((ta_uuid, ta_flags)) = dead { + *entry = SessionEntry::Dead { ta_uuid, ta_flags }; } } } @@ -459,7 +450,7 @@ impl SessionManager { } /// Mark every session currently pointing at `instance` as `Dead`. Must - /// be paired with [`SessionManager::remove_single_instance_if_same`] + /// be paired with [`SessionManager::evict_cached_instance`] /// in the documented order — see that function for the rationale. pub fn mark_sessions_dead_for_instance(&self, instance: &TaInstance) { self.sessions @@ -577,8 +568,8 @@ impl SessionManager { .sessions .get_entry(session_id) .ok_or(OpteeSmcReturnCode::EBadCmd)?; - let snapshot_uuid = entry.ta_uuid; - let snapshot_single = entry.ta_flags.is_single_instance(); + let snapshot_uuid = entry.ta_uuid(); + let snapshot_single = entry.ta_flags().is_single_instance(); if !self.active_sessions.lock().insert(session_id) { return Err(OpteeSmcReturnCode::EThreadLimit); @@ -605,8 +596,8 @@ impl SessionManager { .sessions .get_entry(session_id) .ok_or(OpteeSmcReturnCode::EBadCmd)?; - if entry_now.ta_uuid != snapshot_uuid - || entry_now.ta_flags.is_single_instance() != snapshot_single + if entry_now.ta_uuid() != snapshot_uuid + || entry_now.ta_flags().is_single_instance() != snapshot_single { return Err(OpteeSmcReturnCode::EThreadLimit); } @@ -617,11 +608,12 @@ impl SessionManager { /// Drive an Invoke/Close to completion under the right serialization. /// /// Internally acquires the per-session-id marker (and, for single- - /// instance TAs, the per-UUID lock), passes a validated [`Session`] - /// to `f`, and releases the locks when `f` returns. `f` runs entirely - /// under the token: state mutations it performs on the session manager - /// (e.g. `unregister_session`, `mark_sessions_dead_for_instance`, - /// `remove_single_instance_if_same`) are serialized against other + /// instance TAs, the per-UUID lock), passes `Some(&TaInstance)` to `f` + /// for live sessions or `None` for dead ones, and releases the locks + /// when `f` returns. `f` runs entirely under the token: state + /// mutations it performs on the session manager (e.g. + /// `unregister_session`, `mark_sessions_dead_for_instance`, + /// `evict_cached_instance`) are serialized against other /// cores' Invoke/Close on the same session and (for single-instance) /// the same UUID. /// @@ -630,18 +622,14 @@ impl SessionManager { /// retries `EThreadLimit` transparently. pub fn with_session(&self, session_id: u32, f: F) -> Result<(), OpteeSmcReturnCode> where - F: for<'a> FnOnce(Session<'a>) -> Result<(), OpteeSmcReturnCode>, + F: for<'a> FnOnce(Option<&'a TaInstance>) -> Result<(), OpteeSmcReturnCode>, { let (_token, entry) = self.try_acquire_for_session(session_id)?; - let session = Session { - ta_uuid: entry.ta_uuid, - ta_flags: entry.ta_flags, - instance: match &entry.target { - SessionTarget::Live(arc) => Some(&**arc), - SessionTarget::Dead => None, - }, + let instance = match &entry { + SessionEntry::Live(arc) => Some(&**arc), + SessionEntry::Dead { .. } => None, }; - f(session) + f(instance) } /// Register a session for a freshly-loaded TA. The three parts (`shim`, @@ -660,16 +648,16 @@ impl SessionManager { loaded_program: alloc::boxed::Box, task_page_table_id: usize, ta_uuid: TeeUuid, - ta_flags: TaFlags, ) { + let ta_flags = loaded_program.ta_flags; let arc = Arc::new(TaInstance { shim, loaded_program, task_page_table_id, + ta_uuid, }); self.known_flags.lock().entry(ta_uuid).or_insert(ta_flags); - self.sessions - .insert(session_id, arc.clone(), ta_uuid, ta_flags); + self.sessions.insert_live(session_id, arc.clone()); if ta_flags.is_single_instance() { self.single_instance_cache.insert(ta_uuid, arc); } @@ -691,16 +679,17 @@ impl SessionManager { &self, session_id: u32, instance: &TaInstance, - ta_uuid: TeeUuid, - ta_flags: TaFlags, ) -> Result<(), OpteeSmcReturnCode> { let arc = self .single_instance_cache - .get(&ta_uuid) + .get(&instance.ta_uuid) .filter(|cached| cached.task_page_table_id == instance.task_page_table_id) .ok_or(OpteeSmcReturnCode::EBadCmd)?; - self.known_flags.lock().entry(ta_uuid).or_insert(ta_flags); - self.sessions.insert(session_id, arc, ta_uuid, ta_flags); + self.known_flags + .lock() + .entry(instance.ta_uuid) + .or_insert(instance.loaded_program.ta_flags); + self.sessions.insert_live(session_id, arc); Ok(()) } @@ -713,12 +702,13 @@ impl SessionManager { if entry.is_some() { recycle_session_id(session_id); } - entry.map(|e| e.ta_flags) + entry.map(|e| e.ta_flags()) } - /// Remove a single-instance TA from the cache only if the currently - /// cached instance is the same as `instance` (matched by - /// `task_page_table_id`). + /// Evict `instance` from the single-instance cache. No-op (returns + /// `false`) if the cached entry under `instance.uuid()` is a different + /// instance — matched by `task_page_table_id` to distinguish the + /// caller's instance from a freshly-cached replacement. /// /// Callers tearing down on TA panic must have already called /// [`SessionManager::mark_sessions_dead_for_instance`] before invoking @@ -727,9 +717,9 @@ impl SessionManager { /// the UUID will observe `Dead` on its re-read of the session entry. /// Callers on the last-session-close path may skip the mark step — by /// that point there are no sibling sessions to fence out. - pub fn remove_single_instance_if_same(&self, uuid: &TeeUuid, instance: &TaInstance) -> bool { + pub fn evict_cached_instance(&self, instance: &TaInstance) -> bool { self.single_instance_cache - .remove_if_pt(uuid, instance.task_page_table_id) + .remove_if_pt(&instance.ta_uuid, instance.task_page_table_id) } /// Get the total count of unique TA instances (for limit checking). @@ -749,7 +739,7 @@ impl SessionManager { .inner .lock() .values() - .filter(|e| !e.ta_flags.is_single_instance()) + .filter(|e| !e.ta_flags().is_single_instance()) .count() } From 624830ae79ea5273eda258a188054fb42483530c Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Wed, 27 May 2026 18:19:32 +0000 Subject: [PATCH 12/28] handle invalid uuid --- litebox_shim_optee/src/session.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 495732c184..c5c82bdcbd 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -775,6 +775,9 @@ impl SessionManager { where F: for<'a> FnOnce(Option<&'a TaInstance>) -> Result<(), OpteeSmcReturnCode>, { + // Snapshot: was this UUID already known? If unknown and stays unknown after + // a failed load (i.e., invalid UUID), we'll evict the lock entry below. + let was_unknown = self.get_known_flags(uuid).is_none(); let token = self.try_acquire_for_open(*uuid)?; let is_single_instance = token.uuid_lock.is_some(); @@ -800,6 +803,11 @@ impl SessionManager { *pending = pending.saturating_sub(1); } + if result.is_err() && was_unknown && self.get_known_flags(uuid).is_none() { + // Remove while still holding the lock (via `token`). + self.single_instance_locks.lock().remove(uuid); + } + result } } From 55e8adbad2110108b0681eb0294182a9d46e9aef Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 28 May 2026 02:14:26 +0000 Subject: [PATCH 13/28] add tests --- litebox_runner_lvbs/src/lib.rs | 13 +- litebox_shim_optee/src/session.rs | 236 ++++++++++++++++++++++++++++++ 2 files changed, 244 insertions(+), 5 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 5e8bf40e35..49710e852d 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -356,12 +356,15 @@ unsafe fn delete_task_page_table(task_pt_id: usize) -> Result<(), OpteeSmcReturn } } -/// Guard that restores the base page table when leaving a TA page table scope. +/// Enforces the invariant that the core must be on the base (kernel) page +/// table before returning to VTL0: the guard switches to the TA's task +/// page table on entry and switches back to the base page table on drop, +/// regardless of the path out (early return, `?`, panic). /// -/// `switch_to_base_page_table` is an idempotent CR3 write, so it is fine if -/// teardown paths (which switch to base internally before deleting the task -/// page table) run before this guard's `Drop` — the redundant write at drop -/// time is benign. +/// `switch_to_base_page_table` is an idempotent CR3 write, so teardown +/// paths that switch to base internally before deleting the task page +/// table can run before this guard's `Drop` — the redundant write at +/// drop time is benign. struct TaskPageTableGuard; impl TaskPageTableGuard { diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index c5c82bdcbd..4d6cf2daa3 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -817,3 +817,239 @@ impl Default for SessionManager { Self::new() } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::syscalls::tests::init_platform; + + fn make_shim() -> OpteeShim { + let _ = init_platform(); + crate::OpteeShimBuilder::new().build() + } + + fn make_loaded_program(ta_flags: TaFlags) -> alloc::boxed::Box { + alloc::boxed::Box::new(LoadedProgram { + entrypoints: None, + params_address: None, + ta_flags, + }) + } + + fn make_uuid(seed: u8) -> TeeUuid { + TeeUuid::from_bytes([seed; 16]) + } + + fn single_instance_flags() -> TaFlags { + TaFlags::SINGLE_INSTANCE | TaFlags::MULTI_SESSION + } + + /// Identity is by `task_page_table_id`, not by Arc pointer. After an + /// instance is evicted and a fresh one registered under the same UUID, + /// the stale handle must not evict the new one. + #[test] + fn evict_cached_instance_distinguishes_stale_handle() { + let manager = SessionManager::new(); + let uuid = make_uuid(0xA4); + + manager.register_new_session( + 105, + make_shim(), + make_loaded_program(single_instance_flags()), + 10, + uuid, + ); + let arc_first = manager.single_instance_cache.get(&uuid).unwrap(); + manager.evict_cached_instance(&arc_first); + + manager.register_new_session( + 106, + make_shim(), + make_loaded_program(single_instance_flags()), + 11, + uuid, + ); + assert!(!manager.evict_cached_instance(&arc_first)); + assert!(manager.single_instance_cache.get(&uuid).is_some()); + } + + /// `mark_sessions_dead_for_instance` flips Live entries to Dead — they + /// stop counting for `count_sessions_for_instance`, and `with_session` + /// thereafter sees `None` so cleanup paths run. + #[test] + fn mark_dead_makes_with_session_observe_none() { + let manager = SessionManager::new(); + let uuid = make_uuid(0xA6); + manager.register_new_session( + 108, + make_shim(), + make_loaded_program(single_instance_flags()), + 55, + uuid, + ); + let arc = manager.single_instance_cache.get(&uuid).unwrap(); + assert_eq!(manager.count_sessions_for_instance(&arc), 1); + + manager.mark_sessions_dead_for_instance(&arc); + assert_eq!(manager.count_sessions_for_instance(&arc), 0); + + manager + .with_session(108, |instance| { + assert!(instance.is_none()); + Ok(()) + }) + .unwrap(); + } + + /// Per-session-id marker excludes re-entry on the same id, but releases + /// when the closure returns. + #[test] + fn with_session_marker_excludes_reentry() { + let manager = SessionManager::new(); + let uuid = make_uuid(0xA7); + manager.register_new_session( + 109, + make_shim(), + make_loaded_program(single_instance_flags()), + 6, + uuid, + ); + + manager + .with_session(109, |_| { + assert_eq!( + manager.with_session(109, |_| Ok(())), + Err(OpteeSmcReturnCode::EThreadLimit) + ); + Ok(()) + }) + .unwrap(); + manager.with_session(109, |_| Ok(())).unwrap(); + } + + /// For an unknown UUID whose load fails, the per-UUID lock entry must + /// be evicted so it doesn't leak across malformed-UUID retries. + #[test] + fn with_ta_evicts_lock_entry_on_unknown_failure() { + let manager = SessionManager::new(); + let uuid = make_uuid(0xA9); + assert!(manager.get_known_flags(&uuid).is_none()); + + let _ = manager.with_ta(&uuid, |_| Err(OpteeSmcReturnCode::ENotAvail)); + + assert!(manager.single_instance_locks.lock().get(&uuid).is_none()); + } + + /// For a *known* UUID, the lock entry must be retained even on failure: + /// concurrent threads may already hold the same Arc, and removing it + /// would let a fresh entrant create a parallel mutex (split-brain). + #[test] + fn with_ta_keeps_lock_entry_after_known_uuid_failure() { + let manager = SessionManager::new(); + let uuid = make_uuid(0xAA); + manager.register_new_session( + 111, + make_shim(), + make_loaded_program(single_instance_flags()), + 8, + uuid, + ); + + manager.with_ta(&uuid, |_| Ok(())).unwrap(); + assert!(manager.single_instance_locks.lock().get(&uuid).is_some()); + + let _ = manager.with_ta(&uuid, |_| Err(OpteeSmcReturnCode::EBadCmd)); + assert!(manager.single_instance_locks.lock().get(&uuid).is_some()); + } + + /// `try_acquire_for_session`'s post-marker re-read must catch the case + /// where the entry's identity changed between snapshot and validation + /// (id recycled, re-registered under a different UUID). It should + /// return `EThreadLimit` so the driver retries with a fresh snapshot. + #[test] + fn try_acquire_for_session_rejects_uuid_swap_under_marker() { + let manager = SessionManager::new(); + let uuid_a = make_uuid(0xB0); + let uuid_b = make_uuid(0xB1); + let session_id = 222; + + manager.register_new_session( + session_id, + make_shim(), + make_loaded_program(single_instance_flags()), + 70, + uuid_a, + ); + + // Simulate the swap: unregister and re-register the same session_id + // under a different UUID, then drive try_acquire_for_session by + // hand to validate the snapshot path. (The real-world racing + // version of this is what the post-marker validation defends + // against; here we just verify the validation actually runs.) + let entry_before = manager.sessions.get_entry(session_id).expect("registered"); + assert_eq!(entry_before.ta_uuid(), uuid_a); + + manager.unregister_session(session_id); + manager.register_new_session( + session_id, + make_shim(), + make_loaded_program(single_instance_flags()), + 71, + uuid_b, + ); + + let entry_after = manager + .sessions + .get_entry(session_id) + .expect("re-registered"); + assert_ne!(entry_before.ta_uuid(), entry_after.ta_uuid()); + + // The validation in `try_acquire_for_session` compares snapshot + // uuid/flags against the post-marker re-read; a mismatch returns + // `EThreadLimit`. We exercise that comparison directly: with + // identical snapshots it succeeds, with mismatched it would not. + let (_, validated) = manager.try_acquire_for_session(session_id).unwrap(); + assert_eq!(validated.ta_uuid(), uuid_b); + } + + /// `pending_count` is bumped only on the create path, never on the + /// cache-hit path, and is decremented when the closure returns whether + /// success or failure — across multiple calls it must return to zero. + #[test] + fn pending_count_returns_to_zero_across_paths() { + let manager = SessionManager::new(); + let uuid_multi = make_uuid(0xC0); + let uuid_single = make_uuid(0xC1); + + // Successful create path. + manager + .with_ta(&uuid_multi, |existing| { + assert!(existing.is_none()); + manager.register_new_session( + 301, + make_shim(), + make_loaded_program(TaFlags::default()), + 80, + uuid_multi, + ); + Ok(()) + }) + .unwrap(); + assert_eq!(*manager.pending_count.lock(), 0); + + // Failing create path on an unknown UUID — also evicts the lock entry. + let _ = manager.with_ta(&uuid_single, |_| Err(OpteeSmcReturnCode::ENotAvail)); + assert_eq!(*manager.pending_count.lock(), 0); + + // Cache-hit path doesn't touch pending_count. + manager.register_new_session( + 302, + make_shim(), + make_loaded_program(single_instance_flags()), + 81, + uuid_single, + ); + manager.with_ta(&uuid_single, |_| Ok(())).unwrap(); + assert_eq!(*manager.pending_count.lock(), 0); + } +} From fed7cbf9dc290e2c4aab622aff14282521d4a1b2 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 28 May 2026 02:47:43 +0000 Subject: [PATCH 14/28] update doc --- litebox_runner_lvbs/src/lib.rs | 41 +++++++++++++--------------------- 1 file changed, 15 insertions(+), 26 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 49710e852d..f8a3f635e3 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -410,12 +410,9 @@ unsafe fn teardown_ta_page_table(shim: &litebox_shim_optee::OpteeShim, task_pt_i /// For TA requests (OpenSession, InvokeCommand, CloseSession), it uses `decode_ta_request` /// to extract the TA request information and load/run it using `OpteeShim`. /// -/// OpenSession for multi-instance TA creates: -/// - A new task page table for memory isolation -/// - A new TA instance with its own state -/// - An entry in the global session map -/// -/// OpenSession for single-instance TA reuses existing TA instance if available, +/// OpenSession for multi-instance TAs creates a new task page table and a +/// new TA instance and registers it with the session manager. OpenSession +/// for single-instance TAs reuses the cached instance if available, /// otherwise creates a new one. /// /// InvokeCommand looks up the session and switches to its page table. @@ -616,10 +613,10 @@ fn open_session_single_instance( return_code ); - // Write error response BEFORE switching page tables (accesses user memory). - // The session token held by `with_ta` keeps another core - // from tearing down the active page table while this core is copying - // TA outputs. + // Write error response BEFORE switching page tables — accesses user + // memory, which requires the TA's page table to still be active. + // `with_ta`'s serialization prevents another core from tearing down + // the instance underneath us while we copy TA outputs. let write_result = write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, @@ -697,10 +694,8 @@ fn open_session_single_instance( Ok(()) } -/// Create a new TA instance for a session. -/// -/// The caller must invoke this inside [`SessionManager::with_ta`] -/// to ensure a creation slot is held during execution and released afterward. +/// Create a new TA instance for a session. Must be called from within a +/// [`SessionManager::with_ta`] closure. /// /// If ldelf loading or OpenSession entry point fails, the page table is torn down. /// Per OP-TEE OS semantics: if OpenSession returns non-success, cleanup happens. @@ -889,11 +884,8 @@ fn open_session_new_instance( unsafe { teardown_ta_page_table(&shim, task_pt_id) }; })?; - // Success: hand the three parts to the session manager. It wraps them - // in an `Arc` owned solely by itself (no clone returned), - // and — for single-instance TAs — also caches that `Arc`. The runner - // never holds a `TaInstance` or `Arc`, which is what makes - // the internal `unsafe impl Send/Sync for TaInstance` enforceable. + // Success: hand the three parts to the session manager. The manager + // takes ownership; this runner never retains a handle to the instance. session_manager().register_new_session( runner_session_id, shim, @@ -914,8 +906,8 @@ fn open_session_new_instance( /// Tear down a `Dead` session entry observed at Invoke/Close handler entry. /// -/// Runs inside `with_session`'s closure, so the session token is alive for -/// the duration of this call and released when the closure returns. +/// Must be called from within a `with_session` closure so its serialization +/// covers the cleanup. fn finalize_dead_session( session_id: u32, msg_args: &mut OpteeMsgArgs, @@ -1006,9 +998,8 @@ fn handle_invoke_command( let return_code: u32 = ctx.rax.trunc(); let return_code = TeeResult::try_from(return_code).unwrap_or(TeeResult::GenericError); - // Write response BEFORE switching page tables (accesses user memory). - // The session token prevents another core from tearing down the active - // page table while this core is copying TA outputs. + // Write response BEFORE switching page tables — accesses user memory, + // which requires the TA's page table to still be active. let write_result = write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, @@ -1118,8 +1109,6 @@ fn handle_close_session( None, ); - // Remove the session entry from the map. The session token drops - // when the enclosing `with_session` closure returns. let removed_flags = session_manager().unregister_session(session_id); // Check if this was the last session using the TA instance by counting From de5981173765f83c5c0918a53487d979d2d829d6 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 28 May 2026 03:18:34 +0000 Subject: [PATCH 15/28] revert --- litebox_shim_optee/src/session.rs | 41 +++++-------------------------- 1 file changed, 6 insertions(+), 35 deletions(-) diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 4d6cf2daa3..39dd762400 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -775,9 +775,6 @@ impl SessionManager { where F: for<'a> FnOnce(Option<&'a TaInstance>) -> Result<(), OpteeSmcReturnCode>, { - // Snapshot: was this UUID already known? If unknown and stays unknown after - // a failed load (i.e., invalid UUID), we'll evict the lock entry below. - let was_unknown = self.get_known_flags(uuid).is_none(); let token = self.try_acquire_for_open(*uuid)?; let is_single_instance = token.uuid_lock.is_some(); @@ -803,11 +800,8 @@ impl SessionManager { *pending = pending.saturating_sub(1); } - if result.is_err() && was_unknown && self.get_known_flags(uuid).is_none() { - // Remove while still holding the lock (via `token`). - self.single_instance_locks.lock().remove(uuid); - } - + // `single_instance_locks` entries are never evicted; safe removal + // would need generational tracking. Leak is bounded by distinct UUIDs. result } } @@ -927,38 +921,15 @@ mod tests { manager.with_session(109, |_| Ok(())).unwrap(); } - /// For an unknown UUID whose load fails, the per-UUID lock entry must - /// be evicted so it doesn't leak across malformed-UUID retries. + /// `single_instance_locks` entries are never evicted, even on failure + /// of a previously-unknown UUID. #[test] - fn with_ta_evicts_lock_entry_on_unknown_failure() { + fn with_ta_never_evicts_lock_entry() { let manager = SessionManager::new(); let uuid = make_uuid(0xA9); assert!(manager.get_known_flags(&uuid).is_none()); let _ = manager.with_ta(&uuid, |_| Err(OpteeSmcReturnCode::ENotAvail)); - - assert!(manager.single_instance_locks.lock().get(&uuid).is_none()); - } - - /// For a *known* UUID, the lock entry must be retained even on failure: - /// concurrent threads may already hold the same Arc, and removing it - /// would let a fresh entrant create a parallel mutex (split-brain). - #[test] - fn with_ta_keeps_lock_entry_after_known_uuid_failure() { - let manager = SessionManager::new(); - let uuid = make_uuid(0xAA); - manager.register_new_session( - 111, - make_shim(), - make_loaded_program(single_instance_flags()), - 8, - uuid, - ); - - manager.with_ta(&uuid, |_| Ok(())).unwrap(); - assert!(manager.single_instance_locks.lock().get(&uuid).is_some()); - - let _ = manager.with_ta(&uuid, |_| Err(OpteeSmcReturnCode::EBadCmd)); assert!(manager.single_instance_locks.lock().get(&uuid).is_some()); } @@ -1037,7 +1008,7 @@ mod tests { .unwrap(); assert_eq!(*manager.pending_count.lock(), 0); - // Failing create path on an unknown UUID — also evicts the lock entry. + // Failing create path on an unknown UUID. let _ = manager.with_ta(&uuid_single, |_| Err(OpteeSmcReturnCode::ENotAvail)); assert_eq!(*manager.pending_count.lock(), 0); From c5f52c4565c0ec0bb070a71b772466f113ae1282 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 28 May 2026 03:58:44 +0000 Subject: [PATCH 16/28] doc --- litebox_runner_lvbs/src/lib.rs | 46 ++++++------------- litebox_shim_optee/src/session.rs | 73 ++++++++----------------------- 2 files changed, 33 insertions(+), 86 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index f8a3f635e3..1760b853f0 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -358,8 +358,8 @@ unsafe fn delete_task_page_table(task_pt_id: usize) -> Result<(), OpteeSmcReturn /// Enforces the invariant that the core must be on the base (kernel) page /// table before returning to VTL0: the guard switches to the TA's task -/// page table on entry and switches back to the base page table on drop, -/// regardless of the path out (early return, `?`, panic). +/// page table on entry and switches back on drop, covering early-return +/// and `?` paths. /// /// `switch_to_base_page_table` is an idempotent CR3 write, so teardown /// paths that switch to base internally before deleting the task page @@ -550,12 +550,11 @@ fn open_session_single_instance( let task_pt_id = instance.task_page_table_id(); let ta_uuid = instance.uuid(); - // Allocate session ID BEFORE calling load_ta_context so TA gets correct ID. - // Use SessionIdGuard to ensure the ID is recycled on any error path - // (before it is registered with the session manager). + // Allocate the session ID up front so the TA sees the right one in + // OpenSession. The guard recycles it on every error path until we + // either disarm it (on successful registration) or it drops. let session_id_guard = SessionIdGuard::new(allocate_session_id().ok_or(OpteeSmcReturnCode::EBusy)?); - // Safe to unwrap: guard was just created with Some(id). let runner_session_id = session_id_guard.id().unwrap(); debug_serial_println!( @@ -569,7 +568,7 @@ fn open_session_single_instance( let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; - // Load TA context with parameters for OpenSession - pass actual session_id + // Set up the entry-point parameters for OpenSession. instance .loaded_program() .entrypoints @@ -583,7 +582,6 @@ fn open_session_single_instance( ) .map_err(|_| OpteeSmcReturnCode::EBadCmd)?; - // Run the TA's OpenSession entry point using reference-based reenter let mut ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::reenter_thread_ref( @@ -592,7 +590,6 @@ fn open_session_single_instance( ); } - // Read TA output parameters from the stack buffer let params_address = instance .loaded_program() .params_address @@ -709,9 +706,7 @@ fn open_session_new_instance( ) -> Result<(), OpteeSmcReturnCode> { let ta_bin = find_ta_binary(ta_uuid).ok_or(OpteeSmcReturnCode::ENotAvail)?; - // Create and switch to new page table let task_pt_id = create_task_page_table()?; - debug_serial_println!("Created task page table ID: {}", task_pt_id); let task_pt_guard = TaskPageTableGuard::enter(task_pt_id).inspect_err(|_| { @@ -719,19 +714,14 @@ fn open_session_new_instance( let _ = unsafe { delete_task_page_table(task_pt_id) }; })?; - // Allocate session ID before loading - return EBusy to normal world if exhausted. - // Use SessionIdGuard to ensure the ID is recycled on any error path - // (before it is registered with the session manager). let Some(session_id) = allocate_session_id() else { drop(task_pt_guard); let _ = unsafe { delete_task_page_table(task_pt_id) }; return Err(OpteeSmcReturnCode::EBusy); }; let session_id_guard = SessionIdGuard::new(session_id); - // Safe to unwrap: guard was just created with Some(id). let runner_session_id = session_id_guard.id().unwrap(); - // 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( @@ -756,7 +746,7 @@ fn open_session_new_instance( ta_flags.is_single_instance() ); - // Run ldelf to load the TA using reference-based run to avoid moving the shim + // Run ldelf, which loads the TA and calls TA_CreateEntryPoint. let mut ldelf_ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::run_thread_ref( @@ -792,7 +782,7 @@ fn open_session_new_instance( return Ok(()); } - // Load TA context with parameters for OpenSession - pass actual session_id + // Set up the entry-point parameters for OpenSession. loaded_program.entrypoints.as_ref().ok_or_else(|| { // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(&shim, task_pt_id) }; @@ -814,7 +804,6 @@ fn open_session_new_instance( OpteeSmcReturnCode::EBadCmd })?; - // Run the TA entry function using reference-based reenter to avoid moving the shim let mut ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::reenter_thread_ref( @@ -823,7 +812,6 @@ fn open_session_new_instance( ); } - // Read TA output parameters from the stack buffer let params_address = loaded_program.params_address.ok_or_else(|| { // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(&shim, task_pt_id) }; @@ -884,8 +872,7 @@ fn open_session_new_instance( unsafe { teardown_ta_page_table(&shim, task_pt_id) }; })?; - // Success: hand the three parts to the session manager. The manager - // takes ownership; this runner never retains a handle to the instance. + // Success: register the new session with the manager. session_manager().register_new_session( runner_session_id, shim, @@ -966,7 +953,7 @@ fn handle_invoke_command( cmd_id ); - // Load TA context with parameters and cmd_id - pass actual session_id + // Set up the entry-point parameters for InvokeCommand. let entrypoints_ref = instance.loaded_program().entrypoints.as_ref().unwrap(); entrypoints_ref .load_ta_context( @@ -977,7 +964,6 @@ fn handle_invoke_command( ) .map_err(|_| OpteeSmcReturnCode::EBadCmd)?; - // Run the TA entry function using reference-based reenter to avoid moving the shim let mut ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::reenter_thread_ref( @@ -1076,7 +1062,7 @@ fn handle_close_session( let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; - // Load TA context for CloseSession (no params, no cmd_id) - pass actual session_id + // Set up the entry-point parameters for CloseSession. instance .loaded_program() .entrypoints @@ -1111,17 +1097,13 @@ fn handle_close_session( let removed_flags = session_manager().unregister_session(session_id); - // Check if this was the last session using the TA instance by counting - // remaining sessions that reference this instance. - let remaining_sessions = session_manager() - .count_sessions_for_instance(instance); + let remaining_sessions = session_manager().count_sessions_for_instance(instance); - // If this was the last session using the TA instance, clean up (unless keep_alive is set) + // Last session on this instance — tear it down unless `keep_alive` + // is set (only meaningful for single-instance TAs). if remaining_sessions == 0 && let Some(flags) = removed_flags { - // If this is a single-instance TA with keep_alive flag, don't remove it from memory. - // Note: keep_alive is only meaningful for single-instance TAs. if flags.is_single_instance() && flags.is_keep_alive() { debug_serial_println!( "CloseSession complete: session_id={}, TA kept alive (INSTANCE_KEEP_ALIVE flag)", diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 39dd762400..96c3605dba 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -121,13 +121,6 @@ pub const MAX_TA_INSTANCES: usize = 16; /// A loaded TA instance. /// -/// Fields are private; external callers never construct one (the three -/// constituent parts are passed to [`SessionManager::register_new_session`], -/// which builds the instance internally). Closures running under -/// [`SessionManager::with_ta`] / [`SessionManager::with_session`] observe -/// the instance through `&TaInstance`, with the borrow lifetime pinned to -/// the session token via HRTB on the closure type. -/// /// For single-instance TAs one instance is shared across all sessions; the /// TA stays in memory until the last session closes (or, with /// `TA_FLAG_INSTANCE_KEEP_ALIVE`, until explicit destroy). Each instance @@ -174,17 +167,9 @@ impl TaInstance { unsafe impl Send for TaInstance {} unsafe impl Sync for TaInstance {} -/// Per-session entry in the session map. -/// -/// Module-private. External callers see liveness via the -/// `Option<&TaInstance>` delivered to closures by -/// [`SessionManager::with_ta`] / [`SessionManager::with_session`]: -/// `Some` for live, `None` for dead. -/// -/// Live entries carry the `Arc` (uuid and flags reachable via -/// it). Dead entries retain only the historical `(ta_uuid, ta_flags)` — -/// enough to drive cleanup paths and the `try_acquire_for_session` -/// snapshot check. +/// Per-session entry in the session map. The `Dead` variant retains +/// `(ta_uuid, ta_flags)` so cleanup paths and `try_acquire_for_session`'s +/// snapshot still have them after the instance is gone. #[derive(Clone)] pub(crate) enum SessionEntry { Live(Arc), @@ -372,17 +357,18 @@ impl Drop for SessionIdGuard { /// RAII token bundling the serialization primitives required to safely /// execute an OP-TEE TA operation. /// -/// Held only inside [`SessionManager::with_ta`] (OpenSession) and -/// [`SessionManager::with_session`] (Invoke/Close); never exposed to -/// external callers. Bundles whichever combination of locks is required: +/// Bundles whichever combination of locks the current operation requires: /// /// - **Single-instance TAs**: a per-UUID `SpinMutex` that serializes all /// sessions on the same TA. /// - **Existing-session operations** (Invoke/Close): a per-session-id marker /// that prevents concurrent SMC entry by another core for the same id. /// -/// For multi-instance OpenSession, the token holds nothing (each session -/// gets its own private instance, so no exclusion is required). +/// For multi-instance OpenSession the token holds nothing (each session +/// gets its own private instance, so no exclusion is required). The +/// first-ever OpenSession for an unknown UUID is the exception: until the +/// TA is loaded its flags aren't known, so it's conservatively serialized +/// under the per-UUID lock until flags are observed. /// /// On drop, the per-UUID lock is released first, then the per-session-id /// marker. @@ -634,13 +620,8 @@ impl SessionManager { /// Register a session for a freshly-loaded TA. The three parts (`shim`, /// `loaded_program`, `task_page_table_id`) are taken by value and stored - /// inside the manager (no handle returned to the caller); for - /// single-instance TAs the instance is also cached under `ta_uuid` for - /// later reuse. - /// - /// The caller never retains a `TaInstance`, which is what makes the - /// internal `unsafe impl Send/Sync for TaInstance` invariant enforceable - /// against the public API. + /// inside the manager; for single-instance TAs the instance is also + /// cached under `ta_uuid` for later reuse. pub fn register_new_session( &self, session_id: u32, @@ -933,12 +914,15 @@ mod tests { assert!(manager.single_instance_locks.lock().get(&uuid).is_some()); } - /// `try_acquire_for_session`'s post-marker re-read must catch the case - /// where the entry's identity changed between snapshot and validation - /// (id recycled, re-registered under a different UUID). It should - /// return `EThreadLimit` so the driver retries with a fresh snapshot. + /// `try_acquire_for_session` returns the entry observed by the + /// post-marker re-read, not a stale handle from before the marker was + /// taken. After a recycle+re-register under a new UUID, the returned + /// entry must reflect the current UUID. (The mismatch-rejection branch + /// itself can only be triggered by a concurrent swap between snapshot + /// and re-read, which isn't reproducible in a single-threaded test; + /// this just verifies the re-read is the source of truth.) #[test] - fn try_acquire_for_session_rejects_uuid_swap_under_marker() { + fn try_acquire_for_session_returns_current_entry_after_recycle() { let manager = SessionManager::new(); let uuid_a = make_uuid(0xB0); let uuid_b = make_uuid(0xB1); @@ -951,15 +935,6 @@ mod tests { 70, uuid_a, ); - - // Simulate the swap: unregister and re-register the same session_id - // under a different UUID, then drive try_acquire_for_session by - // hand to validate the snapshot path. (The real-world racing - // version of this is what the post-marker validation defends - // against; here we just verify the validation actually runs.) - let entry_before = manager.sessions.get_entry(session_id).expect("registered"); - assert_eq!(entry_before.ta_uuid(), uuid_a); - manager.unregister_session(session_id); manager.register_new_session( session_id, @@ -969,16 +944,6 @@ mod tests { uuid_b, ); - let entry_after = manager - .sessions - .get_entry(session_id) - .expect("re-registered"); - assert_ne!(entry_before.ta_uuid(), entry_after.ta_uuid()); - - // The validation in `try_acquire_for_session` compares snapshot - // uuid/flags against the post-marker re-read; a mismatch returns - // `EThreadLimit`. We exercise that comparison directly: with - // identical snapshots it succeeds, with mismatched it would not. let (_, validated) = manager.try_acquire_for_session(session_id).unwrap(); assert_eq!(validated.ta_uuid(), uuid_b); } From 099cbb04220b189d8dc8b5cd3b4095259a066119 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 28 May 2026 04:51:01 +0000 Subject: [PATCH 17/28] fix scope --- litebox_shim_optee/src/lib.rs | 2 +- litebox_shim_optee/src/session.rs | 32 +++++++++++++++---------------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index 45acf64dcd..6f68273c73 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -38,7 +38,7 @@ pub mod msg_handler; pub mod ptr; // Re-export session management types for convenience -pub use session::{MAX_TA_INSTANCES, SessionManager, TaInstance, allocate_session_id}; +pub use session::{SessionManager, TaInstance, allocate_session_id}; const MAX_KERNEL_BUF_SIZE: usize = 0x80_000; diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 96c3605dba..c44547e72a 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -117,7 +117,7 @@ use litebox_common_optee::{OpteeSmcReturnCode, TaFlags, TeeUuid}; use spin::mutex::SpinMutex; /// Maximum number of concurrent TA instances to avoid out of memory situations. -pub const MAX_TA_INSTANCES: usize = 16; +const MAX_TA_INSTANCES: usize = 16; /// A loaded TA instance. /// @@ -171,7 +171,7 @@ unsafe impl Sync for TaInstance {} /// `(ta_uuid, ta_flags)` so cleanup paths and `try_acquire_for_session`'s /// snapshot still have them after the instance is gone. #[derive(Clone)] -pub(crate) enum SessionEntry { +enum SessionEntry { Live(Arc), Dead { ta_uuid: TeeUuid, ta_flags: TaFlags }, } @@ -195,33 +195,33 @@ impl SessionEntry { /// Session map for tracking active sessions. /// /// Maps runner-allocated session IDs to session entries. -pub(crate) struct SessionMap { +struct SessionMap { inner: SpinMutex>, } impl SessionMap { - pub(crate) fn new() -> Self { + fn new() -> Self { Self { inner: SpinMutex::new(HashMap::new()), } } - pub(crate) fn get_entry(&self, session_id: u32) -> Option { + fn get_entry(&self, session_id: u32) -> Option { self.inner.lock().get(&session_id).cloned() } - pub(crate) fn insert_live(&self, session_id: u32, instance: Arc) { + fn insert_live(&self, session_id: u32, instance: Arc) { self.inner .lock() .insert(session_id, SessionEntry::Live(instance)); } - pub(crate) fn remove(&self, session_id: u32) -> Option { + fn remove(&self, session_id: u32) -> Option { self.inner.lock().remove(&session_id) } /// Count live sessions whose instance has the given page table id. - pub(crate) fn count_sessions_for_pt(&self, task_page_table_id: usize) -> usize { + fn count_sessions_for_pt(&self, task_page_table_id: usize) -> usize { self.inner .lock() .values() @@ -235,7 +235,7 @@ impl SessionMap { /// Mark all live sessions whose instance has the given page table id /// as `Dead`, capturing the instance's uuid and flags on the way out /// so cleanup paths still have them. - pub(crate) fn mark_sessions_dead_for_pt(&self, task_page_table_id: usize) { + fn mark_sessions_dead_for_pt(&self, task_page_table_id: usize) { for entry in self.inner.lock().values_mut() { let dead = match entry { SessionEntry::Live(arc) if arc.task_page_table_id == task_page_table_id => { @@ -260,22 +260,22 @@ impl Default for SessionMap { /// /// Single-instance TAs (with `TA_FLAG_SINGLE_INSTANCE`) share a single TA instance /// across all sessions. This cache stores instances by UUID for fast reuse lookup. -pub(crate) struct SingleInstanceCache { +struct SingleInstanceCache { inner: SpinMutex>>, } impl SingleInstanceCache { - pub(crate) fn new() -> Self { + fn new() -> Self { Self { inner: SpinMutex::new(HashMap::new()), } } - pub(crate) fn get(&self, uuid: &TeeUuid) -> Option> { + fn get(&self, uuid: &TeeUuid) -> Option> { self.inner.lock().get(uuid).cloned() } - pub(crate) fn insert(&self, uuid: TeeUuid, instance: Arc) { + fn insert(&self, uuid: TeeUuid, instance: Arc) { self.inner.lock().insert(uuid, instance); } @@ -293,7 +293,7 @@ impl SingleInstanceCache { } } - pub(crate) fn len(&self) -> usize { + fn len(&self) -> usize { self.inner.lock().len() } } @@ -372,7 +372,7 @@ impl Drop for SessionIdGuard { /// /// On drop, the per-UUID lock is released first, then the per-session-id /// marker. -pub(crate) struct SessionToken<'a> { +struct SessionToken<'a> { manager: &'a SessionManager, /// Held `Arc` of the per-UUID `SpinMutex`. The guard returned by /// `try_lock()` was [`core::mem::forget`]-ed at acquisition time; this @@ -455,7 +455,7 @@ impl SessionManager { /// /// Returns `None` if this UUID has never been successfully loaded. /// Callers should conservatively assume single-instance when `None`. - pub fn get_known_flags(&self, uuid: &TeeUuid) -> Option { + fn get_known_flags(&self, uuid: &TeeUuid) -> Option { self.known_flags.lock().get(uuid).copied() } From 4c5652d45852db2315e3f72a8878b938a6001546 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 28 May 2026 13:30:08 +0000 Subject: [PATCH 18/28] unknown_uuid_lock --- litebox_runner_lvbs/src/lib.rs | 21 +++-- litebox_shim_optee/src/session.rs | 149 ++++++++++++++++-------------- 2 files changed, 95 insertions(+), 75 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 1760b853f0..8a60fa2c5b 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -5,8 +5,7 @@ extern crate alloc; -use alloc::boxed::Box; -use alloc::vec; +use alloc::{boxed::Box, vec}; use core::{ops::Neg, panic::PanicInfo}; use litebox::{ mm::linux::PAGE_SIZE, @@ -550,11 +549,12 @@ fn open_session_single_instance( let task_pt_id = instance.task_page_table_id(); let ta_uuid = instance.uuid(); - // Allocate the session ID up front so the TA sees the right one in - // OpenSession. The guard recycles it on every error path until we - // either disarm it (on successful registration) or it drops. + // Allocate session ID BEFORE calling load_ta_context so TA gets correct ID. + // Use SessionIdGuard to ensure the ID is recycled on any error path + // (before it is registered with the session manager). let session_id_guard = SessionIdGuard::new(allocate_session_id().ok_or(OpteeSmcReturnCode::EBusy)?); + // Safe to unwrap: guard was just created with Some(id). let runner_session_id = session_id_guard.id().unwrap(); debug_serial_println!( @@ -610,7 +610,7 @@ fn open_session_single_instance( return_code ); - // Write error response BEFORE switching page tables — accesses user + // Write error response BEFORE switching page tables. Accesses user // memory, which requires the TA's page table to still be active. // `with_ta`'s serialization prevents another core from tearing down // the instance underneath us while we copy TA outputs. @@ -706,6 +706,7 @@ fn open_session_new_instance( ) -> Result<(), OpteeSmcReturnCode> { let ta_bin = find_ta_binary(ta_uuid).ok_or(OpteeSmcReturnCode::ENotAvail)?; + // Create and switch to new page table let task_pt_id = create_task_page_table()?; debug_serial_println!("Created task page table ID: {}", task_pt_id); @@ -722,6 +723,7 @@ fn open_session_new_instance( let session_id_guard = SessionIdGuard::new(session_id); let runner_session_id = session_id_guard.id().unwrap(); + // 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( @@ -746,7 +748,7 @@ fn open_session_new_instance( ta_flags.is_single_instance() ); - // Run ldelf, which loads the TA and calls TA_CreateEntryPoint. + // Run ldelf to load the TA using reference-based run to avoid moving the shim let mut ldelf_ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::run_thread_ref( @@ -782,7 +784,7 @@ fn open_session_new_instance( return Ok(()); } - // Set up the entry-point parameters for OpenSession. + // Load TA context with parameters for OpenSession - pass actual session_id loaded_program.entrypoints.as_ref().ok_or_else(|| { // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(&shim, task_pt_id) }; @@ -804,6 +806,8 @@ fn open_session_new_instance( OpteeSmcReturnCode::EBadCmd })?; + // Run the TA's OpenSession entry point using reference-based reenter to + // avoid moving the shim let mut ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::reenter_thread_ref( @@ -812,6 +816,7 @@ fn open_session_new_instance( ); } + // Read TA output parameters from the stack buffer let params_address = loaded_program.params_address.ok_or_else(|| { // SAFETY: no references to user-space memory will be held after this call. unsafe { teardown_ta_page_table(&shim, task_pt_id) }; diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index c44547e72a..b047cf707f 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -10,24 +10,16 @@ //! //! ## Concurrency Model //! -//! TA execution is serialized externally; [`TaInstance`] is shared without an -//! inner mutex. The exclusivity invariant lives in [`SessionManager`] and is -//! acquired through an internal RAII `SessionToken` that bundles whichever -//! locks the current operation requires: -//! -//! - **Single-instance TAs** (with `TA_FLAG_SINGLE_INSTANCE | TA_FLAG_MULTI_SESSION`) -//! share one [`TaInstance`] across all sessions. The token internally holds a -//! per-UUID `SpinMutex` so Open/Invoke/Close serialize on the same UUID. -//! -//! - **Multi-instance TAs** have one [`TaInstance`] per session. The token -//! internally holds a per-`session_id` marker so Invoke/Close cannot -//! re-enter the same session concurrently, while different sessions run -//! in parallel on their own instances. +//! TA execution is serialized externally; [`TaInstance`] is shared without +//! an inner mutex. The exclusivity invariant lives in [`SessionManager`] +//! and is acquired through an internal RAII `SessionToken` that bundles +//! whichever locks the current operation requires — see `SessionToken`'s +//! doc for the per-case breakdown. //! //! Both [`SessionManager::with_ta`] (OpenSession) and //! [`SessionManager::with_session`] (Invoke/Close) acquire the token -//! non-blockingly, run the caller's closure under it, and release on return. -//! On contention they return `EThreadLimit`. +//! non-blockingly, run the caller's closure under it, and release on +//! return. On contention they return `EThreadLimit`. //! //! ### Difference from OP-TEE OS //! @@ -359,19 +351,18 @@ impl Drop for SessionIdGuard { /// /// Bundles whichever combination of locks the current operation requires: /// -/// - **Single-instance TAs**: a per-UUID `SpinMutex` that serializes all -/// sessions on the same TA. +/// - **Known single-instance TAs**: a per-UUID `SpinMutex` that serializes +/// all sessions on the same TA. +/// - **First-ever load of an unknown UUID** (OpenSession only): the shared +/// `unknown_uuid_lock`, used until the TA's flags are observed. /// - **Existing-session operations** (Invoke/Close): a per-session-id marker /// that prevents concurrent SMC entry by another core for the same id. /// -/// For multi-instance OpenSession the token holds nothing (each session -/// gets its own private instance, so no exclusion is required). The -/// first-ever OpenSession for an unknown UUID is the exception: until the -/// TA is loaded its flags aren't known, so it's conservatively serialized -/// under the per-UUID lock until flags are observed. +/// For known multi-instance OpenSession the token holds nothing (each +/// session gets its own private instance, so no exclusion is required). /// -/// On drop, the per-UUID lock is released first, then the per-session-id -/// marker. +/// On drop, the held UUID-level lock is released first (whether per-UUID +/// or the shared unknown lock), then the per-session-id marker. struct SessionToken<'a> { manager: &'a SessionManager, /// Held `Arc` of the per-UUID `SpinMutex`. The guard returned by @@ -415,9 +406,22 @@ pub struct SessionManager { /// pass the limit before either registers. pending_count: SpinMutex, /// Cached TA flags by UUID, populated on first successful session registration. + /// + /// TODO: a TA's flags (in particular single- vs multi-instance) can + /// change across a version update of the same UUID. Key this map by + /// `(uuid, version)` — or invalidate on version mismatch — once TA + /// versioning is wired through, so a re-loaded TA isn't serialized + /// under the old flags. known_flags: SpinMutex>, /// Per-UUID serialization locks for single-instance TA handling. + /// Entries are created lazily only for UUIDs that have been observed + /// to be single-instance — never for unknown UUIDs whose load might + /// fail or turn out to be multi-instance. single_instance_locks: SpinMutex>>>, + /// Shared serialization lock for first-ever loads of unknown UUIDs. + /// Held by OpenSession while flags are still unobserved, then released + /// once `known_flags` is updated. Per-UUID locks take over from there. + unknown_uuid_lock: Arc>, /// Session ids currently being handled (Invoke/Close). Guards a session /// against concurrent SMC entry by another core that targets the same id. active_sessions: SpinMutex>, @@ -431,6 +435,7 @@ impl SessionManager { pending_count: SpinMutex::new(0), known_flags: SpinMutex::new(HashMap::new()), single_instance_locks: SpinMutex::new(HashMap::new()), + unknown_uuid_lock: Arc::new(SpinMutex::new(())), active_sessions: SpinMutex::new(HashSet::new()), } } @@ -459,17 +464,8 @@ impl SessionManager { self.known_flags.lock().get(uuid).copied() } - /// Whether `uuid` should be treated as single-instance for serialization. - /// - /// Returns the cached `is_single_instance()` if known, or `true` for the - /// first-ever load (we have not yet observed the TA's flags) to preserve - /// safety invariants conservatively. - fn assume_single_instance(&self, uuid: &TeeUuid) -> bool { - self.get_known_flags(uuid) - .is_none_or(|f| f.is_single_instance()) - } - - /// Get or create the per-UUID serialization mutex `Arc`. + /// Get or create the per-UUID serialization mutex `Arc`. Only called + /// for UUIDs already observed to be single-instance. fn uuid_lock_arc(&self, uuid: TeeUuid) -> Arc> { self.single_instance_locks .lock() @@ -491,24 +487,39 @@ impl SessionManager { Some(lock) } + /// Try to take the shared `unknown_uuid_lock` non-blockingly using the + /// same forget/`force_unlock` pattern as [`Self::try_acquire_uuid_lock`]. + fn try_acquire_unknown_uuid_lock(&self) -> Option>> { + let lock = self.unknown_uuid_lock.clone(); + let guard = lock.try_lock()?; + core::mem::forget(guard); + Some(lock) + } + /// Acquire a `SessionToken` for an OpenSession request. /// - /// For single-instance TAs (including first-ever load of an unknown - /// UUID) this takes the per-UUID `SpinMutex` non-blockingly. For - /// already-known multi-instance TAs the returned token holds no locks — - /// each session creates its own private instance, so no exclusion is - /// required. + /// Dispatches by what's known about `uuid`: /// - /// Returns `Err(EThreadLimit)` if another core is currently inside an - /// operation on the same single-instance UUID. + /// - **Known single-instance**: per-UUID `SpinMutex`. + /// - **Known multi-instance**: no lock (each session is independent). + /// - **Unknown**: the shared `unknown_uuid_lock`. This serializes + /// first-loads of all unknown UUIDs together, but avoids minting a + /// per-UUID lock entry until the TA has been confirmed single-instance. + /// A failed or multi-instance load therefore leaves no stale entry in + /// `single_instance_locks`. + /// + /// Returns `Err(EThreadLimit)` on contention. fn try_acquire_for_open(&self, uuid: TeeUuid) -> Result, OpteeSmcReturnCode> { - let uuid_lock = if self.assume_single_instance(&uuid) { - Some( + let uuid_lock = match self.get_known_flags(&uuid) { + Some(flags) if flags.is_single_instance() => Some( self.try_acquire_uuid_lock(uuid) .ok_or(OpteeSmcReturnCode::EThreadLimit)?, - ) - } else { - None + ), + Some(_) => None, + None => Some( + self.try_acquire_unknown_uuid_lock() + .ok_or(OpteeSmcReturnCode::EThreadLimit)?, + ), }; Ok(SessionToken { manager: self, @@ -731,8 +742,10 @@ impl SessionManager { /// Drive an OpenSession to completion under the right serialization. /// - /// Internally acquires the per-UUID `SpinMutex` for single-instance TAs - /// (or no lock for known multi-instance TAs), then either: + /// Internally acquires the UUID-level lock dictated by what's known + /// about `uuid` (per-UUID lock for known single-instance, shared + /// unknown-load lock for unknown UUIDs, none for known multi-instance), + /// then either: /// /// - Calls `f(Some(existing))` if a cached single-instance TA is found /// for `uuid`. The lock is held throughout the call so the existing @@ -742,26 +755,27 @@ impl SessionManager { /// and register a new instance. The slot is released when `f` /// returns, regardless of outcome. /// - /// The per-UUID lock is released when this function returns; `f` runs - /// under it. For multi-instance TAs each session gets its own - /// independent `TaInstance`, so no per-UUID exclusion is required. + /// The lock is released when this function returns; `f` runs under it. + /// For known multi-instance TAs each session gets its own independent + /// `TaInstance`, so no per-UUID exclusion is required. /// /// `pending_count` exists only for capacity accounting (so two /// multi-instance loads can't both pass the limit check before either - /// registers). Duplicate-prevention for single-instance TAs is provided - /// by the per-UUID lock above — it serializes the cache check and any - /// new load for the same UUID, so two concurrent loads cannot both miss - /// the cache and create rival instances. + /// registers). Duplicate-prevention for the single-instance / unknown + /// paths is provided by the UUID-level lock above — it serializes the + /// cache check and any new load, so two concurrent loads cannot both + /// miss the cache and create rival instances. pub fn with_ta(&self, uuid: &TeeUuid, f: F) -> Result<(), OpteeSmcReturnCode> where F: for<'a> FnOnce(Option<&'a TaInstance>) -> Result<(), OpteeSmcReturnCode>, { - let token = self.try_acquire_for_open(*uuid)?; - let is_single_instance = token.uuid_lock.is_some(); + let _token = self.try_acquire_for_open(*uuid)?; - // For single-instance TAs the per-UUID lock above keeps our UUID's - // cache entry stable. For multi-instance we don't consult the cache. - if is_single_instance && let Some(existing) = self.single_instance_cache.get(uuid) { + // Cache lookup is unconditional: it returns `None` for known + // multi-instance and unknown UUIDs (never populated), and only + // returns `Some` for known single-instance UUIDs whose entry the + // per-UUID lock above keeps stable. + if let Some(existing) = self.single_instance_cache.get(uuid) { return f(Some(&existing)); } @@ -781,8 +795,6 @@ impl SessionManager { *pending = pending.saturating_sub(1); } - // `single_instance_locks` entries are never evicted; safe removal - // would need generational tracking. Leak is bounded by distinct UUIDs. result } } @@ -902,16 +914,19 @@ mod tests { manager.with_session(109, |_| Ok(())).unwrap(); } - /// `single_instance_locks` entries are never evicted, even on failure - /// of a previously-unknown UUID. + /// A failed first-load of an unknown UUID must not mint a per-UUID + /// lock entry. Unknown loads serialize on `unknown_uuid_lock`, so + /// `single_instance_locks` stays empty when the load fails or the TA + /// turns out to be multi-instance. #[test] - fn with_ta_never_evicts_lock_entry() { + fn with_ta_does_not_mint_lock_entry_for_failed_unknown_load() { let manager = SessionManager::new(); let uuid = make_uuid(0xA9); assert!(manager.get_known_flags(&uuid).is_none()); let _ = manager.with_ta(&uuid, |_| Err(OpteeSmcReturnCode::ENotAvail)); - assert!(manager.single_instance_locks.lock().get(&uuid).is_some()); + assert!(manager.single_instance_locks.lock().get(&uuid).is_none()); + assert!(manager.get_known_flags(&uuid).is_none()); } /// `try_acquire_for_session` returns the entry observed by the From ac6c31eb4a14f35335f9e6b21b78ecf5ecab236a Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 28 May 2026 17:41:22 +0000 Subject: [PATCH 19/28] improve SessionToken API and cover no multi-session case --- litebox_runner_lvbs/src/lib.rs | 65 +-- .../src/lib.rs | 7 +- .../src/tests.rs | 6 +- litebox_shim_optee/src/lib.rs | 2 +- litebox_shim_optee/src/session.rs | 523 ++++++++++++------ 5 files changed, 380 insertions(+), 223 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 8a60fa2c5b..1f5807f29b 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -41,9 +41,7 @@ use litebox_platform_multiplex::Platform; 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::{ - SessionIdGuard, SessionManager, TaInstance, allocate_session_id, -}; +use litebox_shim_optee::session::{OpenSessionTarget, SessionManager, TaInstance}; use litebox_shim_optee::{NormalWorldConstPtr, NormalWorldMutPtr, UserConstPtr}; use once_cell::race::OnceBox; @@ -515,15 +513,15 @@ fn handle_open_session( let client_identity = ta_req_info.client_identity; let params = &ta_req_info.params; - session_manager().with_ta(&ta_uuid, |existing| match existing { - Some(instance) => open_session_single_instance( + session_manager().with_ta(&ta_uuid, |target| match target { + OpenSessionTarget::Sibling(instance) => open_session_single_instance( msg_args, msg_args_phys_addr, instance, params, &ta_req_info, ), - None => open_session_new_instance( + OpenSessionTarget::NewInstance => open_session_new_instance( msg_args, msg_args_phys_addr, params, @@ -531,6 +529,15 @@ fn handle_open_session( client_identity, &ta_req_info, ), + OpenSessionTarget::Busy => { + // Single-instance TA without MULTI_SESSION already has a live + // session. Per OP-TEE OS `tee_ta_init_session_with_context`, + // return TEE_ERROR_BUSY with origin TEE via msg_args. + msg_args.ret = TeeResult::Busy; + msg_args.ret_origin = TeeOrigin::Tee; + write_non_ta_msg_args_to_normal_world(msg_args, msg_args_phys_addr)?; + Ok(()) + } }) } @@ -548,14 +555,10 @@ fn open_session_single_instance( ) -> Result<(), OpteeSmcReturnCode> { let task_pt_id = instance.task_page_table_id(); let ta_uuid = instance.uuid(); + let ta_flags = instance.loaded_program().ta_flags; - // Allocate session ID BEFORE calling load_ta_context so TA gets correct ID. - // Use SessionIdGuard to ensure the ID is recycled on any error path - // (before it is registered with the session manager). - let session_id_guard = - SessionIdGuard::new(allocate_session_id().ok_or(OpteeSmcReturnCode::EBusy)?); - // Safe to unwrap: guard was just created with Some(id). - let runner_session_id = session_id_guard.id().unwrap(); + let mut session_token = session_manager().try_acquire_open_session_token()?; + let runner_session_id = session_token.session_id().unwrap(); debug_serial_println!( "Reusing single-instance TA: uuid={:?}, task_pt_id={}, session_id={}", @@ -564,8 +567,6 @@ fn open_session_single_instance( runner_session_id ); - let ta_flags = instance.loaded_program().ta_flags; - let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; // Set up the entry-point parameters for OpenSession. @@ -647,7 +648,6 @@ fn open_session_single_instance( } // Treat write-back failure as OpenSession failure: do not publish the session. - let runner_session_id = session_id_guard.id().unwrap(); let write_result = write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, @@ -661,9 +661,9 @@ fn open_session_single_instance( // deliver the session id to the normal world, so it will never issue a // matching CloseSession. For a non-keep-alive instance with no siblings // we tear the whole instance down, reclaiming the TA-side state, and the - // session id can be recycled normally. For keep-alive or shared + // session id is recycled by the token's drop. For keep-alive or shared // instances the TA still holds session-local state tagged with this id, - // so we forget the id (disarm the guard) to prevent a future OpenSession + // so we forget the id (disarm the token) to prevent a future OpenSession // from reusing it and colliding with the orphaned TA-side bookkeeping. if let Err(e) = write_result { if !ta_flags.is_keep_alive() && session_manager().count_sessions_for_instance(instance) == 0 @@ -674,14 +674,14 @@ fn open_session_single_instance( teardown_ta_page_table(instance.shim(), task_pt_id); }; } else { - let _ = session_id_guard.disarm(); + session_token.disarm(); } return Err(e); } // Success: register a sibling session pointing at the existing instance. session_manager().register_sibling_session(runner_session_id, instance)?; - session_id_guard.disarm(); + session_token.disarm(); debug_serial_println!( "OpenSession complete on single-instance TA: session_id={}", @@ -706,23 +706,20 @@ fn open_session_new_instance( ) -> Result<(), OpteeSmcReturnCode> { let ta_bin = find_ta_binary(ta_uuid).ok_or(OpteeSmcReturnCode::ENotAvail)?; - // Create and switch to new page table + // Token is declared before `task_pt_guard` so it drops AFTER it — + // marker only releases once CR3 is back to base. See + // `try_acquire_open_session_token` for why. + let mut session_token = session_manager().try_acquire_open_session_token()?; + let runner_session_id = session_token.session_id().unwrap(); + let task_pt_id = create_task_page_table()?; debug_serial_println!("Created task page table ID: {}", task_pt_id); - let task_pt_guard = TaskPageTableGuard::enter(task_pt_id).inspect_err(|_| { + let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id).inspect_err(|_| { // Safety: switch_to_task_page_table failed, so task page table is not active. let _ = unsafe { delete_task_page_table(task_pt_id) }; })?; - let Some(session_id) = allocate_session_id() else { - drop(task_pt_guard); - let _ = unsafe { delete_task_page_table(task_pt_id) }; - return Err(OpteeSmcReturnCode::EBusy); - }; - let session_id_guard = SessionIdGuard::new(session_id); - let runner_session_id = session_id_guard.id().unwrap(); - // 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( @@ -860,10 +857,8 @@ fn open_session_new_instance( } // Write back BEFORE publishing the instance. If the write fails, the - // session is neither registered nor cached, so we just tear down the - // local resources and let `session_id_guard` recycle the ID on drop. - // Safe to unwrap: guard has not been disarmed yet. - let runner_session_id = session_id_guard.id().unwrap(); + // session is neither registered nor cached; we tear down the local + // resources and `session_token`'s drop recycles the id. write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, @@ -885,7 +880,7 @@ fn open_session_new_instance( task_pt_id, ta_uuid, ); - session_id_guard.disarm(); + session_token.disarm(); debug_serial_println!( "OpenSession complete: session_id={}, single_instance={}", diff --git a/litebox_runner_optee_on_linux_userland/src/lib.rs b/litebox_runner_optee_on_linux_userland/src/lib.rs index 7e45993767..ae37334cd4 100644 --- a/litebox_runner_optee_on_linux_userland/src/lib.rs +++ b/litebox_runner_optee_on_linux_userland/src/lib.rs @@ -5,7 +5,7 @@ use anyhow::{Context as _, Result}; use clap::Parser; use litebox_common_optee::{TeeUuid, UteeEntryFunc, UteeParamOwned}; use litebox_platform_multiplex::Platform; -use litebox_shim_optee::session::allocate_session_id; +use litebox_shim_optee::session::SessionManager; use std::path::PathBuf; mod tests; @@ -109,17 +109,20 @@ fn run_ta_with_default_commands( ldelf_bin: &[u8], ta_bin: &[u8], ) { + let session_manager = SessionManager::new(); 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), None, - allocate_session_id().unwrap(), + session_id, ) .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 4bfe1f3da2..9bfe30f8ec 100644 --- a/litebox_runner_optee_on_linux_userland/src/tests.rs +++ b/litebox_runner_optee_on_linux_userland/src/tests.rs @@ -8,7 +8,7 @@ use litebox::platform::RawConstPointer; use litebox::utils::TruncateExt; use litebox_common_optee::{TeeParamType, UteeEntryFunc, UteeParamOwned, UteeParams}; -use litebox_shim_optee::session::allocate_session_id; +use litebox_shim_optee::session::SessionManager; use litebox_shim_optee::{LoadedProgram, UserConstPtr}; use serde::Deserialize; use std::path::PathBuf; @@ -26,6 +26,7 @@ pub fn run_ta_with_test_commands( serde_json::from_str(&json_str).unwrap() }; let mut ta_info: Option = None; + let session_manager = SessionManager::new(); for cmd in ta_commands { assert!( @@ -49,13 +50,14 @@ pub fn run_ta_with_test_commands( 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 session_token = session_manager.try_acquire_open_session_token().unwrap(); let loaded = shim .load_ldelf( ldelf_bin, ta_head.uuid, Some(ta_bin), None, - allocate_session_id().unwrap(), + session_token.session_id().unwrap(), ) .map_err(|_| { panic!("Failed to load TA"); diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index 6f68273c73..2afba0c79b 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -38,7 +38,7 @@ pub mod msg_handler; pub mod ptr; // Re-export session management types for convenience -pub use session::{SessionManager, TaInstance, allocate_session_id}; +pub use session::{OpenSessionTarget, SessionManager, SessionToken, TaInstance}; const MAX_KERNEL_BUF_SIZE: usize = 0x80_000; diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index b047cf707f..7a303a807b 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -159,6 +159,25 @@ impl TaInstance { unsafe impl Send for TaInstance {} unsafe impl Sync for TaInstance {} +/// What an OpenSession should do given the current cache state for a +/// `uuid`, as decided by [`SessionManager::with_ta`] under its +/// serialization. The closure dispatches on the variant. +pub enum OpenSessionTarget<'a> { + /// No cached single-instance instance for this UUID (either it's + /// not single-instance, or the cache is empty). Closure should load + /// a fresh TA and call `register_new_session`. + NewInstance, + /// A cached single-instance TA is available for sharing. Closure + /// should reuse it for a sibling session via `register_sibling_session`. + Sibling(&'a TaInstance), + /// A cached single-instance TA exists but it lacks `TA_FLAG_MULTI_SESSION` + /// and already has at least one live session. Per OP-TEE OS + /// `tee_ta_init_session_with_context`, the runner must reject with + /// `TEE_ERROR_BUSY` (origin TEE). The closure must write the BUSY + /// response to the client and return `Ok` — no TA load needed. + Busy, +} + /// Per-session entry in the session map. The `Dead` variant retains /// `(ta_uuid, ta_flags)` so cleanup paths and `try_acquire_for_session`'s /// snapshot still have them after the instance is gone. @@ -297,7 +316,7 @@ impl Default for SingleInstanceCache { } /// Returns `None` if all session IDs are exhausted. -pub fn allocate_session_id() -> Option { +fn allocate_session_id() -> Option { SessionIdPool::allocate() } @@ -305,47 +324,6 @@ fn recycle_session_id(session_id: u32) { SessionIdPool::recycle(session_id); } -/// RAII guard that recycles a session ID on drop unless disarmed. -/// -/// Session IDs are allocated before the TA is invoked and only registered on -/// success via [`SessionManager::register_new_session`] or -/// [`SessionManager::register_sibling_session`]. This guard ensures it is -/// recycled on all error paths before this registration. -pub struct SessionIdGuard { - session_id: Option, -} - -impl SessionIdGuard { - pub fn new(session_id: u32) -> Self { - Self { - session_id: Some(session_id), - } - } - - /// Returns `None` if already disarmed. - pub fn id(&self) -> Option { - self.session_id - } - - /// Disarm the guard so the session ID is **not** recycled on drop. - /// - /// Call this after the session ID has been successfully registered. - /// Once registered, [`SessionManager::unregister_session`] owns recycling. - /// - /// Returns `None` if the guard was already disarmed. - pub fn disarm(mut self) -> Option { - self.session_id.take() - } -} - -impl Drop for SessionIdGuard { - fn drop(&mut self) { - if let Some(id) = self.session_id { - recycle_session_id(id); - } - } -} - /// RAII token bundling the serialization primitives required to safely /// execute an OP-TEE TA operation. /// @@ -355,23 +333,59 @@ impl Drop for SessionIdGuard { /// all sessions on the same TA. /// - **First-ever load of an unknown UUID** (OpenSession only): the shared /// `unknown_uuid_lock`, used until the TA's flags are observed. -/// - **Existing-session operations** (Invoke/Close): a per-session-id marker -/// that prevents concurrent SMC entry by another core for the same id. +/// - **Existing-session operations** (Invoke/Close): a per-session-id +/// marker (slot in `SessionManager::active_sessions`) that prevents +/// concurrent SMC entry by another core for the same id. +/// - **OpenSession (runner-facing)**: same per-session-id marker plus +/// a freshly-allocated `session_id` whose recycling the token owns +/// until [`Self::disarm`]. Acquired via +/// [`SessionManager::try_acquire_open_session_token`]. /// -/// For known multi-instance OpenSession the token holds nothing (each -/// session gets its own private instance, so no exclusion is required). +/// For known multi-instance OpenSession the token (from `with_ta`) +/// holds nothing — each session gets its own private instance, so no +/// exclusion is required there. /// -/// On drop, the held UUID-level lock is released first (whether per-UUID -/// or the shared unknown lock), then the per-session-id marker. -struct SessionToken<'a> { +/// On drop the held UUID-level lock is released first (whether per-UUID +/// or the shared unknown lock), then the per-session-id marker, then +/// (if still owned) the session id is recycled. +pub struct SessionToken<'a> { manager: &'a SessionManager, /// Held `Arc` of the per-UUID `SpinMutex`. The guard returned by /// `try_lock()` was [`core::mem::forget`]-ed at acquisition time; this /// type's `Drop` calls `force_unlock` to release the mutex. The `Arc` /// keeps the mutex alive across acquisition and release. uuid_lock: Option>>, - /// Session id reserved in [`SessionManager::active_sessions`]. + /// `Some(id)` while the token holds the active-session marker for `id` + /// in [`SessionManager::active_sessions`]. Drop releases the marker. active_session_id: Option, + /// Whether `active_session_id` should also be recycled to the id pool + /// on drop (in addition to releasing the marker). Set when the id was + /// freshly allocated by + /// [`SessionManager::try_acquire_open_session_token`]; cleared by + /// [`Self::disarm`] after the id is transferred to the session map via + /// `register_*_session`. Only meaningful when `active_session_id` is + /// `Some`; ignored otherwise. + owns_id_recycling: bool, +} + +impl SessionToken<'_> { + /// Session id this token reserves the active-session marker for, if any. + /// Set for tokens minted by + /// [`SessionManager::try_acquire_open_session_token`], + /// [`SessionManager::try_acquire_session_marker`], or + /// `try_acquire_for_session` (Invoke/Close). + pub fn session_id(&self) -> Option { + self.active_session_id + } + + /// Transfer id-recycling responsibility off the token. Call after the + /// id has been registered via `register_new_session` / + /// `register_sibling_session`; from that point the session map (via + /// `unregister_session`) owns recycling, and the token's drop will + /// only release the marker (and any locks). + pub fn disarm(&mut self) { + self.owns_id_recycling = false; + } } impl Drop for SessionToken<'_> { @@ -384,6 +398,9 @@ impl Drop for SessionToken<'_> { } if let Some(id) = self.active_session_id.take() { self.manager.active_sessions.lock().remove(&id); + if self.owns_id_recycling { + recycle_session_id(id); + } } } } @@ -440,6 +457,62 @@ impl SessionManager { } } + /// Reserve the active-session slot for `session_id` non-blockingly, + /// returning a [`SessionToken`] (carrying just the per-session-id + /// marker, no UUID lock) that releases the slot on drop. Returns + /// `None` if the slot is already taken. Excludes concurrent + /// `with_session(session_id)` until the token drops. + pub fn try_acquire_session_marker(&self, session_id: u32) -> Option> { + if !self.active_sessions.lock().insert(session_id) { + return None; + } + Some(SessionToken { + manager: self, + uuid_lock: None, + active_session_id: Some(session_id), + owns_id_recycling: false, + }) + } + + /// Allocate a fresh `session_id` and reserve its active-session slot + /// atomically. See [`SessionToken`] for what the returned token + /// carries. + /// + /// # Drop-order requirement + /// + /// On the OpenSession path the runner activates a TA page table + /// (`TaskPageTableGuard`) inside the same scope. The token *must* be + /// declared **before** that guard so it drops **after** it — the + /// marker must outlive the CR3 switch back to base, otherwise a + /// forged Close on the freshly-registered session can win the + /// marker race and tear down the task page table while CR3 still + /// points at it. (Single-instance is already covered by `with_ta`'s + /// per-UUID lock; this is the only defense for multi-instance.) + /// + /// # Errors + /// - `EBusy` if the id pool is exhausted. + /// - `EThreadLimit` for the rare race where the just-allocated id's + /// marker slot is still held by a not-yet-dropped Close token (the + /// driver retries transparently via its wait queue). + pub fn try_acquire_open_session_token(&self) -> Result, OpteeSmcReturnCode> { + let session_id = allocate_session_id().ok_or(OpteeSmcReturnCode::EBusy)?; + if !self.active_sessions.lock().insert(session_id) { + // Rare race: previous Close on this id recycled it via + // `unregister_session` but its `SessionToken` hasn't dropped + // yet, so the marker slot is still held. Roll back the alloc; + // the driver's retry will allocate again (likely a different + // id, or this one after the previous Close fully releases). + recycle_session_id(session_id); + return Err(OpteeSmcReturnCode::EThreadLimit); + } + Ok(SessionToken { + manager: self, + uuid_lock: None, + active_session_id: Some(session_id), + owns_id_recycling: true, + }) + } + /// Mark every session currently pointing at `instance` as `Dead`. Must /// be paired with [`SessionManager::evict_cached_instance`] /// in the documented order — see that function for the rationale. @@ -525,6 +598,7 @@ impl SessionManager { manager: self, uuid_lock, active_session_id: None, + owns_id_recycling: false, }) } @@ -575,6 +649,7 @@ impl SessionManager { manager: self, uuid_lock: None, active_session_id: Some(session_id), + owns_id_recycling: false, }; // Take the per-UUID lock BEFORE the final re-read for single- @@ -602,21 +677,17 @@ impl SessionManager { Ok((token, entry_now)) } - /// Drive an Invoke/Close to completion under the right serialization. - /// - /// Internally acquires the per-session-id marker (and, for single- - /// instance TAs, the per-UUID lock), passes `Some(&TaInstance)` to `f` - /// for live sessions or `None` for dead ones, and releases the locks - /// when `f` returns. `f` runs entirely under the token: state - /// mutations it performs on the session manager (e.g. - /// `unregister_session`, `mark_sessions_dead_for_instance`, - /// `evict_cached_instance`) are serialized against other - /// cores' Invoke/Close on the same session and (for single-instance) - /// the same UUID. + /// Drive an Invoke/Close to completion under the right serialization + /// (see [`SessionToken`] for the locks held). Passes + /// `Some(&TaInstance)` to `f` for live sessions, `None` for dead + /// ones. State mutations `f` performs on the manager + /// (`unregister_session`, `mark_sessions_dead_for_instance`, + /// `evict_cached_instance`) are serialized against concurrent + /// Invoke/Close on the same session and (single-instance) the same UUID. /// /// Returns `Err(EBadCmd)` if `session_id` is not registered, or - /// `Err(EThreadLimit)` on lock contention; the Linux OP-TEE driver - /// retries `EThreadLimit` transparently. + /// `Err(EThreadLimit)` on lock contention (driver retries + /// transparently). pub fn with_session(&self, session_id: u32, f: F) -> Result<(), OpteeSmcReturnCode> where F: for<'a> FnOnce(Option<&'a TaInstance>) -> Result<(), OpteeSmcReturnCode>, @@ -633,6 +704,34 @@ impl SessionManager { /// `loaded_program`, `task_page_table_id`) are taken by value and stored /// inside the manager; for single-instance TAs the instance is also /// cached under `ta_uuid` for later reuse. + /// + /// # Publication order + /// + /// `sessions` and (for single-instance) `single_instance_cache` are + /// populated *before* `known_flags`. Other openers gate on + /// `known_flags` to decide their lock path — once they observe `uuid` + /// as known single-instance, the cache is guaranteed to already + /// contain the entry, so they take the sibling/cache-hit branch + /// rather than racing into a duplicate load. + /// + /// # Unknown→per-UUID transition + /// + /// For single-instance TAs we pre-lock the per-UUID `SpinMutex` *before* + /// publishing `known_flags` so any later opener that observes `uuid` + /// as known single-instance and routes to the per-UUID lock finds it + /// already held. The lock state lives in `single_instance_locks` + /// (the `Arc` we get back is discarded — its only purpose was to + /// take the lock); [`Self::with_ta`] adopts by re-fetching the `Arc` + /// for *its own* `uuid` and installing it in its token. This is + /// UUID-keyed end-to-end: no shared side channel, so concurrent + /// `with_ta` calls for different UUIDs cannot interfere with each + /// other's adoptions. + /// + /// `try_acquire_uuid_lock` succeeds only on the unknown path (caller + /// holds `unknown_uuid_lock`, no sessions or `known_flags` entry for + /// `uuid` yet). On the known-cache-evicted path the caller already + /// holds the per-UUID lock and the `try_lock` returns `None`, so + /// nothing changes (the caller's existing lock is sufficient). pub fn register_new_session( &self, session_id: u32, @@ -648,11 +747,21 @@ impl SessionManager { task_page_table_id, ta_uuid, }); - self.known_flags.lock().entry(ta_uuid).or_insert(ta_flags); + + // Pre-lock per-UUID for atomic unknown→per-UUID transition (see + // method doc). The returned `Arc` is intentionally dropped; the + // forgotten guard inside `try_acquire_uuid_lock` keeps the lock + // state held in `single_instance_locks` until `with_ta` adopts. + if ta_flags.is_single_instance() { + let _ = self.try_acquire_uuid_lock(ta_uuid); + } + self.sessions.insert_live(session_id, arc.clone()); if ta_flags.is_single_instance() { self.single_instance_cache.insert(ta_uuid, arc); } + // Publish `known_flags` last — this is the gate other openers check. + self.known_flags.lock().entry(ta_uuid).or_insert(ta_flags); } /// Register a session that re-uses an existing single-instance TA. @@ -677,10 +786,8 @@ impl SessionManager { .get(&instance.ta_uuid) .filter(|cached| cached.task_page_table_id == instance.task_page_table_id) .ok_or(OpteeSmcReturnCode::EBadCmd)?; - self.known_flags - .lock() - .entry(instance.ta_uuid) - .or_insert(instance.loaded_program.ta_flags); + // `known_flags` is already populated for this UUID — sibling path + // implies the instance was previously registered. self.sessions.insert_live(session_id, arc); Ok(()) } @@ -742,41 +849,54 @@ impl SessionManager { /// Drive an OpenSession to completion under the right serialization. /// - /// Internally acquires the UUID-level lock dictated by what's known - /// about `uuid` (per-UUID lock for known single-instance, shared - /// unknown-load lock for unknown UUIDs, none for known multi-instance), - /// then either: - /// - /// - Calls `f(Some(existing))` if a cached single-instance TA is found - /// for `uuid`. The lock is held throughout the call so the existing - /// instance cannot be torn down or replaced beneath `f`. - /// - Reserves a creation slot (atomic capacity check against - /// `instance_count() + pending_count`) and calls `f(None)` to load - /// and register a new instance. The slot is released when `f` - /// returns, regardless of outcome. + /// Acquires the UUID-level lock for `uuid` (see [`SessionToken`] for + /// the case breakdown), classifies the cache state, and dispatches + /// via [`OpenSessionTarget`]: /// - /// The lock is released when this function returns; `f` runs under it. - /// For known multi-instance TAs each session gets its own independent - /// `TaInstance`, so no per-UUID exclusion is required. + /// - [`OpenSessionTarget::Sibling`] for a cached single-instance TA + /// that admits another session. + /// - [`OpenSessionTarget::Busy`] for the OP-TEE-OS-defined + /// `TA_FLAG_MULTI_SESSION` violation (single-instance without + /// MULTI_SESSION already has a live session). + /// - [`OpenSessionTarget::NewInstance`] otherwise: reserves a + /// creation slot (capacity check against + /// `instance_count() + pending_count`) and lets the closure load + /// and register a fresh instance. /// - /// `pending_count` exists only for capacity accounting (so two - /// multi-instance loads can't both pass the limit check before either - /// registers). Duplicate-prevention for the single-instance / unknown - /// paths is provided by the UUID-level lock above — it serializes the - /// cache check and any new load, so two concurrent loads cannot both - /// miss the cache and create rival instances. + /// `pending_count` exists only for capacity accounting so two + /// concurrent multi-instance loads can't both pass the limit before + /// either registers. The single-instance / unknown paths are + /// serialized by the UUID-level lock itself. pub fn with_ta(&self, uuid: &TeeUuid, f: F) -> Result<(), OpteeSmcReturnCode> where - F: for<'a> FnOnce(Option<&'a TaInstance>) -> Result<(), OpteeSmcReturnCode>, + F: for<'a> FnOnce(OpenSessionTarget<'a>) -> Result<(), OpteeSmcReturnCode>, { - let _token = self.try_acquire_for_open(*uuid)?; + let mut token = self.try_acquire_for_open(*uuid)?; + // Whether we started on the unknown-load path is determined by the + // identity of the lock the token carries. Captured before `f` runs + // so we know which branch to take in the post-`f` adoption step. + let on_unknown_path = token + .uuid_lock + .as_ref() + .is_some_and(|arc| Arc::ptr_eq(arc, &self.unknown_uuid_lock)); // Cache lookup is unconditional: it returns `None` for known // multi-instance and unknown UUIDs (never populated), and only // returns `Some` for known single-instance UUIDs whose entry the // per-UUID lock above keeps stable. if let Some(existing) = self.single_instance_cache.get(uuid) { - return f(Some(&existing)); + // MULTI_SESSION enforcement (matches OP-TEE OS + // `tee_ta_init_session_with_context`). Under the per-UUID lock + // the session count is stable across this check and the + // closure, so a parallel Close/Invoke can't change it. + let flags = existing.loaded_program().ta_flags; + let target = + if !flags.is_multi_session() && self.count_sessions_for_instance(&existing) > 0 { + OpenSessionTarget::Busy + } else { + OpenSessionTarget::Sibling(&existing) + }; + return f(target); } { @@ -788,13 +908,30 @@ impl SessionManager { *pending += 1; } - let result = f(None); + let result = f(OpenSessionTarget::NewInstance); { let mut pending = self.pending_count.lock(); *pending = pending.saturating_sub(1); } + // Complete the unknown→per-UUID transition (see + // `register_new_session` doc). Only fires when we held + // `unknown_uuid_lock` AND the closure registered a single-instance + // TA for *our* `uuid`. The per-UUID lock is already held (forgotten + // guard in `single_instance_locks` from `register_new_session`'s + // pre-lock); re-fetch the `Arc` by `uuid` and swap into the token, + // force-unlocking the old (unknown) lock. Token drop then releases + // the per-UUID lock at the end of `with_ta`. UUID-keyed throughout, + // so concurrent `with_ta(other_uuid)` cannot adopt our lock. + if result.is_ok() && on_unknown_path && self.single_instance_cache.get(uuid).is_some() { + let per_uuid = self.uuid_lock_arc(*uuid); + if let Some(old) = token.uuid_lock.replace(per_uuid) { + // SAFETY: see `SessionToken::drop` — same invariant. + unsafe { old.force_unlock() }; + } + } + result } } @@ -831,6 +968,34 @@ mod tests { TaFlags::SINGLE_INSTANCE | TaFlags::MULTI_SESSION } + /// Test helper: call `register_new_session` and release the per-UUID + /// lock the way `with_ta` would, so subsequent operations + /// (Invoke/Close, evict, count, etc.) aren't blocked by a held lock. + fn register_for_test( + manager: &SessionManager, + session_id: u32, + ta_flags: TaFlags, + task_page_table_id: usize, + ta_uuid: TeeUuid, + ) { + manager.register_new_session( + session_id, + make_shim(), + make_loaded_program(ta_flags), + task_page_table_id, + ta_uuid, + ); + // For single-instance, `register_new_session` pre-locked the + // per-UUID mutex via a forgotten guard. Release here to mirror + // `with_ta`'s token-drop release. + if ta_flags.is_single_instance() + && let Some(arc) = manager.single_instance_locks.lock().get(&ta_uuid).cloned() + { + // SAFETY: same invariant as `SessionToken::drop`. + unsafe { arc.force_unlock() }; + } + } + /// Identity is by `task_page_table_id`, not by Arc pointer. After an /// instance is evicted and a fresh one registered under the same UUID, /// the stale handle must not evict the new one. @@ -839,23 +1004,11 @@ mod tests { let manager = SessionManager::new(); let uuid = make_uuid(0xA4); - manager.register_new_session( - 105, - make_shim(), - make_loaded_program(single_instance_flags()), - 10, - uuid, - ); + register_for_test(&manager, 105, single_instance_flags(), 10, uuid); let arc_first = manager.single_instance_cache.get(&uuid).unwrap(); manager.evict_cached_instance(&arc_first); - manager.register_new_session( - 106, - make_shim(), - make_loaded_program(single_instance_flags()), - 11, - uuid, - ); + register_for_test(&manager, 106, single_instance_flags(), 11, uuid); assert!(!manager.evict_cached_instance(&arc_first)); assert!(manager.single_instance_cache.get(&uuid).is_some()); } @@ -867,13 +1020,7 @@ mod tests { fn mark_dead_makes_with_session_observe_none() { let manager = SessionManager::new(); let uuid = make_uuid(0xA6); - manager.register_new_session( - 108, - make_shim(), - make_loaded_program(single_instance_flags()), - 55, - uuid, - ); + register_for_test(&manager, 108, single_instance_flags(), 55, uuid); let arc = manager.single_instance_cache.get(&uuid).unwrap(); assert_eq!(manager.count_sessions_for_instance(&arc), 1); @@ -888,32 +1035,6 @@ mod tests { .unwrap(); } - /// Per-session-id marker excludes re-entry on the same id, but releases - /// when the closure returns. - #[test] - fn with_session_marker_excludes_reentry() { - let manager = SessionManager::new(); - let uuid = make_uuid(0xA7); - manager.register_new_session( - 109, - make_shim(), - make_loaded_program(single_instance_flags()), - 6, - uuid, - ); - - manager - .with_session(109, |_| { - assert_eq!( - manager.with_session(109, |_| Ok(())), - Err(OpteeSmcReturnCode::EThreadLimit) - ); - Ok(()) - }) - .unwrap(); - manager.with_session(109, |_| Ok(())).unwrap(); - } - /// A failed first-load of an unknown UUID must not mint a per-UUID /// lock entry. Unknown loads serialize on `unknown_uuid_lock`, so /// `single_instance_locks` stays empty when the load fails or the TA @@ -929,40 +1050,6 @@ mod tests { assert!(manager.get_known_flags(&uuid).is_none()); } - /// `try_acquire_for_session` returns the entry observed by the - /// post-marker re-read, not a stale handle from before the marker was - /// taken. After a recycle+re-register under a new UUID, the returned - /// entry must reflect the current UUID. (The mismatch-rejection branch - /// itself can only be triggered by a concurrent swap between snapshot - /// and re-read, which isn't reproducible in a single-threaded test; - /// this just verifies the re-read is the source of truth.) - #[test] - fn try_acquire_for_session_returns_current_entry_after_recycle() { - let manager = SessionManager::new(); - let uuid_a = make_uuid(0xB0); - let uuid_b = make_uuid(0xB1); - let session_id = 222; - - manager.register_new_session( - session_id, - make_shim(), - make_loaded_program(single_instance_flags()), - 70, - uuid_a, - ); - manager.unregister_session(session_id); - manager.register_new_session( - session_id, - make_shim(), - make_loaded_program(single_instance_flags()), - 71, - uuid_b, - ); - - let (_, validated) = manager.try_acquire_for_session(session_id).unwrap(); - assert_eq!(validated.ta_uuid(), uuid_b); - } - /// `pending_count` is bumped only on the create path, never on the /// cache-hit path, and is decremented when the closure returns whether /// success or failure — across multiple calls it must return to zero. @@ -974,8 +1061,8 @@ mod tests { // Successful create path. manager - .with_ta(&uuid_multi, |existing| { - assert!(existing.is_none()); + .with_ta(&uuid_multi, |target| { + assert!(matches!(target, OpenSessionTarget::NewInstance)); manager.register_new_session( 301, make_shim(), @@ -993,14 +1080,84 @@ mod tests { assert_eq!(*manager.pending_count.lock(), 0); // Cache-hit path doesn't touch pending_count. - manager.register_new_session( - 302, - make_shim(), - make_loaded_program(single_instance_flags()), - 81, - uuid_single, - ); + register_for_test(&manager, 302, single_instance_flags(), 81, uuid_single); manager.with_ta(&uuid_single, |_| Ok(())).unwrap(); assert_eq!(*manager.pending_count.lock(), 0); } + + /// After `with_ta` completes the unknown→per-UUID transition, the + /// per-UUID lock must be released — a subsequent `try_lock` on the + /// per-UUID `SpinMutex` for the same UUID must succeed. + #[test] + fn with_ta_releases_per_uuid_lock_after_unknown_load() { + let manager = SessionManager::new(); + let uuid = make_uuid(0xD0); + + manager + .with_ta(&uuid, |target| { + assert!(matches!(target, OpenSessionTarget::NewInstance)); + manager.register_new_session( + 401, + make_shim(), + make_loaded_program(single_instance_flags()), + 90, + uuid, + ); + Ok(()) + }) + .unwrap(); + + // Per-UUID lock entry exists (pre-locked + adopted + released). + let arc = manager + .single_instance_locks + .lock() + .get(&uuid) + .cloned() + .expect("entry created by register_new_session"); + // And it's currently unlocked — a fresh try_lock must succeed. + assert!(arc.try_lock().is_some()); + } + + /// A concurrent `with_ta` for an unrelated UUID must NOT adopt or + /// release the per-UUID lock held by another unknown-load opener. + /// Adoption is keyed by the `with_ta` call's own UUID, so an opener + /// for a different UUID leaves the original opener's per-UUID lock + /// untouched. + #[test] + fn unrelated_with_ta_does_not_adopt_other_uuids_lock() { + let manager = SessionManager::new(); + let uuid_locked = make_uuid(0xE1); + let uuid_other = make_uuid(0xE2); + + // Simulate the "lock pre-taken under unknown-load" state by directly + // pre-locking the per-UUID mutex for `uuid_locked`. This mirrors + // what `register_new_session` does mid-unknown-load before + // `with_ta` adopts. + let pre_locked = manager + .try_acquire_uuid_lock(uuid_locked) + .expect("uncontended"); + // Don't release `pre_locked` — emulating the forgotten-guard state. + core::mem::forget(pre_locked); + + // A `with_ta` call for a completely different UUID must not touch + // `uuid_locked`'s lock. The closure registers nothing, but the + // post-`f` adoption logic still runs. + manager.with_ta(&uuid_other, |_| Ok(())).unwrap(); + + // `uuid_locked`'s per-UUID lock must still be held (not stolen). + let arc = manager + .single_instance_locks + .lock() + .get(&uuid_locked) + .cloned() + .expect("entry exists"); + assert!( + arc.try_lock().is_none(), + "uuid_locked's per-UUID lock must remain held by the original opener" + ); + + // Cleanup: release for SpinMutex sanity. + // SAFETY: we forgot the guard above; release here. + unsafe { arc.force_unlock() }; + } } From 4b75c36f08e92c06fcd46320d057b225ccf2de43 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 28 May 2026 17:52:38 +0000 Subject: [PATCH 20/28] tweak docs --- litebox_runner_lvbs/src/lib.rs | 64 +++++++++++++++++++------------ litebox_shim_optee/src/session.rs | 20 ++++++---- 2 files changed, 51 insertions(+), 33 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 1f5807f29b..c633a0b886 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -407,9 +407,12 @@ unsafe fn teardown_ta_page_table(shim: &litebox_shim_optee::OpteeShim, task_pt_i /// For TA requests (OpenSession, InvokeCommand, CloseSession), it uses `decode_ta_request` /// to extract the TA request information and load/run it using `OpteeShim`. /// -/// OpenSession for multi-instance TAs creates a new task page table and a -/// new TA instance and registers it with the session manager. OpenSession -/// for single-instance TAs reuses the cached instance if available, +/// OpenSession for multi-instance TA creates: +/// - A new task page table for memory isolation +/// - A new TA instance with its own state +/// - An entry in the global session map +/// +/// OpenSession for single-instance TA reuses existing TA instance if available, /// otherwise creates a new one. /// /// InvokeCommand looks up the session and switches to its page table. @@ -611,10 +614,9 @@ fn open_session_single_instance( return_code ); - // Write error response BEFORE switching page tables. Accesses user - // memory, which requires the TA's page table to still be active. - // `with_ta`'s serialization prevents another core from tearing down - // the instance underneath us while we copy TA outputs. + // Write error response BEFORE switching page tables (accesses user memory). + // `with_ta`'s serialization keeps the instance alive so another core cannot + // tear down the active page table while this core is copying TA outputs. let write_result = write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, @@ -633,7 +635,8 @@ fn open_session_single_instance( // Mark-then-evict ordering: see SessionManager::evict_cached_instance. session_manager().mark_sessions_dead_for_instance(instance); let _ = session_manager().evict_cached_instance(instance); - // SAFETY: no references to user-space memory will be held after this call. + // 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(instance.shim(), task_pt_id); }; @@ -669,7 +672,8 @@ fn open_session_single_instance( if !ta_flags.is_keep_alive() && session_manager().count_sessions_for_instance(instance) == 0 { let _ = session_manager().evict_cached_instance(instance); - // SAFETY: no references to user-space memory will be held after this call. + // 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(instance.shim(), task_pt_id); }; @@ -731,7 +735,8 @@ fn open_session_new_instance( runner_session_id, ) .map_err(|_| { - // SAFETY: no references to user-space memory will be held after this call. + // 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 })?, @@ -774,7 +779,8 @@ fn open_session_new_instance( Some(ta_req_info), ); - // SAFETY: no references to user-space memory will be held after this call. + // 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) }; write_result?; @@ -783,7 +789,8 @@ fn open_session_new_instance( // Load TA context with parameters for OpenSession - pass actual session_id loaded_program.entrypoints.as_ref().ok_or_else(|| { - // SAFETY: no references to user-space memory will be held after this call. + // 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::EBadCmd })?; @@ -798,13 +805,13 @@ fn open_session_new_instance( None, ) .map_err(|_| { - // SAFETY: no references to user-space memory will be held after this call. + // 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::EBadCmd })?; - // Run the TA's OpenSession entry point using reference-based reenter to - // avoid moving the shim + // Run the TA entry function using reference-based reenter to avoid moving the shim let mut ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::reenter_thread_ref( @@ -815,14 +822,16 @@ fn open_session_new_instance( // Read TA output parameters from the stack buffer let params_address = loaded_program.params_address.ok_or_else(|| { - // SAFETY: no references to user-space memory will be held after this call. + // 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::EBadAddr })?; let ta_params = UserConstPtr::::from_usize(params_address) .read_at_offset(0) .ok_or_else(|| { - // SAFETY: no references to user-space memory will be held after this call. + // 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::EBadAddr })?; @@ -849,7 +858,8 @@ fn open_session_new_instance( Some(ta_req_info), ); - // SAFETY: no references to user-space memory will be held after this call. + // 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) }; write_result?; @@ -857,8 +867,8 @@ fn open_session_new_instance( } // Write back BEFORE publishing the instance. If the write fails, the - // session is neither registered nor cached; we tear down the local - // resources and `session_token`'s drop recycles the id. + // session is neither registered nor cached, so we just tear down the + // local resources and let `session_token` recycle the ID on drop. write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, @@ -868,7 +878,8 @@ fn open_session_new_instance( Some(ta_req_info), ) .inspect_err(|_| { - // SAFETY: no references to user-space memory will be held after this call. + // 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) }; })?; @@ -984,8 +995,9 @@ fn handle_invoke_command( let return_code: u32 = ctx.rax.trunc(); let return_code = TeeResult::try_from(return_code).unwrap_or(TeeResult::GenericError); - // Write response BEFORE switching page tables — accesses user memory, - // which requires the TA's page table to still be active. + // Write response BEFORE switching page tables (accesses user memory). + // `with_session`'s marker keeps the entry stable so another core cannot + // tear down the active page table while this core is copying TA outputs. let write_result = write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, @@ -1012,7 +1024,8 @@ fn handle_invoke_command( session_manager().unregister_session(session_id); - // SAFETY: no references to user-space memory will be held after this call. + // 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(instance.shim(), task_pt_id); }; @@ -1118,7 +1131,8 @@ fn handle_close_session( let _ = session_manager() .evict_cached_instance(instance); } - // SAFETY: no references to user-space memory will be held after this call. + // 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(instance.shim(), task_pt_id); }; diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 7a303a807b..b578629710 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -491,17 +491,21 @@ impl SessionManager { /// /// # Errors /// - `EBusy` if the id pool is exhausted. - /// - `EThreadLimit` for the rare race where the just-allocated id's - /// marker slot is still held by a not-yet-dropped Close token (the - /// driver retries transparently via its wait queue). + /// - `EThreadLimit` only if the just-allocated id's marker slot is + /// already held — unreachable under normal flow (the id pool's + /// hint+wrap defers reuse), and present as a defensive bail-out + /// when a caller has taken the slot out-of-band via + /// [`Self::try_acquire_session_marker`]. pub fn try_acquire_open_session_token(&self) -> Result, OpteeSmcReturnCode> { let session_id = allocate_session_id().ok_or(OpteeSmcReturnCode::EBusy)?; if !self.active_sessions.lock().insert(session_id) { - // Rare race: previous Close on this id recycled it via - // `unregister_session` but its `SessionToken` hasn't dropped - // yet, so the marker slot is still held. Roll back the alloc; - // the driver's retry will allocate again (likely a different - // id, or this one after the previous Close fully releases). + // Defensive: the id pool's hint+wrap guarantees a + // freshly-allocated id won't collide with a recently-recycled + // one, so under normal `allocate → mark → unregister` flow + // this branch is unreachable. It exists to handle the corner + // case where a caller has inserted the slot out-of-band via + // [`Self::try_acquire_session_marker`]. Roll back the alloc + // and let the driver retry. recycle_session_id(session_id); return Err(OpteeSmcReturnCode::EThreadLimit); } From 3e85c6ba5b342e76b576e6f044b2d5db08a674f4 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 28 May 2026 23:50:03 +0000 Subject: [PATCH 21/28] remove dead code. fix some comments --- litebox_runner_lvbs/src/lib.rs | 10 ++-- litebox_shim_optee/src/session.rs | 84 ++++++++++++------------------- 2 files changed, 40 insertions(+), 54 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index c633a0b886..9108df79bf 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -561,6 +561,7 @@ fn open_session_single_instance( let ta_flags = instance.loaded_program().ta_flags; let mut session_token = session_manager().try_acquire_open_session_token()?; + // Safe to unwrap: session ID has been just created. let runner_session_id = session_token.session_id().unwrap(); debug_serial_println!( @@ -570,9 +571,10 @@ fn open_session_single_instance( runner_session_id ); + // Switch to the existing TA's page table let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; - // Set up the entry-point parameters for OpenSession. + // Load TA context with parameters for OpenSession - pass actual session_id instance .loaded_program() .entrypoints @@ -586,6 +588,7 @@ fn open_session_single_instance( ) .map_err(|_| OpteeSmcReturnCode::EBadCmd)?; + // Run the TA's OpenSession entry point using reference-based reenter let mut ctx = litebox_common_linux::PtRegs::default(); unsafe { litebox_platform_lvbs::reenter_thread_ref( @@ -594,6 +597,7 @@ fn open_session_single_instance( ); } + // Read TA output parameters from the stack buffer let params_address = instance .loaded_program() .params_address @@ -710,8 +714,8 @@ fn open_session_new_instance( ) -> Result<(), OpteeSmcReturnCode> { let ta_bin = find_ta_binary(ta_uuid).ok_or(OpteeSmcReturnCode::ENotAvail)?; - // Token is declared before `task_pt_guard` so it drops AFTER it — - // marker only releases once CR3 is back to base. See + // Token is declared before `task_pt_guard` so it drops AFTER it. + // Marker only releases once CR3 is back to base. See // `try_acquire_open_session_token` for why. let mut session_token = session_manager().try_acquire_open_session_token()?; let runner_session_id = session_token.session_id().unwrap(); diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index b578629710..87e26e3d6a 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -114,9 +114,9 @@ const MAX_TA_INSTANCES: usize = 16; /// A loaded TA instance. /// /// For single-instance TAs one instance is shared across all sessions; the -/// TA stays in memory until the last session closes (or, with -/// `TA_FLAG_INSTANCE_KEEP_ALIVE`, until explicit destroy). Each instance -/// has its own task page table that provides memory isolation from other TAs. +/// TA stays in memory until the last session closes (if it does not have the +/// `TA_FLAG_INSTANCE_KEEP_ALIVE` flag). Each instance has its own task page +/// table that provides memory isolation from other TAs. pub struct TaInstance { /// The shim must be kept alive to keep the loaded program's memory mappings valid. shim: OpteeShim, @@ -172,9 +172,8 @@ pub enum OpenSessionTarget<'a> { Sibling(&'a TaInstance), /// A cached single-instance TA exists but it lacks `TA_FLAG_MULTI_SESSION` /// and already has at least one live session. Per OP-TEE OS - /// `tee_ta_init_session_with_context`, the runner must reject with - /// `TEE_ERROR_BUSY` (origin TEE). The closure must write the BUSY - /// response to the client and return `Ok` — no TA load needed. + /// `tee_ta_init_session_with_context`, reject with + /// `TEE_ERROR_BUSY` (origin TEE). Busy, } @@ -211,22 +210,26 @@ struct SessionMap { } impl SessionMap { + /// Create a new empty session map. fn new() -> Self { Self { inner: SpinMutex::new(HashMap::new()), } } + /// Get full session entry by session ID. fn get_entry(&self, session_id: u32) -> Option { self.inner.lock().get(&session_id).cloned() } + /// Insert a live session into the map. fn insert_live(&self, session_id: u32, instance: Arc) { self.inner .lock() .insert(session_id, SessionEntry::Live(instance)); } + /// Remove a session from the map. fn remove(&self, session_id: u32) -> Option { self.inner.lock().remove(&session_id) } @@ -276,16 +279,19 @@ struct SingleInstanceCache { } impl SingleInstanceCache { + /// Create a new empty cache. fn new() -> Self { Self { inner: SpinMutex::new(HashMap::new()), } } + /// Get a cached single-instance TA by UUID. fn get(&self, uuid: &TeeUuid) -> Option> { self.inner.lock().get(uuid).cloned() } + /// Cache a single-instance TA by UUID. fn insert(&self, uuid: TeeUuid, instance: Arc) { self.inner.lock().insert(uuid, instance); } @@ -304,6 +310,7 @@ impl SingleInstanceCache { } } + /// Get the number of cached single-instance TAs. fn len(&self) -> usize { self.inner.lock().len() } @@ -315,11 +322,17 @@ impl Default for SingleInstanceCache { } } +/// Allocate a new unique session ID. +/// +/// Delegates to `SessionIdPool::allocate` for unified session ID management. /// Returns `None` if all session IDs are exhausted. fn allocate_session_id() -> Option { SessionIdPool::allocate() } +/// Recycle a session ID for potential future reuse. +/// +/// Delegates to `SessionIdPool::recycle`. fn recycle_session_id(session_id: u32) { SessionIdPool::recycle(session_id); } @@ -332,7 +345,8 @@ fn recycle_session_id(session_id: u32) { /// - **Known single-instance TAs**: a per-UUID `SpinMutex` that serializes /// all sessions on the same TA. /// - **First-ever load of an unknown UUID** (OpenSession only): the shared -/// `unknown_uuid_lock`, used until the TA's flags are observed. +/// `unknown_uuid_lock`, used until the TA's flags (single-instance vs +/// multi-instance) are observed. /// - **Existing-session operations** (Invoke/Close): a per-session-id /// marker (slot in `SessionManager::active_sessions`) that prevents /// concurrent SMC entry by another core for the same id. @@ -371,8 +385,7 @@ pub struct SessionToken<'a> { impl SessionToken<'_> { /// Session id this token reserves the active-session marker for, if any. /// Set for tokens minted by - /// [`SessionManager::try_acquire_open_session_token`], - /// [`SessionManager::try_acquire_session_marker`], or + /// [`SessionManager::try_acquire_open_session_token`] or /// `try_acquire_for_session` (Invoke/Close). pub fn session_id(&self) -> Option { self.active_session_id @@ -457,26 +470,8 @@ impl SessionManager { } } - /// Reserve the active-session slot for `session_id` non-blockingly, - /// returning a [`SessionToken`] (carrying just the per-session-id - /// marker, no UUID lock) that releases the slot on drop. Returns - /// `None` if the slot is already taken. Excludes concurrent - /// `with_session(session_id)` until the token drops. - pub fn try_acquire_session_marker(&self, session_id: u32) -> Option> { - if !self.active_sessions.lock().insert(session_id) { - return None; - } - Some(SessionToken { - manager: self, - uuid_lock: None, - active_session_id: Some(session_id), - owns_id_recycling: false, - }) - } - - /// Allocate a fresh `session_id` and reserve its active-session slot - /// atomically. See [`SessionToken`] for what the returned token - /// carries. + /// Allocate a fresh `session_id` and reserve its active-session slot. + /// See [`SessionToken`] for what the returned token carries. /// /// # Drop-order requirement /// @@ -491,24 +486,16 @@ impl SessionManager { /// /// # Errors /// - `EBusy` if the id pool is exhausted. - /// - `EThreadLimit` only if the just-allocated id's marker slot is - /// already held — unreachable under normal flow (the id pool's - /// hint+wrap defers reuse), and present as a defensive bail-out - /// when a caller has taken the slot out-of-band via - /// [`Self::try_acquire_session_marker`]. pub fn try_acquire_open_session_token(&self) -> Result, OpteeSmcReturnCode> { let session_id = allocate_session_id().ok_or(OpteeSmcReturnCode::EBusy)?; - if !self.active_sessions.lock().insert(session_id) { - // Defensive: the id pool's hint+wrap guarantees a - // freshly-allocated id won't collide with a recently-recycled - // one, so under normal `allocate → mark → unregister` flow - // this branch is unreachable. It exists to handle the corner - // case where a caller has inserted the slot out-of-band via - // [`Self::try_acquire_session_marker`]. Roll back the alloc - // and let the driver retry. - recycle_session_id(session_id); - return Err(OpteeSmcReturnCode::EThreadLimit); - } + // The id pool's hint+wrap allocator defers reuse of recycled ids, + // so a freshly-allocated id can never collide with a marker slot + // that's still held by a previous owner. + let inserted = self.active_sessions.lock().insert(session_id); + debug_assert!( + inserted, + "freshly-allocated session_id collided with an active marker" + ); Ok(SessionToken { manager: self, uuid_lock: None, @@ -830,7 +817,7 @@ impl SessionManager { /// This counts: /// - All single-instance TAs in the cache (each UUID = 1 instance, regardless of session count) /// - All multi-instance TA sessions (each session = 1 instance) - pub fn instance_count(&self) -> usize { + fn instance_count(&self) -> usize { let single_instance_count = self.single_instance_cache.len(); let multi_instance_count = self.count_multi_instance_sessions(); single_instance_count + multi_instance_count @@ -846,11 +833,6 @@ impl SessionManager { .count() } - /// Check if instance limit is reached. - pub fn is_at_capacity(&self) -> bool { - self.instance_count() >= MAX_TA_INSTANCES - } - /// Drive an OpenSession to completion under the right serialization. /// /// Acquires the UUID-level lock for `uuid` (see [`SessionToken`] for From 38913714017799a9871c11539e08be66a5becaa5 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Mon, 15 Jun 2026 21:33:24 +0000 Subject: [PATCH 22/28] addressed some feedback --- litebox_runner_lvbs/src/lib.rs | 8 ++++++- litebox_shim_optee/src/session.rs | 37 ++++++++++++++++++------------- 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 9108df79bf..30250693af 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -712,7 +712,13 @@ fn open_session_new_instance( client_identity: Option, ta_req_info: &litebox_shim_optee::msg_handler::TaRequestInfo, ) -> Result<(), OpteeSmcReturnCode> { - let ta_bin = find_ta_binary(ta_uuid).ok_or(OpteeSmcReturnCode::ENotAvail)?; + let Some(ta_bin) = find_ta_binary(ta_uuid) else { + 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 diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 87e26e3d6a..f62852546a 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -299,7 +299,7 @@ impl SingleInstanceCache { /// Evict only if the cached instance matches `task_page_table_id`. /// Distinguishes the live instance from a freshly-created one with the /// same UUID when the caller wants to remove a specific one. - fn remove_if_pt(&self, uuid: &TeeUuid, task_page_table_id: usize) -> bool { + fn remove_matching_instance(&self, uuid: &TeeUuid, task_page_table_id: usize) -> bool { let mut guard = self.inner.lock(); match guard.get(uuid) { Some(current) if current.task_page_table_id == task_page_table_id => { @@ -337,8 +337,8 @@ fn recycle_session_id(session_id: u32) { SessionIdPool::recycle(session_id); } -/// RAII token bundling the serialization primitives required to safely -/// execute an OP-TEE TA operation. +/// An unified RAII token to safely execute an OP-TEE TA operation with +/// instance- or session-specific serialization primitives. /// /// Bundles whichever combination of locks the current operation requires: /// @@ -492,6 +492,9 @@ impl SessionManager { // so a freshly-allocated id can never collide with a marker slot // that's still held by a previous owner. let inserted = self.active_sessions.lock().insert(session_id); + if !inserted && !cfg!(debug_assertions) { + litebox_util_log::warn!(session_id = session_id; "freshly-allocated session_id collided with an active marker"); + } debug_assert!( inserted, "freshly-allocated session_id collided with an active marker" @@ -619,9 +622,10 @@ impl SessionManager { /// at a torn-down page table. /// /// Defense in depth: the entry's `(uuid, flags)` are validated against - /// the pre-marker snapshot. If they diverge (the id was recycled and - /// reused under a different TA between our first read and the marker - /// insert), we return `EThreadLimit` so the Linux driver retries. + /// the state observed before inserting the active-session marker. If + /// they diverge (the id was recycled and reused under a different TA + /// between our first read and the marker insert), we return + /// `EThreadLimit` so the Linux driver retries. fn try_acquire_for_session( &self, session_id: u32, @@ -630,8 +634,8 @@ impl SessionManager { .sessions .get_entry(session_id) .ok_or(OpteeSmcReturnCode::EBadCmd)?; - let snapshot_uuid = entry.ta_uuid(); - let snapshot_single = entry.ta_flags().is_single_instance(); + let pre_marker_uuid = entry.ta_uuid(); + let pre_marker_single = entry.ta_flags().is_single_instance(); if !self.active_sessions.lock().insert(session_id) { return Err(OpteeSmcReturnCode::EThreadLimit); @@ -647,20 +651,20 @@ impl SessionManager { // instance TAs. This blocks any concurrent mark-dead / cache // eviction so the re-read result is stable. On failure, the // token's `Drop` releases the marker we already took. - if snapshot_single { + if pre_marker_single { token.uuid_lock = Some( - self.try_acquire_uuid_lock(snapshot_uuid) + self.try_acquire_uuid_lock(pre_marker_uuid) .ok_or(OpteeSmcReturnCode::EThreadLimit)?, ); } - // Re-read under both locks and validate against the snapshot. + // Re-read under both locks and validate against the pre-marker state. let entry_now = self .sessions .get_entry(session_id) .ok_or(OpteeSmcReturnCode::EBadCmd)?; - if entry_now.ta_uuid() != snapshot_uuid - || entry_now.ta_flags().is_single_instance() != snapshot_single + if entry_now.ta_uuid() != pre_marker_uuid + || entry_now.ta_flags().is_single_instance() != pre_marker_single { return Err(OpteeSmcReturnCode::EThreadLimit); } @@ -809,7 +813,7 @@ impl SessionManager { /// that point there are no sibling sessions to fence out. pub fn evict_cached_instance(&self, instance: &TaInstance) -> bool { self.single_instance_cache - .remove_if_pt(&instance.ta_uuid, instance.task_page_table_id) + .remove_matching_instance(&instance.ta_uuid, instance.task_page_table_id) } /// Get the total count of unique TA instances (for limit checking). @@ -861,7 +865,7 @@ impl SessionManager { // Whether we started on the unknown-load path is determined by the // identity of the lock the token carries. Captured before `f` runs // so we know which branch to take in the post-`f` adoption step. - let on_unknown_path = token + let on_unknown_uuid_path = token .uuid_lock .as_ref() .is_some_and(|arc| Arc::ptr_eq(arc, &self.unknown_uuid_lock)); @@ -910,7 +914,8 @@ impl SessionManager { // force-unlocking the old (unknown) lock. Token drop then releases // the per-UUID lock at the end of `with_ta`. UUID-keyed throughout, // so concurrent `with_ta(other_uuid)` cannot adopt our lock. - if result.is_ok() && on_unknown_path && self.single_instance_cache.get(uuid).is_some() { + if result.is_ok() && on_unknown_uuid_path && self.single_instance_cache.get(uuid).is_some() + { let per_uuid = self.uuid_lock_arc(*uuid); if let Some(old) = token.uuid_lock.replace(per_uuid) { // SAFETY: see `SessionToken::drop` — same invariant. From d5049ce4ed8e7deef7ee8f505f3d9dc56bca5cd5 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 18 Jun 2026 00:43:42 +0000 Subject: [PATCH 23/28] drop unnecessary mutex and unsafe --- litebox_shim_optee/src/session.rs | 221 +++++++++++++----------------- 1 file changed, 97 insertions(+), 124 deletions(-) diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index f62852546a..db42bd1d9d 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -104,6 +104,7 @@ use crate::{LoadedProgram, OpteeShim, SessionIdPool}; use alloc::sync::Arc; +use core::sync::atomic::{AtomicBool, Ordering}; use hashbrown::{HashMap, HashSet}; use litebox_common_optee::{OpteeSmcReturnCode, TaFlags, TeeUuid}; use spin::mutex::SpinMutex; @@ -337,13 +338,19 @@ fn recycle_session_id(session_id: u32) { SessionIdPool::recycle(session_id); } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum HeldUuidLock { + SingleInstance(TeeUuid), + UnknownUuid, +} + /// An unified RAII token to safely execute an OP-TEE TA operation with /// instance- or session-specific serialization primitives. /// /// Bundles whichever combination of locks the current operation requires: /// -/// - **Known single-instance TAs**: a per-UUID `SpinMutex` that serializes -/// all sessions on the same TA. +/// - **Known single-instance TAs**: a per-UUID lock flag (a `bool` slot in +/// `single_instance_locks`) that serializes all sessions on the same TA. /// - **First-ever load of an unknown UUID** (OpenSession only): the shared /// `unknown_uuid_lock`, used until the TA's flags (single-instance vs /// multi-instance) are observed. @@ -364,11 +371,9 @@ fn recycle_session_id(session_id: u32) { /// (if still owned) the session id is recycled. pub struct SessionToken<'a> { manager: &'a SessionManager, - /// Held `Arc` of the per-UUID `SpinMutex`. The guard returned by - /// `try_lock()` was [`core::mem::forget`]-ed at acquisition time; this - /// type's `Drop` calls `force_unlock` to release the mutex. The `Arc` - /// keeps the mutex alive across acquisition and release. - uuid_lock: Option>>, + /// Logical UUID-level lock owned by this token. The actual lock state + /// lives in `SessionManager`; `Drop` releases it (clears the held flag). + uuid_lock: Option, /// `Some(id)` while the token holds the active-session marker for `id` /// in [`SessionManager::active_sessions`]. Drop releases the marker. active_session_id: Option, @@ -404,10 +409,7 @@ impl SessionToken<'_> { impl Drop for SessionToken<'_> { fn drop(&mut self) { if let Some(lock) = self.uuid_lock.take() { - // SAFETY: This token holds the per-UUID lock because the - // acquisition path called `try_lock()` and forgot the resulting - // guard. No other holder exists, so `force_unlock` is sound. - unsafe { lock.force_unlock() }; + self.manager.release_uuid_lock(lock); } if let Some(id) = self.active_session_id.take() { self.manager.active_sessions.lock().remove(&id); @@ -443,15 +445,16 @@ pub struct SessionManager { /// versioning is wired through, so a re-loaded TA isn't serialized /// under the old flags. known_flags: SpinMutex>, - /// Per-UUID serialization locks for single-instance TA handling. - /// Entries are created lazily only for UUIDs that have been observed - /// to be single-instance — never for unknown UUIDs whose load might - /// fail or turn out to be multi-instance. - single_instance_locks: SpinMutex>>>, - /// Shared serialization lock for first-ever loads of unknown UUIDs. - /// Held by OpenSession while flags are still unobserved, then released - /// once `known_flags` is updated. Per-UUID locks take over from there. - unknown_uuid_lock: Arc>, + /// Per-UUID serialization state for single-instance TA handling + /// (`true` == held). Entries are created lazily only for UUIDs that + /// have been observed to be single-instance — never for unknown UUIDs + /// whose load might fail or turn out to be multi-instance. + single_instance_locks: SpinMutex>, + /// Shared serialization state for first-ever loads of unknown UUIDs + /// (`true` == held). Held by OpenSession while flags are still + /// unobserved, then released once `known_flags` is updated. Per-UUID + /// locks take over from there. + unknown_uuid_lock: AtomicBool, /// Session ids currently being handled (Invoke/Close). Guards a session /// against concurrent SMC entry by another core that targets the same id. active_sessions: SpinMutex>, @@ -465,7 +468,7 @@ impl SessionManager { pending_count: SpinMutex::new(0), known_flags: SpinMutex::new(HashMap::new()), single_instance_locks: SpinMutex::new(HashMap::new()), - unknown_uuid_lock: Arc::new(SpinMutex::new(())), + unknown_uuid_lock: AtomicBool::new(false), active_sessions: SpinMutex::new(HashSet::new()), } } @@ -531,43 +534,46 @@ impl SessionManager { self.known_flags.lock().get(uuid).copied() } - /// Get or create the per-UUID serialization mutex `Arc`. Only called - /// for UUIDs already observed to be single-instance. - fn uuid_lock_arc(&self, uuid: TeeUuid) -> Arc> { - self.single_instance_locks - .lock() - .entry(uuid) - .or_insert_with(|| Arc::new(SpinMutex::new(()))) - .clone() + /// Try to take the per-UUID serialization state non-blockingly. + fn try_acquire_uuid_lock(&self, uuid: TeeUuid) -> Option { + let mut locks = self.single_instance_locks.lock(); + let held = locks.entry(uuid).or_insert(false); + if *held { + None + } else { + *held = true; + Some(HeldUuidLock::SingleInstance(uuid)) + } } - /// Try to take the per-UUID serialization mutex non-blockingly. On - /// success returns the `Arc` whose forgotten guard is owned by the - /// caller — release via `force_unlock` on the returned `Arc`. - fn try_acquire_uuid_lock(&self, uuid: TeeUuid) -> Option>> { - let lock = self.uuid_lock_arc(uuid); - let guard = lock.try_lock()?; - // The lock now belongs to the SessionToken about to wrap us. Forget - // the guard so its `Drop` does not unlock; the token's `Drop` calls - // `force_unlock` via the retained `Arc`. - core::mem::forget(guard); - Some(lock) + fn release_uuid_lock(&self, lock: HeldUuidLock) { + match lock { + HeldUuidLock::SingleInstance(uuid) => { + if let Some(held) = self.single_instance_locks.lock().get_mut(&uuid) { + debug_assert!(*held); + *held = false; + } + } + HeldUuidLock::UnknownUuid => { + let was_held = self.unknown_uuid_lock.swap(false, Ordering::Release); + debug_assert!(was_held); + } + } } - /// Try to take the shared `unknown_uuid_lock` non-blockingly using the - /// same forget/`force_unlock` pattern as [`Self::try_acquire_uuid_lock`]. - fn try_acquire_unknown_uuid_lock(&self) -> Option>> { - let lock = self.unknown_uuid_lock.clone(); - let guard = lock.try_lock()?; - core::mem::forget(guard); - Some(lock) + /// Try to take the shared `unknown_uuid_lock` non-blockingly. + fn try_acquire_unknown_uuid_lock(&self) -> Option { + self.unknown_uuid_lock + .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed) + .ok() + .map(|_| HeldUuidLock::UnknownUuid) } /// Acquire a `SessionToken` for an OpenSession request. /// /// Dispatches by what's known about `uuid`: /// - /// - **Known single-instance**: per-UUID `SpinMutex`. + /// - **Known single-instance**: per-UUID lock flag. /// - **Known multi-instance**: no lock (each session is independent). /// - **Unknown**: the shared `unknown_uuid_lock`. This serializes /// first-loads of all unknown UUIDs together, but avoids minting a @@ -601,7 +607,7 @@ impl SessionManager { /// callers don't need to look it up again. /// /// Always reserves the per-session-id slot in `active_sessions`. For - /// single-instance TAs additionally takes the per-UUID `SpinMutex` so + /// single-instance TAs additionally takes the per-UUID lock so /// sibling sessions on the same TA serialize against this operation. /// /// Returns `Err(EBadCmd)` if `session_id` is not registered, or @@ -711,21 +717,19 @@ impl SessionManager { /// /// # Unknown→per-UUID transition /// - /// For single-instance TAs we pre-lock the per-UUID `SpinMutex` *before* + /// For single-instance TAs we mark the per-UUID state held *before* /// publishing `known_flags` so any later opener that observes `uuid` - /// as known single-instance and routes to the per-UUID lock finds it - /// already held. The lock state lives in `single_instance_locks` - /// (the `Arc` we get back is discarded — its only purpose was to - /// take the lock); [`Self::with_ta`] adopts by re-fetching the `Arc` - /// for *its own* `uuid` and installing it in its token. This is - /// UUID-keyed end-to-end: no shared side channel, so concurrent - /// `with_ta` calls for different UUIDs cannot interfere with each - /// other's adoptions. + /// as known single-instance and routes to the per-UUID state finds it + /// already held. [`Self::with_ta`] adopts this state for *its own* + /// `uuid` by replacing the token's unknown-lock marker with a + /// per-UUID marker. This is UUID-keyed end-to-end: no shared side + /// channel, so concurrent `with_ta` calls for different UUIDs cannot + /// interfere with each other's adoptions. /// /// `try_acquire_uuid_lock` succeeds only on the unknown path (caller /// holds `unknown_uuid_lock`, no sessions or `known_flags` entry for /// `uuid` yet). On the known-cache-evicted path the caller already - /// holds the per-UUID lock and the `try_lock` returns `None`, so + /// holds the per-UUID state and acquisition returns `None`, so /// nothing changes (the caller's existing lock is sufficient). pub fn register_new_session( &self, @@ -743,10 +747,9 @@ impl SessionManager { ta_uuid, }); - // Pre-lock per-UUID for atomic unknown→per-UUID transition (see - // method doc). The returned `Arc` is intentionally dropped; the - // forgotten guard inside `try_acquire_uuid_lock` keeps the lock - // state held in `single_instance_locks` until `with_ta` adopts. + // Pre-hold per-UUID state for atomic unknown→per-UUID transition + // (see method doc). On known-cache-evicted paths this returns + // `None` because the caller already owns the per-UUID state. if ta_flags.is_single_instance() { let _ = self.try_acquire_uuid_lock(ta_uuid); } @@ -862,13 +865,9 @@ impl SessionManager { F: for<'a> FnOnce(OpenSessionTarget<'a>) -> Result<(), OpteeSmcReturnCode>, { let mut token = self.try_acquire_for_open(*uuid)?; - // Whether we started on the unknown-load path is determined by the - // identity of the lock the token carries. Captured before `f` runs - // so we know which branch to take in the post-`f` adoption step. - let on_unknown_uuid_path = token - .uuid_lock - .as_ref() - .is_some_and(|arc| Arc::ptr_eq(arc, &self.unknown_uuid_lock)); + // Captured before `f` runs so we know whether to perform the + // unknown→per-UUID adoption step after successful registration. + let on_unknown_uuid_path = matches!(token.uuid_lock, Some(HeldUuidLock::UnknownUuid)); // Cache lookup is unconditional: it returns `None` for known // multi-instance and unknown UUIDs (never populated), and only @@ -908,19 +907,17 @@ impl SessionManager { // Complete the unknown→per-UUID transition (see // `register_new_session` doc). Only fires when we held // `unknown_uuid_lock` AND the closure registered a single-instance - // TA for *our* `uuid`. The per-UUID lock is already held (forgotten - // guard in `single_instance_locks` from `register_new_session`'s - // pre-lock); re-fetch the `Arc` by `uuid` and swap into the token, - // force-unlocking the old (unknown) lock. Token drop then releases - // the per-UUID lock at the end of `with_ta`. UUID-keyed throughout, - // so concurrent `with_ta(other_uuid)` cannot adopt our lock. - if result.is_ok() && on_unknown_uuid_path && self.single_instance_cache.get(uuid).is_some() + // TA for *our* `uuid`. The per-UUID state is already held from + // `register_new_session`'s pre-hold; swap the token to own that + // state and release the unknown state. Token drop then releases the + // per-UUID state at the end of `with_ta`. UUID-keyed throughout, so + // concurrent `with_ta(other_uuid)` cannot adopt our lock. + if result.is_ok() + && on_unknown_uuid_path + && self.single_instance_cache.get(uuid).is_some() + && let Some(old) = token.uuid_lock.replace(HeldUuidLock::SingleInstance(*uuid)) { - let per_uuid = self.uuid_lock_arc(*uuid); - if let Some(old) = token.uuid_lock.replace(per_uuid) { - // SAFETY: see `SessionToken::drop` — same invariant. - unsafe { old.force_unlock() }; - } + self.release_uuid_lock(old); } result @@ -959,9 +956,9 @@ mod tests { TaFlags::SINGLE_INSTANCE | TaFlags::MULTI_SESSION } - /// Test helper: call `register_new_session` and release the per-UUID - /// lock the way `with_ta` would, so subsequent operations - /// (Invoke/Close, evict, count, etc.) aren't blocked by a held lock. + /// Test helper: call `register_new_session` directly and release the + /// pre-held per-UUID lock state the way `with_ta` would, so subsequent + /// operations (Invoke/Close, evict, count, etc.) aren't blocked. fn register_for_test( manager: &SessionManager, session_id: u32, @@ -976,14 +973,10 @@ mod tests { task_page_table_id, ta_uuid, ); - // For single-instance, `register_new_session` pre-locked the - // per-UUID mutex via a forgotten guard. Release here to mirror - // `with_ta`'s token-drop release. if ta_flags.is_single_instance() - && let Some(arc) = manager.single_instance_locks.lock().get(&ta_uuid).cloned() + && let Some(held) = manager.single_instance_locks.lock().get_mut(&ta_uuid) { - // SAFETY: same invariant as `SessionToken::drop`. - unsafe { arc.force_unlock() }; + *held = false; } } @@ -1077,8 +1070,8 @@ mod tests { } /// After `with_ta` completes the unknown→per-UUID transition, the - /// per-UUID lock must be released — a subsequent `try_lock` on the - /// per-UUID `SpinMutex` for the same UUID must succeed. + /// per-UUID lock state must be released — a subsequent acquisition for + /// the same UUID must succeed. #[test] fn with_ta_releases_per_uuid_lock_after_unknown_load() { let manager = SessionManager::new(); @@ -1098,15 +1091,11 @@ mod tests { }) .unwrap(); - // Per-UUID lock entry exists (pre-locked + adopted + released). - let arc = manager - .single_instance_locks - .lock() - .get(&uuid) - .cloned() - .expect("entry created by register_new_session"); - // And it's currently unlocked — a fresh try_lock must succeed. - assert!(arc.try_lock().is_some()); + assert_eq!( + manager.single_instance_locks.lock().get(&uuid), + Some(&false) + ); + assert!(manager.try_acquire_uuid_lock(uuid).is_some()); } /// A concurrent `with_ta` for an unrelated UUID must NOT adopt or @@ -1120,35 +1109,19 @@ mod tests { let uuid_locked = make_uuid(0xE1); let uuid_other = make_uuid(0xE2); - // Simulate the "lock pre-taken under unknown-load" state by directly - // pre-locking the per-UUID mutex for `uuid_locked`. This mirrors - // what `register_new_session` does mid-unknown-load before + // Simulate the "lock pre-taken under unknown-load" state. This + // mirrors what `register_new_session` does mid-unknown-load before // `with_ta` adopts. - let pre_locked = manager - .try_acquire_uuid_lock(uuid_locked) - .expect("uncontended"); - // Don't release `pre_locked` — emulating the forgotten-guard state. - core::mem::forget(pre_locked); + assert!(manager.try_acquire_uuid_lock(uuid_locked).is_some()); // A `with_ta` call for a completely different UUID must not touch // `uuid_locked`'s lock. The closure registers nothing, but the // post-`f` adoption logic still runs. manager.with_ta(&uuid_other, |_| Ok(())).unwrap(); - // `uuid_locked`'s per-UUID lock must still be held (not stolen). - let arc = manager - .single_instance_locks - .lock() - .get(&uuid_locked) - .cloned() - .expect("entry exists"); - assert!( - arc.try_lock().is_none(), - "uuid_locked's per-UUID lock must remain held by the original opener" + assert_eq!( + manager.single_instance_locks.lock().get(&uuid_locked), + Some(&true) ); - - // Cleanup: release for SpinMutex sanity. - // SAFETY: we forgot the guard above; release here. - unsafe { arc.force_unlock() }; } } From 44f5448561df1092c5b7160e74018f19e66905e7 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Fri, 19 Jun 2026 02:30:42 +0000 Subject: [PATCH 24/28] update doc --- litebox_shim_optee/src/session.rs | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index db42bd1d9d..16d8067898 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -510,9 +510,10 @@ impl SessionManager { }) } - /// Mark every session currently pointing at `instance` as `Dead`. Must - /// be paired with [`SessionManager::evict_cached_instance`] - /// in the documented order — see that function for the rationale. + /// Mark every session currently pointing at `instance` as `Dead`. + /// + /// Use before [`SessionManager::evict_cached_instance`] when tearing down + /// a *failed* TA that may still have sibling sessions. pub fn mark_sessions_dead_for_instance(&self, instance: &TaInstance) { self.sessions .mark_sessions_dead_for_pt(instance.task_page_table_id); @@ -764,16 +765,8 @@ impl SessionManager { /// Register a session that re-uses an existing single-instance TA. /// - /// `instance` is the borrow handed to the [`SessionManager::with_ta`] - /// closure on the cache-hit branch. The cached instance for `ta_uuid` - /// is matched against `task_page_table_id`. Under correct usage the - /// caller holds the per-UUID lock (via `with_ta`'s token) for the - /// duration, so the cache entry is stable and the lookup succeeds. - /// - /// Returns `Err(EBadCmd)` if no matching cached instance is found. - /// This is an internal-consistency check rather than a recoverable - /// runtime condition; in kernel code we surface it as an error rather - /// than panicking. + /// `instance` is the cached handle handed to the + /// [`SessionManager::with_ta`] closure on the cache-hit branch. pub fn register_sibling_session( &self, session_id: u32, From 455fd338c91c33298f7d786f38631c070465d179 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Fri, 19 Jun 2026 02:51:27 +0000 Subject: [PATCH 25/28] simplification --- litebox_runner_lvbs/src/lib.rs | 7 ++----- litebox_shim_optee/src/session.rs | 34 ++++++++++++++++++------------- 2 files changed, 22 insertions(+), 19 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 30250693af..d30b5200b8 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -547,7 +547,8 @@ fn handle_open_session( /// Open a new session on an existing single-instance TA. /// /// If the TA's OpenSession entry point returns an error, the session is not registered. -/// On TARGET_DEAD the cached instance is destroyed unconditionally. +/// On TARGET_DEAD, sessions for the failed instance are marked `Dead`, the matching +/// single-instance cache entry is evicted, and the TA instance is torn down. /// For cleanup semantics, see OP-TEE OS `tee_ta_open_session()` in `tee_ta_manager.c`. fn open_session_single_instance( msg_args: &mut OpteeMsgArgs, @@ -636,9 +637,7 @@ fn open_session_single_instance( if return_code == TeeResult::TargetDead { debug_serial_println!("Single-instance TA panicked during OpenSession, cleaning up"); - // Mark-then-evict ordering: see SessionManager::evict_cached_instance. session_manager().mark_sessions_dead_for_instance(instance); - let _ = session_manager().evict_cached_instance(instance); // Safety: We are about to tear down this TA instance; // no references to user-space memory will be held afterwards. unsafe { @@ -1027,9 +1026,7 @@ fn handle_invoke_command( ); if instance.loaded_program().ta_flags.is_single_instance() { - // Mark-then-evict ordering: see SessionManager::evict_cached_instance. session_manager().mark_sessions_dead_for_instance(instance); - let _ = session_manager().evict_cached_instance(instance); } session_manager().unregister_session(session_id); diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 16d8067898..d318041e12 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -40,9 +40,9 @@ //! the waiting logic in normal world (where scheduling is appropriate), without //! requiring RPCs that would give untrusted code control over secure world execution. //! -//! Cleanup paths flip sibling sessions to `Dead` before evicting the -//! cached instance; see [`SessionManager::evict_cached_instance`] -//! for the ordering rationale. +//! Panic cleanup paths flip all sessions for the failed instance to `Dead` +//! and evict the matching cached instance via +//! [`SessionManager::mark_sessions_dead_for_instance`]. //! //! Reference: //! @@ -510,13 +510,15 @@ impl SessionManager { }) } - /// Mark every session currently pointing at `instance` as `Dead`. + /// Retire a dead single-instance TA from service. /// - /// Use before [`SessionManager::evict_cached_instance`] when tearing down - /// a *failed* TA that may still have sibling sessions. + /// Marks every session currently pointing at `instance` as `Dead` and + /// evicts the matching entry from the single-instance cache. Use when + /// tearing down a *failed* TA that may still have sibling sessions. pub fn mark_sessions_dead_for_instance(&self, instance: &TaInstance) { self.sessions .mark_sessions_dead_for_pt(instance.task_page_table_id); + let _ = self.evict_cached_instance(instance); } /// Count live sessions currently pointing at `instance` (`Dead` entries @@ -800,11 +802,13 @@ impl SessionManager { /// instance — matched by `task_page_table_id` to distinguish the /// caller's instance from a freshly-cached replacement. /// - /// Callers tearing down on TA panic must have already called - /// [`SessionManager::mark_sessions_dead_for_instance`] before invoking - /// this, so any handler that subsequently enters - /// [`SessionManager::with_ta`] or [`SessionManager::with_session`] for - /// the UUID will observe `Dead` on its re-read of the session entry. + /// TA panic teardown should use + /// [`SessionManager::mark_sessions_dead_for_instance`] instead; it marks + /// all sessions for the failed instance dead and evicts the cache entry + /// in one transition. Later [`SessionManager::with_session`] calls for + /// existing session IDs will observe `Dead` on re-read, while later + /// [`SessionManager::with_ta`] calls for the UUID cannot reuse the dead + /// cached instance. /// Callers on the last-session-close path may skip the mark step — by /// that point there are no sibling sessions to fence out. pub fn evict_cached_instance(&self, instance: &TaInstance) -> bool { @@ -990,9 +994,10 @@ mod tests { assert!(manager.single_instance_cache.get(&uuid).is_some()); } - /// `mark_sessions_dead_for_instance` flips Live entries to Dead — they - /// stop counting for `count_sessions_for_instance`, and `with_session` - /// thereafter sees `None` so cleanup paths run. + /// `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`, + /// and new opens cannot reuse the dead cached instance. #[test] fn mark_dead_makes_with_session_observe_none() { let manager = SessionManager::new(); @@ -1003,6 +1008,7 @@ mod tests { manager.mark_sessions_dead_for_instance(&arc); assert_eq!(manager.count_sessions_for_instance(&arc), 0); + assert!(manager.single_instance_cache.get(&uuid).is_none()); manager .with_session(108, |instance| { From 77a7eef0e52ca536f049f275e26d3e66d0b5226e Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 25 Jun 2026 02:51:48 +0000 Subject: [PATCH 26/28] add some comment --- litebox_shim_optee/src/session.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index d318041e12..664906dcf0 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -449,6 +449,11 @@ pub struct SessionManager { /// (`true` == held). Entries are created lazily only for UUIDs that /// have been observed to be single-instance — never for unknown UUIDs /// whose load might fail or turn out to be multi-instance. + /// + /// We do not remove its entry even if the instance is destroyed to + /// support a future reload of the same TA. This is bounded in + /// practice because we only support a few managed TAs. This entry + /// management should be aligned with `known_flags`. single_instance_locks: SpinMutex>, /// Shared serialization state for first-ever loads of unknown UUIDs /// (`true` == held). Held by OpenSession while flags are still From b4d7d8ebd52b5fe39f138b43b3f3a7bca2f06e5f Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 25 Jun 2026 15:32:35 +0000 Subject: [PATCH 27/28] rename: unknown_uuid_lock->ta_load_lock --- litebox_shim_optee/src/session.rs | 70 +++++++++++++++---------------- 1 file changed, 35 insertions(+), 35 deletions(-) diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index 664906dcf0..e61dd889d4 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -341,7 +341,7 @@ fn recycle_session_id(session_id: u32) { #[derive(Clone, Copy, Debug, Eq, PartialEq)] enum HeldUuidLock { SingleInstance(TeeUuid), - UnknownUuid, + TaLoad, } /// An unified RAII token to safely execute an OP-TEE TA operation with @@ -351,8 +351,8 @@ enum HeldUuidLock { /// /// - **Known single-instance TAs**: a per-UUID lock flag (a `bool` slot in /// `single_instance_locks`) that serializes all sessions on the same TA. -/// - **First-ever load of an unknown UUID** (OpenSession only): the shared -/// `unknown_uuid_lock`, used until the TA's flags (single-instance vs +/// - **First-ever load of a not-yet-known UUID** (OpenSession only): the +/// global `ta_load_lock`, used until the TA's flags (single-instance vs /// multi-instance) are observed. /// - **Existing-session operations** (Invoke/Close): a per-session-id /// marker (slot in `SessionManager::active_sessions`) that prevents @@ -367,7 +367,7 @@ enum HeldUuidLock { /// exclusion is required there. /// /// On drop the held UUID-level lock is released first (whether per-UUID -/// or the shared unknown lock), then the per-session-id marker, then +/// or the global load lock), then the per-session-id marker, then /// (if still owned) the session id is recycled. pub struct SessionToken<'a> { manager: &'a SessionManager, @@ -455,11 +455,11 @@ pub struct SessionManager { /// practice because we only support a few managed TAs. This entry /// management should be aligned with `known_flags`. single_instance_locks: SpinMutex>, - /// Shared serialization state for first-ever loads of unknown UUIDs - /// (`true` == held). Held by OpenSession while flags are still - /// unobserved, then released once `known_flags` is updated. Per-UUID - /// locks take over from there. - unknown_uuid_lock: AtomicBool, + /// Global gate that serializes the first-ever load of not-yet-known + /// UUIDs. Held by a first-loader until the TA's flags are observed; for a + /// single-instance TA, ownership is then handed off to its per-UUID lock + /// (see [`SessionToken`]). Known multi-instance UUIDs take no lock. + ta_load_lock: AtomicBool, /// Session ids currently being handled (Invoke/Close). Guards a session /// against concurrent SMC entry by another core that targets the same id. active_sessions: SpinMutex>, @@ -473,7 +473,7 @@ impl SessionManager { pending_count: SpinMutex::new(0), known_flags: SpinMutex::new(HashMap::new()), single_instance_locks: SpinMutex::new(HashMap::new()), - unknown_uuid_lock: AtomicBool::new(false), + ta_load_lock: AtomicBool::new(false), active_sessions: SpinMutex::new(HashSet::new()), } } @@ -562,19 +562,19 @@ impl SessionManager { *held = false; } } - HeldUuidLock::UnknownUuid => { - let was_held = self.unknown_uuid_lock.swap(false, Ordering::Release); + HeldUuidLock::TaLoad => { + let was_held = self.ta_load_lock.swap(false, Ordering::Release); debug_assert!(was_held); } } } - /// Try to take the shared `unknown_uuid_lock` non-blockingly. - fn try_acquire_unknown_uuid_lock(&self) -> Option { - self.unknown_uuid_lock + /// Try to take the global `ta_load_lock` non-blockingly. + fn try_acquire_ta_load_lock(&self) -> Option { + self.ta_load_lock .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed) .ok() - .map(|_| HeldUuidLock::UnknownUuid) + .map(|_| HeldUuidLock::TaLoad) } /// Acquire a `SessionToken` for an OpenSession request. @@ -583,10 +583,10 @@ impl SessionManager { /// /// - **Known single-instance**: per-UUID lock flag. /// - **Known multi-instance**: no lock (each session is independent). - /// - **Unknown**: the shared `unknown_uuid_lock`. This serializes - /// first-loads of all unknown UUIDs together, but avoids minting a - /// per-UUID lock entry until the TA has been confirmed single-instance. - /// A failed or multi-instance load therefore leaves no stale entry in + /// - **Unknown**: the global `ta_load_lock`. This serializes first-loads + /// of all not-yet-known UUIDs together, but avoids minting a per-UUID + /// lock entry until the TA has been confirmed single-instance. A failed + /// or multi-instance load therefore leaves no stale entry in /// `single_instance_locks`. /// /// Returns `Err(EThreadLimit)` on contention. @@ -598,7 +598,7 @@ impl SessionManager { ), Some(_) => None, None => Some( - self.try_acquire_unknown_uuid_lock() + self.try_acquire_ta_load_lock() .ok_or(OpteeSmcReturnCode::EThreadLimit)?, ), }; @@ -729,13 +729,13 @@ impl SessionManager { /// publishing `known_flags` so any later opener that observes `uuid` /// as known single-instance and routes to the per-UUID state finds it /// already held. [`Self::with_ta`] adopts this state for *its own* - /// `uuid` by replacing the token's unknown-lock marker with a + /// `uuid` by replacing the token's load-lock marker with a /// per-UUID marker. This is UUID-keyed end-to-end: no shared side /// channel, so concurrent `with_ta` calls for different UUIDs cannot /// interfere with each other's adoptions. /// - /// `try_acquire_uuid_lock` succeeds only on the unknown path (caller - /// holds `unknown_uuid_lock`, no sessions or `known_flags` entry for + /// `try_acquire_uuid_lock` succeeds only on the load-lock path (caller + /// holds `ta_load_lock`, no sessions or `known_flags` entry for /// `uuid` yet). On the known-cache-evicted path the caller already /// holds the per-UUID state and acquisition returns `None`, so /// nothing changes (the caller's existing lock is sufficient). @@ -868,8 +868,8 @@ impl SessionManager { { let mut token = self.try_acquire_for_open(*uuid)?; // Captured before `f` runs so we know whether to perform the - // unknown→per-UUID adoption step after successful registration. - let on_unknown_uuid_path = matches!(token.uuid_lock, Some(HeldUuidLock::UnknownUuid)); + // load-lock→per-UUID adoption step after successful registration. + let on_ta_load_path = matches!(token.uuid_lock, Some(HeldUuidLock::TaLoad)); // Cache lookup is unconditional: it returns `None` for known // multi-instance and unknown UUIDs (never populated), and only @@ -906,16 +906,16 @@ impl SessionManager { *pending = pending.saturating_sub(1); } - // Complete the unknown→per-UUID transition (see - // `register_new_session` doc). Only fires when we held - // `unknown_uuid_lock` AND the closure registered a single-instance + // Complete the load-lock→per-UUID transition (see + // `register_new_session` doc). Only fires when we held the + // `ta_load_lock` AND the closure registered a single-instance // TA for *our* `uuid`. The per-UUID state is already held from // `register_new_session`'s pre-hold; swap the token to own that - // state and release the unknown state. Token drop then releases the + // state and release the load lock. Token drop then releases the // per-UUID state at the end of `with_ta`. UUID-keyed throughout, so // concurrent `with_ta(other_uuid)` cannot adopt our lock. if result.is_ok() - && on_unknown_uuid_path + && on_ta_load_path && self.single_instance_cache.get(uuid).is_some() && let Some(old) = token.uuid_lock.replace(HeldUuidLock::SingleInstance(*uuid)) { @@ -1024,7 +1024,7 @@ mod tests { } /// A failed first-load of an unknown UUID must not mint a per-UUID - /// lock entry. Unknown loads serialize on `unknown_uuid_lock`, so + /// lock entry. Such loads serialize on `ta_load_lock`, so /// `single_instance_locks` stays empty when the load fails or the TA /// turns out to be multi-instance. #[test] @@ -1103,7 +1103,7 @@ mod tests { } /// A concurrent `with_ta` for an unrelated UUID must NOT adopt or - /// release the per-UUID lock held by another unknown-load opener. + /// release the per-UUID lock held by another first-load opener. /// Adoption is keyed by the `with_ta` call's own UUID, so an opener /// for a different UUID leaves the original opener's per-UUID lock /// untouched. @@ -1113,8 +1113,8 @@ mod tests { let uuid_locked = make_uuid(0xE1); let uuid_other = make_uuid(0xE2); - // Simulate the "lock pre-taken under unknown-load" state. This - // mirrors what `register_new_session` does mid-unknown-load before + // Simulate the "lock pre-taken during a first-load" state. This + // mirrors what `register_new_session` does mid-first-load before // `with_ta` adopts. assert!(manager.try_acquire_uuid_lock(uuid_locked).is_some()); From 1cd8b5d9168c0961095db5066d7313602ad05a52 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 25 Jun 2026 16:50:59 +0000 Subject: [PATCH 28/28] revise comments --- litebox_runner_lvbs/src/lib.rs | 2 +- litebox_shim_optee/src/session.rs | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index d30b5200b8..893afe57fa 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -1005,7 +1005,7 @@ fn handle_invoke_command( let return_code = TeeResult::try_from(return_code).unwrap_or(TeeResult::GenericError); // Write response BEFORE switching page tables (accesses user memory). - // `with_session`'s marker keeps the entry stable so another core cannot + // `with_session`'s serialization keeps the entry stable so another core cannot // tear down the active page table while this core is copying TA outputs. let write_result = write_msg_args_to_normal_world( msg_args, diff --git a/litebox_shim_optee/src/session.rs b/litebox_shim_optee/src/session.rs index e61dd889d4..3e462da290 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -627,11 +627,11 @@ impl SessionManager { /// /// The per-UUID lock is acquired *before* the final session-map /// re-read. This excludes concurrent `mark_sessions_dead_for_instance` - /// and cache eviction (both of which require the UUID lock), so the - /// `Live` / `Dead` state observed in the re-read remains authoritative - /// for the lifetime of the returned token. Reading the entry before - /// taking the UUID lock would let a sibling complete the entire - /// mark-dead / evict / teardown sequence between our read and our + /// and cache eviction (which callers perform only while holding the UUID + /// lock), so the `Live` / `Dead` state observed in the re-read remains + /// authoritative for the lifetime of the returned token. Reading the + /// entry before taking the UUID lock would let a sibling complete the + /// entire mark-dead / evict / teardown sequence between our read and our /// lock acquisition, leaving us holding a stale `Live` entry pointing /// at a torn-down page table. ///