Cutieguwu/kramer
1
0
Fork 0
Code Issues Pull Requests 4 Actions Packages Projects Releases Wiki Activity

Clean up and reformat.

Browse Source
This commit is contained in:
Olivia Brooks
2025-12-26 16:51:44 -05:00
parent 3d1273981c
commit 824d01be95
2 changed files with 57 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
src/main.rs
View File

@@ -1,5 +1,5 @@
mod recovery;
mod mapping; mod mapping;
mod recovery;
use clap::Parser; use clap::Parser;
use libc::O_DIRECT; use libc::O_DIRECT;
@@ -12,10 +12,8 @@ use std::{
path::PathBuf, path::PathBuf,
}; };
const FB_SECTOR_SIZE: u16 = 2048; const FB_SECTOR_SIZE: u16 = 2048;
#[derive(Parser, Debug)] #[derive(Parser, Debug)]
struct Args { struct Args {
/// Path to source file or block device /// Path to source file or block device
@@ -43,7 +41,6 @@ struct Args {
sector_size: u16, sector_size: u16,
} }
fn main() { fn main() {
let config = Args::parse(); let config = Args::parse();
@@ -55,23 +52,16 @@ fn main() {
match OpenOptions::new() match OpenOptions::new()
.custom_flags(O_DIRECT) .custom_flags(O_DIRECT)
.read(true) .read(true)
.write(false)
.append(false)
.create(false)
.open(&config.input.as_path()) .open(&config.input.as_path())
{ {
Ok(f) => f, Ok(f) => f,
Err(err) => panic!("Failed to open input file: {:?}", err) Err(err) => panic!("Failed to open input file: {:?}", err),
} }
}; };
let mut output: File = { let mut output: File = {
// Keep this clean, make a short-lived binding. // Keep this clean, make a short-lived binding.
let path = get_path( let path = get_path(&config.output, &config.input.to_str().unwrap(), "iso");
&config.output,
&config.input.to_str().unwrap(),
"iso"
);
match OpenOptions::new() match OpenOptions::new()
.custom_flags(O_DIRECT) .custom_flags(O_DIRECT)
@@ -81,38 +71,31 @@ fn main() {
.open(path) .open(path)
{ {
Ok(f) => f, Ok(f) => f,
Err(err) => panic!("Failed to open/create output file. {:?}", err) Err(err) => panic!("Failed to open/create output file. {:?}", err),
} }
}; };
// Check if output file is shorter than input. // Check if output file is shorter than input.
// If so, autoextend the output file. // If so, autoextend the output file.
{ {
let input_len = get_stream_length(&mut input) let input_len =
.expect("Failed to get the length of the input data."); get_stream_length(&mut input).expect("Failed to get the length of the input data.");
let output_len = get_stream_length(&mut output) let output_len =
.expect("Failed to get the length of the output file."); get_stream_length(&mut output).expect("Failed to get the length of the output file.");
if output_len < input_len { if output_len < input_len {
output.set_len(input_len) output
.set_len(input_len)
.expect("Failed to autofill output file.") .expect("Failed to autofill output file.")
} }
} }
let map: MapFile = { let map: MapFile = {
let path = get_path( let path = get_path(&config.output, &config.input.to_str().unwrap(), "map");
&config.output,
&config.input.to_str().unwrap(),
"map"
);
let file = match OpenOptions::new() let file = match OpenOptions::new().read(true).create(true).open(path) {
.read(true)
.create(true)
.open(path)
{
Ok(f) => f, Ok(f) => f,
Err(err) => panic!("Failed to open/create mapping file. {:?}", err) Err(err) => panic!("Failed to open/create mapping file. {:?}", err),
}; };
if let Ok(map) = MapFile::try_from(file) { if let Ok(map) = MapFile::try_from(file) {
@@ -122,7 +105,7 @@ fn main() {
} }
}; };
let mut recover_tool = Recover::new(config, input, output, map); let mut recover_tool = Recover::new(config, input, output, map);
recover_tool.run(); recover_tool.run();
@@ -131,21 +114,13 @@ fn main() {
/// Generates a file path if one not provided. /// Generates a file path if one not provided.
/// source_name for fallback name. /// source_name for fallback name.
fn get_path( fn get_path(output: &Option<PathBuf>, source_name: &str, extension: &str) -> PathBuf {
output: &Option<PathBuf>,
source_name: &str,
extention: &str
) -> PathBuf {
if let Some(f) = output { if let Some(f) = output {
f.to_owned() f.to_owned()
} else { } else {
PathBuf::from(format!( PathBuf::from(format!("{:?}.{}", source_name, extension,))
"{:?}.{}", .as_path()
source_name, .to_owned()
extention,
))
.as_path()
.to_owned()
} }
} }
@@ -160,7 +135,6 @@ fn get_stream_length<S: Seek>(input: &mut S) -> io::Result<u64> {
Ok(len) Ok(len)
} }
#[cfg(test)] #[cfg(test)]
#[allow(unused)] #[allow(unused)]
mod tests { mod tests {

60
src/mapping.rs
View File

@@ -4,7 +4,6 @@ use std::fs::File;
use crate::FB_SECTOR_SIZE; use crate::FB_SECTOR_SIZE;
/// Domain, in sectors. /// Domain, in sectors.
/// Requires sector_size to be provided elsewhere for conversion to bytes. /// Requires sector_size to be provided elsewhere for conversion to bytes.
#[derive(Clone, Copy, Debug, Deserialize, PartialEq)] #[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
@@ -26,7 +25,6 @@ impl Domain {
} }
} }
/// A map for data stored in memory for processing and saving to disk. /// A map for data stored in memory for processing and saving to disk.
#[derive(Clone, Copy, Debug, Deserialize, PartialEq)] #[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
pub struct Cluster { pub struct Cluster {
@@ -38,7 +36,7 @@ impl Default for Cluster {
fn default() -> Self { fn default() -> Self {
Cluster { Cluster {
domain: Domain::default(), domain: Domain::default(),
stage: Stage::default() stage: Stage::default(),
} }
} }
} }
@@ -80,7 +78,6 @@ impl Cluster {
} }
} }
#[derive(Clone, Copy, Debug, Deserialize, PartialEq, PartialOrd)] #[derive(Clone, Copy, Debug, Deserialize, PartialEq, PartialOrd)]
pub enum Stage { pub enum Stage {
Untested, Untested,
@@ -94,7 +91,6 @@ impl Default for Stage {
} }
} }
#[derive(Clone, Debug, Deserialize, PartialEq)] #[derive(Clone, Debug, Deserialize, PartialEq)]
pub struct MapFile { pub struct MapFile {
pub sector_size: u16, pub sector_size: u16,
@@ -125,9 +121,7 @@ impl Default for MapFile {
impl MapFile { impl MapFile {
pub fn new(sector_size: u16) -> Self { pub fn new(sector_size: u16) -> Self {
MapFile::default() MapFile::default().set_sector_size(sector_size).to_owned()
.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: u16) -> &mut Self {
@@ -144,7 +138,8 @@ impl MapFile {
// If new_cluster doesn't start ahead and ends short, map_cluster is forgotten. // If new_cluster doesn't start ahead and ends short, map_cluster is forgotten.
if new_cluster.domain.start < map_cluster.domain.start if new_cluster.domain.start < map_cluster.domain.start
&& new_cluster.domain.end < map_cluster.domain.end { && new_cluster.domain.end < map_cluster.domain.end
{
/* /*
new_cluster overlaps the start of map_cluster, new_cluster overlaps the start of map_cluster,
but ends short of map_cluster end. but ends short of map_cluster end.
@@ -154,7 +149,6 @@ impl MapFile {
map_cluster.domain.start = new_cluster.domain.end; map_cluster.domain.start = new_cluster.domain.end;
new_map.push(map_cluster); new_map.push(map_cluster);
} else if new_cluster.domain.end < map_cluster.domain.end { } else if new_cluster.domain.end < map_cluster.domain.end {
/* /*
new_cluster starts within map_cluster domain. new_cluster starts within map_cluster domain.
@@ -181,7 +175,7 @@ impl MapFile {
start: new_cluster.domain.end, start: new_cluster.domain.end,
end: domain_end, end: domain_end,
}, },
stage: map_cluster.stage.to_owned() stage: map_cluster.stage.to_owned(),
}); });
} }
} else { } else {
@@ -207,14 +201,13 @@ impl MapFile {
match cluster.stage { match cluster.stage {
Stage::Untested => return Stage::Untested, Stage::Untested => return Stage::Untested,
Stage::ForIsolation(_) => { Stage::ForIsolation(_) => {
if recover_stage == Stage::Damaged if recover_stage == Stage::Damaged || cluster.stage < recover_stage {
|| cluster.stage < recover_stage {
// Note that recover_stage after first condition is // Note that recover_stage after first condition is
// only ever Stage::ForIsolation(_), thus PartialEq, // only ever Stage::ForIsolation(_), thus PartialEq,
// PartialOrd are useful for comparing the internal value. // PartialOrd are useful for comparing the internal value.
recover_stage = cluster.stage recover_stage = cluster.stage
} }
}, }
Stage::Damaged => (), Stage::Damaged => (),
} }
} }
@@ -224,9 +217,14 @@ impl MapFile {
/// Get clusters of common stage. /// Get clusters of common stage.
pub fn get_clusters(&self, stage: Stage) -> Vec<Cluster> { pub fn get_clusters(&self, stage: Stage) -> Vec<Cluster> {
self.map.iter() self.map
.iter()
.filter_map(|mc| { .filter_map(|mc| {
if mc.stage == stage { Some(mc.to_owned()) } else { None } if mc.stage == stage {
Some(mc.to_owned())
} else {
None
}
}) })
.collect() .collect()
} }
@@ -238,9 +236,7 @@ impl MapFile {
let mut new_map: Vec<Cluster> = vec![]; let mut new_map: Vec<Cluster> = vec![];
// Fetch first cluster. // Fetch first cluster.
let mut start_cluster = self.map.iter() let mut start_cluster = self.map.iter().find(|c| c.domain.start == 0).unwrap();
.find(|c| c.domain.start == 0)
.unwrap();
// Even though this would be initialized by its first read, // Even though this would be initialized by its first read,
// the compiler won't stop whining, and idk how to assert that to it. // the compiler won't stop whining, and idk how to assert that to it.
@@ -263,7 +259,9 @@ impl MapFile {
end_cluster = start_cluster.to_owned(); end_cluster = start_cluster.to_owned();
if end_cluster.domain.end != self.domain.end { if end_cluster.domain.end != self.domain.end {
start_cluster = self.map.iter() start_cluster = self
.map
.iter()
.find(|c| end_cluster.domain.end == c.domain.start) .find(|c| end_cluster.domain.end == c.domain.start)
.unwrap(); .unwrap();
@@ -283,7 +281,6 @@ impl MapFile {
} }
} }
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
@@ -304,10 +301,10 @@ mod tests {
assert!( assert!(
mf_stage == Stage::Untested, mf_stage == Stage::Untested,
"Determined stage to be {:?}, when {:?} was expeccted.", "Determined stage to be {:?}, when {:?} was expeccted.",
mf_stage, Stage::Untested mf_stage,
Stage::Untested
); );
let stages = vec![ let stages = vec![
Stage::Damaged, Stage::Damaged,
Stage::ForIsolation(1), Stage::ForIsolation(1),
@@ -325,7 +322,8 @@ mod tests {
assert!( assert!(
stage == mf_stage, stage == mf_stage,
"Expected stage to be {:?}, determined {:?} instead.", "Expected stage to be {:?}, determined {:?} instead.",
stage, mf_stage stage,
mf_stage
) )
} }
} }
@@ -358,12 +356,13 @@ mod tests {
*Cluster::default().set_stage(stage), *Cluster::default().set_stage(stage),
*Cluster::default().set_stage(stage), *Cluster::default().set_stage(stage),
]; ];
let recieved = mf.get_clusters(stage); let received = mf.get_clusters(stage);
assert!( assert!(
expected == recieved, expected == received,
"Expected clusters {:?}, got {:?}.", "Expected clusters {:?}, got {:?}.",
expected, recieved expected,
received
) )
} }
} }
@@ -435,12 +434,13 @@ mod tests {
mf.defrag(); mf.defrag();
let recieved = mf.map; let received = mf.map;
assert!( assert!(
expected == recieved, expected == received,
"Expected {:?} after defragging, got {:?}.", "Expected {:?} after defragging, got {:?}.",
expected, recieved expected,
received
) )
} }
} }