Skip to content

Commit 4c64893

Browse files
Virv12Virv12
committed
Using BlockIO without DiskIO
1 parent 2b91415 commit 4c64893

5 files changed

Lines changed: 111 additions & 68 deletions

File tree

‎pixie-uefi/src/export_cov.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,5 +6,5 @@ pub async fn export(os: UefiOS) {
66
let mut coverage = vec![];
77
// SAFETY: we never create threads anyway.
88
unsafe { minicov::capture_coverage(&mut coverage).unwrap() };
9-
disk.write_(0, &coverage).unwrap();
9+
disk.write_sync(0, &coverage).unwrap();
1010
}

‎pixie-uefi/src/os/disk.rs‎

Lines changed: 74 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,11 @@ use gpt_disk_io::{
1212
};
1313
use uefi::{
1414
boot::{OpenProtocolParams, ScopedProtocol},
15-
proto::media::{block::BlockIO, disk::DiskIo},
15+
proto::media::block::BlockIO,
1616
Handle,
1717
};
1818

19-
fn open_disk(handle: Handle) -> Result<(ScopedProtocol<DiskIo>, ScopedProtocol<BlockIO>)> {
19+
fn open_disk(handle: Handle) -> Result<ScopedProtocol<BlockIO>> {
2020
let image_handle = uefi::boot::image_handle();
2121
let bio = unsafe {
2222
uefi::boot::open_protocol::<BlockIO>(
@@ -28,8 +28,7 @@ fn open_disk(handle: Handle) -> Result<(ScopedProtocol<DiskIo>, ScopedProtocol<B
2828
uefi::boot::OpenProtocolAttributes::GetProtocol,
2929
)?
3030
};
31-
let proto = uefi::boot::open_protocol_exclusive::<DiskIo>(handle)?;
32-
Ok((proto, bio))
31+
Ok(bio)
3332
}
3433

3534
#[derive(Debug)]
@@ -41,7 +40,6 @@ pub struct DiskPartition {
4140
}
4241

4342
pub struct Disk {
44-
disk: ScopedProtocol<DiskIo>,
4543
block: ScopedProtocol<BlockIO>,
4644
os: UefiOS,
4745
}
@@ -50,12 +48,11 @@ pub struct Disk {
5048
// available; support having more than one disk.
5149
impl Disk {
5250
pub fn new(os: UefiOS) -> Disk {
53-
let (_size, handle) = uefi::boot::find_handles::<DiskIo>()
51+
let (_size, handle) = uefi::boot::find_handles::<BlockIO>()
5452
.unwrap()
5553
.into_iter()
5654
.filter_map(|handle| {
57-
let op = open_disk(handle);
58-
let Ok((_, block)) = op else {
55+
let Ok(block) = open_disk(handle) else {
5956
return None;
6057
};
6158
let m = block.media();
@@ -68,18 +65,17 @@ impl Disk {
6865
.max_by_key(|(size, _)| *size)
6966
.expect("Disk not found");
7067

71-
let (disk, block) = open_disk(handle).unwrap();
72-
Disk { disk, block, os }
68+
let block = open_disk(handle).unwrap();
69+
Disk { block, os }
7370
}
7471

7572
#[cfg(feature = "coverage")]
7673
pub fn open_with_size(os: UefiOS, base_size: i64) -> Disk {
77-
let (_size, handle) = uefi::boot::find_handles::<DiskIo>()
74+
let (_size, handle) = uefi::boot::find_handles::<BlockIO>()
7875
.unwrap()
7976
.into_iter()
8077
.filter_map(|handle| {
81-
let op = open_disk(handle);
82-
let Ok((_, block)) = op else {
78+
let Ok(block) = open_disk(handle) else {
8379
return None;
8480
};
8581
let m = block.media();
@@ -92,37 +88,84 @@ impl Disk {
9288
.min_by_key(|(size, _)| *size)
9389
.expect("Disk not found");
9490

95-
let (disk, block) = open_disk(handle).unwrap();
96-
Disk { disk, block, os }
91+
let block = open_disk(handle).unwrap();
92+
Disk { block, os }
9793
}
9894

9995
pub fn size(&self) -> u64 {
10096
self.block.media().block_size() as u64 * (self.block.media().last_block() + 1)
10197
}
10298

103-
pub async fn flush(&mut self) {
104-
self.block.flush_blocks().unwrap();
99+
pub async fn flush(&mut self) -> Result<()> {
100+
self.block.flush_blocks()?;
101+
Ok(())
102+
}
103+
104+
pub fn read_sync(&self, offset: u64, buf: &mut [u8]) -> Result<()> {
105+
let block_size = self.block.media().block_size() as u64;
106+
let media_id = self.block.media().media_id();
107+
let start_block = offset / block_size;
108+
let end_block = (offset + buf.len() as u64).div_ceil(block_size);
109+
let num_blocks = end_block - start_block;
110+
if buf.len() as u64 != num_blocks * block_size
111+
|| !(buf.as_ptr() as usize).is_multiple_of(16)
112+
{
113+
//log::warn!(
114+
// "Unaligned read: offset {}, block size {}, buf addr {:p}, buf len {}",
115+
// offset,
116+
// block_size,
117+
// buf.as_ptr(),
118+
// buf.len()
119+
//);
120+
let mut buf2 = vec![0u8; (num_blocks * block_size) as usize + 15];
121+
let delta = buf2.as_ptr().align_offset(16);
122+
let buf2 = &mut buf2[delta..delta + (num_blocks * block_size) as usize];
123+
self.block.read_blocks(media_id, start_block, buf2)?;
124+
let start_offset = (offset % block_size) as usize;
125+
buf.copy_from_slice(&buf2[start_offset..start_offset + buf.len()]);
126+
} else {
127+
self.block.read_blocks(media_id, start_block, buf)?;
128+
}
129+
Ok(())
105130
}
106131

107132
pub async fn read(&self, offset: u64, buf: &mut [u8]) -> Result<()> {
108133
self.os.schedule().await;
109-
Ok(self
110-
.disk
111-
.read_disk(self.block.media().media_id(), offset, buf)?)
134+
self.read_sync(offset, buf)
112135
}
113136

114-
#[cfg(feature = "coverage")]
115-
pub fn write_(&mut self, offset: u64, buf: &[u8]) -> Result<()> {
116-
Ok(self
117-
.disk
118-
.write_disk(self.block.media().media_id(), offset, buf)?)
137+
pub fn write_sync(&mut self, offset: u64, buf: &[u8]) -> Result<()> {
138+
let block_size = self.block.media().block_size() as u64;
139+
let media_id = self.block.media().media_id();
140+
let start_block = offset / block_size;
141+
let end_block = (offset + buf.len() as u64).div_ceil(block_size);
142+
let num_blocks = end_block - start_block;
143+
if buf.len() as u64 != num_blocks * block_size
144+
|| !(buf.as_ptr() as usize).is_multiple_of(16)
145+
{
146+
//log::warn!(
147+
// "Unaligned write: offset {}, block size {}, buf addr {:p}, buf len {}",
148+
// offset,
149+
// block_size,
150+
// buf.as_ptr(),
151+
// buf.len()
152+
//);
153+
let mut buf2 = vec![0u8; (num_blocks * block_size) as usize + 15];
154+
let delta = buf2.as_ptr().align_offset(16);
155+
let buf2 = &mut buf2[delta..delta + (num_blocks * block_size) as usize];
156+
self.block.read_blocks(media_id, start_block, buf2)?;
157+
let start_offset = (offset % block_size) as usize;
158+
buf2[start_offset..start_offset + buf.len()].copy_from_slice(buf);
159+
self.block.write_blocks(media_id, start_block, buf2)?;
160+
} else {
161+
self.block.write_blocks(media_id, start_block, buf)?;
162+
}
163+
Ok(())
119164
}
120165

121166
pub async fn write(&mut self, offset: u64, buf: &[u8]) -> Result<()> {
122167
self.os.schedule().await;
123-
Ok(self
124-
.disk
125-
.write_disk(self.block.media().media_id(), offset, buf)?)
168+
self.write_sync(offset, buf)
126169
}
127170

128171
pub fn partitions(&mut self) -> Result<Vec<DiskPartition>> {
@@ -175,17 +218,13 @@ impl gpt_disk_io::BlockIo for &mut Disk {
175218
Ok(self.block.media().last_block() + 1)
176219
}
177220
fn read_blocks(&mut self, start_lba: Lba, dst: &mut [u8]) -> Result<()> {
178-
self.disk.read_disk(
179-
self.block.media().media_id(),
180-
self.block.media().block_size() as u64 * start_lba.0,
181-
dst,
182-
)?;
183-
Ok(())
221+
self.read_sync(self.block.media().block_size() as u64 * start_lba.0, dst)
184222
}
185223
fn write_blocks(&mut self, _start_lba: Lba, _src: &[u8]) -> Result<()> {
186224
unreachable!();
187225
}
188226
fn flush(&mut self) -> Result<()> {
189-
Ok(self.block.flush_blocks()?)
227+
// This is a no-op because write_blocks isn't implemented.
228+
Ok(())
190229
}
191230
}

‎pixie-uefi/src/os/executor.rs‎

Lines changed: 22 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,14 @@
11
use super::{sync::SyncRefCell, UefiOS};
2-
use alloc::{boxed::Box, collections::VecDeque, sync::Arc};
3-
use core::{cell::RefCell, pin::Pin, task::Context};
4-
use futures::{
5-
task::{waker_ref, ArcWake},
6-
Future,
2+
use alloc::{boxed::Box, collections::VecDeque, sync::Arc, task::Wake};
3+
use core::{
4+
cell::RefCell,
5+
future::Future,
6+
pin::Pin,
7+
task::{Context, Waker},
78
};
89

910
pub(super) type BoxFuture<T = ()> = Pin<Box<dyn Future<Output = T> + 'static>>;
1011

11-
static EXECUTOR: SyncRefCell<Executor> = SyncRefCell::new(Executor {
12-
tasks: VecDeque::new(),
13-
});
14-
1512
struct TaskInner {
1613
pub in_queue: bool,
1714
pub future: Option<BoxFuture>,
@@ -47,35 +44,39 @@ impl Task {
4744
unsafe impl Send for Task {}
4845
unsafe impl Sync for Task {}
4946

50-
impl ArcWake for Task {
51-
fn wake_by_ref(task: &Arc<Self>) {
52-
if !task.inner.borrow().in_queue {
53-
task.inner.borrow_mut().in_queue = true;
54-
EXECUTOR.borrow_mut().tasks.push_back(task.clone());
47+
impl Wake for Task {
48+
fn wake(self: Arc<Self>) {
49+
if !self.inner.borrow().in_queue {
50+
self.inner.borrow_mut().in_queue = true;
51+
EXECUTOR.borrow_mut().ready_tasks.push_back(self);
5552
}
5653
}
5754
}
5855

56+
static EXECUTOR: SyncRefCell<Executor> = SyncRefCell::new(Executor {
57+
ready_tasks: VecDeque::new(),
58+
});
59+
5960
pub struct Executor {
6061
// TODO(veluca): scheduling.
61-
tasks: VecDeque<Arc<Task>>,
62+
ready_tasks: VecDeque<Arc<Task>>,
6263
}
6364

6465
impl Executor {
6566
pub fn run(os: UefiOS) -> ! {
6667
loop {
6768
let task = EXECUTOR
6869
.borrow_mut()
69-
.tasks
70+
.ready_tasks
7071
.pop_front()
71-
.expect("Executor should never run out of tasks");
72+
.expect("Executor should never run out of ready tasks");
7273

7374
task.inner.borrow_mut().in_queue = false;
74-
let waker = waker_ref(&task);
75-
let context = &mut Context::from_waker(&waker);
75+
let waker = Waker::from(task.clone());
76+
let mut context = Context::from_waker(&waker);
7677
let mut fut = task.inner.borrow_mut().future.take().unwrap();
7778
let begin = os.timer().micros();
78-
let status = fut.as_mut().poll(context);
79+
let status = fut.as_mut().poll(&mut context);
7980
let end = os.timer().micros();
8081
task.inner.borrow_mut().micros += end - begin;
8182
if status.is_pending() {
@@ -85,6 +86,6 @@ impl Executor {
8586
}
8687

8788
pub(super) fn spawn(task: Arc<Task>) {
88-
EXECUTOR.borrow_mut().tasks.push_back(task);
89+
EXECUTOR.borrow_mut().ready_tasks.push_back(task);
8990
}
9091
}

‎pixie-uefi/src/os/mod.rs‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,10 @@ use core::{
1919
cell::{Ref, RefMut},
2020
ffi::c_void,
2121
fmt::Write,
22-
future::{poll_fn, Future, PollFn},
22+
future::{poll_fn, Future},
2323
net::SocketAddrV4,
2424
ptr::NonNull,
25-
task::{Context, Poll},
25+
task::Poll,
2626
};
2727
use pixie_shared::util::BytesFmt;
2828
use uefi::{
@@ -290,7 +290,7 @@ impl UefiOS {
290290
RefMut::map(self.borrow_mut(), |f| f.net.as_mut().unwrap())
291291
}
292292

293-
pub fn wait_for_ip(self) -> PollFn<impl FnMut(&mut Context<'_>) -> Poll<()>> {
293+
pub fn wait_for_ip(self) -> impl Future<Output = ()> {
294294
poll_fn(move |cx| {
295295
if self.net().has_ip() {
296296
Poll::Ready(())
@@ -303,7 +303,7 @@ impl UefiOS {
303303

304304
/// Interrupt task execution.
305305
/// This is useful to yield the CPU to other tasks.
306-
pub fn schedule(&self) -> PollFn<impl FnMut(&mut Context<'_>) -> Poll<()>> {
306+
pub fn schedule(&self) -> impl Future<Output = ()> {
307307
let mut ready = false;
308308
poll_fn(move |cx| {
309309
if ready {
@@ -316,7 +316,7 @@ impl UefiOS {
316316
})
317317
}
318318

319-
pub fn sleep_us(self, us: u64) -> PollFn<impl FnMut(&mut Context<'_>) -> Poll<()>> {
319+
pub fn sleep_us(self, us: u64) -> impl Future<Output = ()> {
320320
let tgt = self.timer().micros() as u64 + us;
321321
poll_fn(move |cx| {
322322
let now = self.timer().micros() as u64;
@@ -409,11 +409,11 @@ impl UefiOS {
409409
UdpHandle::new(*self, port).await
410410
}
411411

412-
pub async fn read_key(&self) -> Result<Key> {
413-
Ok(poll_fn(move |cx| {
412+
pub fn read_key(&self) -> impl Future<Output = Result<Key>> + '_ {
413+
poll_fn(move |cx| {
414414
let key = self.borrow_mut().input.read_key();
415415
if let Err(e) = key {
416-
return Poll::Ready(Err(e));
416+
return Poll::Ready(Err(e.into()));
417417
}
418418
let key = key.unwrap();
419419
if let Some(key) = key {
@@ -422,7 +422,6 @@ impl UefiOS {
422422
cx.waker().wake_by_ref();
423423
Poll::Pending
424424
})
425-
.await?)
426425
}
427426

428427
pub fn write_with_color(&self, msg: &str, fg: Color, bg: Color) {

‎pixie-uefi/src/parse_disk/gpt.rs‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,12 @@ use pixie_shared::util::BytesFmt;
88

99
pub async fn parse_gpt(disk: &mut Disk) -> Result<Option<Vec<ChunkInfo>>> {
1010
let disk_size = disk.size() as usize;
11-
let Ok(partitions) = disk.partitions() else {
12-
return Ok(None);
11+
let partitions = match disk.partitions() {
12+
Ok(partitions) => partitions,
13+
Err(e) => {
14+
log::debug!("Failed to parse GPT partitions: {e:?}");
15+
return Ok(None);
16+
}
1317
};
1418

1519
let mut pos = 0usize;

0 commit comments

Comments
 (0)