#![deny(missing_docs)]
use bincode::config::legacy;
use bincode::error::{DecodeError, EncodeError};
use bincode::serde::{decode_from_std_read, encode_into_std_write};
use bitvec::prelude::*;
use serde::de::{SeqAccess, Visitor};
use serde::ser::SerializeTuple;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_big_array::BigArray;
use std::io::{Read, Seek, SeekFrom, Write};
use std::iter::{once, repeat};
use std::ops::{Index, IndexMut};
use thiserror::Error;
const DEFAULT_ALIGN: u32 = 2048;
const MAX_ALIGN: u32 = 16384;
const FIRST_USABLE_LBA: u32 = 1;
const BOOT_SIGNATURE: [u8; 2] = [0x55, 0xaa];
pub const BOOT_ACTIVE: u8 = 0x80;
pub const BOOT_INACTIVE: u8 = 0x00;
pub type Result<T> = std::result::Result<T, Error>;
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum Error {
#[error("exceeded the maximum limit of CHS")]
LBAExceedsMaximumCHS,
#[error("exceeded the maximum number of cylinders on disk")]
LBAExceedsMaximumCylinders,
#[error("deserialization failed")]
Deserialize(#[from] DecodeError),
#[error("Serialization failed")]
Serialize(#[from] EncodeError),
#[error("generic I/O error")]
Io(#[from] std::io::Error),
#[error("inconsistent extended boot record")]
InconsistentEBR,
#[error("no extended partition")]
NoExtendedPartition,
#[error("EBR starts before the extended partition")]
EBRStartsBeforeExtendedPartition,
#[error("EBR starts too close to the end of the extended partition")]
EBRStartsTooCloseToTheEndOfExtendedPartition,
#[error("EBR ends after the extended partition")]
EBREndsAfterExtendedPartition,
#[error("not enough sectors to create a logical partition")]
NotEnoughSectorsToCreateLogicalPartition,
#[error("partition not found")]
PartitionNotFound,
#[error("no space left")]
NoSpaceLeft,
#[error("invalid MBR signature")]
InvalidSignature,
#[error("partition has invalid boot flag")]
InvalidBootFlag,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MBR {
pub sector_size: u32,
pub header: MBRHeader,
pub logical_partitions: Vec<LogicalPartition>,
pub align: u32,
pub cylinders: u16,
pub heads: u8,
pub sectors: u8,
pub disk_size: u32,
}
impl MBR {
pub fn iter(&self) -> impl Iterator<Item = (usize, &MBRPartitionEntry)> {
self.header.iter().chain(
self.logical_partitions
.iter()
.map(|x| &x.partition)
.enumerate()
.map(|(i, x)| (i + 5, x)),
)
}
pub fn iter_mut(&mut self) -> impl Iterator<Item = (usize, &mut MBRPartitionEntry)> {
let mut partitions: Vec<_> = vec![
&mut self.header.partition_1,
&mut self.header.partition_2,
&mut self.header.partition_3,
&mut self.header.partition_4,
];
partitions.extend(self.logical_partitions.iter_mut().map(|x| &mut x.partition));
partitions.into_iter().enumerate().map(|(i, x)| (i + 1, x))
}
pub fn get(&self, i: usize) -> Option<&MBRPartitionEntry> {
match i {
0 => None,
1 => Some(&self.header.partition_1),
2 => Some(&self.header.partition_2),
3 => Some(&self.header.partition_3),
4 => Some(&self.header.partition_4),
i => self.logical_partitions.get(i - 5).map(|x| &x.partition),
}
}
pub fn get_mut(&mut self, i: usize) -> Option<&mut MBRPartitionEntry> {
match i {
0 => None,
1 => Some(&mut self.header.partition_1),
2 => Some(&mut self.header.partition_2),
3 => Some(&mut self.header.partition_3),
4 => Some(&mut self.header.partition_4),
i => self
.logical_partitions
.get_mut(i - 5)
.map(|x| &mut x.partition),
}
}
pub fn len(&self) -> usize {
4 + self.logical_partitions.len()
}
pub fn is_empty(&self) -> bool {
false
}
pub fn new_from<S>(seeker: &mut S, sector_size: u32, disk_signature: [u8; 4]) -> Result<MBR>
where
S: Seek,
{
let disk_size = u32::try_from(seeker.seek(SeekFrom::End(0))? / u64::from(sector_size))
.unwrap_or(u32::MAX);
let header = MBRHeader::new(disk_signature);
Ok(MBR {
sector_size,
header,
logical_partitions: Vec::new(),
align: DEFAULT_ALIGN,
cylinders: 0,
heads: 0,
sectors: 0,
disk_size,
})
}
pub fn read_from<R>(mut reader: &mut R, sector_size: u32) -> Result<MBR>
where
R: Read + Seek + ?Sized,
{
let disk_size = u32::try_from(reader.seek(SeekFrom::End(0))? / u64::from(sector_size))
.unwrap_or(u32::MAX);
let header = MBRHeader::read_from(&mut reader)?;
let mut logical_partitions = Vec::new();
if let Some(extended) = header.get_extended_partition() {
let mut relative_ebr_lba = 0;
let mut ebr_sectors = None;
let mut ebr_first_chs = extended.first_chs;
let mut ebr_last_chs = None;
loop {
let offset =
(extended.starting_lba as u64 + relative_ebr_lba as u64) * sector_size as u64;
reader.seek(SeekFrom::Start(offset))?;
let (partition, next, bootstrap_code) = match EBRHeader::read_from(&mut reader) {
Ok(ebr) => ebr.unwrap(),
Err(err) => {
if relative_ebr_lba == 0 {
break;
} else {
return Err(err);
}
}
};
let absolute_ebr_lba = extended.starting_lba + relative_ebr_lba;
logical_partitions.push(LogicalPartition {
partition: MBRPartitionEntry {
starting_lba: partition.starting_lba + absolute_ebr_lba,
..partition
},
absolute_ebr_lba,
ebr_sectors,
ebr_first_chs,
ebr_last_chs,
bootstrap_code,
});
if next.starting_lba > 0 && relative_ebr_lba >= next.starting_lba {
return Err(Error::InconsistentEBR);
}
relative_ebr_lba = next.starting_lba;
ebr_sectors = Some(next.sectors);
ebr_first_chs = next.first_chs;
ebr_last_chs = Some(next.last_chs);
if relative_ebr_lba == 0 {
break;
}
}
}
let align = MBR::find_alignment(&header, &logical_partitions);
Ok(MBR {
sector_size,
header,
logical_partitions,
align,
cylinders: 0,
heads: 0,
sectors: 0,
disk_size,
})
}
pub fn update_from<S>(&mut self, seeker: &mut S) -> Result<()>
where
S: Seek + ?Sized,
{
self.disk_size =
u32::try_from(seeker.seek(SeekFrom::End(0))? / u64::from(self.sector_size))
.unwrap_or(u32::MAX);
Ok(())
}
fn find_alignment(header: &MBRHeader, logical_partitions: &[LogicalPartition]) -> u32 {
let lbas = header
.iter()
.map(|(_, x)| x)
.chain(logical_partitions.iter().map(|x| &x.partition))
.filter(|x| x.is_used())
.map(|x| x.starting_lba)
.collect::<Vec<_>>();
if lbas.is_empty() {
return DEFAULT_ALIGN;
}
if lbas.len() == 1 && lbas[0] == FIRST_USABLE_LBA {
return FIRST_USABLE_LBA;
}
(1..=MAX_ALIGN.min(*lbas.iter().max().unwrap_or(&1)))
.filter(|div| lbas.iter().all(|x| x % div == 0))
.max()
.unwrap()
}
pub fn check_geometry(&self) -> bool {
self.cylinders > 0
&& self.cylinders <= 1023
&& self.heads > 0
&& self.sectors > 0
&& self.sectors <= 63
}
pub fn write_into<W>(&mut self, mut writer: &mut W) -> Result<()>
where
W: Write + Seek + ?Sized,
{
self.header.write_into(&mut writer)?;
if let Some(extended) = self.header.get_extended_partition() {
if let Some(first) = self.logical_partitions.get_mut(0) {
first.absolute_ebr_lba = extended.starting_lba;
first.ebr_sectors = None;
first.ebr_last_chs = None;
}
if self.check_geometry() {
for l in self.logical_partitions.iter_mut() {
l.update_chs(self.cylinders, self.heads, self.sectors)?;
}
}
#[derive(Serialize)]
struct BootstrapCode446(#[serde(with = "BigArray")] [u8; 446]);
let next_logical_partitions = self
.logical_partitions
.iter()
.skip(1)
.map(Some)
.chain(once(None));
for (l, next) in self.logical_partitions.iter().zip(next_logical_partitions) {
let partition = MBRPartitionEntry {
starting_lba: l.partition.starting_lba.saturating_sub(l.absolute_ebr_lba),
..l.partition
};
partition.check()?;
writer.seek(SeekFrom::Start(u64::from(
l.absolute_ebr_lba * self.sector_size,
)))?;
encode_into_std_write(BootstrapCode446(l.bootstrap_code), &mut writer, legacy())?;
encode_into_std_write(&partition, &mut writer, legacy())?;
if let Some(next) = next {
encode_into_std_write(
&MBRPartitionEntry {
boot: BOOT_INACTIVE,
first_chs: next.ebr_first_chs,
sys: extended.sys,
last_chs: next.ebr_last_chs.unwrap(),
starting_lba: next
.absolute_ebr_lba
.saturating_sub(extended.starting_lba),
sectors: next.ebr_sectors.unwrap(),
},
&mut writer,
legacy(),
)?;
} else {
encode_into_std_write(MBRPartitionEntry::empty(), &mut writer, legacy())?;
}
writer.write_all(&[0; 16 * 2])?;
encode_into_std_write(BOOT_SIGNATURE, &mut writer, legacy())?;
}
}
Ok(())
}
pub fn get_cylinder_size(&self) -> u32 {
u32::from(self.heads) * u32::from(self.sectors)
}
pub fn find_at_sector(&self, sector: u32) -> Option<usize> {
let between = |sector, start, len| sector >= start && sector < start + len;
let primary = self
.header
.iter()
.find(|(_, x)| x.is_used() && between(sector, x.starting_lba, x.sectors));
match primary {
Some((_, x)) if x.is_extended() => self
.logical_partitions
.iter()
.enumerate()
.find(|(_, x)| {
x.partition.is_used()
&& between(sector, x.partition.starting_lba, x.partition.sectors)
})
.map(|(i, _)| 5 + i),
Some((i, _)) => Some(i),
None => None,
}
}
pub fn remove_at_sector(&mut self, sector: u32) -> Result<()> {
let i = self
.find_at_sector(sector)
.ok_or(Error::PartitionNotFound)?;
if i >= 5 {
self.remove(i);
} else {
if self[i].is_extended() {
self.logical_partitions.clear();
}
self[i] = MBRPartitionEntry::empty();
}
Ok(())
}
pub fn find_free_sectors(&self) -> Vec<(u32, u32)> {
assert!(self.align > 0, "align must be greater than 0");
let collect_free_sectors = |positions: Vec<u32>| {
positions
.chunks(2)
.map(|x| (x[0] + 1, x[1] - x[0] - 1))
.filter(|(_, l)| *l > 0)
.map(|(i, l)| (i, l, ((i - 1) / self.align + 1) * self.align - i))
.map(|(i, l, s)| (i + s, l.saturating_sub(s)))
.filter(|(_, l)| *l > 0)
.collect::<Vec<_>>()
};
let mut positions = vec![0];
for (_, partition) in self.header.iter().filter(|(_, x)| x.is_used()) {
positions.push(partition.starting_lba);
positions.push(partition.starting_lba + partition.sectors - 1);
}
positions.push(self.disk_size);
positions.sort_unstable();
let mut res = collect_free_sectors(positions);
if let Some(extended) = self.header.get_extended_partition() {
let mut positions = vec![extended.starting_lba];
for l in self
.logical_partitions
.iter()
.filter(|x| x.partition.is_used())
{
let starting_lba = l.absolute_ebr_lba + l.partition.starting_lba;
positions.push(starting_lba);
positions.push(starting_lba + l.partition.sectors - 1);
}
positions.push(extended.starting_lba + extended.sectors);
positions.sort_unstable();
res.extend(collect_free_sectors(positions));
}
res
}
pub fn find_first_place(&self, size: u32) -> Option<u32> {
self.find_free_sectors()
.iter()
.find(|(_, l)| *l >= size)
.map(|(i, _)| *i)
}
pub fn find_last_place(&self, size: u32) -> Option<u32> {
self.find_free_sectors()
.iter()
.filter(|(_, l)| *l >= size)
.last()
.map(|(i, l)| (i + l - size) / self.align * self.align)
}
pub fn find_optimal_place(&self, size: u32) -> Option<u32> {
let mut slots = self
.find_free_sectors()
.into_iter()
.filter(|(_, l)| *l >= size)
.collect::<Vec<_>>();
slots.sort_by(|(_, l1), (_, l2)| l1.cmp(l2));
slots.first().map(|&(i, _)| i)
}
pub fn get_maximum_partition_size(&self) -> Result<u32> {
self.find_free_sectors()
.into_iter()
.map(|(_, l)| l / self.align * self.align)
.max()
.ok_or(Error::NoSpaceLeft)
}
pub fn get_maximum_partition_size_for(&mut self, i: usize) -> Result<u32> {
let spots = self.find_free_sectors();
let p = self.get_mut(i).ok_or(Error::PartitionNotFound)?;
let free_sectors = spots
.into_iter()
.find(|(i, _)| *i == p.starting_lba + p.sectors)
.map(|(_, l)| l)
.unwrap_or(0);
Ok(p.sectors + free_sectors)
}
pub fn push(
&mut self,
sys: u8,
mut starting_lba: u32,
mut sectors: u32,
) -> Result<&mut LogicalPartition> {
let extended = self
.header
.get_extended_partition()
.ok_or(Error::NoExtendedPartition)?;
starting_lba = ((starting_lba - 1) / self.align + 1) * self.align;
sectors = ((sectors - 1) / self.align + 1) * self.align;
if sectors < 2 * self.align {
return Err(Error::NotEnoughSectorsToCreateLogicalPartition);
}
let mut l = LogicalPartition {
partition: MBRPartitionEntry {
boot: BOOT_INACTIVE,
first_chs: CHS::empty(),
sys,
last_chs: CHS::empty(),
starting_lba: starting_lba + self.align,
sectors: sectors - self.align,
},
absolute_ebr_lba: starting_lba,
ebr_sectors: if self.logical_partitions.is_empty() {
None
} else {
Some(sectors)
},
ebr_first_chs: CHS::empty(),
ebr_last_chs: if self.logical_partitions.is_empty() {
None
} else {
Some(CHS::empty())
},
bootstrap_code: [0; 446],
};
if l.absolute_ebr_lba < extended.starting_lba {
return Err(Error::EBRStartsBeforeExtendedPartition);
}
if l.absolute_ebr_lba > extended.starting_lba + extended.sectors - 2 * self.align {
return Err(Error::EBRStartsTooCloseToTheEndOfExtendedPartition);
}
if let Some(ebr_sectors) = l.ebr_sectors {
let ending_ebr_lba = l.absolute_ebr_lba + ebr_sectors - 1;
if ending_ebr_lba > extended.starting_lba + extended.sectors - 1 {
return Err(Error::EBREndsAfterExtendedPartition);
}
}
if self.check_geometry() {
l.update_chs(self.cylinders, self.heads, self.sectors)?;
}
self.logical_partitions.push(l);
Ok(self.logical_partitions.last_mut().unwrap())
}
pub fn remove(&mut self, index: usize) -> LogicalPartition {
assert!(index >= 5, "logical partitions start at 5");
self.logical_partitions.remove(index - 5)
}
}
impl Index<usize> for MBR {
type Output = MBRPartitionEntry;
fn index(&self, i: usize) -> &Self::Output {
assert!(i != 0, "invalid partition index: 0");
self.get(i).expect("invalid partition")
}
}
impl IndexMut<usize> for MBR {
fn index_mut(&mut self, i: usize) -> &mut Self::Output {
assert!(i != 0, "invalid partition index: 0");
self.get_mut(i).expect("invalid partition")
}
}
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
pub struct MBRHeader {
#[serde(with = "BigArray")]
pub bootstrap_code: [u8; 440],
pub disk_signature: [u8; 4],
pub copy_protected: [u8; 2],
pub partition_1: MBRPartitionEntry,
pub partition_2: MBRPartitionEntry,
pub partition_3: MBRPartitionEntry,
pub partition_4: MBRPartitionEntry,
pub boot_signature: [u8; 2],
}
impl MBRHeader {
pub fn is_copy_protected(&self) -> Option<bool> {
match self.copy_protected {
[0x00, 0x00] => Some(false),
[0x5a, 0x5a] => Some(true),
_ => None,
}
}
pub fn get(&self, i: usize) -> Option<&MBRPartitionEntry> {
match i {
1 => Some(&self.partition_1),
2 => Some(&self.partition_2),
3 => Some(&self.partition_3),
4 => Some(&self.partition_4),
_ => None,
}
}
pub fn get_mut(&mut self, i: usize) -> Option<&mut MBRPartitionEntry> {
match i {
1 => Some(&mut self.partition_1),
2 => Some(&mut self.partition_2),
3 => Some(&mut self.partition_3),
4 => Some(&mut self.partition_4),
_ => None,
}
}
pub fn iter(&self) -> impl Iterator<Item = (usize, &MBRPartitionEntry)> {
vec![
&self.partition_1,
&self.partition_2,
&self.partition_3,
&self.partition_4,
]
.into_iter()
.enumerate()
.map(|(i, x)| (i + 1, x))
}
pub fn iter_mut(&mut self) -> impl Iterator<Item = (usize, &mut MBRPartitionEntry)> {
vec![
&mut self.partition_1,
&mut self.partition_2,
&mut self.partition_3,
&mut self.partition_4,
]
.into_iter()
.enumerate()
.map(|(i, x)| (i + 1, x))
}
fn get_extended_partition(&self) -> Option<&MBRPartitionEntry> {
self.iter().find(|(_, x)| x.is_extended()).map(|(_, x)| x)
}
pub fn read_from<R>(mut reader: &mut R) -> Result<MBRHeader>
where
R: Read + Seek + ?Sized,
{
reader.seek(SeekFrom::Start(0))?;
let header: Self = decode_from_std_read(&mut reader, legacy())?;
header.check()?;
Ok(header)
}
pub fn new(disk_signature: [u8; 4]) -> MBRHeader {
MBRHeader {
bootstrap_code: [0; 440],
disk_signature,
copy_protected: [0x00, 0x00],
partition_1: MBRPartitionEntry::empty(),
partition_2: MBRPartitionEntry::empty(),
partition_3: MBRPartitionEntry::empty(),
partition_4: MBRPartitionEntry::empty(),
boot_signature: BOOT_SIGNATURE,
}
}
pub fn write_into<W>(&self, mut writer: &mut W) -> Result<()>
where
W: Write + Seek + ?Sized,
{
self.check()?;
writer.seek(SeekFrom::Start(0))?;
encode_into_std_write(self, &mut writer, legacy())?;
Ok(())
}
fn check(&self) -> Result<()> {
if self.boot_signature != BOOT_SIGNATURE {
return Err(Error::InvalidSignature);
}
self.iter().try_for_each(|(_, partition)| partition.check())
}
}
impl Index<usize> for MBRHeader {
type Output = MBRPartitionEntry;
fn index(&self, i: usize) -> &Self::Output {
assert_ne!(i, 0, "invalid partition index: 0");
self.get(i).expect("invalid partition")
}
}
impl IndexMut<usize> for MBRHeader {
fn index_mut(&mut self, i: usize) -> &mut Self::Output {
assert_ne!(i, 0, "invalid partition index: 0");
self.get_mut(i).expect("invalid partition")
}
}
#[derive(Debug, Clone, Deserialize)]
struct EBRHeader {
#[serde(with = "BigArray")]
bootstrap_code: [u8; 446],
partition_1: MBRPartitionEntry,
partition_2: MBRPartitionEntry,
_unused_partition_3: [u8; 16],
_unused_partition_4: [u8; 16],
boot_signature: [u8; 2],
}
impl EBRHeader {
fn read_from<R>(reader: &mut R) -> Result<EBRHeader>
where
R: Read,
{
let header: Self = decode_from_std_read(reader, legacy())?;
header.check()?;
Ok(header)
}
fn check(&self) -> Result<()> {
if self.boot_signature != BOOT_SIGNATURE {
return Err(Error::InvalidSignature);
}
self.partition_1.check()?;
self.partition_2.check()
}
fn unwrap(self) -> (MBRPartitionEntry, MBRPartitionEntry, [u8; 446]) {
(self.partition_1, self.partition_2, self.bootstrap_code)
}
}
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
pub struct MBRPartitionEntry {
pub boot: u8,
pub first_chs: CHS,
pub sys: u8,
pub last_chs: CHS,
pub starting_lba: u32,
pub sectors: u32,
}
impl MBRPartitionEntry {
pub fn empty() -> MBRPartitionEntry {
MBRPartitionEntry {
boot: BOOT_INACTIVE,
first_chs: CHS::empty(),
sys: 0,
last_chs: CHS::empty(),
starting_lba: 0,
sectors: 0,
}
}
pub fn is_used(&self) -> bool {
self.sys > 0
}
pub fn is_unused(&self) -> bool {
!self.is_used()
}
pub fn is_extended(&self) -> bool {
self.sys == 0x05
|| self.sys == 0x0f
|| self.sys == 0x85
|| self.sys == 0xc5
|| self.sys == 0xd5
}
pub fn is_active(&self) -> bool {
self.boot == BOOT_ACTIVE
}
fn check(&self) -> Result<()> {
if self.boot != BOOT_ACTIVE && self.boot != BOOT_INACTIVE {
return Err(Error::InvalidBootFlag);
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LogicalPartition {
pub partition: MBRPartitionEntry,
pub absolute_ebr_lba: u32,
pub ebr_sectors: Option<u32>,
pub ebr_first_chs: CHS,
pub ebr_last_chs: Option<CHS>,
pub bootstrap_code: [u8; 446],
}
impl LogicalPartition {
pub fn update_chs(&mut self, cylinders: u16, heads: u8, sectors: u8) -> Result<()> {
self.partition.first_chs =
CHS::from_lba_exact(self.partition.starting_lba, cylinders, heads, sectors)?;
self.partition.last_chs = CHS::from_lba_exact(
self.partition.starting_lba + self.partition.sectors - 1,
cylinders,
heads,
sectors,
)?;
self.ebr_first_chs = CHS::from_lba_exact(self.absolute_ebr_lba, cylinders, heads, sectors)?;
self.ebr_last_chs = if let Some(ebr_sectors) = self.ebr_sectors {
Some(CHS::from_lba_exact(
self.absolute_ebr_lba + ebr_sectors - 1,
cylinders,
heads,
sectors,
)?)
} else {
None
};
Ok(())
}
}
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct CHS {
pub cylinder: u16,
pub head: u8,
pub sector: u8,
}
impl CHS {
pub fn new(cylinder: u16, head: u8, sector: u8) -> CHS {
CHS {
cylinder,
head,
sector,
}
}
pub fn empty() -> CHS {
CHS {
cylinder: 0,
head: 0,
sector: 0,
}
}
pub fn from_lba_exact(lba: u32, cylinders: u16, heads: u8, sectors: u8) -> Result<CHS> {
let cylinders = u32::from(cylinders);
let heads = u32::from(heads);
let sectors = u32::from(sectors);
let cylinder_size = heads * sectors;
let cylinder = lba / cylinder_size;
let rem = lba % cylinder_size;
let head = rem / sectors;
let sector = rem % sectors + 1;
if cylinder > 1023 {
return Err(Error::LBAExceedsMaximumCHS);
}
if cylinder > cylinders {
return Err(Error::LBAExceedsMaximumCylinders);
}
Ok(CHS {
cylinder: u16::try_from(cylinder).unwrap(),
head: u8::try_from(head).unwrap(),
sector: u8::try_from(sector).unwrap(),
})
}
pub fn from_lba_aligned(lba: u32, cylinders: u16, heads: u8, sectors: u8) -> Result<CHS> {
let cylinders = u32::from(cylinders);
let heads = u32::from(heads);
let sectors = u32::from(sectors);
let cylinder_size = heads * sectors;
let cylinder = ((lba - 1) / cylinder_size) + 1;
if cylinder > 1023 {
return Err(Error::LBAExceedsMaximumCHS);
}
if cylinder > cylinders {
return Err(Error::LBAExceedsMaximumCylinders);
}
Ok(CHS {
cylinder: u16::try_from(cylinder).unwrap(),
head: 0,
sector: 1,
})
}
pub fn to_lba(self, heads: u8, sectors: u8) -> u32 {
let heads = u32::from(heads);
let sectors = u32::from(sectors);
let c = u32::from(self.cylinder);
let h = u32::from(self.head);
let s = u32::from(self.sector);
c * (heads * sectors) + h * sectors + s - 1
}
pub fn is_empty(self) -> bool {
self.cylinder == 0 && self.head == 0 && self.sector == 0
}
pub fn is_valid(self, cylinders: u16, heads: u8, sectors: u8) -> bool {
self.sector > 0 && self.sector <= sectors && self.head < heads && self.cylinder < cylinders
}
}
struct CHSVisitor;
impl<'de> Visitor<'de> for CHSVisitor {
type Value = CHS;
fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
formatter.write_str("CHS addressing")
}
fn visit_seq<A>(self, mut seq: A) -> std::result::Result<CHS, A::Error>
where
A: SeqAccess<'de>,
{
let head = BitVec::<u8, Msb0>::from_vec(vec![seq.next_element::<u8>()?.unwrap()]);
let mut bv = BitVec::<u8, Msb0>::from_vec(vec![seq.next_element::<u8>()?.unwrap()]);
let mut cylinder = BitVec::<u16, Msb0>::with_capacity(10);
cylinder.extend(repeat(false).take(6));
cylinder.extend(bv.drain(..2));
cylinder.extend(BitVec::<u8, Msb0>::from_vec(vec![seq
.next_element::<u8>()?
.unwrap()]));
let mut sector = BitVec::<u8, Msb0>::with_capacity(8);
sector.push(false);
sector.push(false);
sector.extend(bv.drain(..));
Ok(CHS {
cylinder: cylinder.as_raw_slice()[0],
head: head.as_raw_slice()[0],
sector: sector.as_raw_slice()[0],
})
}
}
impl<'de> Deserialize<'de> for CHS {
fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
where
D: Deserializer<'de>,
{
deserializer.deserialize_tuple(3, CHSVisitor)
}
}
impl Serialize for CHS {
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
where
S: Serializer,
{
let mut bv = BitVec::<u8, Msb0>::from_vec(vec![self.head]);
let mut sector = BitVec::<u8, Msb0>::from_vec(vec![self.sector]);
let mut cylinder = BitVec::<u16, Msb0>::from_vec(vec![self.cylinder]);
bv.extend(cylinder.drain(..8).skip(6));
bv.extend(sector.drain(2..));
bv.extend(cylinder.drain(..));
let mut seq = serializer.serialize_tuple(3)?;
for x in bv.as_raw_slice() {
seq.serialize_element(&x)?;
}
seq.end()
}
}
#[cfg(test)]
#[allow(clippy::cognitive_complexity)]
mod tests {
use super::*;
use bincode::serde::{decode_from_slice, encode_into_slice};
use std::fs::File;
use std::io::Cursor;
const DISK1: &str = "tests/fixtures/disk1.img";
const DISK2: &str = "tests/fixtures/disk2.img";
#[test]
fn deserialize_maximum_chs_value() {
let chs: CHS = decode_from_slice(&[0xff, 0xff, 0xff], legacy()).unwrap().0;
assert_eq!(
chs,
CHS {
cylinder: 1023,
head: 255,
sector: 63,
}
);
}
#[test]
fn serialize_maximum_chs_value() {
let chs = CHS {
cylinder: 1023,
head: 255,
sector: 63,
};
let mut slice = [0; 3];
encode_into_slice(chs, &mut slice, legacy()).unwrap();
for element in slice {
assert_eq!(element, 0xff);
}
}
#[test]
fn serialize_and_deserialize_some_chs_value() {
let chs: CHS = decode_from_slice(&[0xaa, 0xaa, 0xaa], legacy()).unwrap().0;
assert_eq!(
chs,
CHS {
cylinder: 682,
head: 170,
sector: 42,
}
);
let mut slice = [0; 3];
encode_into_slice(chs, &mut slice, legacy()).unwrap();
for element in slice {
assert_eq!(element, 0xaa);
}
}
#[test]
fn align_chs_to_cylinder() {
fn lba2c(lba: u32) -> u16 {
let chs = CHS::from_lba_aligned(lba, 100, 2, 2).unwrap();
assert_eq!(chs.head, 0);
assert_eq!(chs.sector, 1);
chs.cylinder
}
assert_eq!(lba2c(12), 3);
assert_eq!(lba2c(10), 3);
}
#[test]
fn convert_chs_to_lba_and_back() {
let chs = CHS::from_lba_exact(666_666, 2484, 16, 63).unwrap();
assert_eq!(chs.to_lba(16, 63), 666_666);
let chs = CHS::from_lba_aligned(666_666, 2484, 16, 63).unwrap();
assert_eq!(chs.to_lba(16, 63), 667_296);
let chs = CHS::from_lba_exact(667_296, 2484, 16, 63).unwrap();
assert_eq!(chs.head, 0);
assert_eq!(chs.sector, 1);
}
#[test]
#[allow(clippy::bool_assert_comparison)]
fn read_disk1() {
let mut mbr = MBR::read_from(&mut File::open(DISK1).unwrap(), 512).unwrap();
assert!(mbr.header.partition_1.is_used());
assert!(mbr.header.partition_2.is_used());
assert!(mbr.header.partition_3.is_unused());
assert!(mbr.header.partition_4.is_unused());
assert!(mbr.header.partition_1.is_active());
assert!(!mbr.header.partition_2.is_active());
assert!(!mbr.header.partition_3.is_active());
assert!(!mbr.header.partition_4.is_active());
assert_eq!(mbr.len(), 4);
assert_eq!(mbr.header.iter().count(), 4);
assert_eq!(mbr.header.iter_mut().count(), 4);
assert_eq!(mbr.iter_mut().count(), 4);
assert_eq!(mbr.header.partition_1.boot, BOOT_ACTIVE);
assert_eq!(mbr.header.partition_1.sys, 0x06);
assert_eq!(mbr.header.partition_1.starting_lba, 1);
assert_eq!(mbr.header.partition_1.sectors, 1);
assert_eq!(mbr.header.partition_2.boot, BOOT_INACTIVE);
assert_eq!(mbr.header.partition_2.sys, 0x0b);
assert_eq!(mbr.header.partition_2.starting_lba, 3);
assert_eq!(mbr.header.partition_2.sectors, 1);
}
#[test]
fn read_disk2() {
let mut mbr = MBR::read_from(&mut File::open(DISK2).unwrap(), 512).unwrap();
assert!(mbr.header.partition_1.is_used());
assert!(mbr.header.partition_2.is_used());
assert!(mbr.header.partition_3.is_unused());
assert!(mbr.header.partition_4.is_unused());
assert!(!mbr.header.partition_1.is_active());
assert!(!mbr.header.partition_2.is_active());
assert!(!mbr.header.partition_3.is_active());
assert!(!mbr.header.partition_4.is_active());
assert_eq!(mbr.header.partition_2.sys, 0x05);
assert_eq!(mbr.header.partition_2.starting_lba, 5);
assert_eq!(mbr.header.partition_2.first_chs, CHS::new(0, 1, 3));
assert_eq!(mbr.header.partition_2.last_chs, CHS::new(3, 0, 2));
assert_eq!(mbr.header.partition_2.sectors, 15);
assert_eq!(mbr.len(), 9);
assert_eq!(mbr.iter().count(), 9);
assert_eq!(mbr.iter_mut().count(), 9);
assert_eq!(mbr.iter().filter(|(_, x)| x.is_used()).count(), 7);
assert!(mbr
.iter()
.filter(|(_, x)| x.is_used() && !x.is_extended())
.all(|(_, x)| x.sys == 0x83));
assert!(mbr.get(9).is_some());
assert!(mbr.get(10).is_none());
assert_eq!(mbr.logical_partitions[0].absolute_ebr_lba, 5);
assert_eq!(mbr.logical_partitions[0].ebr_sectors, None);
assert_eq!(mbr.logical_partitions[1].absolute_ebr_lba, 7);
assert_eq!(mbr.logical_partitions[1].ebr_sectors, Some(3));
assert_eq!(mbr.logical_partitions[2].absolute_ebr_lba, 10);
assert_eq!(mbr.logical_partitions[2].ebr_sectors, Some(4));
assert_eq!(
mbr.logical_partitions[2].partition.first_chs,
CHS::new(1, 1, 3)
);
assert_eq!(
mbr.logical_partitions[2].partition.first_chs.to_lba(2, 3),
mbr.logical_partitions[2].partition.starting_lba
);
assert_eq!(
CHS::from_lba_exact(
mbr.logical_partitions[2].partition.starting_lba,
u16::MAX,
2,
3
)
.unwrap(),
mbr.logical_partitions[2].partition.first_chs
);
assert_eq!(mbr.logical_partitions[3].absolute_ebr_lba, 14);
assert_eq!(mbr.logical_partitions[3].ebr_sectors, Some(2));
assert_eq!(mbr.logical_partitions[4].absolute_ebr_lba, 16);
assert_eq!(mbr.logical_partitions[4].ebr_sectors, Some(2));
assert!(mbr.logical_partitions.get(5).is_none());
}
#[test]
fn read_empty_extended_partition() {
let ss = 512_u32;
let data = vec![0; 10 * ss as usize];
let mut cur = Cursor::new(data);
let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
mbr.header.partition_1.sys = 0x0f;
mbr.header.partition_1.starting_lba = 1;
mbr.header.partition_1.sectors = 10;
mbr.write_into(&mut cur).unwrap();
let mbr = MBR::read_from(&mut cur, ss).unwrap();
assert_eq!(mbr.header.partition_1.sys, 0x0f);
assert_eq!(mbr.header.partition_1.starting_lba, 1);
assert_eq!(mbr.header.partition_1.sectors, 10);
assert!(mbr.logical_partitions.is_empty());
}
#[test]
#[allow(clippy::bool_assert_comparison)]
fn new_mbr_then_write_then_read_twice() {
let ss = 512_u32;
let data = vec![0; 12 * ss as usize];
let mut cur = Cursor::new(data);
let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
mbr.header.partition_1.sys = 0x83;
mbr.header.partition_1.starting_lba = 1;
mbr.header.partition_1.sectors = 4;
mbr.header.partition_3.sys = 0x0f;
mbr.header.partition_3.starting_lba = 5;
mbr.header.partition_3.sectors = 6;
let mut empty = MBRPartitionEntry::empty();
empty.starting_lba = 6;
empty.sectors = 1;
mbr.logical_partitions.push(LogicalPartition {
partition: empty,
absolute_ebr_lba: 5,
ebr_sectors: None,
ebr_first_chs: CHS::empty(),
ebr_last_chs: None,
bootstrap_code: [0; 446],
});
mbr.logical_partitions.push(LogicalPartition {
partition: MBRPartitionEntry {
boot: BOOT_ACTIVE,
first_chs: CHS::empty(),
sys: 0x83,
last_chs: CHS::empty(),
starting_lba: 9,
sectors: 1,
},
absolute_ebr_lba: 7,
ebr_sectors: Some(3),
ebr_first_chs: CHS::empty(),
ebr_last_chs: Some(CHS::empty()),
bootstrap_code: [0; 446],
});
mbr.write_into(&mut cur).unwrap();
let mut mbr = MBR::read_from(&mut cur, ss).unwrap();
assert_eq!(mbr.header.partition_1.boot, BOOT_INACTIVE);
assert_eq!(mbr.header.partition_1.sys, 0x83);
assert_eq!(mbr.header.partition_1.starting_lba, 1);
assert_eq!(mbr.header.partition_1.sectors, 4);
assert_eq!(mbr.logical_partitions.len(), 2);
assert_eq!(mbr.logical_partitions[0].absolute_ebr_lba, 5);
assert_eq!(mbr.logical_partitions[0].partition.boot, BOOT_INACTIVE);
assert_eq!(mbr.logical_partitions[0].partition.starting_lba, 6);
assert_eq!(mbr.logical_partitions[0].partition.sectors, 1);
assert_eq!(mbr.logical_partitions[0].partition.sys, 0);
assert_eq!(mbr.logical_partitions[0].ebr_sectors, None);
assert_eq!(mbr.logical_partitions[1].absolute_ebr_lba, 7);
assert_eq!(mbr.logical_partitions[1].partition.boot, BOOT_ACTIVE);
assert_eq!(mbr.logical_partitions[1].partition.starting_lba, 9);
assert_eq!(mbr.logical_partitions[1].partition.sectors, 1);
assert_eq!(mbr.logical_partitions[1].partition.sys, 0x83);
assert_eq!(mbr.logical_partitions[1].ebr_sectors, Some(3));
mbr.write_into(&mut cur).unwrap();
let mbr = MBR::read_from(&mut cur, ss).unwrap();
assert_eq!(mbr.header.partition_1.boot, BOOT_INACTIVE);
assert_eq!(mbr.header.partition_1.sys, 0x83);
assert_eq!(mbr.header.partition_1.starting_lba, 1);
assert_eq!(mbr.header.partition_1.sectors, 4);
assert_eq!(mbr.logical_partitions.len(), 2);
assert_eq!(mbr.logical_partitions[0].absolute_ebr_lba, 5);
assert_eq!(mbr.logical_partitions[0].partition.boot, BOOT_INACTIVE);
assert_eq!(mbr.logical_partitions[0].partition.starting_lba, 6);
assert_eq!(mbr.logical_partitions[0].partition.sectors, 1);
assert_eq!(mbr.logical_partitions[0].partition.sys, 0);
assert_eq!(mbr.logical_partitions[0].ebr_sectors, None);
assert_eq!(mbr.logical_partitions[1].absolute_ebr_lba, 7);
assert_eq!(mbr.logical_partitions[1].partition.boot, BOOT_ACTIVE);
assert_eq!(mbr.logical_partitions[1].partition.starting_lba, 9);
assert_eq!(mbr.logical_partitions[1].partition.sectors, 1);
assert_eq!(mbr.logical_partitions[1].partition.sys, 0x83);
assert_eq!(mbr.logical_partitions[1].ebr_sectors, Some(3));
}
#[test]
fn find_at_sector() {
let mbr = MBR::read_from(&mut File::open(DISK2).unwrap(), 512).unwrap();
assert_eq!(mbr.find_at_sector(2), Some(1));
assert_eq!(mbr.find_at_sector(4), None);
assert_eq!(mbr.find_at_sector(8), Some(6));
assert_eq!(mbr.find_at_sector(7), None);
}
#[test]
fn find_free_sectors() {
let ss = 512_u32;
let data = vec![0; 10 * ss as usize];
let mut cur = Cursor::new(data);
let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
mbr.align = 1;
mbr.header.partition_1.sys = 0x83;
mbr.header.partition_1.starting_lba = 1;
mbr.header.partition_1.sectors = 1;
mbr.header.partition_3.sys = 0x0f;
mbr.header.partition_3.starting_lba = 5;
mbr.header.partition_3.sectors = 5;
mbr.logical_partitions.push(LogicalPartition {
bootstrap_code: [0; 446],
partition: MBRPartitionEntry {
boot: BOOT_INACTIVE,
first_chs: CHS::empty(),
sys: 0x83,
last_chs: CHS::empty(),
starting_lba: 2,
sectors: 1,
},
absolute_ebr_lba: 5,
ebr_sectors: None,
ebr_first_chs: CHS::empty(),
ebr_last_chs: None,
});
assert_eq!(mbr.find_free_sectors(), vec![(2, 3), (6, 1), (8, 2)]);
}
#[test]
fn push_logical_partition_aligned_to_cylinder() {
let ss = 512_u32;
let data = vec![0; 54 * ss as usize];
let mut cur = Cursor::new(data);
let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
mbr.cylinders = 6;
mbr.heads = 3;
mbr.sectors = 3;
let align = 9;
mbr.align = mbr.get_cylinder_size();
assert_eq!(mbr.align, align);
mbr.header.partition_1.sys = 0x05;
mbr.header.partition_1.first_chs = CHS::new(1, 0, 1);
mbr.header.partition_1.last_chs = CHS::new(5, 0, 1);
mbr.header.partition_1.starting_lba = align;
mbr.header.partition_1.sectors = mbr.disk_size;
let p = mbr.push(0x00, 2, 2 * align).unwrap();
assert_eq!(p.absolute_ebr_lba, align);
assert_eq!(p.partition.starting_lba, 2 * align);
assert_eq!(p.ebr_sectors, None);
assert_eq!(p.partition.sectors, align);
assert_eq!(p.ebr_first_chs, CHS::new(1, 0, 1));
assert_eq!(p.ebr_last_chs, None);
assert_eq!(p.partition.first_chs, CHS::new(2, 0, 1));
assert_eq!(p.partition.last_chs, CHS::new(2, 2, 3));
let p = mbr
.push(
0x83,
CHS::new(3, 0, 1).to_lba(mbr.heads, mbr.sectors),
align * 3,
)
.unwrap();
assert_eq!(p.absolute_ebr_lba, align * 3);
assert_eq!(p.partition.starting_lba, 4 * align);
assert_eq!(p.ebr_sectors, Some(align * 3));
assert_eq!(p.partition.sectors, align * 2);
assert_eq!(p.ebr_first_chs, CHS::new(3, 0, 1));
assert_eq!(p.ebr_last_chs, Some(CHS::new(5, 2, 3)));
assert_eq!(p.partition.first_chs, CHS::new(4, 0, 1));
assert_eq!(p.partition.last_chs, CHS::new(5, 2, 3));
mbr.write_into(&mut cur).unwrap();
let mut same_mbr = MBR::read_from(&mut cur, ss).unwrap();
same_mbr.cylinders = mbr.cylinders;
same_mbr.heads = mbr.heads;
same_mbr.sectors = mbr.sectors;
assert_eq!(mbr, same_mbr);
}
#[test]
fn push_logical_partition_check_within_range() {
let ss = 512_u32;
let data = vec![0; 10 * ss as usize];
let mut cur = Cursor::new(data);
let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
mbr.align = 1;
mbr.header.partition_1.sys = 0x0f;
mbr.header.partition_1.starting_lba = 4;
mbr.header.partition_1.sectors = 2;
assert!(mbr.push(0x83, 3, 2).is_err());
assert!(mbr.push(0x83, 6, 2).is_err());
mbr.push(0x83, 4, 2).unwrap();
assert!(mbr.push(0x83, 4, 3).is_err());
}
#[test]
fn remove_logical_partition() {
let ss = 512_u32;
let data = vec![0; 10 * ss as usize];
let mut cur = Cursor::new(data);
let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
mbr.align = 1;
mbr.header.partition_1.sys = 0x0f;
mbr.header.partition_1.starting_lba = 1;
mbr.header.partition_1.sectors = mbr.disk_size - 1;
mbr.push(0x00, 1, 2).unwrap();
mbr.push(0x83, 4, 3).unwrap();
mbr.logical_partitions.remove(0);
mbr.write_into(&mut cur).unwrap();
let same_mbr = MBR::read_from(&mut cur, ss).unwrap();
assert_eq!(mbr, same_mbr);
assert_eq!(mbr.logical_partitions[0].partition.starting_lba, 5);
assert_eq!(mbr.logical_partitions[0].absolute_ebr_lba, 1);
assert_eq!(mbr.logical_partitions[0].ebr_sectors, None);
}
fn read_corrupt_mbr(path: &str, bad_offset: usize, bad_data: u8) -> Error {
let mut data = Vec::new();
File::open(path).unwrap().read_to_end(&mut data).unwrap();
data[bad_offset] = bad_data;
let mut cur = Cursor::new(&data);
MBR::read_from(&mut cur, 512).unwrap_err()
}
#[test]
fn read_invalid_boot_signature() {
assert!(matches!(
read_corrupt_mbr(DISK2, 0xffe, 0xaa),
Error::InvalidSignature
));
assert!(matches!(
read_corrupt_mbr(DISK2, 0x21fe, 0xaa),
Error::InvalidSignature
));
}
#[test]
fn write_invalid_boot_signature() {
let ss = 512_u32;
let data = vec![0; 10 * ss as usize];
let mut cur = Cursor::new(data);
let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
mbr.header.partition_1.sys = 0x0a;
mbr.header.partition_1.starting_lba = 1;
mbr.header.partition_1.sectors = 10;
mbr.header.boot_signature = [0, 0];
assert!(matches!(
mbr.write_into(&mut cur).unwrap_err(),
Error::InvalidSignature
));
}
#[test]
fn read_invalid_boot_flag() {
assert!(matches!(
read_corrupt_mbr(DISK2, 0x1be, BOOT_ACTIVE | 0x01),
Error::InvalidBootFlag
));
assert!(matches!(
read_corrupt_mbr(DISK2, 0xfbe, BOOT_ACTIVE | 0x01),
Error::InvalidBootFlag
));
assert!(matches!(
read_corrupt_mbr(DISK2, 0xfce, BOOT_ACTIVE | 0x01),
Error::InvalidBootFlag
));
}
#[test]
fn write_invalid_boot_flag() {
let ss = 512_u32;
let data = vec![0; 10 * ss as usize];
let mut cur = Cursor::new(data);
let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
mbr.header.partition_2.boot = BOOT_ACTIVE | 0x01;
mbr.header.partition_2.sys = 0x0a;
mbr.header.partition_2.starting_lba = 1;
mbr.header.partition_2.sectors = 10;
assert!(matches!(
mbr.write_into(&mut cur).unwrap_err(),
Error::InvalidBootFlag
));
let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
mbr.align = 1;
mbr.header.partition_1.sys = 0x0f;
mbr.header.partition_1.starting_lba = 1;
mbr.header.partition_1.sectors = 10;
let partition = mbr.push(0x0f, 1, 9).unwrap();
partition.partition.boot = BOOT_ACTIVE | 0x01;
assert!(matches!(
mbr.write_into(&mut cur).unwrap_err(),
Error::InvalidBootFlag
));
}
#[test]
fn update_from() {
let ss = 512_u32;
let data = vec![0; 10 * ss as usize];
let mut cur = Cursor::new(data);
let mut mbr = MBR::new_from(&mut cur, ss, [0xff; 4]).unwrap();
assert_eq!(mbr.disk_size, 10);
let mut data = cur.into_inner();
data.resize(20 * ss as usize, 0);
let mut cur = Cursor::new(data);
mbr.update_from(&mut cur).unwrap();
assert_eq!(mbr.disk_size, 20);
}
}
#[cfg(doctest)]
mod test_readme {
macro_rules! check_doc {
($x:expr) => {
#[doc = $x]
extern "C" {}
};
}
check_doc!(include_str!("../README.md"));
}