From b9f639f76da910c019b1968862784160df5ad913 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Thu, 10 Sep 2026 23:55:17 +0000 Subject: [PATCH] enable safe TA page table teardown --- .../src/arch/x86/mm/paging.rs | 133 +++++++++++++----- .../src/host/per_cpu_variables.rs | 12 +- litebox_platform_lvbs/src/lib.rs | 122 ++++++---------- litebox_runner_lvbs/src/lib.rs | 131 ++++++----------- litebox_shim_optee/src/lib.rs | 33 +++-- 5 files changed, 204 insertions(+), 227 deletions(-) diff --git a/litebox_platform_lvbs/src/arch/x86/mm/paging.rs b/litebox_platform_lvbs/src/arch/x86/mm/paging.rs index 84686c777..aeb92a76f 100644 --- a/litebox_platform_lvbs/src/arch/x86/mm/paging.rs +++ b/litebox_platform_lvbs/src/arch/x86/mm/paging.rs @@ -18,6 +18,7 @@ use x86_64::{ CleanUp, FlagUpdateError, MapToError, PageTableFrameMapping, TranslateResult, UnmapError as X64UnmapError, }, + page_table::PageTableEntry, }, }, }; @@ -35,12 +36,12 @@ use crate::mm::{ #[cfg(not(test))] const TLB_SINGLE_PAGE_FLUSH_CEILING: usize = 33; -/// Bit position of the PML4 (level-4) index within a virtual address, for -/// x86-64 4-level paging: 12 page-offset bits + 9 bits each for P1-P3. -const PML4_SHIFT: u32 = 39; - -/// Mask for a 9-bit page-table index (512 entries per table). -const PML4_INDEX_MASK: u64 = 0x1FF; +const PAGE_SHIFT: usize = Size4KiB::SIZE.trailing_zeros() as usize; +const PAGE_TABLE_LEVEL_BITS: usize = 9; +const P2_SHIFT: usize = PAGE_SHIFT + PAGE_TABLE_LEVEL_BITS; +const P3_SHIFT: usize = P2_SHIFT + PAGE_TABLE_LEVEL_BITS; +const PML4_SHIFT: usize = P3_SHIFT + PAGE_TABLE_LEVEL_BITS; +const PML4_INDEX_MASK: u64 = (1 << PAGE_TABLE_LEVEL_BITS) - 1; /// PML4 index of the first VTL1-kernel slot (`PA + KERNEL_OFFSET`). /// @@ -291,30 +292,6 @@ impl X64PageTable<'_, M, ALIGN> { Ok(()) } - /// Clean up task-owned intermediate page table frames (P1-P3) for a task - /// page table that is being destroyed. - /// - /// # Safety - /// - /// The caller must ensure that: - /// - All user data frames have been released before calling this function (e.g., using `PageManager::release_memory()`) - /// - The page table is no longer active (not loaded in CR3) - pub(crate) unsafe fn cleanup_page_table_frames(&self) { - let mut allocator = PageTableAllocator::::new(); - // Task-owned slots span the VA below the kernel region, i.e., - // `0 ..= KERNEL_PML4_START * PML4_SLOT_SIZE - 1`. The kernel region at - // and above `KERNEL_PML4_START` is base-owned/shared. - let start = Page::::from_start_address(VirtAddr::new(0)).unwrap(); - let end = Page::::containing_address(VirtAddr::new(crate::KERNEL_OFFSET - 1)); - // Safety: The page table is being destroyed and will not be reused. - // This function crosses the non-canonical hole. - unsafe { - self.inner - .lock() - .clean_up_addr_range(Page::range_inclusive(start, end), &mut allocator); - } - } - pub(crate) unsafe fn remap_pages( &self, old_range: PageRange, @@ -499,8 +476,8 @@ impl X64PageTable<'_, M, ALIGN> { // COW lazy-enable was unimplemented, so granting WRITABLE via a later // fault would land in the unimplemented COW path and kill the task. // Install the writable PTE directly until COW (and shared frames) land. - // FIXME: when COW is implemented, restore the lazy-enable masking that - // was removed here so a R->RW mprotect defers WRITABLE to the fault path. + // FIXME: COW needs lazy write enable and refcounted user frames; + // task-table destruction currently assumes exclusive ownership. if flags != new_flags { match unsafe { inner.update_flags(page, (flags & !Self::MPROTECT_PTE_MASK) | new_flags) @@ -822,16 +799,100 @@ impl X64PageTable<'_, M, ALIGN> { } impl Drop for X64PageTable<'_, M, ALIGN> { - /// Deallocate the physical frame of the top-level page table + /// Reclaims owned user frames and private page tables. + /// + /// No TLB shootdown is needed because it is released after a non-PCID CR3 reload. #[allow(clippy::similar_names)] fn drop(&mut self) { let mut allocator = PageTableAllocator::::new(); - let p4_va = - core::ptr::from_mut::(self.inner.lock().level_4_table_mut()).cast::(); + let mut inner = self.inner.lock(); + let p4 = inner.level_4_table_mut(); + + // Skip shared VTL1-kernel PML4 entries. + for (p4_index, p4_entry) in p4.iter_mut().enumerate().take(KERNEL_PML4_START) { + let Ok(p3_frame) = p4_entry.frame() else { + p4_entry.set_unused(); + continue; + }; + let p3 = unsafe { &mut *frame_to_pointer::(p3_frame) }; + + for (p3_index, p3_entry) in p3.iter_mut().enumerate() { + if p3_entry.flags().contains(PageTableFlags::HUGE_PAGE) { + crate::debug_serial_println!("BUG: LVBS platform does not support huge pages"); + debug_assert!(false, "LVBS platform does not support huge pages"); + p3_entry.set_unused(); + continue; + } + let Ok(p2_frame) = p3_entry.frame() else { + p3_entry.set_unused(); + continue; + }; + let p2 = unsafe { &mut *frame_to_pointer::(p2_frame) }; + + for (p2_index, p2_entry) in p2.iter_mut().enumerate() { + if p2_entry.flags().contains(PageTableFlags::HUGE_PAGE) { + crate::debug_serial_println!( + "BUG: LVBS platform does not support huge pages" + ); + debug_assert!(false, "LVBS platform does not support huge pages"); + p2_entry.set_unused(); + continue; + } + let Ok(p1_frame) = p2_entry.frame() else { + p2_entry.set_unused(); + continue; + }; + let p1 = unsafe { &mut *frame_to_pointer::(p1_frame) }; + + for (p1_index, p1_entry) in p1.iter_mut().enumerate() { + // Indices >= 256 cannot fall in the user range. + let page_address = (p4_index << PML4_SHIFT) + | (p3_index << P3_SHIFT) + | (p2_index << P2_SHIFT) + | (p1_index << PAGE_SHIFT); + if (crate::USER_ADDR_MIN..crate::USER_ADDR_MAX).contains(&page_address) { + match p1_entry.frame() { + Ok(frame) => { + // Safety: task user leaf frames are exclusively owned. + unsafe { allocator.deallocate_frame(frame) }; + } + Err(_) if !p1_entry.is_unused() => { + crate::debug_serial_println!( + "BUG: leaking malformed task leaf during destruction" + ); + debug_assert!(false, "malformed task leaf during destruction"); + } + Err(_) => {} + } + } + p1_entry.set_unused(); + } + } + } + } + + let p4_va = core::ptr::from_mut::(p4).cast::(); let p4_pa = M::va_to_pa(VirtAddr::new(p4_va as u64)); + let start = Page::::containing_address(VirtAddr::new(0)); + let end = Page::::containing_address(VirtAddr::new(crate::KERNEL_OFFSET - 1)); + debug_assert_eq!(usize::from(end.p4_index()), KERNEL_PML4_START - 1); + + // Safety: all private leaves were cleared above. unsafe { - allocator.deallocate_frame(PhysFrame::containing_address(p4_pa)); + inner.clean_up_addr_range(Page::range_inclusive(start, end), &mut allocator); } + debug_assert!( + inner + .level_4_table() + .iter() + .take(KERNEL_PML4_START) + .all(PageTableEntry::is_unused), + "task page table dropped with live private mappings" + ); + drop(inner); + + // Safety: owned P4 frame. + unsafe { allocator.deallocate_frame(PhysFrame::containing_address(p4_pa)) }; } } diff --git a/litebox_platform_lvbs/src/host/per_cpu_variables.rs b/litebox_platform_lvbs/src/host/per_cpu_variables.rs index aec492aeb..f12e3e3e6 100644 --- a/litebox_platform_lvbs/src/host/per_cpu_variables.rs +++ b/litebox_platform_lvbs/src/host/per_cpu_variables.rs @@ -255,17 +255,19 @@ impl PerCpuVariables { .map(|(_, page_table)| Arc::clone(page_table)) } + /// Replaces and returns the active task page table. + /// /// # Safety /// /// CR3 must no longer reference the previous table. A new ID must match /// CR3. Interrupts must be disabled, and this must not run in exception context. - pub(crate) unsafe fn set_active_page_table( + pub(crate) unsafe fn replace_active_page_table( &self, page_table: Option<(usize, Arc>)>, - ) { - // Safety: Only this core accesses the field, interrupts are disabled, - // and the update cannot fault. - unsafe { *self.active_page_table.get() = page_table } + ) -> Option<(usize, Arc>)> { + // Safety: Core-local, IRQs are disabled, and the update cannot fault. + // Return the old owner so the caller can drop it after restoring IRQs. + unsafe { core::mem::replace(&mut *self.active_page_table.get(), page_table) } } } diff --git a/litebox_platform_lvbs/src/lib.rs b/litebox_platform_lvbs/src/lib.rs index 49986b2ae..ed9178d60 100644 --- a/litebox_platform_lvbs/src/lib.rs +++ b/litebox_platform_lvbs/src/lib.rs @@ -198,7 +198,7 @@ impl core::ops::Deref for PageTableHandle<'_> { /// Future work could implement KPTI-style isolation to reduce the kernel attack surface /// exposed to user TAs, mitigating potential side-channel attacks. pub struct PageTableManager { - /// The base page table, containing only VTL1 kernel mappings (no user-space). + /// The base (kernel) page table; lives for the kernel's lifetime and is never dropped. base_page_table: mm::PageTable, /// Cached physical frame of the base page table (for fast CR3 comparison). base_page_table_frame: PhysFrame, @@ -249,6 +249,18 @@ impl PageTableManager { ); } + /// Returns an owning handle to a registered task page table. + /// + /// The table remains alive while this handle exists. + pub fn task_page_table(&self, task_pt_id: usize) -> Result, Errno> { + if task_pt_id == BASE_PAGE_TABLE_ID { + return Err(Errno::EINVAL); + } + let task_pts = self.task_page_tables.read(); + let page_table = Arc::clone(task_pts.get(&task_pt_id).ok_or(Errno::ENOENT)?); + Ok(PageTableHandle::task(page_table)) + } + /// Returns the ID of the current page table based on the CR3 register. /// /// Returns `BASE_PAGE_TABLE_ID` (0) if the base page table is active, @@ -287,15 +299,16 @@ impl PageTableManager { /// after the switch (including the code being executed and stack) /// - No references to user-space memory are held across the switch pub unsafe fn load_base(&self) { - x86_64::instructions::interrupts::without_interrupts(|| { - // Ensure decreasing/dropping `Arc` for the previous page table (`set_active_page_table()`) - // only after switching CR3 (`mm::PageTable::load()`). + let previous = x86_64::instructions::interrupts::without_interrupts(|| { + // Replace the per-CPU owner only after CR3 stops referencing it. self.base_page_table.load(); with_per_cpu_variables(|pcv| { // Safety: CR3 now references the base page table and interrupts are disabled. - unsafe { pcv.set_active_page_table(None) } - }); + unsafe { pcv.replace_active_page_table(None) } + }) }); + // Last-owner reclamation must run with IRQs enabled. + drop(previous); } /// Loads the specified task page table by updating CR3. @@ -323,15 +336,16 @@ impl PageTableManager { Arc::clone(task_pts.get(&task_pt_id).ok_or(Errno::ENOENT)?) }; - x86_64::instructions::interrupts::without_interrupts(|| { - // Ensure decreasing/dropping `Arc` for the previous page table (`set_active_page_table()`) - // only after switching CR3 (`mm::PageTable::load()`). + let previous = x86_64::instructions::interrupts::without_interrupts(|| { + // Replace the per-CPU owner only after CR3 stops referencing it. pt.load(); with_per_cpu_variables(|pcv| { // Safety: CR3 now references `pt` and interrupts are disabled. - unsafe { pcv.set_active_page_table(Some((task_pt_id, pt))) } - }); + unsafe { pcv.replace_active_page_table(Some((task_pt_id, pt))) } + }) }); + // Last-owner reclamation must run with IRQs enabled. + drop(previous); Ok(()) } @@ -363,63 +377,25 @@ impl PageTableManager { Ok(task_pt_id) } - /// Deletes a task page table by its ID. - /// - /// This function: - /// 1. Clean up page table structure frames (P1-P3) - /// 2. Drop the page table (deallocating the top-level P4 frame) - /// - /// # Arguments - /// - /// * `task_pt_id` - The ID of the task page table to delete + /// Unregisters a task table; retained handles keep it alive. /// /// # Safety /// - /// The caller must ensure that: - /// - All user data frames have been released before calling this function - /// - No references or pointers to memory mapped by this page table are held after deletion - /// - /// # Returns + /// The caller must ensure final destruction is safe: user leaf frames are + /// exclusively owned and no access outlives all remaining handles. /// - /// - `Ok(())` if the page table was successfully deleted - /// - `Err(Errno::EINVAL)` if the page table ID is the base page table - /// - `Err(Errno::ENOENT)` if the page table ID does not exist - /// - `Err(Errno::EBUSY)` if the page table is active or has outstanding handles - pub unsafe fn delete_task_page_table(&self, task_pt_id: usize) -> Result<(), Errno> { + /// Returns `EINVAL` for the base ID and `ENOENT` if it is not registered. + pub unsafe fn unregister_task_page_table(&self, task_pt_id: usize) -> Result<(), Errno> { if task_pt_id == BASE_PAGE_TABLE_ID { return Err(Errno::EINVAL); } - let mut task_pts = self.task_page_tables.write(); - - // Fast path for the page table active on this core. - let (cr3_frame, _) = x86_64::registers::control::Cr3::read(); - let cr3_id: usize = cr3_frame.start_address().as_u64().trunc(); - if cr3_id == task_pt_id { - return Err(Errno::EBUSY); - } - - if let Some(pt) = task_pts.remove(&task_pt_id) { - // An active CR3 retains a per-CPU Arc. - let pt = match Arc::try_unwrap(pt) { - Ok(pt) => pt, - Err(pt) => { - task_pts.insert(task_pt_id, pt); - return Err(Errno::EBUSY); - } - }; - drop(task_pts); - - // Safety: successful unwrap proves the table is neither active nor - // borrowed. Kernel slots are base-owned and must not be freed. - unsafe { - pt.cleanup_page_table_frames(); - } - // The PageTable's Drop impl will deallocate the top-level (P4) frame - Ok(()) - } else { - Err(Errno::ENOENT) - } + let pt = { + let mut task_pts = self.task_page_tables.write(); + task_pts.remove(&task_pt_id).ok_or(Errno::ENOENT)? + }; + drop(pt); + Ok(()) } } @@ -725,27 +701,17 @@ impl LinuxKernel { self.page_table_manager.create_task_page_table() } - /// Deletes a task page table by its ID. - /// - /// This function: - /// 1. Cleans up page table structure frames (P1-P3) - /// 2. Drops the page table (deallocating the top-level P4 frame) + /// Unregisters a task table; retained handles keep it alive. /// /// # Safety /// - /// The caller must ensure that: - /// - All user data frames have been released before calling this function - /// - No references or pointers to memory mapped by this page table are held after deletion - /// - /// # Returns - /// - /// - `Ok(())` if successful - /// - `Err(Errno::EINVAL)` if the page table is the base page table - /// - `Err(Errno::ENOENT)` if the page table doesn't exist - /// - `Err(Errno::EBUSY)` if the page table is active or has outstanding handles - pub unsafe fn delete_task_page_table(&self, task_pt_id: usize) -> Result<(), Errno> { - // Safety: caller guarantees no dangling references - unsafe { self.page_table_manager.delete_task_page_table(task_pt_id) } + /// See [`PageTableManager::unregister_task_page_table`]. + pub unsafe fn unregister_task_page_table(&self, task_pt_id: usize) -> Result<(), Errno> { + // Safety: the caller upholds the manager's destruction requirements. + unsafe { + self.page_table_manager + .unregister_task_page_table(task_pt_id) + } } /// Switch to the specified page table. diff --git a/litebox_runner_lvbs/src/lib.rs b/litebox_runner_lvbs/src/lib.rs index 442eba6eb..4a7a7444c 100644 --- a/litebox_runner_lvbs/src/lib.rs +++ b/litebox_runner_lvbs/src/lib.rs @@ -420,25 +420,6 @@ unsafe fn switch_to_task_page_table( } } -/// Deletes a task-specific page table. -/// -/// # Safety -/// -/// The caller must ensure that no references or pointers to memory mapped -/// by this page table are held after deletion. -#[inline] -unsafe fn delete_task_page_table( - platform: &'static Platform, - task_pt_id: usize, -) -> Result<(), OpteeSmcReturnCode> { - // Safety: caller guarantees no dangling references - unsafe { - platform - .delete_task_page_table(task_pt_id) - .map_err(|_| OpteeSmcReturnCode::EBadCmd) - } -} - /// 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 @@ -465,29 +446,16 @@ impl Drop for TaskPageTableGuard { } } -/// Tears down a TA's memory mappings and page table. -/// -/// This performs the following steps in order: -/// 1. Release user-space memory mappings in the TA's page table -/// 2. Switch to the base page table -/// 3. Delete the TA's page table -/// -/// # Safety +/// Switches to base and unregisters the task table. /// -/// The caller must ensure that no references to user-space memory mapped by -/// this task's page table are held after this call. -unsafe fn teardown_ta_page_table( - platform: &'static Platform, - shim: &litebox_shim_optee::OpteeShim, - task_pt_id: usize, -) { - unsafe { - // this function unmaps/deallocates user pages in the **active** page table, so we must - // still be on the TA's page table. - shim.release_user_mappings(); - switch_to_base_page_table(platform); - // Now delete the TA's page table without memory leak. - let _ = delete_task_page_table(platform, task_pt_id); +/// All user-memory accesses on this core must be complete. +fn teardown_ta_page_table(platform: &'static Platform, task_pt_id: usize) { + // Safety: no user-memory references remain. + unsafe { switch_to_base_page_table(platform) }; + // Safety: OP-TEE has no shared/COW user-page mappings, and active access + // retains an owning page-table handle. + if let Err(error) = unsafe { platform.unregister_task_page_table(task_pt_id) } { + debug_serial_println!("Failed to unregister task page table {task_pt_id:#x}: {error:?}"); } } @@ -744,11 +712,7 @@ fn open_session_single_instance( debug_serial_println!("Single-instance TA panicked during OpenSession, cleaning up"); 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(platform, instance.shim(), task_pt_id); - }; + teardown_ta_page_table(platform, 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 @@ -783,11 +747,7 @@ 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: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. - unsafe { - teardown_ta_page_table(platform, instance.shim(), task_pt_id); - }; + teardown_ta_page_table(platform, task_pt_id); } else { // The session id is forgotten (never recycled), so the token's drop // won't clear the recorded identity. Remove the client identity here. @@ -841,16 +801,30 @@ fn open_session_new_instance( let task_pt_id = create_task_page_table(platform)?; debug_serial_println!("Created task page table ID: {}", task_pt_id); - let _task_pt_guard = TaskPageTableGuard::enter(platform, 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(platform, task_pt_id) }; + let page_table = platform + .page_table_manager() + .task_page_table(task_pt_id) + .map_err(|_| OpteeSmcReturnCode::EBadCmd) + .inspect_err(|error| { + debug_serial_println!("Failed to retain task page table: {error:?}"); + teardown_ta_page_table(platform, task_pt_id); + })?; + + let _task_pt_guard = TaskPageTableGuard::enter(platform, task_pt_id).inspect_err(|error| { + debug_serial_println!("Failed to activate task page table: {error:?}"); + teardown_ta_page_table(platform, task_pt_id); + })?; + + let shim = shim.retain_page_table(page_table).ok_or_else(|| { + debug_serial_println!("BUG: failed to retain task page table"); + debug_assert!(false, "failed to retain task page table"); + teardown_ta_page_table(platform, task_pt_id); + OpteeSmcReturnCode::ENotAvail })?; // Load ldelf and TA - Box immediately to keep at fixed heap address let loaded_program = Box::new(shim.load_ldelf(LDELF_BINARY, ta_uuid).map_err(|_| { - // Safety: We are about to tear down this TA instance; - // no references to user-space memory will be held afterwards. - unsafe { teardown_ta_page_table(platform, &shim, task_pt_id) }; + teardown_ta_page_table(platform, task_pt_id); OpteeSmcReturnCode::ENomem })?); @@ -892,10 +866,7 @@ fn open_session_new_instance( Some(ta_req_info), None, ); - - // 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(platform, &shim, task_pt_id) }; + teardown_ta_page_table(platform, task_pt_id); write_result?; return Ok(()); @@ -906,9 +877,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. - unsafe { teardown_ta_page_table(platform, &shim, task_pt_id) }; + teardown_ta_page_table(platform, task_pt_id); OpteeSmcReturnCode::EBadCmd })?; let memref_addresses = loaded_program @@ -923,9 +892,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. - unsafe { teardown_ta_page_table(platform, &shim, task_pt_id) }; + teardown_ta_page_table(platform, task_pt_id); OpteeSmcReturnCode::EBadCmd })?; @@ -940,17 +907,13 @@ 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. - unsafe { teardown_ta_page_table(platform, &shim, task_pt_id) }; + teardown_ta_page_table(platform, 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. - unsafe { teardown_ta_page_table(platform, &shim, task_pt_id) }; + teardown_ta_page_table(platform, task_pt_id); OpteeSmcReturnCode::EBadAddr })?; @@ -977,10 +940,7 @@ fn open_session_new_instance( Some(ta_req_info), Some(&memref_addresses), ); - - // 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(platform, &shim, task_pt_id) }; + teardown_ta_page_table(platform, task_pt_id); write_result?; return Ok(()); @@ -1000,9 +960,7 @@ fn open_session_new_instance( Some(&memref_addresses), ) .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(platform, &shim, task_pt_id) }; + teardown_ta_page_table(platform, task_pt_id); })?; // Success: register the new session with the manager. @@ -1150,12 +1108,7 @@ fn handle_invoke_command( } 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(platform, instance.shim(), task_pt_id); - }; + teardown_ta_page_table(platform, task_pt_id); debug_serial_println!( "InvokeCommand: cleaned up dead TA instance, task_pt_id={}", @@ -1263,11 +1216,7 @@ fn handle_close_session( 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(platform, instance.shim(), task_pt_id); - }; + teardown_ta_page_table(platform, task_pt_id); debug_serial_println!( "CloseSession complete: deleted task_pt_id={} (last session)", diff --git a/litebox_shim_optee/src/lib.rs b/litebox_shim_optee/src/lib.rs index 4d96ccd24..b9d271592 100644 --- a/litebox_shim_optee/src/lib.rs +++ b/litebox_shim_optee/src/lib.rs @@ -11,7 +11,7 @@ extern crate alloc; use crate::loader::elf::ElfLoaderError; use crate::syscalls::pta::PseudoTa; use aes::{Aes128, Aes192, Aes256}; -use alloc::{sync::Arc, vec}; +use alloc::{boxed::Box, sync::Arc, vec}; use core::cell::Cell; use ctr::Ctr128BE; use hashbrown::{HashMap, HashSet}; @@ -184,6 +184,7 @@ impl OpteeShimBuilder { _litebox: self.litebox, ta_uuid_map: ta_uuid_map(), pta_busy: spin::mutex::SpinMutex::new(HashSet::new()), + retained_page_table: None, }); OpteeShim(global) } @@ -213,6 +214,9 @@ struct GlobalState { /// blocking/queuing the caller until the PTA is free. We currently reject /// instead of serialize; revisit if a PTA needs true serialization. pta_busy: spin::mutex::SpinMutex>, + /// Declared last so the retained task table drops after the shim state. + // TODO: Replace type erasure with a typed platform page-table handle. + retained_page_table: Option>, } impl GlobalState { @@ -285,6 +289,17 @@ impl Clone for OpteeShim { } impl OpteeShim { + /// Retains the task page table before the shim is shared. + #[must_use] + pub fn retain_page_table(mut self, page_table: T) -> Option { + let global = Arc::get_mut(&mut self.0)?; + if global.retained_page_table.is_some() { + return None; + } + global.retained_page_table = Some(Box::new(page_table)); + Some(self) + } + /// Load the given `ldelf` binary into memory while making it ready to load the TA binary specified /// by `ta_uuid` (and optionally `ta_bin`). /// @@ -368,22 +383,6 @@ impl OpteeShim { pub fn get_ta_bin(&self, ta_uuid: &TeeUuid) -> Option> { self.0.get_ta_bin(ta_uuid) } - - /// Release all user-space memory mappings owned by this shim instance. - /// - /// This must be called before switching to the base page table and deleting - /// the task page table so that every mapped physical page is properly freed. - /// - /// # Safety - /// - /// The caller must ensure that no references to the released memory regions - /// are held after this call. - pub unsafe fn release_user_mappings(&self) { - let release = |_r: core::ops::Range, _vm: litebox::mm::linux::VmFlags| true; - unsafe { - let _ = self.page_manager().release_memory(release); - } - } } impl OpteeShimEntrypoints {