Have PhysicalBuffer just speak pages

Author: jdreaver

README.md

@@ -93,6 +93,7 @@ make test
 
 - VirtIO improvements:
   - Abstract and improve request/response queue from rng
+  - Ensure we zero out descriptors after using them so we don't accidentally reuse the buffer
   - Allocate buffers on the fly instead of using static buffers, or have client provide buffer
     - Reuse them if we are reusing a descriptor, or free them.
     - If the VirtIO device always owns the memory it uses, it is easier to reason about alloc/free, ensure buffers are the right size, and ensure provided buffers are physically contiguous
@@ -104,8 +105,10 @@ make test
   - Ensure we don't accidentally reuse descriptors while we are waiting for a response from the device. Don't automatically just wrap around! This is what might require a mutex rather than just atomic integers?
   - I think there is a race condition with the interrupts with the current non-locking mechanism. Ensure that if there are concurrent writes while an interrupt, then an interrupt won't miss a read (e.g. there will at least be a followup interrupt)
   - Remember features we negotiate, and ensure we are accounting for the different features in the logic (especially around notifications)
-- Ensure `PhysicalBuffer` allocation and deallocation is safe, and that page and size math is correct (particularly in `drop()`). Make it foolproof.
-  - Perhaps introduce page <-> address translation layer in a single spot.
+- Physical memory allocation
+  - Ensure `PhysicalBuffer` allocation and deallocation is safe, and that page and size math is correct (particularly in `drop()`). Make it foolproof.
+    - Be explicit about it allocating pages. Don't pretend it is just addresses. Perhaps put page <-> address translation layer in a single spot.
+    - In fact, maybe `LockedPhysicalMemoryAllocator` should _not_ implement `Allocator`. If we want a malloc-like behavior on top of contiguous physical memory, we should explicitly use e.g. `LockedHeap` on top of some allocated physical pages (basically an arena).
 - Does `PhysicalBuffer::allocate` need an alignment argument? Does anything ever need more than page alignment?
 - Create `sync` module that has synchronization primitives
   - Wrapper around `spin::Mutex` that also disables interrupts, similar to Linux's `spin_lock_irqsave` (`x86_64::interrupts::without_interrupts` is handy here). Might need our own custom `MutexGuard` wrapper that handles re-enabling interrupts on `drop()`

kernel/src/memory.rs

@@ -1,4 +1,4 @@
-use core::alloc::{AllocError, Allocator, Layout, LayoutError};
+use core::alloc::{AllocError, Allocator, Layout};
 use core::ptr;
 use core::ptr::NonNull;
 
@@ -252,73 +252,74 @@ unsafe impl<S: PageSize> FrameAllocator<S> for LockedPhysicalMemoryAllocator<'_>
     }
 }
 
-/// Error type used in `allocate_zeroed_buffer`.
-#[derive(Debug, Clone)]
-pub(crate) enum AllocZeroedBufferError {
-    LayoutError(LayoutError),
-    AllocError(AllocError),
-}
-
 /// Physically contiguous buffer of memory. Allocates by page, so it can
 /// allocate more memory than requested. Useful for e.g. Direct Memory Access
 /// (DMA) like with VirtIO buffers.
