use crate::archive::Archive;
use crate::detect::{sniff, Format};
use crate::error::{ArchiveError, Result};
const HEAD_PEEK: u64 = 512;
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum Access {
InPlace {
offset: u64,
len: u64,
},
Zran,
SpillToTemp,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum Codec {
Gzip,
Bzip2,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum SegmentKind {
Ewf,
SplitRaw,
SplitVmdk,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Segment {
pub name: String,
pub index: usize,
pub access: Access,
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum AccessPlan {
Direct,
Wrapper {
codec: Codec,
access: Access,
},
Member {
format: Format,
index: usize,
name: String,
access: Access,
},
SegmentSet {
format: Format,
members: Vec<Segment>,
kind: SegmentKind,
},
Collection {
format: Format,
},
}
pub fn detect(data: &[u8]) -> Result<AccessPlan> {
let fmt = sniff(None, data);
match fmt {
Format::Gzip => detect_wrapper(Codec::Gzip, data),
Format::Bzip2 => detect_wrapper(Codec::Bzip2, data),
Format::Zip | Format::SevenZip | Format::Tar | Format::TarGz | Format::TarBz2 => {
detect_archive(fmt, data)
}
_ => Ok(AccessPlan::Direct),
}
}
fn detect_wrapper(codec: Codec, data: &[u8]) -> Result<AccessPlan> {
let Ok(head) = peek_head(codec, data) else {
return Ok(AccessPlan::Direct);
};
match sniff(None, &head) {
Format::Tar => {
let archive_fmt = match codec {
Codec::Gzip => Format::TarGz,
Codec::Bzip2 => Format::TarBz2,
};
detect_archive(archive_fmt, data)
}
Format::Zip
| Format::SevenZip
| Format::Gzip
| Format::Bzip2
| Format::TarGz
| Format::TarBz2 => Ok(AccessPlan::Wrapper {
codec,
access: Access::SpillToTemp,
}),
_ => Ok(AccessPlan::Wrapper {
codec,
access: wrapper_access(codec),
}),
}
}
fn detect_archive(format: Format, data: &[u8]) -> Result<AccessPlan> {
let Some(mut archive) = Archive::open_with_format(format, data)? else {
return Ok(AccessPlan::Collection { format });
};
let files: Vec<(usize, String)> = archive
.entries()
.iter()
.enumerate()
.filter(|(_, e)| !e.is_dir)
.map(|(i, e)| (i, e.name.clone()))
.collect();
if files.len() >= 2 {
if let Some(kind) = classify_segment_kind(&files) {
let ordered = order_segments(&files, kind);
let mut members = Vec::with_capacity(ordered.len());
for (index, name, _seg) in ordered {
let access = archive.member_access(index)?;
members.push(Segment {
name,
index,
access,
});
}
return Ok(AccessPlan::SegmentSet {
format,
members,
kind,
});
}
}
match files.as_slice() {
[(index, name)] => {
let access = archive.member_access(*index)?;
Ok(AccessPlan::Member {
format,
index: *index,
name: name.clone(),
access,
})
}
_ => Ok(AccessPlan::Collection { format }),
}
}
fn peek_head(codec: Codec, data: &[u8]) -> Result<Vec<u8>> {
use std::io::Read;
let reader: Box<dyn Read> = match codec {
Codec::Gzip => Box::new(flate2::read::GzDecoder::new(data)),
Codec::Bzip2 => Box::new(bzip2_rs::DecoderReader::new(data)),
};
let mut out = Vec::new();
reader
.take(HEAD_PEEK)
.read_to_end(&mut out)
.map_err(|e| ArchiveError::Decode {
format: codec_name(codec),
detail: e.to_string(),
})?;
Ok(out)
}
fn wrapper_access(codec: Codec) -> Access {
match codec {
Codec::Gzip => Access::Zran,
Codec::Bzip2 => Access::SpillToTemp,
}
}
fn codec_name(codec: Codec) -> &'static str {
match codec {
Codec::Gzip => "gzip",
Codec::Bzip2 => "bzip2",
}
}
fn segment_number(name: &str, kind: SegmentKind) -> Option<u64> {
match kind {
SegmentKind::Ewf => ewf_segment(name),
SegmentKind::SplitRaw => raw_split(name),
SegmentKind::SplitVmdk => vmdk_segment(name),
}
}
fn classify_segment_kind(files: &[(usize, String)]) -> Option<SegmentKind> {
[
SegmentKind::SplitVmdk,
SegmentKind::Ewf,
SegmentKind::SplitRaw,
]
.into_iter()
.find(|&kind| files.iter().all(|(_, n)| segment_number(n, kind).is_some()))
}
fn order_segments(files: &[(usize, String)], kind: SegmentKind) -> Vec<(usize, String, u64)> {
let mut ordered: Vec<(usize, String, u64)> = files
.iter()
.filter_map(|(i, n)| segment_number(n, kind).map(|seg| (*i, n.clone(), seg)))
.collect();
ordered.sort_by_key(|(_, _, seg)| *seg);
ordered
}
fn ewf_segment(name: &str) -> Option<u64> {
let (_, ext) = name.rsplit_once('.')?;
let ext = ext.to_ascii_lowercase();
let digits = ext
.strip_prefix("ex")
.or_else(|| ext.strip_prefix('e'))
.or_else(|| ext.strip_prefix('s'))?;
if digits.len() == 2 && digits.bytes().all(|b| b.is_ascii_digit()) {
digits.parse::<u64>().ok()
} else {
None
}
}
fn raw_split(name: &str) -> Option<u64> {
let (_, ext) = name.rsplit_once('.')?;
if ext.len() >= 2 && ext.bytes().all(|b| b.is_ascii_digit()) {
ext.parse::<u64>().ok()
} else {
None
}
}
fn vmdk_segment(name: &str) -> Option<u64> {
let lower = name.to_ascii_lowercase();
let stem = lower.strip_suffix(".vmdk")?;
let pos = stem.rfind("-s")?;
let num = stem.get(pos + 2..)?;
if !num.is_empty() && num.bytes().all(|b| b.is_ascii_digit()) {
num.parse::<u64>().ok()
} else {
None
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
const FX: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../tests/data/fixtures/");
fn load(name: &str) -> Vec<u8> {
std::fs::read(format!("{FX}{name}")).unwrap()
}
fn gzip(data: &[u8]) -> Vec<u8> {
let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
e.write_all(data).unwrap();
e.finish().unwrap()
}
fn build_tar(members: &[(&str, Vec<u8>)]) -> Vec<u8> {
let mut b = tar::Builder::new(Vec::new());
for (name, data) in members {
let mut h = tar::Header::new_gnu();
h.set_size(data.len() as u64);
h.set_mode(0o644);
h.set_cksum();
b.append_data(&mut h, name, data.as_slice()).unwrap();
}
b.into_inner().unwrap()
}
#[test]
fn bare_gzip_of_raw_bytes_is_wrapper_zran() {
let gz = gzip(&b"raw disk sector bytes, not an archive at all ".repeat(40));
assert_eq!(
detect(&gz).unwrap(),
AccessPlan::Wrapper {
codec: Codec::Gzip,
access: Access::Zran
}
);
}
#[test]
fn gzip_of_single_member_tar_is_targz_member() {
let tar = build_tar(&[("disk.img", b"RAW-IMAGE-BYTES".to_vec())]);
let gz = gzip(&tar);
match detect(&gz).unwrap() {
AccessPlan::Member {
format,
name,
access,
..
} => {
assert_eq!(format, Format::TarGz);
assert_eq!(name, "disk.img");
assert_eq!(access, Access::SpillToTemp);
}
other => panic!("expected TarGz Member, got {other:?}"),
}
}
#[test]
fn coincidental_gzip_magic_is_direct() {
assert_eq!(
detect(b"\x1f\x8b\x00\x00garbage-not-really-gzip").unwrap(),
AccessPlan::Direct
);
}
#[test]
fn raw_bytes_are_direct() {
assert_eq!(
detect(b"\x00\x01\x02 not a wrapper or archive").unwrap(),
AccessPlan::Direct
);
}
#[test]
fn bare_bzip2_of_raw_bytes_is_wrapper_spill() {
assert_eq!(
detect(&load("payload.bz2")).unwrap(),
AccessPlan::Wrapper {
codec: Codec::Bzip2,
access: Access::SpillToTemp
}
);
}
#[test]
fn coincidental_bzip2_magic_is_direct() {
assert_eq!(
detect(b"BZhnot-a-real-bzip2-stream").unwrap(),
AccessPlan::Direct
);
}
#[test]
fn gzip_of_zip_is_nested_wrapper_spill() {
let gz = gzip(&load("payload.zip"));
assert_eq!(
detect(&gz).unwrap(),
AccessPlan::Wrapper {
codec: Codec::Gzip,
access: Access::SpillToTemp
}
);
}
#[test]
fn zip_single_bzip2_member_spills() {
match detect(&load("bzip2_member.zip")).unwrap() {
AccessPlan::Member {
format,
name,
access,
..
} => {
assert_eq!(format, Format::Zip);
assert_eq!(name, "blob.bin");
assert_eq!(access, Access::SpillToTemp);
}
other => panic!("expected Member, got {other:?}"),
}
}
#[test]
fn zip_single_stored_member_is_inplace() {
match detect(&load("stored_one.zip")).unwrap() {
AccessPlan::Member {
format,
name,
access,
..
} => {
assert_eq!(format, Format::Zip);
assert_eq!(name, "disk.dd");
match access {
Access::InPlace { offset, len } => {
assert_eq!(len, 4096);
assert!(offset > 0, "stored data starts after a local header");
}
other => panic!("expected InPlace, got {other:?}"),
}
}
other => panic!("expected Member, got {other:?}"),
}
}
#[test]
fn zip_single_deflate_member_is_zran() {
match detect(&load("deflate_one.zip")).unwrap() {
AccessPlan::Member {
format,
name,
access,
..
} => {
assert_eq!(format, Format::Zip);
assert_eq!(name, "big.dd");
assert_eq!(access, Access::Zran);
}
other => panic!("expected Member, got {other:?}"),
}
}
#[test]
fn zip_ewf_segments_order_by_number_each_inplace() {
match detect(&load("seg_ewf.zip")).unwrap() {
AccessPlan::SegmentSet {
format,
members,
kind,
} => {
assert_eq!(format, Format::Zip);
assert_eq!(kind, SegmentKind::Ewf);
let names: Vec<&str> = members.iter().map(|s| s.name.as_str()).collect();
assert_eq!(names, ["img.E01", "img.E02", "img.E03"]);
for s in &members {
assert!(
matches!(s.access, Access::InPlace { .. }),
"stored segment → InPlace, got {:?}",
s.access
);
}
}
other => panic!("expected SegmentSet Ewf, got {other:?}"),
}
}
#[test]
fn zip_raw_split_is_segmentset_splitraw() {
match detect(&load("seg_split.zip")).unwrap() {
AccessPlan::SegmentSet {
format,
members,
kind,
} => {
assert_eq!(format, Format::Zip);
assert_eq!(kind, SegmentKind::SplitRaw);
let names: Vec<&str> = members.iter().map(|s| s.name.as_str()).collect();
assert_eq!(names, ["disk.001", "disk.002"]);
}
other => panic!("expected SegmentSet SplitRaw, got {other:?}"),
}
}
#[test]
fn tar_split_vmdk_is_segmentset_splitvmdk() {
let tar = build_tar(&[
("disk-s002.vmdk", b"seg-two".to_vec()),
("disk-s001.vmdk", b"seg-one".to_vec()),
]);
match detect(&tar).unwrap() {
AccessPlan::SegmentSet {
format,
members,
kind,
} => {
assert_eq!(format, Format::Tar);
assert_eq!(kind, SegmentKind::SplitVmdk);
let names: Vec<&str> = members.iter().map(|s| s.name.as_str()).collect();
assert_eq!(names, ["disk-s001.vmdk", "disk-s002.vmdk"]);
assert!(members.iter().all(|s| s.access == Access::SpillToTemp));
}
other => panic!("expected SegmentSet SplitVmdk, got {other:?}"),
}
}
#[test]
fn zip_unrelated_members_is_collection() {
assert_eq!(
detect(&load("payload.zip")).unwrap(),
AccessPlan::Collection {
format: Format::Zip
}
);
}
#[test]
fn sevenzip_unrelated_members_is_collection() {
assert_eq!(
detect(&load("payload.7z")).unwrap(),
AccessPlan::Collection {
format: Format::SevenZip
}
);
}
#[test]
fn bzip2_tar_classified_by_decompressed_content_not_name() {
match detect(&load("payload.tbz2")).unwrap() {
AccessPlan::Collection { format } => assert_eq!(format, Format::TarBz2),
other => panic!("expected TarBz2 Collection from content, got {other:?}"),
}
}
#[test]
fn ewf_segment_matches_e_ex_s_only() {
assert_eq!(ewf_segment("img.E01"), Some(1));
assert_eq!(ewf_segment("img.e12"), Some(12));
assert_eq!(ewf_segment("img.Ex03"), Some(3));
assert_eq!(ewf_segment("img.s07"), Some(7));
assert_eq!(ewf_segment("notes.txt"), None);
assert_eq!(ewf_segment("tool.exe"), None);
assert_eq!(ewf_segment("img.E1"), None); assert_eq!(ewf_segment("noext"), None);
}
#[test]
fn raw_split_matches_all_digit_ext() {
assert_eq!(raw_split("disk.001"), Some(1));
assert_eq!(raw_split("disk.017"), Some(17));
assert_eq!(raw_split("disk.E01"), None);
assert_eq!(raw_split("disk.1"), None); assert_eq!(raw_split("noext"), None);
}
#[test]
fn vmdk_segment_matches_dash_s_only() {
assert_eq!(vmdk_segment("disk-s001.vmdk"), Some(1));
assert_eq!(vmdk_segment("disk-s012.vmdk"), Some(12));
assert_eq!(vmdk_segment("disk.vmdk"), None); assert_eq!(vmdk_segment("disk-flat.vmdk"), None);
assert_eq!(vmdk_segment("disk-s001.bin"), None); }
#[test]
fn vmdk_segment_rejects_malformed_s_index() {
assert_eq!(vmdk_segment("disk-sx.vmdk"), None); assert_eq!(vmdk_segment("disk-s.vmdk"), None); }
}