From ae3b5d8855b93401f6857ec8d03471d4a2597809 Mon Sep 17 00:00:00 2001 From: Cutieguwu Date: Wed, 31 Dec 2025 15:27:41 -0500 Subject: [PATCH] Major fixes. Make git history even messier. --- src/io.rs | 2 +- src/main.rs | 5 +- src/mapping/map.rs | 60 +++++++++++---------- src/mapping/prelude.rs | 2 +- src/mapping/stage.rs | 2 +- src/recovery.rs | 117 ++++++++++++++++++++++++----------------- 6 files changed, 106 insertions(+), 82 deletions(-) diff --git a/src/io.rs b/src/io.rs index 8b51ece..a7a2d48 100644 --- a/src/io.rs +++ b/src/io.rs @@ -32,7 +32,7 @@ where } Ok(PathBuf::from(format!( - "{:?}.{}", + "{}.{}", source_name .as_ref() .to_str() diff --git a/src/main.rs b/src/main.rs index e6cbcfe..66d1de9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -41,12 +41,13 @@ fn main() -> anyhow::Result<()> { }; let map: MapFile = { - let path = crate::io::get_path(&config.output, &config.input, "map") + 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(|| { @@ -58,7 +59,7 @@ fn main() -> anyhow::Result<()> { let mut recover_tool = Recover::new(&config, &mut input, output, map)?; - recover_tool.run(); + recover_tool.run()?; Ok(()) } diff --git a/src/mapping/map.rs b/src/mapping/map.rs index c5d7401..b6f379c 100644 --- a/src/mapping/map.rs +++ b/src/mapping/map.rs @@ -83,7 +83,7 @@ impl MapFile { for cluster in self.map.iter() { match cluster.stage { Stage::Untested => return Stage::Untested, - Stage::ForIsolation(_) => { + Stage::ForIsolation { .. } => { if recover_stage == Stage::Damaged || cluster.stage < recover_stage { // Note that recover_stage after first condition is // only ever Stage::ForIsolation(_), thus PartialEq, @@ -186,20 +186,24 @@ impl MapFile { }); Some(delta) } -} -pub fn write_map_to_file(file: &mut File, map: &MapFile) -> anyhow::Result { - let written_bytes = file.write( - ron::ser::to_string_pretty( - map, - ron::ser::PrettyConfig::new() - .new_line("\n".to_string()) - .struct_names(true), - )? - .as_bytes(), - )?; + /// Writes the map to the provided item implementing `Write` trait. + /// Usually a file. + pub fn write_to(&mut self, file: &mut W) -> anyhow::Result { + self.defrag(); - Ok(written_bytes) + let written_bytes = file.write( + ron::ser::to_string_pretty( + self, + ron::ser::PrettyConfig::new() + .new_line("\n".to_string()) + .struct_names(true), + )? + .as_bytes(), + )?; + + Ok(written_bytes) + } } // This is split out for a shred of readability. @@ -730,8 +734,8 @@ mod tests { let stages = vec![ Stage::Damaged, - Stage::ForIsolation(1), - Stage::ForIsolation(0), + Stage::ForIsolation { level: 1 }, + Stage::ForIsolation { level: 0 }, Stage::Untested, ]; @@ -758,19 +762,19 @@ mod tests { mf.map = vec![ *Cluster::default().set_stage(Stage::Damaged), - *Cluster::default().set_stage(Stage::ForIsolation(0)), - *Cluster::default().set_stage(Stage::ForIsolation(1)), + *Cluster::default().set_stage(Stage::ForIsolation { level: 0 }), + *Cluster::default().set_stage(Stage::ForIsolation { level: 1 }), Cluster::default(), Cluster::default(), - *Cluster::default().set_stage(Stage::ForIsolation(1)), - *Cluster::default().set_stage(Stage::ForIsolation(0)), + *Cluster::default().set_stage(Stage::ForIsolation { level: 1 }), + *Cluster::default().set_stage(Stage::ForIsolation { level: 0 }), *Cluster::default().set_stage(Stage::Damaged), ]; let stages = vec![ Stage::Damaged, - Stage::ForIsolation(1), - Stage::ForIsolation(0), + Stage::ForIsolation { level: 1 }, + Stage::ForIsolation { level: 0 }, Stage::Untested, ]; @@ -811,19 +815,19 @@ mod tests { }, Cluster { domain: Domain { start: 3, end: 4 }, - stage: Stage::ForIsolation(0), + stage: Stage::ForIsolation { level: 0 }, }, Cluster { domain: Domain { start: 4, end: 5 }, - stage: Stage::ForIsolation(0), + stage: Stage::ForIsolation { level: 0 }, }, Cluster { domain: Domain { start: 5, end: 6 }, - stage: Stage::ForIsolation(1), + stage: Stage::ForIsolation { level: 1 }, }, Cluster { domain: Domain { start: 6, end: 7 }, - stage: Stage::ForIsolation(0), + stage: Stage::ForIsolation { level: 0 }, }, Cluster { domain: Domain { start: 7, end: 8 }, @@ -839,15 +843,15 @@ mod tests { }, Cluster { domain: Domain { start: 3, end: 5 }, - stage: Stage::ForIsolation(0), + stage: Stage::ForIsolation { level: 0 }, }, Cluster { domain: Domain { start: 5, end: 6 }, - stage: Stage::ForIsolation(1), + stage: Stage::ForIsolation { level: 1 }, }, Cluster { domain: Domain { start: 6, end: 7 }, - stage: Stage::ForIsolation(0), + stage: Stage::ForIsolation { level: 0 }, }, Cluster { domain: Domain { start: 7, end: 8 }, diff --git a/src/mapping/prelude.rs b/src/mapping/prelude.rs index 2f746eb..9d4b452 100644 --- a/src/mapping/prelude.rs +++ b/src/mapping/prelude.rs @@ -2,5 +2,5 @@ pub use super::cluster::Cluster; pub use super::domain::{Domain, DomainOverlap}; -pub use super::map::{write_map_to_file, MapFile}; +pub use super::map::MapFile; pub use super::stage::Stage; diff --git a/src/mapping/stage.rs b/src/mapping/stage.rs index cb5acd9..85813c9 100644 --- a/src/mapping/stage.rs +++ b/src/mapping/stage.rs @@ -3,7 +3,7 @@ use serde::{Deserialize, Serialize}; #[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)] pub enum Stage { Untested, - ForIsolation(u8), + ForIsolation { level: u8 }, Damaged, Intact, } diff --git a/src/recovery.rs b/src/recovery.rs index 41ca273..3e2c050 100644 --- a/src/recovery.rs +++ b/src/recovery.rs @@ -1,5 +1,6 @@ use std::fs::{File, OpenOptions}; use std::io::{BufWriter, Read, Seek, Write}; +use std::usize; use anyhow::Context; @@ -12,8 +13,6 @@ pub struct Recover<'a> { config: &'a Args, stream_len: u64, - buf: Vec, - input: &'a mut File, output: BufWriter, map: MapFile, @@ -32,7 +31,6 @@ impl<'a> Recover<'a> { let mut r = Recover { config, stream_len, - buf: vec![], input, output: BufWriter::with_capacity(stream_len as usize, output), map, @@ -44,35 +42,41 @@ impl<'a> Recover<'a> { } /// Recover media. - pub fn run(&mut self) -> () { + pub fn run(&mut self) -> anyhow::Result { // From start, read to end or error. // // If all data recovered, return early. // Else, read from end to error. - self.copy_untested(); - - /* let mut is_finished = false; while !is_finished { + self.map.defrag(); + match self.map.get_stage() { - Stage::Untested => { - self.copy_untested(); - } - Stage::ForIsolation(level) => { - self.copy_isolate(level); - } - Stage::Damaged => { + Stage::Untested => self.copy_untested()?, + Stage::ForIsolation { .. } => todo!(), + Stage::Damaged | Stage::Intact => { println!("Cannot recover further."); is_finished = true } - } - } - */ + }; - // return recovered_bytes + // Need to reset seek position between algorithms. + self.input + .rewind() + .context("Failed to reset input seek position.")?; + self.output + .rewind() + .context("Failed to reset output seek position")?; + + dbg!(); + } + + // Temporary. + let recovered_bytes = usize::MIN; + Ok(recovered_bytes) } /// Restore current progress based on MapFile. @@ -83,54 +87,69 @@ impl<'a> Recover<'a> { /// Attempt to copy all untested blocks. fn copy_untested(&mut self) -> anyhow::Result<()> { - // Caching. - let buf_capacity = self.get_buf_capacity(); + for untested in self.map.get_clusters(Stage::Untested) { + // Caching. + 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]; - self.buf = vec![crate::FB_PAD_VALUE; buf_capacity as usize]; + read_position = untested.domain.start; - // Yet more caching. - let mut read_position = 0_u64; - let final_read_position = self.stream_len - buf_capacity; - - while read_position < final_read_position { - if let Err(err) = self.input.read_exact(&mut self.buf) { - println!("Hit error: {:?}", err); - self.input - .seek_relative(buf_capacity as i64) - .context("Failed to seek input by buf_capacity to skip previous error")?; - } else { - self.output - .write_all(self.buf.as_slice()) - .context("Failed to write data to output file")?; - - self.map.update(Cluster { + while read_position < untested.domain.end { + cluster = Cluster { domain: Domain { - start: read_position as usize, - end: (read_position + buf_capacity) as usize, + start: read_position, + end: read_position + buf_capacity, }, stage: Stage::Intact, - }); - } + }; - read_position += buf_capacity; + buf_capacity = buf_capacity.min(untested.domain.end - read_position); + + if let Err(err) = self.input.read_exact(&mut buf) { + // If buf were zeroed out before every read, one could theoretically recover + // part of that read given the assumption that all null values from the end to + // the first non-null value are unread, and some further padding from the last + // values are potentially invalid. + // + // That padding should have a cli arg to control it. + + println!("Hit error: {:?}", err); + self.input + .seek_relative(buf_capacity as i64) + .context("Failed to seek input by buf_capacity to skip previous error")?; + + // I don't remember what level was for. + cluster.stage = Stage::ForIsolation { level: 1 }; + } + + if cluster.stage == Stage::Intact { + self.output + .write_all(&buf[0..buf_capacity]) + .context("Failed to write data to output file")?; + } + + self.map.update(cluster); + read_position += buf_capacity; + } } - write_map_to_file( - { + self.map + .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() - .create(true) .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()) })? - }, - &self.map, - ) - .context("Failed to write map file")?; + }) + .context("Failed to write map file")?; Ok(()) }