+///
+/// NOTE: This type implements `Drop` and will free the allocated memory when
+/// it goes out of scope.
 #[derive(Debug)]
 pub(crate) struct PhysicalBuffer {
-    ptr: NonNull<[u8]>,
+    start_page: usize,
+    num_pages: usize,
 }
 
 impl PhysicalBuffer {
-    pub(crate) fn allocate(
-        len_bytes: usize,
-        alignment: usize,
-    ) -> Result<Self, AllocZeroedBufferError> {
-        let layout = Layout::from_size_align(len_bytes, alignment)
-            .map_err(AllocZeroedBufferError::LayoutError)?;
-        let ptr = KERNEL_PHYSICAL_ALLOCATOR
-            .allocate_zeroed(layout)
-            .map_err(AllocZeroedBufferError::AllocError)?;
-        Ok(Self { ptr })
+    // Don't need to expose this b/c allocate_zeroed is safer.
+    fn allocate(min_bytes: usize) -> Result<Self, AllocError> {
+        let num_pages = min_bytes.div_ceil(PhysicalMemoryAllocator::PAGE_SIZE);
+        let start_page = KERNEL_PHYSICAL_ALLOCATOR.with_lock(|allocator| {
+            allocator
+                .allocator
+                .allocate_contiguous(num_pages)
+                .ok_or(AllocError)
+        })?;
+        Ok(Self {
+            start_page,
+            num_pages,
+        })
     }
 
-    pub(crate) fn allocate_value<T>(val: T) -> Result<Self, AllocZeroedBufferError> {
-        let layout = Layout::new::<T>();
-        let ptr = KERNEL_PHYSICAL_ALLOCATOR
-            .allocate_zeroed(layout)
-            .map_err(AllocZeroedBufferError::AllocError)?;
+    pub(crate) fn allocate_zeroed(min_bytes: usize) -> Result<Self, AllocError> {
+        let buffer = Self::allocate(min_bytes)?;
+        let ptr = buffer.address().as_u64() as *mut u8;
         unsafe {
-            ptr::write_volatile(ptr.as_ptr().cast::<T>(), val);
+            ptr::write_bytes(ptr, 0, buffer.len_bytes());
         }
-        Ok(Self { ptr })
+        Ok(buffer)
     }
 
-    /// Consumes the buffer and returns the underlying physical address and
-    /// length in bytes. NOTE: It is up to the caller to free this memory,
-    /// ideally by constructing a new buffer with
-    /// `PhysicalBuffer::from_raw_parts` and letting that `Drop`.
-    pub(crate) fn into_raw_parts(self) -> (PhysAddr, usize) {
-        let buf = core::mem::ManuallyDrop::new(self);
-        let addr = PhysAddr::new(buf.ptr.addr().get() as u64);
-        let len_bytes = buf.ptr.len();
-        (addr, len_bytes)
+    pub(crate) fn address(&self) -> PhysAddr {
+        PhysAddr::new(self.start_page as u64 * PhysicalMemoryAllocator::PAGE_SIZE as u64)
     }
 
-    pub(crate) unsafe fn from_raw_parts(addr: PhysAddr, len_bytes: usize) -> Self {
-        let ptr = unsafe { nonnull_ptr_slice_from_addr_len(addr.as_u64() as usize, len_bytes) };
-        Self { ptr }
+    pub(crate) fn len_bytes(&self) -> usize {
+        self.num_pages * PhysicalMemoryAllocator::PAGE_SIZE
+    }
+
+    pub(crate) unsafe fn write_offset<T>(&mut self, offset: usize, val: T) {
+        let buffer_len = self.len_bytes();
+        assert!(
+            offset < self.len_bytes(),
+            "tried to write value at offset {offset} but buffer only has {buffer_len} bytes"
+        );
+        let ptr = (self.address().as_u64() + offset as u64) as *mut T;
+        ptr::write_volatile(ptr, val);
+    }
+
+    // TODO: Don't allow leaking. We are only doing this temporarily.
+    pub(crate) fn leak(self) -> PhysAddr {
+        let buf = core::mem::ManuallyDrop::new(self);
+        buf.address()
     }
 }
 
 impl Drop for PhysicalBuffer {
     fn drop(&mut self) {
-        // TODO: Is this correct? DRY with the `deallocate`, and perhaps add
-        // some types to ensure that we are converting to pages correctly. Also
-        // ensure that we do indeed "own" the entire page we are de-allocating.
-        let layout = unsafe {
-            Layout::from_size_align_unchecked(self.ptr.len(), PhysicalMemoryAllocator::PAGE_SIZE)
-        };
-        let u8_ptr = self.ptr.cast::<u8>();
-        unsafe {
-            KERNEL_PHYSICAL_ALLOCATOR.deallocate(u8_ptr, layout);
-        };
+        KERNEL_PHYSICAL_ALLOCATOR.with_lock(|allocator| {
+            allocator
+                .allocator
+                .free_contiguous(self.start_page, self.num_pages);
+        });
     }
 }

kernel/src/virtio/block.rs

@@ -58,13 +58,15 @@ pub(crate) fn virtio_block_get_id(device_index: usize) {
         .get_virtqueue(VirtQueueIndex(0))
         .unwrap();
 
-    let (data_addr, _) = PhysicalBuffer::allocate(BlockRequest::GET_ID_DATA_LEN as usize, 8)
+    let data_addr = PhysicalBuffer::allocate_zeroed(BlockRequest::GET_ID_DATA_LEN as usize)
         .expect("failed to allocate GetID buffer")
-        .into_raw_parts();
+        // TODO: Don't leak here
+        .leak();
 
     let request = BlockRequest::GetID { data_addr };
     let raw_request = request.to_raw();
-    virtq.add_buffer(&raw_request.to_descriptor_chain());
+    let desc_chain = raw_request.to_descriptor_chain();
+    virtq.add_buffer(&desc_chain);
     virtq.notify_device();
 }
 
@@ -375,17 +377,26 @@ impl RawBlockRequest {
         // TODO: Ensure all of these descriptor components are dropped!
 
         // Allocate header
-        let header_buffer = PhysicalBuffer::allocate_value(RawBlockRequestHeader {
-            request_type: self.request_type as u32,
-            reserved: 0,
-            sector: self.sector,
-        })
-        .expect("failed to allocate RawBlockRequest header");
-        let header_desc = ChainedVirtQueueDescriptorElem::from_buffer(
-            header_buffer,
-            core::mem::size_of::<RawBlockRequestHeader>() as u32,
-            VirtQueueDescriptorFlags::new().with_device_write(false),
-        );
+        let header_size = core::mem::size_of::<RawBlockRequestHeader>() as u32;
+        let mut header_buffer = PhysicalBuffer::allocate_zeroed(header_size as usize)
+            .expect("failed to allocate RawBlockRequest header");
+        unsafe {
+            header_buffer.write_offset(
+                0,
+                RawBlockRequestHeader {
+                    request_type: self.request_type as u32,
+                    reserved: 0,
+                    sector: self.sector,
+                },
+            );
+        };
+
+        let header_desc = ChainedVirtQueueDescriptorElem {
+            // TODO: Don't leak
+            addr: header_buffer.leak(),
+            len: header_size,
+            flags: VirtQueueDescriptorFlags::new().with_device_write(false),
+        };
 
         // Buffer descriptor
         //
@@ -401,13 +412,19 @@ impl RawBlockRequest {
         // TODO: This is a waste of an entire page! We should allocate this
         // along with the header. We should probably also just use a slab
         // allocator or a sort of physically contiguous heap for this.
-        let status_buffer = PhysicalBuffer::allocate_value(self.status as u8)
+        let status_size = core::mem::size_of::<BlockRequestStatus>() as u32;
+        let mut status_buffer = PhysicalBuffer::allocate_zeroed(status_size as usize)
             .expect("failed to allocate RawBlockRequest status");
-        let status_desc = ChainedVirtQueueDescriptorElem::from_buffer(
-            status_buffer,
-            core::mem::size_of::<u8>() as u32,
-            VirtQueueDescriptorFlags::new().with_device_write(true),
-        );
+        unsafe {
+            status_buffer.write_offset(0, self.status);
+        };
+
+        let status_desc = ChainedVirtQueueDescriptorElem {
+            // TODO: Don't leak
+            addr: status_buffer.leak(),
+            len: status_size,
+            flags: VirtQueueDescriptorFlags::new().with_device_write(true),
+        };
 
         [header_desc, buffer_desc, status_desc]
     }