mod args;
mod bufcache;
mod disktest;
mod drop_caches;
mod generator;
mod kdf;
mod seed;
mod stream;
mod stream_aggregator;
mod util;
use anyhow as ah;
use args::{Args, parse_args};
use crate::seed::print_generated_seed;
use disktest::{Disktest, DisktestFile};
use std::env::args_os;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
fn install_abort_handlers() -> ah::Result<Arc<AtomicBool>> {
let abort = Arc::new(AtomicBool::new(false));
for sig in &[signal_hook::consts::signal::SIGTERM,
signal_hook::consts::signal::SIGINT] {
if let Err(e) = signal_hook::flag::register(*sig, Arc::clone(&abort)) {
return Err(ah::format_err!("Failed to register signal {}: {}", sig, e));
}
}
Ok(abort)
}
fn new_disktest(args: &Args,
write: bool,
abort: &Arc<AtomicBool>) -> ah::Result<(Disktest, DisktestFile)> {
Ok((
Disktest::new(args.algorithm,
args.seed.as_bytes().to_vec(),
args.invert_pattern,
args.threads,
Some(Arc::clone(abort))),
DisktestFile::open(&args.device,
!write,
write,
args.quiet)?,
))
}
fn main() -> ah::Result<()> {
let args = parse_args(args_os())?;
let abort = install_abort_handlers()?;
if !args.user_seed && args.quiet < 2 {
print_generated_seed(&args.seed, true);
}
let mut result = Ok(());
if args.write {
let (mut disktest, file) = new_disktest(&args, true, &abort)?;
result = match disktest.write(file, args.seek, args.max_bytes) {
Ok(_) => Ok(()),
Err(e) => Err(e),
};
}
if args.verify && result.is_ok() {
let (mut disktest, file) = new_disktest(&args, false, &abort)?;
result = match disktest.verify(file, args.seek, args.max_bytes) {
Ok(_) => Ok(()),
Err(e) => Err(e),
};
}
if !args.user_seed && args.quiet < 2 {
print_generated_seed(&args.seed, false);
}
if result.is_ok() && args.quiet < 1 {
println!("Success!");
}
result
}