Using BlockIO without DiskIO

pixie-uefi/src/export_cov.rs

@@ -6,5 +6,5 @@ pub async fn export(os: UefiOS) {
     let mut coverage = vec![];
     // SAFETY: we never create threads anyway.
     unsafe { minicov::capture_coverage(&mut coverage).unwrap() };
-    disk.write_(0, &coverage).unwrap();
+    disk.write_sync(0, &coverage).unwrap();
 }

pixie-uefi/src/os/disk.rs

@@ -12,11 +12,11 @@ use gpt_disk_io::{
 };
 use uefi::{
     boot::{OpenProtocolParams, ScopedProtocol},
-    proto::media::{block::BlockIO, disk::DiskIo},
+    proto::media::block::BlockIO,
     Handle,
 };
 
-fn open_disk(handle: Handle) -> Result<(ScopedProtocol<DiskIo>, ScopedProtocol<BlockIO>)> {
+fn open_disk(handle: Handle) -> Result<ScopedProtocol<BlockIO>> {
     let image_handle = uefi::boot::image_handle();
     let bio = unsafe {
         uefi::boot::open_protocol::<BlockIO>(
@@ -28,8 +28,7 @@ fn open_disk(handle: Handle) -> Result<(ScopedProtocol<DiskIo>, ScopedProtocol<B
             uefi::boot::OpenProtocolAttributes::GetProtocol,
         )?
     };
-    let proto = uefi::boot::open_protocol_exclusive::<DiskIo>(handle)?;
-    Ok((proto, bio))
+    Ok(bio)
 }
 
 #[derive(Debug)]
@@ -41,7 +40,6 @@ pub struct DiskPartition {
 }
 
 pub struct Disk {
-    disk: ScopedProtocol<DiskIo>,
     block: ScopedProtocol<BlockIO>,
     os: UefiOS,
 }
@@ -50,12 +48,11 @@ pub struct Disk {
 // available; support having more than one disk.
 impl Disk {
     pub fn new(os: UefiOS) -> Disk {
-        let (_size, handle) = uefi::boot::find_handles::<DiskIo>()
+        let (_size, handle) = uefi::boot::find_handles::<BlockIO>()
             .unwrap()
             .into_iter()
             .filter_map(|handle| {
-                let op = open_disk(handle);
-                let Ok((_, block)) = op else {
+                let Ok(block) = open_disk(handle) else {
                     return None;
                 };
                 let m = block.media();
@@ -68,18 +65,17 @@ impl Disk {
             .max_by_key(|(size, _)| *size)
             .expect("Disk not found");
 
-        let (disk, block) = open_disk(handle).unwrap();
-        Disk { disk, block, os }
+        let block = open_disk(handle).unwrap();
+        Disk { block, os }
     }
 
     #[cfg(feature = "coverage")]
     pub fn open_with_size(os: UefiOS, base_size: i64) -> Disk {
-        let (_size, handle) = uefi::boot::find_handles::<DiskIo>()
+        let (_size, handle) = uefi::boot::find_handles::<BlockIO>()
             .unwrap()
             .into_iter()
             .filter_map(|handle| {
-                let op = open_disk(handle);
-                let Ok((_, block)) = op else {
+                let Ok(block) = open_disk(handle) else {
                     return None;
                 };
                 let m = block.media();
@@ -92,37 +88,84 @@ impl Disk {
             .min_by_key(|(size, _)| *size)
             .expect("Disk not found");
 
-        let (disk, block) = open_disk(handle).unwrap();
-        Disk { disk, block, os }
+        let block = open_disk(handle).unwrap();
+        Disk { block, os }
     }
 
     pub fn size(&self) -> u64 {
         self.block.media().block_size() as u64 * (self.block.media().last_block() + 1)
     }
 
-    pub async fn flush(&mut self) {
-        self.block.flush_blocks().unwrap();
+    pub async fn flush(&mut self) -> Result<()> {
+        self.block.flush_blocks()?;
+        Ok(())
+    }
+
+    pub fn read_sync(&self, offset: u64, buf: &mut [u8]) -> Result<()> {
+        let block_size = self.block.media().block_size() as u64;
+        let media_id = self.block.media().media_id();
+        let start_block = offset / block_size;
+        let end_block = (offset + buf.len() as u64).div_ceil(block_size);
+        let num_blocks = end_block - start_block;
+        if buf.len() as u64 != num_blocks * block_size
+            || !(buf.as_ptr() as usize).is_multiple_of(16)
+        {
+            //log::warn!(
+            //    "Unaligned read: offset {}, block size {}, buf addr {:p}, buf len {}",
+            //    offset,
+            //    block_size,
+            //    buf.as_ptr(),
+            //    buf.len()
+            //);
+            let mut buf2 = vec![0u8; (num_blocks * block_size) as usize + 15];
+            let delta = buf2.as_ptr().align_offset(16);
+            let buf2 = &mut buf2[delta..delta + (num_blocks * block_size) as usize];
+            self.block.read_blocks(media_id, start_block, buf2)?;
+            let start_offset = (offset % block_size) as usize;
+            buf.copy_from_slice(&buf2[start_offset..start_offset + buf.len()]);
+        } else {
+            self.block.read_blocks(media_id, start_block, buf)?;
+        }
+        Ok(())
     }
 
     pub async fn read(&self, offset: u64, buf: &mut [u8]) -> Result<()> {
         self.os.schedule().await;
-        Ok(self
-            .disk
-            .read_disk(self.block.media().media_id(), offset, buf)?)
+        self.read_sync(offset, buf)
     }
 
-    #[cfg(feature = "coverage")]
-    pub fn write_(&mut self, offset: u64, buf: &[u8]) -> Result<()> {
-        Ok(self
-            .disk
-            .write_disk(self.block.media().media_id(), offset, buf)?)
+    pub fn write_sync(&mut self, offset: u64, buf: &[u8]) -> Result<()> {
+        let block_size = self.block.media().block_size() as u64;
+        let media_id = self.block.media().media_id();
+        let start_block = offset / block_size;
+        let end_block = (offset + buf.len() as u64).div_ceil(block_size);
+        let num_blocks = end_block - start_block;
+        if buf.len() as u64 != num_blocks * block_size
+            || !(buf.as_ptr() as usize).is_multiple_of(16)
+        {
+            //log::warn!(
+            //    "Unaligned write: offset {}, block size {}, buf addr {:p}, buf len {}",
+            //    offset,
+            //    block_size,
+            //    buf.as_ptr(),
+            //    buf.len()
+            //);
+            let mut buf2 = vec![0u8; (num_blocks * block_size) as usize + 15];
+            let delta = buf2.as_ptr().align_offset(16);
+            let buf2 = &mut buf2[delta..delta + (num_blocks * block_size) as usize];
+            self.block.read_blocks(media_id, start_block, buf2)?;
+            let start_offset = (offset % block_size) as usize;
+            buf2[start_offset..start_offset + buf.len()].copy_from_slice(buf);
+            self.block.write_blocks(media_id, start_block, buf2)?;
+        } else {
+            self.block.write_blocks(media_id, start_block, buf)?;
+        }
+        Ok(())
     }
 
     pub async fn write(&mut self, offset: u64, buf: &[u8]) -> Result<()> {
         self.os.schedule().await;
-        Ok(self
-            .disk
-            .write_disk(self.block.media().media_id(), offset, buf)?)
+        self.write_sync(offset, buf)
     }
 
     pub fn partitions(&mut self) -> Result<Vec<DiskPartition>> {
@@ -175,17 +218,13 @@ impl gpt_disk_io::BlockIo for &mut Disk {
         Ok(self.block.media().last_block() + 1)
     }
     fn read_blocks(&mut self, start_lba: Lba, dst: &mut [u8]) -> Result<()> {
-        self.disk.read_disk(
-            self.block.media().media_id(),
-            self.block.media().block_size() as u64 * start_lba.0,
-            dst,
-        )?;
-        Ok(())
+        self.read_sync(self.block.media().block_size() as u64 * start_lba.0, dst)
     }
     fn write_blocks(&mut self, _start_lba: Lba, _src: &[u8]) -> Result<()> {
         unreachable!();
     }
     fn flush(&mut self) -> Result<()> {
-        Ok(self.block.flush_blocks()?)
+        // This is a no-op because write_blocks isn't implemented.
+        Ok(())
     }
 }

pixie-uefi/src/os/executor.rs

@@ -1,17 +1,14 @@
 use super::{sync::SyncRefCell, UefiOS};
-use alloc::{boxed::Box, collections::VecDeque, sync::Arc};
-use core::{cell::RefCell, pin::Pin, task::Context};
-use futures::{
-    task::{waker_ref, ArcWake},
-    Future,
+use alloc::{boxed::Box, collections::VecDeque, sync::Arc, task::Wake};
+use core::{
+    cell::RefCell,
+    future::Future,
+    pin::Pin,
+    task::{Context, Waker},
 };
 
 pub(super) type BoxFuture<T = ()> = Pin<Box<dyn Future<Output = T> + 'static>>;
 
-static EXECUTOR: SyncRefCell<Executor> = SyncRefCell::new(Executor {
-    tasks: VecDeque::new(),
-});
-
 struct TaskInner {
     pub in_queue: bool,
     pub future: Option<BoxFuture>,
@@ -47,35 +44,39 @@ impl Task {
 unsafe impl Send for Task {}
 unsafe impl Sync for Task {}
 
-impl ArcWake for Task {
-    fn wake_by_ref(task: &Arc<Self>) {
-        if !task.inner.borrow().in_queue {
-            task.inner.borrow_mut().in_queue = true;
-            EXECUTOR.borrow_mut().tasks.push_back(task.clone());
+impl Wake for Task {
+    fn wake(self: Arc<Self>) {
+        if !self.inner.borrow().in_queue {
+            self.inner.borrow_mut().in_queue = true;
+            EXECUTOR.borrow_mut().ready_tasks.push_back(self);
         }
     }
 }
 
+static EXECUTOR: SyncRefCell<Executor> = SyncRefCell::new(Executor {
+    ready_tasks: VecDeque::new(),
+});
+
 pub struct Executor {
     // TODO(veluca): scheduling.
-    tasks: VecDeque<Arc<Task>>,
+    ready_tasks: VecDeque<Arc<Task>>,
 }
 
 impl Executor {
     pub fn run(os: UefiOS) -> ! {
         loop {
             let task = EXECUTOR
                 .borrow_mut()
-                .tasks
+                .ready_tasks
                 .pop_front()
-                .expect("Executor should never run out of tasks");
+                .expect("Executor should never run out of ready tasks");
 
             task.inner.borrow_mut().in_queue = false;
-            let waker = waker_ref(&task);
-            let context = &mut Context::from_waker(&waker);
+            let waker = Waker::from(task.clone());
+            let mut context = Context::from_waker(&waker);
             let mut fut = task.inner.borrow_mut().future.take().unwrap();
             let begin = os.timer().micros();
-            let status = fut.as_mut().poll(context);
+            let status = fut.as_mut().poll(&mut context);
             let end = os.timer().micros();
             task.inner.borrow_mut().micros += end - begin;
             if status.is_pending() {
@@ -85,6 +86,6 @@ impl Executor {
     }
 
     pub(super) fn spawn(task: Arc<Task>) {
-        EXECUTOR.borrow_mut().tasks.push_back(task);
+        EXECUTOR.borrow_mut().ready_tasks.push_back(task);
     }
 }

pixie-uefi/src/os/mod.rs

@@ -19,10 +19,10 @@ use core::{
     cell::{Ref, RefMut},
     ffi::c_void,
     fmt::Write,
-    future::{poll_fn, Future, PollFn},
+    future::{poll_fn, Future},
     net::SocketAddrV4,
     ptr::NonNull,
-    task::{Context, Poll},
+    task::Poll,
 };
 use pixie_shared::util::BytesFmt;
 use uefi::{
@@ -290,7 +290,7 @@ impl UefiOS {
         RefMut::map(self.borrow_mut(), |f| f.net.as_mut().unwrap())
     }
 
-    pub fn wait_for_ip(self) -> PollFn<impl FnMut(&mut Context<'_>) -> Poll<()>> {
+    pub fn wait_for_ip(self) -> impl Future<Output = ()> {
         poll_fn(move |cx| {
             if self.net().has_ip() {
                 Poll::Ready(())
@@ -303,7 +303,7 @@ impl UefiOS {
 
     /// Interrupt task execution.
     /// This is useful to yield the CPU to other tasks.
-    pub fn schedule(&self) -> PollFn<impl FnMut(&mut Context<'_>) -> Poll<()>> {
+    pub fn schedule(&self) -> impl Future<Output = ()> {
         let mut ready = false;
         poll_fn(move |cx| {
             if ready {
@@ -316,7 +316,7 @@ impl UefiOS {
         })
     }
 
-    pub fn sleep_us(self, us: u64) -> PollFn<impl FnMut(&mut Context<'_>) -> Poll<()>> {
+    pub fn sleep_us(self, us: u64) -> impl Future<Output = ()> {
         let tgt = self.timer().micros() as u64 + us;
         poll_fn(move |cx| {
             let now = self.timer().micros() as u64;
@@ -409,11 +409,11 @@ impl UefiOS {
         UdpHandle::new(*self, port).await
     }
 
-    pub async fn read_key(&self) -> Result<Key> {
-        Ok(poll_fn(move |cx| {
+    pub fn read_key(&self) -> impl Future<Output = Result<Key>> + '_ {
+        poll_fn(move |cx| {
             let key = self.borrow_mut().input.read_key();
             if let Err(e) = key {
-                return Poll::Ready(Err(e));
+                return Poll::Ready(Err(e.into()));
             }
             let key = key.unwrap();
             if let Some(key) = key {
@@ -422,7 +422,6 @@ impl UefiOS {
             cx.waker().wake_by_ref();
             Poll::Pending
         })
-        .await?)
     }
 
     pub fn write_with_color(&self, msg: &str, fg: Color, bg: Color) {

pixie-uefi/src/parse_disk/gpt.rs

@@ -8,8 +8,12 @@ use pixie_shared::util::BytesFmt;
 
 pub async fn parse_gpt(disk: &mut Disk) -> Result<Option<Vec<ChunkInfo>>> {
     let disk_size = disk.size() as usize;
-    let Ok(partitions) = disk.partitions() else {
-        return Ok(None);
+    let partitions = match disk.partitions() {
+        Ok(partitions) => partitions,
+        Err(e) => {
+            log::debug!("Failed to parse GPT partitions: {e:?}");
+            return Ok(None);
+        }
     };
 
     let mut pos = 0usize;