Move the SharedAllocator to allocator.rs and delete sev.rs.

Everything has moved our of `sev.rs` to either the HAL or some other
generic place, so there's nothing left in there.

This should conclude the main points of HALification of stage0. There is
some clean-up work still to be done as there are unconditional
dependencies on some SEV-related crates, especially for data structures
that should be genericised, but that's not going to block the TDX work.

Bug: 350496083
Change-Id: I9a0204e8effe18522f4b29b8003a56f13110b5de

stage0/src/allocator.rs

@@ -14,16 +14,23 @@
 // limitations under the License.
 //
 
+use alloc::boxed::Box;
 use core::{
     alloc::{AllocError, Allocator, Layout},
     mem::MaybeUninit,
+    ops::{Deref, DerefMut},
     ptr::NonNull,
     sync::atomic::{AtomicUsize, Ordering},
 };
 
 use oak_linux_boot_params::{BootE820Entry, E820EntryType};
 use spinning_top::Spinlock;
-use x86_64::VirtAddr;
+use x86_64::{
+    structures::paging::{Page, PageSize, Size4KiB},
+    VirtAddr,
+};
+
+use crate::paging::{share_page, unshare_page};
 
 struct Inner<const N: usize> {
     index: AtomicUsize,
@@ -93,6 +100,116 @@ unsafe impl<const N: usize> Allocator for BumpAllocator<N> {
         // Bump allocator never deallocates.
     }
 }
+/// Allocator that forces allocations to be 4K-aligned (and sized) and marks the
+/// pages as shared.
+///
+/// This allocator is inefficient as it will only allocate 4K chunks,
+/// potentially wasting memory. For example, if you allocate two u32-s, although
+/// they could well fit on one page, currently that'd use 8K of memory.
+/// That, however, is an implementation detail, and may change in the future.
+#[repr(transparent)]
+struct SharedAllocator<A: Allocator> {
+    inner: A,
+}
+
+impl<A: Allocator> SharedAllocator<A> {
+    fn new(allocator: A) -> Self {
+        Self { inner: allocator }
+    }
+}
+
+unsafe impl<A: Allocator> Allocator for SharedAllocator<A> {
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        let layout =
+            layout.align_to(Size4KiB::SIZE as usize).map_err(|_| AllocError)?.pad_to_align();
+        let allocation = self.inner.allocate(layout)?;
+        for offset in (0..allocation.len()).step_by(Size4KiB::SIZE as usize) {
+            // Safety: the allocation has succeeded and the offset won't exceed the size of
+            // the allocation.
+            share_page(Page::containing_address(VirtAddr::from_ptr(unsafe {
+                allocation.as_non_null_ptr().as_ptr().add(offset)
+            })))
+        }
+        Ok(allocation)
+    }
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        let layout = layout
+            .align_to(Size4KiB::SIZE as usize)
+            .map_err(|_| AllocError)
+            .unwrap()
+            .pad_to_align();
+        for offset in (0..layout.size()).step_by(Size4KiB::SIZE as usize) {
+            // Safety: the allocation has succeeded and the offset won't exceed the size of
+            // the allocation.
+            unshare_page(Page::containing_address(VirtAddr::from_ptr(unsafe {
+                ptr.as_ptr().add(offset)
+            })))
+        }
+        self.inner.deallocate(ptr, layout)
+    }
+}
+
+/// Stores a data structure on a shared page.
+pub struct Shared<T: 'static, A: Allocator> {
+    inner: Box<T, SharedAllocator<A>>,
+}
+
+impl<T, A: Allocator> Shared<T, A> {
+    pub fn new_in(t: T, alloc: A) -> Self
+    where
+        A: 'static,
+    {
+        Self { inner: Box::new_in(t, SharedAllocator::new(alloc)) }
+    }
+
+    /// See `Box::from_raw_in` for documentation.
+    ///
+    /// # Safety
+    ///
+    /// The caller needs to guarantee that (a) the pointer was obtained by
+    /// `Shared::leak` and (b) the allocator you pass in is exactly the same as
+    /// was used for the original allocation of the `Shared`.
+    ///
+    /// Again, see `Box::from_raw_in` for more details.
+    pub unsafe fn from_raw_in(raw: *mut T, alloc: A) -> Shared<T, A> {
+        Self { inner: Box::from_raw_in(raw, SharedAllocator::new(alloc)) }
+    }
+
+    /// See `Box::leak` for documentation.
+    pub fn leak<'a>(s: Shared<T, A>) -> &'a mut T
+    where
+        A: 'a,
+    {
+        Box::leak(s.inner)
+    }
+}
+
+impl<T, A: Allocator> Deref for Shared<T, A> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl<T, A: Allocator> DerefMut for Shared<T, A> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.inner
+    }
+}
+
+impl<T, A: Allocator> AsRef<T> for Shared<T, A> {
+    fn as_ref(&self) -> &T {
+        &self.inner
+    }
+}
+
+impl<T, A: Allocator> AsMut<T> for Shared<T, A> {
+    fn as_mut(&mut self) -> &mut T {
+        &mut self.inner
+    }
+}
 
 pub fn init_global_allocator(e820_table: &[BootE820Entry]) {
     // Create the heap between 1MiB and 2MiB.

stage0/src/fw_cfg.rs

@@ -29,7 +29,7 @@ use x86_64::{
 };
 use zerocopy::{AsBytes, FromBytes, FromZeroes};
 
-use crate::{hal::Port, sev::Shared, BootAllocator};
+use crate::{allocator::Shared, hal::Port, BootAllocator};
 
 // See https://www.qemu.org/docs/master/specs/fw_cfg.html for documentation about the various data structures and constants.
 const FWCFG_PORT_SELECTOR: u16 = 0x510;

stage0/src/hal/sev/dice_attestation.rs

@@ -30,7 +30,7 @@ use zerocopy::{AsBytes, FromBytes};
 use zeroize::Zeroize;
 
 use super::{GHCB_WRAPPER, SEV_SECRETS};
-use crate::sev::Shared;
+use crate::allocator::Shared;
 
 /// Cryptographic helper to encrypt and decrypt messages for the GHCB guest
 /// message protocol.

stage0/src/hal/sev/mod.rs

@@ -37,7 +37,9 @@ pub use port::*;
 use spinning_top::{lock_api::MutexGuard, RawSpinlock, Spinlock};
 use x86_64::{PhysAddr, VirtAddr};
 
-use crate::{paging::PageEncryption, sev::Shared, zero_page::ZeroPage, BootAllocator, BOOT_ALLOC};
+use crate::{
+    allocator::Shared, paging::PageEncryption, zero_page::ZeroPage, BootAllocator, BOOT_ALLOC,
+};
 
 #[link_section = ".boot"]
 #[no_mangle]

stage0/src/lib.rs

@@ -57,7 +57,6 @@ mod logging;
 mod msr;
 pub mod paging;
 mod pic;
-mod sev;
 mod smp;
 mod zero_page;