1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
|
// -*- coding: utf-8 -*-
//
// disktest - Hard drive tester
//
// Copyright 2020-2023 Michael Buesch <m@bues.ch>
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
//
mod args;
mod bufcache;
mod disktest;
mod generator;
mod kdf;
mod rawio;
mod seed;
mod stream;
mod stream_aggregator;
mod util;
use crate::seed::print_generated_seed;
use anyhow as ah;
use args::{parse_args, Args};
use disktest::{Disktest, DisktestFile, DisktestQuiet};
use std::env::args_os;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
/// Install abort signal handlers and return
/// the abort-flag that is written to true by these handlers.
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)
}
/// Create a new disktest core instance.
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,
args.quiet,
Some(Arc::clone(abort)),
),
DisktestFile::open(&args.device, !write, write)?,
))
}
/// Main program entry point.
fn main() -> ah::Result<()> {
let args = parse_args(args_os())?;
let abort = install_abort_handlers()?;
if !args.user_seed && args.quiet < DisktestQuiet::NoInfo {
print_generated_seed(&args.seed, true);
}
// Run write-mode, if requested.
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),
};
}
// Run verify-mode, if requested.
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 < DisktestQuiet::NoInfo {
print_generated_seed(&args.seed, false);
}
if result.is_ok() && args.quiet == DisktestQuiet::Normal {
println!("Success!");
}
result
}
// vim: ts=4 sw=4 expandtab
|
