diff --git a/Cargo.lock b/Cargo.lock index 0944759..2a6e76b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -233,7 +233,6 @@ name = "kramer" version = "0.1.0" dependencies = [ "clap", - "libc", "ron", "rust-i18n", "serde", @@ -245,12 +244,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "libc" -version = "0.2.171" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c19937216e9d3aa9956d9bb8dfc0b0c8beb6058fc4f7a4dc4d850edf86a237d6" - [[package]] name = "libyml" version = "0.0.5" diff --git a/Cargo.toml b/Cargo.toml index 44023bc..efd9b56 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ edition = "2021" # # For clap info, see [dependencies.clap] # For serde info, see [dependencies.serde] -libc = "0.2.171, ~0.2.169" + ron = "0.8.1, >=0.8, <0.9" rust-i18n = "3.1.3, ~3.1.3" diff --git a/src/io.rs b/src/io.rs new file mode 100644 index 0000000..cc75e80 --- /dev/null +++ b/src/io.rs @@ -0,0 +1,15 @@ +use std::io::{self, Seek, SeekFrom}; + +/// Get length of data stream. +/// Physical length of data stream in bytes +/// (multiple of sector_size, rather than actual). +/// +/// This will attempt to return the stream to its current read position. +pub fn get_stream_length(stream: &mut S) -> io::Result { + let pos = stream.stream_position()?; + let len = stream.seek(SeekFrom::End(0)); + + stream.seek(SeekFrom::Start(pos))?; + + len +} diff --git a/src/main.rs b/src/main.rs index b05c176..208f96d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,18 +1,16 @@ +mod io; mod mapping; mod recovery; +use std::fs::{File, OpenOptions}; +use std::path::{Path, PathBuf}; + use clap::Parser; -use libc::O_DIRECT; use mapping::MapFile; use recovery::Recover; -use std::{ - fs::{File, OpenOptions}, - io::{self, Seek, SeekFrom}, - os::unix::fs::OpenOptionsExt, - path::PathBuf, -}; const FB_SECTOR_SIZE: u16 = 2048; +const FB_PAD_VALUE: u8 = 0; #[derive(Parser, Debug)] struct Args { @@ -48,35 +46,28 @@ fn main() { // I'm lazy and don't want to mess around with comparing error types. // Thus, any error in I/O here should be treated as fatal. - let mut input: File = { - match OpenOptions::new() - .custom_flags(O_DIRECT) - .read(true) - .open(&config.input.as_path()) - { - Ok(f) => f, - Err(err) => panic!("Failed to open input file: {:?}", err), - } - }; + let mut input: File = OpenOptions::new() + .read(true) + .open(&config.input.as_path()) + .expect("Failed to open input file"); - let mut output: File = { + let output: File = { // Keep this clean, make a short-lived binding. let path = get_path(&config.output, &config.input.to_str().unwrap(), "iso"); - match OpenOptions::new() - .custom_flags(O_DIRECT) + OpenOptions::new() .read(true) .write(true) .create(true) .open(path) - { - Ok(f) => f, - Err(err) => panic!("Failed to open/create output file. {:?}", err), - } + .expect("Failed to open/create output file") }; // Check if output file is shorter than input. // If so, autoextend the output file. + // + // Is this actually needed? I don't think I need to preallocate the space. + /* { let input_len = get_stream_length(&mut input).expect("Failed to get the length of the input data."); @@ -89,14 +80,24 @@ fn main() { .expect("Failed to autofill output file.") } } + */ + /* let map: MapFile = { - let path = get_path(&config.output, &config.input.to_str().unwrap(), "map"); + let path = get_path( + &config.output, + &config + .input + .to_str() + .expect("Input path is not UTF-8 valid"), + "map", + ); - let file = match OpenOptions::new().read(true).create(true).open(path) { - Ok(f) => f, - Err(err) => panic!("Failed to open/create mapping file. {:?}", err), - }; + let file = OpenOptions::new() + .read(true) + .create(true) + .open(path) + .expect("Failed to open/create mapping file"); if let Ok(map) = MapFile::try_from(file) { map @@ -104,37 +105,53 @@ fn main() { MapFile::new(config.sector_size) } }; + */ - let mut recover_tool = Recover::new(config, input, output, map); + let buf_capacity = + crate::io::get_stream_length(&mut input).expect("Failed to get buffer capacity from input"); + + let mut recover_tool = Recover::new(config, input, output, MapFile::default(), buf_capacity); recover_tool.run(); - todo!("Recovery, Map saving, and closure of all files."); + //todo!("Recovery, Map saving, and closure of all files."); + + /* + let mut buf: Vec = Vec::with_capacity( + get_stream_length(&mut input).expect("Failed to get the length of the input data.") + as usize, + ); + println!( + "Read {} bytes", + input + .read_to_end(&mut buf) + .expect("Failed to read complete input stream.") + ); + + println!("Wrote {} bytes", { + output + .write_all(&buf) + .expect("Failed to write complete output stream."); + &buf.len() + }); + */ } /// Generates a file path if one not provided. /// source_name for fallback name. -fn get_path(output: &Option, source_name: &str, extension: &str) -> PathBuf { - if let Some(f) = output { - f.to_owned() +fn get_path

(file_path: &Option

, source_name: &str, extension: &str) -> PathBuf +where + P: AsRef, +{ + if let Some(f) = file_path { + f.as_ref().to_path_buf() } else { - PathBuf::from(format!("{:?}.{}", source_name, extension,)) + PathBuf::from(format!("{:?}.{}", source_name, extension)) .as_path() .to_owned() } } -/// Get length of data stream. -/// Physical length of data stream in bytes -/// (multiple of sector_size, rather than actual). -fn get_stream_length(input: &mut S) -> io::Result { - let len = input.seek(SeekFrom::End(0))?; - - let _ = input.seek(SeekFrom::Start(0)); - - Ok(len) -} - #[cfg(test)] #[allow(unused)] mod tests { diff --git a/src/mapping.rs b/src/mapping.rs deleted file mode 100644 index 5b8f318..0000000 --- a/src/mapping.rs +++ /dev/null @@ -1,446 +0,0 @@ -use ron::de::{from_reader, SpannedError}; -use serde::Deserialize; -use std::fs::File; - -use crate::FB_SECTOR_SIZE; - -/// Domain, in sectors. -/// Requires sector_size to be provided elsewhere for conversion to bytes. -#[derive(Clone, Copy, Debug, Deserialize, PartialEq)] -pub struct Domain { - pub start: usize, - pub end: usize, -} - -impl Default for Domain { - fn default() -> Self { - Domain { start: 0, end: 1 } - } -} - -impl Domain { - /// Return length of domain in sectors. - pub fn len(self) -> usize { - self.end - self.start - } -} - -/// A map for data stored in memory for processing and saving to disk. -#[derive(Clone, Copy, Debug, Deserialize, PartialEq)] -pub struct Cluster { - domain: Domain, - stage: Stage, -} - -impl Default for Cluster { - fn default() -> Self { - Cluster { - domain: Domain::default(), - stage: Stage::default(), - } - } -} - -impl Cluster { - /// Breaks apart into a vec of clusters, - /// each of cluster_size, excepting last. - pub fn subdivide(&mut self, cluster_len: usize) -> Vec { - let domain_len = self.domain.len(); - let mut start = self.domain.start; - let mut clusters: Vec = vec![]; - - for _ in 0..(domain_len as f64 / cluster_len as f64).floor() as usize { - clusters.push(Cluster { - domain: Domain { - start, - end: start + cluster_len, - }, - stage: self.stage, - }); - - start += cluster_len; - } - - clusters.push(Cluster { - domain: Domain { - start, - end: self.domain.end, - }, - stage: self.stage, - }); - - clusters - } - - pub fn set_stage(&mut self, stage: Stage) -> &mut Self { - self.stage = stage; - self - } -} - -#[derive(Clone, Copy, Debug, Deserialize, PartialEq, PartialOrd)] -pub enum Stage { - Untested, - ForIsolation(u8), - Damaged, -} - -impl Default for Stage { - fn default() -> Self { - Stage::Untested - } -} - -#[derive(Clone, Debug, Deserialize, PartialEq)] -pub struct MapFile { - pub sector_size: u16, - pub domain: Domain, - pub map: Vec, -} - -impl TryFrom for MapFile { - type Error = SpannedError; - - fn try_from(file: File) -> Result { - from_reader(file) - } -} - -impl Default for MapFile { - fn default() -> Self { - MapFile { - sector_size: FB_SECTOR_SIZE, - domain: Domain::default(), - map: vec![Cluster { - domain: Domain::default(), - stage: Stage::Untested, - }], - } - } -} - -impl MapFile { - pub fn new(sector_size: u16) -> Self { - MapFile::default().set_sector_size(sector_size).to_owned() - } - - pub fn set_sector_size(&mut self, sector_size: u16) -> &mut Self { - self.sector_size = sector_size; - self - } - - /// Recalculate cluster mappings. - fn update(&mut self, new_cluster: Cluster) -> &mut Self { - let mut new_map: Vec = vec![Cluster::from(new_cluster.to_owned())]; - - for map_cluster in self.map.iter() { - let mut map_cluster = *map_cluster; - - // If new_cluster doesn't start ahead and ends short, map_cluster is forgotten. - if new_cluster.domain.start < map_cluster.domain.start - && new_cluster.domain.end < map_cluster.domain.end - { - /* - new_cluster overlaps the start of map_cluster, - but ends short of map_cluster end. - - ACTION: Crop map_cluster to start at end of new_cluster. - */ - - map_cluster.domain.start = new_cluster.domain.end; - new_map.push(map_cluster); - } else if new_cluster.domain.end < map_cluster.domain.end { - /* - new_cluster starts within map_cluster domain. - - ACTION: Crop - */ - - let domain_end = map_cluster.domain.end; - - // Crop current object. - map_cluster.domain.end = new_cluster.domain.start; - new_map.push(map_cluster); - - if new_cluster.domain.end < map_cluster.domain.end { - /* - new_cluster is within map_cluster. - - ACTION: Crop & Fracture map_cluster - NOTE: Crop completed above. - */ - - new_map.push(Cluster { - domain: Domain { - start: new_cluster.domain.end, - end: domain_end, - }, - stage: map_cluster.stage.to_owned(), - }); - } - } else { - /* - No overlap. - - ACTION: Transfer - */ - - new_map.push(map_cluster); - } - } - - self.map = new_map; - self - } - - /// Get current recovery stage. - pub fn get_stage(&self) -> Stage { - let mut recover_stage = Stage::Damaged; - - for cluster in self.map.iter() { - match cluster.stage { - Stage::Untested => return Stage::Untested, - 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, - // PartialOrd are useful for comparing the internal value. - recover_stage = cluster.stage - } - } - Stage::Damaged => (), - } - } - - recover_stage - } - - /// Get clusters of common stage. - pub fn get_clusters(&self, stage: Stage) -> Vec { - self.map - .iter() - .filter_map(|mc| { - if mc.stage == stage { - Some(mc.to_owned()) - } else { - None - } - }) - .collect() - } - - /// Defragments cluster groups. - /// I.E. check forwards every cluster from current until stage changes, - /// then group at once. - fn defrag(&mut self) -> &mut Self { - let mut new_map: Vec = vec![]; - - // Fetch first cluster. - let mut start_cluster = self.map.iter().find(|c| c.domain.start == 0).unwrap(); - - // Even though this would be initialized by its first read, - // the compiler won't stop whining, and idk how to assert that to it. - let mut end_cluster = Cluster::default(); - let mut new_cluster: Cluster; - - let mut stage_common: bool; - let mut is_finished = false; - - while !is_finished { - stage_common = true; - - // Start a new cluster based on the cluster following - // the end of last new_cluster. - new_cluster = start_cluster.to_owned(); - - // While stage is common, and not finished, - // find each trailing cluster. - while stage_common && !is_finished { - end_cluster = start_cluster.to_owned(); - - if end_cluster.domain.end != self.domain.end { - start_cluster = self - .map - .iter() - .find(|c| end_cluster.domain.end == c.domain.start) - .unwrap(); - - stage_common = new_cluster.stage == start_cluster.stage - } else { - is_finished = true; - } - } - - // Set the new ending, encapsulating any clusters of common stage. - new_cluster.domain.end = end_cluster.domain.end; - new_map.push(new_cluster); - } - - self.map = new_map; - self - } -} - -#[cfg(test)] -mod tests { - use super::*; - - // Test for Cluster::subdivide() - - // Test for MapFile::update() - - // Test for MapFile::get_stage() - #[test] - fn test_get_stage() { - use std::vec; - - let mut mf = MapFile::default(); - let mut mf_stage = mf.get_stage(); - - // If this fails here, there's something SERIOUSLY wrong. - assert!( - mf_stage == Stage::Untested, - "Determined stage to be {:?}, when {:?} was expeccted.", - mf_stage, - Stage::Untested - ); - - let stages = vec![ - Stage::Damaged, - Stage::ForIsolation(1), - Stage::ForIsolation(0), - Stage::Untested, - ]; - - mf.map = vec![]; - - for stage in stages { - mf.map.push(*Cluster::default().set_stage(stage)); - - mf_stage = mf.get_stage(); - - assert!( - stage == mf_stage, - "Expected stage to be {:?}, determined {:?} instead.", - stage, - mf_stage - ) - } - } - - // Test for MapFile::get_clusters() - #[test] - fn test_get_clusters() { - let mut mf = MapFile::default(); - - 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(), - Cluster::default(), - *Cluster::default().set_stage(Stage::ForIsolation(1)), - *Cluster::default().set_stage(Stage::ForIsolation(0)), - *Cluster::default().set_stage(Stage::Damaged), - ]; - - let stages = vec![ - Stage::Damaged, - Stage::ForIsolation(1), - Stage::ForIsolation(0), - Stage::Untested, - ]; - - for stage in stages { - let expected = vec![ - *Cluster::default().set_stage(stage), - *Cluster::default().set_stage(stage), - ]; - let received = mf.get_clusters(stage); - - assert!( - expected == received, - "Expected clusters {:?}, got {:?}.", - expected, - received - ) - } - } - - // Test for MapFile::defrag() - #[test] - fn test_defrag() { - let mut mf = MapFile { - sector_size: 1, - domain: Domain { start: 0, end: 8 }, - map: vec![ - Cluster { - domain: Domain { start: 0, end: 1 }, - stage: Stage::Untested, - }, - Cluster { - domain: Domain { start: 1, end: 2 }, - stage: Stage::Untested, - }, - Cluster { - domain: Domain { start: 2, end: 3 }, - stage: Stage::Untested, - }, - Cluster { - domain: Domain { start: 3, end: 4 }, - stage: Stage::ForIsolation(0), - }, - Cluster { - domain: Domain { start: 4, end: 5 }, - stage: Stage::ForIsolation(0), - }, - Cluster { - domain: Domain { start: 5, end: 6 }, - stage: Stage::ForIsolation(1), - }, - Cluster { - domain: Domain { start: 6, end: 7 }, - stage: Stage::ForIsolation(0), - }, - Cluster { - domain: Domain { start: 7, end: 8 }, - stage: Stage::Damaged, - }, - ], - }; - - let expected = vec![ - Cluster { - domain: Domain { start: 0, end: 3 }, - stage: Stage::Untested, - }, - Cluster { - domain: Domain { start: 3, end: 5 }, - stage: Stage::ForIsolation(0), - }, - Cluster { - domain: Domain { start: 5, end: 6 }, - stage: Stage::ForIsolation(1), - }, - Cluster { - domain: Domain { start: 6, end: 7 }, - stage: Stage::ForIsolation(0), - }, - Cluster { - domain: Domain { start: 7, end: 8 }, - stage: Stage::Damaged, - }, - ]; - - mf.defrag(); - - let received = mf.map; - - assert!( - expected == received, - "Expected {:?} after defragging, got {:?}.", - expected, - received - ) - } -} diff --git a/src/mapping/cluster.rs b/src/mapping/cluster.rs new file mode 100644 index 0000000..0a68709 --- /dev/null +++ b/src/mapping/cluster.rs @@ -0,0 +1,56 @@ +use super::{Domain, Stage}; + +use serde::{Deserialize, Serialize}; + +/// A map for data stored in memory for processing and saving to disk. +#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq)] +pub struct Cluster { + pub domain: Domain, + pub stage: Stage, +} + +impl Default for Cluster { + fn default() -> Self { + Cluster { + domain: Domain::default(), + stage: Stage::default(), + } + } +} + +impl Cluster { + /// Breaks apart into a vec of clusters, + /// each of cluster_size, excepting last. + pub fn subdivide(&mut self, cluster_len: usize) -> Vec { + let domain_len = self.domain.len(); + let mut start = self.domain.start; + let mut clusters: Vec = vec![]; + + for _ in 0..(domain_len / cluster_len) { + clusters.push(Cluster { + domain: Domain { + start, + end: start + cluster_len, + }, + stage: self.stage, + }); + + start += cluster_len; + } + + clusters.push(Cluster { + domain: Domain { + start, + end: self.domain.end, + }, + stage: self.stage, + }); + + clusters + } + + pub fn set_stage(&mut self, stage: Stage) -> &mut Self { + self.stage = stage; + self + } +} diff --git a/src/mapping/domain.rs b/src/mapping/domain.rs new file mode 100644 index 0000000..aff5163 --- /dev/null +++ b/src/mapping/domain.rs @@ -0,0 +1,22 @@ +use serde::{Deserialize, Serialize}; + +/// Domain, in sectors. +/// Requires sector_size to be provided elsewhere for conversion to bytes. +#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq)] +pub struct Domain { + pub start: usize, + pub end: usize, +} + +impl Default for Domain { + fn default() -> Self { + Domain { start: 0, end: 1 } + } +} + +impl Domain { + /// Return length of domain in sectors. + pub fn len(self) -> usize { + self.end - self.start + } +} diff --git a/src/mapping/map.rs b/src/mapping/map.rs new file mode 100644 index 0000000..65c3f8f --- /dev/null +++ b/src/mapping/map.rs @@ -0,0 +1,207 @@ +use std::fs::File; + +use super::{Cluster, Domain, Stage}; + +use ron::de::from_reader; +use ron::error::SpannedError; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] +pub struct MapFile { + pub sector_size: u16, + pub domain: Domain, + pub map: Vec, +} + +impl TryFrom for MapFile { + type Error = SpannedError; + + fn try_from(file: File) -> Result { + from_reader(file) + } +} + +impl Default for MapFile { + fn default() -> Self { + MapFile { + sector_size: crate::FB_SECTOR_SIZE, + domain: Domain::default(), + map: vec![Cluster { + domain: Domain::default(), + stage: Stage::Untested, + }], + } + } +} + +impl MapFile { + pub fn new(sector_size: u16) -> Self { + MapFile::default().set_sector_size(sector_size).to_owned() + } + + pub fn set_sector_size(&mut self, sector_size: u16) -> &mut Self { + self.sector_size = sector_size; + self + } + + /// Recalculate cluster mappings. + pub fn update(&mut self, new_cluster: Cluster) -> &mut Self { + let mut new_map: Vec = vec![Cluster::from(new_cluster.to_owned())]; + + for map_cluster in self.map.iter() { + let mut map_cluster = *map_cluster; + + // If new_cluster doesn't start ahead and ends short, map_cluster is forgotten. + if new_cluster.domain.start < map_cluster.domain.start + && new_cluster.domain.end < map_cluster.domain.end + { + /* + new_cluster overlaps the start of map_cluster, + but ends short of map_cluster end. + + ACTION: Crop map_cluster to start at end of new_cluster. + */ + + map_cluster.domain.start = new_cluster.domain.end; + new_map.push(map_cluster); + } else if new_cluster.domain.end < map_cluster.domain.end { + /* + new_cluster starts within map_cluster domain. + + ACTION: Crop + */ + + let domain_end = map_cluster.domain.end; + + // Crop current object. + map_cluster.domain.end = new_cluster.domain.start; + new_map.push(map_cluster); + + if new_cluster.domain.end < map_cluster.domain.end { + /* + new_cluster is within map_cluster. + + ACTION: Crop & Fracture map_cluster + NOTE: Crop completed above. + */ + + new_map.push(Cluster { + domain: Domain { + start: new_cluster.domain.end, + end: domain_end, + }, + stage: map_cluster.stage.to_owned(), + }); + } + } else { + /* + No overlap. + + ACTION: Transfer + */ + + new_map.push(map_cluster); + } + } + + self.map = new_map; + self + } + + /// Get current recovery stage. + pub fn get_stage(&self) -> Stage { + let mut recover_stage = Stage::Damaged; + + for cluster in self.map.iter() { + match cluster.stage { + Stage::Untested => return Stage::Untested, + 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, + // PartialOrd are useful for comparing the internal value. + recover_stage = cluster.stage + } + } + Stage::Damaged => (), + Stage::Intact => unreachable!(), + } + } + + recover_stage + } + + /// Get clusters of common stage. + pub fn get_clusters(&self, stage: Stage) -> Vec { + self.map + .iter() + .filter_map(|mc| { + if mc.stage == stage { + Some(mc.to_owned()) + } else { + None + } + }) + .collect() + } + + /// Defragments cluster groups. + /// I.E. check forwards every cluster from current until stage changes, + /// then group at once. + pub fn defrag(&mut self) -> &mut Self { + let mut new_map: Vec = vec![]; + + // Fetch first cluster. + let mut start_cluster = self.map.iter().find(|c| c.domain.start == 0).unwrap(); + + // Even though this would be initialized by its first read, + // the compiler won't stop whining, and idk how to assert that to it. + let mut end_cluster = Cluster::default(); + let mut new_cluster: Cluster; + + let mut stage_common: bool; + let mut is_finished = false; + + while !is_finished { + stage_common = true; + + // Start a new cluster based on the cluster following + // the end of last new_cluster. + new_cluster = start_cluster.to_owned(); + + // While stage is common, and not finished, + // find each trailing cluster. + while stage_common && !is_finished { + end_cluster = start_cluster.to_owned(); + + if end_cluster.domain.end != self.domain.end { + start_cluster = self + .map + .iter() + .find(|c| end_cluster.domain.end == c.domain.start) + .unwrap(); + + stage_common = new_cluster.stage == start_cluster.stage + } else { + is_finished = true; + } + } + + // Set the new ending, encapsulating any clusters of common stage. + new_cluster.domain.end = end_cluster.domain.end; + new_map.push(new_cluster); + } + + self.map = new_map; + self + } +} + +pub fn write_map_to_file(file: File, map: &MapFile) -> ron::error::Result { + ron::ser::to_string_pretty( + map, + ron::ser::PrettyConfig::new() + .new_line("\n".to_string()) + .struct_names(true), + ) +} diff --git a/src/mapping/mod.rs b/src/mapping/mod.rs new file mode 100644 index 0000000..1791df0 --- /dev/null +++ b/src/mapping/mod.rs @@ -0,0 +1,187 @@ +pub mod cluster; +pub mod domain; +pub mod map; + +use serde::{Deserialize, Serialize}; + +pub use cluster::Cluster; +pub use domain::Domain; +pub use map::MapFile; + +#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, PartialOrd)] +pub enum Stage { + Intact, + Untested, + ForIsolation(u8), + Damaged, +} + +impl Default for Stage { + fn default() -> Self { + Stage::Untested + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // Test for Cluster::subdivide() + + // Test for MapFile::update() + + // Test for MapFile::get_stage() + #[test] + fn test_get_stage() { + use std::vec; + + let mut mf = MapFile::default(); + let mut mf_stage = mf.get_stage(); + + // If this fails here, there's something SERIOUSLY wrong. + assert!( + mf_stage == Stage::Untested, + "Determined stage to be {:?}, when {:?} was expeccted.", + mf_stage, + Stage::Untested + ); + + let stages = vec![ + Stage::Damaged, + Stage::ForIsolation(1), + Stage::ForIsolation(0), + Stage::Untested, + ]; + + mf.map = vec![]; + + for stage in stages { + mf.map.push(*Cluster::default().set_stage(stage)); + + mf_stage = mf.get_stage(); + + assert!( + stage == mf_stage, + "Expected stage to be {:?}, determined {:?} instead.", + stage, + mf_stage + ) + } + } + + // Test for MapFile::get_clusters() + #[test] + fn test_get_clusters() { + let mut mf = MapFile::default(); + + 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(), + Cluster::default(), + *Cluster::default().set_stage(Stage::ForIsolation(1)), + *Cluster::default().set_stage(Stage::ForIsolation(0)), + *Cluster::default().set_stage(Stage::Damaged), + ]; + + let stages = vec![ + Stage::Damaged, + Stage::ForIsolation(1), + Stage::ForIsolation(0), + Stage::Untested, + ]; + + for stage in stages { + let expected = vec![ + *Cluster::default().set_stage(stage), + *Cluster::default().set_stage(stage), + ]; + let received = mf.get_clusters(stage); + + assert!( + expected == received, + "Expected clusters {:?}, got {:?}.", + expected, + received + ) + } + } + + // Test for MapFile::defrag() + #[test] + fn test_defrag() { + let mut mf = MapFile { + sector_size: 1, + domain: Domain { start: 0, end: 8 }, + map: vec![ + Cluster { + domain: Domain { start: 0, end: 1 }, + stage: Stage::Untested, + }, + Cluster { + domain: Domain { start: 1, end: 2 }, + stage: Stage::Untested, + }, + Cluster { + domain: Domain { start: 2, end: 3 }, + stage: Stage::Untested, + }, + Cluster { + domain: Domain { start: 3, end: 4 }, + stage: Stage::ForIsolation(0), + }, + Cluster { + domain: Domain { start: 4, end: 5 }, + stage: Stage::ForIsolation(0), + }, + Cluster { + domain: Domain { start: 5, end: 6 }, + stage: Stage::ForIsolation(1), + }, + Cluster { + domain: Domain { start: 6, end: 7 }, + stage: Stage::ForIsolation(0), + }, + Cluster { + domain: Domain { start: 7, end: 8 }, + stage: Stage::Damaged, + }, + ], + }; + + let expected = vec![ + Cluster { + domain: Domain { start: 0, end: 3 }, + stage: Stage::Untested, + }, + Cluster { + domain: Domain { start: 3, end: 5 }, + stage: Stage::ForIsolation(0), + }, + Cluster { + domain: Domain { start: 5, end: 6 }, + stage: Stage::ForIsolation(1), + }, + Cluster { + domain: Domain { start: 6, end: 7 }, + stage: Stage::ForIsolation(0), + }, + Cluster { + domain: Domain { start: 7, end: 8 }, + stage: Stage::Damaged, + }, + ]; + + mf.defrag(); + + let received = mf.map; + + assert!( + expected == received, + "Expected {:?} after defragging, got {:?}.", + expected, + received + ) + } +} diff --git a/src/recovery.rs b/src/recovery.rs index 010ea8a..aae1371 100644 --- a/src/recovery.rs +++ b/src/recovery.rs @@ -1,46 +1,32 @@ -use std::{ - io::{BufReader, BufWriter}, - fs::File, -}; - -use crate::{ - Args, - mapping::{Cluster, MapFile, Stage}, -}; +use std::fs::{File, OpenOptions}; +use std::io::{BufWriter, Read, Seek, SeekFrom, Write}; +use std::ptr::read; +use crate::mapping::{Cluster, Domain, MapFile, Stage}; +use crate::Args; #[derive(Debug)] +#[allow(dead_code)] pub struct Recover { - buf_capacity: usize, + /// Buffer capacity in bytes. + buf_capacity: u64, config: Args, - input: BufReader, + + input: File, output: BufWriter, map: MapFile, stage: Stage, } impl Recover { - pub fn new( - config: Args, - input: File, - output: File, - map: MapFile, - ) -> Self { + pub fn new(config: Args, input: File, output: File, map: MapFile, buf_capacity: u64) -> Self { let stage = map.get_stage(); - // Temporarily make buffer length one sector. - let buf_capacity = config.sector_size as usize; let mut r = Recover { buf_capacity, config, - input: BufReader::with_capacity( - buf_capacity, - input, - ), - output: BufWriter::with_capacity( - buf_capacity, - output, - ), + input: input, + output: BufWriter::with_capacity(buf_capacity as usize, output), map, stage: stage, }; @@ -51,27 +37,82 @@ impl Recover { } /// Recover media. - pub fn run(&mut self) -> &mut Self { + pub fn run(&mut self) -> () { + self.copy_untested(); + + /* let mut is_finished = false; while !is_finished { match self.map.get_stage() { - Stage::Untested => { self.copy_untested(); }, - Stage::ForIsolation(level) => { self.copy_isolate(level); }, + Stage::Untested => { + self.copy_untested(); + } + Stage::ForIsolation(level) => { + self.copy_isolate(level); + } Stage::Damaged => { println!("Cannot recover further."); is_finished = true - }, + } } } + */ - self + // return recovered_bytes } /// Attempt to copy all untested blocks. fn copy_untested(&mut self) -> &mut Self { + let mut buf = vec![crate::FB_PAD_VALUE; self.buf_capacity as usize]; + // Purely caching. + let mut read_position = 0_u64; + let last_read_position = crate::io::get_stream_length(&mut self.input) + .expect("Failed to get length of input stream") + - self.buf_capacity; + + while read_position < last_read_position { + dbg!(read_position); + + if let Err(err) = self.input.read_exact(&mut buf) { + println!("Hit error: {:?}", err); + self.input + .seek_relative(self.buf_capacity as i64) + .expect("Failed to seek input by buf_capacity to skip previous error"); + } else { + self.output + .write_all(buf.as_slice()) + .expect("Failed to write data to output file"); + + self.map.update(Cluster { + domain: Domain { + start: read_position as usize, + end: (read_position + self.buf_capacity) as usize, + }, + stage: Stage::Intact, + }); + } + + read_position += self.buf_capacity; + } + + crate::mapping::map::write_map_to_file( + OpenOptions::new() + .create(true) + .write(true) + .open(crate::get_path( + &self.config.map, + self.config.input.to_str().unwrap(), + "map", + )) + .expect("Failed to open map file"), + &self.map, + ) + .expect("Failed to write map file"); + + /* let mut untested: Vec = vec![]; for cluster in self.map.get_clusters(Stage::Untested).iter_mut() { @@ -80,12 +121,13 @@ impl Recover { todo!("Read and save data."); + */ + self } /// Attempt to copy blocks via isolation at pass level. fn copy_isolate(&mut self, level: u8) -> &mut Self { - todo!(); self @@ -94,17 +136,16 @@ impl Recover { /// Set buffer capacities as cluster length in bytes. /// Varies depending on the recovery stage. fn set_buf_capacity(&mut self) -> &mut Self { - self.buf_capacity = (self.config.sector_size * self.config.cluster_length) as usize; + self.buf_capacity = self.config.sector_size as u64 * self.config.cluster_length as u64; self } } - #[cfg(test)] #[allow(unused)] mod tests { use super::*; // Test for Recover::set_buf_capacity -} \ No newline at end of file +}