From 163dfad9074632540152d1b309cfa2ba052da820 Mon Sep 17 00:00:00 2001 From: Filippo Casarin Date: Fri, 14 Nov 2025 02:30:43 +0100 Subject: [PATCH] Using BlockIO without DiskIO --- pixie-uefi/src/export_cov.rs | 2 +- pixie-uefi/src/os/disk.rs | 109 +++++++++++++++++++++---------- pixie-uefi/src/os/executor.rs | 43 ++++++------ pixie-uefi/src/os/mod.rs | 17 +++-- pixie-uefi/src/parse_disk/gpt.rs | 8 ++- 5 files changed, 111 insertions(+), 68 deletions(-) diff --git a/pixie-uefi/src/export_cov.rs b/pixie-uefi/src/export_cov.rs index b665af8d..e02bc46d 100644 --- a/pixie-uefi/src/export_cov.rs +++ b/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(); } diff --git a/pixie-uefi/src/os/disk.rs b/pixie-uefi/src/os/disk.rs index 76b892eb..6185a551 100644 --- a/pixie-uefi/src/os/disk.rs +++ b/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, ScopedProtocol)> { +fn open_disk(handle: Handle) -> Result> { let image_handle = uefi::boot::image_handle(); let bio = unsafe { uefi::boot::open_protocol::( @@ -28,8 +28,7 @@ fn open_disk(handle: Handle) -> Result<(ScopedProtocol, ScopedProtocol(handle)?; - Ok((proto, bio)) + Ok(bio) } #[derive(Debug)] @@ -41,7 +40,6 @@ pub struct DiskPartition { } pub struct Disk { - disk: ScopedProtocol, block: ScopedProtocol, 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::() + let (_size, handle) = uefi::boot::find_handles::() .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::() + let (_size, handle) = uefi::boot::find_handles::() .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> { @@ -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(()) } } diff --git a/pixie-uefi/src/os/executor.rs b/pixie-uefi/src/os/executor.rs index 20ae0408..b4a782a6 100644 --- a/pixie-uefi/src/os/executor.rs +++ b/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 = Pin + 'static>>; -static EXECUTOR: SyncRefCell = SyncRefCell::new(Executor { - tasks: VecDeque::new(), -}); - struct TaskInner { pub in_queue: bool, pub future: Option, @@ -47,18 +44,22 @@ impl Task { unsafe impl Send for Task {} unsafe impl Sync for Task {} -impl ArcWake for Task { - fn wake_by_ref(task: &Arc) { - 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) { + if !self.inner.borrow().in_queue { + self.inner.borrow_mut().in_queue = true; + EXECUTOR.borrow_mut().ready_tasks.push_back(self); } } } +static EXECUTOR: SyncRefCell = SyncRefCell::new(Executor { + ready_tasks: VecDeque::new(), +}); + pub struct Executor { // TODO(veluca): scheduling. - tasks: VecDeque>, + ready_tasks: VecDeque>, } impl Executor { @@ -66,16 +67,16 @@ impl Executor { 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) { - EXECUTOR.borrow_mut().tasks.push_back(task); + EXECUTOR.borrow_mut().ready_tasks.push_back(task); } } diff --git a/pixie-uefi/src/os/mod.rs b/pixie-uefi/src/os/mod.rs index ddc8e294..27af756e 100644 --- a/pixie-uefi/src/os/mod.rs +++ b/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) -> Poll<()>> { + pub fn wait_for_ip(self) -> impl Future { 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) -> Poll<()>> { + pub fn schedule(&self) -> impl Future { let mut ready = false; poll_fn(move |cx| { if ready { @@ -316,7 +316,7 @@ impl UefiOS { }) } - pub fn sleep_us(self, us: u64) -> PollFn) -> Poll<()>> { + pub fn sleep_us(self, us: u64) -> impl Future { 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 { - Ok(poll_fn(move |cx| { + pub fn read_key(&self) -> impl Future> + '_ { + 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) { diff --git a/pixie-uefi/src/parse_disk/gpt.rs b/pixie-uefi/src/parse_disk/gpt.rs index 123392a0..d6bcffd1 100644 --- a/pixie-uefi/src/parse_disk/gpt.rs +++ b/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>> { 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;