From 03329430b457b4e631ee5d3f51a66ce9affc96ab Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Fri, 7 Mar 2025 20:34:33 +0100
Subject: [PATCH 01/10] storage/driver-block: Handle block_read read buffer
 internally

---
 storage/driver-block/src/lib.rs   | 31 ++++++++++---------------------
 storage/nvmed/src/scheme.rs       | 11 +++--------
 storage/virtio-blkd/src/scheme.rs |  5 +----
 3 files changed, 14 insertions(+), 33 deletions(-)

diff --git a/storage/driver-block/src/lib.rs b/storage/driver-block/src/lib.rs
index 2d40c69d57..316478f61a 100644
--- a/storage/driver-block/src/lib.rs
+++ b/storage/driver-block/src/lib.rs
@@ -6,7 +6,6 @@ use partitionlib::{LogicalBlockSize, PartitionTable};
 
 /// Split the read operation into a series of block reads.
 /// `read_fn` will be called with a block number to be read, and a buffer to be filled.
-/// The buffer must be large enough to hold `blksize` of data.
 /// `read_fn` must return a full block of data.
 /// Result will be the number of bytes read.
 // FIXME make private once nvmed uses the DiskWrapper defined in this crate
@@ -14,7 +13,6 @@ pub fn block_read(
     offset: u64,
     blksize: u32,
     buf: &mut [u8],
-    block_bytes: &mut [u8],
     mut read_fn: impl FnMut(u64, &mut [u8]) -> Result<(), Error>,
 ) -> Result<usize, Error> {
     // TODO: Yield sometimes, perhaps after a few blocks or something.
@@ -32,6 +30,9 @@ pub fn block_read(
     let blk_size = usize::try_from(blksize).expect("blksize larger than usize");
     let mut total_read = 0;
 
+    let mut block_bytes = [0u8; 4096];
+    let block_bytes = &mut block_bytes[..blk_size];
+
     while curr_buf.len() > 0 {
         // TODO: Async/await? I mean, shouldn't AHCI be async?
 
@@ -72,13 +73,12 @@ impl DiskWrapper {
             Ok(512) => LogicalBlockSize::Lb512,
             _ => return None,
         };
-        struct Device<'a, 'b> {
+        struct Device<'a> {
             disk: &'a mut dyn Disk,
             offset: u64,
-            block_bytes: &'b mut [u8],
         }
 
-        impl<'a, 'b> Seek for Device<'a, 'b> {
+        impl<'a> Seek for Device<'a> {
             fn seek(&mut self, from: SeekFrom) -> io::Result<u64> {
                 let size = i64::try_from(self.disk.size()).or(Err(io::Error::new(
                     io::ErrorKind::Other,
@@ -97,7 +97,7 @@ impl DiskWrapper {
             }
         }
         // TODO: Perhaps this impl should be used in the rest of the scheme.
-        impl<'a, 'b> Read for Device<'a, 'b> {
+        impl<'a> Read for Device<'a> {
             fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
                 let blksize = self
                     .disk
@@ -115,7 +115,6 @@ impl DiskWrapper {
                         match disk.read(block, block_bytes) {
                             Ok(Some(bytes)) => {
                                 assert_eq!(bytes, block_bytes.len());
-                                assert_eq!(bytes, blksize as usize);
                                 return Ok(());
                             }
                             Ok(None) => {
@@ -126,26 +125,16 @@ impl DiskWrapper {
                         }
                     }
                 };
-                let bytes_read =
-                    block_read(self.offset, blksize, buf, self.block_bytes, read_block)?;
+                let bytes_read = block_read(self.offset, blksize, buf, read_block)?;
 
                 self.offset += bytes_read as u64;
                 Ok(bytes_read)
             }
         }
 
-        let mut block_bytes = [0u8; 4096];
-
-        partitionlib::get_partitions(
-            &mut Device {
-                disk,
-                offset: 0,
-                block_bytes: &mut block_bytes[..bs.into()],
-            },
-            bs,
-        )
-        .ok()
-        .flatten()
+        partitionlib::get_partitions(&mut Device { disk, offset: 0 }, bs)
+            .ok()
+            .flatten()
     }
 
     pub fn new(mut disk: Box<dyn Disk>) -> Self {
diff --git a/storage/nvmed/src/scheme.rs b/storage/nvmed/src/scheme.rs
index 493a08293b..977d03328b 100644
--- a/storage/nvmed/src/scheme.rs
+++ b/storage/nvmed/src/scheme.rs
@@ -41,14 +41,13 @@ impl DiskWrapper {
             4096 => LogicalBlockSize::Lb4096,
             _ => return None,
         };
-        struct Device<'a, 'b> {
+        struct Device<'a> {
             disk: &'a mut NvmeNamespace,
             nvme: &'a Nvme,
             offset: u64,
-            block_bytes: &'b mut [u8],
         }
 
-        impl<'a, 'b> Seek for Device<'a, 'b> {
+        impl<'a> Seek for Device<'a> {
             fn seek(&mut self, from: io::SeekFrom) -> io::Result<u64> {
                 let size_u = self.disk.blocks * self.disk.block_size;
                 let size = i64::try_from(size_u).or(Err(io::Error::new(
@@ -70,7 +69,7 @@ impl DiskWrapper {
             }
         }
 
-        impl<'a, 'b> Read for Device<'a, 'b> {
+        impl<'a> Read for Device<'a> {
             fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
                 let blksize = self.disk.block_size;
                 let size_in_blocks = self.disk.blocks;
@@ -105,7 +104,6 @@ impl DiskWrapper {
                         .try_into()
                         .expect("Unreasonable block size above 2^32 bytes"),
                     buf,
-                    self.block_bytes,
                     read_block,
                 )?;
                 self.offset += bytes_read as u64;
@@ -113,14 +111,11 @@ impl DiskWrapper {
             }
         }
 
-        let mut block_bytes = [0u8; 4096];
-
         partitionlib::get_partitions(
             &mut Device {
                 disk,
                 nvme,
                 offset: 0,
-                block_bytes: &mut block_bytes[..bs.into()],
             },
             bs,
         )
diff --git a/storage/virtio-blkd/src/scheme.rs b/storage/virtio-blkd/src/scheme.rs
index 609f5f5059..8880779c59 100644
--- a/storage/virtio-blkd/src/scheme.rs
+++ b/storage/virtio-blkd/src/scheme.rs
@@ -124,7 +124,6 @@ impl<'a> DiskScheme<'a> {
         struct VirtioShim<'a, 'b> {
             scheme: &'b DiskScheme<'a>,
             offset: u64,
-            block_bytes: &'b mut [u8],
         }
 
         impl<'a, 'b> Read for VirtioShim<'a, 'b> {
@@ -156,8 +155,7 @@ impl<'a> DiskScheme<'a> {
                     };
 
                 let bytes_read =
-                    driver_block::block_read(self.offset, 512, buf, self.block_bytes, read_block)
-                        .unwrap();
+                    driver_block::block_read(self.offset, 512, buf, read_block).unwrap();
                 self.offset += bytes_read as u64;
                 Ok(bytes_read)
             }
@@ -188,7 +186,6 @@ impl<'a> DiskScheme<'a> {
         let mut shim = VirtioShim {
             scheme: &this,
             offset: 0,
-            block_bytes: &mut [0u8; 4096],
         };
 
         //driver_block::DiskWrapper::new(disk)

From aacf09cef4f630d06e2e4e1d206eaad90e0d22e2 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Fri, 7 Mar 2025 20:42:41 +0100
Subject: [PATCH 02/10] storage/driver-block: Remove id method from Disk trait

---
 storage/ahcid/src/ahci/disk_ata.rs   | 4 ----
 storage/ahcid/src/ahci/disk_atapi.rs | 4 ----
 storage/bcm2835-sdhcid/src/sd/mod.rs | 4 ----
 storage/driver-block/src/lib.rs      | 1 -
 storage/ided/src/ide.rs              | 4 ----
 5 files changed, 17 deletions(-)

diff --git a/storage/ahcid/src/ahci/disk_ata.rs b/storage/ahcid/src/ahci/disk_ata.rs
index 2dc2a1f190..d5173e2f9e 100644
--- a/storage/ahcid/src/ahci/disk_ata.rs
+++ b/storage/ahcid/src/ahci/disk_ata.rs
@@ -154,10 +154,6 @@ impl DiskATA {
 }
 
 impl Disk for DiskATA {
-    fn id(&self) -> usize {
-        self.id
-    }
-
     fn size(&mut self) -> u64 {
         self.size
     }
diff --git a/storage/ahcid/src/ahci/disk_atapi.rs b/storage/ahcid/src/ahci/disk_atapi.rs
index 29ae47996c..33b35f92b1 100644
--- a/storage/ahcid/src/ahci/disk_atapi.rs
+++ b/storage/ahcid/src/ahci/disk_atapi.rs
@@ -72,10 +72,6 @@ impl DiskATAPI {
 }
 
 impl Disk for DiskATAPI {
-    fn id(&self) -> usize {
-        self.id
-    }
-
     fn size(&mut self) -> u64 {
         match self.read_capacity() {
             Ok((blk_count, blk_size)) => (blk_count as u64) * (blk_size as u64),
diff --git a/storage/bcm2835-sdhcid/src/sd/mod.rs b/storage/bcm2835-sdhcid/src/sd/mod.rs
index 8c718ba105..b371833773 100644
--- a/storage/bcm2835-sdhcid/src/sd/mod.rs
+++ b/storage/bcm2835-sdhcid/src/sd/mod.rs
@@ -733,10 +733,6 @@ impl SdHostCtrl {
 }
 
 impl Disk for SdHostCtrl {
-    fn id(&self) -> usize {
-        0xdead_dead
-    }
-
     fn size(&mut self) -> u64 {
         //assert 512MiB
         self.size
diff --git a/storage/driver-block/src/lib.rs b/storage/driver-block/src/lib.rs
index 316478f61a..6b10c14b26 100644
--- a/storage/driver-block/src/lib.rs
+++ b/storage/driver-block/src/lib.rs
@@ -54,7 +54,6 @@ pub fn block_read(
 }
 
 pub trait Disk {
-    fn id(&self) -> usize;
     fn block_length(&mut self) -> syscall::error::Result<u32>;
     fn size(&mut self) -> u64;
 
diff --git a/storage/ided/src/ide.rs b/storage/ided/src/ide.rs
index ad2b771446..b03213a464 100644
--- a/storage/ided/src/ide.rs
+++ b/storage/ided/src/ide.rs
@@ -169,10 +169,6 @@ pub struct AtaDisk {
 }
 
 impl Disk for AtaDisk {
-    fn id(&self) -> usize {
-        self.chan_i << 1 | self.dev as usize
-    }
-
     fn size(&mut self) -> u64 {
         self.size
     }

From ebdfc89e7a56baac86cefc26604827cdce80a531 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Fri, 7 Mar 2025 20:46:33 +0100
Subject: [PATCH 03/10] storage/nvmed: Pass NvmeNamespace by value and remove
 the separate nsid arguments

---
 storage/nvmed/src/nvme/mod.rs | 19 ++++++++-----------
 storage/nvmed/src/scheme.rs   | 22 ++++++++++------------
 2 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/storage/nvmed/src/nvme/mod.rs b/storage/nvmed/src/nvme/mod.rs
index 127dd5ae9f..93a1c03b9b 100644
--- a/storage/nvmed/src/nvme/mod.rs
+++ b/storage/nvmed/src/nvme/mod.rs
@@ -170,7 +170,7 @@ pub struct NvmeRegs {
     cmbsz: Mmio<u32>,
 }
 
-#[derive(Debug)]
+#[derive(Copy, Clone, Debug)]
 pub struct NvmeNamespace {
     pub id: u32,
     pub blocks: u64,
@@ -641,8 +641,7 @@ impl Nvme {
 
     fn namespace_rw(
         &self,
-        namespace: &NvmeNamespace,
-        nsid: u32,
+        namespace: NvmeNamespace,
         lba: u64,
         blocks_1: u16,
         write: bool,
@@ -666,9 +665,9 @@ impl Nvme {
         let mut cmd = NvmeCmd::default();
         let comp = self.submit_and_complete_command(1, |cid| {
             cmd = if write {
-                NvmeCmd::io_write(cid, nsid, lba, blocks_1, ptr0, ptr1)
+                NvmeCmd::io_write(cid, namespace.id, lba, blocks_1, ptr0, ptr1)
             } else {
-                NvmeCmd::io_read(cid, nsid, lba, blocks_1, ptr0, ptr1)
+                NvmeCmd::io_read(cid, namespace.id, lba, blocks_1, ptr0, ptr1)
             };
             cmd.clone()
         });
@@ -683,8 +682,7 @@ impl Nvme {
 
     pub fn namespace_read(
         &self,
-        namespace: &NvmeNamespace,
-        nsid: u32,
+        namespace: NvmeNamespace,
         mut lba: u64,
         buf: &mut [u8],
     ) -> Result<Option<usize>> {
@@ -698,7 +696,7 @@ impl Nvme {
             assert!(blocks > 0);
             assert!(blocks <= 0x1_0000);
 
-            self.namespace_rw(namespace, nsid, lba, (blocks - 1) as u16, false)?;
+            self.namespace_rw(namespace, lba, (blocks - 1) as u16, false)?;
 
             chunk.copy_from_slice(&buffer_guard[..chunk.len()]);
 
@@ -710,8 +708,7 @@ impl Nvme {
 
     pub fn namespace_write(
         &self,
-        namespace: &NvmeNamespace,
-        nsid: u32,
+        namespace: NvmeNamespace,
         mut lba: u64,
         buf: &[u8],
     ) -> Result<Option<usize>> {
@@ -727,7 +724,7 @@ impl Nvme {
 
             buffer_guard[..chunk.len()].copy_from_slice(chunk);
 
-            self.namespace_rw(namespace, nsid, lba, (blocks - 1) as u16, true)?;
+            self.namespace_rw(namespace, lba, (blocks - 1) as u16, true)?;
 
             lba += blocks as u64;
         }
diff --git a/storage/nvmed/src/scheme.rs b/storage/nvmed/src/scheme.rs
index 977d03328b..30dd5d82d5 100644
--- a/storage/nvmed/src/scheme.rs
+++ b/storage/nvmed/src/scheme.rs
@@ -35,14 +35,14 @@ impl AsRef<NvmeNamespace> for DiskWrapper {
 }
 
 impl DiskWrapper {
-    fn pt(disk: &mut NvmeNamespace, nvme: &Nvme) -> Option<PartitionTable> {
+    fn pt(disk: NvmeNamespace, nvme: &Nvme) -> Option<PartitionTable> {
         let bs = match disk.block_size {
             512 => LogicalBlockSize::Lb512,
             4096 => LogicalBlockSize::Lb4096,
             _ => return None,
         };
         struct Device<'a> {
-            disk: &'a mut NvmeNamespace,
+            disk: NvmeNamespace,
             nvme: &'a Nvme,
             offset: u64,
         }
@@ -74,7 +74,7 @@ impl DiskWrapper {
                 let blksize = self.disk.block_size;
                 let size_in_blocks = self.disk.blocks;
 
-                let disk = &mut self.disk;
+                let disk = self.disk;
                 let nvme = &mut self.nvme;
 
                 let read_block = |block: u64, block_bytes: &mut [u8]| {
@@ -83,7 +83,7 @@ impl DiskWrapper {
                     }
                     loop {
                         match nvme
-                            .namespace_read(disk, disk.id, block, block_bytes)
+                            .namespace_read(disk, block, block_bytes)
                             .map_err(|err| io::Error::from_raw_os_error(err.errno))?
                         {
                             Some(bytes) => {
@@ -122,9 +122,9 @@ impl DiskWrapper {
         .ok()
         .flatten()
     }
-    fn new(mut inner: NvmeNamespace, nvme: &Nvme) -> Self {
+    fn new(inner: NvmeNamespace, nvme: &Nvme) -> Self {
         Self {
-            pt: Self::pt(&mut inner, nvme),
+            pt: Self::pt(inner, nvme),
             inner,
         }
     }
@@ -372,7 +372,7 @@ impl SchemeBlock for DiskScheme {
                 let disk = self.disks.get_mut(&number).ok_or(Error::new(EBADF))?;
                 let block_size = disk.as_ref().block_size;
                 self.nvme
-                    .namespace_read(disk.as_ref(), disk.as_ref().id, offset / block_size, buf)
+                    .namespace_read(*disk.as_ref(), offset / block_size, buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(&disk_num).ok_or(Error::new(EBADF))?;
@@ -392,8 +392,7 @@ impl SchemeBlock for DiskScheme {
 
                 let abs_block = part.start_lba + rel_block;
 
-                self.nvme
-                    .namespace_read(disk.as_ref(), disk.as_ref().id, abs_block, buf)
+                self.nvme.namespace_read(*disk.as_ref(), abs_block, buf)
             }
         }
     }
@@ -411,7 +410,7 @@ impl SchemeBlock for DiskScheme {
                 let disk = self.disks.get_mut(&number).ok_or(Error::new(EBADF))?;
                 let block_size = disk.as_ref().block_size;
                 self.nvme
-                    .namespace_write(disk.as_ref(), disk.as_ref().id, offset / block_size, buf)
+                    .namespace_write(*disk.as_ref(), offset / block_size, buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(&disk_num).ok_or(Error::new(EBADF))?;
@@ -431,8 +430,7 @@ impl SchemeBlock for DiskScheme {
 
                 let abs_block = part.start_lba + rel_block;
 
-                self.nvme
-                    .namespace_write(disk.as_ref(), disk.as_ref().id, abs_block, buf)
+                self.nvme.namespace_write(*disk.as_ref(), abs_block, buf)
             }
         }
     }

From f01a8c8d3e20ae90b79742c76d725a7af44dab25 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Fri, 7 Mar 2025 20:58:01 +0100
Subject: [PATCH 04/10] storage: Use DiskWrapper::pt in nvmed and virtio-blkd

---
 storage/driver-block/src/lib.rs   |   6 +-
 storage/nvmed/src/scheme.rs       | 114 +++++++-----------------------
 storage/virtio-blkd/src/scheme.rs | 101 ++++++--------------------
 3 files changed, 49 insertions(+), 172 deletions(-)

diff --git a/storage/driver-block/src/lib.rs b/storage/driver-block/src/lib.rs
index 6b10c14b26..02f771e475 100644
--- a/storage/driver-block/src/lib.rs
+++ b/storage/driver-block/src/lib.rs
@@ -8,8 +8,7 @@ use partitionlib::{LogicalBlockSize, PartitionTable};
 /// `read_fn` will be called with a block number to be read, and a buffer to be filled.
 /// `read_fn` must return a full block of data.
 /// Result will be the number of bytes read.
-// FIXME make private once nvmed uses the DiskWrapper defined in this crate
-pub fn block_read(
+fn block_read(
     offset: u64,
     blksize: u32,
     buf: &mut [u8],
@@ -67,9 +66,10 @@ pub struct DiskWrapper {
 }
 
 impl DiskWrapper {
-    fn pt(disk: &mut dyn Disk) -> Option<PartitionTable> {
+    pub fn pt(disk: &mut dyn Disk) -> Option<PartitionTable> {
         let bs = match disk.block_length() {
             Ok(512) => LogicalBlockSize::Lb512,
+            Ok(4096) => LogicalBlockSize::Lb4096,
             _ => return None,
         };
         struct Device<'a> {
diff --git a/storage/nvmed/src/scheme.rs b/storage/nvmed/src/scheme.rs
index 30dd5d82d5..ee27527d9e 100644
--- a/storage/nvmed/src/scheme.rs
+++ b/storage/nvmed/src/scheme.rs
@@ -1,11 +1,11 @@
 use std::collections::BTreeMap;
 use std::convert::{TryFrom, TryInto};
 use std::fmt::Write;
-use std::io;
-use std::io::prelude::*;
+use std::str;
 use std::sync::Arc;
-use std::{cmp, str};
 
+use driver_block::Disk;
+use partitionlib::PartitionTable;
 use redox_scheme::{CallerCtx, OpenResult, SchemeBlock};
 use syscall::schemev2::NewFdFlags;
 use syscall::{
@@ -15,14 +15,32 @@ use syscall::{
 
 use crate::nvme::{Nvme, NvmeNamespace};
 
-use partitionlib::{LogicalBlockSize, PartitionTable};
-
 enum Handle {
     List(Vec<u8>),       // entries
     Disk(u32),           // disk num
     Partition(u32, u32), // disk num, part num
 }
 
+struct NamespaceAndNvme<'a>(&'a Nvme, NvmeNamespace);
+
+impl Disk for NamespaceAndNvme<'_> {
+    fn block_length(&mut self) -> syscall::error::Result<u32> {
+        Ok(self.1.block_size.try_into().unwrap())
+    }
+
+    fn size(&mut self) -> u64 {
+        self.1.blocks * self.1.block_size
+    }
+
+    fn read(&mut self, block: u64, buffer: &mut [u8]) -> syscall::Result<Option<usize>> {
+        self.0.namespace_read(self.1, block, buffer)
+    }
+
+    fn write(&mut self, block: u64, buffer: &[u8]) -> syscall::Result<Option<usize>> {
+        self.0.namespace_write(self.1, block, buffer)
+    }
+}
+
 pub struct DiskWrapper {
     inner: NvmeNamespace,
     pt: Option<PartitionTable>,
@@ -36,91 +54,7 @@ impl AsRef<NvmeNamespace> for DiskWrapper {
 
 impl DiskWrapper {
     fn pt(disk: NvmeNamespace, nvme: &Nvme) -> Option<PartitionTable> {
-        let bs = match disk.block_size {
-            512 => LogicalBlockSize::Lb512,
-            4096 => LogicalBlockSize::Lb4096,
-            _ => return None,
-        };
-        struct Device<'a> {
-            disk: NvmeNamespace,
-            nvme: &'a Nvme,
-            offset: u64,
-        }
-
-        impl<'a> Seek for Device<'a> {
-            fn seek(&mut self, from: io::SeekFrom) -> io::Result<u64> {
-                let size_u = self.disk.blocks * self.disk.block_size;
-                let size = i64::try_from(size_u).or(Err(io::Error::new(
-                    io::ErrorKind::Other,
-                    "Disk larger than 2^63 - 1 bytes",
-                )))?;
-
-                self.offset = match from {
-                    io::SeekFrom::Start(new_pos) => cmp::min(size_u, new_pos),
-                    io::SeekFrom::Current(new_pos) => {
-                        cmp::max(0, cmp::min(size, self.offset as i64 + new_pos)) as u64
-                    }
-                    io::SeekFrom::End(new_pos) => {
-                        cmp::max(0, cmp::min(size + new_pos, size)) as u64
-                    }
-                };
-
-                Ok(self.offset)
-            }
-        }
-
-        impl<'a> Read for Device<'a> {
-            fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-                let blksize = self.disk.block_size;
-                let size_in_blocks = self.disk.blocks;
-
-                let disk = self.disk;
-                let nvme = &mut self.nvme;
-
-                let read_block = |block: u64, block_bytes: &mut [u8]| {
-                    if block >= size_in_blocks {
-                        return Err(io::Error::from_raw_os_error(syscall::EOVERFLOW));
-                    }
-                    loop {
-                        match nvme
-                            .namespace_read(disk, block, block_bytes)
-                            .map_err(|err| io::Error::from_raw_os_error(err.errno))?
-                        {
-                            Some(bytes) => {
-                                assert_eq!(bytes, block_bytes.len());
-                                assert_eq!(bytes, blksize as usize);
-                                return Ok(());
-                            }
-                            None => {
-                                std::thread::yield_now();
-                                continue;
-                            } // TODO: Does this driver have (internal) error handling at all?
-                        }
-                    }
-                };
-                let bytes_read = driver_block::block_read(
-                    self.offset,
-                    blksize
-                        .try_into()
-                        .expect("Unreasonable block size above 2^32 bytes"),
-                    buf,
-                    read_block,
-                )?;
-                self.offset += bytes_read as u64;
-                Ok(bytes_read)
-            }
-        }
-
-        partitionlib::get_partitions(
-            &mut Device {
-                disk,
-                nvme,
-                offset: 0,
-            },
-            bs,
-        )
-        .ok()
-        .flatten()
+        driver_block::DiskWrapper::pt(&mut NamespaceAndNvme(nvme, disk))
     }
     fn new(inner: NvmeNamespace, nvme: &Nvme) -> Self {
         Self {
diff --git a/storage/virtio-blkd/src/scheme.rs b/storage/virtio-blkd/src/scheme.rs
index 8880779c59..d4727b9278 100644
--- a/storage/virtio-blkd/src/scheme.rs
+++ b/storage/virtio-blkd/src/scheme.rs
@@ -1,13 +1,10 @@
 use std::collections::BTreeMap;
-use std::io::Read;
-use std::io::Result as IoResult;
-use std::io::Seek;
 
 use std::fmt::Write;
 use std::sync::Arc;
 
 use common::dma::Dma;
-use partitionlib::LogicalBlockSize;
+use driver_block::DiskWrapper;
 use partitionlib::PartitionTable;
 
 use redox_scheme::CallerCtx;
@@ -121,85 +118,31 @@ impl<'a> DiskScheme<'a> {
             part_table: None,
         };
 
-        struct VirtioShim<'a, 'b> {
-            scheme: &'b DiskScheme<'a>,
-            offset: u64,
-        }
-
-        impl<'a, 'b> Read for VirtioShim<'a, 'b> {
-            fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
-                let read_block =
-                    |block: u64, block_bytes: &mut [u8]| -> Result<(), std::io::Error> {
-                        let req = Dma::new(BlockVirtRequest {
-                            ty: BlockRequestTy::In,
-                            reserved: 0,
-                            sector: block,
-                        })
-                        .unwrap();
-
-                        let result = Dma::new([0u8; 512]).unwrap();
-                        let status = Dma::new(u8::MAX).unwrap();
-
-                        let chain = ChainBuilder::new()
-                            .chain(Buffer::new(&req))
-                            .chain(Buffer::new(&result).flags(DescriptorFlags::WRITE_ONLY))
-                            .chain(Buffer::new(&status).flags(DescriptorFlags::WRITE_ONLY))
-                            .build();
-
-                        futures::executor::block_on(self.scheme.queue.send(chain));
-                        assert_eq!(*status, 0);
-
-                        let size = core::cmp::min(block_bytes.len(), result.len());
-                        block_bytes[..size].copy_from_slice(&result.as_slice()[..size]);
-                        Ok(())
-                    };
-
-                let bytes_read =
-                    driver_block::block_read(self.offset, 512, buf, read_block).unwrap();
-                self.offset += bytes_read as u64;
-                Ok(bytes_read)
-            }
-        }
-
-        impl<'a, 'b> Seek for VirtioShim<'a, 'b> {
-            fn seek(&mut self, from: std::io::SeekFrom) -> IoResult<u64> {
-                let size_u = self.scheme.cfg.capacity() * self.scheme.cfg.block_size() as u64;
-                let size = i64::try_from(size_u).or(Err(std::io::Error::new(
-                    std::io::ErrorKind::Other,
-                    "Disk larger than 2^63 - 1 bytes",
-                )))?;
-
-                self.offset = match from {
-                    std::io::SeekFrom::Start(new_pos) => std::cmp::min(size_u, new_pos),
-                    std::io::SeekFrom::Current(new_pos) => {
-                        std::cmp::max(0, std::cmp::min(size, self.offset as i64 + new_pos)) as u64
-                    }
-                    std::io::SeekFrom::End(new_pos) => {
-                        std::cmp::max(0, std::cmp::min(size + new_pos, size)) as u64
-                    }
-                };
-
-                Ok(self.offset)
-            }
-        }
-
-        let mut shim = VirtioShim {
-            scheme: &this,
-            offset: 0,
-        };
-
-        //driver_block::DiskWrapper::new(disk)
-
-        // FIXME use DiskWrapper instead
-        let part_table = partitionlib::get_partitions(&mut shim, LogicalBlockSize::Lb512)
-            .ok()
-            .flatten();
-
-        this.part_table = part_table;
+        this.part_table = DiskWrapper::pt(&mut this);
         this
     }
 }
 
+impl driver_block::Disk for DiskScheme<'_> {
+    fn block_length(&mut self) -> syscall::error::Result<u32> {
+        Ok(self.cfg.block_size())
+    }
+
+    fn size(&mut self) -> u64 {
+        self.cfg.capacity() * self.cfg.block_size() as u64
+    }
+
+    fn read(&mut self, block: u64, buffer: &mut [u8]) -> syscall::Result<Option<usize>> {
+        Ok(Some(futures::executor::block_on(
+            self.queue.read(block, buffer),
+        )))
+    }
+
+    fn write(&mut self, _block: u64, _buffer: &[u8]) -> syscall::Result<Option<usize>> {
+        todo!()
+    }
+}
+
 impl<'a> SchemeBlock for DiskScheme<'a> {
     fn xopen(
         &mut self,

From dfe23b06834700847891c69f4831fda3fb8a9a03 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Fri, 7 Mar 2025 21:35:49 +0100
Subject: [PATCH 05/10] storage/driver-block: Make Disk::block_length
 infallible

It should be read when creating the struct implementing Disk and every
disk needs to have some block length.
---
 storage/ahcid/src/ahci/disk_ata.rs   |  4 +--
 storage/ahcid/src/ahci/disk_atapi.rs | 44 ++++++++++++----------------
 storage/ahcid/src/scheme.rs          | 16 +++++-----
 storage/bcm2835-sdhcid/src/scheme.rs | 16 +++++-----
 storage/bcm2835-sdhcid/src/sd/mod.rs |  4 +--
 storage/driver-block/src/lib.rs      | 11 +++----
 storage/ided/src/ide.rs              |  4 +--
 storage/ided/src/scheme.rs           | 16 +++++-----
 storage/nvmed/src/scheme.rs          |  4 +--
 storage/virtio-blkd/src/scheme.rs    |  4 +--
 10 files changed, 57 insertions(+), 66 deletions(-)

diff --git a/storage/ahcid/src/ahci/disk_ata.rs b/storage/ahcid/src/ahci/disk_ata.rs
index d5173e2f9e..ee3b32a8b6 100644
--- a/storage/ahcid/src/ahci/disk_ata.rs
+++ b/storage/ahcid/src/ahci/disk_ata.rs
@@ -178,7 +178,7 @@ impl Disk for DiskATA {
         }
     }
 
-    fn block_length(&mut self) -> Result<u32> {
-        Ok(512)
+    fn block_length(&mut self) -> u32 {
+        512
     }
 }
diff --git a/storage/ahcid/src/ahci/disk_atapi.rs b/storage/ahcid/src/ahci/disk_atapi.rs
index 33b35f92b1..ce3ce0b962 100644
--- a/storage/ahcid/src/ahci/disk_atapi.rs
+++ b/storage/ahcid/src/ahci/disk_atapi.rs
@@ -25,6 +25,8 @@ pub struct DiskATAPI {
     // Just using the same buffer size as DiskATA
     // Although the sector size is different (and varies)
     buf: Dma<[u8; 256 * 512]>,
+    blk_count: u32,
+    blk_size: u32,
 }
 
 impl DiskATAPI {
@@ -38,12 +40,20 @@ impl DiskATAPI {
             .unwrap_or_else(|_| unreachable!());
 
         let mut fb = unsafe { Dma::zeroed()?.assume_init() };
-        let buf = unsafe { Dma::zeroed()?.assume_init() };
+        let mut buf = unsafe { Dma::zeroed()?.assume_init() };
 
         port.init(&mut clb, &mut ctbas, &mut fb);
 
         let size = unsafe { port.identify_packet(&mut clb, &mut ctbas).unwrap_or(0) };
 
+        let mut cmd = [0; 16];
+        cmd[0] = SCSI_READ_CAPACITY;
+        port.atapi_dma(&cmd, 8, &mut clb, &mut ctbas, &mut buf)?;
+
+        // Instead of a count, contains number of last LBA, so add 1
+        let blk_count = BigEndian::read_u32(&buf[0..4]) + 1;
+        let blk_size = BigEndian::read_u32(&buf[4..8]);
+
         Ok(DiskATAPI {
             id,
             port,
@@ -52,37 +62,25 @@ impl DiskATAPI {
             ctbas,
             _fb: fb,
             buf,
+            blk_count,
+            blk_size,
         })
     }
-
-    fn read_capacity(&mut self) -> Result<(u32, u32)> {
-        // TODO: only query when needed (disk changed)
-
-        let mut cmd = [0; 16];
-        cmd[0] = SCSI_READ_CAPACITY;
-        self.port
-            .atapi_dma(&cmd, 8, &mut self.clb, &mut self.ctbas, &mut self.buf)?;
-
-        // Instead of a count, contains number of last LBA, so add 1
-        let blk_count = BigEndian::read_u32(&self.buf[0..4]) + 1;
-        let blk_size = BigEndian::read_u32(&self.buf[4..8]);
-
-        Ok((blk_count, blk_size))
-    }
 }
 
 impl Disk for DiskATAPI {
+    fn block_length(&mut self) -> u32 {
+        self.blk_size
+    }
+
     fn size(&mut self) -> u64 {
-        match self.read_capacity() {
-            Ok((blk_count, blk_size)) => (blk_count as u64) * (blk_size as u64),
-            Err(_) => 0, // XXX
-        }
+        u64::from(self.blk_count) * u64::from(self.blk_size)
     }
 
     fn read(&mut self, block: u64, buffer: &mut [u8]) -> Result<Option<usize>> {
         // TODO: Handle audio CDs, which use special READ CD command
 
-        let blk_len = self.block_length()?;
+        let blk_len = self.blk_size;
         let sectors = buffer.len() as u32 / blk_len;
 
         fn read10_cmd(block: u32, count: u16) -> [u8; 16] {
@@ -147,8 +145,4 @@ impl Disk for DiskATAPI {
     fn write(&mut self, _block: u64, _buffer: &[u8]) -> Result<Option<usize>> {
         Err(Error::new(EBADF)) // TODO: Implement writing
     }
-
-    fn block_length(&mut self) -> Result<u32> {
-        Ok(self.read_capacity()?.1)
-    }
 }
diff --git a/storage/ahcid/src/scheme.rs b/storage/ahcid/src/scheme.rs
index 7ed30f2de1..4655529a6b 100644
--- a/storage/ahcid/src/scheme.rs
+++ b/storage/ahcid/src/scheme.rs
@@ -175,7 +175,7 @@ impl SchemeBlock for DiskScheme {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
                 stat.st_mode = MODE_FILE;
                 stat.st_size = disk.size();
-                stat.st_blksize = disk.block_length()?;
+                stat.st_blksize = disk.block_length();
                 Ok(Some(0))
             }
             Handle::Partition(disk_id, part_num) => {
@@ -190,8 +190,8 @@ impl SchemeBlock for DiskScheme {
                 };
 
                 stat.st_mode = MODE_FILE; // TODO: Block device?
-                stat.st_size = size * u64::from(disk.block_length()?);
-                stat.st_blksize = disk.block_length()?;
+                stat.st_size = size * u64::from(disk.block_length());
+                stat.st_blksize = disk.block_length();
                 stat.st_blocks = size;
                 Ok(Some(0))
             }
@@ -263,12 +263,12 @@ impl SchemeBlock for DiskScheme {
             }
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_length()?;
+                let blk_len = disk.block_length();
                 disk.read(offset / u64::from(blk_len), buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_length()?;
+                let blksize = disk.block_length();
 
                 // validate that we're actually reading within the bounds of the partition
                 let rel_block = offset / u64::from(blksize);
@@ -304,12 +304,12 @@ impl SchemeBlock for DiskScheme {
             Handle::List(_) => Err(Error::new(EBADF)),
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_length()?;
+                let blk_len = disk.block_length();
                 disk.write(offset / u64::from(blk_len), buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_length()?;
+                let blksize = disk.block_length();
 
                 // validate that we're actually reading within the bounds of the partition
                 let rel_block = offset / u64::from(blksize);
@@ -351,7 +351,7 @@ impl SchemeBlock for DiskScheme {
                         .get(part_num as usize)
                         .ok_or(Error::new(EBADF))?
                         .size;
-                    u64::from(disk.block_length()?) * block_count
+                    u64::from(disk.block_length()) * block_count
                 }
             },
         ))
diff --git a/storage/bcm2835-sdhcid/src/scheme.rs b/storage/bcm2835-sdhcid/src/scheme.rs
index 70bf4772e0..7b5b195c61 100644
--- a/storage/bcm2835-sdhcid/src/scheme.rs
+++ b/storage/bcm2835-sdhcid/src/scheme.rs
@@ -164,7 +164,7 @@ impl SchemeBlock for DiskScheme {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
                 stat.st_mode = MODE_FILE;
                 stat.st_size = disk.size();
-                stat.st_blksize = disk.block_length()?;
+                stat.st_blksize = disk.block_length();
                 Ok(Some(0))
             }
             Handle::Partition(disk_id, part_num) => {
@@ -179,8 +179,8 @@ impl SchemeBlock for DiskScheme {
                 };
 
                 stat.st_mode = MODE_FILE; // TODO: Block device?
-                stat.st_size = size * u64::from(disk.block_length()?);
-                stat.st_blksize = disk.block_length()?;
+                stat.st_size = size * u64::from(disk.block_length());
+                stat.st_blksize = disk.block_length();
                 stat.st_blocks = size;
                 Ok(Some(0))
             }
@@ -251,12 +251,12 @@ impl SchemeBlock for DiskScheme {
             }
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_length()?;
+                let blk_len = disk.block_length();
                 disk.read(offset / u64::from(blk_len), buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_length()?;
+                let blksize = disk.block_length();
 
                 // validate that we're actually reading within the bounds of the partition
                 let rel_block = offset / u64::from(blksize);
@@ -285,12 +285,12 @@ impl SchemeBlock for DiskScheme {
             Handle::List(_) => Err(Error::new(EBADF)),
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_length()?;
+                let blk_len = disk.block_length();
                 disk.write(offset / u64::from(blk_len), buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_length()?;
+                let blksize = disk.block_length();
 
                 // validate that we're actually reading within the bounds of the partition
                 let rel_block = offset / u64::from(blksize as u64);
@@ -331,7 +331,7 @@ impl SchemeBlock for DiskScheme {
                     .get(part_num as usize)
                     .ok_or(Error::new(EBADF))?
                     .size;
-                Ok(Some(u64::from(disk.block_length()?) * block_count))
+                Ok(Some(u64::from(disk.block_length()) * block_count))
             }
         }
     }
diff --git a/storage/bcm2835-sdhcid/src/sd/mod.rs b/storage/bcm2835-sdhcid/src/sd/mod.rs
index b371833773..deb3839f81 100644
--- a/storage/bcm2835-sdhcid/src/sd/mod.rs
+++ b/storage/bcm2835-sdhcid/src/sd/mod.rs
@@ -774,7 +774,7 @@ impl Disk for SdHostCtrl {
         }
     }
 
-    fn block_length(&mut self) -> Result<u32> {
-        Ok(512)
+    fn block_length(&mut self) -> u32 {
+        512
     }
 }
diff --git a/storage/driver-block/src/lib.rs b/storage/driver-block/src/lib.rs
index 02f771e475..24e45b8a1b 100644
--- a/storage/driver-block/src/lib.rs
+++ b/storage/driver-block/src/lib.rs
@@ -53,7 +53,7 @@ fn block_read(
 }
 
 pub trait Disk {
-    fn block_length(&mut self) -> syscall::error::Result<u32>;
+    fn block_length(&mut self) -> u32;
     fn size(&mut self) -> u64;
 
     fn read(&mut self, block: u64, buffer: &mut [u8]) -> syscall::Result<Option<usize>>;
@@ -68,8 +68,8 @@ pub struct DiskWrapper {
 impl DiskWrapper {
     pub fn pt(disk: &mut dyn Disk) -> Option<PartitionTable> {
         let bs = match disk.block_length() {
-            Ok(512) => LogicalBlockSize::Lb512,
-            Ok(4096) => LogicalBlockSize::Lb4096,
+            512 => LogicalBlockSize::Lb512,
+            4096 => LogicalBlockSize::Lb4096,
             _ => return None,
         };
         struct Device<'a> {
@@ -98,10 +98,7 @@ impl DiskWrapper {
         // TODO: Perhaps this impl should be used in the rest of the scheme.
         impl<'a> Read for Device<'a> {
             fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-                let blksize = self
-                    .disk
-                    .block_length()
-                    .map_err(|err| io::Error::from_raw_os_error(err.errno))?;
+                let blksize = self.disk.block_length();
                 let size_in_blocks = self.disk.size() / u64::from(blksize);
 
                 let disk = &mut self.disk;
diff --git a/storage/ided/src/ide.rs b/storage/ided/src/ide.rs
index b03213a464..94584a9dfb 100644
--- a/storage/ided/src/ide.rs
+++ b/storage/ided/src/ide.rs
@@ -461,7 +461,7 @@ impl Disk for AtaDisk {
         Ok(Some(count))
     }
 
-    fn block_length(&mut self) -> Result<u32> {
-        Ok(512)
+    fn block_length(&mut self) -> u32 {
+        512
     }
 }
diff --git a/storage/ided/src/scheme.rs b/storage/ided/src/scheme.rs
index ff0a63e58a..aca709013e 100644
--- a/storage/ided/src/scheme.rs
+++ b/storage/ided/src/scheme.rs
@@ -178,7 +178,7 @@ impl SchemeBlock for DiskScheme {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
                 stat.st_mode = MODE_FILE;
                 stat.st_size = disk.size();
-                stat.st_blksize = disk.block_length()?;
+                stat.st_blksize = disk.block_length();
                 Ok(Some(0))
             }
             Handle::Partition(disk_id, part_num) => {
@@ -193,8 +193,8 @@ impl SchemeBlock for DiskScheme {
                 };
 
                 stat.st_mode = MODE_FILE; // TODO: Block device?
-                stat.st_size = size * u64::from(disk.block_length()?);
-                stat.st_blksize = disk.block_length()?;
+                stat.st_size = size * u64::from(disk.block_length());
+                stat.st_blksize = disk.block_length();
                 stat.st_blocks = size;
                 Ok(Some(0))
             }
@@ -265,12 +265,12 @@ impl SchemeBlock for DiskScheme {
             }
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_length()?;
+                let blk_len = disk.block_length();
                 disk.read(offset / u64::from(blk_len), buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_length()?;
+                let blksize = disk.block_length();
 
                 // validate that we're actually reading within the bounds of the partition
                 let rel_block = offset / u64::from(blksize);
@@ -299,12 +299,12 @@ impl SchemeBlock for DiskScheme {
             Handle::List(_) => Err(Error::new(EBADF)),
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_length()?;
+                let blk_len = disk.block_length();
                 disk.write(offset / u64::from(blk_len), buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_length()?;
+                let blksize = disk.block_length();
 
                 // validate that we're actually reading within the bounds of the partition
                 let rel_block = offset / u64::from(blksize as u64);
@@ -345,7 +345,7 @@ impl SchemeBlock for DiskScheme {
                     .get(part_num as usize)
                     .ok_or(Error::new(EBADF))?
                     .size;
-                Ok(Some(u64::from(disk.block_length()?) * block_count))
+                Ok(Some(u64::from(disk.block_length()) * block_count))
             }
         }
     }
diff --git a/storage/nvmed/src/scheme.rs b/storage/nvmed/src/scheme.rs
index ee27527d9e..1a730920c9 100644
--- a/storage/nvmed/src/scheme.rs
+++ b/storage/nvmed/src/scheme.rs
@@ -24,8 +24,8 @@ enum Handle {
 struct NamespaceAndNvme<'a>(&'a Nvme, NvmeNamespace);
 
 impl Disk for NamespaceAndNvme<'_> {
-    fn block_length(&mut self) -> syscall::error::Result<u32> {
-        Ok(self.1.block_size.try_into().unwrap())
+    fn block_length(&mut self) -> u32 {
+        self.1.block_size.try_into().unwrap()
     }
 
     fn size(&mut self) -> u64 {
diff --git a/storage/virtio-blkd/src/scheme.rs b/storage/virtio-blkd/src/scheme.rs
index d4727b9278..17d7ae62b5 100644
--- a/storage/virtio-blkd/src/scheme.rs
+++ b/storage/virtio-blkd/src/scheme.rs
@@ -124,8 +124,8 @@ impl<'a> DiskScheme<'a> {
 }
 
 impl driver_block::Disk for DiskScheme<'_> {
-    fn block_length(&mut self) -> syscall::error::Result<u32> {
-        Ok(self.cfg.block_size())
+    fn block_length(&mut self) -> u32 {
+        self.cfg.block_size()
     }
 
     fn size(&mut self) -> u64 {

From d0a7aed2fde4c5ab12e9406f2998652213d93669 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Fri, 7 Mar 2025 21:52:19 +0100
Subject: [PATCH 06/10] storage/virtio-blk: Introduce VirtioDisk type

---
 storage/virtio-blkd/src/scheme.rs | 155 +++++++++++++++---------------
 1 file changed, 78 insertions(+), 77 deletions(-)

diff --git a/storage/virtio-blkd/src/scheme.rs b/storage/virtio-blkd/src/scheme.rs
index 17d7ae62b5..710a43ba4d 100644
--- a/storage/virtio-blkd/src/scheme.rs
+++ b/storage/virtio-blkd/src/scheme.rs
@@ -4,7 +4,7 @@ use std::fmt::Write;
 use std::sync::Arc;
 
 use common::dma::Dma;
-use driver_block::DiskWrapper;
+use driver_block::{Disk, DiskWrapper};
 use partitionlib::PartitionTable;
 
 use redox_scheme::CallerCtx;
@@ -87,6 +87,39 @@ impl BlkExtension for Queue<'_> {
     }
 }
 
+struct VirtioDisk<'a> {
+    queue: Arc<Queue<'a>>,
+    cfg: BlockDeviceConfig,
+}
+
+impl<'a> VirtioDisk<'a> {
+    pub fn new(queue: Arc<Queue<'a>>, cfg: BlockDeviceConfig) -> Self {
+        Self { queue, cfg }
+    }
+}
+
+impl driver_block::Disk for VirtioDisk<'_> {
+    fn block_length(&mut self) -> u32 {
+        self.cfg.block_size()
+    }
+
+    fn size(&mut self) -> u64 {
+        self.cfg.capacity() * u64::from(self.cfg.block_size())
+    }
+
+    fn read(&mut self, block: u64, buffer: &mut [u8]) -> syscall::Result<Option<usize>> {
+        Ok(Some(futures::executor::block_on(
+            self.queue.read(block, buffer),
+        )))
+    }
+
+    fn write(&mut self, block: u64, buffer: &[u8]) -> syscall::Result<Option<usize>> {
+        Ok(Some(futures::executor::block_on(
+            self.queue.write(block, buffer),
+        )))
+    }
+}
+
 pub enum Handle {
     Partition {
         /// Partition Number
@@ -101,45 +134,22 @@ pub enum Handle {
 }
 
 pub struct DiskScheme<'a> {
-    queue: Arc<Queue<'a>>,
+    disk: VirtioDisk<'a>,
     next_id: usize,
-    cfg: BlockDeviceConfig,
     handles: BTreeMap<usize, Handle>,
     part_table: Option<PartitionTable>,
 }
 
 impl<'a> DiskScheme<'a> {
     pub fn new(queue: Arc<Queue<'a>>, cfg: BlockDeviceConfig) -> Self {
-        let mut this = Self {
-            queue,
+        let mut disk = VirtioDisk::new(queue, cfg);
+
+        Self {
+            part_table: DiskWrapper::pt(&mut disk),
+            disk,
             next_id: 0,
-            cfg,
             handles: BTreeMap::new(),
-            part_table: None,
-        };
-
-        this.part_table = DiskWrapper::pt(&mut this);
-        this
-    }
-}
-
-impl driver_block::Disk for DiskScheme<'_> {
-    fn block_length(&mut self) -> u32 {
-        self.cfg.block_size()
-    }
-
-    fn size(&mut self) -> u64 {
-        self.cfg.capacity() * self.cfg.block_size() as u64
-    }
-
-    fn read(&mut self, block: u64, buffer: &mut [u8]) -> syscall::Result<Option<usize>> {
-        Ok(Some(futures::executor::block_on(
-            self.queue.read(block, buffer),
-        )))
-    }
-
-    fn write(&mut self, _block: u64, _buffer: &[u8]) -> syscall::Result<Option<usize>> {
-        todo!()
+        }
     }
 }
 
@@ -224,45 +234,40 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
         offset: u64,
         _fcntl_flags: u32,
     ) -> syscall::Result<Option<usize>> {
-        Ok(Some(
-            match *self.handles.get_mut(&id).ok_or(Error::new(EBADF))? {
-                Handle::List { ref mut entries } => {
-                    let src = usize::try_from(offset)
-                        .ok()
-                        .and_then(|o| entries.get(o..))
-                        .unwrap_or(&[]);
-                    let count = core::cmp::min(src.len(), buf.len());
-                    buf[..count].copy_from_slice(&src[..count]);
-                    count
-                }
+        match *self.handles.get_mut(&id).ok_or(Error::new(EBADF))? {
+            Handle::List { ref mut entries } => {
+                let src = usize::try_from(offset)
+                    .ok()
+                    .and_then(|o| entries.get(o..))
+                    .unwrap_or(&[]);
+                let count = core::cmp::min(src.len(), buf.len());
+                buf[..count].copy_from_slice(&src[..count]);
+                Ok(Some(count))
+            }
 
-                Handle::Partition { number } => {
-                    let part_table = self.part_table.as_ref().unwrap();
-                    let part = part_table
-                        .partitions
-                        .get(number as usize)
-                        .ok_or(Error::new(EBADF))?;
+            Handle::Partition { number } => {
+                let part_table = self.part_table.as_ref().unwrap();
+                let part = part_table
+                    .partitions
+                    .get(number as usize)
+                    .ok_or(Error::new(EBADF))?;
 
-                    // Get the offset in sectors.
-                    let rel_block = offset / BLK_SIZE;
-                    // if rel_block >= part.size {
-                    //     return Err(Error::new(EOVERFLOW));
-                    // }
+                // Get the offset in sectors.
+                let rel_block = offset / BLK_SIZE;
+                // if rel_block >= part.size {
+                //     return Err(Error::new(EOVERFLOW));
+                // }
 
-                    let abs_block = part.start_lba + rel_block;
+                let abs_block = part.start_lba + rel_block;
 
-                    futures::executor::block_on(self.queue.read(abs_block, buf))
-                }
+                self.disk.read(abs_block, buf)
+            }
 
-                Handle::Disk => {
-                    let block_size = self.cfg.block_size();
-
-                    futures::executor::block_on(
-                        self.queue.read(offset / u64::from(block_size), buf),
-                    )
-                }
-            },
-        ))
+            Handle::Disk => {
+                let block = offset / u64::from(self.disk.block_length());
+                self.disk.read(block, buf)
+            }
+        }
     }
 
     fn write(
@@ -272,18 +277,14 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
         offset: u64,
         _fcntl_flags: u32,
     ) -> syscall::Result<Option<usize>> {
-        Ok(Some(
-            match *self.handles.get_mut(&id).ok_or(Error::new(EBADF))? {
-                Handle::Disk => {
-                    let block_size = self.cfg.block_size();
-                    futures::executor::block_on(
-                        self.queue.write(offset / u64::from(block_size), buf),
-                    )
-                }
+        match *self.handles.get_mut(&id).ok_or(Error::new(EBADF))? {
+            Handle::Disk => {
+                let block = offset / u64::from(self.disk.block_length());
+                self.disk.write(block, buf)
+            }
 
-                _ => todo!(),
-            },
-        ))
+            _ => todo!(),
+        }
     }
 
     fn fsize(&mut self, id: usize) -> syscall::Result<Option<u64>> {
@@ -311,7 +312,7 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
                     len
                 }
 
-                Handle::Disk => self.cfg.capacity() * u64::from(self.cfg.block_size()),
+                Handle::Disk => self.disk.size(),
             },
         ))
     }

From 9d27fad3d56e7ea6a3b2a2158c1e064332e3b9f5 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Fri, 7 Mar 2025 22:15:28 +0100
Subject: [PATCH 07/10] storage/driver-block: Make DiskWrapper generic over the
 disk type

This allows it to contain non-'static disks.
---
 storage/ahcid/src/scheme.rs          |  2 +-
 storage/bcm2835-sdhcid/src/scheme.rs |  2 +-
 storage/driver-block/src/lib.rs      | 42 ++++++++++++++++++++--------
 storage/ided/src/scheme.rs           |  2 +-
 4 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/storage/ahcid/src/scheme.rs b/storage/ahcid/src/scheme.rs
index 4655529a6b..437f70bd2c 100644
--- a/storage/ahcid/src/scheme.rs
+++ b/storage/ahcid/src/scheme.rs
@@ -22,7 +22,7 @@ enum Handle {
 pub struct DiskScheme {
     scheme_name: String,
     hba_mem: &'static mut HbaMem,
-    disks: Box<[DiskWrapper]>,
+    disks: Box<[DiskWrapper<Box<dyn Disk>>]>,
     handles: BTreeMap<usize, Handle>,
     next_id: usize,
 }
diff --git a/storage/bcm2835-sdhcid/src/scheme.rs b/storage/bcm2835-sdhcid/src/scheme.rs
index 7b5b195c61..66f469a943 100644
--- a/storage/bcm2835-sdhcid/src/scheme.rs
+++ b/storage/bcm2835-sdhcid/src/scheme.rs
@@ -19,7 +19,7 @@ enum Handle {
 
 pub struct DiskScheme {
     scheme_name: String,
-    disks: Box<[DiskWrapper]>,
+    disks: Box<[DiskWrapper<Box<dyn Disk>>]>,
     handles: BTreeMap<usize, Handle>,
     next_id: usize,
 }
diff --git a/storage/driver-block/src/lib.rs b/storage/driver-block/src/lib.rs
index 24e45b8a1b..12d0926204 100644
--- a/storage/driver-block/src/lib.rs
+++ b/storage/driver-block/src/lib.rs
@@ -60,13 +60,31 @@ pub trait Disk {
     fn write(&mut self, block: u64, buffer: &[u8]) -> syscall::Result<Option<usize>>;
 }
 
-pub struct DiskWrapper {
-    pub disk: Box<dyn Disk>,
+impl<T: Disk + ?Sized> Disk for Box<T> {
+    fn block_length(&mut self) -> u32 {
+        (**self).block_length()
+    }
+
+    fn size(&mut self) -> u64 {
+        (**self).size()
+    }
+
+    fn read(&mut self, block: u64, buffer: &mut [u8]) -> syscall::Result<Option<usize>> {
+        (**self).read(block, buffer)
+    }
+
+    fn write(&mut self, block: u64, buffer: &[u8]) -> syscall::Result<Option<usize>> {
+        (**self).write(block, buffer)
+    }
+}
+
+pub struct DiskWrapper<T> {
+    pub disk: T,
     pub pt: Option<PartitionTable>,
 }
 
-impl DiskWrapper {
-    pub fn pt(disk: &mut dyn Disk) -> Option<PartitionTable> {
+impl<T: Disk> DiskWrapper<T> {
+    pub fn pt(disk: &mut T) -> Option<PartitionTable> {
         let bs = match disk.block_length() {
             512 => LogicalBlockSize::Lb512,
             4096 => LogicalBlockSize::Lb4096,
@@ -133,23 +151,23 @@ impl DiskWrapper {
             .flatten()
     }
 
-    pub fn new(mut disk: Box<dyn Disk>) -> Self {
+    pub fn new(mut disk: T) -> Self {
         Self {
-            pt: Self::pt(&mut *disk),
+            pt: Self::pt(&mut disk),
             disk,
         }
     }
 }
 
-impl std::ops::Deref for DiskWrapper {
-    type Target = dyn Disk;
+impl<T> std::ops::Deref for DiskWrapper<T> {
+    type Target = T;
 
     fn deref(&self) -> &Self::Target {
-        &*self.disk
+        &self.disk
     }
 }
-impl std::ops::DerefMut for DiskWrapper {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut *self.disk
+impl<T> std::ops::DerefMut for DiskWrapper<T> {
+    fn deref_mut(&mut self) -> &mut T {
+        &mut self.disk
     }
 }
diff --git a/storage/ided/src/scheme.rs b/storage/ided/src/scheme.rs
index aca709013e..fb0fa26169 100644
--- a/storage/ided/src/scheme.rs
+++ b/storage/ided/src/scheme.rs
@@ -22,7 +22,7 @@ enum Handle {
 pub struct DiskScheme {
     scheme_name: String,
     chans: Box<[Arc<Mutex<Channel>>]>,
-    disks: Box<[DiskWrapper]>,
+    disks: Box<[DiskWrapper<Box<dyn Disk>>]>,
     handles: BTreeMap<usize, Handle>,
     next_id: usize,
 }

From 371e8c0326bdfb918726bb22791f7fa640bcb4a0 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Fri, 7 Mar 2025 22:15:56 +0100
Subject: [PATCH 08/10] storage/virtio-blkd: Use DiskWrapper

---
 Cargo.lock                        |  1 -
 storage/virtio-blkd/Cargo.toml    |  1 -
 storage/virtio-blkd/src/scheme.rs | 17 ++++++-----------
 3 files changed, 6 insertions(+), 13 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 0bc2322768..537432f9cd 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1672,7 +1672,6 @@ dependencies = [
  "futures",
  "libredox",
  "log",
- "partitionlib",
  "pcid",
  "redox-daemon",
  "redox-scheme 0.3.0",
diff --git a/storage/virtio-blkd/Cargo.toml b/storage/virtio-blkd/Cargo.toml
index de7aba04a1..2b2a663565 100644
--- a/storage/virtio-blkd/Cargo.toml
+++ b/storage/virtio-blkd/Cargo.toml
@@ -15,7 +15,6 @@ spin = "*"
 redox-daemon = "0.1"
 redox_syscall = { version = "0.5", features = ["std"] }
 redox-scheme = { git = "https://gitlab.redox-os.org/redox-os/redox-scheme.git" }
-partitionlib = { path = "../partitionlib" }
 
 common = { path = "../../common" }
 driver-block = { path = "../driver-block" }
diff --git a/storage/virtio-blkd/src/scheme.rs b/storage/virtio-blkd/src/scheme.rs
index 710a43ba4d..e498798556 100644
--- a/storage/virtio-blkd/src/scheme.rs
+++ b/storage/virtio-blkd/src/scheme.rs
@@ -5,7 +5,6 @@ use std::sync::Arc;
 
 use common::dma::Dma;
 use driver_block::{Disk, DiskWrapper};
-use partitionlib::PartitionTable;
 
 use redox_scheme::CallerCtx;
 use redox_scheme::OpenResult;
@@ -134,19 +133,15 @@ pub enum Handle {
 }
 
 pub struct DiskScheme<'a> {
-    disk: VirtioDisk<'a>,
+    disk: DiskWrapper<VirtioDisk<'a>>,
     next_id: usize,
     handles: BTreeMap<usize, Handle>,
-    part_table: Option<PartitionTable>,
 }
 
 impl<'a> DiskScheme<'a> {
     pub fn new(queue: Arc<Queue<'a>>, cfg: BlockDeviceConfig) -> Self {
-        let mut disk = VirtioDisk::new(queue, cfg);
-
         Self {
-            part_table: DiskWrapper::pt(&mut disk),
-            disk,
+            disk: DiskWrapper::new(VirtioDisk::new(queue, cfg)),
             next_id: 0,
             handles: BTreeMap::new(),
         }
@@ -170,7 +165,7 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
                 //            to the namespace id).
                 write!(list, "{}\n", 0).unwrap();
 
-                if let Some(part_table) = &self.part_table {
+                if let Some(part_table) = &self.disk.pt {
                     for part_num in 0..part_table.partitions.len() {
                         write!(list, "{}p{}\n", 0, part_num).unwrap();
                     }
@@ -201,7 +196,7 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
             let part_num_str = &path_str[p_pos + 1..];
             let part_num = part_num_str.parse::<u32>().unwrap();
 
-            let part_table = self.part_table.as_ref().unwrap();
+            let part_table = self.disk.pt.as_ref().unwrap();
             let _part = part_table.partitions.get(part_num as usize).unwrap();
 
             let id = self.next_id;
@@ -246,7 +241,7 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
             }
 
             Handle::Partition { number } => {
-                let part_table = self.part_table.as_ref().unwrap();
+                let part_table = self.disk.pt.as_ref().unwrap();
                 let part = part_table
                     .partitions
                     .get(number as usize)
@@ -298,7 +293,7 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
                 }
 
                 Handle::Partition { number } => {
-                    let part_table = self.part_table.as_ref().unwrap();
+                    let part_table = self.disk.pt.as_ref().unwrap();
                     let part = part_table
                         .partitions
                         .get(number as usize)

From b1dcda4cf770110b30b9af2591bdfef3a4239188 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Fri, 7 Mar 2025 22:35:01 +0100
Subject: [PATCH 09/10] storage/nvmed: Use DiskWrapper

---
 Cargo.lock                           |  1 -
 storage/ahcid/src/ahci/disk_ata.rs   | 10 ++--
 storage/ahcid/src/ahci/disk_atapi.rs |  4 +-
 storage/ahcid/src/scheme.rs          | 16 +++---
 storage/bcm2835-sdhcid/src/scheme.rs | 16 +++---
 storage/bcm2835-sdhcid/src/sd/mod.rs | 10 ++--
 storage/driver-block/src/lib.rs      | 14 ++---
 storage/ided/src/ide.rs              | 10 ++--
 storage/ided/src/scheme.rs           | 16 +++---
 storage/nvmed/Cargo.toml             |  1 -
 storage/nvmed/src/scheme.rs          | 80 ++++++++--------------------
 storage/virtio-blkd/src/scheme.rs    |  8 +--
 12 files changed, 74 insertions(+), 112 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 537432f9cd..7f304690ff 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -854,7 +854,6 @@ dependencies = [
  "futures",
  "libredox",
  "log",
- "partitionlib",
  "pcid",
  "redox-daemon",
  "redox-scheme 0.3.0",
diff --git a/storage/ahcid/src/ahci/disk_ata.rs b/storage/ahcid/src/ahci/disk_ata.rs
index ee3b32a8b6..bb52a4571d 100644
--- a/storage/ahcid/src/ahci/disk_ata.rs
+++ b/storage/ahcid/src/ahci/disk_ata.rs
@@ -154,7 +154,11 @@ impl DiskATA {
 }
 
 impl Disk for DiskATA {
-    fn size(&mut self) -> u64 {
+    fn block_size(&self) -> u32 {
+        512
+    }
+
+    fn size(&self) -> u64 {
         self.size
     }
 
@@ -177,8 +181,4 @@ impl Disk for DiskATA {
             }
         }
     }
-
-    fn block_length(&mut self) -> u32 {
-        512
-    }
 }
diff --git a/storage/ahcid/src/ahci/disk_atapi.rs b/storage/ahcid/src/ahci/disk_atapi.rs
index ce3ce0b962..b92112c36e 100644
--- a/storage/ahcid/src/ahci/disk_atapi.rs
+++ b/storage/ahcid/src/ahci/disk_atapi.rs
@@ -69,11 +69,11 @@ impl DiskATAPI {
 }
 
 impl Disk for DiskATAPI {
-    fn block_length(&mut self) -> u32 {
+    fn block_size(&self) -> u32 {
         self.blk_size
     }
 
-    fn size(&mut self) -> u64 {
+    fn size(&self) -> u64 {
         u64::from(self.blk_count) * u64::from(self.blk_size)
     }
 
diff --git a/storage/ahcid/src/scheme.rs b/storage/ahcid/src/scheme.rs
index 437f70bd2c..63d432e9a8 100644
--- a/storage/ahcid/src/scheme.rs
+++ b/storage/ahcid/src/scheme.rs
@@ -175,7 +175,7 @@ impl SchemeBlock for DiskScheme {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
                 stat.st_mode = MODE_FILE;
                 stat.st_size = disk.size();
-                stat.st_blksize = disk.block_length();
+                stat.st_blksize = disk.block_size();
                 Ok(Some(0))
             }
             Handle::Partition(disk_id, part_num) => {
@@ -190,8 +190,8 @@ impl SchemeBlock for DiskScheme {
                 };
 
                 stat.st_mode = MODE_FILE; // TODO: Block device?
-                stat.st_size = size * u64::from(disk.block_length());
-                stat.st_blksize = disk.block_length();
+                stat.st_size = size * u64::from(disk.block_size());
+                stat.st_blksize = disk.block_size();
                 stat.st_blocks = size;
                 Ok(Some(0))
             }
@@ -263,12 +263,12 @@ impl SchemeBlock for DiskScheme {
             }
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_length();
+                let blk_len = disk.block_size();
                 disk.read(offset / u64::from(blk_len), buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_length();
+                let blksize = disk.block_size();
 
                 // validate that we're actually reading within the bounds of the partition
                 let rel_block = offset / u64::from(blksize);
@@ -304,12 +304,12 @@ impl SchemeBlock for DiskScheme {
             Handle::List(_) => Err(Error::new(EBADF)),
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_length();
+                let blk_len = disk.block_size();
                 disk.write(offset / u64::from(blk_len), buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_length();
+                let blksize = disk.block_size();
 
                 // validate that we're actually reading within the bounds of the partition
                 let rel_block = offset / u64::from(blksize);
@@ -351,7 +351,7 @@ impl SchemeBlock for DiskScheme {
                         .get(part_num as usize)
                         .ok_or(Error::new(EBADF))?
                         .size;
-                    u64::from(disk.block_length()) * block_count
+                    u64::from(disk.block_size()) * block_count
                 }
             },
         ))
diff --git a/storage/bcm2835-sdhcid/src/scheme.rs b/storage/bcm2835-sdhcid/src/scheme.rs
index 66f469a943..ebe0612346 100644
--- a/storage/bcm2835-sdhcid/src/scheme.rs
+++ b/storage/bcm2835-sdhcid/src/scheme.rs
@@ -164,7 +164,7 @@ impl SchemeBlock for DiskScheme {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
                 stat.st_mode = MODE_FILE;
                 stat.st_size = disk.size();
-                stat.st_blksize = disk.block_length();
+                stat.st_blksize = disk.block_size();
                 Ok(Some(0))
             }
             Handle::Partition(disk_id, part_num) => {
@@ -179,8 +179,8 @@ impl SchemeBlock for DiskScheme {
                 };
 
                 stat.st_mode = MODE_FILE; // TODO: Block device?
-                stat.st_size = size * u64::from(disk.block_length());
-                stat.st_blksize = disk.block_length();
+                stat.st_size = size * u64::from(disk.block_size());
+                stat.st_blksize = disk.block_size();
                 stat.st_blocks = size;
                 Ok(Some(0))
             }
@@ -251,12 +251,12 @@ impl SchemeBlock for DiskScheme {
             }
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_length();
+                let blk_len = disk.block_size();
                 disk.read(offset / u64::from(blk_len), buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_length();
+                let blksize = disk.block_size();
 
                 // validate that we're actually reading within the bounds of the partition
                 let rel_block = offset / u64::from(blksize);
@@ -285,12 +285,12 @@ impl SchemeBlock for DiskScheme {
             Handle::List(_) => Err(Error::new(EBADF)),
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_length();
+                let blk_len = disk.block_size();
                 disk.write(offset / u64::from(blk_len), buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_length();
+                let blksize = disk.block_size();
 
                 // validate that we're actually reading within the bounds of the partition
                 let rel_block = offset / u64::from(blksize as u64);
@@ -331,7 +331,7 @@ impl SchemeBlock for DiskScheme {
                     .get(part_num as usize)
                     .ok_or(Error::new(EBADF))?
                     .size;
-                Ok(Some(u64::from(disk.block_length()) * block_count))
+                Ok(Some(u64::from(disk.block_size()) * block_count))
             }
         }
     }
diff --git a/storage/bcm2835-sdhcid/src/sd/mod.rs b/storage/bcm2835-sdhcid/src/sd/mod.rs
index deb3839f81..f8e9ab351b 100644
--- a/storage/bcm2835-sdhcid/src/sd/mod.rs
+++ b/storage/bcm2835-sdhcid/src/sd/mod.rs
@@ -733,7 +733,11 @@ impl SdHostCtrl {
 }
 
 impl Disk for SdHostCtrl {
-    fn size(&mut self) -> u64 {
+    fn block_size(&self) -> u32 {
+        512
+    }
+
+    fn size(&self) -> u64 {
         //assert 512MiB
         self.size
     }
@@ -773,8 +777,4 @@ impl Disk for SdHostCtrl {
             Err(err) => Err(err),
         }
     }
-
-    fn block_length(&mut self) -> u32 {
-        512
-    }
 }
diff --git a/storage/driver-block/src/lib.rs b/storage/driver-block/src/lib.rs
index 12d0926204..1eb165d185 100644
--- a/storage/driver-block/src/lib.rs
+++ b/storage/driver-block/src/lib.rs
@@ -53,19 +53,19 @@ fn block_read(
 }
 
 pub trait Disk {
-    fn block_length(&mut self) -> u32;
-    fn size(&mut self) -> u64;
+    fn block_size(&self) -> u32;
+    fn size(&self) -> u64;
 
     fn read(&mut self, block: u64, buffer: &mut [u8]) -> syscall::Result<Option<usize>>;
     fn write(&mut self, block: u64, buffer: &[u8]) -> syscall::Result<Option<usize>>;
 }
 
 impl<T: Disk + ?Sized> Disk for Box<T> {
-    fn block_length(&mut self) -> u32 {
-        (**self).block_length()
+    fn block_size(&self) -> u32 {
+        (**self).block_size()
     }
 
-    fn size(&mut self) -> u64 {
+    fn size(&self) -> u64 {
         (**self).size()
     }
 
@@ -85,7 +85,7 @@ pub struct DiskWrapper<T> {
 
 impl<T: Disk> DiskWrapper<T> {
     pub fn pt(disk: &mut T) -> Option<PartitionTable> {
-        let bs = match disk.block_length() {
+        let bs = match disk.block_size() {
             512 => LogicalBlockSize::Lb512,
             4096 => LogicalBlockSize::Lb4096,
             _ => return None,
@@ -116,7 +116,7 @@ impl<T: Disk> DiskWrapper<T> {
         // TODO: Perhaps this impl should be used in the rest of the scheme.
         impl<'a> Read for Device<'a> {
             fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-                let blksize = self.disk.block_length();
+                let blksize = self.disk.block_size();
                 let size_in_blocks = self.disk.size() / u64::from(blksize);
 
                 let disk = &mut self.disk;
diff --git a/storage/ided/src/ide.rs b/storage/ided/src/ide.rs
index 94584a9dfb..9748b46b93 100644
--- a/storage/ided/src/ide.rs
+++ b/storage/ided/src/ide.rs
@@ -169,7 +169,11 @@ pub struct AtaDisk {
 }
 
 impl Disk for AtaDisk {
-    fn size(&mut self) -> u64 {
+    fn block_size(&self) -> u32 {
+        512
+    }
+
+    fn size(&self) -> u64 {
         self.size
     }
 
@@ -460,8 +464,4 @@ impl Disk for AtaDisk {
 
         Ok(Some(count))
     }
-
-    fn block_length(&mut self) -> u32 {
-        512
-    }
 }
diff --git a/storage/ided/src/scheme.rs b/storage/ided/src/scheme.rs
index fb0fa26169..167eb78a80 100644
--- a/storage/ided/src/scheme.rs
+++ b/storage/ided/src/scheme.rs
@@ -178,7 +178,7 @@ impl SchemeBlock for DiskScheme {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
                 stat.st_mode = MODE_FILE;
                 stat.st_size = disk.size();
-                stat.st_blksize = disk.block_length();
+                stat.st_blksize = disk.block_size();
                 Ok(Some(0))
             }
             Handle::Partition(disk_id, part_num) => {
@@ -193,8 +193,8 @@ impl SchemeBlock for DiskScheme {
                 };
 
                 stat.st_mode = MODE_FILE; // TODO: Block device?
-                stat.st_size = size * u64::from(disk.block_length());
-                stat.st_blksize = disk.block_length();
+                stat.st_size = size * u64::from(disk.block_size());
+                stat.st_blksize = disk.block_size();
                 stat.st_blocks = size;
                 Ok(Some(0))
             }
@@ -265,12 +265,12 @@ impl SchemeBlock for DiskScheme {
             }
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_length();
+                let blk_len = disk.block_size();
                 disk.read(offset / u64::from(blk_len), buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_length();
+                let blksize = disk.block_size();
 
                 // validate that we're actually reading within the bounds of the partition
                 let rel_block = offset / u64::from(blksize);
@@ -299,12 +299,12 @@ impl SchemeBlock for DiskScheme {
             Handle::List(_) => Err(Error::new(EBADF)),
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_length();
+                let blk_len = disk.block_size();
                 disk.write(offset / u64::from(blk_len), buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_length();
+                let blksize = disk.block_size();
 
                 // validate that we're actually reading within the bounds of the partition
                 let rel_block = offset / u64::from(blksize as u64);
@@ -345,7 +345,7 @@ impl SchemeBlock for DiskScheme {
                     .get(part_num as usize)
                     .ok_or(Error::new(EBADF))?
                     .size;
-                Ok(Some(u64::from(disk.block_length()) * block_count))
+                Ok(Some(u64::from(disk.block_size()) * block_count))
             }
         }
     }
diff --git a/storage/nvmed/Cargo.toml b/storage/nvmed/Cargo.toml
index 04cf4bb27c..8a6d1c6663 100644
--- a/storage/nvmed/Cargo.toml
+++ b/storage/nvmed/Cargo.toml
@@ -13,7 +13,6 @@ redox-daemon = "0.1"
 redox_syscall = { version = "0.5", features = ["std"] }
 redox-scheme = { git = "https://gitlab.redox-os.org/redox-os/redox-scheme.git" }
 redox_event = "0.4"
-partitionlib = { path = "../partitionlib" }
 smallvec = "1"
 
 common = { path = "../../common" }
diff --git a/storage/nvmed/src/scheme.rs b/storage/nvmed/src/scheme.rs
index 1a730920c9..48018b33e5 100644
--- a/storage/nvmed/src/scheme.rs
+++ b/storage/nvmed/src/scheme.rs
@@ -4,8 +4,7 @@ use std::fmt::Write;
 use std::str;
 use std::sync::Arc;
 
-use driver_block::Disk;
-use partitionlib::PartitionTable;
+use driver_block::{Disk, DiskWrapper};
 use redox_scheme::{CallerCtx, OpenResult, SchemeBlock};
 use syscall::schemev2::NewFdFlags;
 use syscall::{
@@ -21,14 +20,14 @@ enum Handle {
     Partition(u32, u32), // disk num, part num
 }
 
-struct NamespaceAndNvme<'a>(&'a Nvme, NvmeNamespace);
+struct NvmeDisk(Arc<Nvme>, NvmeNamespace);
 
-impl Disk for NamespaceAndNvme<'_> {
-    fn block_length(&mut self) -> u32 {
+impl Disk for NvmeDisk {
+    fn block_size(&self) -> u32 {
         self.1.block_size.try_into().unwrap()
     }
 
-    fn size(&mut self) -> u64 {
+    fn size(&self) -> u64 {
         self.1.blocks * self.1.block_size
     }
 
@@ -41,33 +40,9 @@ impl Disk for NamespaceAndNvme<'_> {
     }
 }
 
-pub struct DiskWrapper {
-    inner: NvmeNamespace,
-    pt: Option<PartitionTable>,
-}
-
-impl AsRef<NvmeNamespace> for DiskWrapper {
-    fn as_ref(&self) -> &NvmeNamespace {
-        &self.inner
-    }
-}
-
-impl DiskWrapper {
-    fn pt(disk: NvmeNamespace, nvme: &Nvme) -> Option<PartitionTable> {
-        driver_block::DiskWrapper::pt(&mut NamespaceAndNvme(nvme, disk))
-    }
-    fn new(inner: NvmeNamespace, nvme: &Nvme) -> Self {
-        Self {
-            pt: Self::pt(inner, nvme),
-            inner,
-        }
-    }
-}
-
 pub struct DiskScheme {
     scheme_name: String,
-    nvme: Arc<Nvme>,
-    disks: BTreeMap<u32, DiskWrapper>,
+    disks: BTreeMap<u32, DiskWrapper<NvmeDisk>>,
     handles: BTreeMap<usize, Handle>,
     next_id: usize,
 }
@@ -82,9 +57,8 @@ impl DiskScheme {
             scheme_name,
             disks: disks
                 .into_iter()
-                .map(|(k, v)| (k, DiskWrapper::new(v, &nvme)))
+                .map(|(k, ns)| (k, DiskWrapper::new(NvmeDisk(nvme.clone(), ns))))
                 .collect(),
-            nvme,
             handles: BTreeMap::new(),
             next_id: 0,
         }
@@ -209,13 +183,9 @@ impl SchemeBlock for DiskScheme {
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(&number).ok_or(Error::new(EBADF))?;
                 stat.st_mode = MODE_FILE;
-                stat.st_blocks = disk.as_ref().blocks;
-                stat.st_blksize = disk
-                    .as_ref()
-                    .block_size
-                    .try_into()
-                    .expect("Unreasonable block size of over 2^32 bytes");
-                stat.st_size = disk.as_ref().blocks * disk.as_ref().block_size;
+                stat.st_blocks = disk.1.blocks;
+                stat.st_blksize = disk.block_size();
+                stat.st_size = disk.size();
                 Ok(Some(0))
             }
             Handle::Partition(disk_num, part_num) => {
@@ -228,13 +198,9 @@ impl SchemeBlock for DiskScheme {
                     .get(part_num as usize)
                     .ok_or(Error::new(EBADF))?;
                 stat.st_mode = MODE_FILE;
-                stat.st_size = part.size * disk.as_ref().block_size;
+                stat.st_size = part.size * u64::from(disk.block_size());
                 stat.st_blocks = part.size;
-                stat.st_blksize = disk
-                    .as_ref()
-                    .block_size
-                    .try_into()
-                    .expect("Unreasonable block size of over 2^32 bytes");
+                stat.st_blksize = disk.block_size();
                 Ok(Some(0))
             }
         }
@@ -304,9 +270,8 @@ impl SchemeBlock for DiskScheme {
             }
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(&number).ok_or(Error::new(EBADF))?;
-                let block_size = disk.as_ref().block_size;
-                self.nvme
-                    .namespace_read(*disk.as_ref(), offset / block_size, buf)
+                let block_size = u64::from(disk.block_size());
+                disk.read(offset / block_size, buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(&disk_num).ok_or(Error::new(EBADF))?;
@@ -318,7 +283,7 @@ impl SchemeBlock for DiskScheme {
                     .get(part_num as usize)
                     .ok_or(Error::new(EBADF))?;
 
-                let block_size = disk.as_ref().block_size;
+                let block_size = u64::from(disk.block_size());
                 let rel_block = offset / block_size;
                 if rel_block >= part.size {
                     return Err(Error::new(EOVERFLOW));
@@ -326,7 +291,7 @@ impl SchemeBlock for DiskScheme {
 
                 let abs_block = part.start_lba + rel_block;
 
-                self.nvme.namespace_read(*disk.as_ref(), abs_block, buf)
+                disk.read(abs_block, buf)
             }
         }
     }
@@ -342,9 +307,8 @@ impl SchemeBlock for DiskScheme {
             Handle::List(_) => Err(Error::new(EBADF)),
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(&number).ok_or(Error::new(EBADF))?;
-                let block_size = disk.as_ref().block_size;
-                self.nvme
-                    .namespace_write(*disk.as_ref(), offset / block_size, buf)
+                let block_size = u64::from(disk.block_size());
+                disk.write(offset / block_size, buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(&disk_num).ok_or(Error::new(EBADF))?;
@@ -356,7 +320,7 @@ impl SchemeBlock for DiskScheme {
                     .get(part_num as usize)
                     .ok_or(Error::new(EBADF))?;
 
-                let block_size = disk.as_ref().block_size;
+                let block_size = u64::from(disk.block_size());
                 let rel_block = offset / block_size;
                 if rel_block >= part.size {
                     return Err(Error::new(EOVERFLOW));
@@ -364,7 +328,7 @@ impl SchemeBlock for DiskScheme {
 
                 let abs_block = part.start_lba + rel_block;
 
-                self.nvme.namespace_write(*disk.as_ref(), abs_block, buf)
+                disk.write(abs_block, buf)
             }
         }
     }
@@ -375,7 +339,7 @@ impl SchemeBlock for DiskScheme {
                 Handle::List(ref handle) => handle.len() as u64,
                 Handle::Disk(number) => {
                     let disk = self.disks.get_mut(&number).ok_or(Error::new(EBADF))?;
-                    disk.as_ref().blocks * disk.as_ref().block_size
+                    disk.size()
                 }
                 Handle::Partition(disk_num, part_num) => {
                     let disk = self.disks.get_mut(&disk_num).ok_or(Error::new(EBADF))?;
@@ -387,7 +351,7 @@ impl SchemeBlock for DiskScheme {
                         .get(part_num as usize)
                         .ok_or(Error::new(EBADF))?;
 
-                    part.size * disk.as_ref().block_size
+                    part.size * u64::from(disk.block_size())
                 }
             },
         ))
diff --git a/storage/virtio-blkd/src/scheme.rs b/storage/virtio-blkd/src/scheme.rs
index e498798556..f47abc33c6 100644
--- a/storage/virtio-blkd/src/scheme.rs
+++ b/storage/virtio-blkd/src/scheme.rs
@@ -98,11 +98,11 @@ impl<'a> VirtioDisk<'a> {
 }
 
 impl driver_block::Disk for VirtioDisk<'_> {
-    fn block_length(&mut self) -> u32 {
+    fn block_size(&self) -> u32 {
         self.cfg.block_size()
     }
 
-    fn size(&mut self) -> u64 {
+    fn size(&self) -> u64 {
         self.cfg.capacity() * u64::from(self.cfg.block_size())
     }
 
@@ -259,7 +259,7 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
             }
 
             Handle::Disk => {
-                let block = offset / u64::from(self.disk.block_length());
+                let block = offset / u64::from(self.disk.block_size());
                 self.disk.read(block, buf)
             }
         }
@@ -274,7 +274,7 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
     ) -> syscall::Result<Option<usize>> {
         match *self.handles.get_mut(&id).ok_or(Error::new(EBADF))? {
             Handle::Disk => {
-                let block = offset / u64::from(self.disk.block_length());
+                let block = offset / u64::from(self.disk.block_size());
                 self.disk.write(block, buf)
             }
 

From 60a141b51c27219984f2d4e1712e8e544d9f2d52 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Fri, 7 Mar 2025 22:49:22 +0100
Subject: [PATCH 10/10] storage: Move most partition handling into DiskWrapper

---
 storage/ahcid/src/scheme.rs          | 57 ++++-----------------
 storage/bcm2835-sdhcid/src/scheme.rs | 56 ++++----------------
 storage/driver-block/src/lib.rs      | 76 +++++++++++++++++++++++++---
 storage/ided/src/scheme.rs           | 56 ++++----------------
 storage/nvmed/src/scheme.rs          | 52 ++++---------------
 storage/virtio-blkd/src/scheme.rs    | 36 +++++--------
 6 files changed, 121 insertions(+), 212 deletions(-)

diff --git a/storage/ahcid/src/scheme.rs b/storage/ahcid/src/scheme.rs
index 63d432e9a8..74dc53067e 100644
--- a/storage/ahcid/src/scheme.rs
+++ b/storage/ahcid/src/scheme.rs
@@ -7,8 +7,8 @@ use driver_block::{Disk, DiskWrapper};
 use redox_scheme::{CallerCtx, OpenResult, SchemeBlock};
 use syscall::schemev2::NewFdFlags;
 use syscall::{
-    Error, Result, Stat, EACCES, EBADF, EISDIR, ENOENT, ENOLCK, EOVERFLOW, MODE_DIR, MODE_FILE,
-    O_DIRECTORY, O_STAT,
+    Error, Result, Stat, EACCES, EBADF, EISDIR, ENOENT, ENOLCK, MODE_DIR, MODE_FILE, O_DIRECTORY,
+    O_STAT,
 };
 
 use crate::ahci::hba::HbaMem;
@@ -263,32 +263,13 @@ impl SchemeBlock for DiskScheme {
             }
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_size();
-                disk.read(offset / u64::from(blk_len), buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.read(None, block, buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_size();
-
-                // validate that we're actually reading within the bounds of the partition
-                let rel_block = offset / u64::from(blksize);
-
-                let abs_block = {
-                    let pt = disk.pt.as_ref().ok_or(Error::new(EBADF))?;
-                    let partition = pt
-                        .partitions
-                        .get(part_num as usize)
-                        .ok_or(Error::new(EBADF))?;
-
-                    let abs_block = partition.start_lba + rel_block;
-                    // TODO: This shouldn't return EOVERFLOW?
-                    if rel_block >= partition.size {
-                        return Err(Error::new(EOVERFLOW));
-                    }
-                    abs_block
-                };
-
-                disk.read(abs_block, buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.read(Some(part_num as usize), block, buf)
             }
         }
     }
@@ -304,31 +285,13 @@ impl SchemeBlock for DiskScheme {
             Handle::List(_) => Err(Error::new(EBADF)),
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_size();
-                disk.write(offset / u64::from(blk_len), buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.write(None, block, buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_size();
-
-                // validate that we're actually reading within the bounds of the partition
-                let rel_block = offset / u64::from(blksize);
-
-                let abs_block = {
-                    let pt = disk.pt.as_ref().ok_or(Error::new(EBADF))?;
-                    let partition = pt
-                        .partitions
-                        .get(part_num as usize)
-                        .ok_or(Error::new(EBADF))?;
-
-                    let abs_block = partition.start_lba + rel_block;
-                    if rel_block >= partition.size {
-                        return Err(Error::new(EOVERFLOW));
-                    }
-                    abs_block
-                };
-
-                disk.write(abs_block, buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.write(Some(part_num as usize), block, buf)
             }
         }
     }
diff --git a/storage/bcm2835-sdhcid/src/scheme.rs b/storage/bcm2835-sdhcid/src/scheme.rs
index ebe0612346..dfd7264786 100644
--- a/storage/bcm2835-sdhcid/src/scheme.rs
+++ b/storage/bcm2835-sdhcid/src/scheme.rs
@@ -5,8 +5,8 @@ use std::str;
 use driver_block::{Disk, DiskWrapper};
 use syscall::schemev2::NewFdFlags;
 use syscall::{
-    Error, Result, Stat, EACCES, EBADF, EISDIR, ENOENT, ENOLCK, EOVERFLOW, MODE_DIR, MODE_FILE,
-    O_DIRECTORY, O_STAT,
+    Error, Result, Stat, EACCES, EBADF, EISDIR, ENOENT, ENOLCK, MODE_DIR, MODE_FILE, O_DIRECTORY,
+    O_STAT,
 };
 
 use redox_scheme::{CallerCtx, OpenResult, SchemeBlock};
@@ -251,31 +251,13 @@ impl SchemeBlock for DiskScheme {
             }
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_size();
-                disk.read(offset / u64::from(blk_len), buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.read(None, block, buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_size();
-
-                // validate that we're actually reading within the bounds of the partition
-                let rel_block = offset / u64::from(blksize);
-
-                let abs_block = {
-                    let pt = disk.pt.as_ref().ok_or(Error::new(EBADF))?;
-                    let partition = pt
-                        .partitions
-                        .get(part_num as usize)
-                        .ok_or(Error::new(EBADF))?;
-
-                    let abs_block = partition.start_lba + rel_block;
-                    if rel_block >= partition.size {
-                        return Err(Error::new(EOVERFLOW));
-                    }
-                    abs_block
-                };
-
-                disk.read(abs_block, buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.read(Some(part_num as usize), block, buf)
             }
         }
     }
@@ -285,31 +267,13 @@ impl SchemeBlock for DiskScheme {
             Handle::List(_) => Err(Error::new(EBADF)),
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_size();
-                disk.write(offset / u64::from(blk_len), buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.write(None, block, buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_size();
-
-                // validate that we're actually reading within the bounds of the partition
-                let rel_block = offset / u64::from(blksize as u64);
-
-                let abs_block = {
-                    let pt = disk.pt.as_ref().ok_or(Error::new(EBADF))?;
-                    let partition = pt
-                        .partitions
-                        .get(part_num as usize)
-                        .ok_or(Error::new(EBADF))?;
-
-                    let abs_block = partition.start_lba + rel_block;
-                    if rel_block >= partition.size {
-                        return Err(Error::new(EOVERFLOW));
-                    }
-                    abs_block
-                };
-
-                disk.write(abs_block, buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.write(Some(part_num as usize), block, buf)
             }
         }
     }
diff --git a/storage/driver-block/src/lib.rs b/storage/driver-block/src/lib.rs
index 1eb165d185..a2910c7b11 100644
--- a/storage/driver-block/src/lib.rs
+++ b/storage/driver-block/src/lib.rs
@@ -3,6 +3,7 @@ use std::io::Error;
 use std::io::{self, Read, Seek, SeekFrom};
 
 use partitionlib::{LogicalBlockSize, PartitionTable};
+use syscall::{EBADF, EOVERFLOW};
 
 /// Split the read operation into a series of block reads.
 /// `read_fn` will be called with a block number to be read, and a buffer to be filled.
@@ -157,17 +158,76 @@ impl<T: Disk> DiskWrapper<T> {
             disk,
         }
     }
-}
 
-impl<T> std::ops::Deref for DiskWrapper<T> {
-    type Target = T;
-
-    fn deref(&self) -> &Self::Target {
+    pub fn disk(&self) -> &T {
         &self.disk
     }
-}
-impl<T> std::ops::DerefMut for DiskWrapper<T> {
-    fn deref_mut(&mut self) -> &mut T {
+
+    pub fn disk_mut(&mut self) -> &mut T {
         &mut self.disk
     }
+
+    pub fn block_size(&self) -> u32 {
+        self.disk.block_size()
+    }
+
+    pub fn size(&self) -> u64 {
+        self.disk.size()
+    }
+
+    pub fn read(
+        &mut self,
+        part_num: Option<usize>,
+        block: u64,
+        buf: &mut [u8],
+    ) -> syscall::Result<Option<usize>> {
+        if let Some(part_num) = part_num {
+            let part = self
+                .pt
+                .as_ref()
+                .ok_or(syscall::Error::new(EBADF))?
+                .partitions
+                .get(part_num)
+                .ok_or(syscall::Error::new(EBADF))?;
+
+            let block_size = u64::from(self.block_size());
+            if block >= part.size / block_size {
+                return Err(syscall::Error::new(EOVERFLOW));
+            }
+
+            let abs_block = part.start_lba + block;
+
+            self.disk.read(abs_block, buf)
+        } else {
+            self.disk.read(block, buf)
+        }
+    }
+
+    pub fn write(
+        &mut self,
+        part_num: Option<usize>,
+        block: u64,
+        buf: &[u8],
+    ) -> syscall::Result<Option<usize>> {
+        if let Some(part_num) = part_num {
+            let part = self
+                .pt
+                .as_ref()
+                .ok_or(syscall::Error::new(EBADF))?
+                .partitions
+                .get(part_num)
+                .ok_or(syscall::Error::new(EBADF))?;
+
+            let block_size = u64::from(self.block_size());
+            if block >= part.size / block_size {
+                return Err(syscall::Error::new(EOVERFLOW));
+            }
+
+            let abs_block = part.start_lba + block;
+
+            self.disk.write(abs_block, buf)
+        } else {
+            self.disk.write(block, buf)
+        }
+    }
 }
diff --git a/storage/ided/src/scheme.rs b/storage/ided/src/scheme.rs
index 167eb78a80..f41a227579 100644
--- a/storage/ided/src/scheme.rs
+++ b/storage/ided/src/scheme.rs
@@ -6,8 +6,8 @@ use std::sync::{Arc, Mutex};
 use driver_block::{Disk, DiskWrapper};
 use syscall::schemev2::NewFdFlags;
 use syscall::{
-    Error, Result, Stat, EACCES, EBADF, EISDIR, ENOENT, ENOLCK, EOVERFLOW, MODE_DIR, MODE_FILE,
-    O_DIRECTORY, O_STAT,
+    Error, Result, Stat, EACCES, EBADF, EISDIR, ENOENT, ENOLCK, MODE_DIR, MODE_FILE, O_DIRECTORY,
+    O_STAT,
 };
 
 use crate::ide::Channel;
@@ -265,31 +265,13 @@ impl SchemeBlock for DiskScheme {
             }
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_size();
-                disk.read(offset / u64::from(blk_len), buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.read(None, block, buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_size();
-
-                // validate that we're actually reading within the bounds of the partition
-                let rel_block = offset / u64::from(blksize);
-
-                let abs_block = {
-                    let pt = disk.pt.as_ref().ok_or(Error::new(EBADF))?;
-                    let partition = pt
-                        .partitions
-                        .get(part_num as usize)
-                        .ok_or(Error::new(EBADF))?;
-
-                    let abs_block = partition.start_lba + rel_block;
-                    if rel_block >= partition.size {
-                        return Err(Error::new(EOVERFLOW));
-                    }
-                    abs_block
-                };
-
-                disk.read(abs_block, buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.read(Some(part_num as usize), block, buf)
             }
         }
     }
@@ -299,31 +281,13 @@ impl SchemeBlock for DiskScheme {
             Handle::List(_) => Err(Error::new(EBADF)),
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(number).ok_or(Error::new(EBADF))?;
-                let blk_len = disk.block_size();
-                disk.write(offset / u64::from(blk_len), buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.write(None, block, buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(disk_num).ok_or(Error::new(EBADF))?;
-                let blksize = disk.block_size();
-
-                // validate that we're actually reading within the bounds of the partition
-                let rel_block = offset / u64::from(blksize as u64);
-
-                let abs_block = {
-                    let pt = disk.pt.as_ref().ok_or(Error::new(EBADF))?;
-                    let partition = pt
-                        .partitions
-                        .get(part_num as usize)
-                        .ok_or(Error::new(EBADF))?;
-
-                    let abs_block = partition.start_lba + rel_block;
-                    if rel_block >= partition.size {
-                        return Err(Error::new(EOVERFLOW));
-                    }
-                    abs_block
-                };
-
-                disk.write(abs_block, buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.write(Some(part_num as usize), block, buf)
             }
         }
     }
diff --git a/storage/nvmed/src/scheme.rs b/storage/nvmed/src/scheme.rs
index 48018b33e5..bd864441dd 100644
--- a/storage/nvmed/src/scheme.rs
+++ b/storage/nvmed/src/scheme.rs
@@ -8,8 +8,8 @@ use driver_block::{Disk, DiskWrapper};
 use redox_scheme::{CallerCtx, OpenResult, SchemeBlock};
 use syscall::schemev2::NewFdFlags;
 use syscall::{
-    Error, Result, Stat, EACCES, EBADF, EISDIR, ENOENT, ENOLCK, EOVERFLOW, MODE_DIR, MODE_FILE,
-    O_DIRECTORY, O_STAT,
+    Error, Result, Stat, EACCES, EBADF, EISDIR, ENOENT, ENOLCK, MODE_DIR, MODE_FILE, O_DIRECTORY,
+    O_STAT,
 };
 
 use crate::nvme::{Nvme, NvmeNamespace};
@@ -183,7 +183,7 @@ impl SchemeBlock for DiskScheme {
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(&number).ok_or(Error::new(EBADF))?;
                 stat.st_mode = MODE_FILE;
-                stat.st_blocks = disk.1.blocks;
+                stat.st_blocks = disk.disk().1.blocks;
                 stat.st_blksize = disk.block_size();
                 stat.st_size = disk.size();
                 Ok(Some(0))
@@ -270,28 +270,13 @@ impl SchemeBlock for DiskScheme {
             }
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(&number).ok_or(Error::new(EBADF))?;
-                let block_size = u64::from(disk.block_size());
-                disk.read(offset / block_size, buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.read(None, block, buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(&disk_num).ok_or(Error::new(EBADF))?;
-                let part = disk
-                    .pt
-                    .as_ref()
-                    .ok_or(Error::new(EBADF))?
-                    .partitions
-                    .get(part_num as usize)
-                    .ok_or(Error::new(EBADF))?;
-
-                let block_size = u64::from(disk.block_size());
-                let rel_block = offset / block_size;
-                if rel_block >= part.size {
-                    return Err(Error::new(EOVERFLOW));
-                }
-
-                let abs_block = part.start_lba + rel_block;
-
-                disk.read(abs_block, buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.read(Some(part_num as usize), block, buf)
             }
         }
     }
@@ -307,28 +292,13 @@ impl SchemeBlock for DiskScheme {
             Handle::List(_) => Err(Error::new(EBADF)),
             Handle::Disk(number) => {
                 let disk = self.disks.get_mut(&number).ok_or(Error::new(EBADF))?;
-                let block_size = u64::from(disk.block_size());
-                disk.write(offset / block_size, buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.write(None, block, buf)
             }
             Handle::Partition(disk_num, part_num) => {
                 let disk = self.disks.get_mut(&disk_num).ok_or(Error::new(EBADF))?;
-                let part = disk
-                    .pt
-                    .as_ref()
-                    .ok_or(Error::new(EBADF))?
-                    .partitions
-                    .get(part_num as usize)
-                    .ok_or(Error::new(EBADF))?;
-
-                let block_size = u64::from(disk.block_size());
-                let rel_block = offset / block_size;
-                if rel_block >= part.size {
-                    return Err(Error::new(EOVERFLOW));
-                }
-
-                let abs_block = part.start_lba + rel_block;
-
-                disk.write(abs_block, buf)
+                let block = offset / u64::from(disk.block_size());
+                disk.write(Some(part_num as usize), block, buf)
             }
         }
     }
diff --git a/storage/virtio-blkd/src/scheme.rs b/storage/virtio-blkd/src/scheme.rs
index f47abc33c6..974573da51 100644
--- a/storage/virtio-blkd/src/scheme.rs
+++ b/storage/virtio-blkd/src/scheme.rs
@@ -4,7 +4,7 @@ use std::fmt::Write;
 use std::sync::Arc;
 
 use common::dma::Dma;
-use driver_block::{Disk, DiskWrapper};
+use driver_block::DiskWrapper;
 
 use redox_scheme::CallerCtx;
 use redox_scheme::OpenResult;
@@ -239,28 +239,13 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
                 buf[..count].copy_from_slice(&src[..count]);
                 Ok(Some(count))
             }
-
-            Handle::Partition { number } => {
-                let part_table = self.disk.pt.as_ref().unwrap();
-                let part = part_table
-                    .partitions
-                    .get(number as usize)
-                    .ok_or(Error::new(EBADF))?;
-
-                // Get the offset in sectors.
-                let rel_block = offset / BLK_SIZE;
-                // if rel_block >= part.size {
-                //     return Err(Error::new(EOVERFLOW));
-                // }
-
-                let abs_block = part.start_lba + rel_block;
-
-                self.disk.read(abs_block, buf)
-            }
-
             Handle::Disk => {
                 let block = offset / u64::from(self.disk.block_size());
-                self.disk.read(block, buf)
+                self.disk.read(None, block, buf)
+            }
+            Handle::Partition { number } => {
+                let block = offset / u64::from(self.disk.block_size());
+                self.disk.read(Some(number as usize), block, buf)
             }
         }
     }
@@ -273,12 +258,15 @@ impl<'a> SchemeBlock for DiskScheme<'a> {
         _fcntl_flags: u32,
     ) -> syscall::Result<Option<usize>> {
         match *self.handles.get_mut(&id).ok_or(Error::new(EBADF))? {
+            Handle::List { .. } => Err(Error::new(EBADF)),
             Handle::Disk => {
                 let block = offset / u64::from(self.disk.block_size());
-                self.disk.write(block, buf)
+                self.disk.write(None, block, buf)
+            }
+            Handle::Partition { number } => {
+                let block = offset / u64::from(self.disk.block_size());
+                self.disk.write(Some(number as usize), block, buf)
             }
-
-            _ => todo!(),
         }
     }