#![deny(missing_docs)]
use log::*;
use std::collections::BTreeMap;
use std::io::Write;
use std::{fs, io, path};
#[macro_use]
mod macros;
pub mod disk;
pub mod header;
pub mod mbr;
pub mod partition;
pub mod partition_types;
#[derive(Debug, Eq, PartialEq)]
pub struct GptConfig {
lb_size: disk::LogicalBlockSize,
writable: bool,
initialized: bool,
}
impl GptConfig {
pub fn new() -> Self {
GptConfig::default()
}
pub fn writable(mut self, writable: bool) -> Self {
self.writable = writable;
self
}
pub fn initialized(mut self, initialized: bool) -> Self {
self.initialized = initialized;
self
}
pub fn logical_block_size(mut self, lb_size: disk::LogicalBlockSize) -> Self {
self.lb_size = lb_size;
self
}
pub fn open(self, diskpath: &path::Path) -> io::Result<GptDisk> {
if !self.initialized {
let file = fs::OpenOptions::new()
.write(self.writable)
.read(true)
.open(diskpath)?;
let empty = GptDisk {
config: self,
file,
guid: uuid::Uuid::new_v4(),
path: diskpath.to_path_buf(),
primary_header: None,
backup_header: None,
partitions: BTreeMap::new(),
};
return Ok(empty);
}
let mut file = fs::OpenOptions::new()
.write(self.writable)
.read(true)
.open(diskpath)?;
let h1 = header::read_primary_header(&mut file, self.lb_size)?;
let h2 = header::read_backup_header(&mut file, self.lb_size)?;
let table = partition::file_read_partitions(&mut file, &h1, self.lb_size)?;
let disk = GptDisk {
config: self,
file,
guid: h1.disk_guid,
path: diskpath.to_path_buf(),
primary_header: Some(h1),
backup_header: Some(h2),
partitions: table,
};
debug!("disk: {:?}", disk);
Ok(disk)
}
}
impl Default for GptConfig {
fn default() -> Self {
Self {
lb_size: disk::DEFAULT_SECTOR_SIZE,
initialized: true,
writable: false,
}
}
}
#[derive(Debug)]
pub struct GptDisk {
config: GptConfig,
file: fs::File,
guid: uuid::Uuid,
path: path::PathBuf,
primary_header: Option<header::Header>,
backup_header: Option<header::Header>,
partitions: BTreeMap<u32, partition::Partition>,
}
impl GptDisk {
pub fn add_partition(
&mut self,
name: &str,
size: u64,
part_type: partition_types::Type,
flags: u64,
) -> io::Result<u32> {
let size_lba = match size.checked_div(self.config.lb_size.into()) {
Some(s) => s,
None => {
return Err(io::Error::new(
io::ErrorKind::Other,
format!(
"size must be greater than {} which is the logical block size.",
self.config.lb_size
),
));
}
};
let free_sections = self.find_free_sectors();
for (starting_lba, length) in free_sections {
if length >= size_lba {
let partition_id = self.find_next_partition_id();
debug!(
"Adding partition id: {} {:?}. first_lba: {} last_lba: {}",
partition_id,
part_type,
starting_lba,
starting_lba + size_lba as u64
);
let part = partition::Partition {
part_type_guid: part_type,
part_guid: uuid::Uuid::new_v4(),
first_lba: starting_lba,
last_lba: starting_lba + size_lba as u64,
flags,
name: name.to_string(),
};
if let Some(p) = self.partitions.insert(partition_id, part.clone()) {
debug!("Replacing\n{}\nwith\n{}", p, part);
}
return Ok(partition_id);
}
}
Err(io::Error::new(
io::ErrorKind::Other,
"Unable to find enough space on drive",
))
}
pub fn find_free_sectors(&self) -> Vec<(u64, u64)> {
if let Some(header) = self.primary_header().or_else(|| self.backup_header()) {
trace!("first_usable: {}", header.first_usable);
let mut disk_positions = vec![header.first_usable + 1];
for part in self.partitions().iter().filter(|p| p.1.is_used()) {
trace!("partition: ({}, {})", part.1.first_lba, part.1.last_lba);
disk_positions.push(part.1.first_lba);
disk_positions.push(part.1.last_lba);
}
disk_positions.push(header.last_usable - 1);
trace!("last_usable: {}", header.last_usable);
disk_positions.sort();
return disk_positions
.chunks(2)
.map(|p| (p[0] + 1, p[1].saturating_sub(p[0])))
.collect();
}
vec![]
}
pub fn find_next_partition_id(&self) -> u32 {
match self
.partitions()
.iter()
.filter(|p| p.1.is_used())
.max_by_key(|x| x.0)
{
Some(i) => i.0 + 1,
None => 1,
}
}
pub fn primary_header(&self) -> Option<&header::Header> {
self.primary_header.as_ref()
}
pub fn backup_header(&self) -> Option<&header::Header> {
self.backup_header.as_ref()
}
pub fn partitions(&self) -> &BTreeMap<u32, partition::Partition> {
&self.partitions
}
pub fn guid(&self) -> &uuid::Uuid {
&self.guid
}
pub fn logical_block_size(&self) -> &disk::LogicalBlockSize {
&self.config.lb_size
}
pub fn update_guid(&mut self, uuid: Option<uuid::Uuid>) -> io::Result<&Self> {
let guid = match uuid {
Some(u) => u,
None => {
let u = uuid::Uuid::new_v4();
debug!("Generated random uuid: {}", u);
u
}
};
self.guid = guid;
Ok(self)
}
pub fn update_partitions(
&mut self,
pp: BTreeMap<u32, partition::Partition>,
) -> io::Result<&Self> {
let bak = header::find_backup_lba(&mut self.file, self.config.lb_size)?;
let h1 = header::Header::compute_new(true, &pp, self.guid, bak)?;
let h2 = header::Header::compute_new(false, &pp, self.guid, bak)?;
self.primary_header = Some(h1);
self.backup_header = Some(h2);
self.partitions = pp;
self.config.initialized = true;
Ok(self)
}
pub fn write(mut self) -> io::Result<fs::File> {
if !self.config.writable {
return Err(io::Error::new(
io::ErrorKind::Other,
"disk not opened in writable mode",
));
}
if !self.config.initialized {
return Err(io::Error::new(io::ErrorKind::Other, "disk not initialized"));
}
debug!("Computing new headers");
trace!("old primary header: {:?}", self.primary_header);
trace!("old backup header: {:?}", self.backup_header);
let bak = header::find_backup_lba(&mut self.file, self.config.lb_size)?;
trace!("old backup lba: {}", bak);
for partition in self.partitions().iter().filter(|p| p.1.is_used()) {
partition.1.write(
&self.path,
u64::from(partition.0.checked_sub(1).unwrap_or_else(|| 0)),
self.primary_header.clone().unwrap().part_start,
self.config.lb_size,
)?;
}
let new_backup_header =
header::Header::compute_new(false, &self.partitions, self.guid, bak)?;
let new_primary_header =
header::Header::compute_new(true, &self.partitions, self.guid, bak)?;
debug!("Writing backup header");
new_backup_header.write_backup(&mut self.file, self.config.lb_size)?;
debug!("Writing primary header");
new_primary_header.write_primary(&mut self.file, self.config.lb_size)?;
trace!("new primary header: {:?}", new_primary_header);
trace!("new backup header: {:?}", new_backup_header);
self.file.flush()?;
self.primary_header = Some(new_primary_header.clone());
self.backup_header = Some(new_backup_header);
Ok(self.file)
}
}