summaryrefslogtreecommitdiffstats
path: root/src/main.rs
blob: f21711ec86231246b85847dedebc7463d4128778 (plain)
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
// -*- coding: utf-8 -*-
//
// disktest - Hard drive tester
//
// Copyright 2020-2022 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 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;

/// 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 < 2 {
        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 < 2 {
        print_generated_seed(&args.seed, false);
    }
    if result.is_ok() && args.quiet < 1 {
        println!("Success!");
    }

    result
}

// vim: ts=4 sw=4 expandtab
bues.ch cgit interface