use std::io::{self, Read, Seek, SeekFrom};
use std::sync::Arc;
use crate::detect::Format;
use crate::error::Result;
use crate::execute::{execute_member_access, offset_from, PeeledSource};
use crate::plan::{detect, AccessPlan, SegmentKind};
use crate::resolve::Limits;
#[derive(Debug)]
#[non_exhaustive]
pub struct SegmentSources {
pub format: Format,
pub kind: SegmentKind,
pub sources: Vec<PeeledSource>,
}
#[derive(Debug)]
pub struct ConcatSource {
segs: Vec<ConcatSeg>,
total: u64,
pos: u64,
}
#[derive(Debug)]
struct ConcatSeg {
source: PeeledSource,
start: u64,
len: u64,
}
impl ConcatSource {
pub(crate) fn new(sources: Vec<PeeledSource>) -> Self {
let mut segs = Vec::with_capacity(sources.len());
let mut start = 0u64;
for source in sources {
let len = source.len();
segs.push(ConcatSeg { source, start, len });
start = start.saturating_add(len);
}
ConcatSource {
segs,
total: start,
pos: 0,
}
}
#[must_use]
pub fn len(&self) -> u64 {
self.total
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.total == 0
}
fn seg_at(&self, pos: u64) -> Option<usize> {
let idx = self.segs.partition_point(|s| s.start + s.len <= pos);
(idx < self.segs.len()).then_some(idx)
}
}
impl Read for ConcatSource {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if buf.is_empty() || self.pos >= self.total {
return Ok(0);
}
let pos = self.pos;
let Some(i) = self.seg_at(pos) else {
return Ok(0); };
let seg = &mut self.segs[i];
let local = pos - seg.start;
seg.source.seek(SeekFrom::Start(local))?;
let avail = seg.len - local;
let want = (buf.len() as u64).min(avail) as usize;
let n = seg.source.read(&mut buf[..want])?;
self.pos = self.pos.saturating_add(n as u64);
Ok(n)
}
}
impl Seek for ConcatSource {
fn seek(&mut self, from: SeekFrom) -> io::Result<u64> {
let target = match from {
SeekFrom::Start(o) => o,
SeekFrom::End(o) => offset_from(self.total, o)?,
SeekFrom::Current(o) => offset_from(self.pos, o)?,
};
self.pos = target.min(self.total);
Ok(self.pos)
}
}
#[derive(Debug)]
#[non_exhaustive]
pub enum Reassembled {
Concat(ConcatSource),
Segments(SegmentSources),
NotSegmented,
}
pub fn segment_sources(data: Vec<u8>, limits: &Limits) -> Result<Option<SegmentSources>> {
let AccessPlan::SegmentSet {
format,
members,
kind,
} = detect(&data)?
else {
return Ok(None);
};
let data = Arc::new(data);
let mut sources = Vec::with_capacity(members.len());
for seg in &members {
sources.push(execute_member_access(
format,
&data,
seg.index,
&seg.name,
&seg.access,
limits,
)?);
}
Ok(Some(SegmentSources {
format,
kind,
sources,
}))
}
pub fn reassemble(data: Vec<u8>, limits: &Limits) -> Result<Reassembled> {
let Some(ss) = segment_sources(data, limits)? else {
return Ok(Reassembled::NotSegmented);
};
Ok(match ss.kind {
SegmentKind::SplitRaw => Reassembled::Concat(ConcatSource::new(ss.sources)),
SegmentKind::Ewf | SegmentKind::SplitVmdk => Reassembled::Segments(ss),
})
}
#[cfg(test)]
mod tests {
use super::*;
fn pattern(n: usize, seed: usize) -> Vec<u8> {
(0..n).map(|i| ((i + seed) % 251) as u8).collect()
}
fn stored_zip(members: &[(&str, &[u8])]) -> Vec<u8> {
let mut z = Vec::new();
let mut central = Vec::new();
let mut count = 0u16;
for (name, payload) in members {
let nb = name.as_bytes();
let mut crc = flate2::Crc::new();
crc.update(payload);
let crc = crc.sum();
let lho = z.len() as u32;
let (sz, nlen) = (payload.len() as u32, nb.len() as u16);
z.extend_from_slice(&0x0403_4b50u32.to_le_bytes());
z.extend_from_slice(&20u16.to_le_bytes()); z.extend_from_slice(&0u16.to_le_bytes()); z.extend_from_slice(&0u16.to_le_bytes()); z.extend_from_slice(&0u16.to_le_bytes()); z.extend_from_slice(&0u16.to_le_bytes()); z.extend_from_slice(&crc.to_le_bytes());
z.extend_from_slice(&sz.to_le_bytes()); z.extend_from_slice(&sz.to_le_bytes()); z.extend_from_slice(&nlen.to_le_bytes());
z.extend_from_slice(&0u16.to_le_bytes()); z.extend_from_slice(nb);
z.extend_from_slice(payload);
central.extend_from_slice(&0x0201_4b50u32.to_le_bytes());
central.extend_from_slice(&20u16.to_le_bytes()); central.extend_from_slice(&20u16.to_le_bytes()); central.extend_from_slice(&0u16.to_le_bytes()); central.extend_from_slice(&0u16.to_le_bytes()); central.extend_from_slice(&0u16.to_le_bytes()); central.extend_from_slice(&0u16.to_le_bytes()); central.extend_from_slice(&crc.to_le_bytes());
central.extend_from_slice(&sz.to_le_bytes());
central.extend_from_slice(&sz.to_le_bytes());
central.extend_from_slice(&nlen.to_le_bytes());
central.extend_from_slice(&0u16.to_le_bytes()); central.extend_from_slice(&0u16.to_le_bytes()); central.extend_from_slice(&0u16.to_le_bytes()); central.extend_from_slice(&0u16.to_le_bytes()); central.extend_from_slice(&0u32.to_le_bytes()); central.extend_from_slice(&lho.to_le_bytes()); central.extend_from_slice(nb);
count += 1;
}
let cd_offset = z.len() as u32;
let cd_size = central.len() as u32;
z.extend_from_slice(¢ral);
z.extend_from_slice(&0x0605_4b50u32.to_le_bytes());
z.extend_from_slice(&0u16.to_le_bytes()); z.extend_from_slice(&0u16.to_le_bytes()); z.extend_from_slice(&count.to_le_bytes()); z.extend_from_slice(&count.to_le_bytes()); z.extend_from_slice(&cd_size.to_le_bytes());
z.extend_from_slice(&cd_offset.to_le_bytes());
z.extend_from_slice(&0u16.to_le_bytes()); z
}
const FX: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../tests/data/fixtures/");
fn load(name: &str) -> Vec<u8> {
std::fs::read(format!("{FX}{name}")).unwrap()
}
#[test]
fn split_raw_concatenates_in_order_with_random_access() {
let p1 = pattern(1000, 0);
let p2 = pattern(1500, 100);
let p3 = pattern(777, 200);
let data = stored_zip(&[("img.002", &p2), ("img.003", &p3), ("img.001", &p1)]);
let Reassembled::Concat(mut c) = reassemble(data, &Limits::default()).unwrap() else {
panic!("expected a SplitRaw ConcatSource");
};
let total = (p1.len() + p2.len() + p3.len()) as u64;
assert_eq!(c.len(), total);
assert!(!c.is_empty());
let mut whole = Vec::new();
c.seek(SeekFrom::Start(0)).unwrap();
c.read_to_end(&mut whole).unwrap();
let mut expected = Vec::new();
expected.extend_from_slice(&p1);
expected.extend_from_slice(&p2);
expected.extend_from_slice(&p3);
assert_eq!(whole, expected);
c.seek(SeekFrom::Start(995)).unwrap();
let mut span = [0u8; 12];
c.read_exact(&mut span).unwrap();
assert_eq!(&span[..5], &p1[995..1000]);
assert_eq!(&span[5..], &p2[..7]);
c.seek(SeekFrom::Start(2495)).unwrap();
let mut span2 = [0u8; 10];
c.read_exact(&mut span2).unwrap();
assert_eq!(&span2[..5], &p2[p2.len() - 5..]);
assert_eq!(&span2[5..], &p3[..5]);
c.seek(SeekFrom::End(-4)).unwrap();
let mut tail = [0u8; 4];
c.read_exact(&mut tail).unwrap();
assert_eq!(tail, p3[p3.len() - 4..]);
assert_eq!(c.seek(SeekFrom::Start(u64::MAX)).unwrap(), total);
assert_eq!(c.read(&mut [0u8; 8]).unwrap(), 0);
}
#[test]
fn split_raw_segment_sources_are_inplace_and_ordered() {
let p1 = pattern(300, 1);
let p2 = pattern(400, 2);
let data = stored_zip(&[("d.002", &p2), ("d.001", &p1)]);
let ss = segment_sources(data, &Limits::default())
.unwrap()
.expect("a segment set");
assert_eq!(ss.kind, SegmentKind::SplitRaw);
assert_eq!(ss.sources.len(), 2);
let expected = [p1, p2];
for (src, want) in ss.sources.into_iter().zip(expected) {
assert!(
matches!(src, PeeledSource::InPlace(_)),
"a stored segment is a zero-copy InPlace window"
);
let mut got = Vec::new();
let mut src = src;
src.read_to_end(&mut got).unwrap();
assert_eq!(got, want, "segment bytes read in order");
}
}
#[test]
fn ewf_segment_sources_are_ordered_per_segment_handles() {
let data = load("seg_ewf.zip");
let mut a = crate::Archive::open(&data, Some("seg_ewf.zip"))
.unwrap()
.unwrap();
let order = ["img.E01", "img.E02", "img.E03"];
let oracle: Vec<Vec<u8>> = order
.iter()
.map(|n| {
let idx = a.entries().iter().position(|e| e.name == *n).unwrap();
a.read(idx).unwrap()
})
.collect();
let ss = segment_sources(data, &Limits::default())
.unwrap()
.expect("a segment set");
assert_eq!(ss.kind, SegmentKind::Ewf);
assert_eq!(ss.sources.len(), 3);
for (src, want) in ss.sources.into_iter().zip(oracle) {
let mut got = Vec::new();
let mut src = src;
src.read_to_end(&mut got).unwrap();
assert_eq!(got, want, "ordered E0N segment bytes");
}
let data = load("seg_ewf.zip");
match reassemble(data, &Limits::default()).unwrap() {
Reassembled::Segments(s) => {
assert_eq!(s.kind, SegmentKind::Ewf);
assert_eq!(s.sources.len(), 3);
}
other => panic!("expected Segments seam, got {other:?}"),
}
}
#[test]
fn single_member_is_not_segmented() {
let data = load("stored_one.zip");
assert!(matches!(
reassemble(data.clone(), &Limits::default()).unwrap(),
Reassembled::NotSegmented
));
assert!(segment_sources(data, &Limits::default()).unwrap().is_none());
}
}