Recover a block from a glitchy thumb

main.rs

use std::{
    env::args,
    fs::{remove_file, File, OpenOptions},
    io::{prelude::*, SeekFrom},
    thread::sleep,
    time::Duration,
};

const BLOCK_SIZE: usize = 4 * 1024 * 1024;
const SECTOR_SIZE: usize = 512;
const SECTORS_PER_BLOCK: usize = BLOCK_SIZE / SECTOR_SIZE;

fn main() {
    let nth_block: usize = args().nth(1).unwrap().parse().unwrap();
    let offset_bytes = nth_block * BLOCK_SIZE;

    dbg!(offset_bytes);

    let mut drive = File::open("/dev/sdb1").unwrap();
    let mut block_file = OpenOptions::new()
        .write(true)
        .create(true)
        .open(format!("block-{}", nth_block))
        .unwrap();

    let state_file_name = format!("block-{}.state", nth_block);
    let mut state_file = OpenOptions::new()
        .write(true)
        .read(true)
        .create(true)
        .open(&state_file_name)
        .unwrap();
    state_file.set_len(SECTORS_PER_BLOCK as u64).unwrap();

    let mut read_sectors: Vec<bool> = {
        let mut buffer = vec![0; SECTORS_PER_BLOCK];
        state_file.read_exact(&mut buffer).unwrap();

        buffer.into_iter().map(|x| x == 1).collect()
    };

    let mut parts = 2;

    while SECTORS_PER_BLOCK >= parts {
        let sectors_per_part = SECTORS_PER_BLOCK / parts;
        println!(
            "Splitting the block into {} parts, reading {} sectors at once.",
            parts, sectors_per_part,
        );

        let mut buffer = vec![0; sectors_per_part * SECTOR_SIZE];
        for (part, sectors_to_be_read) in
            read_sectors.chunks_exact_mut(sectors_per_part).enumerate()
        {
            if sectors_to_be_read.iter().all(|&x| x) {
                println!("{}th part is already read. skipping it...", part + 1);
                continue;
            }

            let first_sector_offset = part * sectors_per_part * SECTOR_SIZE;
            println!(
                "Reading {}th part, sectors {}..{}",
                part + 1,
                part * sectors_per_part,
                (part + 1) * sectors_per_part,
            );

            drive
                .seek(SeekFrom::Start((offset_bytes + first_sector_offset) as u64))
                .unwrap();

            if let Err(error) = drive.read_exact(&mut buffer) {
                println!(
                    "failed to read {}th part: {}. waiting for the drive...",
                    part + 1,
                    error
                );

                drop(drive);
                drive = loop {
                    if let Ok(file) = File::open("/dev/sdb1") {
                        break file;
                    }

                    sleep(Duration::from_secs(1));
                };

                continue;
            }

            block_file
                .seek(SeekFrom::Start(first_sector_offset as u64))
                .unwrap();
            block_file.write_all(&buffer).unwrap();
            block_file.flush().unwrap();
            sectors_to_be_read.fill(true);
        }

        if read_sectors.iter().all(|&x| x) {
            println!("Read the entire block. Exiting.");
            remove_file(&state_file_name).unwrap();
            return;
        }

        parts *= 2;

        let bytes: Vec<u8> = read_sectors.iter().copied().map(|x| x as u8).collect();
        state_file.seek(SeekFrom::Start(0)).unwrap();
        state_file.write_all(&bytes).unwrap();
        state_file.flush().unwrap();
    }

    remove_file(&state_file_name).unwrap();
    print!("Unread sectors: ");

    for (nth, is_read) in read_sectors.into_iter().enumerate() {
        if !is_read {
            print!("{} ", nth);
        }
    }

    println!();
}