diff --git a/src/recovery.rs b/src/recovery.rs index fb08b3c..702d642 100644 --- a/src/recovery.rs +++ b/src/recovery.rs @@ -1,11 +1,36 @@ //! Recovery block reading from PAR2 volume files. //! -//! Scans all `.par2` files in a directory and extracts RecoverySlice packets, -//! which contain the actual Reed-Solomon recovery data needed for repair. +//! [`RecoveryReader`] streams RecoverySlice packets out of any [`Read`] source +//! (a file, a network body, `&[u8]`), one volume at a time, keeping at most a +//! budgeted number of rows and preferring the lowest exponents. Packets it will +//! not keep are skipped without being buffered, so a volume much larger than +//! the budget streams past in constant memory. +//! +//! [`load_recovery_blocks`] is the directory-scanning convenience wrapper. +//! +//! ``` +//! use rust_par2::recovery::{RecoveryBudget, RecoveryReader}; +//! +//! # fn demo(set_id: [u8; 16], slice_size: u64, volumes: Vec>, missing: usize) { +//! let mut reader = RecoveryReader::new(set_id, slice_size) +//! .with_budget(RecoveryBudget::rows(missing)); +//! for volume in &volumes { +//! let scan = reader.read_volume(&volume[..]); +//! if !scan.read_to_end() { +//! // This volume may carry more rows than we saw: don't conclude the +//! // set is unrepairable from it. +//! } +//! } +//! let rows = reader.into_blocks(); // feed to `Repair::new_owning` +//! # let _ = rows; } +//! ``` -use std::io::{Read, Seek, SeekFrom}; +use std::collections::BTreeSet; +use std::fmt; +use std::io::{self, Read}; use std::path::Path; +use md5::{Digest, Md5}; use tracing::{debug, info, warn}; use crate::packets::{HEADER_SIZE, MAGIC}; @@ -13,6 +38,9 @@ use crate::packets::{HEADER_SIZE, MAGIC}; /// RecoverySlice packet type identifier. const TYPE_RECOVERY: &[u8; 16] = b"PAR 2.0\x00RecvSlic"; +/// Default cap on bytes read from one volume by a [`RecoveryReader`]. +pub const DEFAULT_MAX_VOLUME_BYTES: u64 = 512 * 1024 * 1024; + /// A recovery block extracted from a RecoverySlice packet. #[derive(Debug, Clone)] pub struct RecoveryBlock { @@ -22,13 +50,403 @@ pub struct RecoveryBlock { pub data: Vec, } +/// How many recovery rows a [`RecoveryReader`] may hold at once. +/// +/// Both limits apply when both are set; the tighter one wins. Rows over the +/// budget are dropped highest exponent first. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct RecoveryBudget { + /// Keep at most this many rows. + pub max_rows: Option, + /// Keep at most this many bytes of row data (`rows × slice_size`). + pub max_bytes: Option, +} + +impl RecoveryBudget { + /// No limit: keep every row seen. + pub const UNLIMITED: Self = Self { + max_rows: None, + max_bytes: None, + }; + + /// Keep at most `n` rows. + pub fn rows(n: usize) -> Self { + Self { + max_rows: Some(n), + max_bytes: None, + } + } + + /// Keep at most `bytes` of row data. + pub fn bytes(bytes: u64) -> Self { + Self { + max_rows: None, + max_bytes: Some(bytes), + } + } + + /// Row limit for a set with this slice size. + pub fn row_limit(&self, slice_size: u64) -> usize { + let by_bytes = match self.max_bytes { + Some(b) if slice_size > 0 => usize::try_from(b / slice_size).unwrap_or(usize::MAX), + _ => usize::MAX, + }; + self.max_rows.unwrap_or(usize::MAX).min(by_bytes) + } +} + +/// How reading one volume ended. +#[derive(Debug)] +pub enum VolumeEnd { + /// Read to the end of the stream: the volume carries no more rows than + /// were seen. + Complete, + /// The stream ended inside a packet. Rows after the cut are unknown. + Truncated, + /// Stopped at the per-volume byte cap with data left unread. + CapReached, + /// A read failed. Rows after the failure are unknown. + Io(io::Error), +} + +impl fmt::Display for VolumeEnd { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Complete => write!(f, "read to end"), + Self::Truncated => write!(f, "truncated inside a packet"), + Self::CapReached => write!(f, "stopped at the per-volume byte cap"), + Self::Io(e) => write!(f, "read error: {e}"), + } + } +} + +/// What [`RecoveryReader::read_volume`] found in one volume. +#[derive(Debug)] +pub struct VolumeScan { + /// Intact RecoverySlice packets for this set found in the volume, + /// duplicates of exponents already held included. Packets skipped for + /// budget are counted without their hash being checked. + pub seen: usize, + /// Rows from this volume that were added to the reader. A later volume can + /// evict some of them again. + pub kept: usize, + /// RecoverySlice packets for this set dropped as damaged: a body too + /// short for the slice size, or a packet hash mismatch. + pub corrupt: usize, + /// Bytes consumed from the stream. + pub bytes_read: u64, + /// Whether the volume was read to its end. + pub end: VolumeEnd, +} + +impl VolumeScan { + /// Whether the whole volume was read. When `false`, the volume may carry + /// rows that were never seen, so a shortfall must not be taken to mean + /// the set is unrepairable. + pub fn read_to_end(&self) -> bool { + matches!(self.end, VolumeEnd::Complete) + } +} + +/// Streaming RecoverySlice reader for one recovery set, accumulating rows +/// across volumes under a [`RecoveryBudget`]. +/// +/// Rows are held in the order their packets were read. When the budget is +/// full, a new row is kept only if its exponent is lower than the highest one +/// held, which is then dropped. Peak memory is one slice over the budget. +#[derive(Debug, Clone)] +pub struct RecoveryReader { + set_id: [u8; 16], + slice_size: u64, + limit: usize, + max_volume_bytes: u64, + blocks: Vec, + exponents: BTreeSet, + seen: usize, +} + +impl RecoveryReader { + /// Reader for recovery set `set_id` with slices of `slice_size` bytes, an + /// unlimited budget and the [`DEFAULT_MAX_VOLUME_BYTES`] cap. + pub fn new(set_id: [u8; 16], slice_size: u64) -> Self { + Self { + set_id, + slice_size, + limit: usize::MAX, + max_volume_bytes: DEFAULT_MAX_VOLUME_BYTES, + blocks: Vec::new(), + exponents: BTreeSet::new(), + seen: 0, + } + } + + /// Set the budget (builder form of [`set_budget`](Self::set_budget)). + pub fn with_budget(mut self, budget: RecoveryBudget) -> Self { + self.set_budget(budget); + self + } + + /// Cap the bytes read from one volume. A volume longer than this reports + /// [`VolumeEnd::CapReached`]. + pub fn with_max_volume_bytes(mut self, max: u64) -> Self { + self.max_volume_bytes = max; + self + } + + /// Change the budget. Shrinking it drops the highest-exponent rows now. + pub fn set_budget(&mut self, budget: RecoveryBudget) { + self.limit = budget.row_limit(self.slice_size); + while self.blocks.len() > self.limit { + self.evict_highest(); + } + } + + /// Current row limit. + pub fn row_limit(&self) -> usize { + self.limit + } + + /// Intact RecoverySlice packets for this set seen across all volumes. + pub fn seen(&self) -> usize { + self.seen + } + + /// Rows held now. + pub fn kept(&self) -> usize { + self.blocks.len() + } + + /// Rows held, in packet order. + pub fn blocks(&self) -> &[RecoveryBlock] { + &self.blocks + } + + /// Take the held rows, in packet order. + pub fn into_blocks(self) -> Vec { + self.blocks + } + + /// Read RecoverySlice packets from one volume. + /// + /// Packets for other sets, other packet types and rows outside the budget + /// are read past without being buffered. Never fails: how the read ended + /// is in [`VolumeScan::end`], and rows found before a failure are kept. + pub fn read_volume(&mut self, reader: R) -> VolumeScan { + let mut src = Source::new(reader, self.max_volume_bytes); + let mut scan = VolumeScan { + seen: 0, + kept: 0, + corrupt: 0, + bytes_read: 0, + end: VolumeEnd::Complete, + }; + scan.end = match self.scan(&mut src, &mut scan) { + Ok(()) => VolumeEnd::Complete, + Err(Stop::Eof) => VolumeEnd::Truncated, + Err(Stop::Cap) => VolumeEnd::CapReached, + Err(Stop::Io(e)) => VolumeEnd::Io(e), + }; + scan.bytes_read = src.pos; + self.seen += scan.seen; + scan + } + + fn scan(&mut self, src: &mut Source, scan: &mut VolumeScan) -> Result<(), Stop> { + let mut hdr = [0u8; HEADER_SIZE]; + let mut have = 0; + loop { + let got = src.read_full(&mut hdr[have..])?; + have += got; + if have < HEADER_SIZE { + // End of stream between packets. A partial packet header is + // a cut; anything else is trailing garbage. + return if have > 0 && starts_like_magic(&hdr[..have]) { + Err(Stop::Eof) + } else { + Ok(()) + }; + } + + if &hdr[0..8] != MAGIC { + // Resync: keep the window from the next possible magic start. + let next = find_magic_start(&hdr[1..]).map_or(HEADER_SIZE, |i| i + 1); + hdr.copy_within(next.., 0); + have = HEADER_SIZE - next; + continue; + } + let packet_len = u64::from_le_bytes(hdr[8..16].try_into().unwrap()); + if packet_len < HEADER_SIZE as u64 || packet_len % 4 != 0 { + let next = find_magic_start(&hdr[1..]).map_or(HEADER_SIZE, |i| i + 1); + hdr.copy_within(next.., 0); + have = HEADER_SIZE - next; + continue; + } + have = 0; + let body_len = packet_len - HEADER_SIZE as u64; + + if hdr[32..48] != self.set_id || &hdr[48..64] != TYPE_RECOVERY { + src.skip(body_len)?; + continue; + } + if body_len < 4 + self.slice_size { + scan.corrupt += 1; + src.skip(body_len)?; + continue; + } + + let mut exp_buf = [0u8; 4]; + src.read_exact(&mut exp_buf)?; + let exponent = u32::from_le_bytes(exp_buf); + let rest = body_len - 4; + + if !self.wants(exponent) { + scan.seen += 1; + src.skip(rest)?; + continue; + } + + let mut hasher = Md5::new(); + hasher.update(&hdr[32..64]); + hasher.update(exp_buf); + let mut data = vec![0u8; self.slice_size as usize]; + src.read_exact(&mut data)?; + hasher.update(&data); + src.skip_hashing(rest - self.slice_size, &mut hasher)?; + let digest: [u8; 16] = hasher.finalize().into(); + if digest[..] != hdr[16..32] { + scan.corrupt += 1; + continue; + } + + scan.seen += 1; + scan.kept += 1; + if self.blocks.len() >= self.limit { + self.evict_highest(); + } + self.exponents.insert(exponent); + self.blocks.push(RecoveryBlock { exponent, data }); + } + } + + /// Whether a new row with this exponent would be kept. + fn wants(&self, exponent: u32) -> bool { + if self.limit == 0 || self.exponents.contains(&exponent) { + return false; + } + self.blocks.len() < self.limit || self.exponents.last().is_some_and(|&max| exponent < max) + } + + fn evict_highest(&mut self) { + let Some(max) = self.exponents.pop_last() else { + return; + }; + if let Some(i) = self.blocks.iter().position(|b| b.exponent == max) { + self.blocks.remove(i); + } + } +} + +/// Why a volume scan stopped before a clean end. +enum Stop { + Eof, + Cap, + Io(io::Error), +} + +/// A reader that counts bytes and refuses to go past a cap. +struct Source { + inner: R, + pos: u64, + cap: u64, +} + +impl Source { + fn new(inner: R, cap: u64) -> Self { + Self { inner, pos: 0, cap } + } + + /// Read until `buf` is full or the stream ends; returns bytes read. + fn read_full(&mut self, buf: &mut [u8]) -> Result { + let mut n = 0; + while n < buf.len() { + let room = self.cap - self.pos; + if room == 0 { + return if self.at_eof()? { + Ok(n) + } else { + Err(Stop::Cap) + }; + } + let want = (buf.len() - n).min(usize::try_from(room).unwrap_or(usize::MAX)); + match self.inner.read(&mut buf[n..n + want]) { + Ok(0) => break, + Ok(k) => { + n += k; + self.pos += k as u64; + } + Err(e) if e.kind() == io::ErrorKind::Interrupted => {} + Err(e) => return Err(Stop::Io(e)), + } + } + Ok(n) + } + + /// At the cap: probe whether the stream has anything left. + fn at_eof(&mut self) -> Result { + let mut probe = [0u8; 1]; + loop { + match self.inner.read(&mut probe) { + Ok(0) => return Ok(true), + Ok(_) => return Ok(false), + Err(e) if e.kind() == io::ErrorKind::Interrupted => {} + Err(e) => return Err(Stop::Io(e)), + } + } + } + + fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), Stop> { + if self.read_full(buf)? < buf.len() { + return Err(Stop::Eof); + } + Ok(()) + } + + fn skip(&mut self, n: u64) -> Result<(), Stop> { + self.skip_with(n, |_| {}) + } + + fn skip_hashing(&mut self, n: u64, hasher: &mut Md5) -> Result<(), Stop> { + self.skip_with(n, |chunk| hasher.update(chunk)) + } + + fn skip_with(&mut self, mut n: u64, mut f: impl FnMut(&[u8])) -> Result<(), Stop> { + let mut buf = [0u8; 8192]; + while n > 0 { + let want = n.min(buf.len() as u64) as usize; + self.read_exact(&mut buf[..want])?; + f(&buf[..want]); + n -= want as u64; + } + Ok(()) + } +} + +/// Offset of the first position in `buf` where the packet magic could start +/// (a full match, or a prefix of it running to the end of `buf`). +fn find_magic_start(buf: &[u8]) -> Option { + (0..buf.len()).find(|&i| starts_like_magic(&buf[i..])) +} + +fn starts_like_magic(buf: &[u8]) -> bool { + let n = buf.len().min(MAGIC.len()); + buf[..n] == MAGIC[..n] +} + /// Load all recovery blocks from PAR2 files in a directory. /// /// Scans all `.par2` files (index and volumes) for RecoverySlice packets /// matching the given recovery set ID. Returns blocks sorted by exponent. pub fn load_recovery_blocks(dir: &Path, set_id: &[u8; 16], slice_size: u64) -> Vec { - let mut blocks = Vec::new(); - // Find all .par2 files in the directory let mut par2_files: Vec<_> = match std::fs::read_dir(dir) { Ok(entries) => entries @@ -41,31 +459,31 @@ pub fn load_recovery_blocks(dir: &Path, set_id: &[u8; 16], slice_size: u64) -> V .collect(), Err(e) => { warn!(error = %e, "Failed to read directory for recovery files"); - return blocks; + return Vec::new(); } }; par2_files.sort(); + let mut reader = RecoveryReader::new(*set_id, slice_size).with_max_volume_bytes(u64::MAX); for par2_path in &par2_files { - match read_recovery_packets(par2_path, set_id, slice_size) { - Ok(mut file_blocks) => { - debug!( - file = %par2_path.display(), - count = file_blocks.len(), - "Loaded recovery blocks" - ); - blocks.append(&mut file_blocks); - } + let file = match std::fs::File::open(par2_path) { + Ok(f) => f, Err(e) => { - debug!( - file = %par2_path.display(), - error = %e, - "Skipping file (no recovery blocks or read error)" - ); + debug!(file = %par2_path.display(), error = %e, "Skipping unreadable file"); + continue; } - } + }; + let scan = reader.read_volume(io::BufReader::with_capacity(64 * 1024, file)); + debug!( + file = %par2_path.display(), + count = scan.kept, + corrupt = scan.corrupt, + end = %scan.end, + "Loaded recovery blocks" + ); } + let mut blocks = reader.into_blocks(); blocks.sort_by_key(|b| b.exponent); info!(total_blocks = blocks.len(), "Recovery blocks loaded"); @@ -73,83 +491,185 @@ pub fn load_recovery_blocks(dir: &Path, set_id: &[u8; 16], slice_size: u64) -> V blocks } -/// Read RecoverySlice packets from a single PAR2 file. -fn read_recovery_packets( - path: &Path, - set_id: &[u8; 16], - slice_size: u64, -) -> std::io::Result> { - let mut file = std::fs::File::open(path)?; - let file_size = file.metadata()?.len(); - let mut blocks = Vec::new(); - let mut pos: u64 = 0; +#[cfg(test)] +mod tests { + use super::*; - let mut header_buf = [0u8; HEADER_SIZE]; + const SET: [u8; 16] = [7u8; 16]; - while pos + HEADER_SIZE as u64 <= file_size { - file.seek(SeekFrom::Start(pos))?; + fn packet(set_id: &[u8; 16], ptype: &[u8; 16], body: &[u8]) -> Vec { + let mut hashed = Vec::new(); + hashed.extend_from_slice(set_id); + hashed.extend_from_slice(ptype); + hashed.extend_from_slice(body); + let md5: [u8; 16] = Md5::digest(&hashed).into(); + let mut p = Vec::new(); + p.extend_from_slice(MAGIC); + p.extend_from_slice(&((HEADER_SIZE + body.len()) as u64).to_le_bytes()); + p.extend_from_slice(&md5); + p.extend_from_slice(&hashed); + p + } - // Read packet header - if file.read_exact(&mut header_buf).is_err() { - break; - } + fn recv(exponent: u32, slice_size: usize) -> Vec { + let mut body = exponent.to_le_bytes().to_vec(); + body.extend((0..slice_size).map(|i| (i as u32 ^ exponent) as u8)); + packet(&SET, TYPE_RECOVERY, &body) + } - // Check magic - if &header_buf[0..8] != MAGIC { - pos += 4; // Scan forward - continue; - } + fn volume(exponents: &[u32], slice_size: usize) -> Vec { + exponents + .iter() + .flat_map(|&e| recv(e, slice_size)) + .collect() + } - // Parse packet length - let packet_len = u64::from_le_bytes(header_buf[8..16].try_into().unwrap()); - if packet_len < HEADER_SIZE as u64 || packet_len % 4 != 0 { - pos += 4; - continue; - } + fn exps(r: &RecoveryReader) -> Vec { + r.blocks().iter().map(|b| b.exponent).collect() + } - // Check recovery set ID - let pkt_set_id: [u8; 16] = header_buf[32..48].try_into().unwrap(); - if &pkt_set_id != set_id { - pos += packet_len; - continue; - } + #[test] + fn test_load_empty_dir() { + let dir = tempfile::tempdir().unwrap(); + let set_id = [0u8; 16]; + let blocks = load_recovery_blocks(dir.path(), &set_id, 768000); + assert!(blocks.is_empty()); + } - // Check packet type - let pkt_type = &header_buf[48..64]; - if pkt_type == TYPE_RECOVERY { - // RecoverySlice packet body layout: - // Offset 0..4 (within body): exponent (u32 LE) - // Offset 4..4+slice_size: recovery data - let body_len = packet_len - HEADER_SIZE as u64; - let expected_body = 4 + slice_size; + #[test] + fn packet_order_and_data() { + let vol = volume(&[5, 2, 9, 0], 16); + let mut r = RecoveryReader::new(SET, 16); + let scan = r.read_volume(&vol[..]); + assert!(scan.read_to_end()); + assert_eq!((scan.seen, scan.kept, scan.corrupt), (4, 4, 0)); + assert_eq!(scan.bytes_read, vol.len() as u64); + assert_eq!(exps(&r), [5, 2, 9, 0]); + assert_eq!(r.blocks()[1].data, recv(2, 16)[HEADER_SIZE + 4..]); + } - if body_len >= expected_body { - let mut body = vec![0u8; expected_body as usize]; - file.seek(SeekFrom::Start(pos + HEADER_SIZE as u64))?; - file.read_exact(&mut body)?; + #[test] + fn budget_keeps_lowest_exponents_in_packet_order() { + let vol = volume(&[5, 2, 9, 0, 7, 1], 16); + let mut r = RecoveryReader::new(SET, 16).with_budget(RecoveryBudget::rows(3)); + let scan = r.read_volume(&vol[..]); + assert!(scan.read_to_end()); + assert_eq!(scan.seen, 6); + assert_eq!(r.seen(), 6); + assert_eq!(r.kept(), 3); + assert_eq!(exps(&r), [2, 0, 1]); + } - let exponent = u32::from_le_bytes(body[0..4].try_into().unwrap()); - let data = body[4..4 + slice_size as usize].to_vec(); + #[test] + fn byte_budget_and_shrinking() { + let vol = volume(&[3, 1, 2, 0], 16); + let mut r = RecoveryReader::new(SET, 16).with_budget(RecoveryBudget::bytes(16 * 3 + 5)); + r.read_volume(&vol[..]); + assert_eq!(exps(&r), [1, 2, 0]); + r.set_budget(RecoveryBudget::rows(1)); + assert_eq!(exps(&r), [0]); + } - blocks.push(RecoveryBlock { exponent, data }); - } + #[test] + fn zero_budget_counts_without_keeping() { + let vol = volume(&[1, 2, 3], 16); + let mut r = RecoveryReader::new(SET, 16).with_budget(RecoveryBudget::rows(0)); + let scan = r.read_volume(&vol[..]); + assert!(scan.read_to_end()); + assert_eq!((scan.seen, scan.kept), (3, 0)); + } + + #[test] + fn accumulates_across_volumes_and_skips_duplicates() { + let mut r = RecoveryReader::new(SET, 16).with_budget(RecoveryBudget::rows(3)); + r.read_volume(&volume(&[4, 5, 6], 16)[..]); + let scan = r.read_volume(&volume(&[5, 1], 16)[..]); + assert_eq!((scan.seen, scan.kept), (2, 1)); + assert_eq!(exps(&r), [4, 5, 1]); + assert_eq!(r.seen(), 5); + } + + #[test] + fn truncated_volume_is_not_complete() { + let vol = volume(&[0, 1, 2], 16); + for cut in [vol.len() - 1, vol.len() - 10, 2 * vol.len() / 3 + 3] { + let mut r = RecoveryReader::new(SET, 16); + let scan = r.read_volume(&vol[..cut]); + assert!(!scan.read_to_end(), "cut at {cut}"); + assert!(matches!(scan.end, VolumeEnd::Truncated)); + assert_eq!(exps(&r), [0, 1]); } + // Cut inside a skipped (over-budget) packet is still a truncation. + let mut r = RecoveryReader::new(SET, 16).with_budget(RecoveryBudget::rows(1)); + let scan = r.read_volume(&vol[..vol.len() - 3]); + assert!(matches!(scan.end, VolumeEnd::Truncated)); + // Cut inside the next packet's header. + let mut r = RecoveryReader::new(SET, 16); + let one = recv(0, 16); + let scan = r.read_volume(&vol[..one.len() + 5]); + assert!(matches!(scan.end, VolumeEnd::Truncated)); + } - pos += packet_len; + #[test] + fn cap_reached_is_not_complete() { + let vol = volume(&[0, 1, 2], 16); + let one = recv(0, 16).len() as u64; + let mut r = RecoveryReader::new(SET, 16).with_max_volume_bytes(one * 2); + let scan = r.read_volume(&vol[..]); + assert!(matches!(scan.end, VolumeEnd::CapReached)); + assert_eq!(scan.bytes_read, one * 2); + assert_eq!(exps(&r), [0, 1]); + + // A volume exactly at the cap is complete. + let mut r = RecoveryReader::new(SET, 16).with_max_volume_bytes(vol.len() as u64); + assert!(r.read_volume(&vol[..]).read_to_end()); } - Ok(blocks) -} + #[test] + fn io_error_is_reported() { + struct Failing<'a>(&'a [u8]); + impl Read for Failing<'_> { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + if self.0.is_empty() { + return Err(io::Error::other("boom")); + } + let n = self.0.read(buf)?; + Ok(n) + } + } + let vol = volume(&[0, 1], 16); + let mut r = RecoveryReader::new(SET, 16); + let scan = r.read_volume(Failing(&vol[..vol.len() - 4])); + assert!(matches!(scan.end, VolumeEnd::Io(_))); + assert_eq!(exps(&r), [0]); + } -#[cfg(test)] -mod tests { - use super::*; + #[test] + fn skips_other_sets_types_garbage_and_corrupt() { + let mut vol = b"junk!".to_vec(); + vol.extend(packet(&[9u8; 16], TYPE_RECOVERY, &[0u8; 20])); + vol.extend(packet(&SET, b"PAR 2.0\x00Creator\x00", b"test")); + vol.extend(recv(3, 16)); + let mut bad = recv(4, 16); + let last = bad.len() - 1; + bad[last] ^= 1; + vol.extend(bad); + vol.extend(packet(&SET, TYPE_RECOVERY, &[0u8; 8])); // too short + vol.extend(recv(1, 16)); + vol.extend(b"tail"); + + let mut r = RecoveryReader::new(SET, 16); + let scan = r.read_volume(&vol[..]); + assert!(scan.read_to_end(), "{}", scan.end); + assert_eq!((scan.seen, scan.kept, scan.corrupt), (2, 2, 2)); + assert_eq!(exps(&r), [3, 1]); + } #[test] - fn test_load_empty_dir() { - let dir = tempfile::tempdir().unwrap(); - let set_id = [0u8; 16]; - let blocks = load_recovery_blocks(dir.path(), &set_id, 768000); - assert!(blocks.is_empty()); + fn empty_input_is_complete() { + let mut r = RecoveryReader::new(SET, 16); + let scan = r.read_volume(&[][..]); + assert!(scan.read_to_end()); + assert_eq!(scan.seen, 0); } } diff --git a/tests/recovery_reader.rs b/tests/recovery_reader.rs new file mode 100644 index 0000000..7b77458 --- /dev/null +++ b/tests/recovery_reader.rs @@ -0,0 +1,170 @@ +//! Streaming recovery reader (`RecoveryReader`) tests against the par2cmdline +//! `multi` fixture (12 input slices of 4096 bytes, 8 recovery slices in +//! `multi.vol0+8.par2`). + +use std::path::Path; + +use rust_par2::recovery::{RecoveryBudget, RecoveryReader, VolumeEnd, load_recovery_blocks}; +use rust_par2::{BlockMap, Repair}; + +const SLICE: usize = 4096; + +fn dir() -> &'static Path { + Path::new("tests/fixtures/multi") +} + +fn set() -> rust_par2::Par2FileSet { + rust_par2::parse(&dir().join("multi.par2")).unwrap() +} + +fn volume() -> Vec { + std::fs::read(dir().join("multi.vol0+8.par2")).unwrap() +} + +/// Split a volume into whole packets by their length fields. +fn packets(vol: &[u8]) -> Vec<&[u8]> { + let mut out = Vec::new(); + let mut pos = 0; + while pos + 64 <= vol.len() { + let len = u64::from_le_bytes(vol[pos + 8..pos + 16].try_into().unwrap()) as usize; + out.push(&vol[pos..pos + len]); + pos += len; + } + assert_eq!(pos, vol.len()); + out +} + +fn exponent_of(packet: &[u8]) -> Option { + (&packet[48..64] == b"PAR 2.0\x00RecvSlic") + .then(|| u32::from_le_bytes(packet[64..68].try_into().unwrap())) +} + +fn input_slices(fs: &rust_par2::Par2FileSet) -> Vec> { + let map = BlockMap::new(fs); + (0..map.total_blocks as usize) + .map(|g| { + let loc = map.global_to_file(g).unwrap(); + let data = std::fs::read(dir().join(&loc.file.filename)).unwrap(); + data[loc.offset as usize..loc.offset as usize + loc.len].to_vec() + }) + .collect() +} + +fn padded(data: &[u8]) -> Vec { + let mut v = data.to_vec(); + v.resize(SLICE, 0); + v +} + +/// Solve `missing` from `rows` and check against the real slices. +fn assert_repairs(rows: Vec, missing: &[usize]) { + let fs = set(); + let slices = input_slices(&fs); + let mut r = Repair::new_owning(rows, SLICE, slices.len()).unwrap(); + for (i, s) in slices.iter().enumerate() { + if !missing.contains(&i) { + r.feed(i, s).unwrap(); + } + } + let solved = r.solve(missing).unwrap(); + for &m in missing { + assert_eq!(solved[&m], padded(&slices[m]), "slice {m}"); + } +} + +#[test] +fn reader_matches_directory_loader() { + let fs = set(); + let from_dir = load_recovery_blocks(dir(), &fs.recovery_set_id, fs.slice_size); + + let mut reader = RecoveryReader::new(fs.recovery_set_id, fs.slice_size); + let index = std::fs::read(dir().join("multi.par2")).unwrap(); + assert!(reader.read_volume(&index[..]).read_to_end()); + let scan = reader.read_volume(std::fs::File::open(dir().join("multi.vol0+8.par2")).unwrap()); + assert!(scan.read_to_end()); + assert_eq!((scan.seen, scan.kept, scan.corrupt), (8, 8, 0)); + assert_eq!(scan.bytes_read, volume().len() as u64); + + let mut streamed = reader.into_blocks(); + streamed.sort_by_key(|b| b.exponent); + assert_eq!(streamed.len(), from_dir.len()); + for (a, b) in streamed.iter().zip(&from_dir) { + assert_eq!(a.exponent, b.exponent); + assert_eq!(a.data, b.data); + } +} + +/// Recovery packets reordered so exponents arrive out of order; a budget of +/// three keeps the three lowest, in packet order, and they repair three +/// missing slices. +#[test] +fn out_of_order_exponents_under_budget() { + let fs = set(); + let vol = volume(); + let mut pkts = packets(&vol); + pkts.sort_by_key(|p| exponent_of(p).map(|e| (e * 5 + 3) % 8)); + let order: Vec = pkts.iter().filter_map(|p| exponent_of(p)).collect(); + assert_eq!(order, [1, 6, 3, 0, 5, 2, 7, 4]); + let shuffled: Vec = pkts.concat(); + + let mut reader = + RecoveryReader::new(fs.recovery_set_id, fs.slice_size).with_budget(RecoveryBudget::rows(3)); + let scan = reader.read_volume(&shuffled[..]); + assert!(scan.read_to_end()); + assert_eq!(scan.seen, 8); + assert_eq!(reader.kept(), 3); + let kept: Vec = reader.blocks().iter().map(|b| b.exponent).collect(); + assert_eq!(kept, [1, 0, 2]); + + assert_repairs(reader.into_blocks(), &[0, 5, 11]); +} + +#[test] +fn byte_budget_bounds_rows() { + let fs = set(); + let mut reader = RecoveryReader::new(fs.recovery_set_id, fs.slice_size) + .with_budget(RecoveryBudget::bytes(2 * SLICE as u64 + 100)); + assert!(reader.read_volume(&volume()[..]).read_to_end()); + let kept: Vec = reader.blocks().iter().map(|b| b.exponent).collect(); + assert_eq!(kept, [0, 1]); +} + +/// A volume cut short reports `Truncated`, not a clean end, and still yields +/// the rows before the cut. +#[test] +fn truncated_volume() { + let fs = set(); + let vol = volume(); + let pkts = packets(&vol); + let recv_ends: Vec = pkts + .iter() + .scan(0, |end, p| { + *end += p.len(); + Some((*end, exponent_of(p).is_some())) + }) + .filter_map(|(end, is_recv)| is_recv.then_some(end)) + .collect(); + // Cut half way through the fourth recovery packet. + let cut = recv_ends[3] - SLICE / 2; + + let mut reader = RecoveryReader::new(fs.recovery_set_id, fs.slice_size); + let scan = reader.read_volume(&vol[..cut]); + assert!(!scan.read_to_end()); + assert!(matches!(scan.end, VolumeEnd::Truncated), "{}", scan.end); + assert_eq!(scan.seen, 3); + assert_eq!(scan.bytes_read, cut as u64); + + assert_repairs(reader.into_blocks(), &[2, 3, 9]); +} + +#[test] +fn per_volume_cap() { + let fs = set(); + let vol = volume(); + let mut reader = + RecoveryReader::new(fs.recovery_set_id, fs.slice_size).with_max_volume_bytes(10_000); + let scan = reader.read_volume(&vol[..]); + assert!(matches!(scan.end, VolumeEnd::CapReached), "{}", scan.end); + assert_eq!(scan.bytes_read, 10_000); + assert!(scan.seen < 8); +}