diff --git a/Cargo.lock b/Cargo.lock index d32914691..573bb6e1e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1528,6 +1528,7 @@ dependencies = [ "litebox", "litebox_common_linux", "litebox_common_optee", + "litebox_util_log", "seccompiler", "spin 0.9.8", "syscalls", diff --git a/Cargo.toml b/Cargo.toml index 258dd51e9..a0c346387 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -47,6 +47,12 @@ default-members = [ # We exclude `litebox_runner_lvbs` from `default-members` because it requires # a custom target and `-Z build-std`, which requires a nightly toolchain. +# Note: `panic = "abort"` is intentionally not set; the userland platforms rely +# on unwinding across `extern "C-unwind"` boundaries. +[profile.release] +lto = "thin" +codegen-units = 1 + # Introduce all the pedantic clippy lints and remove ones I (jayb) think are # pushing it too far; this way we get something even further than default clippy # but not _too_ ridiculous. diff --git a/dev_tests/src/ratchet.rs b/dev_tests/src/ratchet.rs index 0e0b17f1e..0c11a580d 100644 --- a/dev_tests/src/ratchet.rs +++ b/dev_tests/src/ratchet.rs @@ -36,7 +36,12 @@ fn ratchet_globals() -> Result<()> { ("dev_bench/", 1), ("litebox/", 9), ("litebox_platform_linux_kernel/", 6), - ("litebox_platform_linux_userland/", 5), + // `litebox_platform_linux_userland` includes the process-wide + // alternate-signal-stack pool (`ALT_STACK_POOL`), the + // `HAVE_FSGSBASE` capability flag read from the guest/host + // transition assembly, and the `NUM_CPUS` cache seeded before + // seccomp installation — all deliberate process-wide state. + ("litebox_platform_linux_userland/", 8), ("litebox_platform_lvbs/", 23), ("litebox_platform_multiplex/", 1), ("litebox_platform_windows_userland/", 8), diff --git a/litebox/src/fd/mod.rs b/litebox/src/fd/mod.rs index 6f988ce76..9661858fb 100644 --- a/litebox/src/fd/mod.rs +++ b/litebox/src/fd/mod.rs @@ -3,9 +3,12 @@ //! File descriptors used in LiteBox -#![expect( - dead_code, - reason = "still under development, remove before merging PR" +#![cfg_attr( + not(test), + expect( + dead_code, + reason = "still under development, remove before merging PR" + ) )] use alloc::sync::Arc; @@ -22,9 +25,54 @@ use crate::utilities::anymap::AnyMap; #[cfg(test)] mod tests; +/// A bitmap tracking which slots of a `Vec>`-style table are occupied, so that the +/// lowest free slot (POSIX lowest-fd semantics) can be found via an O(n/64) word scan instead of +/// an O(n) entry scan. +/// +/// Invariant: bit `i` is set iff slot `i` is occupied. In particular, bits at indices at or beyond +/// the table's length must be clear, so that [`Self::allocate`] never returns an index more than +/// one past the table's current length. +#[derive(Default)] +struct FreeSlotBitmap { + words: Vec, +} + +impl FreeSlotBitmap { + /// Returns the lowest free slot index, marking it as occupied. + fn allocate(&mut self) -> usize { + for (i, word) in self.words.iter_mut().enumerate() { + if *word != u64::MAX { + let bit = (!*word).trailing_zeros() as usize; + *word |= 1u64 << bit; + return i * 64 + bit; + } + } + self.words.push(1); + (self.words.len() - 1) * 64 + } + + /// Marks `idx` as occupied, growing the bitmap as needed. + fn set_used(&mut self, idx: usize) { + let word = idx / 64; + if word >= self.words.len() { + self.words.resize(word + 1, 0); + } + self.words[word] |= 1u64 << (idx % 64); + } + + /// Marks `idx` as free. + /// + /// `idx` must have previously been marked occupied (which guarantees the backing word exists). + fn set_free(&mut self, idx: usize) { + self.words[idx / 64] &= !(1u64 << (idx % 64)); + } +} + /// Storage of file descriptors and their entries. pub struct Descriptors { entries: Vec>>, + /// Tracks which `entries` slots are occupied; see [`FreeSlotBitmap`]. + free_slots: FreeSlotBitmap, } impl Descriptors { @@ -33,7 +81,20 @@ impl Descriptors { /// This is expected to be invoked only by [`crate::LiteBox`]'s creation method, and should not /// be invoked anywhere else in the codebase. pub(crate) fn new_from_litebox_creation() -> Self { - Self { entries: vec![] } + Self { + entries: vec![], + free_slots: FreeSlotBitmap::default(), + } + } + + /// Allocate the lowest free slot in `entries`, extending the table if needed. + fn allocate_slot(&mut self) -> usize { + let idx = self.free_slots.allocate(); + debug_assert!(idx <= self.entries.len()); + if idx == self.entries.len() { + self.entries.push(None); + } + idx } /// Insert `entry` into the descriptor table, returning an `OwnedFd` to this entry. @@ -50,15 +111,10 @@ impl Descriptors { entry: alloc::boxed::Box::new(entry.into()), metadata: AnyMap::new(), }; - let idx = self - .entries - .iter() - .position(Option::is_none) - .unwrap_or_else(|| { - self.entries.push(None); - self.entries.len() - 1 - }); - let old = self.entries[idx].replace(IndividualEntry::new(Arc::new(RwLock::new(entry)))); + let idx = self.allocate_slot(); + let old = self.entries[idx].replace(IndividualEntry::new::(Arc::new( + RwLock::new(entry), + ))); assert!(old.is_none()); TypedFd { _phantom: PhantomData, @@ -84,17 +140,10 @@ impl Descriptors { &mut self, fd: &TypedFd, ) -> Option> { - let idx = self - .entries - .iter() - .position(Option::is_none) - .unwrap_or_else(|| { - self.entries.push(None); - self.entries.len() - 1 - }); - let new_ind_entry = IndividualEntry::new(Arc::clone( + let new_ind_entry = IndividualEntry::new::(Arc::clone( &self.entries[fd.x.as_usize()?].as_ref().unwrap().x, )); + let idx = self.allocate_slot(); let old = self.entries[idx].replace(new_ind_entry); assert!(old.is_none()); Some(TypedFd { @@ -113,9 +162,11 @@ impl Descriptors { &mut self, fd: &TypedFd, ) -> Option { - let Some(old) = self.entries[fd.x.as_usize()?].take() else { + let idx = fd.x.as_usize()?; + let Some(old) = self.entries[idx].take() else { unreachable!(); }; + self.free_slots.set_free(idx); fd.x.mark_as_closed(); Arc::into_inner(old.x) .map(RwLock::into_inner) @@ -143,6 +194,7 @@ impl Descriptors { if Arc::strong_count(&old.x) == 1 { // Unique, so we can just return it if allowed. if can_close_immediately(old.x.read().as_subsystem::()) { + self.free_slots.set_free(idx); fd.x.mark_as_closed(); let entry = Arc::into_inner(old.x) .map(RwLock::into_inner) @@ -241,17 +293,16 @@ impl Descriptors { ) -> impl Iterator)> { self.entries.iter().enumerate().filter_map(|(i, entry)| { entry.as_ref().and_then(|e| { - let entry = e.read(); - if entry.matches_subsystem::() { - Some(( - InternalFd { - raw: i.try_into().unwrap(), - }, - crate::sync::RwLockReadGuard::map(entry, |e| e.as_subsystem::()), - )) - } else { - None + if !e.matches_subsystem::() { + return None; } + let entry = e.read(); + Some(( + InternalFd { + raw: i.try_into().unwrap(), + }, + crate::sync::RwLockReadGuard::map(entry, |e| e.as_subsystem::()), + )) }) }) } @@ -270,11 +321,10 @@ impl Descriptors { > { self.entries.iter().enumerate().filter_map(|(i, entry)| { entry.as_ref().and_then(|e| { - if !e.read().matches_subsystem::() { + if !e.matches_subsystem::() { return None; } let entry = e.write(); - assert!(entry.matches_subsystem::()); Some(( InternalFd { raw: i.try_into().unwrap(), @@ -356,12 +406,11 @@ impl Descriptors { Subsystem: FdEnabledSubsystem, F: FnOnce(&mut Subsystem::Entry) -> R, { - let mut entry = self.entries[usize::try_from(internal_fd.raw).unwrap()] + let ind_entry = self.entries[usize::try_from(internal_fd.raw).unwrap()] .as_ref() - .unwrap() - .write(); - if entry.matches_subsystem::() { - Some(f(entry.as_subsystem_mut::())) + .unwrap(); + if ind_entry.matches_subsystem::() { + Some(f(ind_entry.write().as_subsystem_mut::())) } else { None } @@ -573,6 +622,8 @@ pub(crate) enum CloseResult { pub struct RawDescriptorStorage { /// Stored FDs are used to provide raw integer values in a safer way. stored_fds: Vec>, + /// Tracks which `stored_fds` slots are occupied; see [`FreeSlotBitmap`]. + free_slots: FreeSlotBitmap, } struct StoredFd { @@ -596,7 +647,10 @@ impl RawDescriptorStorage { #[expect(clippy::new_without_default)] /// Create a new raw descriptor store. pub fn new() -> Self { - Self { stored_fds: vec![] } + Self { + stored_fds: vec![], + free_slots: FreeSlotBitmap::default(), + } } /// Get the corresponding integer value of the provided `fd`. @@ -610,11 +664,9 @@ impl RawDescriptorStorage { &mut self, fd: TypedFd, ) -> usize { - let ret = self - .stored_fds - .iter() - .position(Option::is_none) - .unwrap_or(self.stored_fds.len()); + // `allocate` already marks the slot as used; `fd_into_specific_raw_integer` marking it + // again is a harmless no-op. + let ret = self.free_slots.allocate(); let success = self.fd_into_specific_raw_integer(fd, ret); assert!(success); ret @@ -625,9 +677,11 @@ impl RawDescriptorStorage { /// This is similar to [`Self::fd_into_raw_integer`] except that it specifies a specific FD to /// be stored into. /// - /// Will return with `true` iff it succeeds (i.e., nothing else was using that raw integer FD). - /// If you want to replace a used slot, you must first consume that slot via - /// [`Self::fd_consume_raw_integer`]. + /// Will return with `true` iff it succeeds. Returns `false` both when the slot is already + /// occupied (if you want to replace a used slot, you must first consume that slot via + /// [`Self::fd_consume_raw_integer`]) and when `raw_fd` is far enough beyond the current + /// backing storage that granting it would force a large, guest-triggerable allocation (see + /// below) — callers must treat both as an ordinary failure to insert, not assume success. #[must_use] #[expect( clippy::missing_panics_doc, @@ -641,12 +695,13 @@ impl RawDescriptorStorage { // TODO(jayb): Should we be storing things via a HashMap to make sure this operation cannot // be too expensive if someone tries to store into a large raw FD? // - // If this assertion failure is hit in practice, we might need to be more defensive via the - // HashMap, rather than just silently allow big growth - assert!( - raw_fd < self.stored_fds.len() + 256, - "explicit upper bound restriction for now; see implementation details" - ); + // `stored_fds` is a dense, index-by-fd array; growing it to fit an arbitrary raw_fd + // (e.g. a guest calling `dup2(fd, 1_000_000)`, valid up to `RLIMIT_NOFILE`) would force + // an allocation sized to the requested fd number. Decline rather than growing unbounded + // or panicking — this is a real, guest-reachable path, not just a defensive invariant. + if raw_fd >= self.stored_fds.len() + 256 { + return false; + } if self.stored_fds.get(raw_fd).is_some_and(Option::is_some) { // There's already something at this slot. return false; @@ -656,6 +711,7 @@ impl RawDescriptorStorage { } let old = self.stored_fds[raw_fd].replace(StoredFd::new(fd)); assert!(old.is_none()); + self.free_slots.set_used(raw_fd); true } @@ -684,6 +740,7 @@ impl RawDescriptorStorage { let underlying = self.stored_fds[fd].take(); debug_assert!(underlying.is_some()); drop(underlying); + self.free_slots.set_free(fd); Ok(ret) } @@ -807,6 +864,10 @@ pub enum MetadataError { struct IndividualEntry { x: Arc>, metadata: AnyMap, + /// The `TypeId` of the subsystem's entry type, cached at insertion time so that subsystem + /// filtering (e.g., in [`Descriptors::iter`]) does not need to lock the entry. This is valid + /// because an entry's type can never change after insertion. + subsystem_entry_type_id: core::any::TypeId, } impl core::ops::Deref for IndividualEntry { type Target = Arc>; @@ -815,12 +876,19 @@ impl core::ops::Deref for IndividualEntry

IndividualEntry { - fn new(x: Arc>) -> Self { + fn new(x: Arc>) -> Self { Self { x, metadata: AnyMap::new(), + subsystem_entry_type_id: core::any::TypeId::of::(), } } + + /// Check if this entry matches the specified subsystem, without locking the entry. + #[must_use] + fn matches_subsystem(&self) -> bool { + self.subsystem_entry_type_id == core::any::TypeId::of::() + } } /// A crate-internal entry for a descriptor. @@ -830,12 +898,6 @@ pub(crate) struct DescriptorEntry { } impl DescriptorEntry { - /// Check if this entry matches the specified subsystem - #[must_use] - fn matches_subsystem(&self) -> bool { - core::any::TypeId::of::() == core::any::Any::type_id(self.entry.as_ref()) - } - /// Obtains `self` as the subsystem's entry type. /// /// # Panics diff --git a/litebox/src/fd/tests.rs b/litebox/src/fd/tests.rs index 04a482b44..4f0621112 100644 --- a/litebox/src/fd/tests.rs +++ b/litebox/src/fd/tests.rs @@ -116,6 +116,86 @@ fn test_with_entry() { }); } +#[test] +fn test_lowest_slot_reuse() { + let litebox = litebox(); + let mut descriptors = litebox.descriptor_table_mut(); + + // Span more than one bitmap word to exercise the word-scan logic. + let fds: Vec> = (0..70) + .map(|i| { + let fd = descriptors.insert(MockEntry { + data: i.to_string(), + }); + assert_eq!(fd.as_internal_fd().raw, i); + fd + }) + .collect(); + + // Free a couple of slots (one per bitmap word) and check that the lowest + // free slot is always reused first (POSIX lowest-fd semantics). + for &idx in &[3usize, 65] { + assert!(descriptors.remove(&fds[idx]).is_some()); + } + let fd: TypedFd = descriptors.insert(MockEntry { + data: "reused".to_string(), + }); + assert_eq!(fd.as_internal_fd().raw, 3); + let fd2 = descriptors.duplicate(&fd).unwrap(); + assert_eq!(fd2.as_internal_fd().raw, 65); + // All lower slots are in use again, so the next insert extends the table. + let fd3: TypedFd = descriptors.insert(MockEntry { + data: "appended".to_string(), + }); + assert_eq!(fd3.as_internal_fd().raw, 70); + + // Clean up: close everything out (note `fd`/`fd2` share an entry, so only the + // last removal returns it). + assert!(descriptors.remove(&fd).is_none()); + assert!(descriptors.remove(&fd2).is_some()); + assert!(descriptors.remove(&fd3).is_some()); + for (i, fd) in fds.into_iter().enumerate() { + if i != 3 && i != 65 { + assert!(descriptors.remove(&fd).is_some()); + } + } +} + +#[test] +fn test_raw_integer_lowest_reuse() { + let litebox = litebox(); + let mut descriptors = litebox.descriptor_table_mut(); + let mut rds = super::RawDescriptorStorage::new(); + + let raw_fds: Vec = (0..3) + .map(|i| { + let fd: TypedFd = descriptors.insert(MockEntry { + data: i.to_string(), + }); + rds.fd_into_raw_integer(fd) + }) + .collect(); + assert_eq!(raw_fds, vec![0, 1, 2]); + + // Consuming a raw fd frees its integer for reuse, lowest-first. + let _ = rds.fd_consume_raw_integer::(1).unwrap(); + let fd: TypedFd = descriptors.insert(MockEntry { + data: "reused".to_string(), + }); + assert_eq!(rds.fd_into_raw_integer(fd), 1); + + // A specifically-placed fd is skipped over by subsequent allocations. + let fd: TypedFd = descriptors.insert(MockEntry { + data: "specific".to_string(), + }); + assert!(rds.fd_into_specific_raw_integer(fd, 4)); + let fd: TypedFd = descriptors.insert(MockEntry { + data: "next".to_string(), + }); + assert_eq!(rds.fd_into_raw_integer(fd), 3); + assert!(rds.is_alive(4)); +} + #[test] fn test_fd_raw_integer() { let litebox = litebox(); diff --git a/litebox/src/fs/resolver.rs b/litebox/src/fs/resolver.rs index 967e589c8..23c05420a 100644 --- a/litebox/src/fs/resolver.rs +++ b/litebox/src/fs/resolver.rs @@ -256,6 +256,8 @@ impl>, ) -> Result<(), PathError> { + // `idx` is only used when `debug_assertions` is enabled. + #[cfg_attr(not(debug_assertions), expect(unused_variables))] for (idx, walked) in outcome.components.iter().enumerate() { match &walked.permissions { PermissionCheck::ByBackend => {} diff --git a/litebox/src/mm/linux.rs b/litebox/src/mm/linux.rs index 951519d46..4b7349220 100644 --- a/litebox/src/mm/linux.rs +++ b/litebox/src/mm/linux.rs @@ -237,7 +237,7 @@ impl core::ops::Add for NonZeroPageSize { type Output = Option; fn add(self, rhs: usize) -> Self::Output { - NonZeroPageSize::new(self.size + rhs) + NonZeroPageSize::new(self.size.checked_add(rhs)?) } } @@ -387,8 +387,10 @@ impl + 'static, const ALIGN: usize> Vmem /// If `anonymous_only` is true and any part of the range is non‑anonymous (i.e., file‑backed), /// returns `Err(VmemResetError::FileBacked)`. /// - /// The current implementation effectively re-inserts the mapping with the same - /// `VmArea` properties, which will cause the pages to be unmapped and mapped again. + /// The current implementation discards the pages in place when the platform supports it + /// (see [`PageManagementProvider::discard_pages`]); otherwise it re-inserts the mapping + /// with the same `VmArea` properties, which will cause the pages to be unmapped and + /// mapped again. /// /// # Panics /// @@ -420,6 +422,13 @@ impl + 'static, const ALIGN: usize> Vmem } let start = r.start.max(range.start); let end = r.end.min(range.end); + // Fast path: discard the pages in place if the platform supports it. Anonymous + // guest mappings are backed by anonymous platform pages, so an in-place discard + // (e.g., host `madvise(MADV_DONTNEED)`) yields the same zero-fill semantics as + // re-creating the mapping, without any change to the mapping itself. + if unsafe { self.platform.discard_pages(start..end) }.is_ok() { + continue; + } let new_range = PageRange::new(start, end).unwrap(); unsafe { self.insert_mapping(new_range, vma, false, FixedAddressBehavior::Replace) } .expect("failed to reset pages"); @@ -559,13 +568,16 @@ impl + 'static, const ALIGN: usize> Vmem vma: VmArea, flags: CreatePagesFlags, ) -> Result, AllocationError> { + // `length` is guest-controlled (e.g. `mmap`'s requested length, + // page-rounded); reject rather than silently wrap when it is too + // large to reserve space after, or to form a valid range at all. let total_length = (length + if flags.contains(CreatePagesFlags::ENSURE_SPACE_AFTER) { DEFAULT_RESERVED_SPACE_SIZE } else { 0 }) - .unwrap(); + .ok_or(AllocationError::OutOfMemory)?; let new_addr = self .get_unmmaped_area( suggested_address, @@ -574,7 +586,13 @@ impl + 'static, const ALIGN: usize> Vmem ) .ok_or(AllocationError::OutOfMemory)?; // new_addr must be ALIGN aligned - let new_range = PageRange::new(new_addr, new_addr + length.as_usize()).unwrap(); + let new_range = PageRange::new( + new_addr, + new_addr + .checked_add(length.as_usize()) + .ok_or(AllocationError::OutOfMemory)?, + ) + .ok_or(AllocationError::OutOfMemory)?; unsafe { self.insert_mapping( new_range, @@ -618,7 +636,21 @@ impl + 'static, const ALIGN: usize> Vmem .get_key_value(&range.start) .ok_or(VmemResizeError::NotExist(range.start))?; - let new_end = range.start + new_size.as_usize(); + // `new_size` is guest-controlled (from `mremap`'s `new_size` argument, + // only rounded up to a page boundary — an already-page-aligned huge + // value passes through unchanged). A raw add here would wrap and let + // a guest pick `new_end` anywhere in the address space by choosing + // `new_size` relative to `range.start`, causing the shrink branch + // below to unmap an attacker-chosen range instead of a suffix of this + // mapping. + let new_end = + range + .start + .checked_add(new_size.as_usize()) + .ok_or(VmemResizeError::InvalidAddr { + range: range.clone(), + addr: usize::MAX, + })?; match new_end.cmp(&range.end) { core::cmp::Ordering::Equal => { // no change @@ -755,37 +787,69 @@ impl + 'static, const ALIGN: usize> Vmem return Err(VmemProtectError::InvalidRange(range)); } + // All intersections change to the same `permissions`, so coalesce adjacent areas into + // contiguous runs, each applied with a single (expensive) platform permission update. + let mut run: Vec<(usize, usize, VmArea)> = Vec::new(); for (start, end, vma) in mappings_to_change { if vma.flags & VmFlags::VM_ACCESS_FLAGS == flags { + // Already has the target permissions; ends the current contiguous run. + unsafe { self.commit_protect_run(&mut run, &range, permissions, flags) }?; continue; } // flags >> 4 shift VM_MAY% in place of VM_% // turning on VM_% requires VM_MAY% if (!(vma.flags.bits() >> 4) & flags.bits()) & VmFlags::VM_ACCESS_FLAGS.bits() != 0 { + // Areas preceding the failing one are still changed, as Linux does. + unsafe { self.commit_protect_run(&mut run, &range, permissions, flags) }?; return Err(VmemProtectError::NoAccess { old: vma.flags, new: flags, }); } + if run + .last() + .is_some_and(|&(_, prev_end, _)| prev_end != start) + { + // Not adjacent to the pending run; apply the pending run first. + unsafe { self.commit_protect_run(&mut run, &range, permissions, flags) }?; + } + run.push((start, end, vma)); + } + unsafe { self.commit_protect_run(&mut run, &range, permissions, flags) }?; + Ok(()) + } + + /// Apply a pending contiguous run of areas collected by [`Self::protect_mapping`]: issue one + /// platform permission update covering the whole run, then update the bookkeeping of each + /// affected area. Drains `run`; does nothing if it is empty. + /// + /// # Safety + /// + /// Same as [`Self::protect_mapping`]. + unsafe fn commit_protect_run( + &mut self, + run: &mut Vec<(usize, usize, VmArea)>, + request: &Range, + permissions: MemoryRegionPermissions, + flags: VmFlags, + ) -> Result<(), VmemProtectError> { + let (Some(&(first_start, ..)), Some(&(_, last_end, _))) = (run.first(), run.last()) else { + return Ok(()); + }; + // The run's intersection with the request is page aligned. + let host_range = request.start.max(first_start)..request.end.min(last_end); + unsafe { self.platform.update_permissions(host_range, permissions) } + .map_err(VmemProtectError::ProtectError)?; + + for (start, end, vma) in run.drain(..) { self.vmas.remove(start..end); - let intersection = range.start.max(start)..range.end.min(end); - // split r into three parts: before, intersection, and after + let intersection = request.start.max(start)..request.end.min(end); + // split the area into three parts: before, intersection, and after let before = start..intersection.start; let after = intersection.end..end; let new_flags = (vma.flags & !VmFlags::VM_ACCESS_FLAGS) | flags; - // `intersection` is page aligned. - unsafe { - self.platform - .update_permissions(intersection.clone(), permissions) - } - .map_err(|e| { - // restore the original mapping - self.vmas.insert(start..end, vma); - VmemProtectError::ProtectError(e) - })?; - self.vmas.insert( intersection, VmArea { @@ -800,7 +864,6 @@ impl + 'static, const ALIGN: usize> Vmem self.vmas.insert(after, vma); } } - Ok(()) } diff --git a/litebox/src/mm/mod.rs b/litebox/src/mm/mod.rs index a46b3c855..5f050428e 100644 --- a/litebox/src/mm/mod.rs +++ b/litebox/src/mm/mod.rs @@ -253,6 +253,35 @@ where } } + /// Create pages with the given permissions, without an initialization callback. + /// + /// Unlike the other `create_*_pages` methods, no initialization `op` runs, so the pages are + /// mapped directly with their final `perms` instead of being created writable and protected + /// afterwards, saving a platform permission update. + /// + /// `suggested_address` is the hint address for where to create the pages if it is not `None`. + /// Otherwise, let the kernel choose an available memory region. + /// + /// `length` is the size of the pages to be created. + /// + /// Set `flags` to control options such as fixed address, stack, and populate pages. + /// + /// # Safety + /// + /// If the suggested start address is given (i.e., not zero) and `fixed_addr` is set to `true`, + /// the kernel uses it directly without checking if it is available, causing overlapping + /// mappings to be unmapped. Caller must ensure any overlapping mappings are not used by any other. + pub unsafe fn create_pages_no_init( + &self, + suggested_address: Option>, + length: NonZeroPageSize, + flags: CreatePagesFlags, + perms: MemoryRegionPermissions, + ) -> Result, MappingError> { + let mut vmem = self.vmem.write(); + unsafe { vmem.create_pages(suggested_address, length, flags, perms) } + } + /// Create stack pages. /// /// `suggested_address` is the hint address for where to create the pages if it is not `None`. @@ -275,7 +304,7 @@ where ) -> Result, MappingError> { let perms = MemoryRegionPermissions::READ | MemoryRegionPermissions::WRITE; let flags = CreatePagesFlags::IS_STACK | flags; - unsafe { self.create_pages(suggested_address, length, flags, perms, perms, |_| Ok(0)) } + unsafe { self.create_pages_no_init(suggested_address, length, flags, perms) } } /// Set the initial program break address. @@ -368,11 +397,21 @@ where &self, releasable: fn(Range, VmFlags) -> bool, ) -> Result<(), VmemUnmapError> { + // Coalesce adjacent releasable ranges so that each contiguous run is released with a + // single platform deallocation. + let mut ranges: Vec> = Vec::new(); for (r, vma) in self.mappings() { if !releasable(r.clone(), vma) { continue; } - let mut vmem = self.vmem.write(); + match ranges.last_mut() { + Some(last) if last.end == r.start => last.end = r.end, + _ => ranges.push(r), + } + } + + let mut vmem = self.vmem.write(); + for r in ranges { let Some(range) = PageRange::new(r.start, r.end) else { unreachable!() }; @@ -380,7 +419,6 @@ where } // reset brk - let mut vmem = self.vmem.write(); vmem.brk = 0; Ok(()) diff --git a/litebox/src/mm/tests.rs b/litebox/src/mm/tests.rs index 31ad722fa..9aa9b3a57 100644 --- a/litebox/src/mm/tests.rs +++ b/litebox/src/mm/tests.rs @@ -295,3 +295,50 @@ fn test_vmm_mapping() { ] ); } + +#[test] +fn test_vmm_protect_coalescing() { + let start: usize = 0x1_0000; + let mut vmm = Vmem::new(&DummyVmemBackend); + let range = PageRange::new(start, start + 4 * PAGE_SIZE).unwrap(); + unsafe { + vmm.insert_mapping( + range, + VmArea::new( + VmFlags::VM_READ | VmFlags::VM_MAYREAD | VmFlags::VM_MAYWRITE, + false, + ), + false, + crate::platform::page_mgmt::FixedAddressBehavior::Replace, + ) + } + .unwrap(); + + // Split the mapping: the middle page becomes READ | WRITE. + unsafe { + vmm.protect_mapping( + PageRange::new(start + PAGE_SIZE, start + 2 * PAGE_SIZE).unwrap(), + MemoryRegionPermissions::READ | MemoryRegionPermissions::WRITE, + ) + } + .unwrap(); + assert_eq!( + collect_mappings(&vmm), + vec![ + start..start + PAGE_SIZE, + start + PAGE_SIZE..start + 2 * PAGE_SIZE, + start + 2 * PAGE_SIZE..start + 4 * PAGE_SIZE, + ] + ); + + // Protect the whole range: the areas needing change are coalesced around the already-RW + // middle page, and the resulting equal-flag areas merge back into a single mapping. + unsafe { + vmm.protect_mapping( + range, + MemoryRegionPermissions::READ | MemoryRegionPermissions::WRITE, + ) + } + .unwrap(); + assert_eq!(collect_mappings(&vmm), vec![start..start + 4 * PAGE_SIZE]); +} diff --git a/litebox/src/net/mod.rs b/litebox/src/net/mod.rs index de31c04ed..24ebbd843 100644 --- a/litebox/src/net/mod.rs +++ b/litebox/src/net/mod.rs @@ -80,6 +80,10 @@ where queued_for_closure: Vec>, /// Sockets that are closing in the background closing_in_background: Vec, + /// Whether any socket may be in the deferred-close state (`consider_closed`), allowing + /// [`Self::close_pending_sockets`] to skip its scan of the descriptor table. May be stale + /// towards `true` (a scan then finds nothing and resets it), but never towards `false`. + has_deferred_closes: bool, } impl Network @@ -123,6 +127,7 @@ where platform_interaction: PlatformInteraction::Automatic, queued_for_closure: vec![], closing_in_background: vec![], + has_deferred_closes: false, } } } @@ -518,6 +523,11 @@ where /// Close all finished sockets that are marked as closed but waiting for pending data to be sent fn close_pending_sockets(&mut self) { + if !self.has_deferred_closes { + // fast path: no socket is in the deferred-close state + return; + } + let mut any_remaining = false; let table = self.litebox.descriptor_table(); for (_, mut handle) in table.iter_mut::>() { let socket_handle = &mut handle.entry; @@ -526,6 +536,7 @@ where if let Some(proxy) = &socket_handle.proxy && proxy.has_pending_tx() { + any_remaining = true; continue; } @@ -548,9 +559,13 @@ where ); if closed { socket_handle.consider_closed = false; + } else { + any_remaining = true; } } } + drop(table); + self.has_deferred_closes = any_remaining; } /// Drain all socket channel buffers @@ -727,9 +742,13 @@ where /// [`set_socket_proxy`](Self::set_socket_proxy). pub fn socket(&mut self, protocol: Protocol) -> Result, SocketError> { let handle = match protocol { + // TCP sockets start with zero-capacity buffers; the real (large) rx/tx rings are + // allocated lazily on `connect` (see `ensure_tcp_socket_buffers`). This avoids the + // allocation entirely for sockets that never carry data themselves--most notably + // listening sockets, whose data flows through the backlog sockets instead. Protocol::Tcp => self.socket_set.add(tcp::Socket::new( - smoltcp::storage::RingBuffer::new(vec![0u8; SOCKET_BUFFER_SIZE]), - smoltcp::storage::RingBuffer::new(vec![0u8; SOCKET_BUFFER_SIZE]), + smoltcp::storage::RingBuffer::new(vec![]), + smoltcp::storage::RingBuffer::new(vec![]), )), Protocol::Udp => self.socket_set.add(udp::Socket::new( smoltcp::storage::PacketBuffer::new( @@ -801,6 +820,33 @@ where self.litebox.descriptor_table_mut().insert(socket_handle) } + /// Lazily allocate the rx/tx buffers of a TCP socket that was created with zero-capacity + /// buffers (see [`Self::socket`]), before it is first connected. + /// + /// Since smoltcp provides no way to replace a socket's buffers in place, this swaps in a + /// freshly-buffered socket, carrying over any socket options that may have been set already. + fn ensure_tcp_socket_buffers( + socket_set: &mut smoltcp::iface::SocketSet<'static>, + socket_handle: &mut SocketHandle, + ) { + let socket = socket_set.get_mut::(socket_handle.handle); + if socket.recv_capacity() != 0 || socket.send_capacity() != 0 { + return; + } + let mut new_socket = tcp::Socket::new( + smoltcp::storage::RingBuffer::new(vec![0u8; SOCKET_BUFFER_SIZE]), + smoltcp::storage::RingBuffer::new(vec![0u8; SOCKET_BUFFER_SIZE]), + ); + new_socket.set_nagle_enabled(socket.nagle_enabled()); + new_socket.set_keep_alive(socket.keep_alive()); + new_socket.set_timeout(socket.timeout()); + new_socket.set_ack_delay(socket.ack_delay()); + new_socket.set_hop_limit(socket.hop_limit()); + new_socket.set_congestion_control(socket.congestion_control()); + socket_set.remove(socket_handle.handle); + socket_handle.handle = socket_set.add(new_socket); + } + /// Set the network proxy for the socket at `fd` /// /// Associating a proxy enables event notification and sending/receiving data without accessing @@ -880,6 +926,7 @@ where else { unreachable!() }; + self.has_deferred_closes = true; return Err(CloseError::DataPending); } } @@ -989,6 +1036,11 @@ where } }; + if !check_progress { + // The socket is about to be connected; make sure its rx/tx buffers have + // been allocated. + Self::ensure_tcp_socket_buffers(&mut self.socket_set, socket_handle); + } let socket: &mut tcp::Socket = self.socket_set.get_mut(socket_handle.handle); if check_progress { check_state(socket.state()) diff --git a/litebox/src/net/phy.rs b/litebox/src/net/phy.rs index ff4d85c6e..2cb90c56b 100644 --- a/litebox/src/net/phy.rs +++ b/litebox/src/net/phy.rs @@ -95,9 +95,13 @@ impl smoltcp::phy::TxToken for TxToken< { let packet = &mut self.buffer[..len]; let res = f(packet); - self.platform - .send_ip_packet(packet) - .expect("Sending IP packet failed"); + // `WouldBlock` (e.g. a full transmit queue) is ordinary backpressure, + // not a fatal condition: drop the packet, exactly as a real network + // device driver would when its TX ring is full. Higher layers (TCP + // retransmission, etc.) are responsible for recovery. + match self.platform.send_ip_packet(packet) { + Ok(()) | Err(platform::SendError::WouldBlock) => {} + } res } } diff --git a/litebox/src/path.rs b/litebox/src/path.rs index 32f098c6d..5344ef7a6 100644 --- a/litebox/src/path.rs +++ b/litebox/src/path.rs @@ -65,28 +65,29 @@ pub trait Arg: private::Sealed { /// This is similar to [`Self::components`] except with normalization. Look at the tests for /// details on normalization. fn normalized_components(&self) -> Result> { - let mut parent_count = 0; - let mut rev_norm_components = self - .as_rust_str()? - .rsplit('/') - .filter(|&component| match component { - "" | "." => false, - ".." => { - parent_count += 1; - false - } - _ if parent_count > 0 => { - parent_count -= 1; - false - } - _ => true, - }) - .collect::>(); - rev_norm_components.extend(core::iter::repeat_n("..", parent_count)); - if self.as_rust_str()?.starts_with('/') { - rev_norm_components.push(""); + let path = self.as_rust_str()?; + // A single forward scan, using the components collected so far as a stack: + // ordinary components are pushed, and `..` pops the most recent ordinary + // component (or is kept, when there is nothing left to pop--i.e., leading + // `..`s and the root marker of an absolute path). + let mut components = alloc::vec::Vec::new(); + if path.starts_with('/') { + // Marker for the root of an absolute path; never popped by `..`. + components.push(""); } - Ok(rev_norm_components.into_iter().rev()) + for component in path.split('/') { + match component { + "" | "." => {} + ".." => match components.last() { + None | Some(&("" | "..")) => components.push(".."), + Some(_) => { + components.pop(); + } + }, + _ => components.push(component), + } + } + Ok(components.into_iter()) } /// Convenience wrapper around [`Self::normalized_components`] diff --git a/litebox/src/platform/common_providers/userspace_pointers.rs b/litebox/src/platform/common_providers/userspace_pointers.rs index 25b61c129..42bc85df4 100644 --- a/litebox/src/platform/common_providers/userspace_pointers.rs +++ b/litebox/src/platform/common_providers/userspace_pointers.rs @@ -213,6 +213,31 @@ fn to_owned_slice( Some(unsafe { data.assume_init() }) } +fn copy_to_slice( + ptr: *const T, + start_offset: usize, + buf: &mut [T], +) -> Option<()> { + if buf.is_empty() { + return Some(()); + } + let src = ptr.wrapping_add(start_offset); + let src = + V::validate_slice(core::ptr::slice_from_raw_parts(src, buf.len()).cast_mut())?.cast_const(); + // SAFETY: The FromBytes bound on T guarantees that any byte pattern is valid for T, + // and `buf` is a valid kernel buffer. The memcpy_fallible operation returns an error + // on invalid memory access. + V::with_user_memory_access(|| unsafe { + memcpy_fallible( + buf.as_mut_ptr().cast(), + src.cast(), + core::mem::size_of_val(buf), + ) + }) + .ok()?; + Some(()) +} + impl RawConstPointer for UserConstPtr { fn read_at_offset(self, count: isize) -> Option { read_at_offset::(self.as_ptr(), count) @@ -222,6 +247,10 @@ impl RawConstPointer for UserConstPtr to_owned_slice::(self.as_ptr(), len) } + fn copy_to_slice(self, start_offset: usize, buf: &mut [T]) -> Option<()> { + copy_to_slice::(self.as_ptr(), start_offset, buf) + } + fn as_usize(&self) -> usize { self.inner } @@ -287,6 +316,10 @@ impl RawConstPointer for UserMutPtr { to_owned_slice::(self.as_ptr().cast_const(), len) } + fn copy_to_slice(self, start_offset: usize, buf: &mut [T]) -> Option<()> { + copy_to_slice::(self.as_ptr().cast_const(), start_offset, buf) + } + fn as_usize(&self) -> usize { self.inner } diff --git a/litebox/src/platform/mod.rs b/litebox/src/platform/mod.rs index 3b354a4db..b895a564a 100644 --- a/litebox/src/platform/mod.rs +++ b/litebox/src/platform/mod.rs @@ -77,6 +77,15 @@ pub trait ThreadProvider: RawPointerProvider { /// [`EnterShim::interrupt`]: crate::shim::EnterShim::interrupt fn interrupt_thread(&self, thread: &Self::ThreadHandle); + /// Returns the number of logical CPUs available for running threads, if + /// the platform can determine it. + /// + /// The default implementation returns `None`, letting callers choose a + /// fallback. Platforms that can query the host should override this. + fn num_cpus(&self) -> Option { + None + } + /// Runs `f` on the current thread after performing any platform-specific /// thread registration needed for [`current_thread`](Self::current_thread) /// and related functionality to work. @@ -173,6 +182,15 @@ pub trait RawMutex: Send + Sync + 'static { /// value of zero. const INIT: Self; + /// The number of times lock acquisition spins on the lock word before + /// blocking when the lock is contended (see [`crate::sync::Mutex`]). + /// + /// Platforms where spinning is likely counterproductive—e.g., + /// paravirtualized guests where the lock holder's vCPU may be preempted, + /// making even short critical sections appear arbitrarily long—should + /// override this with a small value (or zero, to block immediately). + const CONTENDED_LOCK_SPIN_COUNT: u32 = 100; + /// Returns a reference to the underlying atomic value fn underlying_atomic(&self) -> &core::sync::atomic::AtomicU32; @@ -252,7 +270,15 @@ pub trait IPInterfaceProvider { /// A non-exhaustive list of errors that can be thrown by [`IPInterfaceProvider::send_ip_packet`]. #[derive(Error, Debug)] #[non_exhaustive] -pub enum SendError {} +pub enum SendError { + /// The send operation would block (e.g. the underlying device's transmit + /// queue is full). Ordinary network backpressure, not a fatal condition: + /// the caller should drop the packet and rely on higher-layer + /// retransmission, exactly as a real network device driver would when its + /// TX ring is full. + #[error("Send operation would block")] + WouldBlock, +} /// A non-exhaustive list of errors that can be thrown by [`IPInterfaceProvider::receive_ip_packet`]. #[derive(Error, Debug)] @@ -351,6 +377,24 @@ where /// be invalid. fn to_owned_slice(self, len: usize) -> Option>; + /// Copy `buf.len()` values, starting `start_offset` values past the pointer, into the provided + /// kernel buffer. + /// + /// Unlike [`Self::to_owned_slice`], this does not allocate, making it suitable for copying + /// through a reusable kernel buffer. + /// + /// Returns `None` if the provided pointer is invalid, or such a slice is known (in advance) to + /// be invalid; in that case there are no guarantees about how much of `buf` has been + /// overwritten. + #[must_use] + fn copy_to_slice(self, start_offset: usize, buf: &mut [T]) -> Option<()> { + let start: isize = start_offset.try_into().ok()?; + for (offset, slot) in (start..).zip(buf.iter_mut()) { + *slot = self.read_at_offset(offset)?; + } + Some(()) + } + /// Read the pointer as an owned C string. /// /// Returns `None` if the provided pointer is invalid, or such a string is known (in advance) to diff --git a/litebox/src/platform/page_mgmt.rs b/litebox/src/platform/page_mgmt.rs index a321ce7aa..0c0672a90 100644 --- a/litebox/src/platform/page_mgmt.rs +++ b/litebox/src/platform/page_mgmt.rs @@ -174,6 +174,27 @@ pub trait PageManagementProvider: RawPointerProvider { /// Note that the returned ranges should be `ALIGN`-aligned. fn reserved_pages(&self) -> impl Iterator>; + /// Discard the contents of already-mapped pages in `range` without unmapping them. + /// + /// On success the pages remain mapped with unchanged permissions, but their contents are + /// reset as if freshly allocated (i.e., zero-filled on next access). This enables an + /// in-place fast path for guest `madvise(MADV_DONTNEED)`-style operations, avoiding the + /// page-table churn of unmapping and re-creating the mapping. + /// + /// The default implementation returns [`DiscardError::Unsupported`], in which case the + /// caller must fall back to re-creating the mapping. Platforms should only override this + /// when discarding pages allocated via [`Self::allocate_pages`] yields zero-fill semantics. + /// + /// # Safety + /// + /// The caller must ensure that the previous contents of `range` are no longer accessed or + /// relied upon, and that `range` only covers pages allocated via [`Self::allocate_pages`] + /// (i.e., not file-backed mappings, whose discard semantics may differ). + #[expect(unused_variables, reason = "default body, non-underscored param names")] + unsafe fn discard_pages(&self, range: Range) -> Result<(), DiscardError> { + Err(DiscardError::Unsupported) + } + /// Attempt to allocate pages with copy-on-write semantics backed by static data. /// /// This method allows platforms that support it to create CoW mappings instead of performing @@ -238,6 +259,16 @@ pub enum DeallocationError { AlreadyUnallocated, } +/// Possible errors for [`PageManagementProvider::discard_pages`] +#[derive(Error, Debug)] +#[non_exhaustive] +pub enum DiscardError { + #[error("in-place page discard is not supported by this platform")] + Unsupported, + #[error("provided range contains unallocated pages")] + Unallocated, +} + /// Possible errors for [`PageManagementProvider::remap_pages`] #[derive(Error, Debug)] #[non_exhaustive] diff --git a/litebox/src/platform/trivial_providers.rs b/litebox/src/platform/trivial_providers.rs index bda4d01eb..9057d5197 100644 --- a/litebox/src/platform/trivial_providers.rs +++ b/litebox/src/platform/trivial_providers.rs @@ -85,6 +85,23 @@ impl RawConstPointer for TransparentConstPtr { } } + fn copy_to_slice(self, start_offset: usize, buf: &mut [T]) -> Option<()> { + let ptr = self.as_ptr(); + if ptr.is_null() || !ptr.is_aligned() { + return None; + } + // SAFETY: We checked the pointer is non-null and aligned. The FromBytes bound + // on T guarantees that any byte pattern is valid for T. + unsafe { + core::ptr::copy_nonoverlapping( + ptr.wrapping_add(start_offset), + buf.as_mut_ptr(), + buf.len(), + ); + } + Some(()) + } + fn as_usize(&self) -> usize { self.inner } @@ -161,6 +178,23 @@ impl RawConstPointer for TransparentMutPtr { } } + fn copy_to_slice(self, start_offset: usize, buf: &mut [T]) -> Option<()> { + let ptr = self.as_ptr(); + if ptr.is_null() || !ptr.is_aligned() { + return None; + } + // SAFETY: We checked the pointer is non-null and aligned. The FromBytes bound + // on T guarantees that any byte pattern is valid for T. + unsafe { + core::ptr::copy_nonoverlapping( + ptr.wrapping_add(start_offset).cast_const(), + buf.as_mut_ptr(), + buf.len(), + ); + } + Some(()) + } + fn as_usize(&self) -> usize { self.inner } diff --git a/litebox/src/sync/futex.rs b/litebox/src/sync/futex.rs index 5e262b1a5..b41d8368f 100644 --- a/litebox/src/sync/futex.rs +++ b/litebox/src/sync/futex.rs @@ -13,7 +13,7 @@ use core::hash::BuildHasher as _; use core::num::NonZeroU32; use core::pin::pin; -use core::sync::atomic::{AtomicBool, Ordering}; +use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use super::RawSyncPrimitivesProvider; use crate::event::wait::{WaitContext, WaitError, Waker}; @@ -31,6 +31,21 @@ pub struct FutexManager { /// Chaining hash table to map from futex address to waiter lists. table: alloc::boxed::Box<[LoanList>; HASH_TABLE_ENTRIES]>, hash_builder: hashbrown::DefaultHashBuilder, + /// Number of live waiter entries whose current futex word hashes to a + /// bucket other than the one they physically reside in (a result of + /// [`FutexManager::requeue`]). When non-zero, wake-ups must scan every + /// bucket to find such entries. + /// + /// This is a single global counter, so while *any* cross-bucket-requeued + /// waiter is parked, *every* wake/requeue scans all [`HASH_TABLE_ENTRIES`] + /// buckets rather than one. That degradation is bounded to the (uncommon) + /// window in which such a waiter stays blocked — cross-bucket requeue only + /// arises from `FUTEX_REQUEUE`/`CMP_REQUEUE`, i.e. glibc/musl condvar + /// broadcast whose target mutex word hashes elsewhere. Correctness is + /// unaffected. FUTURE: track displacement per target bucket so the fast + /// path survives, but only with an accounting scheme proven not to skip a + /// bucket holding a displaced waiter (a miscount would lose a wake-up). + displaced: AtomicUsize, } /// The number of buckets in the hash table. @@ -40,10 +55,15 @@ pub struct FutexManager { const HASH_TABLE_ENTRIES: usize = 256; struct FutexEntry { - addr: usize, + /// The futex word address this entry is waiting on. Only mutated (under + /// the resident bucket's lock) when the waiter is requeued to another + /// word; the entry stays in the bucket it was originally inserted into. + addr: AtomicUsize, waker: Waker, bitset: u32, done: AtomicBool, + /// Whether this entry is currently counted in [`FutexManager::displaced`]. + displaced: AtomicBool, } const ALL_BITS: NonZeroU32 = NonZeroU32::new(u32::MAX).unwrap(); @@ -59,13 +79,19 @@ impl Self { table: alloc::boxed::Box::new(core::array::from_fn(|_| LoanList::new())), hash_builder: hashbrown::DefaultHashBuilder::default(), + displaced: AtomicUsize::new(0), } } + /// Returns the hash table bucket index for the given futex address. + fn bucket_index(&self, addr: usize) -> usize { + let hash: usize = self.hash_builder.hash_one(addr).trunc(); + hash % HASH_TABLE_ENTRIES + } + /// Returns the hash table bucket for the given futex address. fn bucket(&self, addr: usize) -> &LoanList> { - let hash: usize = self.hash_builder.hash_one(addr).trunc(); - &self.table[hash % HASH_TABLE_ENTRIES] + &self.table[self.bucket_index(addr)] } /// Performs a futex wait. @@ -93,29 +119,64 @@ impl return Err(FutexError::NotAligned); } + // Optimistically check the futex word before touching the bucket: if the + // value already mismatches, report that without paying for the bucket's + // lock. This check alone is insufficient (the value could change right + // after it), so the authoritative check below is done after inserting + // into the bucket, ensuring no wakeup is missed. + let value = futex_addr.read_at_offset(0).ok_or(FutexError::Fault)?; + if value != expected_value { + return Err(FutexError::ImmediatelyWokenBecauseValueMismatch); + } + let bucket = self.bucket(addr); let mut entry = pin!(LoanListEntry::new(FutexEntry { - addr, + addr: AtomicUsize::new(addr), waker: cx.waker().clone(), bitset, done: AtomicBool::new(false), + displaced: AtomicBool::new(false), },)); // Insert into the bucket's list. It will be removed when woken or the // entry goes out of scope. entry.as_mut().insert(bucket); - // Check the value once. Do this only after inserting into the list so + // Check the value again. Do this after inserting into the list so // that we don't miss a wakeup. - let value = futex_addr.read_at_offset(0).ok_or(FutexError::Fault)?; - if value != expected_value { - return Err(FutexError::ImmediatelyWokenBecauseValueMismatch); + let result = match futex_addr.read_at_offset(0) { + None => Err(FutexError::Fault), + Some(value) if value != expected_value => { + Err(FutexError::ImmediatelyWokenBecauseValueMismatch) + } + // Only return when woken--don't reevaluate the futex word. This + // ensures that the rate control mechanisms provided by the futex + // interface are effective. + Some(_) => cx + .wait_until(|| entry.get().done.load(Ordering::Acquire)) + .map_err(FutexError::WaitError), + }; + + // Remove the entry before reading its `displaced` flag: once removed, + // no concurrent requeue can touch the entry, making the flag stable. + entry.as_mut().remove(); + // Resolve the wake-vs-timeout/interrupt race in favor of the wake. A + // concurrent `wake`/`requeue` may have selected this waiter — counting + // it as woken and consuming that wakeup — at the same moment + // `wait_until` reported a timeout or interruption. `remove()` above + // synchronizes with the waker's release of the entry (see + // `LoanedEntry::drop`), so `done` is now stable; if it is set, report + // success rather than dropping a wakeup the waker already accounted for + // (matching Linux `futex_wait`, where a failed dequeue returns 0). + let result = if entry.get().done.load(Ordering::Acquire) { + Ok(()) + } else { + result + }; + if entry.get().displaced.load(Ordering::Relaxed) { + self.displaced.fetch_sub(1, Ordering::SeqCst); } - // Only return when woken--don't reevaluate the futex word. This - // ensures that the rate control mechanisms provided by the futex - // interface are effective. - cx.wait_until(|| entry.get().done.load(Ordering::Acquire)) - .map_err(FutexError::WaitError) + result } /// Wakes waiters on the given futex word. @@ -144,27 +205,162 @@ impl } let bitset = bitset.unwrap_or(ALL_BITS).get(); let mut woken = 0; - let bucket = self.bucket(addr); - // Extract matching entries from the bucket until we've woken enough. - let entries = bucket.extract_if(|entry| { - if entry.addr != addr || entry.bitset & bitset == 0 { - return core::ops::ControlFlow::Continue(false); + // Entries requeued across buckets stay in their original bucket, so + // when any such entry exists the wake target may reside in any bucket + // and they must all be scanned (rare). + let buckets = if self.displaced.load(Ordering::SeqCst) == 0 { + core::slice::from_ref(self.bucket(addr)) + } else { + &self.table[..] + }; + for bucket in buckets { + // Extract matching entries from the bucket until we've woken enough. + let entries = bucket.extract_if(|entry| { + if entry.addr.load(Ordering::Relaxed) != addr || entry.bitset & bitset == 0 { + return core::ops::ControlFlow::Continue(false); + } + woken += 1; + if woken >= num_to_wake_up.get() { + core::ops::ControlFlow::Break(true) + } else { + core::ops::ControlFlow::Continue(true) + } + }); + // Wake the waiters outside the `extract_if` closure to minimize the list's lock hold + // time. + for entry in entries { + // Clone the waker and publish `done` so the entry loan can be + // returned *before* waking: `wake` may issue an expensive host wake + // call, and holding the loan across it would block a concurrent + // owner-side removal of the entry for that entire duration. + let waker = entry.waker.clone(); + entry.done.store(true, Ordering::Release); + drop(entry); + waker.wake(); } - woken += 1; if woken >= num_to_wake_up.get() { - core::ops::ControlFlow::Break(true) - } else { - core::ops::ControlFlow::Continue(true) + break; } - }); - // Wake the waiters outside the `extract_if` closure to minimize the list's lock hold - // time. - for entry in entries { - entry.done.store(true, Ordering::Relaxed); - entry.waker.wake(); } Ok(woken) } + + /// Wakes and requeues waiters on the given futex word, implementing the + /// semantics of `FUTEX_REQUEUE` and `FUTEX_CMP_REQUEUE`. + /// + /// If `expected_value` is `Some`, the current value of the futex word at + /// `futex_addr` is first compared against it, failing with + /// [`FutexError::ImmediatelyWokenBecauseValueMismatch`] on a mismatch + /// (Linux reports this as `EAGAIN`). + /// + /// Up to `num_to_wake_up` waiters on `futex_addr` are woken. Up to + /// `num_to_requeue` of the remaining waiters are moved to wait on + /// `target_addr` instead, without waking them; they become eligible for + /// subsequent [`FutexManager::wake`] calls on `target_addr`. + /// + /// Returns `(woken, requeued)`. + pub fn requeue( + &self, + futex_addr: Platform::RawMutPointer, + target_addr: Platform::RawMutPointer, + num_to_wake_up: u32, + num_to_requeue: u32, + expected_value: Option, + ) -> Result<(u32, u32), FutexError> { + let addr = futex_addr.as_usize(); + let target = target_addr.as_usize(); + if !addr.is_multiple_of(align_of::()) || !target.is_multiple_of(align_of::()) { + return Err(FutexError::NotAligned); + } + let target_bucket_index = self.bucket_index(target); + let source_bucket_index = self.bucket_index(addr); + // As in `wake`, a displaced (cross-bucket requeued) waiter forces a + // full scan; otherwise only the source word's bucket holds candidates. + let scan_all = self.displaced.load(Ordering::SeqCst) != 0; + + // Validate the futex word against `expected_value` while holding the + // source bucket's lock, before any wake/requeue side effect, matching + // Linux's locked value check for `FUTEX_CMP_REQUEUE`. An unlocked + // compare would let a waiter that enqueues on the source word (with the + // new value) between the check and the scan be requeued onto the target + // word based on a stale comparison; because requeue does not mark the + // target word contended, a later uncontended unlock would never wake + // that waiter and it could hang. The source bucket is always scanned + // first so this guard gates the entire operation. + let source_guard = move || match expected_value { + None => Ok(()), + Some(expected) => match futex_addr.read_at_offset(0) { + None => Err(FutexError::Fault), + Some(value) if value != expected => { + Err(FutexError::ImmediatelyWokenBecauseValueMismatch) + } + Some(_) => Ok(()), + }, + }; + + let mut woken = 0; + let mut requeued = 0; + let mut source_guard = Some(source_guard); + let bucket_order = core::iter::once(source_bucket_index) + .chain((0..HASH_TABLE_ENTRIES).filter(move |&i| scan_all && i != source_bucket_index)); + for resident in bucket_order { + let bucket = &self.table[resident]; + let scan = |entry: &FutexEntry| { + use core::ops::ControlFlow::{Break, Continue}; + if entry.addr.load(Ordering::Relaxed) != addr { + return Continue(false); + } + if woken < num_to_wake_up { + woken += 1; + return if woken >= num_to_wake_up && num_to_requeue == 0 { + Break(true) + } else { + Continue(true) + }; + } + if requeued < num_to_requeue { + requeued += 1; + // Retarget the entry in place; it stays in its resident + // bucket, so track (while holding this bucket's lock) + // whether it is now displaced. + entry.addr.store(target, Ordering::Relaxed); + if target_bucket_index != resident { + if !entry.displaced.swap(true, Ordering::Relaxed) { + self.displaced.fetch_add(1, Ordering::SeqCst); + } + } else if entry.displaced.swap(false, Ordering::Relaxed) { + self.displaced.fetch_sub(1, Ordering::SeqCst); + } + return if requeued >= num_to_requeue { + Break(false) + } else { + Continue(false) + }; + } + Break(false) + }; + // The source bucket (first) carries the locked value-check guard; + // on a mismatch it returns before any entry is touched. + let entries = match source_guard.take() { + Some(guard) => bucket.extract_if_guarded(guard, scan)?, + None => bucket.extract_if(scan), + }; + // Wake the waiters outside the `extract_if` closure to minimize + // the list's lock hold time. + for entry in entries { + // As in `wake`: publish `done` and return the entry loan + // before issuing the (potentially expensive) host wake. + let waker = entry.waker.clone(); + entry.done.store(true, Ordering::Release); + drop(entry); + waker.wake(); + } + if woken >= num_to_wake_up && requeued >= num_to_requeue { + break; + } + } + Ok((woken, requeued)) + } } /// Potential errors that can be returned by [`FutexManager`]'s operations. @@ -358,4 +554,77 @@ mod tests { assert!((1..=3).contains(&woken)); } + + #[test] + fn test_futex_requeue() { + let platform = MockPlatform::new(); + let _litebox = LiteBox::new(platform); + let futex_manager = Arc::new(FutexManager::new()); + + let word_a = Arc::new(AtomicU32::new(0)); + let word_b = Arc::new(AtomicU32::new(0)); + let barrier = Arc::new(Barrier::new(4)); // 3 waiters + 1 requeuer + + let addr_of = |word: &Arc| { + ::RawMutPointer::from_usize( + word.as_ptr() as usize, + ) + }; + + let mut waiters = std::vec::Vec::new(); + for _ in 0..3 { + let futex_manager_clone = Arc::clone(&futex_manager); + let word_a_clone = Arc::clone(&word_a); + let barrier_clone = Arc::clone(&barrier); + waiters.push(thread::spawn(move || { + let futex_addr = + ::RawMutPointer::from_usize( + word_a_clone.as_ptr() as usize, + ); + barrier_clone.wait(); // Sync with the requeuer + futex_manager_clone.wait(&WaitState::new(platform).context(), futex_addr, 0, None) + })); + } + + barrier.wait(); // Wait for the waiters to be ready + thread::sleep(Duration::from_millis(100)); // Give the waiters time to block + + // A mismatched expected value must fail with EAGAIN semantics without + // waking or requeueing anyone. + assert!(matches!( + futex_manager.requeue(addr_of(&word_a), addr_of(&word_b), 1, u32::MAX, Some(1)), + Err(FutexError::ImmediatelyWokenBecauseValueMismatch) + )); + + // Wake one waiter and requeue the rest onto word B. + let (woken, requeued) = futex_manager + .requeue(addr_of(&word_a), addr_of(&word_b), 1, u32::MAX, Some(0)) + .unwrap(); + assert_eq!(woken, 1); + assert_eq!(requeued, 2); + + // No one is left waiting on word A... + let woken_a = futex_manager + .wake(addr_of(&word_a), NonZeroU32::new(u32::MAX).unwrap(), None) + .unwrap(); + assert_eq!(woken_a, 0); + + // ...and the requeued waiters must be reachable via word B, even + // though they physically reside in word A's bucket. + let woken_b = futex_manager + .wake(addr_of(&word_b), NonZeroU32::new(u32::MAX).unwrap(), None) + .unwrap(); + assert_eq!(woken_b, 2); + + for waiter in waiters { + assert!(waiter.join().unwrap().is_ok()); + } + + // Every displaced entry must have been un-counted on wake-up. + assert_eq!( + futex_manager.displaced.load(Ordering::SeqCst), + 0, + "displaced counter must return to zero" + ); + } } diff --git a/litebox/src/sync/mutex.rs b/litebox/src/sync/mutex.rs index 1c0def2b3..f9351d532 100644 --- a/litebox/src/sync/mutex.rs +++ b/litebox/src/sync/mutex.rs @@ -92,7 +92,7 @@ impl SpinEnabledRawMutex { /// Returns the state of the raw lock as soon as it is in unlocked (0) or contended (2), or when /// it has spun for long enough. fn spin(&self) -> u32 { - let mut spin = 100; + let mut spin = ::CONTENDED_LOCK_SPIN_COUNT; loop { // We only use `load` (and not `swap` or `compare_exchange`) // while spinning, to be easier on the caches. diff --git a/litebox/src/utilities/loan_list.rs b/litebox/src/utilities/loan_list.rs index 5ca868eea..bdea8b269 100644 --- a/litebox/src/utilities/loan_list.rs +++ b/litebox/src/utilities/loan_list.rs @@ -110,7 +110,6 @@ impl<'a, Platform: RawSyncPrimitivesProvider, T> LoanListEntry<'a, Platform, T> /// completes and the entry is fully returned. /// /// If the entry is not currently inserted, this method does nothing. - #[cfg_attr(not(test), expect(dead_code))] pub fn remove(self: Pin<&mut Self>) { if let Some(list) = self.list { list.remove_node(&self.node); @@ -201,7 +200,12 @@ impl LoanList { let _ = node.data.state.block(s.0); } EntryState::REMOVED_WAKING => { - // Spin until the remover finishes waking us. + // Spin until the remover finishes waking us. This wait is + // bounded to the remover executing a single store after its + // wake call returns (see `LoanedEntry::drop`), so spinning + // is acceptable. Note that `block`ing here instead would be + // unsound: the remover issues no further wake after its + // final store to `REMOVED`. core::hint::spin_loop(); } state => panic!("invalid state waiting for entry removal: {state:?}"), @@ -250,9 +254,29 @@ impl LoanList { /// ``` pub fn extract_if( &self, - mut f: impl FnMut(&T) -> ControlFlow, + f: impl FnMut(&T) -> ControlFlow, ) -> ExtractIf { + match self.extract_if_guarded(|| Ok::<(), core::convert::Infallible>(()), f) { + Ok(entries) => entries, + } + } + + /// Like [`extract_if`](Self::extract_if), but first runs `guard` while + /// holding the list lock and **before** any entry is examined or removed. + /// If `guard` returns `Err`, the list is left untouched and the error is + /// returned; no entries are extracted and no predicate side effects occur. + /// + /// This lets a caller atomically validate a precondition against state that + /// other threads mutate under the same lock (e.g. re-reading a futex word + /// for `FUTEX_CMP_REQUEUE`) and the subsequent extraction, matching the + /// kernel's "check the value under the bucket lock, then requeue" ordering. + pub fn extract_if_guarded( + &self, + guard: impl FnOnce() -> Result<(), E>, + mut f: impl FnMut(&T) -> ControlFlow, + ) -> Result, E> { let mut this = self.0.lock(); + guard()?; let mut removed = LinkedList::new(); let mut current = this.head; while !current.is_null() { @@ -283,9 +307,9 @@ impl LoanList { break; } } - ExtractIf { + Ok(ExtractIf { head: unsafe { removed.into_head() }, - } + }) } } diff --git a/litebox_common_linux/src/errno/mod.rs b/litebox_common_linux/src/errno/mod.rs index c5397ab4c..2e4f0ecda 100644 --- a/litebox_common_linux/src/errno/mod.rs +++ b/litebox_common_linux/src/errno/mod.rs @@ -284,7 +284,10 @@ impl From for Errno { litebox::mm::linux::MappingError::UnAligned => Errno::EINVAL, litebox::mm::linux::MappingError::OutOfMemory => Errno::ENOMEM, litebox::mm::linux::MappingError::BadFD(_) => Errno::EBADF, - litebox::mm::linux::MappingError::NotAFile => Errno::EISDIR, + // `mmap` of a directory or a non-regular fd (pipe/socket/eventfd/ + // epoll) reports `ENODEV` on Linux ("mmap not supported by this + // file"), unlike `read`/`write` which report `EISDIR`. + litebox::mm::linux::MappingError::NotAFile => Errno::ENODEV, litebox::mm::linux::MappingError::NotForReading => Errno::EACCES, litebox::mm::linux::MappingError::MapError(e) => e.into(), _ => unimplemented!(), diff --git a/litebox_common_linux/src/lib.rs b/litebox_common_linux/src/lib.rs index 67c904275..f8538cc1d 100644 --- a/litebox_common_linux/src/lib.rs +++ b/litebox_common_linux/src/lib.rs @@ -1680,6 +1680,8 @@ bitflags::bitflags! { pub enum FutexOperation { Wait = 0, Wake = 1, + Requeue = 3, + CmpRequeue = 4, WaitBitset = 9, } @@ -1719,6 +1721,25 @@ pub enum FutexArgs { flags: FutexFlags, count: u32, }, + /// `FUTEX_REQUEUE`: wake up to `wake_count` waiters on `addr`, then move + /// up to `requeue_count` of the remaining waiters to wait on `addr2`. + Requeue { + addr: Platform::RawMutPointer, + flags: FutexFlags, + wake_count: u32, + addr2: Platform::RawMutPointer, + requeue_count: u32, + }, + /// `FUTEX_CMP_REQUEUE`: like [`FutexArgs::Requeue`], but fails with + /// `EAGAIN` unless `*addr == expected`. + CmpRequeue { + addr: Platform::RawMutPointer, + flags: FutexFlags, + wake_count: u32, + addr2: Platform::RawMutPointer, + requeue_count: u32, + expected: u32, + }, } #[repr(u32)] @@ -3021,6 +3042,24 @@ impl SyscallRequest { flags, count: val, }, + // For the requeue operations the timeout slot carries `val2`, the + // maximum number of waiters to requeue, and arg 4 is the second + // futex word. + FutexOperation::Requeue => FutexArgs::Requeue { + addr, + flags, + wake_count: val, + addr2: ctx.sys_req_ptr(4), + requeue_count: ctx.sys_req_arg(3), + }, + FutexOperation::CmpRequeue => FutexArgs::CmpRequeue { + addr, + flags, + wake_count: val, + addr2: ctx.sys_req_ptr(4), + requeue_count: ctx.sys_req_arg(3), + expected: ctx.sys_req_arg(5), + }, }; Ok(SyscallRequest::Futex { args }) } diff --git a/litebox_common_linux/src/mm.rs b/litebox_common_linux/src/mm.rs index 0c66c4016..8be273ea7 100644 --- a/litebox_common_linux/src/mm.rs +++ b/litebox_common_linux/src/mm.rs @@ -7,13 +7,44 @@ use litebox::{ mm::linux::{ CreatePagesFlags, MappingError, NonZeroAddress, NonZeroPageSize, PAGE_SIZE, VmemUnmapError, }, - platform::{RawConstPointer, page_mgmt::DeallocationError}, + platform::{ + RawConstPointer, + page_mgmt::{DeallocationError, MemoryRegionPermissions}, + }, }; use crate::{MRemapFlags, MapFlags, ProtFlags, errno::Errno}; const PAGE_MASK: usize = !(PAGE_SIZE - 1); +/// Convert `mmap` flags into [`CreatePagesFlags`]. +fn create_pages_flags(flags: MapFlags, ensure_space_after: bool) -> CreatePagesFlags { + let mut create_flags = CreatePagesFlags::empty(); + // MAP_FIXED_NOREPLACE implies MAP_FIXED behavior (exact address, not a hint) + create_flags.set( + CreatePagesFlags::FIXED_ADDR, + flags.intersects(MapFlags::MAP_FIXED | MapFlags::MAP_FIXED_NOREPLACE), + ); + create_flags.set( + CreatePagesFlags::NOREPLACE, + flags.contains(MapFlags::MAP_FIXED_NOREPLACE), + ); + create_flags.set( + CreatePagesFlags::POPULATE_PAGES_IMMEDIATELY, + flags.contains(MapFlags::MAP_POPULATE), + ); + create_flags.set(CreatePagesFlags::ENSURE_SPACE_AFTER, ensure_space_after); + create_flags.set( + CreatePagesFlags::MAP_FILE, + !flags.contains(MapFlags::MAP_ANONYMOUS), + ); + create_flags.set( + CreatePagesFlags::SHARED, + flags.contains(MapFlags::MAP_SHARED), + ); + create_flags +} + pub fn do_mmap< Platform: litebox::platform::RawPointerProvider + litebox::sync::RawSyncPrimitivesProvider @@ -27,32 +58,7 @@ pub fn do_mmap< ensure_space_after: bool, op: impl FnOnce(Platform::RawMutPointer) -> Result, ) -> Result, litebox::mm::linux::MappingError> { - let flags = { - let mut create_flags = CreatePagesFlags::empty(); - // MAP_FIXED_NOREPLACE implies MAP_FIXED behavior (exact address, not a hint) - create_flags.set( - CreatePagesFlags::FIXED_ADDR, - flags.intersects(MapFlags::MAP_FIXED | MapFlags::MAP_FIXED_NOREPLACE), - ); - create_flags.set( - CreatePagesFlags::NOREPLACE, - flags.contains(MapFlags::MAP_FIXED_NOREPLACE), - ); - create_flags.set( - CreatePagesFlags::POPULATE_PAGES_IMMEDIATELY, - flags.contains(MapFlags::MAP_POPULATE), - ); - create_flags.set(CreatePagesFlags::ENSURE_SPACE_AFTER, ensure_space_after); - create_flags.set( - CreatePagesFlags::MAP_FILE, - !flags.contains(MapFlags::MAP_ANONYMOUS), - ); - create_flags.set( - CreatePagesFlags::SHARED, - flags.contains(MapFlags::MAP_SHARED), - ); - create_flags - }; + let flags = create_pages_flags(flags, ensure_space_after); let suggested_addr = match suggested_addr { Some(addr) => Some(NonZeroAddress::new(addr).ok_or(MappingError::UnAligned)?), None => None, @@ -84,6 +90,48 @@ pub fn do_mmap< } } +/// Like [`do_mmap`], but for mappings that need no initialization callback (e.g., anonymous +/// mappings). +/// +/// The pages are mapped directly with their final permissions, saving the platform permission +/// update that [`do_mmap`] performs after running the initialization `op`. +pub fn do_mmap_no_init< + Platform: litebox::platform::RawPointerProvider + + litebox::sync::RawSyncPrimitivesProvider + + litebox::platform::PageManagementProvider<{ litebox::mm::linux::PAGE_SIZE }>, +>( + pm: &litebox::mm::PageManager, + suggested_addr: Option, + len: usize, + prot: ProtFlags, + flags: MapFlags, + ensure_space_after: bool, +) -> Result, litebox::mm::linux::MappingError> { + let flags = create_pages_flags(flags, ensure_space_after); + let suggested_addr = match suggested_addr { + Some(addr) => Some(NonZeroAddress::new(addr).ok_or(MappingError::UnAligned)?), + None => None, + }; + let length = NonZeroPageSize::new(len).ok_or(MappingError::UnAligned)?; + let perms = match prot { + ProtFlags::PROT_READ_EXEC => MemoryRegionPermissions::READ | MemoryRegionPermissions::EXEC, + ProtFlags::PROT_READ_WRITE => { + MemoryRegionPermissions::READ | MemoryRegionPermissions::WRITE + } + ProtFlags::PROT_READ => MemoryRegionPermissions::READ, + ProtFlags::PROT_NONE => MemoryRegionPermissions::empty(), + _ => { + #[cfg(debug_assertions)] + todo!("Unsupported prot flags {:?}", prot); + // TODO: create inaccessible pages for now. Creating mapping + // for both executable and writable might be needed for JIT. + #[cfg(not(debug_assertions))] + MemoryRegionPermissions::empty() + } + }; + unsafe { pm.create_pages_no_init(suggested_addr, length, flags, perms) } +} + /// Handle syscall `munmap` pub fn sys_munmap< Platform: litebox::platform::RawPointerProvider diff --git a/litebox_platform_linux_kernel/src/arch/x86/mm/paging.rs b/litebox_platform_linux_kernel/src/arch/x86/mm/paging.rs index ddc6d3e68..3e23e6422 100644 --- a/litebox_platform_linux_kernel/src/arch/x86/mm/paging.rs +++ b/litebox_platform_linux_kernel/src/arch/x86/mm/paging.rs @@ -65,12 +65,24 @@ impl FrameDeallocator for PageTableAllocator { } pub(crate) fn vmflags_to_pteflags(values: VmFlags) -> PageTableFlags { - let mut flags = PageTableFlags::empty(); + // Pre-set ACCESSED on every faulted-in leaf PTE. Otherwise the CPU's page-table + // walker sets the A bit itself on first access with a locked, microcoded write + // that takes the PTE cache line for exclusive ownership (an RFO). Filling in the + // software-known result up front avoids that write, cutting cache-coherency + // traffic and slightly speeding up the walk. Mirrors Linux's pte_mkyoung() on + // fault-in. ACCESSED/DIRTY are excluded from MPROTECT_PTE_MASK, so mprotect never + // observes or clears them, and no code in this platform reads A/D for data pages, + // so this is a pure hardware-walk optimization with no semantic effect. + let mut flags = PageTableFlags::ACCESSED; if values.intersects(VmFlags::VM_READ | VmFlags::VM_WRITE) { flags |= PageTableFlags::USER_ACCESSIBLE; } if values.contains(VmFlags::VM_WRITE) { - flags |= PageTableFlags::WRITABLE; + // Writable leaves are additionally pre-marked DIRTY (matching Linux's + // pte_mkdirty()): the first write would otherwise cost a second locked walker + // write to set D. DIRTY is only ever set together with WRITABLE, never on a + // read-only PTE. + flags |= PageTableFlags::WRITABLE | PageTableFlags::DIRTY; } if !values.contains(VmFlags::VM_EXEC) { flags |= PageTableFlags::NO_EXECUTE; @@ -324,9 +336,15 @@ impl PageTableImpl for X64PageTabl let mut allocator = PageTableAllocator::::new(); // TODO: if it is file-backed, we need to read the page from file let frame = PageTableAllocator::::allocate_frame(true).unwrap(); + // Intermediate (non-leaf) entries the mapper may create for this leaf. + // ACCESSED is pre-set for the same reason as on the leaf: it spares the + // walker the locked A-bit write on these entries (Linux's _PAGE_TABLE + // includes _PAGE_ACCESSED). DIRTY is meaningless on non-leaf entries and + // is deliberately omitted. let table_flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE - | PageTableFlags::USER_ACCESSIBLE; + | PageTableFlags::USER_ACCESSIBLE + | PageTableFlags::ACCESSED; match unsafe { inner.map_to_with_table_flags( page, diff --git a/litebox_platform_linux_userland/Cargo.toml b/litebox_platform_linux_userland/Cargo.toml index a4e5ef9a2..224f56e59 100644 --- a/litebox_platform_linux_userland/Cargo.toml +++ b/litebox_platform_linux_userland/Cargo.toml @@ -11,6 +11,7 @@ libc = { version = "0.2.169", default-features = false } litebox = { path = "../litebox/", version = "0.1.0" } litebox_common_linux = { path = "../litebox_common_linux", version = "0.1.0" } litebox_common_optee = { path = "../litebox_common_optee", version = "0.1.0", default-features = false, optional = true } +litebox_util_log = { path = "../litebox_util_log", version = "0.1.0" } spin = "0.9.8" syscalls = { version = "0.6", default-features = false } zerocopy = { version = "0.8", default-features = false } diff --git a/litebox_platform_linux_userland/src/lib.rs b/litebox_platform_linux_userland/src/lib.rs index 7b0c825f4..aa264e278 100644 --- a/litebox_platform_linux_userland/src/lib.rs +++ b/litebox_platform_linux_userland/src/lib.rs @@ -112,10 +112,11 @@ impl core::fmt::Debug for LinuxUserland { } /// Information about a CoW-eligible memory region backed by a file. -#[derive(Debug, Clone)] +#[derive(Debug)] struct CowRegionInfo { - /// The path to the backing file on the host filesystem. - file_path: PathBuf, + /// The backing file, opened read-only once at registration time so that + /// each CoW mapping is a single host `mmap` (no per-mapping open/close). + file: std::fs::File, /// Length of the backing file. file_length: usize, } @@ -175,6 +176,10 @@ impl LinuxUserland { /// /// Panics if the tun device could not be successfully opened. pub fn new(tun_device_name: Option<&str>) -> &'static Self { + init_fsgsbase_mode(); + // Probe the CPU count now, before the seccomp filter is installed, so + // the query never trips the sandbox (see `num_cpus`). + NUM_CPUS.get_or_init(|| std::thread::available_parallelism().ok()); register_exception_handlers(); let tun_socket_fd = tun_device_name @@ -266,11 +271,20 @@ impl LinuxUserland { /// /// # Panics /// + /// Panics if the backing file cannot be opened read-only. + /// /// Panics if an overlapping region is already registered. pub fn register_cow_region(&self, data: &'static [u8], file_path: impl Into) { let start = data.as_ptr() as usize; + let file_path = file_path.into(); + let file = std::fs::File::open(&file_path).unwrap_or_else(|err| { + panic!( + "failed to open CoW backing file {}: {err}", + file_path.display() + ) + }); let info = CowRegionInfo { - file_path: file_path.into(), + file, file_length: data.len(), }; @@ -285,9 +299,11 @@ impl LinuxUserland { /// Look up the file backing a static slice for CoW mapping. /// - /// Returns `Some((file_path, offset_in_file))` if the slice is backed by a registered - /// CoW region, `None` otherwise. - fn lookup_cow_region(&self, source_data: &'static [u8]) -> Option<(PathBuf, usize)> { + /// Returns `Some((fd, offset_in_file))` if the slice is backed by a registered + /// CoW region, `None` otherwise. The returned raw fd stays valid for the + /// process lifetime: regions are never unregistered and the platform is + /// leaked (`&'static`). + fn lookup_cow_region(&self, source_data: &'static [u8]) -> Option<(std::os::fd::RawFd, usize)> { let slice_start = source_data.as_ptr() as usize; let slice_len = source_data.len(); @@ -298,7 +314,7 @@ impl LinuxUserland { let slice_end = slice_start.checked_add(slice_len).unwrap(); if slice_start >= region_start && slice_end <= region_end { - return Some((info.file_path.clone(), slice_start - region_start)); + return Some((info.file.as_raw_fd(), slice_start - region_start)); } } None @@ -451,6 +467,17 @@ impl LinuxUserland { (libc::SYS_clone3, vec![]), // sync (libc::SYS_futex, vec![]), + // time: glibc normally services these via the vDSO, but it falls + // back to the real syscalls when the clocksource is not + // vDSO-capable (e.g. an unstable TSC in PVM guests), so they must + // be allowed. They are read-only and sandbox-safe. + (libc::SYS_clock_gettime, vec![]), + (libc::SYS_clock_getres, vec![]), + (libc::SYS_gettimeofday, vec![]), + // segment-base switching on hosts without user-mode FSGSBASE + // (see `init_fsgsbase_mode`); affects only the calling thread's + // fs/gs bases, so it is sandbox-neutral. + (libc::SYS_arch_prctl, vec![]), // misc (libc::SYS_getrandom, vec![]), // required by std spawn @@ -463,7 +490,9 @@ impl LinuxUserland { // required by libc allocator (libc::SYS_brk, vec![]), (libc::SYS_getpid, vec![]), - // TODO: could be removed if we pre-open files (see `try_allocate_cow_pages`) + // CoW backing files are pre-opened at registration time (see + // `register_cow_region`); this read-only rule remains for host + // libc internals. ( libc::SYS_open, vec![ @@ -518,6 +547,20 @@ impl litebox::platform::SignalProvider for LinuxUserland { /// Atomically takes the per-thread pending host signal bitmask. fn take_pending_host_signals() -> litebox_common_linux::signal::SigSet { + // Fast path: a plain load first. Only this thread consumes the mask + // (signal handlers merely OR bits in), so a zero read means there is + // nothing to take and the locked `xchg` can be skipped. + let pending: u32; + unsafe { + core::arch::asm!( + concat!("mov {tmp:e}, DWORD PTR ", tls!("pending_host_signals")), + tmp = out(reg) pending, + options(nostack, preserves_flags, readonly) + ); + } + if pending == 0 { + return litebox_common_linux::signal::SigSet::from_u64(0); + } // Atomically swap the per-thread pending signals with zero. // Only the low 32 bits are used (covers traditional signals 1-31). let lo: u32; @@ -621,6 +664,8 @@ pending_host_signals: .globl wait_waker_addr wait_waker_addr: .quad 0 +host_fsbase: + .quad 0 " ); @@ -685,10 +730,37 @@ unsafe extern "C-unwind" fn run_thread_arch( lea r8, [rsi + {GUEST_CONTEXT_SIZE}] mov fs:guest_context_top@tpoff, r8 - // Save host fs base in gs base. This will stay set for the lifetime + // Save host fs base in gs base (and mirror it in the host_fsbase TLS + // slot for the no-FSGSBASE fallback). This will stay set for the lifetime // of this call stack. + cmp BYTE PTR [rip + {have_fsgsbase}], {FSGSBASE_NATIVE} + jne .Lrta_no_fsgsbase rdfsbase r8 wrgsbase r8 + jmp .Lrta_gs_done +.Lrta_no_fsgsbase: + // arch_prctl(ARCH_GET_FS) then arch_prctl(ARCH_SET_GS). Raw syscalls + // clobber rax/rcx/r11 (dead here) plus the registers we load; rdx still + // holds the live `reenter` flag, so preserve it in r9. + mov r9, rdx + sub rsp, 8 + mov edi, {ARCH_GET_FS} + mov rsi, rsp + mov eax, {SYS_arch_prctl} + syscall + pop r8 // host fs base + mov edi, {ARCH_SET_GS} + mov rsi, r8 + mov eax, {SYS_arch_prctl} + syscall + mov rdx, r9 +.Lrta_gs_done: + mov fs:host_fsbase@tpoff, r8 + + // Reload the thread-context pointer (first handler argument): the + // no-FSGSBASE path above clobbers rdi/rsi with arch_prctl arguments. + // The pushed `rdi` (thread_ctx) is at the top of the stack. + mov rdi, [rsp] // Call init_handler or reenter_handler based on reenter flag (in dl). test dl, dl @@ -711,15 +783,12 @@ syscall_callback: // expectations of `interrupt_signal_handler`. mov BYTE PTR gs:in_guest@tpoff, 0 - // Restore host fs base. - rdfsbase r11 - mov gs:guest_fsbase@tpoff, r11 - rdgsbase r11 - wrfsbase r11 - - // Switch to the top of the guest context. + // Switch to the top of the guest context. The guest context is saved + // before the host fs base is restored so that the no-FSGSBASE fallback + // below may clobber registers and flags freely; until then only + // gs-relative TLS (host base) and flag-preserving instructions are used. mov r11, rsp - mov rsp, fs:guest_context_top@tpoff + mov rsp, gs:guest_context_top@tpoff // TODO: save float and vector registers (xsave or fxsave) // Save caller-saved registers @@ -746,6 +815,26 @@ syscall_callback: push r14 // pt_regs->r14 push r15 // pt_regs->r15 + // Restore host fs base, capturing the guest's current fs base first. + // The full guest register frame is saved at this point, so both paths + // may clobber caller-saved registers and flags. + cmp BYTE PTR [rip + {have_fsgsbase}], {FSGSBASE_NATIVE} + jne .Lsc_no_fsgsbase + rdfsbase r11 + mov gs:guest_fsbase@tpoff, r11 + rdgsbase r11 + wrfsbase r11 + jmp .Lsc_fs_done +.Lsc_no_fsgsbase: + // Without FSGSBASE the guest cannot have changed its fs base directly + // (wrfsbase would fault), so gs:guest_fsbase is already current; just + // restore the host fs base saved at entry via arch_prctl(ARCH_SET_FS). + mov rsi, gs:host_fsbase@tpoff + mov edi, {ARCH_SET_FS} + mov eax, {SYS_arch_prctl} + syscall +.Lsc_fs_done: + // Restore the stack and frame pointer. mov rsp, fs:host_sp@tpoff mov rbp, fs:host_bp@tpoff @@ -793,6 +882,12 @@ interrupt_callback: syscall_handler = sym syscall_handler, exception_handler = sym exception_handler, interrupt_handler = sym interrupt_handler, + have_fsgsbase = sym HAVE_FSGSBASE, + FSGSBASE_NATIVE = const FSGSBASE_NATIVE, + ARCH_GET_FS = const ARCH_GET_FS, + ARCH_SET_GS = const ARCH_SET_GS, + ARCH_SET_FS = const ARCH_SET_FS, + SYS_arch_prctl = const libc::SYS_arch_prctl, ); } @@ -822,9 +917,20 @@ unsafe extern "C" fn switch_to_guest(ctx: &litebox_common_linux::PtRegs) -> ! { "jne interrupt_callback", // Restore guest context from ctx. "mov rsp, rdi", - // Switch to the guest fsbase + // Switch to the guest fsbase. Every register clobbered below (rax, + // rcx, rdx, rsi, rdi, r11, rflags) is restored from the guest context + // by the pops that follow. + "cmp BYTE PTR [rip + {have_fsgsbase}], {FSGSBASE_NATIVE}", + "jne .Lstg_no_fsgsbase", "mov rdx, fs:guest_fsbase@tpoff", "wrfsbase rdx", + "jmp .Lstg_fs_done", + ".Lstg_no_fsgsbase:", + "mov rsi, fs:guest_fsbase@tpoff", + "mov edi, {ARCH_SET_FS}", + "mov eax, {SYS_arch_prctl}", + "syscall", + ".Lstg_fs_done:", "pop r15", "pop r14", "pop r13", @@ -847,6 +953,10 @@ unsafe extern "C" fn switch_to_guest(ctx: &litebox_common_linux::PtRegs) -> ! { "pop rsp", "jmp gs:scratch@tpoff", // jump to the guest "switch_to_guest_end:", + have_fsgsbase = sym HAVE_FSGSBASE, + FSGSBASE_NATIVE = const FSGSBASE_NATIVE, + ARCH_SET_FS = const ARCH_SET_FS, + SYS_arch_prctl = const libc::SYS_arch_prctl, ); } @@ -980,19 +1090,31 @@ impl litebox::platform::ThreadProvider for LinuxUserland { thread.interrupt(); } + fn num_cpus(&self) -> Option { + // Cache the CPU count: it does not change over the process lifetime, + // and `available_parallelism` probes the host (and may consult files + // the seccomp filter blocks). Seeded pre-seccomp in `new`; the + // `get_or_init` fallback here keeps it correct if queried first. + *NUM_CPUS.get_or_init(|| std::thread::available_parallelism().ok()) + } + #[cfg(debug_assertions)] fn run_test_thread(f: impl FnOnce() -> R) -> R { // Sets `gsbase = fsbase` (x86_64) or `fs = gs` (x86) on the current thread // to mirror the TLS base used in guest context, so that test threads can use the // same TLS access code as guest threads. #[cfg(target_arch = "x86_64")] - unsafe { - core::arch::asm!( - "rdfsbase {tmp}", - "wrgsbase {tmp}", - tmp = out(reg) _, - options(nostack, preserves_flags), - ); + if have_fsgsbase() { + unsafe { + core::arch::asm!( + "rdfsbase {tmp}", + "wrgsbase {tmp}", + tmp = out(reg) _, + options(nostack, preserves_flags), + ); + } + } else { + arch_prctl_set_gs(arch_prctl_get_fs()); } ThreadHandle::run_with_handle(f) @@ -1187,9 +1309,17 @@ impl litebox::platform::IPInterfaceProvider for LinuxUserland { } Ok(()) } - Err(errno) => { - unimplemented!("unexpected error {errno}") - } + Err(errno) => match errno { + // The TUN fd is opened O_NONBLOCK, so under ordinary host-side + // backpressure (a full transmit queue) this is expected, + // not fatal — mirrors the EWOULDBLOCK/EAGAIN handling in + // `receive_ip_packet` below. + #[allow(unreachable_patterns, reason = "EAGAIN == EWOULDBLOCK")] + syscalls::Errno::EWOULDBLOCK | syscalls::Errno::EAGAIN => { + Err(litebox::platform::SendError::WouldBlock) + } + _ => unimplemented!("unexpected error {errno}"), + }, } } @@ -1225,7 +1355,12 @@ impl litebox::platform::TimeProvider for LinuxUserland { fn now(&self) -> Self::Instant { let mut t = core::mem::MaybeUninit::::uninit(); - unsafe { libc::clock_gettime(libc::CLOCK_MONOTONIC, t.as_mut_ptr()) }; + let ret = unsafe { libc::clock_gettime(libc::CLOCK_MONOTONIC, t.as_mut_ptr()) }; + assert!( + ret == 0, + "clock_gettime(CLOCK_MONOTONIC) failed: {}", + std::io::Error::last_os_error() + ); let t = unsafe { t.assume_init() }; Instant { #[cfg_attr(target_arch = "x86_64", expect(clippy::useless_conversion))] @@ -1238,7 +1373,12 @@ impl litebox::platform::TimeProvider for LinuxUserland { fn current_time(&self) -> Self::SystemTime { let mut t = core::mem::MaybeUninit::::uninit(); - unsafe { libc::clock_gettime(libc::CLOCK_REALTIME, t.as_mut_ptr()) }; + let ret = unsafe { libc::clock_gettime(libc::CLOCK_REALTIME, t.as_mut_ptr()) }; + assert!( + ret == 0, + "clock_gettime(CLOCK_REALTIME) failed: {}", + std::io::Error::last_os_error() + ); let t = unsafe { t.assume_init() }; SystemTime { #[cfg_attr(target_arch = "x86_64", expect(clippy::useless_conversion))] @@ -1349,7 +1489,6 @@ enum FutexOperation { } /// Safer invocation of the Linux futex syscall, with the "timeout" variant of the arguments. -#[expect(clippy::similar_names, reason = "sec/nsec are as needed by libc")] fn futex_timeout( uaddr: &AtomicU32, futex_op: FutexOperation, @@ -1359,20 +1498,9 @@ fn futex_timeout( ) -> Result { let uaddr: *const AtomicU32 = std::ptr::from_ref(uaddr); let futex_op: i32 = futex_op as _; - let timeout = timeout.map(|t| { - const TEN_POWER_NINE: u128 = 1_000_000_000; - let nanos: u128 = t.as_nanos(); - let tv_sec = nanos - .checked_div(TEN_POWER_NINE) - .unwrap() - .try_into() - .unwrap(); - let tv_nsec = nanos - .checked_rem(TEN_POWER_NINE) - .unwrap() - .try_into() - .unwrap(); - litebox_common_linux::Timespec { tv_sec, tv_nsec } + let timeout = timeout.map(|t| litebox_common_linux::Timespec { + tv_sec: t.as_secs().try_into().unwrap(), + tv_nsec: t.subsec_nanos().into(), }); let uaddr2: *const AtomicU32 = uaddr2.map_or(std::ptr::null(), |u| u); unsafe { @@ -1556,6 +1684,25 @@ impl litebox::platform::PageManagementProvider for Li Ok(()) } + unsafe fn discard_pages( + &self, + range: core::ops::Range, + ) -> Result<(), litebox::platform::page_mgmt::DiscardError> { + // A single host `madvise(MADV_DONTNEED)` resets anonymous private pages to + // zero-fill-on-demand without remapping. On failure, report `Unsupported` so + // that the caller falls back to re-creating the mapping. + unsafe { + syscalls::syscall3( + syscalls::Sysno::madvise, + range.start, + range.len(), + usize::try_from(libc::MADV_DONTNEED).unwrap(), + ) + } + .map(|_| ()) + .map_err(|_| litebox::platform::page_mgmt::DiscardError::Unsupported) + } + fn reserved_pages(&self) -> impl Iterator> { self.reserved_pages.iter() } @@ -1567,26 +1714,13 @@ impl litebox::platform::PageManagementProvider for Li permissions: MemoryRegionPermissions, fixed_address_behavior: FixedAddressBehavior, ) -> Result, CowAllocationError> { - let Some((file_path, file_offset)) = self.lookup_cow_region(source_data) else { + let Some((fd, file_offset)) = self.lookup_cow_region(source_data) else { return Err(CowAllocationError::UnsupportedSourceRegion); }; if !file_offset.is_multiple_of(ALIGN) { return Err(CowAllocationError::Unaligned); } - let file_path_cstr = - std::ffi::CString::new(file_path.as_os_str().as_encoded_bytes()).unwrap(); - // TODO(jb): We should likely be storing pre-opened FDs, right? - let fd = unsafe { - syscalls::syscall3( - syscalls::Sysno::open, - file_path_cstr.as_ptr() as usize, - OFlags::RDONLY.bits() as usize, - 0, - ) - }; - let fd = fd.expect("file should remain unchanged on host"); - let mut flags = MapFlags::MAP_PRIVATE; match fixed_address_behavior { FixedAddressBehavior::Hint => {} @@ -1606,7 +1740,7 @@ impl litebox::platform::PageManagementProvider for Li source_data.len(), prot_flags(permissions).bits().reinterpret_as_unsigned() as usize, flags.bits().reinterpret_as_unsigned() as usize, - fd, + usize::try_from(fd).unwrap(), { #[cfg(target_arch = "x86_64")] { @@ -1616,8 +1750,6 @@ impl litebox::platform::PageManagementProvider for Li ) }; - let _ = unsafe { syscalls::syscall1(syscalls::Sysno::close, fd) }; - match result { Ok(ptr) => Ok(UserMutPtr::from_usize(ptr)), Err(_) => Err(CowAllocationError::InternalFailure), @@ -1754,7 +1886,20 @@ impl ThreadContext<'_> { } let op = f(self.shim, self.ctx); match op { - ContinueOperation::Resume => unsafe { switch_to_guest(self.ctx) }, + ContinueOperation::Resume => { + // Guest-controlled paths (notably `rt_sigreturn`, which copies + // a guest-memory-supplied `rip`/`rsp`/`eflags` into `ctx` + // unchecked) could otherwise resume the guest at an arbitrary + // address with an arbitrary stack and privileged RFLAGS bits. + // Validate and normalize before every resume, exactly as the + // `lvbs`/`snp` platforms do (see + // `PtRegs::sanitize_for_user_return`'s doc comment); a thread + // that fails this check is terminated instead of resumed. + if self.ctx.sanitize_for_user_return() { + unsafe { switch_to_guest(self.ctx) } + } + litebox_util_log::warn!("terminating thread with invalid user return context"); + } ContinueOperation::Terminate => {} } } @@ -1793,9 +1938,135 @@ unsafe impl litebox::platform::ThreadLocalStorageProvider for LinuxUserland { } } +/// Whether user-mode FSGSBASE instructions are usable on this host. +/// +/// Values: `0` = not yet detected, `1` = unavailable (use the `arch_prctl` +/// fallback), `2` = available (use `rdfsbase`/`wrgsbase` etc. natively). +/// +/// The guest/host transition assembly compares this byte against +/// [`FSGSBASE_NATIVE`] and takes the fallback path otherwise, so the +/// not-yet-detected state is safe on any hardware (merely slower). +static HAVE_FSGSBASE: core::sync::atomic::AtomicU8 = core::sync::atomic::AtomicU8::new(0); + +/// [`HAVE_FSGSBASE`] value meaning FSGSBASE instructions may be used. +const FSGSBASE_NATIVE: u8 = 2; + +/// `arch_prctl` operation codes from `asm/prctl.h` (not exposed by `libc`). +const ARCH_SET_GS: usize = 0x1001; +const ARCH_SET_FS: usize = 0x1002; +const ARCH_GET_FS: usize = 0x1003; +const ARCH_GET_GS: usize = 0x1004; + +/// Detects whether user-mode FSGSBASE instructions are available. +/// +/// The guest/host transition code prefers `rdfsbase`/`wrgsbase` etc., which +/// raise `#UD` (delivered as `SIGILL`) unless the kernel has set +/// `CR4.FSGSBASE` (CPU too old, kernel older than 5.9, or booted with +/// `nofsgsbase` — including PVM guests on such hosts). When unavailable, the +/// transition code falls back to `arch_prctl` syscalls with the host FS base +/// cached in the `host_fsbase` TLS slot. +/// +/// Setting `LITEBOX_NO_FSGSBASE` in the environment forces the fallback even +/// on capable hosts (primarily for testing the fallback path). +fn init_fsgsbase_mode() { + /// `HWCAP2_FSGSBASE` from `asm/hwcap2.h` (not exposed by the `libc` crate). + const HWCAP2_FSGSBASE: libc::c_ulong = 1 << 1; + let hwcap2 = unsafe { libc::getauxval(libc::AT_HWCAP2) }; + let native = hwcap2 & HWCAP2_FSGSBASE != 0 && std::env::var_os("LITEBOX_NO_FSGSBASE").is_none(); + HAVE_FSGSBASE.store( + if native { FSGSBASE_NATIVE } else { 1 }, + core::sync::atomic::Ordering::Release, + ); +} + +/// Returns true if FSGSBASE instructions may be used, detecting on first use. +/// +/// Signal-handler callers are safe: detection happens at platform creation +/// (and this function is idempotent), and both `getauxval` and `env::var_os` +/// here only run outside signal context via `LinuxUserland::new` or +/// `run_test_thread`. +fn have_fsgsbase() -> bool { + match HAVE_FSGSBASE.load(core::sync::atomic::Ordering::Acquire) { + 0 => { + init_fsgsbase_mode(); + HAVE_FSGSBASE.load(core::sync::atomic::Ordering::Acquire) == FSGSBASE_NATIVE + } + v => v == FSGSBASE_NATIVE, + } +} + +/// Reads the current thread's GS base without FSGSBASE instructions. +/// +/// Uses a raw `arch_prctl(ARCH_GET_GS)` syscall, which is async-signal-safe +/// and does not touch `errno` (the caller may be running with the guest's FS +/// base, where libc's TLS-based `errno` would corrupt guest memory). +fn arch_prctl_get_gs() -> u64 { + let mut value = 0u64; + let r = unsafe { + syscalls::syscall2( + syscalls::Sysno::arch_prctl, + ARCH_GET_GS, + core::ptr::from_mut(&mut value) as usize, + ) + }; + assert!(r.is_ok(), "arch_prctl(ARCH_GET_GS) failed"); + value +} + +/// Sets the current thread's FS base without FSGSBASE instructions. +/// +/// Async-signal-safe; see [`arch_prctl_get_gs`]. +fn arch_prctl_set_fs(value: u64) { + let r = unsafe { + syscalls::syscall2( + syscalls::Sysno::arch_prctl, + ARCH_SET_FS, + usize::try_from(value).unwrap(), + ) + }; + assert!(r.is_ok(), "arch_prctl(ARCH_SET_FS) failed"); +} + +/// Reads the current thread's FS base without FSGSBASE instructions. +/// +/// Async-signal-safe; see [`arch_prctl_get_gs`]. +#[cfg(debug_assertions)] +fn arch_prctl_get_fs() -> u64 { + let mut value = 0u64; + let r = unsafe { + syscalls::syscall2( + syscalls::Sysno::arch_prctl, + ARCH_GET_FS, + core::ptr::from_mut(&mut value) as usize, + ) + }; + assert!(r.is_ok(), "arch_prctl(ARCH_GET_FS) failed"); + value +} + +/// Sets the current thread's GS base without FSGSBASE instructions. +/// +/// Async-signal-safe; see [`arch_prctl_get_gs`]. +#[cfg(debug_assertions)] +fn arch_prctl_set_gs(value: u64) { + let r = unsafe { + syscalls::syscall2( + syscalls::Sysno::arch_prctl, + ARCH_SET_GS, + usize::try_from(value).unwrap(), + ) + }; + assert!(r.is_ok(), "arch_prctl(ARCH_SET_GS) failed"); +} + static mut NEXT_SA: [libc::sigaction; 64] = unsafe { core::mem::zeroed() }; static INTERRUPT_SIGNAL_NUMBER: AtomicI32 = AtomicI32::new(0); +/// Cached host CPU count, seeded before the seccomp filter is installed so +/// that querying it never trips the sandbox. See `LinuxUserland`'s +/// `ThreadProvider::num_cpus` implementation. +static NUM_CPUS: std::sync::OnceLock> = std::sync::OnceLock::new(); + fn register_exception_handlers() { static ONCE: std::sync::Once = std::sync::Once::new(); ONCE.call_once(|| { @@ -1888,41 +2159,53 @@ fn register_exception_handlers() { }); } +/// Free list of alternate signal stacks (base addresses of the guard page). +/// +/// Stacks are uniform in size, so they can be reused by any thread. Reuse +/// avoids an mmap/mprotect/munmap round trip per guest thread; the pool is +/// never shrunk (stacks are reclaimed only at process exit). +static ALT_STACK_POOL: std::sync::Mutex> = std::sync::Mutex::new(Vec::new()); + /// Runs `f` with an alternate signal stack set up. fn with_signal_alt_stack(f: impl FnOnce() -> R) -> R { let alt_stack_size = libc::SIGSTKSZ * 2; let guard_page_size = 0x1000; - let stack_base = unsafe { - libc::mmap( - std::ptr::null_mut(), - guard_page_size + alt_stack_size, - libc::PROT_READ | libc::PROT_WRITE, - libc::MAP_PRIVATE | libc::MAP_ANONYMOUS, - -1, - 0, - ) - }; - assert!( - stack_base != libc::MAP_FAILED, - "failed to allocate memory for alternate signal stack: {}", - std::io::Error::last_os_error() - ); - let _unmap_guard = litebox::utils::defer(|| { - let r = unsafe { libc::munmap(stack_base, guard_page_size + alt_stack_size) }; - assert!( - r == 0, - "failed to free memory for alternate signal stack: {}", - std::io::Error::last_os_error() - ); - }); + let pooled = ALT_STACK_POOL.lock().unwrap().pop(); + let stack_base = pooled.map_or_else( + || { + let stack_base = unsafe { + libc::mmap( + std::ptr::null_mut(), + guard_page_size + alt_stack_size, + libc::PROT_READ | libc::PROT_WRITE, + libc::MAP_PRIVATE | libc::MAP_ANONYMOUS, + -1, + 0, + ) + }; + assert!( + stack_base != libc::MAP_FAILED, + "failed to allocate memory for alternate signal stack: {}", + std::io::Error::last_os_error() + ); - // Set up a guard page to catch stack overflows. - let r = unsafe { libc::mprotect(stack_base, guard_page_size, libc::PROT_NONE) }; - assert!( - r == 0, - "failed to set guard page for alternate signal stack: {}", - std::io::Error::last_os_error() + // Set up a guard page to catch stack overflows. + let r = unsafe { libc::mprotect(stack_base, guard_page_size, libc::PROT_NONE) }; + assert!( + r == 0, + "failed to set guard page for alternate signal stack: {}", + std::io::Error::last_os_error() + ); + stack_base + }, + |addr| addr as *mut libc::c_void, ); + // Return the stack to the pool once we are no longer running on it. This + // runs after `_restore_guard` below (drop order is reverse declaration + // order), i.e. only after the original signal stack has been restored. + let _pool_guard = litebox::utils::defer(|| { + ALT_STACK_POOL.lock().unwrap().push(stack_base as usize); + }); let alt_stack = libc::stack_t { ss_sp: stack_base.cast(), @@ -1966,9 +2249,14 @@ fn signal_handler_exit_guest( set_interrupt: bool, ) -> Option<*mut litebox_common_linux::PtRegs> { unsafe { - let gsbase: u64; - core::arch::asm! { - "rdgsbase {}", out(reg) gsbase + let gsbase: u64 = if have_fsgsbase() { + let gsbase: u64; + core::arch::asm! { + "rdgsbase {}", out(reg) gsbase + }; + gsbase + } else { + arch_prctl_get_gs() }; let is_in_guest = if gsbase == 0 { false @@ -1992,11 +2280,18 @@ fn signal_handler_exit_guest( return None; } + if have_fsgsbase() { + core::arch::asm! { + "wrfsbase {gsbase}", + gsbase = in(reg) gsbase, + options(nostack, preserves_flags) + }; + } else { + arch_prctl_set_fs(gsbase); + } let guest_context_top: *mut litebox_common_linux::PtRegs; core::arch::asm! { - "wrfsbase {gsbase}", "mov {guest_context_top}, fs:guest_context_top@tpoff", - gsbase = in(reg) gsbase, guest_context_top = out(reg) guest_context_top, options(nostack, preserves_flags) }; @@ -2272,8 +2567,13 @@ unsafe fn interrupt_signal_handler( let is_guest_thread; #[cfg(target_arch = "x86_64")] { - let gsbase: u64; - unsafe { core::arch::asm!("rdgsbase {}", out(reg) gsbase) }; + let gsbase: u64 = if have_fsgsbase() { + let gsbase: u64; + unsafe { core::arch::asm!("rdgsbase {}", out(reg) gsbase) }; + gsbase + } else { + arch_prctl_get_gs() + }; is_guest_thread = gsbase != 0; } diff --git a/litebox_platform_lvbs/src/arch/x86/mm/paging.rs b/litebox_platform_lvbs/src/arch/x86/mm/paging.rs index 4e2866e85..c732890cd 100644 --- a/litebox_platform_lvbs/src/arch/x86/mm/paging.rs +++ b/litebox_platform_lvbs/src/arch/x86/mm/paging.rs @@ -134,12 +134,24 @@ impl FrameDeallocator for PageTableAllocator { } pub(crate) fn vmflags_to_pteflags(values: VmFlags) -> PageTableFlags { - let mut flags = PageTableFlags::empty(); + // Pre-set ACCESSED on every faulted-in leaf PTE. Otherwise the CPU's page-table + // walker sets the A bit itself on first access with a locked, microcoded write + // that takes the PTE cache line for exclusive ownership (an RFO). Filling in the + // software-known result up front avoids that write, cutting cache-coherency + // traffic and slightly speeding up the walk. Mirrors Linux's pte_mkyoung() on + // fault-in. ACCESSED/DIRTY are excluded from MPROTECT_PTE_MASK, so mprotect never + // observes or clears them, and no code in this platform reads A/D for data pages, + // so this is a pure hardware-walk optimization with no semantic effect. + let mut flags = PageTableFlags::ACCESSED; if values.intersects(VmFlags::VM_READ | VmFlags::VM_WRITE) { flags |= PageTableFlags::USER_ACCESSIBLE; } if values.contains(VmFlags::VM_WRITE) { - flags |= PageTableFlags::WRITABLE; + // Writable leaves are additionally pre-marked DIRTY (matching Linux's + // pte_mkdirty()): the first write would otherwise cost a second locked walker + // write to set D. DIRTY is only ever set together with WRITABLE, never on a + // read-only PTE. + flags |= PageTableFlags::WRITABLE | PageTableFlags::DIRTY; } if !values.contains(VmFlags::VM_EXEC) { flags |= PageTableFlags::NO_EXECUTE; @@ -925,9 +937,15 @@ impl PageTableImpl for X64PageTabl let mut allocator = PageTableAllocator::::new(); // TODO: if it is file-backed, we need to read the page from file let frame = PageTableAllocator::::allocate_frame(true).unwrap(); + // Intermediate (non-leaf) entries the mapper may create for this leaf. + // ACCESSED is pre-set for the same reason as on the leaf: it spares the + // walker the locked A-bit write on these entries (Linux's _PAGE_TABLE + // includes _PAGE_ACCESSED). DIRTY is meaningless on non-leaf entries and + // is deliberately omitted. let table_flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE - | PageTableFlags::USER_ACCESSIBLE; + | PageTableFlags::USER_ACCESSIBLE + | PageTableFlags::ACCESSED; match unsafe { inner.map_to_with_table_flags( page, diff --git a/litebox_shim_linux/src/lib.rs b/litebox_shim_linux/src/lib.rs index 50fe14a2f..76be9093c 100644 --- a/litebox_shim_linux/src/lib.rs +++ b/litebox_shim_linux/src/lib.rs @@ -257,6 +257,7 @@ impl LinuxShim { fs: Arc::new(syscalls::file::FsState::new()).into(), files: files.into(), signals: syscalls::signal::SignalState::new_process(), + scratch_buf: RefCell::new(Vec::new()), }, }; @@ -474,6 +475,26 @@ impl ToSyscallResult for Result { } impl Task { + /// Runs `f` with the per-task scratch buffer, sized to `len` (capped at + /// [`MAX_KERNEL_BUF_SIZE`]) bytes. + /// + /// This avoids a fresh (and zeroed) allocation on every I/O syscall. The + /// buffer contents are unspecified on entry; callers must only rely on + /// bytes they have written themselves. + /// + /// # Panics + /// + /// Panics if used reentrantly (i.e., if `f` itself borrows the scratch + /// buffer). + pub(crate) fn with_scratch_buf(&self, len: usize, f: impl FnOnce(&mut [u8]) -> R) -> R { + let len = len.min(MAX_KERNEL_BUF_SIZE); + let mut buf = self.scratch_buf.borrow_mut(); + if buf.len() < len { + buf.resize(len, 0); + } + f(&mut buf[..len]) + } + /// A wrapper function around `sys_pread64` that copies data in chunks to avoid OOMing. fn pread_with_user_buf( &self, @@ -482,26 +503,72 @@ impl Task { count: usize, offset: i64, ) -> Result { - let mut kernel_buf = vec![0u8; count.min(MAX_KERNEL_BUF_SIZE)]; - let mut read_total = 0; - while read_total < count { - let to_read = (count - read_total).min(kernel_buf.len()); - match self.sys_pread64( - fd, - &mut kernel_buf[..to_read], - offset + (read_total.reinterpret_as_signed() as i64), - ) { - Ok(0) => break, // EOF - Ok(size) => { - buf.copy_from_slice(read_total, &kernel_buf[..size]) - .ok_or(Errno::EFAULT)?; - read_total += size; + self.with_scratch_buf(count, |kernel_buf| { + let mut read_total = 0; + while read_total < count { + let to_read = (count - read_total).min(kernel_buf.len()); + match self.sys_pread64( + fd, + &mut kernel_buf[..to_read], + offset + (read_total.reinterpret_as_signed() as i64), + ) { + Ok(0) => break, // EOF + Ok(size) => { + buf.copy_from_slice(read_total, &kernel_buf[..size]) + .ok_or(Errno::EFAULT)?; + read_total += size; + } + Err(e) => return Err(e), } - Err(e) => return Err(e), } + assert!(read_total <= count); + Ok(read_total) + }) + } + + /// A wrapper around `sys_write` that copies the user buffer through the + /// bounded per-task scratch buffer in chunks, to avoid an unbounded + /// allocation for large writes. + fn write_with_user_buf( + &self, + fd: i32, + buf: ConstPtr, + count: usize, + ) -> Result { + // Splitting a single write into several `sys_write` calls would break + // message-boundary and `O_APPEND` atomicity (a datagram socket would + // fragment one message; concurrent appends could interleave). Chunking + // is only safe within one `sys_write`, i.e. up to the scratch buffer's + // cap, so larger writes fall back to a single owned copy and a single + // call — the same guard `sys_sendto` uses. + if count > MAX_KERNEL_BUF_SIZE { + let owned = buf.to_owned_slice(count).ok_or(Errno::EFAULT)?; + return self.sys_write(fd, &owned, None); } - assert!(read_total <= count); - Ok(read_total) + self.with_scratch_buf(count, |kernel_buf| { + // Once any bytes have been delivered, an error collapses to + // `Ok(written_total)` so partial progress is reported to user space. + let bail = |total: usize, e: Errno| if total > 0 { Ok(total) } else { Err(e) }; + let mut written_total = 0; + loop { + let to_write = (count - written_total).min(kernel_buf.len()); + if buf + .copy_to_slice(written_total, &mut kernel_buf[..to_write]) + .is_none() + { + return bail(written_total, Errno::EFAULT); + } + let size = match self.sys_write(fd, &kernel_buf[..to_write], None) { + Ok(size) => size, + Err(e) => return bail(written_total, e), + }; + written_total += size; + if written_total >= count || size < to_write { + // Okay to transfer fewer bytes than requested. + return Ok(written_total); + } + } + }) } /// Handle Linux syscalls and dispatch them to LiteBox implementations. @@ -551,11 +618,12 @@ impl Task { // Note some applications (e.g., `node`) seem to assume that getting fewer bytes than // requested indicates EOF. if count <= MAX_KERNEL_BUF_SIZE { - let mut kernel_buf = vec![0u8; count.min(MAX_KERNEL_BUF_SIZE)]; - self.sys_read(fd, &mut kernel_buf, None).and_then(|size| { - buf.copy_from_slice(0, &kernel_buf[..size]) - .map(|()| size) - .ok_or(Errno::EFAULT) + self.with_scratch_buf(count, |kernel_buf| { + self.sys_read(fd, kernel_buf, None).and_then(|size| { + buf.copy_from_slice(0, &kernel_buf[..size]) + .map(|()| size) + .ok_or(Errno::EFAULT) + }) }) } else { // If the read size is too large, we need to do some extra work to avoid OOMing. @@ -590,10 +658,7 @@ impl Task { }) } } - SyscallRequest::Write { fd, buf, count } => match buf.to_owned_slice(count) { - Some(buf) => self.sys_write(fd, &buf, None), - None => Err(Errno::EFAULT), - }, + SyscallRequest::Write { fd, buf, count } => self.write_with_user_buf(fd, buf, count), SyscallRequest::Close { fd } => syscall!(sys_close(fd)), SyscallRequest::Lseek { fd, offset, whence } => { use litebox::utils::TruncateExt as _; @@ -1139,6 +1204,8 @@ struct Task { files: RefCell>>, /// Signal state signals: syscalls::signal::SignalState, + /// Reusable kernel-side I/O buffer (see [`Task::with_scratch_buf`]). + scratch_buf: RefCell>, } impl Drop for Task { @@ -1176,6 +1243,7 @@ mod test_utils { fs: Arc::new(syscalls::file::FsState::new()).into(), files: files.into(), signals: syscalls::signal::SignalState::new_process(), + scratch_buf: RefCell::new(Vec::new()), global: self, } } @@ -1200,6 +1268,7 @@ mod test_utils { fs: self.fs.clone(), files: self.files.clone(), signals: self.signals.clone_for_new_task(), + scratch_buf: RefCell::new(Vec::new()), }; Some(task) } diff --git a/litebox_shim_linux/src/syscalls/file.rs b/litebox_shim_linux/src/syscalls/file.rs index 06020e270..958ad431b 100644 --- a/litebox_shim_linux/src/syscalls/file.rs +++ b/litebox_shim_linux/src/syscalls/file.rs @@ -765,8 +765,13 @@ impl Task { Ok(fd) => ConsumedFd::Fs(fd), Err(litebox::fd::ErrRawIntFd::NotFound) => { if let Some(new_fd) = replace { - let success = rds.fd_into_specific_raw_integer(new_fd, raw_fd); - assert!(success, "raw_fd slot is empty, so insert must succeed"); + // The slot is confirmed empty, so a failure here means + // `raw_fd` exceeds the backing store's growth allowance + // (see `fd_into_specific_raw_integer`), not that it is + // occupied. + if !rds.fd_into_specific_raw_integer(new_fd, raw_fd) { + return Err(Errno::ENOMEM); + } } return Err(Errno::EBADF); } @@ -799,11 +804,13 @@ impl Task { // Insert the replacement into the now-vacated slot while still holding the lock. if let Some(new_fd) = replace { - let success = rds.fd_into_specific_raw_integer(new_fd, raw_fd); - assert!( - success, - "we just consumed this raw_fd, so it must be available" - ); + // As above: we just consumed this raw_fd, so a failure here means + // it exceeds the backing store's growth allowance, not that it is + // occupied. + if !rds.fd_into_specific_raw_integer(new_fd, raw_fd) { + drop(rds); + return Err(Errno::ENOMEM); + } } drop(rds); @@ -866,11 +873,15 @@ impl Task { self.check_raw_fd_exists(fd)?; check_iovcnt(iovcnt)?; let iovs: &[IoReadVec>] = &iovec.to_owned_slice(iovcnt).ok_or(Errno::EFAULT)?; - let mut kernel_buffer = vec![0u8; PAGE_SIZE]; - read_from_iovec(iovs, &mut kernel_buffer, |buf, total| { - let cur_offset = base_offset.checked_add(total).ok_or(Errno::EOVERFLOW)?; - self.sys_read(fd, buf, Some(cur_offset)) - }) + self.with_scratch_buf( + total_iov_len(iovs.iter().map(|iov| iov.iov_len)), + |kernel_buffer| { + read_from_iovec(iovs, kernel_buffer, |buf, total| { + let cur_offset = base_offset.checked_add(total).ok_or(Errno::EOVERFLOW)?; + self.sys_read(fd, buf, Some(cur_offset)) + }) + }, + ) } /// Handle syscall `pwritev` @@ -903,13 +914,17 @@ impl Task { self.check_raw_fd_exists(fd)?; check_iovcnt(iovcnt)?; let iovs: &[IoReadVec>] = &iovec.to_owned_slice(iovcnt).ok_or(Errno::EFAULT)?; - let mut kernel_buffer = vec![0u8; PAGE_SIZE]; // TODO: The data transfers performed by readv() and writev() are atomic: the data // written by writev() is written as a single block that is not intermingled with // output from writes in other processes - read_from_iovec(iovs, &mut kernel_buffer, |buf, _total| { - self.sys_read(fd, buf, None) - }) + self.with_scratch_buf( + total_iov_len(iovs.iter().map(|iov| iov.iov_len)), + |kernel_buffer| { + read_from_iovec(iovs, kernel_buffer, |buf, _total| { + self.sys_read(fd, buf, None) + }) + }, + ) } } @@ -954,6 +969,12 @@ fn check_iov_lens(iov_lens: impl IntoIterator) -> Result<(), Errno Ok(()) } +/// Total byte count of a set of iovecs, saturated; used only to size kernel +/// buffers (overlong totals are rejected separately by [`check_iov_lens`]). +fn total_iov_len(iov_lens: impl IntoIterator) -> usize { + iov_lens.into_iter().fold(0usize, usize::saturating_add) +} + /// Drain reads into a sequence of user iovecs. fn read_from_iovec( iovs: &[IoReadVec

], @@ -1022,17 +1043,19 @@ where continue; } if kernel_buffer.is_empty() { - kernel_buffer.resize(PAGE_SIZE, 0); + kernel_buffer.resize( + total_iov_len(iovs.iter().map(|iov| iov.iov_len)).min(crate::MAX_KERNEL_BUF_SIZE), + 0, + ); } let mut iov_written = 0; while iov_written < iov_len { let to_write = (iov_len - iov_written).min(kernel_buffer.len()); - let base_offset = isize::try_from(iov_written).unwrap(); - for (byte_offset, byte) in (0_isize..).zip(kernel_buffer[..to_write].iter_mut()) { - let Some(value) = iov_base.read_at_offset(base_offset + byte_offset) else { - return bail(total_written, Errno::EFAULT); - }; - *byte = value; + if iov_base + .copy_to_slice(iov_written, &mut kernel_buffer[..to_write]) + .is_none() + { + return bail(total_written, Errno::EFAULT); } let size = match write_fn(&kernel_buffer[..to_write], total_written) { Ok(size) => size, @@ -2348,7 +2371,19 @@ impl Task { let new_fd = match target { DupFdRequest::Exact(target) => { - let _ = task.do_close_and_replace(target, Some(fd)); + // `do_close_and_replace` returns `Err(EBADF)` whenever + // `target` had nothing open to close — the common case + // for dup2 to a fresh fd — even though it still performs + // the replacement insert; that outcome is intentionally + // ignored, matching `close()`'s own semantics. It can + // separately fail with `Err(ENOMEM)` if `target` is far + // beyond the descriptor store's current growth allowance + // (see `fd_into_specific_raw_integer`), in which case the + // replacement was NOT inserted and the dup must fail. + match task.do_close_and_replace(target, Some(fd)) { + Ok(()) | Err(Errno::EBADF) => {} + Err(_) => return Err(DupFdError::TooManyFiles), + } target } DupFdRequest::LowestAvailable => { @@ -2364,8 +2399,12 @@ impl Task { } raw_fd += 1; } - let success = rds.fd_into_specific_raw_integer(fd, raw_fd); - assert!(success); + // As in the `Exact` case: `raw_fd` (bounded only by + // `RLIMIT_NOFILE` via the `max_fd` check above) can still + // exceed the descriptor store's growth allowance. + if !rds.fd_into_specific_raw_integer(fd, raw_fd) { + return Err(DupFdError::TooManyFiles); + } raw_fd } }; diff --git a/litebox_shim_linux/src/syscalls/mm.rs b/litebox_shim_linux/src/syscalls/mm.rs index c5f67791e..3fcbcc59f 100644 --- a/litebox_shim_linux/src/syscalls/mm.rs +++ b/litebox_shim_linux/src/syscalls/mm.rs @@ -59,7 +59,10 @@ pub(crate) struct ElfPatchState { } /// Per-process collection of ELF patching state, keyed by fd number. -pub(crate) type ElfPatchCache = BTreeMap; +/// +/// `None` is a negative-cache entry: the fd was probed and is not a patchable +/// ELF, so subsequent mmaps skip re-reading and re-parsing its headers. +pub(crate) type ElfPatchCache = BTreeMap>; #[inline] fn align_up(addr: usize, align: usize) -> usize { @@ -73,6 +76,43 @@ fn align_down(addr: usize, align: usize) -> usize { addr & !(align - 1) } +/// Write back only the byte ranges of `patched` that differ from `original`, +/// coalescing spans separated by small gaps. Patch sites are sparse, so this +/// keeps pages without any site clean instead of dirtying the whole segment +/// (CoW faults are especially expensive under shadow-paging hypervisors). +/// +/// Returns `None` if a write faults (e.g. the mapping was removed). +fn write_back_patched_spans( + mapped_addr: crate::MutPtr, + original: &[u8], + patched: &[u8], +) -> Option<()> { + // Merging nearby sites bounds the number of fault-handled copies without + // dirtying extra pages (each patch site is only a few bytes). + const MERGE_GAP: usize = 64; + debug_assert_eq!(original.len(), patched.len()); + let len = original.len().min(patched.len()); + let mut i = 0; + while i < len { + if original[i] == patched[i] { + i += 1; + continue; + } + let start = i; + let mut end = i + 1; // exclusive end of the last differing byte + let mut j = i + 1; + while j < len && j - end <= MERGE_GAP { + if original[j] != patched[j] { + end = j + 1; + } + j += 1; + } + mapped_addr.copy_from_slice(start, &patched[start..end])?; + i = j; + } + Some(()) +} + impl Task { #[inline] fn do_mmap( @@ -103,8 +143,15 @@ impl Task { prot: ProtFlags, flags: MapFlags, ) -> Result, MappingError> { - let op = |_| Ok(0); - self.do_mmap(suggested_addr, len, prot, flags, false, op) + // No initialization needed, so map directly with the final permissions. + litebox_common_linux::mm::do_mmap_no_init( + &self.global.pm, + suggested_addr, + len, + prot, + flags, + false, + ) } fn do_mmap_file( @@ -146,7 +193,7 @@ impl Task { // Track non-exec file mappings so we can patch them if they later // gain PROT_EXEC via mprotect. let mut cache = self.global.elf_patch_cache.lock(); - if let Some(state) = cache.get_mut(&fd) { + if let Some(Some(state)) = cache.get_mut(&fd) { let mapping_key = (result.as_usize(), len); // Overlapping entries are safe here: file_mappings is only used // to know which (addr, len) ranges belong to this fd so we can @@ -278,30 +325,53 @@ impl Task { ) -> Result, MappingError> { let op = |ptr: MutPtr| -> Result { // Note a malicious user may unmap ptr while we are reading. - // `sys_read` does not handle page faults, so we need to use a + // Filesystem reads do not handle page faults, so we need to use a // temporary buffer to read the data from fs (without worrying page // faults) and write it to the user buffer with page fault handling. - let mut file_offset = offset; - let mut buffer = [0; PAGE_SIZE]; - let mut copied = 0; - while copied < len { - let size = - self.sys_read(fd, &mut buffer, Some(file_offset)) - .map_err(|e| match e { - Errno::EBADF => MappingError::BadFD(fd), - Errno::EISDIR => MappingError::NotAFile, - Errno::EACCES => MappingError::NotForReading, - _ => unimplemented!(), - })?; - if size == 0 { - break; - } - // ptr is a valid pointer returned by do_mmap. - ptr.copy_from_slice(copied, &buffer[..size]).unwrap(); - copied += size; - file_offset += size; - } - Ok(copied) + // + // This runs for every program/library load, so use a large heap + // chunk and resolve the fd to its file object once, instead of a + // page-sized buffer with a per-page fd-table lookup. + const CHUNK_SIZE: usize = 128 * 1024; + let Ok(raw_fd) = u32::try_from(fd).and_then(usize::try_from) else { + return Err(MappingError::BadFD(fd)); + }; + let files = self.files.borrow(); + files + .run_on_raw_fd( + raw_fd, + |typed_fd| { + let mut buffer = alloc::vec![0u8; CHUNK_SIZE.min(len)]; + let mut file_offset = offset; + let mut copied = 0; + while copied < len { + let want = buffer.len().min(len - copied); + let size = files + .fs + .read(typed_fd, &mut buffer[..want], Some(file_offset)) + .map_err(|e| match Errno::from(e) { + Errno::EBADF => MappingError::BadFD(fd), + Errno::EISDIR => MappingError::NotAFile, + Errno::EACCES => MappingError::NotForReading, + _ => unimplemented!(), + })?; + if size == 0 { + break; + } + // ptr is a valid pointer returned by do_mmap. + ptr.copy_from_slice(copied, &buffer[..size]).unwrap(); + copied += size; + file_offset += size; + } + Ok(copied) + }, + |_| Err(MappingError::NotAFile), + |_| Err(MappingError::NotAFile), + |_| Err(MappingError::NotAFile), + |_| Err(MappingError::NotAFile), + |_| Err(MappingError::NotAFile), + ) + .map_err(|_| MappingError::BadFD(fd))? }; let fixed_addr = flags.intersects(MapFlags::MAP_FIXED | MapFlags::MAP_FIXED_NOREPLACE); self.do_mmap( @@ -398,7 +468,7 @@ impl Task { fn clear_file_mappings_for_range(&self, unmap_start: usize, unmap_len: usize) { let unmap_end = unmap_start.saturating_add(unmap_len); let mut cache = self.global.elf_patch_cache.lock(); - for state in cache.values_mut() { + for state in cache.values_mut().flatten() { state.file_mappings.retain(|&(vaddr, seg_len)| { let seg_end = vaddr.saturating_add(seg_len); seg_end <= unmap_start || vaddr >= unmap_end @@ -421,8 +491,28 @@ impl Task { // Intercept transitions to PROT_EXEC: patch unpatched file mappings. if prot.contains(ProtFlags::PROT_EXEC) { let syscall_entry = self.global.platform.get_syscall_entry_point(); - if syscall_entry != 0 { - self.maybe_patch_on_mprotect_exec(addr, len, syscall_entry); + let file_backed = + syscall_entry != 0 && self.maybe_patch_on_mprotect_exec(addr, len, syscall_entry); + if !file_backed { + // No tracked file mapping accounts for this range, so it is + // anonymous (or otherwise untracked) memory becoming + // executable. Unlike file-backed segments, anonymous memory + // is guest-writable RAM the guest can fill with arbitrary + // bytes before this call — without this scan, `mmap(RW, + // ANON)` + write a raw `syscall` opcode + `mprotect(RX)` + // would make that opcode directly executable, bypassing + // syscall interception entirely (a full sandbox escape). + // Neutralize any embedded `syscall` instructions the same + // way an unpatchable file segment is handled. + // + // This does not defend memory that is simultaneously + // writable and executable (`PROT_READ_WRITE_EXEC`): bytes + // written after this scan are never re-scanned, since no + // further permission transition occurs to catch them. See + // `litebox::mm::PageManager::make_pages_rwx`'s safety + // documentation — that combination is inherently unsafe and + // is only retained for legitimate JIT use cases. + self.apply_trap_fallback(addr, len, false); } } self.sys_mprotect_raw(addr, len, prot) @@ -481,12 +571,17 @@ impl Task { /// Check all tracked file mappings for unpatched regions that overlap the /// mprotect range. If found, run the runtime rewriter before the region /// becomes executable. + /// + /// Returns `true` if at least one tracked file mapping overlapped the + /// range (regardless of whether patching found any syscalls to rewrite), + /// so the caller knows whether any part of the range is *not* accounted + /// for by this fd-tracked path and must be handled separately. fn maybe_patch_on_mprotect_exec( &self, addr: crate::MutPtr, len: usize, syscall_entry: usize, - ) { + ) -> bool { let mprotect_start = addr.as_usize(); let mprotect_end = mprotect_start.saturating_add(len); @@ -497,6 +592,9 @@ impl Task { let cache = self.global.elf_patch_cache.lock(); let mut result = alloc::vec::Vec::new(); for (&fd, state) in cache.iter() { + let Some(state) = state else { + continue; // Negative-cached — not an ELF + }; if state.pre_patched { continue; } @@ -523,6 +621,7 @@ impl Task { } } + let found_any = !to_patch.is_empty(); for (fd, seg_start, seg_len) in to_patch { // Clamp to the intersection of the tracked mapping and the // mprotect range — only patch the portion becoming executable. @@ -538,6 +637,7 @@ impl Task { let mapped_addr = MutPtr::::from_usize(patch_start); self.maybe_patch_exec_segment(mapped_addr, patch_len, fd, syscall_entry, None); } + found_any } /// Initialize ELF patch state for an fd on its first mmap. @@ -559,21 +659,41 @@ impl Task { return; } + // Probe outside the lock; a failed probe is cached as `None` so + // non-ELF fds are not re-probed on every subsequent mmap. + let state = self.probe_elf_patch_state(fd, mapped_addr, file_offset); + + // Insert under lock (re-check for races). + self.global + .elf_patch_cache + .lock() + .entry(fd) + .or_insert(state); + } + + /// Probe an fd for ELF patch state. Returns `None` if the fd is not a + /// readable ELF with PT_LOAD segments. See [`Self::init_elf_patch_state`]. + fn probe_elf_patch_state( + &self, + fd: i32, + mapped_addr: usize, + file_offset: usize, + ) -> Option { // Read the ELF header (64 bytes for Elf64). let mut ehdr_buf = [0u8; core::mem::size_of::>()]; match self.sys_read(fd, &mut ehdr_buf, Some(0)) { Ok(n) if n == ehdr_buf.len() => {} - _ => return, // Not readable or short read, skip + _ => return None, // Not readable or short read, skip } // Parse as typed ELF64 header. let Ok((ehdr, _)) = object::from_bytes::>(&ehdr_buf) else { - return; + return None; }; // Verify ELF magic if &ehdr.e_ident.magic != b"\x7fELF" { - return; + return None; } let e_type = ehdr.e_type.get(ENDIAN); @@ -583,20 +703,18 @@ impl Task { // Validate e_phentsize: must be at least sizeof(Elf64_Phdr). if e_phentsize < core::mem::size_of::>() { - return; + return None; } // Read program headers. - let Some(phdrs_size) = e_phentsize.checked_mul(e_phnum) else { - return; - }; + let phdrs_size = e_phentsize.checked_mul(e_phnum)?; if phdrs_size == 0 || phdrs_size > 0x10000 { - return; // Sanity check + return None; // Sanity check } let mut phdrs_buf = alloc::vec![0u8; phdrs_size]; match self.sys_read(fd, &mut phdrs_buf, Some(e_phoff)) { Ok(n) if n == phdrs_buf.len() => {} - _ => return, + _ => return None, } // Find highest PT_LOAD end (p_vaddr + p_memsz) and compute base_addr @@ -633,7 +751,7 @@ impl Task { } if max_load_end == 0 { - return; // No PT_LOAD segments + return None; // No PT_LOAD segments } // Check if file is pre-patched by reading the last 32 bytes for magic @@ -668,9 +786,7 @@ impl Task { base + max_end.next_multiple_of(PAGE_SIZE) }; - // Insert under lock (re-check for races). - let mut cache = self.global.elf_patch_cache.lock(); - cache.entry(fd).or_insert(ElfPatchState { + Some(ElfPatchState { pre_patched, trampoline_file_offset: tramp_file_offset, trampoline_file_size: tramp_file_size.trunc(), @@ -681,7 +797,7 @@ impl Task { runtime_patches_committed: false, file_mappings: BTreeSet::new(), patched_ranges: BTreeSet::new(), - }); + }) } /// Check if a file has the LITEBOX trampoline magic at its tail. @@ -786,7 +902,7 @@ impl Task { // patching operation. In practice this is fine because the dynamic // linker loads shared libraries sequentially. let mut cache = self.global.elf_patch_cache.lock(); - let Some(state) = cache.get_mut(&fd) else { + let Some(Some(state)) = cache.get_mut(&fd) else { return true; // No patch state — not an ELF we're tracking }; @@ -977,6 +1093,8 @@ impl Task { panic!("fatal: failed to read code segment for patching"); }; let mut code_buf = code_owned.into_vec(); + // Snapshot of the unpatched code, used to compute which spans the + // rewriter modified so only those are written back. let original_code = code_buf.clone(); let code_vaddr = addr_usize as u64; @@ -1046,9 +1164,10 @@ impl Task { panic!("fatal: failed to write trampoline stubs"); } - // Write patched code back to the mapped region. - if mapped_addr.copy_from_slice(0, &code_buf).is_none() { - let _ = mapped_addr.copy_from_slice(0, &original_code); + // Write the patched spans back to the mapped region. Writing + // only the modified bytes keeps untouched pages clean (no CoW + // faults for pages without syscall sites). + if write_back_patched_spans(mapped_addr, &original_code, &code_buf).is_none() { let _ = self.sys_mprotect_raw( mapped_addr, len, @@ -1063,12 +1182,11 @@ impl Task { Ok(_) => { // No trampoline stubs were generated, but the rewriter may // have replaced unpatchable syscalls with trap instructions. - // Write back the modified code if it changed. - if code_buf != original_code && mapped_addr.copy_from_slice(0, &code_buf).is_none() - { - let _ = mapped_addr.copy_from_slice(0, &original_code); - panic!("fatal: failed to write trap bytes back to code segment"); - } + // Write back only the modified spans, if any. + assert!( + write_back_patched_spans(mapped_addr, &original_code, &code_buf).is_some(), + "fatal: failed to write trap bytes back to code segment" + ); // Fall through to restore RX protections below. } Err(e) => { @@ -1094,7 +1212,7 @@ impl Task { /// Removes the cache entry (preventing stale state if the fd is reused) /// and unmaps any trampoline that was allocated but never used. pub(crate) fn finalize_elf_patch(&self, fd: i32) { - let state = self.global.elf_patch_cache.lock().remove(&fd); + let state = self.global.elf_patch_cache.lock().remove(&fd).flatten(); if let Some(state) = state && state.trampoline_mapped && !state.pre_patched diff --git a/litebox_shim_linux/src/syscalls/net.rs b/litebox_shim_linux/src/syscalls/net.rs index 997ccb632..8797fb633 100644 --- a/litebox_shim_linux/src/syscalls/net.rs +++ b/litebox_shim_linux/src/syscalls/net.rs @@ -1150,15 +1150,21 @@ pub(crate) fn write_sockaddr_to_user( } UnixSocketAddr::Path(path) => { let offset = offset_of!(CSockUnixAddr, path); - let max_len = addrlen_val as usize - offset; - let name = &path.as_bytes()[..path.len().min(max_len)]; - addr.write_slice_at_offset(isize::try_from(offset).unwrap(), name) - .ok_or(Errno::EFAULT)?; - let null_offset = offset + name.len(); - // write null terminator if there is space - if addrlen_val as usize > null_offset { - addr.write_at_offset(isize::try_from(null_offset).unwrap(), 0) + // `addrlen_val` is guest-supplied (the caller's declared + // buffer size); a value at or below `offset` means there + // is no room for any path bytes at all, matching the + // `Abstract` arm's guard above. + if addrlen_val as usize > offset { + let max_len = addrlen_val as usize - offset; + let name = &path.as_bytes()[..path.len().min(max_len)]; + addr.write_slice_at_offset(isize::try_from(offset).unwrap(), name) .ok_or(Errno::EFAULT)?; + let null_offset = offset + name.len(); + // write null terminator if there is space + if addrlen_val as usize > null_offset { + addr.write_at_offset(isize::try_from(null_offset).unwrap(), 0) + .ok_or(Errno::EFAULT)?; + } } offset + path.len() + 1 } @@ -1369,8 +1375,18 @@ impl Task { let sockaddr = addr .map(|addr| read_sockaddr_from_user(addr, addrlen as usize)) .transpose()?; - let buf = buf.to_owned_slice(len).ok_or(Errno::EFAULT)?; - self.do_sendto(fd, &buf, flags, sockaddr) + if len <= crate::MAX_KERNEL_BUF_SIZE { + self.with_scratch_buf(len, |kernel_buf| { + buf.copy_to_slice(0, kernel_buf).ok_or(Errno::EFAULT)?; + self.do_sendto(fd, kernel_buf, flags, sockaddr) + }) + } else { + // Larger sends cannot be blindly chunked, as datagram sockets must + // preserve message boundaries; they are rare enough that a single + // owned copy is acceptable. + let buf = buf.to_owned_slice(len).ok_or(Errno::EFAULT)?; + self.do_sendto(fd, &buf, flags, sockaddr) + } } fn do_sendto( &self, diff --git a/litebox_shim_linux/src/syscalls/process.rs b/litebox_shim_linux/src/syscalls/process.rs index bdcfa19b8..549a26411 100644 --- a/litebox_shim_linux/src/syscalls/process.rs +++ b/litebox_shim_linux/src/syscalls/process.rs @@ -735,6 +735,7 @@ impl Task { fs: fs.into(), files: self.files.clone(), // TODO: !CLONE_FILES support signals: self.signals.clone_for_new_task(), + scratch_buf: core::cell::RefCell::new(alloc::vec::Vec::new()), }, }), ) @@ -1277,7 +1278,7 @@ impl Task { } } -/// Number of CPUs +/// Number of CPUs to report when the platform cannot tell us. const NR_CPUS: usize = 2; pub(crate) struct CpuSet { @@ -1296,9 +1297,15 @@ impl CpuSet { impl Task { /// Handle syscall `sched_getaffinity`. /// - /// Note this is a dummy implementation that always returns the same CPU set + /// Note this is a dummy implementation that reports every online CPU as + /// available. pub(crate) fn sys_sched_getaffinity(&self, _pid: Option) -> CpuSet { - let mut cpuset = bitvec::bitvec![u8, bitvec::order::Lsb0; 0; NR_CPUS]; + let ncpus = self + .global + .platform + .num_cpus() + .map_or(NR_CPUS, core::num::NonZeroUsize::get); + let mut cpuset = bitvec::bitvec![u8, bitvec::order::Lsb0; 0; ncpus]; cpuset.iter_mut().for_each(|mut b| *b = true); CpuSet { bits: cpuset } } @@ -1371,7 +1378,45 @@ impl Task { )?; 0 } - _ => unimplemented!("Unsupported futex operation"), + FutexArgs::Requeue { + addr, + flags, + wake_count, + addr2, + requeue_count, + } => { + warn_shared_futex!(flags); + let (woken, _requeued) = self.global.futex_manager.requeue( + addr, + addr2, + wake_count, + requeue_count, + None, + )?; + woken as usize + } + FutexArgs::CmpRequeue { + addr, + flags, + wake_count, + addr2, + requeue_count, + expected, + } => { + warn_shared_futex!(flags); + let (woken, requeued) = self.global.futex_manager.requeue( + addr, + addr2, + wake_count, + requeue_count, + Some(expected), + )?; + woken as usize + requeued as usize + } + other => { + log_unsupported!("futex operation {other:?}"); + return Err(Errno::ENOSYS); + } }; Ok(res) } @@ -1693,17 +1738,24 @@ mod tests { #[test] fn test_sched_getaffinity() { + use litebox::platform::ThreadProvider as _; + let task = crate::syscalls::tests::init_platform(None); + let expected = task + .global + .platform + .num_cpus() + .map_or(super::NR_CPUS, core::num::NonZeroUsize::get); let cpuset = task.sys_sched_getaffinity(None); - assert_eq!(cpuset.bits.len(), super::NR_CPUS); + assert_eq!(cpuset.bits.len(), expected); cpuset.bits.iter().for_each(|b| assert!(*b)); let ones: usize = cpuset .as_bytes() .iter() .map(|b| b.count_ones() as usize) .sum(); - assert_eq!(ones, super::NR_CPUS); + assert_eq!(ones, expected); } #[test] diff --git a/litebox_shim_linux/src/transport.rs b/litebox_shim_linux/src/transport.rs index ce450f53e..9473a4f09 100644 --- a/litebox_shim_linux/src/transport.rs +++ b/litebox_shim_linux/src/transport.rs @@ -6,6 +6,9 @@ use alloc::boxed::Box; use alloc::sync::Arc; +use litebox::event::polling::TryOpError; +use litebox::event::wait::WaitState; +use litebox::event::{Events, IOPollable as _}; use litebox::fs::nine_p::transport; use litebox::net::socket_channel::{ChannelWriteError, NetworkProxy}; use litebox::net::{ReceiveFlags, SendFlags}; @@ -14,6 +17,55 @@ use litebox_common_linux::{SockFlags, SockType, errno::Errno}; use crate::syscalls::net::SocketFd; use crate::{GlobalState, Platform, ShimFS}; +/// Non-blocking retries before parking the calling thread. +const SPIN_LIMIT: usize = 128; + +/// Timeout safety net while parked, in case a pollee notification is missed. +const POLL_INTERVAL: core::time::Duration = core::time::Duration::from_millis(1); + +/// Runs `try_op` until it succeeds: first with a short bounded spin, then by +/// parking the thread until `events` fire on the proxy's pollee. +/// +/// A pure `spin_loop()` wait is toxic when the network worker's vCPU is +/// preempted (e.g., in a ring-3 PVM guest, spinning burns the whole scheduling +/// quantum), so after [`SPIN_LIMIT`] retries this blocks via the platform's +/// wait facilities, with a [`POLL_INTERVAL`] timeout as a safety net. +fn spin_then_wait( + wait_state: &WaitState, + proxy: &NetworkProxy, + events: Events, + mut try_op: impl FnMut() -> Result>, +) -> Result { + for _ in 0..SPIN_LIMIT { + match try_op() { + Ok(r) => return Ok(r), + Err(TryOpError::TryAgain) => core::hint::spin_loop(), + Err(TryOpError::Other(e)) => return Err(e), + Err(TryOpError::WaitError(_)) => unreachable!("try_op does not wait"), + } + } + loop { + match wait_state + .context() + .with_timeout(POLL_INTERVAL) + .wait_on_events( + false, + events, + |observer, filter| { + proxy.register_observer(observer, filter); + Ok(()) + }, + &mut try_op, + ) { + Ok(r) => return Ok(r), + // Timed out: retry in case a notification was missed. + Err(TryOpError::WaitError(_)) => {} + Err(TryOpError::Other(e)) => return Err(e), + Err(TryOpError::TryAgain) => unreachable!("blocking wait cannot return TryAgain"), + } + } +} + /// Handles socket cleanup on drop without exposing the `FS` generic. /// /// This is stored as `Box` inside [`ShimTransport`] so that the @@ -38,19 +90,22 @@ impl DropGuard for SocketDropGuard { } } -/// A spin-polling TCP transport backed by a raw `SocketFd` and its [`NetworkProxy`]. +/// A TCP transport backed by a raw `SocketFd` and its [`NetworkProxy`]. /// /// The socket lives in the litebox descriptor table (for metadata / proxy) but is /// **not** registered in the guest's file-descriptor table, keeping it invisible /// to the guest program. /// /// All I/O goes through the non-blocking [`NetworkProxy`] methods directly -/// (`try_read` / `try_write`), with spin-polling when data is not yet available. -/// This avoids the need for a `WaitState` or any association with a particular -/// guest `Task`. +/// (`try_read` / `try_write`), briefly spinning and then parking the calling +/// thread (via a transport-private [`WaitState`], with no association with a +/// particular guest `Task`) when data is not yet available. pub struct ShimTransport { drop_guard: Box, proxy: Arc>, + /// Wait state used to park whichever thread is driving the transport. + /// I/O calls take `&mut self`, so it is never used concurrently. + wait_state: WaitState, } impl ShimTransport { @@ -61,7 +116,8 @@ impl ShimTransport { /// set up, but the socket is **not** assigned a guest fd number. /// /// Connection and all subsequent I/O use the [`NetworkProxy`] directly, - /// spin-polling when the operation cannot complete immediately. + /// spinning briefly and then parking when the operation cannot complete + /// immediately. pub(crate) fn connect( global: Arc>, addr: core::net::SocketAddr, @@ -77,21 +133,29 @@ impl ShimTransport { let proxy = global.initialize_socket(&sockfd, SockType::Stream, SockFlags::empty()); // 3. Initiate the TCP connection. + let wait_state = WaitState::new(global.platform); let mut check_progress = false; - loop { - match global.net.lock().connect(&sockfd, &addr, check_progress) { - Ok(()) => break, + spin_then_wait( + &wait_state, + &proxy, + Events::IN | Events::OUT, + || match global.net.lock().connect(&sockfd, &addr, check_progress) { + Ok(()) => Ok(()), Err(litebox::net::errors::ConnectError::InProgress) => { - core::hint::spin_loop(); check_progress = true; + Err(TryOpError::TryAgain) } - Err(e) => return Err(Errno::from(e)), - } - } + Err(e) => Err(TryOpError::Other(Errno::from(e))), + }, + )?; let drop_guard = Box::new(SocketDropGuard { global, sockfd }); - Ok(Self { drop_guard, proxy }) + Ok(Self { + drop_guard, + proxy, + wait_state, + }) } } @@ -103,31 +167,27 @@ impl Drop for ShimTransport { impl transport::Read for ShimTransport { fn read(&mut self, buf: &mut [u8]) -> Result { - loop { + spin_then_wait(&self.wait_state, &self.proxy, Events::IN, || { match self.proxy.try_read(buf, ReceiveFlags::empty(), None) { - Ok(0) => { - // No data yet — spin until something arrives. - core::hint::spin_loop(); - } - Ok(n) => return Ok(n), - Err(_) => return Err(transport::ReadError), + // No data yet — wait until something arrives. + Ok(0) => Err(TryOpError::TryAgain), + Ok(n) => Ok(n), + Err(_) => Err(TryOpError::Other(transport::ReadError)), } - } + }) } } impl transport::Write for ShimTransport { fn write(&mut self, buf: &[u8]) -> Result { - loop { + spin_then_wait(&self.wait_state, &self.proxy, Events::OUT, || { match self.proxy.try_write(buf, SendFlags::empty(), None) { - Ok(n) => return Ok(n), - Err(ChannelWriteError::BufferFull) => { - // TX ring full — spin until space opens up. - core::hint::spin_loop(); - } - Err(_) => return Err(transport::WriteError), + Ok(n) => Ok(n), + // TX ring full — wait until space opens up. + Err(ChannelWriteError::BufferFull) => Err(TryOpError::TryAgain), + Err(_) => Err(TryOpError::Other(transport::WriteError)), } - } + }) } } diff --git a/litebox_shim_optee/src/syscalls/mm.rs b/litebox_shim_optee/src/syscalls/mm.rs index 116421b8f..7e7d3d2a5 100644 --- a/litebox_shim_optee/src/syscalls/mm.rs +++ b/litebox_shim_optee/src/syscalls/mm.rs @@ -23,15 +23,14 @@ impl Task { prot: ProtFlags, flags: MapFlags, ) -> Result, MappingError> { - let op = |_| Ok(0); - litebox_common_linux::mm::do_mmap( + // No initialization needed, so map directly with the final permissions. + litebox_common_linux::mm::do_mmap_no_init( &self.global.pm, suggested_addr, len, prot, flags, false, - op, ) } diff --git a/litebox_syscall_rewriter/src/lib.rs b/litebox_syscall_rewriter/src/lib.rs index 9dcc10c28..43e9b5547 100644 --- a/litebox_syscall_rewriter/src/lib.rs +++ b/litebox_syscall_rewriter/src/lib.rs @@ -17,7 +17,6 @@ #![cfg_attr(not(feature = "std"), no_std)] extern crate alloc; -use alloc::collections::BTreeSet; use alloc::format; use alloc::string::{String, ToString}; use alloc::vec; @@ -153,7 +152,7 @@ pub fn hook_syscalls_in_elf(input_binary: &[u8], trampoline: Option) -> Res fixup_phdr_alignment(buf); // Parse the ELF and extract all metadata we need, then drop the borrow so we can mutate buf. - let (arch, text_sections, control_transfer_targets, trampoline_base_addr) = { + let (arch, text_sections, section_instructions, control_transfer_targets, trampoline_base_addr) = { let file = object::File::parse(&*buf).map_err(|e| Error::ParseError(e.to_string()))?; let arch = match file { @@ -172,13 +171,24 @@ pub fn hook_syscalls_in_elf(input_binary: &[u8], trampoline: Option) -> Res return Ok(input_binary.to_vec()); } - let control_transfer_targets = get_control_transfer_targets(arch, &*buf, &text_sections)?; + // Decode each section exactly once: the same pass yields both the + // control-transfer targets and the instructions used for patching. + let mut section_instructions = Vec::with_capacity(text_sections.len()); + let mut control_transfer_targets = Vec::new(); + for s in &text_sections { + let section_data = section_slice(&*buf, s)?; + let instructions = decode_section_instructions(arch, section_data, s.vaddr)?; + collect_control_transfer_targets(&instructions, &mut control_transfer_targets); + section_instructions.push(instructions); + } + sort_control_transfer_targets(&mut control_transfer_targets); let trampoline_base_addr = find_addr_for_trampoline_code(&file)?; ( arch, text_sections, + section_instructions, control_transfer_targets, trampoline_base_addr, ) @@ -192,11 +202,12 @@ pub fn hook_syscalls_in_elf(input_binary: &[u8], trampoline: Option) -> Res // Patch syscalls in-place in buf let mut skipped_addrs = Vec::new(); let mut syscall_insns_found = false; - for s in &text_sections { + for (s, instructions) in text_sections.iter().zip(§ion_instructions) { let section_data = section_slice_mut(buf, s)?; match hook_syscalls_in_section( arch, &control_transfer_targets, + instructions, s.vaddr, section_data, trampoline_base_addr, @@ -231,10 +242,14 @@ pub fn hook_syscalls_in_elf(input_binary: &[u8], trampoline: Option) -> Res return Ok(out); } - // Build output: [patched ELF][padding to page boundary][trampoline code][header] - let mut out = buf.to_vec(); - let remain = out.len() % 0x1000; - out.extend_from_slice(&vec![0; if remain == 0 { 0 } else { 0x1000 - remain }]); + // Build output: [patched ELF][padding to page boundary][trampoline code][header]. + // Reserve the exact final size up front so the patched ELF is copied once, + // with no reallocation when appending padding, trampoline code, or header. + let padded_len = buf.len().next_multiple_of(0x1000); + let mut out = + Vec::with_capacity(padded_len + trampoline_data.len() + size_of::()); + out.extend_from_slice(buf); + out.resize(padded_len, 0); // Calculate file offset where trampoline code starts let trampoline_file_offset = out.len() as u64; @@ -339,19 +354,23 @@ enum Arch { /// (private) Hook all syscalls in `section`, possibly extending `trampoline_data` to do so. /// -/// `trampoline_base_addr` is the virtual address corresponding to `trampoline_data[0]`. -/// `syscall_entry_addr` is the address of the 8-byte entry-point value that each trampoline -/// stub jumps to (via `JMP [RIP+disp32]` on x86-64). +/// `instructions` must be the decoded instructions of `section_data` (see +/// [`decode_section_instructions`]). `control_transfer_targets` must be sorted +/// and deduplicated. `trampoline_base_addr` is the virtual address +/// corresponding to `trampoline_data[0]`. `syscall_entry_addr` is the address +/// of the 8-byte entry-point value that each trampoline stub jumps to (via +/// `JMP [RIP+disp32]` on x86-64). +#[allow(clippy::too_many_arguments)] fn hook_syscalls_in_section( arch: Arch, - control_transfer_targets: &BTreeSet, + control_transfer_targets: &[u64], + instructions: &[iced_x86::Instruction], section_base_addr: u64, section_data: &mut [u8], trampoline_base_addr: u64, syscall_entry_addr: u64, trampoline_data: &mut Vec, ) -> core::result::Result, InternalError> { - let instructions = decode_section_instructions(arch, section_data, section_base_addr)?; let mut found_any = false; let mut skipped_addrs = Vec::new(); for (i, inst) in instructions.iter().enumerate() { @@ -373,7 +392,7 @@ fn hook_syscalls_in_section( // the replaced range backward (a jump landing on the syscall would hit // NOPs instead). Skip the backward scan and fall through to the // forward-only path (hook_syscall_and_after). - if !control_transfer_targets.contains(&inst.ip()) { + if !is_control_transfer_target(control_transfer_targets, inst.ip()) { for inst_id in (0..i).rev() { let prev_inst = &instructions[inst_id]; if prev_inst.flow_control() != iced_x86::FlowControl::Next { @@ -383,7 +402,7 @@ fn hook_syscalls_in_section( replace_start = Some(prev_inst.ip()); replace_start_idx = inst_id; break; - } else if control_transfer_targets.contains(&prev_inst.ip()) { + } else if is_control_transfer_target(control_transfer_targets, prev_inst.ip()) { // If the previous instruction is a control transfer target, we don't want to cross it break; } @@ -398,7 +417,7 @@ fn hook_syscalls_in_section( trampoline_base_addr, syscall_entry_addr, trampoline_data, - &instructions, + instructions, i, ) { Ok(()) => {} @@ -437,7 +456,7 @@ fn hook_syscalls_in_section( trampoline_base_addr, syscall_entry_addr, trampoline_data, - &instructions, + instructions, i, ) { Ok(()) => {} @@ -691,18 +710,15 @@ pub fn patch_code_segment( ) -> Result<(Vec, Vec)> { // Build control-transfer targets for this segment. let instructions = decode_section_instructions(Arch::X86_64, code, code_vaddr)?; - let mut control_transfer_targets = BTreeSet::new(); - for inst in &instructions { - let target = inst.near_branch_target(); - if target != 0 { - control_transfer_targets.insert(target); - } - } + let mut control_transfer_targets = Vec::new(); + collect_control_transfer_targets(&instructions, &mut control_transfer_targets); + sort_control_transfer_targets(&mut control_transfer_targets); let mut trampoline_data = Vec::new(); match hook_syscalls_in_section( Arch::X86_64, &control_transfer_targets, + &instructions, code_vaddr, code, trampoline_write_vaddr, @@ -763,22 +779,29 @@ where .max() } -fn get_control_transfer_targets( - arch: Arch, - input_binary: &[u8], - text_sections: &[TextSectionInfo], -) -> Result> { - let mut control_transfer_targets = BTreeSet::new(); - for s in text_sections { - let section_data = section_slice(input_binary, s)?; - let instructions = decode_section_instructions(arch, section_data, s.vaddr)?; - control_transfer_targets.extend(instructions.into_iter().filter_map(|inst| { - let target = inst.near_branch_target(); - (target != 0).then_some(target) - })); - } +/// Appends the near-branch targets of `instructions` to `targets`. The result +/// must be passed through [`sort_control_transfer_targets`] before it can be +/// queried with [`is_control_transfer_target`]. +fn collect_control_transfer_targets( + instructions: &[iced_x86::Instruction], + targets: &mut Vec, +) { + targets.extend(instructions.iter().filter_map(|inst| { + let target = inst.near_branch_target(); + (target != 0).then_some(target) + })); +} + +/// Sorts and deduplicates a control-transfer-target list so that membership +/// can be tested with a binary search. +fn sort_control_transfer_targets(targets: &mut Vec) { + targets.sort_unstable(); + targets.dedup(); +} - Ok(control_transfer_targets) +/// Membership test on a sorted, deduplicated control-transfer-target list. +fn is_control_transfer_target(targets: &[u64], ip: u64) -> bool { + targets.binary_search(&ip).is_ok() } const MAX_X86_INSTRUCTION_LEN: usize = 15; @@ -804,7 +827,9 @@ fn decode_section_instructions( Arch::X86_64 => 64, }; - let mut instructions = Vec::new(); + // Average x86-64 instruction length is a bit over 4 bytes; reserving + // len/4 slots avoids repeated reallocation of a very large Vec. + let mut instructions = Vec::with_capacity(section_data.len() / 4); let mut offset = 0usize; while offset < section_data.len() { @@ -819,7 +844,7 @@ fn decode_section_instructions( let chunk_start_ip = section_base_addr + offset as u64; let chunk_end_ip = chunk_start_ip + chunk_advance_len as u64; - append_decoded_instructions( + let resume_ip = append_decoded_instructions( bitness, &remaining[..decode_window_len], chunk_start_ip, @@ -827,19 +852,28 @@ fn decode_section_instructions( &mut instructions, )?; - offset = offset.checked_add(chunk_advance_len).unwrap(); + // An instruction may straddle chunk_end_ip (the overlap window lets it + // decode fully). Resuming exactly at chunk_end_ip would then decode + // phantom instructions from its tail bytes, so resume at the ip after + // the last fully decoded instruction instead. + let consumed = usize::try_from(resume_ip - chunk_start_ip).unwrap(); + assert!(consumed > 0); + offset = offset.checked_add(consumed).unwrap(); } Ok(instructions) } +/// Decodes `window` starting at `chunk_start_ip`, keeping instructions whose +/// ip is below `chunk_end_ip`. Returns the ip immediately after the last kept +/// instruction — the point where the next chunk must resume decoding. fn append_decoded_instructions( bitness: u32, window: &[u8], chunk_start_ip: u64, chunk_end_ip: u64, instructions: &mut Vec, -) -> Result<()> { +) -> Result { if bytes_until_next_4g_boundary(window.as_ptr()) > window.len() { return append_decoded_non_crossing_window( bitness, @@ -887,10 +921,11 @@ fn append_decoded_non_crossing_window( chunk_start_ip: u64, chunk_end_ip: u64, instructions: &mut Vec, -) -> Result<()> { +) -> Result { let mut decoder = iced_x86::Decoder::new(bitness, window, iced_x86::DecoderOptions::NONE); decoder.set_ip(chunk_start_ip); + let mut resume_ip = chunk_start_ip; for inst in &mut decoder { if inst.len() == 0 { return Err(Error::DisassemblyFailure(format!( @@ -903,10 +938,11 @@ fn append_decoded_non_crossing_window( break; } + resume_ip = inst.next_ip(); instructions.push(inst); } - Ok(()) + Ok(resume_ip) } /// Returns the section data slice from `buf` corresponding to `section`, or an error if out of bounds. @@ -962,7 +998,7 @@ fn reencode_instructions( #[allow(clippy::too_many_arguments)] fn hook_syscall_and_after( - control_transfer_targets: &BTreeSet, + control_transfer_targets: &[u64], section_base_addr: u64, section_data: &mut [u8], trampoline_base_addr: u64, @@ -978,7 +1014,7 @@ fn hook_syscall_and_after( let mut replace_end_idx = inst_index; for (idx, next_inst) in instructions.iter().enumerate().skip(inst_index + 1) { - if control_transfer_targets.contains(&next_inst.ip()) { + if is_control_transfer_target(control_transfer_targets, next_inst.ip()) { // If the next instruction is a control transfer target, we don't want to cross it break; } @@ -1086,6 +1122,30 @@ fn hook_syscall_and_after( mod tests { use super::*; + #[test] + fn chunked_decode_does_not_resync_mid_instruction() { + // A section longer than one decode chunk, filled with NOPs, with a + // 5-byte `MOV EAX, imm32` straddling the chunk boundary. The last two + // immediate bytes (0F 05) decode as a phantom SYSCALL if the next + // chunk resumes exactly at the boundary (mid-instruction) instead of + // at the ip after the last fully decoded instruction. + let len = TARGET_DECODE_CHUNK_LEN + 16; + let mut code = vec![0x90u8; len]; + let boundary = TARGET_DECODE_CHUNK_LEN; + code[boundary - 3..boundary + 2].copy_from_slice(&[0xB8, 0x90, 0x90, 0x0F, 0x05]); + + let base_addr = 0x1000u64; + let instructions = decode_section_instructions(Arch::X86_64, &code, base_addr).unwrap(); + + let mut expected_ip = base_addr; + for inst in &instructions { + assert_eq!(inst.ip(), expected_ip, "instructions must not overlap"); + assert_ne!(inst.code(), iced_x86::Code::Syscall, "phantom syscall"); + expected_ip = inst.next_ip(); + } + assert_eq!(expected_ip, base_addr + len as u64); + } + #[cfg(target_pointer_width = "64")] #[test] #[ignore = "allocates over 4GiB to reproduce the iced-x86 host-pointer bug without mmap"] diff --git a/litebox_syscall_rewriter/src/main.rs b/litebox_syscall_rewriter/src/main.rs index 7ef8eef14..c5d2fdecb 100644 --- a/litebox_syscall_rewriter/src/main.rs +++ b/litebox_syscall_rewriter/src/main.rs @@ -45,7 +45,8 @@ fn copy_file_permissions( fn main() -> anyhow::Result<()> { let cli_args = CliArgs::parse(); let mut input_binary = std::fs::File::open(&cli_args.input_binary)?; - let mut input_binary_bytes = vec![]; + let mut input_binary_bytes = + Vec::with_capacity(usize::try_from(input_binary.metadata()?.len()).unwrap_or(0)); input_binary.read_to_end(&mut input_binary_bytes)?; let output_binary = litebox_syscall_rewriter::hook_syscalls_in_elf( &input_binary_bytes,