Tons of cleanup. More useless git notices.

src/cli.rs

@@ -1,10 +1,13 @@
 use std::path::PathBuf;
+use std::sync::LazyLock;
 
 use crate::FB_SECTOR_SIZE;
 
-use clap::Parser;
+use clap::{ArgAction, Parser};
 
-#[derive(Parser, Debug)]
+pub static CONFIG: LazyLock<Args> = LazyLock::new(|| Args::parse());
+
+#[derive(Parser, Debug, Clone)]
 pub struct Args {
     /// Path to source file or block device
     #[arg(short, long, value_hint = clap::ValueHint::DirPath)]
@@ -20,7 +23,7 @@ pub struct Args {
 
     /// Max number of consecutive sectors to test as a group
     #[arg(short, long, default_value_t = 128)]
-    pub cluster_length: u16,
+    pub cluster_length: usize,
 
     /// Number of brute force read passes
     #[arg(short, long, default_value_t = 2)]
@@ -28,5 +31,12 @@ pub struct Args {
 
     /// Sector size
     #[arg(short, long, default_value_t = FB_SECTOR_SIZE)]
-    pub sector_size: u16,
+    pub sector_size: usize,
+
+    // Behaviour is backwards.
+    // ArgAction::SetFalse by default evaluates to true,
+    // ArgAction::SetTrue by default evaluates to false.
+    /// Whether to reopen the file on a read error or not.
+    #[arg(short, long, action=ArgAction::SetTrue)]
+    pub reopen_on_error: bool,
 }

src/io.rs

@@ -1,5 +1,7 @@
+use std::fs::{File, OpenOptions};
 use std::io::{self, Seek, SeekFrom};
-use std::path::{Path, PathBuf};
+
+use crate::cli::CONFIG;
 
 use anyhow::Context;
 
@@ -17,28 +19,49 @@ pub fn get_stream_length<S: Seek>(stream: &mut S) -> io::Result<u64> {
     len
 }
 
-/// Generates a file path if one not provided.
+pub fn load_input() -> anyhow::Result<File> {
-/// source_name for fallback name.
+    OpenOptions::new()
-pub fn get_path<P>(
+        .read(true)
-    file_path: &Option<P>,
+        .open(&CONFIG.input)
-    source_name: &P,
+        .with_context(|| format!("Failed to open input file: {}", &CONFIG.input.display()))
-    extension: &str,
-) -> anyhow::Result<PathBuf>
-where
-    P: AsRef<Path>,
-{
-    if let Some(f) = file_path {
-        return Ok(f.as_ref().to_path_buf());
 }
 
-    Ok(PathBuf::from(format!(
+pub fn load_output() -> anyhow::Result<File> {
-        "{}.{}",
+    OpenOptions::new()
-        source_name
+        .read(true)
-            .as_ref()
+        .write(true)
-            .to_str()
+        .create(true)
-            .context("source_name path was not UTF-8 valid.")?,
+        .open(crate::path::OUTPUT_PATH.clone())
-        extension
+        .with_context(|| {
-    ))
+            format!(
-    .as_path()
+                "Failed to open/create output file at: {}",
-    .to_owned())
+                crate::path::OUTPUT_PATH.display()
+            )
+        })
+}
+
+pub fn load_map_read() -> anyhow::Result<File> {
+    OpenOptions::new()
+        .read(true)
+        .open(crate::path::MAP_PATH.clone())
+        .with_context(|| {
+            format!(
+                "Failed to open/create mapping file at: {}",
+                crate::path::MAP_PATH.display()
+            )
+        })
+}
+
+pub fn load_map_write() -> anyhow::Result<File> {
+    OpenOptions::new()
+        .write(true)
+        .create(true)
+        .truncate(true) // Wipe old map. Should really make a backup first.
+        .open(crate::path::MAP_PATH.clone())
+        .with_context(|| {
+            format!(
+                "Failed to open map file at: {}",
+                crate::path::MAP_PATH.display()
+            )
+        })
 }

src/main.rs

@@ -1,79 +1,19 @@
 mod cli;
 mod io;
 mod mapping;
+mod path;
 mod recovery;
 
-use std::fs::{File, OpenOptions};
-
-use cli::Args;
-use mapping::prelude::*;
 use recovery::Recover;
 
-use anyhow::{self, Context};
+use anyhow;
-use clap::Parser;
 
-const FB_SECTOR_SIZE: u16 = 2048;
+const FB_SECTOR_SIZE: usize = 2048;
-const FB_PAD_VALUE: u8 = 0;
+const FB_NULL_VALUE: u8 = 0;
 
 fn main() -> anyhow::Result<()> {
-    let config = Args::parse();
+    let mut recover_tool = Recover::new()?;
-
-    let mut input: File = {
-        let input_path = &config.input.as_path();
-
-        OpenOptions::new()
-            .read(true)
-            .open(&input_path)
-            .with_context(|| format!("Failed to open input file: {}", input_path.display()))?
-    };
-
-    let output: File = {
-        // Keep this clean, make a short-lived binding.
-        let path = crate::io::get_path(&config.output, &config.input, "iso")
-            .context("Failed to generate output path.")?;
-
-        OpenOptions::new()
-            .read(true)
-            .write(true)
-            .create(true)
-            .open(&path)
-            .with_context(|| format!("Failed to open/create output file at: {}", path.display()))?
-    };
-
-    let map: MapFile = {
-        let path = crate::io::get_path(&config.map, &config.input, "map")
-            .context("Failed to generate map path.")?;
-
-        MapFile::try_from(
-            OpenOptions::new()
-                .read(true)
-                .write(true)
-                .create(true)
-                .open(&path)
-                .with_context(|| {
-                    format!("Failed to open/create mapping file at: {}", path.display())
-                })?,
-        )
-        .unwrap_or(MapFile::new(config.sector_size))
-    };
-
-    let mut recover_tool = Recover::new(&config, &mut input, output, map)?;
-
     recover_tool.run()?;
 
     Ok(())
 }
-
-#[cfg(test)]
-#[allow(unused)]
-mod tests {
-    use super::*;
-
-    // Test for get_path
-    // Need to determine how to package files to test with, or at least
-    // how to test with PathBuf present.
-    // Test must also check unwrapping of file name, not just generation.
-
-    // Test for get_stream_length
-    // Need to determine how to test with Seek-able objects.
-}

src/mapping/map.rs

@@ -10,7 +10,7 @@ use serde::{Deserialize, Serialize};
 
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
 pub struct MapFile {
-    pub sector_size: u16,
+    pub sector_size: usize,
     pub domain: Domain,
     pub map: Vec<Cluster>,
 }
@@ -37,11 +37,11 @@ impl Default for MapFile {
 }
 
 impl MapFile {
-    pub fn new(sector_size: u16) -> Self {
+    pub fn new(sector_size: usize) -> Self {
         MapFile::default().set_sector_size(sector_size).to_owned()
     }
 
-    pub fn set_sector_size(&mut self, sector_size: u16) -> &mut Self {
+    pub fn set_sector_size(&mut self, sector_size: usize) -> &mut Self {
         self.sector_size = sector_size;
         self
     }

src/path.rs

@@ -0,0 +1,40 @@
+use std::path::{Path, PathBuf};
+use std::sync::LazyLock;
+
+use crate::cli::CONFIG;
+
+use anyhow::{self, Context};
+
+/// Generates a file path if one not provided.
+/// root_path for fallback name.
+pub fn get_path<P>(path: &Option<P>, root_path: &P, extension: &str) -> anyhow::Result<PathBuf>
+where
+    P: AsRef<Path>,
+{
+    if let Some(f) = path {
+        return Ok(f.as_ref().to_path_buf());
+    }
+
+    Ok(PathBuf::from(format!(
+        "{}.{}",
+        root_path
+            .as_ref()
+            .to_str()
+            .context("source_name path was not UTF-8 valid.")?,
+        extension
+    ))
+    .as_path()
+    .to_owned())
+}
+
+pub static MAP_PATH: LazyLock<PathBuf> = LazyLock::new(|| {
+    get_path(&CONFIG.map, &CONFIG.input, "map")
+        .context("Failed to generate map path.")
+        .unwrap()
+});
+
+pub static OUTPUT_PATH: LazyLock<PathBuf> = LazyLock::new(|| {
+    get_path(&CONFIG.output, &CONFIG.input, "iso")
+        .context("Failed to generate output path.")
+        .unwrap()
+});

src/recovery.rs

@@ -1,42 +1,39 @@
-use std::fs::{File, OpenOptions};
+use std::fs::File;
-use std::io::{BufWriter, Read, Seek, Write};
+use std::io::{BufWriter, Read, Seek, SeekFrom, Write};
 use std::usize;
 
 use anyhow::Context;
 
-use crate::cli::Args;
+use crate::cli::CONFIG;
 use crate::mapping::prelude::*;
 
 #[derive(Debug)]
 #[allow(dead_code)]
-pub struct Recover<'a> {
+pub struct Recover {
-    config: &'a Args,
+    input: File,
-    stream_len: u64,
-
-    input: &'a mut File,
     output: BufWriter<File>,
     map: MapFile,
 }
 
-impl<'a> Recover<'a> {
+impl Recover {
-    pub fn new(
+    pub fn new() -> anyhow::Result<Self> {
-        config: &'a Args,
+        let input: File = crate::io::load_input()?;
-        input: &'a mut File,
+
-        output: File,
+        let output: File = crate::io::load_output()?;
-        map: MapFile,
+
-    ) -> anyhow::Result<Self> {
+        let map: MapFile = {
-        let stream_len =
+            crate::io::load_map_read()?
-            crate::io::get_stream_length(input).context("Failed to get input stream length.")?;
+                .try_into()
+                .unwrap_or(MapFile::new(CONFIG.sector_size))
+        };
 
         let mut r = Recover {
-            config,
-            stream_len,
             input,
-            output: BufWriter::with_capacity(stream_len as usize, output),
+            output: BufWriter::with_capacity(map.domain.end as usize, output),
             map,
         };
 
-        r.restore();
+        r.restore()?;
 
         Ok(r)
     }
@@ -70,8 +67,6 @@ impl<'a> Recover<'a> {
             self.output
                 .rewind()
                 .context("Failed to reset output seek position")?;
-
-            dbg!();
         }
 
         // Temporary.
@@ -81,8 +76,13 @@ impl<'a> Recover<'a> {
 
     /// Restore current progress based on MapFile.
     /// Also updates MapFile if needed, such as to extend the MapFile domain.
-    pub fn restore(&mut self) {
+    pub fn restore(&mut self) -> anyhow::Result<()> {
-        self.map.extend(self.stream_len as usize);
+        self.map.extend(
+            crate::io::get_stream_length(&mut self.input)
+                .context("Failed to get input stream length.")? as usize,
+        );
+
+        Ok(())
     }
 
     /// Attempt to copy all untested blocks.
@@ -92,11 +92,13 @@ impl<'a> Recover<'a> {
             let mut read_position: usize;
             let mut cluster: Cluster;
             let mut buf_capacity = self.get_buf_capacity() as usize;
-            let mut buf = vec![crate::FB_PAD_VALUE; buf_capacity];
+            let mut buf = vec![crate::FB_NULL_VALUE; buf_capacity];
 
             read_position = untested.domain.start;
 
             while read_position < untested.domain.end {
+                dbg!(read_position);
+
                 cluster = Cluster {
                     domain: Domain {
                         start: read_position,
@@ -116,8 +118,13 @@ impl<'a> Recover<'a> {
                     // That padding should have a cli arg to control it.
 
                     println!("Hit error: {:?}", err);
+                    if CONFIG.reopen_on_error {
+                        self.reload_input()
+                            .context("Failed to reload input file after previous error.")?;
+                    }
+
                     self.input
-                        .seek_relative(buf_capacity as i64)
+                        .seek_relative((read_position + buf_capacity) as i64)
                         .context("Failed to seek input by buf_capacity to skip previous error")?;
 
                     // I don't remember what level was for.
@@ -135,21 +142,7 @@ impl<'a> Recover<'a> {
             }
         }
 
-        self.map
+        self.map.write_to(&mut crate::io::load_map_write()?)?;
-            .write_to({
-                let map_path = crate::io::get_path(&self.config.map, &self.config.input, "map")
-                    .context("Failed to generate map path.")?;
-
-                &mut OpenOptions::new()
-                    .write(true)
-                    .create(true)
-                    .truncate(true) // Wipe old map. Should really make a backup first.
-                    .open(&map_path)
-                    .with_context(|| {
-                        format!("Failed to open map file at: {}", map_path.display())
-                    })?
-            })
-            .context("Failed to write map file")?;
 
         Ok(())
     }
@@ -157,7 +150,16 @@ impl<'a> Recover<'a> {
     /// Set buffer capacity as cluster length in bytes.
     /// Varies depending on the recovery stage.
     fn get_buf_capacity(&mut self) -> u64 {
-        self.config.sector_size as u64 * self.config.cluster_length as u64
+        CONFIG.sector_size as u64 * CONFIG.cluster_length as u64
+    }
+
+    /// Reloads the input and restores the seek position.
+    fn reload_input(&mut self) -> anyhow::Result<()> {
+        let seek_pos = self.input.stream_position()?;
+
+        self.input = crate::io::load_input()?;
+        self.input.seek(SeekFrom::Start(seek_pos))?;
+        Ok(())
     }
 }