restic-integrity/src/main.rs

use clap::Parser;
use indicatif::{ParallelProgressIterator, ProgressBar, ProgressStyle, ProgressFinish};
use rayon::ThreadPoolBuilder;
use rayon::iter::{ParallelIterator, IntoParallelRefIterator};
use std::fs::File;
use std::path::Path;
use std::io::{self, Result, Error, ErrorKind};
use sha2::{Sha256, Digest};

#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
    /// The path to the repository
    path: String,

    /// The number of jobs to use for hashing
    #[clap(short, long, default_value_t = 0)]
    jobs: u8,

    #[clap(short, long, default_value_t = 0)]
    start: usize,

    /// If the path argument should be used to check all files in a directory it points to instead of a restic repository (use for testing)
    #[clap(long)]
    single_directory: bool
}

fn run() -> Result<()> {
    let args = Args::parse();
    let path = Path::new(&args.path);
    let start = args.start;

    let mut directories_with_hashed_files = Vec::new();

    if args.single_directory {
        directories_with_hashed_files.push(path.to_path_buf());
    } else {
        directories_with_hashed_files.push(path.join(Path::new("index")));
        directories_with_hashed_files.push(path.join(Path::new("keys")));
        directories_with_hashed_files.push(path.join(Path::new("snapshots")));

        for subdirectory in path.join(Path::new("data")).read_dir()? {
            let subdirectory = subdirectory?;

            directories_with_hashed_files.push(subdirectory.path());
        }
    }

    let mut hashed_files = Vec::new();

    for directory in directories_with_hashed_files {
        for file in directory.read_dir()? {
            let file = file?;
            let path = file.path();

            hashed_files.push(path);
        }
    }

    let progress_bar = ProgressBar::new(hashed_files.len() as u64);
    progress_bar.set_style(ProgressStyle::default_bar()
        .template("Hashing files... {bar} {pos:>7}/{len:7} [{elapsed_precise}] ")
        .on_finish(ProgressFinish::AndLeave));

    let pool = ThreadPoolBuilder::new().num_threads(args.jobs as usize).build().unwrap();

    hashed_files.sort();

    let hashed_files_len = hashed_files.len();

    if start != 0 {
        match hashed_files_len.checked_sub(start) {
            None => return Err(Error::new(ErrorKind::Other, format!("Specified start at {} is larger than the total number of files {}", start, hashed_files_len))),
            Some(remaining) => {
                println!("Starting at {} files, {} remaining...", start, remaining);

                hashed_files.drain(0..start);
            }
        }
    }

    pool.install(|| {
        return hashed_files.par_iter().progress_with(progress_bar).try_for_each(|path| -> Result<()> {
            // Just opening the file and hashing using io::copy is roughly ~2.5x fater compared to sha256::digest_file
            let mut hasher = Sha256::new();
            let mut file = File::open(path)?;

            io::copy(&mut file, &mut hasher)?;

            let hash = format!("{:x}", hasher.finalize());
            let filename = path.file_name().unwrap().to_str().unwrap();

            if filename != hash {
                let path_string = path.to_str().unwrap();

                return Err(Error::new(ErrorKind::Other, format!("Integrity check failed for {}: Expected {}, got {}", path_string, filename, hash)));
            }

            Ok(())
        });
    })
}

fn main() {
    let result = run();
    if result.is_err() {
        eprintln!("{}", result.unwrap_err());
        std::process::exit(1);
    }
}