extern crate failure;
#[macro_use]
extern crate failure_derive;
extern crate libc;
mod mount;
pub use self::mount::Mount;
use std::cmp;
use std::ffi::OsString;
use std::fs::{canonicalize, read_dir, File, OpenOptions};
use std::io::{self, Read, Seek, SeekFrom, Write};
use std::os::unix::ffi::OsStrExt;
use std::os::unix::fs::{FileTypeExt, OpenOptionsExt};
use std::path::{Path, PathBuf};
use std::process::{Command, ExitStatus};
const BUFFER_SIZE: usize = 4 * 1024 * 1024;
#[derive(Debug, Fail)]
#[cfg_attr(rustfmt, rustfmt_skip)]
pub enum ImageError {
#[fail(display = "image could not be opened: {}", why)]
Open { why: io::Error },
#[fail(display = "unable to get image metadata: {}", why)]
Metadata { why: io::Error },
#[fail(display = "image was not a file")]
NotAFile,
#[fail(display = "unable to read image: {}", why)]
ReadError { why: io::Error },
#[fail(display = "reached EOF prematurely")]
Eof,
}
pub struct Image {
path: PathBuf,
file: File,
size: u64,
}
impl Image {
pub fn new<P: AsRef<Path>>(path: P) -> Result<Image, ImageError> {
let path = path.as_ref();
File::open(path)
.map_err(|why| ImageError::Open { why })
.and_then(|file| {
file.metadata()
.map_err(|why| ImageError::Metadata { why })
.and_then(|metadata| {
if metadata.file_type().is_file() {
Ok(Image {
path: path.to_path_buf(),
file,
size: metadata.len(),
})
} else {
Err(ImageError::NotAFile)
}
})
})
}
pub fn get_path(&self) -> &Path { &self.path }
pub fn get_size(&self) -> u64 { self.size }
pub fn read<P: FnMut(u64)>(
&mut self,
data: &mut Vec<u8>,
mut progress_callback: P,
) -> Result<(), ImageError> {
if data.capacity() < self.size as usize {
let capacity = self.size as usize - data.capacity();
data.reserve_exact(capacity);
data.append(&mut vec![0; capacity])
} else if data.capacity() > self.size as usize {
data.truncate(self.size as usize);
data.shrink_to_fit();
}
let mut total = 0;
while total < data.len() {
let end = cmp::min(data.len(), total + BUFFER_SIZE);
let count = self.file
.read(&mut data[total..end])
.map_err(|why| ImageError::ReadError { why })?;
if count == 0 {
return Err(ImageError::Eof);
}
total += count;
progress_callback(total as u64);
}
Ok(())
}
}
#[derive(Debug, Fail)]
#[cfg_attr(rustfmt, rustfmt_skip)]
pub enum DiskError {
#[fail(display = "unable to open directory at '{}': {}", dir, why)]
Directory { dir: &'static str, why: io::Error },
#[fail(display = "unable to read directory entry at '{:?}': invalid UTF-8", dir)]
UTF8 { dir: PathBuf },
#[fail(display = "unable to find disk '{}': {}", disk, why)]
NoDisk { disk: String, why: io::Error },
#[fail(display = "failed to unmount {:?}: exit status {}", path, status)]
UnmountStatus { path: OsString, status: ExitStatus },
#[fail(display = "failed to unmount {:?}: {}", path, why)]
UnmountCommand { path: OsString, why: io::Error },
#[fail(display = "error using disk '{}': {:?} already mounted at {:?}", arg, source, dest)]
AlreadyMounted { arg: String, source: OsString, dest: OsString },
#[fail(display = "'{}' is not a block device", arg)]
NotABlock { arg: String },
#[fail(display = "unable to get metadata of disk '{}': {}", arg, why)]
Metadata { arg: String, why: io::Error },
#[fail(display = "unable to open disk '{}': {}", disk, why)]
Open { disk: String, why: io::Error },
#[fail(display = "error writing disk '{}': {}", disk, why)]
Write { disk: String, why: io::Error },
#[fail(display = "error writing disk '{}': reached EOF", disk)]
WriteEOF { disk: String },
#[fail(display = "unable to flush disk '{}': {}", disk, why)]
Flush { disk: String, why: io::Error },
#[fail(display = "error seeking disk '{}': seeked to {} instead of 0", disk, invalid)]
SeekInvalid { disk: String, invalid: u64 },
#[fail(display = "error seeking disk '{}': {}", disk, why)]
Seek { disk: String, why: io::Error },
#[fail(display = "error verifying disk '{}': {}", disk, why)]
Verify { disk: String, why: io::Error },
#[fail(display = "error verifying disk '{}': reached EOF", disk)]
VerifyEOF { disk: String },
#[fail(display = "error verifying disk '{}': mismatch at {}:{}", disk, x, y)]
VerifyMismatch { disk: String, x: usize, y: usize },
}
fn is_usb(filename: &str) -> bool {
filename.starts_with("pci-") && filename.contains("-usb-") && filename.ends_with("-0:0:0:0")
}
const DISK_DIR: &str = "/dev/disk/by-path/";
pub fn get_disk_args(disks: &mut Vec<String>) -> Result<(), DiskError> {
let readdir = read_dir(DISK_DIR).map_err(|why| DiskError::Directory { dir: DISK_DIR, why })?;
for entry_res in readdir {
let entry = entry_res.map_err(|why| DiskError::Directory { dir: DISK_DIR, why })?;
let path = entry.path();
if let Some(filename_os) = path.file_name() {
if is_usb(filename_os.to_str().unwrap()) {
disks.push(
path.to_str()
.ok_or_else(|| DiskError::UTF8 { dir: path.clone() })?
.into(),
);
}
}
}
Ok(())
}
pub fn disks_from_args<D: Iterator<Item = String>>(
disk_args: D,
mounts: &[Mount],
unmount: bool,
) -> Result<Vec<(String, File)>, DiskError> {
let mut disks = Vec::new();
for disk_arg in disk_args {
let canonical_path = canonicalize(&disk_arg).map_err(|why| DiskError::NoDisk {
disk: disk_arg.clone(),
why,
})?;
for mount in mounts.iter() {
if mount
.source
.as_bytes()
.starts_with(canonical_path.as_os_str().as_bytes())
{
if unmount {
eprintln!(
"unmounting '{}': {:?} is mounted at {:?}",
disk_arg, mount.source, mount.dest
);
Command::new("umount")
.arg(&mount.source)
.status()
.map_err(|why| DiskError::UnmountCommand {
path: mount.source.clone(),
why,
})
.and_then(|status| {
if !status.success() {
Err(DiskError::UnmountStatus {
path: mount.source.clone(),
status,
})
} else {
Ok(())
}
})?;
} else {
return Err(DiskError::AlreadyMounted {
arg: disk_arg.clone(),
source: mount.source.clone(),
dest: mount.dest.clone(),
});
}
}
}
let metadata = canonical_path
.metadata()
.map_err(|why| DiskError::Metadata {
arg: disk_arg.clone(),
why,
})?;
if !metadata.file_type().is_block_device() {
return Err(DiskError::NotABlock {
arg: disk_arg.clone(),
});
}
let disk = OpenOptions::new()
.read(true)
.write(true)
.custom_flags(libc::O_SYNC)
.open(&canonical_path)
.map_err(|why| DiskError::Open {
disk: disk_arg.clone(),
why,
})?;
disks.push((disk_arg, disk));
}
Ok(disks)
}
pub fn write_to_disk<M, F, S>(
mut message: M,
finish: F,
mut set: S,
mut disk: File,
disk_path: String,
image_size: u64,
image_data: &[u8],
check: bool,
) -> Result<(), DiskError>
where
M: FnMut(&str),
F: Fn(),
S: FnMut(u64),
{
let mut total = 0;
while total < image_data.len() {
let end = cmp::min(image_size as usize, total + BUFFER_SIZE);
let count = disk.write(&image_data[total..end]).map_err(|why| {
message(&format!("! {}: ", disk_path));
finish();
DiskError::Write {
disk: disk_path.clone(),
why,
}
})?;
if count == 0 {
message(&format!("! {}: ", disk_path));
finish();
return Err(DiskError::WriteEOF {
disk: disk_path.clone(),
});
}
total += count;
set(total as u64);
}
disk.flush().map_err(|why| {
message(&format!("! {}: ", disk_path));
finish();
DiskError::Flush {
disk: disk_path.clone(),
why,
}
})?;
if check {
match disk.seek(SeekFrom::Start(0)) {
Ok(0) => (),
Ok(invalid) => {
message(&format!("! {}: ", disk_path));
finish();
return Err(DiskError::SeekInvalid {
disk: disk_path.clone(),
invalid,
});
}
Err(why) => {
message(&format!("! {}: ", disk_path));
finish();
return Err(DiskError::Seek {
disk: disk_path.clone(),
why,
});
}
}
message(&format!("V {}: ", disk_path));
set(0);
total = 0;
let mut buf = vec![0; BUFFER_SIZE];
while total < image_data.len() {
let end = cmp::min(image_size as usize, total + BUFFER_SIZE);
let count = match disk.read(&mut buf[..end - total]) {
Ok(count) => count,
Err(why) => {
message(&format!("! {}: ", disk_path));
finish();
return Err(DiskError::Verify {
disk: disk_path.clone(),
why,
});
}
};
if count == 0 {
message(&format!("! {}: ", disk_path));
finish();
return Err(DiskError::VerifyEOF {
disk: disk_path.clone(),
});
}
if buf[..count] != image_data[total..total + count] {
message(&format!("! {}: ", disk_path));
finish();
return Err(DiskError::VerifyMismatch {
disk: disk_path.clone(),
x: total,
y: total + count,
});
}
total += count;
set(total as u64);
}
}
finish();
Ok(())
}