use crate::BlockIo;
use bytemuck::{bytes_of, from_bytes};
use core::fmt::{self, Debug, Display, Formatter};
use core::mem;
use gpt_disk_types::{
GptHeader, GptPartitionEntry, GptPartitionEntryArray,
GptPartitionEntryArrayError, GptPartitionEntryArrayLayout, Lba,
MasterBootRecord,
};
struct GptPartitionEntryIter<'disk, 'buf, Io: BlockIo> {
disk: &'disk mut Disk<Io>,
block_buf: &'buf mut [u8],
layout: GptPartitionEntryArrayLayout,
next_index: u32,
current_lba: Lba,
byte_offset_within_lba: usize,
entry_size: usize,
}
impl<'disk, 'buf, Io: BlockIo> GptPartitionEntryIter<'disk, 'buf, Io> {
fn new(
disk: &'disk mut Disk<Io>,
layout: GptPartitionEntryArrayLayout,
block_buf: &'buf mut [u8],
) -> Result<Self, DiskError<Io::Error>> {
let mut iter = Self {
disk,
block_buf,
next_index: 0,
current_lba: layout.start_lba,
byte_offset_within_lba: 0,
layout,
entry_size: layout
.entry_size
.to_usize()
.ok_or(DiskError::Overflow)?,
};
iter.set_current_lba(iter.current_lba)?;
Ok(iter)
}
fn set_current_lba(
&mut self,
lba: Lba,
) -> Result<(), DiskError<Io::Error>> {
self.current_lba = lba;
self.byte_offset_within_lba = 0;
Ok(self.disk.io.read_blocks(self.current_lba, self.block_buf)?)
}
fn read_current_entry(&mut self) -> Option<<Self as Iterator>::Item> {
let entry_bytes = self.block_buf.get(
self.byte_offset_within_lba
..self.byte_offset_within_lba + self.entry_size,
)?;
self.byte_offset_within_lba += self.entry_size;
self.next_index += 1;
Some(Ok(*from_bytes::<GptPartitionEntry>(
&entry_bytes[..mem::size_of::<GptPartitionEntry>()],
)))
}
}
impl<'disk, 'buf, Io: BlockIo> Iterator
for GptPartitionEntryIter<'disk, 'buf, Io>
{
type Item = Result<GptPartitionEntry, DiskError<Io::Error>>;
fn next(&mut self) -> Option<Self::Item> {
if self.next_index >= self.layout.num_entries {
return None;
}
if let Some(entry) = self.read_current_entry() {
Some(entry)
} else {
let next_lba = Lba(self.current_lba.to_u64() + 1);
if let Err(err) = self.set_current_lba(next_lba) {
Some(Err(err))
} else {
self.read_current_entry()
}
}
}
}
pub trait Captures<'a, 'b> {}
impl<'a, 'b, T: ?Sized> Captures<'a, 'b> for T {}
#[allow(clippy::module_name_repetitions)]
#[derive(Debug)]
pub enum DiskError<IoError: Debug + Display> {
BufferTooSmall,
Overflow,
BlockSizeSmallerThanPartitionEntry,
Io(IoError),
}
impl<IoError> From<IoError> for DiskError<IoError>
where
IoError: Debug + Display,
{
fn from(err: IoError) -> Self {
DiskError::Io(err)
}
}
impl<IoError> Display for DiskError<IoError>
where
IoError: Debug + Display,
{
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
match self {
Self::BufferTooSmall => f.write_str("storage buffer is too small"),
Self::Overflow => f.write_str("numeric overflow occurred"),
Self::BlockSizeSmallerThanPartitionEntry => {
f.write_str("partition entries are larger than a single block")
}
Self::Io(io) => Display::fmt(io, f),
}
}
}
pub struct Disk<Io: BlockIo> {
io: Io,
}
impl<Io: BlockIo> Disk<Io> {
pub fn new(io: Io) -> Result<Self, DiskError<Io::Error>> {
Ok(Self { io })
}
fn clip_block_buf_size<'buf>(
&self,
block_buf: &'buf mut [u8],
) -> Result<&'buf mut [u8], DiskError<Io::Error>> {
if let Some(block_size) = self.io.block_size().to_usize() {
block_buf
.get_mut(..block_size)
.ok_or(DiskError::BufferTooSmall)
} else {
Err(DiskError::BufferTooSmall)
}
}
pub fn read_primary_gpt_header(
&mut self,
block_buf: &mut [u8],
) -> Result<GptHeader, DiskError<Io::Error>> {
self.read_gpt_header(Lba(1), block_buf)
}
pub fn read_secondary_gpt_header(
&mut self,
block_buf: &mut [u8],
) -> Result<GptHeader, DiskError<Io::Error>> {
let num_blocks = self.io.num_blocks()?;
let last_block =
Lba(num_blocks.checked_sub(1).ok_or(DiskError::Overflow)?);
self.read_gpt_header(last_block, block_buf)
}
pub fn read_gpt_header(
&mut self,
lba: Lba,
mut block_buf: &mut [u8],
) -> Result<GptHeader, DiskError<Io::Error>> {
block_buf = self.clip_block_buf_size(block_buf)?;
self.io.read_blocks(lba, block_buf)?;
let bytes = block_buf
.get(..mem::size_of::<GptHeader>())
.unwrap();
Ok(*from_bytes(bytes))
}
pub fn read_gpt_partition_entry_array<'buf>(
&mut self,
layout: GptPartitionEntryArrayLayout,
storage: &'buf mut [u8],
) -> Result<GptPartitionEntryArray<'buf>, DiskError<Io::Error>> {
let mut entry_array =
GptPartitionEntryArray::new(layout, self.io.block_size(), storage)
.map_err(|err| match err {
GptPartitionEntryArrayError::BufferTooSmall => {
DiskError::BufferTooSmall
}
GptPartitionEntryArrayError::Overflow => {
DiskError::Overflow
}
})?;
self.io
.read_blocks(layout.start_lba, entry_array.storage_mut())?;
Ok(entry_array)
}
pub fn write_gpt_partition_entry_array(
&mut self,
entry_array: &GptPartitionEntryArray,
) -> Result<(), DiskError<Io::Error>> {
Ok(self.io.write_blocks(
entry_array.layout().start_lba,
entry_array.storage(),
)?)
}
pub fn gpt_partition_entry_array_iter<'disk, 'buf>(
&'disk mut self,
layout: GptPartitionEntryArrayLayout,
mut block_buf: &'buf mut [u8],
) -> Result<
impl Iterator<Item = Result<GptPartitionEntry, DiskError<Io::Error>>>
+ Captures<'disk, 'buf>,
DiskError<Io::Error>,
> {
block_buf = self.clip_block_buf_size(block_buf)?;
let entry_size =
layout.entry_size.to_usize().ok_or(DiskError::Overflow)?;
if entry_size > block_buf.len() {
return Err(DiskError::BlockSizeSmallerThanPartitionEntry);
}
GptPartitionEntryIter::<'disk, 'buf>::new(self, layout, block_buf)
}
pub fn write_protective_mbr(
&mut self,
block_buf: &mut [u8],
) -> Result<(), DiskError<Io::Error>> {
let mbr = MasterBootRecord::protective_mbr(self.io.num_blocks()?);
self.write_mbr(&mbr, block_buf)
}
pub fn write_mbr(
&mut self,
mbr: &MasterBootRecord,
mut block_buf: &mut [u8],
) -> Result<(), DiskError<Io::Error>> {
block_buf = self.clip_block_buf_size(block_buf)?;
let mbr_bytes = bytes_of(mbr);
assert!(block_buf.len() >= mbr_bytes.len());
{
let (left, right) = block_buf.split_at_mut(mbr_bytes.len());
left.copy_from_slice(mbr_bytes);
right.fill(0);
}
self.io.write_blocks(Lba(0), block_buf)?;
Ok(())
}
pub fn write_primary_gpt_header(
&mut self,
header: &GptHeader,
block_buf: &mut [u8],
) -> Result<(), DiskError<Io::Error>> {
self.write_gpt_header(Lba(1), header, block_buf)
}
pub fn write_secondary_gpt_header(
&mut self,
header: &GptHeader,
block_buf: &mut [u8],
) -> Result<(), DiskError<Io::Error>> {
let num_blocks = self.io.num_blocks()?;
let last_block =
Lba(num_blocks.checked_sub(1).ok_or(DiskError::Overflow)?);
self.write_gpt_header(last_block, header, block_buf)
}
pub fn write_gpt_header(
&mut self,
lba: Lba,
header: &GptHeader,
mut block_buf: &mut [u8],
) -> Result<(), DiskError<Io::Error>> {
block_buf = self.clip_block_buf_size(block_buf)?;
let header_bytes = bytes_of(header);
assert!(block_buf.len() >= header_bytes.len());
{
let (left, right) = block_buf.split_at_mut(header_bytes.len());
left.copy_from_slice(header_bytes);
right.fill(0);
}
self.io.write_blocks(lba, block_buf)?;
Ok(())
}
pub fn flush(&mut self) -> Result<(), DiskError<Io::Error>> {
Ok(self.io.flush()?)
}
}
impl<Io: BlockIo> Drop for Disk<Io> {
fn drop(&mut self) {
let _r = self.flush();
}
}