diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 0ae319e59d..893afe57fa 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -5,9 +5,7 @@ extern crate alloc; -use alloc::boxed::Box; -use alloc::sync::Arc; -use alloc::vec; +use alloc::{boxed::Box, vec}; use core::{ops::Neg, panic::PanicInfo}; use litebox::{ mm::linux::PAGE_SIZE, @@ -43,12 +41,9 @@ 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::{ - CreationReservation, 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; -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 +353,30 @@ 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 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 +/// table can 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) + } +} + +impl Drop for TaskPageTableGuard { + fn drop(&mut self) { + unsafe { switch_to_base_page_table() }; + } +} + /// Tears down a TA's memory mappings and page table. /// /// This performs the following steps in order: @@ -497,98 +516,54 @@ 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()); - - // 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, || { - open_session_new_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, + ), + OpenSessionTarget::NewInstance => open_session_new_instance( msg_args, msg_args_phys_addr, params, ta_uuid, 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) - } - } + ), + 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(()) } - CreationReservation::SlotReserved => Ok(()), - } -} - -/// 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, + }) } /// 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, 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`. -#[allow(clippy::type_complexity)] fn open_session_single_instance( msg_args: &mut OpteeMsgArgs, msg_args_phys_addr: u64, - instance_arc: Arc>, + instance: &TaInstance, 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.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()?; + // Safe to unwrap: session ID has been just created. + let runner_session_id = session_token.session_id().unwrap(); debug_serial_println!( "Reusing single-instance TA: uuid={:?}, task_pt_id={}, session_id={}", @@ -597,14 +572,12 @@ fn open_session_single_instance( runner_session_id ); - let ta_flags = instance.loaded_program.ta_flags; - // Switch to the existing TA's page table - unsafe { switch_to_task_page_table(task_pt_id)? }; + let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; // Load TA context with parameters for OpenSession - pass actual session_id instance - .loaded_program + .loaded_program() .entrypoints .as_ref() .ok_or(OpteeSmcReturnCode::EBadCmd)? @@ -620,14 +593,14 @@ 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.loaded_program().entrypoints.as_ref().unwrap(), &mut ctx, ); } // Read TA output parameters from the stack buffer let params_address = instance - .loaded_program + .loaded_program() .params_address .ok_or(OpteeSmcReturnCode::EBadAddr)?; let ta_params = UserConstPtr::::from_usize(params_address) @@ -647,7 +620,7 @@ 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 + // `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, @@ -664,25 +637,23 @@ fn open_session_single_instance( if return_code == TeeResult::TargetDead { debug_serial_println!("Single-instance TA panicked during OpenSession, cleaning up"); - let _ = session_manager().remove_single_instance_if_same(&ta_uuid, &instance_arc); - instance.closed = true; - + session_manager().mark_sessions_dead_for_instance(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.shim, task_pt_id) }; + unsafe { + teardown_ta_page_table(instance.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. - let runner_session_id = session_id_guard.id().unwrap(); let write_result = write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, @@ -696,48 +667,39 @@ 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() - .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); - instance.closed = true; - + 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 { teardown_ta_page_table(&instance.shim, task_pt_id) }; + unsafe { + teardown_ta_page_table(instance.shim(), task_pt_id); + }; } else { - let _ = session_id_guard.disarm(); + session_token.disarm(); } - drop(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); - session_id_guard.disarm(); - - drop(instance); + // Success: register a sibling session pointing at the existing instance. + session_manager().register_sibling_session(runner_session_id, instance)?; + session_token.disarm(); 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. -/// -/// The caller must invoke this inside [`SessionManager::with_creation_slot`] -/// 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. @@ -749,32 +711,27 @@ 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(()); + }; - // Create and switch to new page table - let task_pt_id = create_task_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); - 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); - })?; - } - - // 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 _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) }; - OpteeSmcReturnCode::EBusy - })?); - // 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(); @@ -920,9 +877,7 @@ 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(); + // local resources and let `session_token` recycle the ID on drop. write_msg_args_to_normal_world( msg_args, msg_args_phys_addr, @@ -937,22 +892,15 @@ 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(SpinMutex::new(TaInstance { + // Success: register the new session with the manager. + session_manager().register_new_session( + runner_session_id, 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); - 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()); - } + task_pt_id, + ta_uuid, + ); + session_token.disarm(); debug_serial_println!( "OpenSession complete: session_id={}, single_instance={}", @@ -963,6 +911,29 @@ fn open_session_new_instance( Ok(()) } +/// Tear down a `Dead` session entry observed at Invoke/Close handler entry. +/// +/// 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, + msg_args_phys_addr: u64, + return_code: TeeResult, + log_prefix: &str, +) -> Result<(), OpteeSmcReturnCode> { + session_manager().unregister_session(session_id); + 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. @@ -981,129 +952,103 @@ 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)?; - // 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); - }; - // `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); - session_manager().unregister_session(session_id); - 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", - session_id - ); - return Ok(()); - } - let task_pt_id = instance.task_page_table_id; - - // Switch to the TA instance's page table - unsafe { switch_to_task_page_table(task_pt_id)? }; + session_manager().with_session(session_id, |instance| { + let Some(instance) = instance else { + return finalize_dead_session( + session_id, + msg_args, + msg_args_phys_addr, + TeeResult::TargetDead, + "InvokeCommand", + ); + }; + let task_pt_id = instance.task_page_table_id(); - debug_serial_println!( - "InvokeCommand: session_id={}, task_pt_id={}, cmd_id={}", - session_id, - task_pt_id, - cmd_id - ); + let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; - // Load TA context with parameters and cmd_id - pass actual session_id - let entrypoints_ref = instance.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.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 - .loaded_program - .params_address - .ok_or(OpteeSmcReturnCode::EBadAddr)?; - let ta_params = UserConstPtr::::from_usize(params_address) - .read_at_offset(0) - .ok_or(OpteeSmcReturnCode::EBadAddr)?; + // Set up the entry-point parameters for InvokeCommand. + let entrypoints_ref = instance.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 mut ctx = litebox_common_linux::PtRegs::default(); + unsafe { + litebox_platform_lvbs::reenter_thread_ref( + instance.loaded_program().entrypoints.as_ref().unwrap(), + &mut ctx, + ); + } - let return_code: u32 = ctx.rax.trunc(); - let return_code = TeeResult::try_from(return_code).unwrap_or(TeeResult::GenericError); + // params_address is constant - stack buffer is reused across invocations + let params_address = instance + .loaded_program() + .params_address + .ok_or(OpteeSmcReturnCode::EBadAddr)?; + let ta_params = UserConstPtr::::from_usize(params_address) + .read_at_offset(0) + .ok_or(OpteeSmcReturnCode::EBadAddr)?; - // 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. - 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 return_code: u32 = ctx.rax.trunc(); + let return_code = TeeResult::try_from(return_code).unwrap_or(TeeResult::GenericError); - // 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 + // Write response BEFORE switching page tables (accesses user memory). + // `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, + 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; - - // 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); - } + // 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 + ); - instance.closed = true; + if instance.loaded_program().ta_flags.is_single_instance() { + session_manager().mark_sessions_dead_for_instance(instance); + } - // 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) }; + session_manager().unregister_session(session_id); - drop(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.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. @@ -1123,87 +1068,63 @@ 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)?; - // 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); - }; - // `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); - session_manager().unregister_session(session_id); - 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", - session_id - ); - return Ok(()); - } - let task_pt_id = instance.task_page_table_id; - - // Switch to the TA instance's page table - unsafe { switch_to_task_page_table(task_pt_id)? }; + session_manager().with_session(session_id, |instance| { + let Some(instance) = instance else { + return finalize_dead_session( + session_id, + msg_args, + msg_args_phys_addr, + TeeResult::Success, + "CloseSession", + ); + }; + let task_pt_id = instance.task_page_table_id(); + + let _task_pt_guard = TaskPageTableGuard::enter(task_pt_id)?; + + // Set up the entry-point parameters for CloseSession. + instance + .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.loaded_program().entrypoints.as_ref().unwrap(), + &mut ctx, + ); + } - // Load TA context for CloseSession (no params, no cmd_id) - pass actual session_id - instance - .loaded_program - .entrypoints - .as_ref() - .unwrap() - .load_ta_context( - &[], - Some(session_id), - UteeEntryFunc::CloseSession as u32, + // 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, - ) - .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.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, - None, - None, - None, - ); + let removed_flags = session_manager().unregister_session(session_id); + + let remaining_sessions = session_manager().count_sessions_for_instance(instance); - // 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); - - // 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 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); + // 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 flags.is_single_instance() && flags.is_keep_alive() { debug_serial_println!( "CloseSession complete: session_id={}, TA kept alive (INSTANCE_KEEP_ALIVE flag)", session_id @@ -1211,35 +1132,31 @@ fn handle_close_session( return write_result; } - // Clear single-instance cache if this was a single-instance TA - if entry.ta_flags.is_single_instance() { - let _ = - session_manager().remove_single_instance_if_same(&entry.ta_uuid, &instance_arc); + // 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 flags.is_single_instance() { + let _ = session_manager() + .evict_cached_instance(instance); } - 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_ta_page_table(instance.shim(), task_pt_id); + }; debug_serial_println!( "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 { - drop(instance); - 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_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 d4f411f0f4..2afba0c79b 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::{ - CreationReservation, MAX_TA_INSTANCES, SessionEntry, SessionManager, SessionMap, - SingleInstanceCache, TaInstance, allocate_session_id, -}; +pub use session::{OpenSessionTarget, SessionManager, SessionToken, TaInstance}; const MAX_KERNEL_BUF_SIZE: usize = 0x80_000; @@ -1449,6 +1446,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 5d21b404a8..3e462da290 100644 --- a/litebox_shim_optee/src/session.rs +++ b/litebox_shim_optee/src/session.rs @@ -10,10 +10,16 @@ //! //! ## 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 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`. //! //! ### Difference from OP-TEE OS //! @@ -34,6 +40,10 @@ //! the waiting logic in normal world (where scheduling is appropriate), without //! requiring RPCs that would give untrusted code control over secure world execution. //! +//! 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: //! //! ## OP-TEE OS Thread IDs and RPC @@ -94,122 +104,165 @@ 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; /// 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 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). -/// -/// Each instance has its own task page table that provides memory isolation from other TAs. +/// For single-instance TAs one instance is shared across all sessions; the +/// 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. - 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, - /// The task page table ID associated with this TA instance. Valid only - /// while `closed == false`. - 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. + loaded_program: alloc::boxed::Box, + /// The task page table ID associated with this TA instance. /// - /// The per-instance lock must be held when setting `closed = true` and across - /// the subsequent `teardown_ta_page_table`. - pub closed: bool, + /// 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, + ta_uuid: TeeUuid, } -// SAFETY: TaInstance is protected by SpinMutex and try_lock (`SessionEntry`) +impl TaInstance { + pub fn task_page_table_id(&self) -> usize { + self.task_page_table_id + } + + pub fn shim(&self) -> &OpteeShim { + &self.shim + } + + 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 +// 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 {} -/// Per-session entry in the session map. +/// 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`, reject with + /// `TEE_ERROR_BUSY` (origin TEE). + 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. #[derive(Clone)] -pub struct SessionEntry { - /// The TA instance (may be shared with other sessions for single-instance TAs). - pub instance: Arc>, - /// 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, +enum SessionEntry { + Live(Arc), + Dead { ta_uuid: TeeUuid, ta_flags: TaFlags }, +} + +impl SessionEntry { + fn ta_uuid(&self) -> TeeUuid { + match self { + SessionEntry::Live(arc) => arc.ta_uuid, + SessionEntry::Dead { ta_uuid, .. } => *ta_uuid, + } + } + + 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. /// /// Maps runner-allocated session IDs to session entries. -pub struct SessionMap { +struct SessionMap { inner: SpinMutex>, } impl SessionMap { /// Create a new empty session map. - pub fn new() -> Self { + 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>> { - self.inner - .lock() - .get(&session_id) - .map(|e| e.instance.clone()) - } - /// Get full session entry by session ID. - pub fn get_entry(&self, session_id: u32) -> Option { + fn get_entry(&self, session_id: u32) -> Option { self.inner.lock().get(&session_id).cloned() } - /// Insert a session into the map. - pub fn insert( - &self, - session_id: u32, - instance: Arc>, - ta_uuid: TeeUuid, - ta_flags: TaFlags, - ) { - self.inner.lock().insert( - session_id, - SessionEntry { - instance, - ta_uuid, - ta_flags, - }, - ); + /// 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. - pub fn remove(&self, session_id: u32) -> Option { + 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 { + /// Count live sessions whose instance has the given page table id. + fn count_sessions_for_pt(&self, task_page_table_id: usize) -> usize { self.inner .lock() .values() - .filter(|e| Arc::ptr_eq(&e.instance, instance)) + .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`, capturing the instance's uuid and flags on the way out + /// so cleanup paths still have them. + 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 => { + 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 }; + } + } + } } impl Default for SessionMap { @@ -222,33 +275,35 @@ 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>>>, +struct SingleInstanceCache { + inner: SpinMutex>>, } impl SingleInstanceCache { /// Create a new empty cache. - pub fn new() -> Self { + fn new() -> Self { Self { inner: SpinMutex::new(HashMap::new()), } } /// Get a cached single-instance TA by UUID. - pub fn get(&self, uuid: &TeeUuid) -> Option>> { + 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>) { + 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_matching_instance(&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 } @@ -257,14 +312,9 @@ impl SingleInstanceCache { } /// Get the number of cached single-instance TAs. - pub fn len(&self) -> usize { + 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 { @@ -277,177 +327,498 @@ impl Default for SingleInstanceCache { /// /// Delegates to `SessionIdPool::allocate` for unified session ID management. /// Returns `None` if all session IDs are exhausted. -pub fn allocate_session_id() -> Option { +fn allocate_session_id() -> Option { SessionIdPool::allocate() } /// 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 -/// recycled on all error paths before this registration. -pub struct SessionIdGuard { - session_id: Option, +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum HeldUuidLock { + SingleInstance(TeeUuid), + TaLoad, } -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), - } - } +/// 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 lock flag (a `bool` slot in +/// `single_instance_locks`) that serializes all sessions on the same TA. +/// - **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 +/// 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 (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 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, + /// 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, + /// 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, +} - /// Return the guarded session ID, or `None` if already disarmed. - pub fn id(&self) -> Option { - self.session_id +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`] or + /// `try_acquire_for_session` (Invoke/Close). + pub fn session_id(&self) -> Option { + self.active_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() + /// 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 SessionIdGuard { +impl Drop for SessionToken<'_> { fn drop(&mut self) { - if let Some(id) = self.session_id { - recycle_session_id(id); + if let Some(lock) = self.uuid_lock.take() { + self.manager.release_uuid_lock(lock); + } + 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); + } } } } -/// 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 and -/// duplicate-UUID prevention. -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. - /// 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: -/// - 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, /// 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. + /// + /// 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 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. + /// + /// 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>, + /// 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>, } impl SessionManager { - /// Create a new session manager. pub fn new() -> Self { Self { sessions: SessionMap::new(), single_instance_cache: SingleInstanceCache::new(), - creation_state: SpinMutex::new(CreationState { - pending_uuids: HashSet::new(), - pending_count: 0, - }), + pending_count: SpinMutex::new(0), known_flags: SpinMutex::new(HashMap::new()), + single_instance_locks: SpinMutex::new(HashMap::new()), + ta_load_lock: AtomicBool::new(false), + active_sessions: SpinMutex::new(HashSet::new()), } } - /// 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); + /// Allocate a fresh `session_id` and reserve its active-session slot. + /// 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. + pub fn try_acquire_open_session_token(&self) -> Result, OpteeSmcReturnCode> { + let session_id = allocate_session_id().ok_or(OpteeSmcReturnCode::EBusy)?; + // 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); + 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" + ); + Ok(SessionToken { + manager: self, + uuid_lock: None, + active_session_id: Some(session_id), + owns_id_recycling: true, + }) } - /// Get a session by ID. - pub fn get_session(&self, session_id: u32) -> Option>> { - self.sessions.get(session_id) + /// Retire a dead single-instance TA from service. + /// + /// 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); } - /// Get full session entry by ID. - pub fn get_session_entry(&self, session_id: u32) -> Option { - self.sessions.get_entry(session_id) + /// 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. /// /// 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() } - /// Register a new session. - pub fn register_session( + /// 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)) + } + } + + 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::TaLoad => { + let was_held = self.ta_load_lock.swap(false, Ordering::Release); + debug_assert!(was_held); + } + } + } + + /// 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::TaLoad) + } + + /// Acquire a `SessionToken` for an OpenSession request. + /// + /// Dispatches by what's known about `uuid`: + /// + /// - **Known single-instance**: per-UUID lock flag. + /// - **Known multi-instance**: no lock (each session is independent). + /// - **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. + fn try_acquire_for_open(&self, uuid: TeeUuid) -> Result, OpteeSmcReturnCode> { + 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)?, + ), + Some(_) => None, + None => Some( + self.try_acquire_ta_load_lock() + .ok_or(OpteeSmcReturnCode::EThreadLimit)?, + ), + }; + Ok(SessionToken { + manager: self, + uuid_lock, + active_session_id: None, + owns_id_recycling: false, + }) + } + + /// Acquire a token + validated entry for an Invoke/Close on an existing + /// 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 + /// 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 + /// `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`. + /// + /// # 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 (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. + /// + /// Defense in depth: the entry's `(uuid, flags)` are validated against + /// 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, + ) -> Result<(SessionToken<'_>, SessionEntry), OpteeSmcReturnCode> { + let entry = self + .sessions + .get_entry(session_id) + .ok_or(OpteeSmcReturnCode::EBadCmd)?; + 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); + } + let mut token = SessionToken { + 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- + // 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 pre_marker_single { + token.uuid_lock = Some( + self.try_acquire_uuid_lock(pre_marker_uuid) + .ok_or(OpteeSmcReturnCode::EThreadLimit)?, + ); + } + + // 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() != pre_marker_uuid + || entry_now.ta_flags().is_single_instance() != pre_marker_single + { + return Err(OpteeSmcReturnCode::EThreadLimit); + } + + Ok((token, entry_now)) + } + + /// 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 (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>, + { + let (_token, entry) = self.try_acquire_for_session(session_id)?; + let instance = match &entry { + SessionEntry::Live(arc) => Some(&**arc), + SessionEntry::Dead { .. } => None, + }; + f(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; 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 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 state finds it + /// already held. [`Self::with_ta`] adopts this state for *its own* + /// `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 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). + pub fn register_new_session( &self, session_id: u32, - instance: Arc>, + shim: OpteeShim, + 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, + }); + + // 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); + } + + 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); - self.sessions - .insert(session_id, instance, ta_uuid, ta_flags); } - /// Unregister a session, recycle its session ID, and return the entry. - pub fn unregister_session(&self, session_id: u32) -> Option { + /// Register a session that re-uses an existing single-instance TA. + /// + /// `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, + instance: &TaInstance, + ) -> Result<(), OpteeSmcReturnCode> { + let arc = self + .single_instance_cache + .get(&instance.ta_uuid) + .filter(|cached| cached.task_page_table_id == instance.task_page_table_id) + .ok_or(OpteeSmcReturnCode::EBadCmd)?; + // `known_flags` is already populated for this UUID — sibling path + // implies the instance was previously registered. + self.sessions.insert_live(session_id, arc); + Ok(()) + } + + /// 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`. - pub fn remove_single_instance_if_same( - &self, - uuid: &TeeUuid, - expected: &Arc>, - ) -> bool { - self.single_instance_cache.remove_if_same(uuid, expected) + /// 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. + /// + /// 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 { + self.single_instance_cache + .remove_matching_instance(&instance.ta_uuid, instance.task_page_table_id) } /// Get the total count of unique TA instances (for limit checking). @@ -455,7 +826,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 @@ -467,78 +838,91 @@ impl SessionManager { .inner .lock() .values() - .filter(|e| !e.ta_flags.is_single_instance()) + .filter(|e| !e.ta_flags().is_single_instance()) .count() } - /// Check if instance limit is reached. - pub fn is_at_capacity(&self) -> bool { - self.instance_count() >= MAX_TA_INSTANCES - } - - /// Atomically reserve a creation slot and run `f` to create a new TA instance. + /// Drive an OpenSession to completion under the right serialization. /// - /// Behavior depends on whether the TA is: + /// Acquires the UUID-level lock for `uuid` (see [`SessionToken`] for + /// the case breakdown), classifies the cache state, and dispatches + /// via [`OpenSessionTarget`]: /// - /// - **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`. + /// - [`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. /// - /// - **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. - pub fn with_creation_slot( - &self, - uuid: &TeeUuid, - is_single_instance: bool, - f: F, - ) -> Result + /// `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: FnOnce() -> Result<(), OpteeSmcReturnCode>, + F: for<'a> FnOnce(OpenSessionTarget<'a>) -> Result<(), OpteeSmcReturnCode>, { - { - 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)); - } + let mut token = self.try_acquire_for_open(*uuid)?; + // Captured before `f` runs so we know whether to perform the + // load-lock→per-UUID adoption step after successful registration. + let on_ta_load_path = matches!(token.uuid_lock, Some(HeldUuidLock::TaLoad)); - // 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); - } - } + // 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) { + // 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); + } + { + 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); } - - if is_single_instance { - state.pending_uuids.insert(*uuid); - } - state.pending_count += 1; + *pending += 1; } - let result = f(); + let result = f(OpenSessionTarget::NewInstance); { - 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); + let mut pending = self.pending_count.lock(); + *pending = pending.saturating_sub(1); } - result.map(|()| CreationReservation::SlotReserved) + // 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 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_ta_load_path + && self.single_instance_cache.get(uuid).is_some() + && let Some(old) = token.uuid_lock.replace(HeldUuidLock::SingleInstance(*uuid)) + { + self.release_uuid_lock(old); + } + + result } } @@ -547,3 +931,201 @@ 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 + } + + /// 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, + 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, + ); + if ta_flags.is_single_instance() + && let Some(held) = manager.single_instance_locks.lock().get_mut(&ta_uuid) + { + *held = false; + } + } + + /// 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); + + 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); + + 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()); + } + + /// `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(); + let uuid = make_uuid(0xA6); + 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); + + 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| { + assert!(instance.is_none()); + Ok(()) + }) + .unwrap(); + } + + /// A failed first-load of an unknown UUID must not mint a per-UUID + /// 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] + 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_none()); + assert!(manager.get_known_flags(&uuid).is_none()); + } + + /// `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, |target| { + assert!(matches!(target, OpenSessionTarget::NewInstance)); + 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. + let _ = manager.with_ta(&uuid_single, |_| Err(OpteeSmcReturnCode::ENotAvail)); + assert_eq!(*manager.pending_count.lock(), 0); + + // Cache-hit path doesn't touch pending_count. + 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 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(); + 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(); + + 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 + /// 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. + #[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 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()); + + // 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(); + + assert_eq!( + manager.single_instance_locks.lock().get(&uuid_locked), + Some(&true) + ); + } +}