diff --git a/stage0/src/allocator.rs b/stage0/src/allocator.rs index c1d5e176a2..743b4717c2 100644 --- a/stage0/src/allocator.rs +++ b/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 { index: AtomicUsize, @@ -93,6 +100,116 @@ unsafe impl Allocator for BumpAllocator { // 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 { + inner: A, +} + +impl SharedAllocator { + fn new(allocator: A) -> Self { + Self { inner: allocator } + } +} + +unsafe impl Allocator for SharedAllocator { + fn allocate(&self, layout: Layout) -> Result, 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, 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 { + inner: Box>, +} + +impl Shared { + 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 { + Self { inner: Box::from_raw_in(raw, SharedAllocator::new(alloc)) } + } + + /// See `Box::leak` for documentation. + pub fn leak<'a>(s: Shared) -> &'a mut T + where + A: 'a, + { + Box::leak(s.inner) + } +} + +impl Deref for Shared { + type Target = T; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +impl DerefMut for Shared { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.inner + } +} + +impl AsRef for Shared { + fn as_ref(&self) -> &T { + &self.inner + } +} + +impl AsMut for Shared { + 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. diff --git a/stage0/src/fw_cfg.rs b/stage0/src/fw_cfg.rs index b4d6fabdb9..1dc600825b 100644 --- a/stage0/src/fw_cfg.rs +++ b/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; diff --git a/stage0/src/hal/sev/dice_attestation.rs b/stage0/src/hal/sev/dice_attestation.rs index 6ee76fe4f4..8db437ec16 100644 --- a/stage0/src/hal/sev/dice_attestation.rs +++ b/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. diff --git a/stage0/src/hal/sev/mod.rs b/stage0/src/hal/sev/mod.rs index 3e15e7b5ed..1f61f87ee7 100644 --- a/stage0/src/hal/sev/mod.rs +++ b/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] diff --git a/stage0/src/lib.rs b/stage0/src/lib.rs index fe19f90db4..9300a61eee 100644 --- a/stage0/src/lib.rs +++ b/stage0/src/lib.rs @@ -57,7 +57,6 @@ mod logging; mod msr; pub mod paging; mod pic; -mod sev; mod smp; mod zero_page; diff --git a/stage0/src/sev.rs b/stage0/src/sev.rs deleted file mode 100644 index dc3a3d1a00..0000000000 --- a/stage0/src/sev.rs +++ /dev/null @@ -1,140 +0,0 @@ -// -// Copyright 2022 The Project Oak Authors -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -use alloc::boxed::Box; -use core::{ - alloc::{AllocError, Allocator, Layout}, - ops::{Deref, DerefMut}, - ptr::NonNull, -}; - -use x86_64::{ - structures::paging::{Page, PageSize, Size4KiB}, - VirtAddr, -}; - -use crate::paging::{share_page, unshare_page}; - -/// 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 { - inner: A, -} - -impl SharedAllocator { - fn new(allocator: A) -> Self { - Self { inner: allocator } - } -} - -unsafe impl Allocator for SharedAllocator { - fn allocate(&self, layout: Layout) -> Result, 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, 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 { - inner: Box>, -} - -impl Shared { - 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 { - Self { inner: Box::from_raw_in(raw, SharedAllocator::new(alloc)) } - } - - /// See `Box::leak` for documentation. - pub fn leak<'a>(s: Shared) -> &'a mut T - where - A: 'a, - { - Box::leak(s.inner) - } -} - -impl Deref for Shared { - type Target = T; - - fn deref(&self) -> &Self::Target { - &self.inner - } -} - -impl DerefMut for Shared { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.inner - } -} - -impl AsRef for Shared { - fn as_ref(&self) -> &T { - &self.inner - } -} - -impl AsMut for Shared { - fn as_mut(&mut self) -> &mut T { - &mut self.inner - } -}