diff --git a/Cargo.toml b/Cargo.toml index 10098cf..9526142 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,6 @@ authors = ["Hsy-Intel ", "Edmund Song () + self.bitmap_size_bytes`. -/// The bitmap begins at the first byte following this structure. -/// -/// ### Bitmap Semantics -/// - Each bit in the trailing bitmap represents a memory region of -/// `unit_size_bytes` bytes. -/// - Bit 0 corresponds to the physical address specified by `phys_base`. -/// - A **set bit (1)** indicates memory is unaccepted (pending); -/// a **cleared bit (0)** indicates it has been accepted. -/// -/// ### Concurrency Contract -/// This type does not provide internal synchronization for bitmap mutation. -/// Callers must serialize mutating operations (for example with a spinlock) -/// before invoking methods like `register_range` or `accept_range`. -#[derive(Debug)] -#[repr(C, packed)] -pub struct EfiUnacceptedMemory { - /// The version of the table. Currently, only version 1 is defined. - version: u32, - /// The size of the memory region represented by a single bit in the bitmap. - /// Typically set to 2MiB (0x200000) to align with huge page boundaries. - unit_size_bytes: u32, - /// The start physical address of the memory range covered by this bitmap. - /// Bit 0 of the bitmap corresponds to this address. - phys_base: u64, - /// The bitmap payload length in bytes, excluding this header. - bitmap_size_bytes: u64, -} - -impl EfiUnacceptedMemory { - /// Initializes the table header fields for EFI installation. - pub fn init_header( - &mut self, - unit_size_bytes: u32, - phys_base: u64, - bitmap_size_bytes: u64, - ) -> Result<(), AcceptError> { - if unit_size_bytes == 0 || !unit_size_bytes.is_power_of_two() { - return Err(AcceptError::InvalidAlignment); - } - - self.version = LINUX_EFI_UNACCEPTED_MEM_TABLE_VERSION; - self.unit_size_bytes = unit_size_bytes; - self.phys_base = phys_base; - self.bitmap_size_bytes = bitmap_size_bytes; - Ok(()) - } - - /// Returns the version of the table header. - pub const fn version(&self) -> u32 { - self.version - } - - /// Returns the unit size represented by one bitmap bit, in bytes. - pub const fn unit_size_bytes(&self) -> u32 { - self.unit_size_bytes - } - - /// Returns the start physical address covered by the bitmap. - pub const fn phys_base(&self) -> u64 { - self.phys_base - } - - /// Returns the trailing bitmap payload length, in bytes. - pub const fn bitmap_size_bytes(&self) -> u64 { - self.bitmap_size_bytes - } - - /// Returns whether `(start, size)` overlaps any pending bitmap unit. - /// - /// # Safety - /// - /// The caller must ensure this header is followed in memory by at least - /// `self.bitmap_size_bytes` readable bitmap bytes. - pub unsafe fn is_range_pending_by_size( - &self, - start: u64, - size: u64, - ) -> Result { - let Some(end) = start.checked_add(size) else { - return Err(AcceptError::ArithmeticOverflow); - }; - - // SAFETY: Caller guarantees the table header is followed by readable bitmap bytes. - unsafe { self.is_range_pending(start, end) } - } - - /// Returns whether `[start, end)` overlaps any pending bitmap unit. - /// - /// # Safety - /// - /// The caller must ensure this header is followed in memory by at least - /// `self.bitmap_size_bytes` readable bitmap bytes. - pub unsafe fn is_range_pending(&self, start: u64, end: u64) -> Result { - let Some((range_start, range_end, unit_size)) = - self.clamp_gpa_range_to_bitmap_coverage(start, end)? - else { - return Ok(false); - }; - - let (first_bit, last_bit) = self.addr_to_bit_range(range_start, range_end, unit_size)?; - // SAFETY: Caller guarantees the table header is followed by readable bitmap bytes. - let bitmap = unsafe { self.as_bitmap_slice() }; - BitmapRef::new(bitmap).has_set_bit(first_bit, last_bit) - } - - /// Returns whether every bitmap unit overlapping `[start, end)` is accepted. - /// - /// Ranges outside bitmap coverage are considered accepted by definition, - /// because this table only tracks deferred acceptance inside its own coverage. - /// - /// # Safety - /// - /// The caller must ensure this header is followed in memory by at least - /// `self.bitmap_size_bytes` readable bitmap bytes. - pub unsafe fn is_fully_accepted(&self, start: u64, end: u64) -> Result { - // SAFETY: Caller guarantees the table header is followed by readable bitmap bytes. - Ok(!unsafe { self.is_range_pending(start, end) }?) - } - - /// Convenience wrapper for - /// [`EfiUnacceptedMemory::accept_range`] using - /// `(start, size)` instead of `(start, end)`. - /// - /// Computes `end = start + size` and forwards to the range-based API. - /// - /// # Safety - /// - /// The caller must ensure `self` is uniquely borrowed for in-place bitmap updates and - /// points to a valid unaccepted-memory table with writable bitmap memory. - /// - /// # Errors - /// - /// Returns [`AcceptError::ArithmeticOverflow`] if `start + size` overflows. - /// Propagates any error from [`EfiUnacceptedMemory::accept_range`]. - pub unsafe fn accept_by_size(&mut self, start: u64, size: u64) -> Result<(), AcceptError> { - let Some(end) = start.checked_add(size) else { - return Err(AcceptError::ArithmeticOverflow); - }; - - // SAFETY: Caller guarantees table/bitmap validity and target range correctness. - unsafe { self.accept_range(start, end) } - } - - /// Accepts bitmap-marked units that overlap `start..end`, then clears accepted bits. - /// - /// The input is interpreted as a half-open GPA interval `[start, end)`. - /// - /// Behavior summary: - /// - The requested range is first clamped to bitmap coverage. - /// - Any bitmap bit set to `1` and overlapping the clamped range is accepted. - /// - Successfully accepted bits are cleared to `0` in-place. - /// - If the clamped range is empty, this is a no-op. - /// - /// # Safety - /// - /// The caller must ensure this table and bitmap describe pending private-memory units, - /// and the target GPA ranges are valid for TDX acceptance. - /// - /// # Errors - /// - /// Returns [`AcceptError::InvalidAlignment`] for invalid unit configuration. - /// Returns [`AcceptError::ArithmeticOverflow`] for address/index arithmetic overflow. - /// Returns [`AcceptError::OutOfBounds`] for bitmap index out-of-range accesses. - /// Returns hardware-originated failures from `accept_memory` via - /// [`AcceptError::TdCall`]. - pub unsafe fn accept_range(&mut self, start: u64, end: u64) -> Result<(), AcceptError> { - // SAFETY: Caller guarantees table/bitmap validity and target range correctness. - let _ = unsafe { self.accept_if_needed_range(start, end) }?; - Ok(()) - } - - /// Returns the end GPA (exclusive) covered by the bitmap. - /// - /// This is equivalent to `phys_base + total_coverage_size()`. - pub fn bitmap_coverage_end(&self) -> Option { - let base = PhysAddr::new(self.phys_base); - Some(base.checked_add(self.total_coverage_size()?).ok()?.raw()) - } - - /// Returns an immutable slice view of the trailing bitmap payload. - /// - /// # Safety - /// - /// The caller must ensure that this header is followed by at least - /// `self.bitmap_size_bytes` - /// readable bytes in memory. - pub unsafe fn as_bitmap_slice(&self) -> &[u8] { - debug_assert!(self.byte_len().is_ok()); - let bitmap_ptr = core::ptr::from_ref(self) - .cast::() - .wrapping_add(core::mem::size_of::()); - let bitmap_len = self - .byte_len() - .expect("bitmap size must fit usize on this platform"); - // SAFETY: `bitmap_ptr` points to the trailing bitmap bytes immediately - // after `self`; `bitmap_len` is validated from `self.bitmap_size_bytes`; - // caller guarantees - // readable backing memory for the returned slice. - unsafe { core::slice::from_raw_parts(bitmap_ptr, bitmap_len) } - } - - /// Returns a mutable slice view of the trailing bitmap payload. - /// - /// # Safety - /// - /// The caller must ensure that this header is followed by at least - /// `self.bitmap_size_bytes` - /// writable bytes in memory, and that no aliased mutable reference exists - /// while the returned slice is in use. - pub unsafe fn as_bitmap_slice_mut(&mut self) -> &mut [u8] { - debug_assert!(self.byte_len().is_ok()); - let bitmap_ptr_mut = core::ptr::from_mut(self) - .cast::() - .wrapping_add(core::mem::size_of::()); - debug_assert!(!bitmap_ptr_mut.is_null()); - let bitmap_len = self - .byte_len() - .expect("bitmap size must fit usize on this platform"); - // SAFETY: `bitmap_ptr_mut` points to the trailing bitmap bytes immediately - // after `self`; `bitmap_len` is validated from `self.bitmap_size_bytes`; - // caller guarantees - // writable backing memory and unique mutable access for the returned slice. - unsafe { core::slice::from_raw_parts_mut(bitmap_ptr_mut, bitmap_len) } - } - - /// Processes `start..end` by eagerly accepting required parts and deferring the rest in bitmap. - /// - /// This method applies a hybrid policy: - /// - edge fragments that are not `unit_size`-aligned are accepted immediately; - /// - aligned interior regions within bitmap coverage are marked as unaccepted bits; - /// - aligned regions outside bitmap coverage are accepted immediately. - /// - /// # Safety - /// - /// The caller must ensure the range is valid guest-private memory in pending/acceptable state. - /// - /// # Errors - /// - /// Returns [`AcceptError::InvalidAlignment`] for invalid unit configuration. - /// Returns [`AcceptError::ArithmeticOverflow`] for address arithmetic overflow. - /// Returns [`AcceptError::Overlap`] if the same bitmap unit is registered twice. - /// Returns hardware-originated errors from `accept_memory` - /// via [`AcceptError::TdCall`]. - pub unsafe fn register_range(&mut self, start: u64, end: u64) -> Result<(), AcceptError> { - let table_phys_base = self.phys_base; - let unit_size = self.validated_unit_size()?; - if start >= end { - return Ok(()); - } - - let unit_mask = unit_size - 1; - - if end - start < 2 * unit_size { - // SAFETY: Caller guarantees the physical range is valid for TDX acceptance. - return unsafe { Self::try_accept_range(start, end) }; - } - - let mut current_start = start; - let mut current_end = end; - - if current_start & unit_mask != 0 { - let Some(aligned_start) = align_up(current_start, unit_size) else { - return Err(AcceptError::ArithmeticOverflow); - }; - // SAFETY: Caller guarantees the physical subrange is valid for TDX acceptance. - unsafe { Self::try_accept_range(current_start, aligned_start)? }; - current_start = aligned_start; - } - - if current_end & unit_mask != 0 { - let aligned_end = align_down(current_end, unit_size); - // SAFETY: Caller guarantees the physical subrange is valid for TDX acceptance. - unsafe { Self::try_accept_range(aligned_end, current_end)? }; - current_end = aligned_end; - } - - let Some(bitmap_coverage) = self.total_coverage_size() else { - return Err(AcceptError::ArithmeticOverflow); - }; - - let Some(bitmap_end) = table_phys_base.checked_add(bitmap_coverage) else { - return Err(AcceptError::ArithmeticOverflow); - }; - - // 1) Process aligned range before bitmap coverage. - if current_start < table_phys_base { - let accept_end = current_end.min(table_phys_base); - // SAFETY: Caller guarantees the physical subrange is valid for TDX acceptance. - unsafe { Self::try_accept_range(current_start, accept_end)? }; - current_start = accept_end; - } - - if current_start >= current_end { - return Ok(()); - } - - // 2) Process aligned range within bitmap coverage. - if current_start < bitmap_end { - let bitmap_range_end = current_end.min(bitmap_end); - if current_start < bitmap_range_end { - // SAFETY: GPA range is unit-aligned and within bitmap coverage. - unsafe { - self.mark_range_as_unaccepted(current_start, bitmap_range_end, unit_size)? - }; - } - current_start = bitmap_range_end; - } - - // 3) Process aligned range after bitmap coverage. - if current_start < current_end { - // SAFETY: Caller guarantees the physical subrange is valid for TDX acceptance. - unsafe { Self::try_accept_range(current_start, current_end)? }; - } - - Ok(()) - } - - pub fn total_coverage_size(&self) -> Option { - let unit_size = u64::from(self.unit_size_bytes); - self.bitmap_size_bytes - .checked_mul(unit_size)? - .checked_mul(8) - } - - /// Conditionally accepts bitmap-marked units that overlap `start..end`. - /// - /// Returns [`AcceptOutcome::AlreadyAccepted`] when the range overlaps the - /// bitmap but there are no pending bits to process. - /// - /// # Safety - /// - /// Same requirements as [`Self::accept_range`]. - unsafe fn accept_if_needed_range( - &mut self, - start: u64, - end: u64, - ) -> Result { - if start >= end { - return Ok(AcceptOutcome::AlreadyAccepted); - } - - let (range_start, range_end, unit_size) = - match self.clamp_gpa_range_to_bitmap_coverage(start, end)? { - Some(vals) => vals, - None => return Ok(AcceptOutcome::OutOfCoverage), - }; - - let (first_bit, last_bit) = self.addr_to_bit_range(range_start, range_end, unit_size)?; - let phys_base = PhysAddr::new(self.phys_base); - - let bit_to_gpa = |bit: BitIndex| -> Result { - Ok(phys_base.checked_add_units(bit, unit_size)?.raw()) - }; - - // SAFETY: Caller guarantees table/bitmap validity and exclusive mutable access. - let mut bitmap = BitmapMut::new(unsafe { self.as_bitmap_slice_mut() }); - let mut accepted_units = 0u64; - - let mut scan = first_bit; - while let Some(run_start) = bitmap.find_next_set(scan, last_bit)? { - let run_end = bitmap - .find_next_zero(run_start, last_bit)? - .unwrap_or(last_bit); - - let run_gpa_start = bit_to_gpa(run_start)?; - let run_gpa_end = bit_to_gpa(run_end)?; - - // SAFETY: Caller guarantees bitmap/GPA mapping validity for pending private pages. - unsafe { accept_memory(run_gpa_start, run_gpa_end)? }; - accepted_units = accepted_units - .checked_add(run_end.raw() - run_start.raw()) - .ok_or(AcceptError::ArithmeticOverflow)?; - - bitmap.clear_range(run_start, run_end)?; - - scan = run_end; - } - - match accepted_units { - 0 => Ok(AcceptOutcome::AlreadyAccepted), - n => Ok(AcceptOutcome::AcceptedNow { accepted_units: n }), - } - } - - /// Marks `start..end` (bitmap-relative) as unaccepted bits. - /// - /// # Safety - /// - /// The caller must ensure `start..end` is already converted to bitmap-relative offsets - /// (i.e., based on `self.phys_base`) and does not violate bitmap ownership/aliasing rules. - unsafe fn set_unaccepted_bits(&mut self, start: u64, end: u64) -> Result<(), AcceptError> { - let unit_size = self.validated_unit_size()?; - - let abs_start = self - .phys_base - .checked_add(start) - .ok_or(AcceptError::ArithmeticOverflow)?; - let abs_end = self - .phys_base - .checked_add(end) - .ok_or(AcceptError::ArithmeticOverflow)?; - - // SAFETY: Caller guarantees bitmap-relative range correctness and exclusive access. - unsafe { self.mark_range_as_unaccepted(abs_start, abs_end, unit_size) } - } - - fn total_bits(&self) -> Result { - self.bitmap_size_bytes - .checked_mul(8) - .ok_or(AcceptError::ArithmeticOverflow) - } - - fn byte_len(&self) -> Result { - usize::try_from(self.bitmap_size_bytes).map_err(|_| AcceptError::OutOfBounds) - } - - fn validated_unit_size(&self) -> Result { - let unit_size = u64::from(self.unit_size_bytes); - if unit_size == 0 || !unit_size.is_power_of_two() { - return Err(AcceptError::InvalidAlignment); - } - Ok(unit_size) - } - - fn max_phys_addr_exclusive(&self, unit_size: u64) -> Result { - let total_bits = self.total_bits()?; - let coverage_len = total_bits - .checked_mul(unit_size) - .ok_or(AcceptError::ArithmeticOverflow)?; - self.phys_base - .checked_add(coverage_len) - .ok_or(AcceptError::ArithmeticOverflow) - } - - fn clamp_gpa_range_to_bitmap_coverage( - &self, - start: u64, - end: u64, - ) -> Result, AcceptError> { - if start >= end { - return Ok(None); - } - - let unit_size = self.validated_unit_size()?; - let coverage_end = self.max_phys_addr_exclusive(unit_size)?; - - let range_start = start.max(self.phys_base); - let range_end = end.min(coverage_end); - if range_start >= range_end { - return Ok(None); - } - - Ok(Some((range_start, range_end, unit_size))) - } - - fn addr_to_bit_range( - &self, - start: u64, - end: u64, - unit_size: u64, - ) -> Result<(BitIndex, BitIndex), AcceptError> { - debug_assert!(start >= self.phys_base); - debug_assert!(start < end); - debug_assert!(unit_size.is_power_of_two()); - - let rel_start = start - self.phys_base; - let rel_end = end - self.phys_base; - - let first_bit = rel_start / unit_size; - // NOTE: last_bit is exclusive and uses ceil-div for overlap semantics. - // Any unit intersecting [start, end) is considered. - let last_bit = rel_end - .checked_add(unit_size - 1) - .ok_or(AcceptError::ArithmeticOverflow)? - / unit_size; - - Ok((BitIndex::new(first_bit), BitIndex::new(last_bit))) - } - - /// Marks unit-aligned bitmap-covered GPA range `start..end` as unaccepted bits. - /// - /// # Safety - /// - /// The caller must ensure: - /// - `self` points to a valid unaccepted-memory table whose trailing bitmap memory is writable; - /// - mutable access to `self`/bitmap is unique for the duration of this call (no aliasing); - /// - `start..end` is unit-aligned for `unit_size` and corresponds to this table's - /// bitmap coverage semantics. - unsafe fn mark_range_as_unaccepted( - &mut self, - start: u64, - end: u64, - unit_size: u64, - ) -> Result<(), AcceptError> { - if start >= end { - return Ok(()); - } - - debug_assert_eq!(start % unit_size, 0); - debug_assert_eq!(end % unit_size, 0); - - let start_bit = BitIndex::new((start - self.phys_base) / unit_size); - let end_bit = BitIndex::new((end - self.phys_base) / unit_size); - let total_bits = BitIndex::new(self.total_bits()?); - - let clamped_start_bit = BitIndex::new(start_bit.raw().min(total_bits.raw())); - let clamped_end_bit = BitIndex::new(end_bit.raw().min(total_bits.raw())); - if clamped_start_bit >= clamped_end_bit { - return Ok(()); - } - - // SAFETY: Caller guarantees bitmap memory is writable and uniquely accessible. - let mut bitmap = BitmapMut::new(unsafe { self.as_bitmap_slice_mut() }); - for bit in clamped_start_bit.raw()..clamped_end_bit.raw() { - bitmap.set_bit(BitIndex::new(bit))?; - } - - Ok(()) - } - - /// Accepts physical memory in `start..end` if the range is non-empty. - /// - /// # Safety - /// - /// The caller must ensure `start..end` is a valid GPA range for TDX acceptance, - /// and that accepting this range does not race with other concurrent acceptance or - /// access operations on the same memory. - unsafe fn try_accept_range(start: u64, end: u64) -> Result<(), AcceptError> { - if start >= end { - return Ok(()); - } - // SAFETY: Caller guarantees the physical range is valid for TDX acceptance. - unsafe { accept_memory(start, end) } - } -} - -/// Mutable bitmap view used by registration/acceptance paths. -/// -/// This helper is intentionally non-atomic; callers must provide external -/// synchronization when multiple CPUs could touch the same bitmap. -struct BitmapMut<'a> { - bits: &'a mut [u8], -} - -impl<'a> BitmapMut<'a> { - fn new(bits: &'a mut [u8]) -> Self { - Self { bits } - } - - fn capacity(&self) -> Result { - let len = u64::try_from(self.bits.len()).map_err(|_| AcceptError::OutOfBounds)?; - len.checked_mul(8).ok_or(AcceptError::ArithmeticOverflow) - } - - fn get_pos_mask(&self, bit_index: BitIndex) -> Result<(usize, u8), AcceptError> { - if bit_index.raw() >= self.capacity()? { - return Err(AcceptError::OutOfBounds); - } - - let byte_index = - usize::try_from(bit_index.raw() >> 3).map_err(|_| AcceptError::OutOfBounds)?; - let mask = 1u8 << (bit_index.raw() & 7); - Ok((byte_index, mask)) - } - - fn is_set(&self, bit_index: BitIndex) -> Result { - let (byte_index, mask) = self.get_pos_mask(bit_index)?; - Ok((self.bits[byte_index] & mask) != 0) - } - - fn set_bit(&mut self, bit_index: BitIndex) -> Result<(), AcceptError> { - if self.is_set(bit_index)? { - return Err(AcceptError::Overlap); - } - let (byte_index, mask) = self.get_pos_mask(bit_index)?; - self.bits[byte_index] |= mask; - Ok(()) - } - - fn clear_bit(&mut self, bit_index: BitIndex) -> Result<(), AcceptError> { - let (byte_index, mask) = self.get_pos_mask(bit_index)?; - self.bits[byte_index] &= !mask; - Ok(()) - } - - fn clear_range(&mut self, start_bit: BitIndex, end_bit: BitIndex) -> Result<(), AcceptError> { - let bit_len = self.capacity()?; - if start_bit.raw() > end_bit.raw() || end_bit.raw() > bit_len { - return Err(AcceptError::OutOfBounds); - } - if start_bit == end_bit { - return Ok(()); - } - - let start_byte = - usize::try_from(start_bit.raw() >> 3).map_err(|_| AcceptError::OutOfBounds)?; - let end_exclusive_byte = - usize::try_from((end_bit.raw() + 7) >> 3).map_err(|_| AcceptError::OutOfBounds)?; - let start_off = u8::try_from(start_bit.raw() & 7).map_err(|_| AcceptError::OutOfBounds)?; - let end_off = u8::try_from(end_bit.raw() & 7).map_err(|_| AcceptError::OutOfBounds)?; - - if start_byte + 1 == end_exclusive_byte { - // Entire range is inside one byte. - let end_off_eff = if end_off == 0 { 8 } else { end_off }; - let clear_mask = bit_range_mask(start_off, end_off_eff); - self.bits[start_byte] &= !clear_mask; - return Ok(()); - } - - // Leading partial byte. - if start_off != 0 { - self.bits[start_byte] &= low_bits_mask(start_off); - } else { - self.bits[start_byte] = 0; - } - - // Middle full bytes. - let middle_start = start_byte + 1; - let middle_end = if end_off == 0 { - end_exclusive_byte - } else { - end_exclusive_byte - 1 - }; - if middle_start < middle_end { - self.bits[middle_start..middle_end].fill(0); - } - - // Trailing partial byte. - if end_off != 0 { - let keep_high = !low_bits_mask(end_off); - let last = end_exclusive_byte - 1; - self.bits[last] &= keep_high; - } - - Ok(()) - } - - fn find_next_set( - &self, - start_bit: BitIndex, - end_bit: BitIndex, - ) -> Result, AcceptError> { - self.find_next_matching(start_bit, end_bit, true) - } - - fn find_next_zero( - &self, - start_bit: BitIndex, - end_bit: BitIndex, - ) -> Result, AcceptError> { - self.find_next_matching(start_bit, end_bit, false) - } - - fn find_next_matching( - &self, - start_bit: BitIndex, - end_bit: BitIndex, - target: bool, - ) -> Result, AcceptError> { - let bit_len = self.capacity()?; - if start_bit.raw() > end_bit.raw() || end_bit.raw() > bit_len { - return Err(AcceptError::OutOfBounds); - } - - if start_bit == end_bit { - return Ok(None); - } - - let mut scan_bit = start_bit.raw(); - let end_bit_raw = end_bit.raw(); - - // Scan leading bits until the index is 64-bit aligned. - while scan_bit < end_bit_raw && (scan_bit & 63) != 0 { - if self.is_set(BitIndex::new(scan_bit))? == target { - return Ok(Some(BitIndex::new(scan_bit))); - } - scan_bit += 1; - } - - // Bulk scan by 64-bit words, then use trailing_zeros for first matching bit. - while end_bit_raw - scan_bit >= 64 { - let next = scan_bit + 64; - - let byte_index = - usize::try_from(scan_bit >> 3).map_err(|_| AcceptError::OutOfBounds)?; - - // SAFETY: `next <= end_bit <= bit_len` guarantees we can read exactly 8 bytes here. - let word = unsafe { - let ptr = self.bits.as_ptr().add(byte_index).cast::(); - u64::from_le(ptr.read_unaligned()) - }; - - let match_word = if target { word } else { !word }; - if match_word != 0 { - let delta = u64::from(match_word.trailing_zeros()); - let found = scan_bit + delta; - return Ok(Some(BitIndex::new(found))); - } - - scan_bit = next; - } - - // Scan remaining tail bits (< 64). - while scan_bit < end_bit_raw { - if self.is_set(BitIndex::new(scan_bit))? == target { - return Ok(Some(BitIndex::new(scan_bit))); - } - scan_bit += 1; - } - - Ok(None) - } -} - -struct BitmapRef<'a> { - bits: &'a [u8], -} - -impl<'a> BitmapRef<'a> { - fn new(bits: &'a [u8]) -> Self { - Self { bits } - } - - fn has_set_bit(&self, start_bit: BitIndex, end_bit: BitIndex) -> Result { - if start_bit >= end_bit { - return Ok(false); - } - - let bit_len = self - .bits - .len() - .checked_mul(8) - .ok_or(AcceptError::ArithmeticOverflow)?; - let start_bit = usize::try_from(start_bit.raw()).map_err(|_| AcceptError::OutOfBounds)?; - let end_bit = usize::try_from(end_bit.raw()).map_err(|_| AcceptError::OutOfBounds)?; - if end_bit > bit_len { - return Err(AcceptError::OutOfBounds); - } - - let mut bit = start_bit; - - // Scan head until byte alignment. - while bit < end_bit && (bit & 7) != 0 { - let byte_idx = bit >> 3; - let mask = 1u8 << (bit & 7); - if (self.bits[byte_idx] & mask) != 0 { - return Ok(true); - } - bit += 1; - } - - let mut byte_idx = bit >> 3; - let end_full_byte = end_bit >> 3; - - // Bulk scan by u64 for full bytes. - while byte_idx + 8 <= end_full_byte { - // SAFETY: loop condition guarantees 8 readable bytes. - let word = unsafe { - let ptr = self.bits.as_ptr().add(byte_idx).cast::(); - u64::from_le(ptr.read_unaligned()) - }; - if word != 0 { - return Ok(true); - } - byte_idx += 8; - } - - while byte_idx < end_full_byte { - if self.bits[byte_idx] != 0 { - return Ok(true); - } - byte_idx += 1; - } - - // Check tail bits in one masked-byte test. - let tail_bits = (end_bit & 7) as u8; - if tail_bits != 0 { - let tail_mask = low_bits_mask(tail_bits); - if (self.bits[end_full_byte] & tail_mask) != 0 { - return Ok(true); - } - } - - Ok(false) - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -struct PhysAddr(u64); - -impl PhysAddr { - const fn new(raw: u64) -> Self { - Self(raw) - } - - const fn raw(self) -> u64 { - self.0 - } - - fn checked_add(self, bytes: u64) -> Result { - self.0 - .checked_add(bytes) - .map(Self) - .ok_or(AcceptError::ArithmeticOverflow) - } - - fn checked_add_units(self, bits: BitIndex, unit_size: u64) -> Result { - let bytes = bits - .raw() - .checked_mul(unit_size) - .ok_or(AcceptError::ArithmeticOverflow)?; - self.checked_add(bytes) - } -} - -/// Strongly-typed bitmap index used to avoid mixing bit position with GPA. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -struct BitIndex(u64); - -impl BitIndex { - const fn new(raw: u64) -> Self { - Self(raw) - } - - const fn raw(self) -> u64 { - self.0 - } -} - -fn low_bits_mask(count: u8) -> u8 { - debug_assert!(count <= 8); - if count == 0 { - 0 - } else { - u8::MAX >> (8 - count) - } -} - -fn bit_range_mask(start_off: u8, end_off: u8) -> u8 { - debug_assert!(start_off <= end_off && end_off <= 8); - if start_off == end_off { - return 0; - } - - let width = end_off - start_off; - low_bits_mask(width) << start_off -} - -fn align_down(addr: u64, align: u64) -> u64 { - addr & !(align - 1) -} - -fn align_up(addr: u64, align: u64) -> Option { - addr.checked_add(align - 1).map(|v| v & !(align - 1)) -} - -/// Result of a conditional accept operation. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum AcceptOutcome { - /// The range overlaps bitmap coverage but no unit is pending (already accepted). - AlreadyAccepted, - /// At least one pending unit was accepted and the corresponding bitmap bits were cleared. - AcceptedNow { - /// Number of bitmap units accepted by this call. - accepted_units: u64, - }, - /// The range does not overlap bitmap coverage. - OutOfCoverage, -} diff --git a/src/unaccepted_memory/accept.rs b/src/unaccepted_memory/accept.rs new file mode 100644 index 0000000..2d44b7e --- /dev/null +++ b/src/unaccepted_memory/accept.rs @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: BSD-3-Clause +// Copyright(c) 2026 Intel Corporation. + +//! Acceptance-path operations for unaccepted-memory bitmap ranges. +//! +//! This file contains range acceptance flows, plus pending-run claim/restore helpers. + +use super::{bitmap::BitIndex, EfiUnacceptedMemory}; +use crate::{accept_memory, AcceptError}; + +impl EfiUnacceptedMemory { + /// Accepts bitmap-marked units that overlap `start..end`, then clears accepted bits. + /// + /// # Safety + /// + /// The caller must ensure this table and bitmap describe pending private-memory units, + /// and the target GPA ranges are valid for TDX acceptance. + pub unsafe fn accept_range(&self, start: u64, end: u64) -> Result<(), AcceptError> { + let Some((first_bit, last_bit, unit_size)) = self.overlapping_bit_range(start, end) else { + return Ok(()); + }; + + let phys_base = self.phys_base; + let bitmap = self.bitmap_ref(); + let mut scan = first_bit; + while let Some(run_start) = bitmap.find_next_set(scan, last_bit) { + let run_end = bitmap + .find_next_zero(run_start, last_bit) + .unwrap_or(last_bit); + + let run_gpa_start = Self::bit_to_gpa(phys_base, run_start, unit_size)?; + let run_gpa_end = Self::bit_to_gpa(phys_base, run_end, unit_size)?; + + // SAFETY: Caller guarantees bitmap/GPA mapping validity for pending private pages. + unsafe { accept_memory(run_gpa_start, run_gpa_end)? }; + bitmap.clear_range(run_start, run_end); + + scan = run_end; + } + + Ok(()) + } + + /// Finds the first contiguous run of set bits overlapping `[start, end)`, + /// clears those bits, and returns the corresponding GPA range. + /// clears those bits, and returns the corresponding GPA range. + /// + /// # Safety + /// + /// The caller must ensure: + /// - No concurrent operation touches the same bitmap bits. + pub unsafe fn claim_next_pending_run( + &self, + start: u64, + end: u64, + ) -> Result, AcceptError> { + let Some((first_bit, last_bit, unit_size)) = self.overlapping_bit_range(start, end) else { + return Ok(None); + }; + + // SAFETY: Public concurrent API contract guarantees valid writable bitmap + // payload and atomic-access discipline for overlapping ranges. + let bitmap = self.bitmap_ref(); + let Some(run_start) = bitmap.find_next_set(first_bit, last_bit) else { + return Ok(None); + }; + let run_end = bitmap + .find_next_zero(run_start, last_bit) + .unwrap_or(last_bit); + + bitmap.clear_range(run_start, run_end); + + let gpa_start = Self::bit_to_gpa(self.phys_base, run_start, unit_size)?; + let gpa_end = Self::bit_to_gpa(self.phys_base, run_end, unit_size)?; + Ok(Some((gpa_start, gpa_end))) + } + + /// Re-sets bitmap bits for a GPA range whose TDX accept failed. + /// + /// # Safety + /// + /// The caller must ensure: + /// - No concurrent operation touches the same bitmap bits. + /// - `start..end` is exactly a unit-aligned range previously returned by + /// [`Self::claim_next_pending_run`] and has not been accepted or restored. + pub unsafe fn restore_pending_range(&self, start: u64, end: u64) { + let Some((first_bit, last_bit, _unit_size)) = self.overlapping_bit_range(start, end) else { + return; + }; + + // SAFETY: Public concurrent API contract guarantees valid writable bitmap + // payload and atomic-access discipline for overlapping ranges. + let bitmap = self.bitmap_ref(); + bitmap.set_range(first_bit, last_bit); + } + + fn bit_to_gpa(phys_base: u64, bit: BitIndex, unit_size: u64) -> Result { + let offset = bit + .checked_mul(unit_size) + .ok_or(AcceptError::ArithmeticOverflow)?; + phys_base + .checked_add(offset) + .ok_or(AcceptError::ArithmeticOverflow) + } +} diff --git a/src/unaccepted_memory/bitmap.rs b/src/unaccepted_memory/bitmap.rs new file mode 100644 index 0000000..aced81f --- /dev/null +++ b/src/unaccepted_memory/bitmap.rs @@ -0,0 +1,191 @@ +// SPDX-License-Identifier: BSD-3-Clause +// Copyright(c) 2026 Intel Corporation. + +//! Bitmap data structures for tracking unaccepted memory. +//! +//! This module provides [`BitmapRef`], a view backed by a slice of atomic 64-bit +//! words (`&'a [AtomicU64]`). + +use core::sync::atomic::{AtomicU64, Ordering}; + +pub type BitIndex = u64; + +/// Bitmap view backed by a slice of [`AtomicU64`]. +/// +/// Supports lock-free queries (`has_set_bit`, `find_next_set`, etc.) as well +/// as atomic range updates (`set_range`, `clear_range`, `clear_all`) via +/// `AtomicU64`'s relaxed atomic operations. +#[derive(Clone, Copy)] +pub struct BitmapRef<'a> { + bits: &'a [AtomicU64], +} + +impl<'a> BitmapRef<'a> { + /// Creates a bitmap view from a slice of atomic 64-bit words. + pub const fn new(bits: &'a [AtomicU64]) -> Self { + Self { bits } + } + + /// Creates a bitmap view from a `u64`-aligned raw pointer. + /// + /// # Safety + /// + /// - `ptr` must be non-null and aligned to `align_of::()`. + /// - `ptr` must point to at least `len_bytes` bytes of valid memory. + /// - `len_bytes` must be a multiple of `size_of::()`. + /// - The memory must remain valid for lifetime `'a`. + pub(super) unsafe fn from_raw(ptr: *const u8, len_bytes: usize) -> Self { + debug_assert_eq!(len_bytes % core::mem::size_of::(), 0); + let len_words = len_bytes / core::mem::size_of::(); + // SAFETY: Caller guarantees alignment, validity, and length constraints. + let bits = unsafe { core::slice::from_raw_parts(ptr.cast::(), len_words) }; + Self { bits } + } + + /// Returns the underlying atomic words. + pub const fn words(&self) -> &'a [AtomicU64] { + self.bits + } + + /// Returns the total capacity in bits. + pub const fn capacity(&self) -> u64 { + (self.bits.len() as u64) * 64 + } + + /// Returns `true` if any bit in `[start_bit, end_bit)` is set. + pub fn has_set_bit(&self, start_bit: BitIndex, end_bit: BitIndex) -> bool { + let total_bits = self.capacity(); + if start_bit >= end_bit || start_bit >= total_bits { + return false; + } + let end_bit = end_bit.min(total_bits); + + let start = start_bit as usize; + let end = end_bit as usize; + + let start_word = start / 64; + let end_word = (end - 1) / 64; + + for word_idx in start_word..=end_word { + let word_bit_start = word_idx * 64; + let lo = start.saturating_sub(word_bit_start); + let hi = end.min(word_bit_start + 64) - word_bit_start; + let mask = word_range_mask(lo, hi); + if self.bits[word_idx].load(Ordering::Relaxed) & mask != 0 { + return true; + } + } + + false + } + + /// Returns the total number of set bits (count of ones) across the bitmap. + pub fn pending_unit_count(&self) -> u64 { + self.bits + .iter() + .map(|word| word.load(Ordering::Relaxed).count_ones() as u64) + .sum() + } + + /// Clears all bits in the bitmap. + pub fn clear_all(&self) { + for word in self.bits { + word.store(0, Ordering::Relaxed); + } + } + + /// Finds the first set bit (1) in `[start_bit, end_bit)`. + pub fn find_next_set(&self, start_bit: BitIndex, end_bit: BitIndex) -> Option { + self.find_next_matching(start_bit, end_bit, true) + } + + /// Finds the first cleared bit (0) in `[start_bit, end_bit)`. + pub fn find_next_zero(&self, start_bit: BitIndex, end_bit: BitIndex) -> Option { + self.find_next_matching(start_bit, end_bit, false) + } + + /// Sets all bits in `[start_bit, end_bit)` to `1`. + pub fn set_range(&self, start_bit: BitIndex, end_bit: BitIndex) { + self.update_range(start_bit, end_bit, true); + } + + /// Clears all bits in `[start_bit, end_bit)` to `0`. + pub fn clear_range(&self, start_bit: BitIndex, end_bit: BitIndex) { + self.update_range(start_bit, end_bit, false); + } + + fn find_next_matching( + &self, + start_bit: BitIndex, + end_bit: BitIndex, + target: bool, + ) -> Option { + let total_bits = self.capacity(); + if start_bit >= end_bit || start_bit >= total_bits { + return None; + } + let end_bit = end_bit.min(total_bits); + + let start = start_bit as usize; + let end = end_bit as usize; + + let start_word = start / 64; + let end_word = (end - 1) / 64; + + for word_idx in start_word..=end_word { + let word_bit_start = word_idx * 64; + let lo = start.saturating_sub(word_bit_start); + let hi = end.min(word_bit_start + 64) - word_bit_start; + let mask = word_range_mask(lo, hi); + + let word = self.bits[word_idx].load(Ordering::Relaxed); + let match_bits = (if target { word } else { !word }) & mask; + if match_bits != 0 { + let delta = match_bits.trailing_zeros() as usize; + let found = (word_bit_start + delta) as u64; + return Some(found); + } + } + + None + } + + fn update_range(&self, start_bit: BitIndex, end_bit: BitIndex, set_bits: bool) { + let total_bits = self.capacity(); + if start_bit >= end_bit || start_bit >= total_bits { + return; + } + let end_bit = end_bit.min(total_bits); + + let start = start_bit as usize; + let end = end_bit as usize; + + let start_word = start / 64; + let end_word = (end - 1) / 64; + + for word_idx in start_word..=end_word { + let word_bit_start = word_idx * 64; + let lo = start.saturating_sub(word_bit_start); + let hi = end.min(word_bit_start + 64) - word_bit_start; + let mask = word_range_mask(lo, hi); + + if set_bits { + self.bits[word_idx].fetch_or(mask, Ordering::Relaxed); + } else { + self.bits[word_idx].fetch_and(!mask, Ordering::Relaxed); + } + } + } +} + +/// Returns a 64-bit mask with bits in `[lo, hi)` set to `1` and all other bits cleared. +fn word_range_mask(lo: usize, hi: usize) -> u64 { + debug_assert!(lo <= hi && hi <= 64); + if lo >= hi { + 0 + } else { + let mask_hi = if hi == 64 { !0u64 } else { (1u64 << hi) - 1 }; + let mask_lo = (1u64 << lo) - 1; + mask_hi & !mask_lo + } +} diff --git a/src/unaccepted_memory/layout.rs b/src/unaccepted_memory/layout.rs new file mode 100644 index 0000000..4fb8d35 --- /dev/null +++ b/src/unaccepted_memory/layout.rs @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: BSD-3-Clause +// Copyright(c) 2026 Intel Corporation. + +//! Layout and mapping helpers for unaccepted-memory metadata. +//! +//! This file provides address/bit conversions, coverage ranges, and +//! trailing-bitmap views based on the [`EfiUnacceptedMemory`] type invariants. + +use core::{mem::size_of, ops::Range}; + +use super::{ + bitmap::{BitIndex, BitmapRef}, + EfiUnacceptedMemory, EFI_UNACCEPTED_UNIT_SIZE, +}; +use crate::AcceptError; + +impl EfiUnacceptedMemory { + /// Returns the physical address range covered by the bitmap. + pub fn coverage_range(&self) -> Range { + let end = self.phys_base + self.coverage_size(); + self.phys_base..end + } + + /// Returns the total physical address range size (in bytes) covered by the bitmap. + pub fn coverage_size(&self) -> u64 { + self.total_bits() * u64::from(self.unit_size_bytes) + } + + /// Alias for [`Self::coverage_range`]. + #[inline] + pub fn bitmap_coverage_range(&self) -> Range { + self.coverage_range() + } + + /// Returns an atomic view of the trailing bitmap. + pub fn bitmap_ref(&self) -> BitmapRef<'_> { + let bitmap_len = self.bitmap_size_bytes as usize; + let bitmap_ptr = core::ptr::from_ref(self) + .cast::() + .wrapping_add(core::mem::size_of::()); + // SAFETY: The type invariant of `EfiUnacceptedMemory` guarantees that the header + // is immediately followed by a valid trailing bitmap of `self.bitmap_size_bytes` bytes, + // aligned to `AtomicU64`, and with length a multiple of `size_of::()`. + unsafe { BitmapRef::from_raw(bitmap_ptr, bitmap_len) } + } + + pub(super) fn total_bits(&self) -> u64 { + self.bitmap_size_bytes * 8 + } + + /// Converts a GPA range into an overlapping bitmap bit range. + /// + /// Returns `None` when there is no overlap with bitmap coverage. + pub(super) fn overlapping_bit_range( + &self, + start: u64, + end: u64, + ) -> Option<(BitIndex, BitIndex, u64)> { + if start >= end { + return None; + } + + let unit_size = u64::from(self.unit_size_bytes); + let coverage = self.coverage_range(); + + let range_start = start.max(coverage.start); + let range_end = end.min(coverage.end); + if range_start >= range_end { + return None; + } + + let rel_start = range_start - self.phys_base; + let rel_end = range_end - self.phys_base; + + let first_bit = rel_start / unit_size; + let last_bit = rel_end.div_ceil(unit_size); + + Some((first_bit, last_bit, unit_size)) + } +} + +pub(super) struct BitmapLayout { + pub(super) coverage_phys_base: u64, + pub(super) size_bytes: usize, +} + +impl BitmapLayout { + /// Computes the physical base address and required trailing bitmap size (in bytes) + /// to cover memory within `range`. + pub(super) fn from_range(range: Range) -> Result { + if range.start >= range.end { + return Err(AcceptError::OutOfBounds); + } + + let phys_base = align_down(range.start, EFI_UNACCEPTED_UNIT_SIZE); + let coverage_end = + align_up(range.end, EFI_UNACCEPTED_UNIT_SIZE).ok_or(AcceptError::ArithmeticOverflow)?; + let coverage_size = coverage_end - phys_base; + let bitmap_bits = coverage_size / EFI_UNACCEPTED_UNIT_SIZE; + let bitmap_words = bitmap_bits.div_ceil(64); + let bitmap_size_bytes = bitmap_words * size_of::() as u64; + + Ok(Self { + coverage_phys_base: phys_base, + size_bytes: usize::try_from(bitmap_size_bytes).map_err(|_| AcceptError::OutOfBounds)?, + }) + } +} + +pub(super) fn min_max_from_ranges(ranges: &[Range]) -> Result, AcceptError> { + let first = ranges.first().ok_or(AcceptError::OutOfBounds)?; + if first.start >= first.end { + return Err(AcceptError::OutOfBounds); + } + + let mut min_addr = first.start; + let mut max_addr = first.end; + for range in &ranges[1..] { + if range.start >= range.end { + return Err(AcceptError::OutOfBounds); + } + min_addr = min_addr.min(range.start); + max_addr = max_addr.max(range.end); + } + + Ok(min_addr..max_addr) +} + +pub(super) fn align_down(addr: u64, align: u64) -> u64 { + addr & !(align - 1) +} + +pub(super) fn align_up(addr: u64, align: u64) -> Option { + addr.checked_add(align - 1).map(|v| v & !(align - 1)) +} diff --git a/src/unaccepted_memory/mod.rs b/src/unaccepted_memory/mod.rs new file mode 100644 index 0000000..384d5df --- /dev/null +++ b/src/unaccepted_memory/mod.rs @@ -0,0 +1,243 @@ +// SPDX-License-Identifier: BSD-3-Clause +// Copyright(c) 2026 Intel Corporation. + +//! Support for unaccepted memory in TDX guest environments. +//! +//! This module provides mechanisms to manage and accept +//! unaccepted memory regions in TDX guests. +//! The core data structure is [`EfiUnacceptedMemory`], +//! which represents the EFI table header +//! and provides methods to manipulate the unaccepted memory bitmap +//! and perform acceptance operations. + +mod accept; +mod bitmap; +mod layout; +mod parse; +mod query; +mod register; +#[cfg(test)] +mod tests; + +use core::{ + marker::PhantomPinned, mem::size_of, ops::Range, pin::Pin, ptr::NonNull, + sync::atomic::AtomicU64, +}; + +pub use bitmap::BitmapRef; +use layout::{min_max_from_ranges, BitmapLayout}; + +use crate::AcceptError; + +/// Builder for constructing an [`EfiUnacceptedMemory`] table incrementally without heap allocation. +pub struct EfiUnacceptedMemoryBuilder { + table: Pin<&'static mut EfiUnacceptedMemory>, +} + +impl EfiUnacceptedMemoryBuilder { + /// Registers a single unaccepted memory range. + /// + /// # Safety + /// + /// The caller must ensure `start..end` is valid guest-private memory in pending/acceptable state. + pub unsafe fn register_range(&mut self, start: u64, end: u64) -> Result<(), AcceptError> { + unsafe { self.table.as_mut().register_range(start, end) } + } + + /// Consumes the builder and returns the completed pinned [`EfiUnacceptedMemory`] table. + pub fn build(self) -> Pin<&'static mut EfiUnacceptedMemory> { + self.table + } +} + +/// GUID of the Linux-compatible unaccepted-memory EFI table. +pub const LINUX_EFI_UNACCEPTED_MEM_TABLE_GUID: uefi_raw::Guid = + uefi_raw::guid!("d5d1de3c-105c-44f9-9ea9-bcef98120031"); + +/// Version of the Linux-compatible unaccepted-memory EFI table supported here. +pub const LINUX_EFI_UNACCEPTED_MEM_TABLE_VERSION: u32 = 1; + +/// Unit size for unaccepted-memory bitmap entries (2 MiB). +pub const EFI_UNACCEPTED_UNIT_SIZE: u64 = 2 * 1024 * 1024; + +/// Header of the Linux-compatible EFI unaccepted-memory table. +/// +/// This type describes only the fixed-size header. The bitmap payload is stored +/// immediately after the header in memory (C-style trailing data): +/// +/// ### Memory Layout +/// The total memory footprint is +/// `size_of::() + self.bitmap_size_bytes`. +/// The bitmap begins at the first byte following this structure. +/// +/// ### Type Invariants +/// Any valid reference (`&Self` or `Pin<&mut Self>`) satisfies the following invariants: +/// - The header is immediately followed in memory by a valid trailing bitmap of +/// `self.bitmap_size_bytes` bytes. +/// - The trailing bitmap starts at a pointer aligned to `align_of::()` (8 bytes), +/// and `self.bitmap_size_bytes` is a non-zero multiple of `size_of::()` (8 bytes). +/// - `self.unit_size_bytes` is a non-zero power of two. +/// - The physical address range covered by the bitmap (`self.phys_base..self.phys_base + coverage_size`) +/// does not overflow `u64`. +/// +/// Because moving an `EfiUnacceptedMemory` would detach it from its trailing bitmap, +/// mutable access requires pinning ([`Pin<&mut Self>`]). +/// +/// ### Bitmap Semantics +/// - Each bit in the trailing bitmap represents a memory region of +/// `unit_size_bytes` bytes. +/// - Bit 0 corresponds to the physical address specified by `phys_base`. +/// - A **set bit (1)** indicates memory is unaccepted (pending); +/// a **cleared bit (0)** indicates it has been accepted. +#[derive(Debug)] +#[repr(C, packed)] +pub struct EfiUnacceptedMemory { + /// The version of the table. Currently, only version 1 is defined. + version: u32, + /// The size of the memory region represented by a single bit in the bitmap. + /// Typically set to 2MiB (0x200000) to align with huge page boundaries. + unit_size_bytes: u32, + /// The start physical address of the memory range covered by this bitmap. + /// Bit 0 of the bitmap corresponds to this address. + phys_base: u64, + /// The bitmap payload length in bytes, excluding this header. + bitmap_size_bytes: u64, + _pinned: PhantomPinned, +} + +impl EfiUnacceptedMemory { + /// Returns the allocation size required for a table covering `ranges`. + /// + /// Returns `Err(AcceptError::OutOfBounds)` when `ranges` is empty or contains an empty range. + pub fn required_size(ranges: &[Range]) -> Result { + let coverage = min_max_from_ranges(ranges)?; + Self::required_size_for_range(coverage) + } + + /// Returns the allocation size required for a table covering `coverage`. + pub fn required_size_for_range(coverage: Range) -> Result { + let layout = BitmapLayout::from_range(coverage)?; + + size_of::() + .checked_add(layout.size_bytes) + .ok_or(AcceptError::ArithmeticOverflow) + } + + /// Starts building a table in caller-provided memory covering `coverage`. + /// + /// This initializes the header and clears the bitmap, returning an [`EfiUnacceptedMemoryBuilder`] + /// so the caller can register unaccepted regions one at a time without heap allocation. + /// + /// # Safety + /// + /// The caller must ensure that: + /// - `table_addr` points to a writable allocation of `allocation_size` bytes; + /// - `allocation_size` is at least [`Self::required_size_for_range(coverage)`]; + /// - the allocation remains valid for the returned `'static` reference; + /// - no other reference accesses the allocation while the table exists. + pub unsafe fn builder( + table_addr: NonNull, + allocation_size: usize, + coverage: Range, + ) -> Result { + let bitmap_layout = BitmapLayout::from_range(coverage)?; + let required_size = size_of::() + .checked_add(bitmap_layout.size_bytes) + .ok_or(AcceptError::ArithmeticOverflow)?; + + let table_raw = table_addr.as_ptr(); + if !table_raw.is_aligned() + || !table_raw + .addr() + .is_multiple_of(core::mem::align_of::()) + { + return Err(AcceptError::InvalidAlignment); + } + let bitmap_addr = table_raw + .addr() + .checked_add(size_of::()) + .ok_or(AcceptError::ArithmeticOverflow)?; + if !bitmap_addr.is_multiple_of(core::mem::align_of::()) { + return Err(AcceptError::InvalidAlignment); + } + if allocation_size < required_size { + return Err(AcceptError::OutOfBounds); + } + + let table_ptr = table_addr.cast::(); + table_ptr.as_ptr().write(Self { + version: LINUX_EFI_UNACCEPTED_MEM_TABLE_VERSION, + unit_size_bytes: EFI_UNACCEPTED_UNIT_SIZE as u32, + phys_base: bitmap_layout.coverage_phys_base, + bitmap_size_bytes: bitmap_layout.size_bytes as u64, + _pinned: PhantomPinned, + }); + + // SAFETY: The caller provides unique access to a sufficiently large, + // writable allocation that remains valid for the returned reference. + // The type invariants for header and bitmap layout have been established. + let table = unsafe { Pin::new_unchecked(&mut *table_ptr.as_ptr()) }; + table.bitmap_ref().clear_all(); + + Ok(EfiUnacceptedMemoryBuilder { table }) + } + + /// Initializes a complete table in a caller-provided allocation. + /// + /// This initializes the header and bitmap, then registers every range before + /// returning the table. `ranges` must not be empty. + /// + /// This operation is not transactional. Registering a range may accept its + /// unaligned edges immediately. If a later registration or TDX operation + /// fails, some memory may already have been accepted and the allocation may + /// contain a partially initialized table; the returned error does not roll + /// back those effects. + /// + /// # Safety + /// + /// The caller must ensure that: + /// - `table_addr` points to a writable allocation of `allocation_size` bytes; + /// - the allocation remains valid for the returned `'static` reference; + /// - `ranges` is non-empty and its entries are non-empty, mutually + /// non-overlapping, 4-KiB-aligned GPA ranges; + /// - every range is valid guest-private memory in the pending state and may + /// be accepted by this function; + /// - no other reference accesses the allocation while the returned mutable + /// reference exists. + pub unsafe fn new( + table_addr: NonNull, + allocation_size: usize, + ranges: &[Range], + ) -> Result, AcceptError> { + let coverage = min_max_from_ranges(ranges)?; + let mut builder = unsafe { Self::builder(table_addr, allocation_size, coverage)? }; + + for range in ranges { + // SAFETY: The caller guarantees every input range is valid + // guest-private memory in pending/acceptable state. + unsafe { builder.register_range(range.start, range.end)? }; + } + + Ok(builder.build()) + } + + /// Returns the version of the table header. + pub const fn version(&self) -> u32 { + self.version + } + + /// Returns the unit size represented by one bitmap bit, in bytes. + pub const fn unit_size_bytes(&self) -> u32 { + self.unit_size_bytes + } + + /// Returns the start physical address covered by the bitmap. + pub const fn phys_base(&self) -> u64 { + self.phys_base + } + + /// Returns the trailing bitmap payload length, in bytes. + pub const fn bitmap_size_bytes(&self) -> u64 { + self.bitmap_size_bytes + } +} diff --git a/src/unaccepted_memory/parse.rs b/src/unaccepted_memory/parse.rs new file mode 100644 index 0000000..550aa68 --- /dev/null +++ b/src/unaccepted_memory/parse.rs @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: BSD-3-Clause +// Copyright(c) 2026 Intel Corporation. + +//! Functions for finding and validating the unaccepted-memory table from EFI tables. + +use core::{ + mem::{align_of, size_of}, + ptr::NonNull, + sync::atomic::AtomicU64, +}; + +use super::{ + EfiUnacceptedMemory, LINUX_EFI_UNACCEPTED_MEM_TABLE_GUID, + LINUX_EFI_UNACCEPTED_MEM_TABLE_VERSION, +}; + +impl EfiUnacceptedMemory { + /// Locates and validates the unaccepted-memory table from an EFI system table. + /// + /// # Safety + /// + /// The caller must ensure that: + /// - if `systab` is non-null, it points to a valid, readable EFI system table; + /// - if the system table's configuration table pointer is non-null, it points to a readable + /// array containing `number_of_configuration_table_entries` valid entries; + /// - any non-null vendor table pointer in a matching configuration table entry points to a + /// valid, readable [`EfiUnacceptedMemory`] header. + pub unsafe fn from_system_table( + systab: *const uefi_raw::table::system::SystemTable, + ) -> Option> { + if systab.is_null() || !systab.is_aligned() { + log::warn!("EFI system table is null or misaligned"); + return None; + } + + // SAFETY: Caller guarantees `systab` points to a valid, accessible EFI System Table. + let systab = unsafe { &*systab }; + + let configuration_table = systab.configuration_table; + if configuration_table.is_null() || !configuration_table.is_aligned() { + log::warn!("EFI configuration table is null or misaligned"); + return None; + } + + let configuration_table_size = systab + .number_of_configuration_table_entries + .checked_mul(size_of::()); + if configuration_table_size.is_none_or(|size| size > isize::MAX as usize) { + log::warn!("EFI configuration table is too large"); + return None; + } + + // SAFETY: `configuration_table` is non-null, suitably aligned, and its total byte length fits in a slice. + let entries = unsafe { + core::slice::from_raw_parts( + configuration_table, + systab.number_of_configuration_table_entries, + ) + }; + + // SAFETY: Caller guarantees accessible memory for table pointers. + unsafe { Self::from_configuration_tables(entries) } + } + + /// Locates and validates the unaccepted-memory table from EFI configuration table entries. + /// + /// # Safety + /// + /// The vendor table pointers within matching configuration table entries must point to + /// valid and accessible memory if present. + pub unsafe fn from_configuration_tables( + entries: &[uefi_raw::table::configuration::ConfigurationTable], + ) -> Option> { + let table_ptr = entries + .iter() + .find(|entry| entry.vendor_guid == LINUX_EFI_UNACCEPTED_MEM_TABLE_GUID)? + .vendor_table + .cast::(); + + let non_null_table = NonNull::new(table_ptr)?; + if !table_ptr.is_aligned() { + return None; + } + + // SAFETY: The pointer is non-null and aligned. Caller ensures it points to accessible memory. + let table = unsafe { non_null_table.as_ref() }; + + if table.version() != LINUX_EFI_UNACCEPTED_MEM_TABLE_VERSION { + log::warn!( + "Unknown unaccepted memory table version: {}", + table.version() + ); + return None; + } + + if table.unit_size_bytes() == 0 || !table.unit_size_bytes().is_power_of_two() { + log::warn!( + "Invalid unaccepted memory table unit size: {}", + table.unit_size_bytes() + ); + return None; + } + + let bitmap_addr = table_ptr.addr().checked_add(size_of::())?; + if !bitmap_addr.is_multiple_of(align_of::()) + || table.bitmap_size_bytes() == 0 + || !table + .bitmap_size_bytes() + .is_multiple_of(size_of::() as u64) + { + log::warn!( + "Invalid unaccepted memory table bitmap size: {}", + table.bitmap_size_bytes() + ); + return None; + } + + let total_bits = table.bitmap_size_bytes().checked_mul(8)?; + let total_size = total_bits.checked_mul(u64::from(table.unit_size_bytes()))?; + if table.phys_base().checked_add(total_size).is_none() { + log::warn!("Unaccepted memory table coverage overflows"); + return None; + } + + Some(non_null_table) + } +} diff --git a/src/unaccepted_memory/query.rs b/src/unaccepted_memory/query.rs new file mode 100644 index 0000000..62fec2b --- /dev/null +++ b/src/unaccepted_memory/query.rs @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: BSD-3-Clause +// Copyright(c) 2026 Intel Corporation. + +//! Read-only query APIs for unaccepted-memory status. +//! +//! This file implements pending-state checks over GPA ranges using +//! lock-free atomic-word reads backed by [`BitmapRef`]. + +use super::EfiUnacceptedMemory; + +impl EfiUnacceptedMemory { + /// Returns the number of pending (set) bitmap units in the whole table. + pub fn pending_unit_count(&self) -> u64 { + self.bitmap_ref().pending_unit_count() + } + + /// Checks whether `[start, end)` overlaps any pending (unaccepted) bitmap unit. + pub fn is_range_pending(&self, start: u64, end: u64) -> bool { + let Some((first_bit, last_bit, _unit_size)) = self.overlapping_bit_range(start, end) else { + return false; + }; + self.bitmap_ref().has_set_bit(first_bit, last_bit) + } + + /// Returns `true` if all bitmap units overlapping `[start, end)` have been accepted. + pub fn is_fully_accepted(&self, start: u64, end: u64) -> bool { + !self.is_range_pending(start, end) + } +} diff --git a/src/unaccepted_memory/register.rs b/src/unaccepted_memory/register.rs new file mode 100644 index 0000000..06d4df4 --- /dev/null +++ b/src/unaccepted_memory/register.rs @@ -0,0 +1,141 @@ +// SPDX-License-Identifier: BSD-3-Clause +// Copyright(c) 2026 Intel Corporation. + +//! Registration-path logic for unaccepted-memory ranges. +//! +//! This file handles `register_range` splitting: it accepts edges/out-of-coverage +//! segments eagerly and marks in-coverage aligned segments as pending in the bitmap. + +use core::pin::Pin; + +use super::{ + layout::{align_down, align_up}, + EfiUnacceptedMemory, +}; +use crate::AcceptError; + +impl EfiUnacceptedMemory { + /// Processes `start..end` by eagerly accepting required parts and deferring the rest in bitmap. + /// + /// # Safety + /// + /// The caller must ensure the range is valid guest-private memory in pending/acceptable state. + pub(super) unsafe fn register_range( + mut self: Pin<&mut Self>, + start: u64, + end: u64, + ) -> Result<(), AcceptError> { + let table_phys_base = self.phys_base; + let unit_size = u64::from(self.unit_size_bytes); + if start >= end { + return Ok(()); + } + + let unit_mask = unit_size - 1; + + if end - start < 2 * unit_size { + // SAFETY: Caller guarantees the physical range is valid for TDX acceptance. + return unsafe { Self::try_accept_range(start, end) }; + } + + let mut current_start = start; + let mut current_end = end; + + if current_start & unit_mask != 0 { + let Some(aligned_start) = align_up(current_start, unit_size) else { + return Err(AcceptError::ArithmeticOverflow); + }; + // SAFETY: Caller guarantees the physical subrange is valid for TDX acceptance. + unsafe { Self::try_accept_range(current_start, aligned_start)? }; + current_start = aligned_start; + } + + if current_end & unit_mask != 0 { + let aligned_end = align_down(current_end, unit_size); + // SAFETY: Caller guarantees the physical subrange is valid for TDX acceptance. + unsafe { Self::try_accept_range(aligned_end, current_end)? }; + current_end = aligned_end; + } + + let coverage = self.bitmap_coverage_range(); + let bitmap_end = coverage.end; + + // 1) Process aligned range before bitmap coverage. + if current_start < table_phys_base { + let accept_end = current_end.min(table_phys_base); + // SAFETY: Caller guarantees the physical subrange is valid for TDX acceptance. + unsafe { Self::try_accept_range(current_start, accept_end)? }; + current_start = accept_end; + } + + if current_start >= current_end { + return Ok(()); + } + + // 2) Process aligned range within bitmap coverage. + if current_start < bitmap_end { + let bitmap_range_end = current_end.min(bitmap_end); + self.as_mut() + .mark_range_as_unaccepted(current_start, bitmap_range_end, unit_size)?; + current_start = bitmap_range_end; + } + + // 3) Process aligned range after bitmap coverage. + if current_start < current_end { + // SAFETY: Caller guarantees the physical subrange is valid for TDX acceptance. + unsafe { Self::try_accept_range(current_start, current_end)? }; + } + + Ok(()) + } + + /// Accepts physical memory in `start..end` if the range is non-empty. + /// + /// # Safety + /// + /// The caller must ensure `start..end` is a valid GPA range for TDX acceptance. + unsafe fn try_accept_range(start: u64, end: u64) -> Result<(), AcceptError> { + if start < end { + // SAFETY: Caller guarantees the physical range is valid for TDX acceptance. + return unsafe { crate::accept_memory(start, end) }; + } + Ok(()) + } + + /// Marks unit-aligned bitmap-covered GPA range `start..end` as unaccepted bits. + fn mark_range_as_unaccepted( + self: Pin<&mut Self>, + start: u64, + end: u64, + unit_size: u64, + ) -> Result<(), AcceptError> { + if start >= end { + return Ok(()); + } + + if !start.is_multiple_of(unit_size) || !end.is_multiple_of(unit_size) { + return Err(AcceptError::InvalidAlignment); + } + + let start_bit = (start - self.phys_base) / unit_size; + let end_bit = (end - self.phys_base) / unit_size; + let total_bits = self.total_bits(); + + let clamped_start_bit = start_bit.min(total_bits); + let clamped_end_bit = end_bit.min(total_bits); + if clamped_start_bit >= clamped_end_bit { + return Ok(()); + } + + let bitmap = self.bitmap_ref(); + if bitmap + .find_next_set(clamped_start_bit, clamped_end_bit) + .is_some() + { + return Err(AcceptError::Overlap); + } + bitmap.set_range(clamped_start_bit, clamped_end_bit); + + Ok(()) + } +} diff --git a/src/unaccepted_memory/tests.rs b/src/unaccepted_memory/tests.rs new file mode 100644 index 0000000..7a808e7 --- /dev/null +++ b/src/unaccepted_memory/tests.rs @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: BSD-3-Clause +// Copyright(c) 2026 Intel Corporation. + +use alloc::vec; +use core::{mem::size_of, ptr::NonNull, sync::atomic::AtomicU64}; + +use super::*; +use crate::AcceptError; + +#[test] +#[allow(clippy::reversed_empty_ranges, clippy::single_range_in_vec_init)] +fn test_required_size() { + assert!(matches!( + EfiUnacceptedMemory::required_size(&[]), + Err(AcceptError::OutOfBounds) + )); + + let inverted = Range { + start: 100, + end: 50, + }; + assert!(matches!( + EfiUnacceptedMemory::required_size(core::slice::from_ref(&inverted)), + Err(AcceptError::OutOfBounds) + )); + + // 2 MiB to 6 MiB is 4 MiB = 2 units. 2 bits fit in 1 u64 (8 bytes). + let range = 0x200_000..0x600_000; + let size = EfiUnacceptedMemory::required_size(core::slice::from_ref(&range)).unwrap(); + assert_eq!(size, size_of::() + 8); +} + +#[test] +fn test_new_and_type_invariants() { + let range = 0x200_000..0x600_000; + let required_size = EfiUnacceptedMemory::required_size(core::slice::from_ref(&range)).unwrap(); + + let mut buf = vec![0u8; required_size + 16]; + // Ensure 8-byte alignment + let align_offset = buf + .as_ptr() + .align_offset(core::mem::align_of::()); + let aligned_slice = &mut buf[align_offset..align_offset + required_size]; + let table_addr = NonNull::new(aligned_slice.as_mut_ptr()).unwrap(); + + let table = unsafe { + EfiUnacceptedMemory::new( + table_addr, + aligned_slice.len(), + core::slice::from_ref(&range), + ) + .unwrap() + }; + + assert_eq!(table.version(), LINUX_EFI_UNACCEPTED_MEM_TABLE_VERSION); + assert_eq!(table.unit_size_bytes(), EFI_UNACCEPTED_UNIT_SIZE as u32); + assert_eq!(table.phys_base(), 0x200_000); + assert_eq!(table.bitmap_size_bytes(), 8); + assert_eq!(table.coverage_range(), 0x200_000..0x8_200_000); // 8 bytes * 8 bits = 64 bits * 2 MiB = 128 MiB + assert_eq!(table.bitmap_coverage_range(), table.coverage_range()); + assert_eq!(table.coverage_size(), 128 * 1024 * 1024); + + // Both units in [0x200_000, 0x600_000) are pending + assert_eq!(table.pending_unit_count(), 2); + assert!(table.is_range_pending(0x200_000, 0x600_000)); + assert!(!table.is_fully_accepted(0x200_000, 0x600_000)); + + // Range outside coverage or without pending bits + assert!(!table.is_range_pending(0x600_000, 0x800_000)); + assert!(table.is_fully_accepted(0x600_000, 0x800_000)); + + // Clear unit 0 via bitmap_ref + table.bitmap_ref().clear_range(0, 1); + assert_eq!(table.pending_unit_count(), 1); + assert!(!table.is_range_pending(0x200_000, 0x400_000)); + assert!(table.is_range_pending(0x400_000, 0x600_000)); + + // Clear all + table.bitmap_ref().clear_all(); + assert_eq!(table.pending_unit_count(), 0); + assert!(!table.is_range_pending(0x200_000, 0x600_000)); + assert!(table.is_fully_accepted(0x200_000, 0x600_000)); +} + +#[test] +fn test_builder_incremental() { + let coverage = 0x200_000..0x800_000; + let required_size = EfiUnacceptedMemory::required_size_for_range(coverage.clone()).unwrap(); + let mut buf = vec![0u8; required_size + 16]; + + let align_offset = buf + .as_ptr() + .align_offset(core::mem::align_of::()); + let aligned_slice = &mut buf[align_offset..align_offset + required_size]; + let table_addr = NonNull::new(aligned_slice.as_mut_ptr()).unwrap(); + + let mut builder = + unsafe { EfiUnacceptedMemory::builder(table_addr, aligned_slice.len(), coverage).unwrap() }; + + // Register a range of at least 2 * unit_size (>= 4 MiB) to avoid eager accept in unit test + unsafe { + builder.register_range(0x200_000, 0x600_000).unwrap(); + } + + let table = builder.build(); + assert_eq!(table.pending_unit_count(), 2); + assert!(table.is_range_pending(0x200_000, 0x600_000)); + assert!(!table.is_range_pending(0x600_000, 0x800_000)); +} + +#[test] +fn test_new_invalid_inputs() { + let range = 0x200_000..0x600_000; + let required_size = EfiUnacceptedMemory::required_size(core::slice::from_ref(&range)).unwrap(); + let mut buf = vec![0u8; required_size + 16]; + + let align_offset = buf + .as_ptr() + .align_offset(core::mem::align_of::()); + let aligned_slice = &mut buf[align_offset..align_offset + required_size]; + + // Too small allocation + let table_addr = NonNull::new(aligned_slice.as_mut_ptr()).unwrap(); + let res = unsafe { + EfiUnacceptedMemory::new(table_addr, required_size - 1, core::slice::from_ref(&range)) + }; + assert!(matches!(res, Err(AcceptError::OutOfBounds))); + + // Misaligned pointer + let unaligned_addr = NonNull::new(unsafe { aligned_slice.as_mut_ptr().add(1) }).unwrap(); + let res = unsafe { + EfiUnacceptedMemory::new(unaligned_addr, required_size, core::slice::from_ref(&range)) + }; + assert!(matches!(res, Err(AcceptError::InvalidAlignment))); +} + +#[test] +fn test_bitmap_word_boundaries() { + let words = [const { AtomicU64::new(0) }; 4]; // 256 bits + let ptr = words.as_ptr().cast::(); + let bitmap = unsafe { BitmapRef::from_raw(ptr, words.len() * 8) }; + + assert_eq!(bitmap.capacity(), 256); + assert!(!bitmap.has_set_bit(0, 256)); + assert_eq!(bitmap.find_next_set(0, 256), None); + + // Set a range spanning word 0 and word 1 (e.g. bits 60..70) + bitmap.set_range(60, 70); + assert!(bitmap.has_set_bit(60, 70)); + assert!(bitmap.has_set_bit(0, 65)); + assert!(!bitmap.has_set_bit(0, 60)); + assert!(!bitmap.has_set_bit(70, 256)); + assert_eq!(bitmap.find_next_set(0, 256), Some(60)); + assert_eq!(bitmap.find_next_zero(60, 256), Some(70)); + assert_eq!(bitmap.pending_unit_count(), 10); + + // Check underlying raw word values + // Bits 60..64 in word 0 are 0xF000_0000_0000_0000 + assert_eq!( + bitmap.words()[0].load(core::sync::atomic::Ordering::Relaxed), + 0xF000_0000_0000_0000 + ); + // Bits 64..70 (bits 0..6 of word 1) are 0x3F + assert_eq!( + bitmap.words()[1].load(core::sync::atomic::Ordering::Relaxed), + 0x3F + ); + + // Clear spanning range + bitmap.clear_range(62, 68); + assert_eq!(bitmap.pending_unit_count(), 4); // 60, 61 and 68, 69 + assert_eq!(bitmap.find_next_set(0, 256), Some(60)); + assert_eq!(bitmap.find_next_set(62, 256), Some(68)); + assert_eq!(bitmap.find_next_set(70, 256), None); +} diff --git a/src/ve.rs b/src/ve.rs index 8fa7a23..ded92bc 100644 --- a/src/ve.rs +++ b/src/ve.rs @@ -45,9 +45,7 @@ pub(crate) fn handle_io(trapframe: &mut dyn TdxTrapFrame, ve_info: &TdgVeInfo) - let value = io_read(size, port).unwrap() as usize; match size { IoSize::Size1 => trapframe.set_rax((trapframe.rax() & !0xFF) | (value & 0xFF)), - IoSize::Size2 => { - trapframe.set_rax((trapframe.rax() & !0xFFFF) | (value & 0xFFFF)) - } + IoSize::Size2 => trapframe.set_rax((trapframe.rax() & !0xFFFF) | (value & 0xFFFF)), IoSize::Size4 => trapframe.set_rax(value & 0xFFFF_FFFF), _ => unreachable!(), }