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

68
src/main.rs
View File

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

72
src/mapping.rs
View File

@@ -4,7 +4,6 @@ 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)]
@@ -26,7 +25,6 @@ impl Domain {
}
}
/// A map for data stored in memory for processing and saving to disk.
#[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
pub struct Cluster {
@@ -38,7 +36,7 @@ impl Default for Cluster {
fn default() -> Self {
Cluster {
domain: Domain::default(),
stage: Stage::default()
stage: Stage::default(),
}
}
}
@@ -80,7 +78,6 @@ impl Cluster {
}
}
#[derive(Clone, Copy, Debug, Deserialize, PartialEq, PartialOrd)]
pub enum Stage {
Untested,
@@ -94,7 +91,6 @@ impl Default for Stage {
}
}
#[derive(Clone, Debug, Deserialize, PartialEq)]
pub struct MapFile {
pub sector_size: u16,
@@ -125,9 +121,7 @@ impl Default for MapFile {
impl MapFile {
pub fn new(sector_size: u16) -> Self {
MapFile::default()
.set_sector_size(sector_size)
.to_owned()
MapFile::default().set_sector_size(sector_size).to_owned()
}
pub fn set_sector_size(&mut self, sector_size: u16) -> &mut Self {
@@ -143,10 +137,11 @@ impl MapFile {
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,
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.
@@ -154,7 +149,6 @@ impl MapFile {
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.
@@ -181,7 +175,7 @@ impl MapFile {
start: new_cluster.domain.end,
end: domain_end,
},
stage: map_cluster.stage.to_owned()
stage: map_cluster.stage.to_owned(),
});
}
} else {
@@ -207,14 +201,13 @@ impl MapFile {
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
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 => (),
}
}
@@ -224,9 +217,14 @@ impl MapFile {
/// Get clusters of common stage.
pub fn get_clusters(&self, stage: Stage) -> Vec<Cluster> {
self.map.iter()
self.map
.iter()
.filter_map(|mc| {
if mc.stage == stage { Some(mc.to_owned()) } else { None }
if mc.stage == stage {
Some(mc.to_owned())
} else {
None
}
})
.collect()
}
@@ -238,10 +236,8 @@ impl MapFile {
let mut new_map: Vec<Cluster> = vec![];
// Fetch first cluster.
let mut start_cluster = self.map.iter()
.find(|c| c.domain.start == 0)
.unwrap();
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();
@@ -263,7 +259,9 @@ impl MapFile {
end_cluster = start_cluster.to_owned();
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)
.unwrap();
@@ -283,7 +281,6 @@ impl MapFile {
}
}
#[cfg(test)]
mod tests {
use super::*;
@@ -304,10 +301,10 @@ mod tests {
assert!(
mf_stage == Stage::Untested,
"Determined stage to be {:?}, when {:?} was expeccted.",
mf_stage, Stage::Untested
mf_stage,
Stage::Untested
);
let stages = vec![
Stage::Damaged,
Stage::ForIsolation(1),
@@ -325,7 +322,8 @@ mod tests {
assert!(
stage == mf_stage,
"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),
];
let recieved = mf.get_clusters(stage);
let received = mf.get_clusters(stage);
assert!(
expected == recieved,
expected == received,
"Expected clusters {:?}, got {:?}.",
expected, recieved
expected,
received
)
}
}
@@ -435,12 +434,13 @@ mod tests {
mf.defrag();
let recieved = mf.map;
let received = mf.map;
assert!(
expected == recieved,
expected == received,
"Expected {:?} after defragging, got {:?}.",
expected, recieved
expected,
received
)
}
}
}