From bf39a451645d62b70b2ec3fb33a151c7e2a0355d Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Tue, 4 Aug 2026 23:40:41 -0700 Subject: [PATCH] feat(bgzf): add slice decompression and caller-driven buffer recycling Three additive entry points that `fgumi-sort`'s arena engine and the pipeline BGZF steps need, landed ahead of their consumers so the crate arrives complete rather than growing under each later port. `decompress_into_slice` is the fixed-slice analogue of `decompress_block_slice_into`: it decompresses a block straight into a caller-sized `&mut [u8]` instead of appending to a `Vec`, which lets the sort ingest path decompress directly into an arena slot rather than into a staging buffer it then copies out of. It keeps the integrity guarantee its siblings provide -- the payload must exactly fill `out` and match the footer CRC32 -- and takes the same deflate stored-block fast path for level-0 input. `uncompressed_size` is the accessor that makes that contract safe to satisfy. A caller has to size the slot before calling, and ISIZE is a `u32` sitting in the file, so sizing straight off the footer means allocating from unvalidated input -- up to 4 GiB from a corrupt block. This reads the footer off the same `&[u8]` (no owned `Vec`, so no copy of the block just to read four bytes) and bounds the claim to `MAX_UNCOMPRESSED_BLOCK_SIZE`. `decompress_into_slice` resolves the size through it too, so what a caller allocates and what the decompressor accepts cannot diverge. The stored fast path is now shared three ways. `copy_stored_and_verify`'s framing checks were inlined in its body, so the slice variant would have had to duplicate them; they move to `parse_stored_frame`. The BTYPE dispatch predicate becomes `is_stored_block`, and the inflate-then-verify tail becomes `deflate_into_slice_and_verify`, both shared with `decompress_and_verify`. Extracting only the framing check would have left the same drift exposure one level up. `InlineBgzfCompressor::recycle_buffer` lets a `take_blocks` consumer hand a drained block buffer back to the pool. Only `write_blocks_to` recycled before, so a consumer driving the compressor with `write_all` + `flush` + `take_blocks` left the pool permanently empty and allocated a fresh output `Vec` for every block. `write_blocks_to` now routes through the same method. The pool is bounded on both axes -- count and buffer capacity -- because bounding the count alone would let one oversized `Vec` handed in by a caller sit there for the compressor's lifetime. Rewriting `write_blocks_to` to call `recycle_buffer` (a `&mut self` method) means draining into a temporary rather than holding a borrow on `completed_blocks`. That makes the write-error path recoverable, so it is handled rather than left as-is: the failing block and its tail are restored to the queue instead of being dropped by the `Drain` guard. They are documented as safe to inspect, not to replay -- `write_all` can commit part of the failing block before erroring, and the whole block is re-queued, so writing the queue again would repeat those bytes inside a gzip member. --- crates/fgumi-bgzf/src/lib.rs | 9 +- crates/fgumi-bgzf/src/reader.rs | 571 +++++++++++++++++++++++++++++--- crates/fgumi-bgzf/src/writer.rs | 263 ++++++++++++++- 3 files changed, 779 insertions(+), 64 deletions(-) diff --git a/crates/fgumi-bgzf/src/lib.rs b/crates/fgumi-bgzf/src/lib.rs index 889699863..1decaaa6b 100644 --- a/crates/fgumi-bgzf/src/lib.rs +++ b/crates/fgumi-bgzf/src/lib.rs @@ -12,12 +12,19 @@ pub mod reader; pub mod writer; // Re-export commonly used types + pub use header::{ BGZF_HEADER_SIZE, HeaderRejection, MIN_BLOCK_SIZE, block_size, block_size_checked, is_bgzf_header, }; +/// The libdeflater decompressor every `decompress_*` entry point here takes. +/// +/// Re-exported so a consumer can name the type it must construct and reuse +/// without also declaring a direct `libdeflater` dependency — and so the +/// version it names is necessarily the one this crate decompresses with. +pub use libdeflater::Decompressor; pub use reader::{ BGZF_EOF, BGZF_FOOTER_SIZE, RawBgzfBlock, decompress_block, decompress_block_into, - decompress_block_slice_into, read_raw_blocks, + decompress_block_slice_into, decompress_into_slice, read_raw_blocks, uncompressed_size, }; pub use writer::{BGZF_MAX_BLOCK_SIZE, CompressedBlock, InlineBgzfCompressor}; diff --git a/crates/fgumi-bgzf/src/reader.rs b/crates/fgumi-bgzf/src/reader.rs index a9373d87f..4423565d1 100644 --- a/crates/fgumi-bgzf/src/reader.rs +++ b/crates/fgumi-bgzf/src/reader.rs @@ -47,6 +47,21 @@ pub use crate::header::BGZF_HEADER_SIZE; /// Size of the BGZF block footer (CRC32 + ISIZE). pub const BGZF_FOOTER_SIZE: usize = 8; +/// Largest uncompressed payload a single BGZF block can hold. +/// +/// A block's `BSIZE` is a `u16` holding the total size minus one, which caps a +/// block at 64 KiB; the SAM/BGZF spec and htslib's `BGZF_MAX_BLOCK_SIZE` both +/// use this bound, and the `bgzf` crate refuses to *write* a larger one +/// (`BlockSizeExceeded`). The footer's ISIZE is a `u32`, so a corrupt or +/// hostile block can claim up to 4 GiB; anything a caller sizes from that value +/// has to be bounded first. +/// +/// Note this is deliberately 64 KiB rather than [`crate::writer::BGZF_MAX_BLOCK_SIZE`] +/// (65280), which is the size *we* fill blocks to. Other writers legitimately +/// emit up to the spec limit, so validating against our own chunk size would +/// reject valid input. +pub const MAX_UNCOMPRESSED_BLOCK_SIZE: usize = 64 * 1024; + /// BGZF EOF marker block (empty block signaling end of file). pub const BGZF_EOF: [u8; 28] = [ 0x1f, 0x8b, 0x08, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0x06, 0x00, 0x42, 0x43, 0x02, 0x00, @@ -268,6 +283,59 @@ pub fn read_raw_blocks( // Decompression Helpers // ============================================================================ +/// The uncompressed size a BGZF block's footer claims, validated against the +/// per-block maximum. +/// +/// This is the size a caller must give [`decompress_into_slice`]'s `out`, and +/// it is the number a caller allocates from, so it is checked rather than +/// returned raw: ISIZE is a `u32` sitting in the file, and a corrupt or hostile +/// footer claiming 4 GiB would otherwise become a 4 GiB allocation in the +/// caller before this crate ever saw the block. +/// +/// Takes the block as a `&[u8]` so reading the footer costs nothing — +/// [`RawBgzfBlock::uncompressed_size`] answers the same question but needs an +/// owned `Vec`, which would mean copying the block to size a slot for it. +/// +/// # Errors +/// +/// Returns `io::ErrorKind::InvalidData` if `block` is too short to hold a BGZF +/// header + footer, or if the claimed size exceeds +/// [`MAX_UNCOMPRESSED_BLOCK_SIZE`]. +pub fn uncompressed_size(block: &[u8]) -> io::Result { + if block.len() < BGZF_HEADER_SIZE + BGZF_FOOTER_SIZE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "BGZF block too short to contain header + footer", + )); + } + let claimed = uncompressed_size_from_slice(block); + if claimed > MAX_UNCOMPRESSED_BLOCK_SIZE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "BGZF block claims an uncompressed size of {claimed} bytes, above the \ + {MAX_UNCOMPRESSED_BLOCK_SIZE}-byte maximum for a single block" + ), + )); + } + Ok(claimed) +} + +/// Whether a BGZF block's deflate payload is a **stored** (uncompressed) frame, +/// i.e. `BTYPE = 00`. +/// +/// `BTYPE` lives in bits 1-2 of the deflate stream's first byte, so masking with +/// `0b110` isolates it. An empty payload is not a stored frame — there is no +/// first byte to read, and the framing parser needs five. +/// +/// Shared by every decompress entry point so the dispatch predicate cannot +/// drift between them; [`parse_stored_frame`] validates the rest of the framing +/// once this returns `true`. +#[must_use] +fn is_stored_block(compressed: &[u8]) -> bool { + !compressed.is_empty() && compressed[0] & 0b110 == 0 +} + /// Verify that decompressed data matches the expected size and CRC32 checksum. /// /// # Arguments @@ -434,6 +502,126 @@ pub fn decompress_block_slice_into( ) } +/// Decompress a full BGZF block's DEFLATE payload into a caller-provided, +/// pre-sized slice. `out.len()` must equal the block's ISIZE (uncompressed +/// size), which callers should take from [`uncompressed_size`] — it reads the +/// footer straight off the same `&[u8]` and bounds the claim, so the slot a +/// caller allocates can never come from an unvalidated `u32`. +/// +/// This is the fixed-slice analogue of [`decompress_block_slice_into`] (which +/// appends to a `Vec`); it lets a caller decompress straight into an arena +/// slot rather than into a buffer it then has to copy out of. The two names are +/// close and their `slice` refers to opposite operands: in +/// [`decompress_block_slice_into`] it is the *input* block, given as a slice +/// rather than a [`RawBgzfBlock`]; here it is the *output*. +/// +/// Like the sibling decompressors ([`decompress_block_into`] / +/// [`decompress_block_slice_into`]), the decompressed payload is verified +/// against the BGZF footer's ISIZE (it must exactly fill `out`) and CRC32, so +/// a short fill or a silently-corrupt-but-decodable block is caught here +/// rather than fed into the arena. +/// +/// # Returns +/// +/// The number of bytes written, which on success is always `out.len()`: the +/// decompressing paths are held to it by the exact-fill check, and a block with +/// an ISIZE of zero (the BGZF EOF marker) returns `0` against the zero-length +/// slot its ISIZE requires. It is returned so the call reads like the +/// `Read`-style APIs it sits beside, not because a short result is reachable. +/// +/// # Errors +/// +/// Returns an `io::Error` if `block` is shorter than a BGZF header + footer, if +/// the footer's ISIZE exceeds [`MAX_UNCOMPRESSED_BLOCK_SIZE`], if `out` is not +/// sized to that ISIZE, if the DEFLATE stream is invalid, if it does not +/// exactly fill `out`, or if the CRC32 does not match the footer. +/// +/// Note this validates the block's *framing*, not its full header: a payload +/// whose header [`crate::header::validate`] would reject can still reach the +/// decompressor and fail there. Callers reading from a stream get the header +/// check from [`read_raw_blocks`]; this entry point is for callers that already +/// hold a framed block. +/// +/// On error `out` is left clobbered — unlike the `Vec` siblings, which roll +/// their output back. There is nothing to roll back to here: the buffer belongs +/// to the caller, who must treat its contents as undefined unless this returns +/// `Ok`. +pub fn decompress_into_slice( + block: &[u8], + decompressor: &mut Decompressor, + out: &mut [u8], +) -> io::Result { + // Same accessor the caller sizes `out` with, so the two cannot disagree + // about either the value or the bound. It carries the length and ISIZE + // checks, which is why neither is repeated here. + let uncompressed_size = uncompressed_size(block)?; + // Check the caller's slot against the footer up front rather than letting a + // mis-sized `out` surface downstream. Both paths below compare against + // `out.len()` on the assumption it *is* the ISIZE, so without this a wrongly + // sized slot is reported as a fault in the block: the stored path would say + // `ISIZE mismatch: footer=`, naming a value the footer does not + // contain, and send a reader looking for file corruption instead of at the + // arena that sized the slot. + if out.len() != uncompressed_size { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "BGZF output slice is {} bytes, but the block's ISIZE is {uncompressed_size}", + out.len() + ), + )); + } + // A block that carries nothing — the BGZF EOF marker is the one every + // reader meets — is done, and both `Vec` siblings short-circuit here too. + // libdeflater does return `Ok(0)` for the EOF marker's payload against a + // zero-length slice, so this is not repairing a failure; it is making the + // answer ours rather than resting on an undocumented edge of the C library. + // + // Placed *after* the size check on purpose: ahead of it, a caller passing a + // wrongly-sized slot for a zero-ISIZE block would get a silent `Ok(0)` + // instead of being told the slot is wrong. + if uncompressed_size == 0 { + return Ok(0); + } + let compressed = compressed_data_from_slice(block); + // Stored-block fast path — level-0 blocks (`samtools view -u`, htsjdk's + // level-0 writer, [`InlineBgzfCompressor::new(0)`]) skip the libdeflater + // round-trip and get the stored-framing-specific LEN/ISIZE diagnostics. + if is_stored_block(compressed) { + return copy_stored_and_verify_slice(compressed, out, crc32_from_slice(block), block.len()); + } + deflate_into_slice_and_verify( + compressed, + crc32_from_slice(block), + block.len(), + decompressor, + out, + ) +} + +/// Inflate `compressed` into the whole of `out` and verify the result against +/// the BGZF footer, returning the number of bytes written. +/// +/// `out.len()` is taken as the expected uncompressed size, so +/// [`verify_decompression`] checks both the exact fill +/// (`bytes_written == out.len()`) and the CRC32. Shared by +/// [`decompress_into_slice`] and [`decompress_and_verify`]'s non-stored branch +/// so the inflate-then-verify invariant lives in one place — the same reason +/// [`parse_stored_frame`] exists for the stored branch. +fn deflate_into_slice_and_verify( + compressed: &[u8], + expected_crc: u32, + block_len: usize, + decompressor: &mut Decompressor, + out: &mut [u8], +) -> io::Result { + let bytes_written = decompressor.deflate_decompress(compressed, out).map_err(|e| { + io::Error::new(io::ErrorKind::InvalidData, format!("BGZF decompression failed: {e:?}")) + })?; + verify_decompression(&out[..bytes_written], out.len(), expected_crc, block_len)?; + Ok(bytes_written) +} + /// Decompress (or copy, for deflate stored blocks) BGZF block data into the /// output buffer and verify the result. /// @@ -477,11 +665,10 @@ fn decompress_and_verify( // bytes 3-4: NLEN (one's complement of LEN; not checked here — the // BGZF footer's CRC32/ISIZE are authoritative) // bytes 5..: LEN bytes of uncompressed payload - // BTYPE lives in bits 1-2 of the first byte, so `b & 0b110 == 0` means - // "stored". `payload_len == LEN + 5` is the structural guarantee that - // there's exactly one stored sub-block spanning the BGZF payload — the - // form every real level-0 producer emits. - if !compressed.is_empty() && compressed[0] & 0b110 == 0 { + // `payload_len == LEN + 5` is the structural guarantee that there's exactly + // one stored sub-block spanning the BGZF payload — the form every real + // level-0 producer emits. + if is_stored_block(compressed) { return copy_stored_and_verify( compressed, uncompressed_size, @@ -492,37 +679,33 @@ fn decompress_and_verify( } let start = output.len(); + // Sizing the tail to `uncompressed_size` is what lets the shared helper take + // `out.len()` as the expected size, so the exact-fill check it performs is + // the same one this function used to do inline. output.resize(start + uncompressed_size, 0); - let result = (|| { - let bytes_written = - decompressor.deflate_decompress(compressed, &mut output[start..]).map_err(|e| { - io::Error::new( - io::ErrorKind::InvalidData, - format!("BGZF decompression failed: {e:?}"), - ) - })?; - - verify_decompression( - &output[start..start + bytes_written], - uncompressed_size, - expected_crc, - block_len, - ) - })(); + let result = deflate_into_slice_and_verify( + compressed, + expected_crc, + block_len, + decompressor, + &mut output[start..], + ); if result.is_err() { output.truncate(start); } - result + result.map(|_| ()) } -/// Decompress a deflate **stored** sub-block by copying its payload directly -/// into `output`, skipping libdeflater. +/// Parse and validate a deflate **stored** sub-block frame, returning the +/// LEN-byte payload slice (`&compressed[5..]`). Shared by +/// [`copy_stored_and_verify`] and [`copy_stored_and_verify_slice`] so the +/// stored-framing checks cannot drift between the `Vec` and fixed-slice +/// decompress entry points. /// /// The caller is responsible for confirming that `compressed[0]` has -/// `BTYPE = 00` before calling — this function then validates the rest of -/// the stored framing: +/// `BTYPE = 00` before calling. This validates the rest of the framing: /// /// * `compressed.len() >= 5` (room for the BFINAL/BTYPE byte + LEN + NLEN). /// * `LEN + 5 == compressed.len()` — exactly one stored sub-block fills the @@ -532,27 +715,20 @@ fn decompress_and_verify( /// this check fails: the input is either malformed or uses a multi-sub- /// block stored stream we have no real-world reason to support, and a /// loud error beats silently papering over corruption. -/// * `LEN == uncompressed_size` — the deflate frame's LEN agrees with the -/// BGZF footer's ISIZE field. +/// * `LEN == expected_len` — the deflate frame's LEN agrees with the BGZF +/// footer's ISIZE (the caller passes the footer ISIZE, or the caller-sized +/// output length that equals it). /// /// NLEN (one's complement of LEN) is not checked. NLEN doesn't cover the /// payload bytes, so a corrupt NLEN with an intact payload would pass the /// framing check anyway; the BGZF footer's CRC32 is the authoritative /// integrity check on the data itself. /// -/// CRC32 verification against the BGZF footer is still performed on the -/// copied bytes (the BGZF spec mandates it, and the bypass is intended to -/// skip the libdeflater call, not the integrity check). +/// # Errors /// -/// On any framing or verification failure, `output` is rolled back to its -/// original length. -fn copy_stored_and_verify( - compressed: &[u8], - uncompressed_size: usize, - expected_crc: u32, - block_len: usize, - output: &mut Vec, -) -> io::Result<()> { +/// Returns `io::ErrorKind::InvalidData` if the stored framing is truncated, +/// spans more than one sub-block, or its LEN disagrees with `expected_len`. +fn parse_stored_frame(compressed: &[u8], expected_len: usize) -> io::Result<&[u8]> { // Deflate framing is 5 bytes (BFINAL/BTYPE byte + LEN + NLEN). if compressed.len() < 5 { return Err(io::Error::new( @@ -574,29 +750,68 @@ fn copy_stored_and_verify( format!("BGZF stored block size mismatch: LEN={len}, payload={}", compressed.len()), )); } - if len != uncompressed_size { + if len != expected_len { return Err(io::Error::new( io::ErrorKind::InvalidData, - format!("BGZF stored block ISIZE mismatch: footer={uncompressed_size}, LEN={len}"), + format!("BGZF stored block ISIZE mismatch: footer={expected_len}, LEN={len}"), )); } + // `len + 5 == compressed.len()`, so this is exactly the LEN payload bytes. + Ok(&compressed[5..]) +} +/// Decompress a deflate **stored** sub-block by copying its payload directly +/// into `output`, skipping libdeflater. The stored framing is validated by +/// [`parse_stored_frame`]; CRC32 verification against the BGZF footer is still +/// performed on the copied bytes (the BGZF spec mandates it, and the bypass is +/// intended to skip the libdeflater call, not the integrity check). +/// +/// The caller must have confirmed `compressed[0]` has `BTYPE = 00` first. +/// +/// On any framing or verification failure, `output` is rolled back to its +/// original length. +fn copy_stored_and_verify( + compressed: &[u8], + uncompressed_size: usize, + expected_crc: u32, + block_len: usize, + output: &mut Vec, +) -> io::Result<()> { + let payload = parse_stored_frame(compressed, uncompressed_size)?; let start = output.len(); - // `len + 5 == compressed.len()` is checked above, so `&compressed[5..]` - // is exactly the LEN payload bytes. - output.extend_from_slice(&compressed[5..]); - let result = verify_decompression( - &output[start..start + len], - uncompressed_size, - expected_crc, - block_len, - ); + output.extend_from_slice(payload); + let result = verify_decompression(&output[start..], uncompressed_size, expected_crc, block_len); if result.is_err() { output.truncate(start); } result } +/// Slice-writing sibling of [`copy_stored_and_verify`]: copy a deflate +/// **stored** sub-block's payload straight into a caller-sized `out` slice +/// (whose length is the block's ISIZE), skipping libdeflater, then verify the +/// BGZF footer CRC32 over the copied bytes. Returns the number of bytes written +/// (`== out.len()` on success). Used by [`decompress_into_slice`]'s fast path. +/// +/// The caller must have confirmed `compressed[0]` has `BTYPE = 00` first. +/// Unlike the `Vec` sibling there is nothing to roll back: `out` is the +/// caller's buffer, and — matching [`decompress_into_slice`]'s non-stored path +/// — a verification failure leaves `out` holding the (CRC-rejected) bytes, +/// which the caller discards along with the returned error. +fn copy_stored_and_verify_slice( + compressed: &[u8], + out: &mut [u8], + expected_crc: u32, + block_len: usize, +) -> io::Result { + // `out` is caller-sized to the footer's ISIZE, so passing `out.len()` also + // enforces LEN == ISIZE; the returned payload is then exactly `out.len()`. + let payload = parse_stored_frame(compressed, out.len())?; + out.copy_from_slice(payload); + verify_decompression(out, out.len(), expected_crc, block_len)?; + Ok(out.len()) +} + // ============================================================================ // Tests // ============================================================================ @@ -946,7 +1161,7 @@ mod tests { /// A stored block whose `LEN` field disagrees with the BGZF payload size /// is malformed; the fast path should reject it (rather than silently /// truncating or falling through to libdeflater). This exercises the - /// `len + 5 != compressed.len()` branch in `copy_stored_and_verify`. + /// `len + 5 != compressed.len()` branch in `parse_stored_frame`. #[test] fn test_decompress_stored_block_rejects_bad_len() { use crate::writer::InlineBgzfCompressor; @@ -1008,8 +1223,7 @@ mod tests { /// A stored block whose deflate LEN agrees with the wire payload size /// but disagrees with the BGZF footer's ISIZE must be rejected. - /// Exercises the `len != uncompressed_size` branch in - /// `copy_stored_and_verify`. + /// Exercises the `len != expected_len` branch in `parse_stored_frame`. #[test] fn test_decompress_stored_block_rejects_isize_mismatch() { use crate::writer::InlineBgzfCompressor; @@ -1082,4 +1296,253 @@ mod tests { ); assert!(out.is_empty(), "output should be rolled back on failure"); } + + // ── `uncompressed_size` ───────────────────────────────────────────────── + + /// The public accessor callers size their slot with must agree with the + /// footer for a real block, and must reject the two inputs that would + /// otherwise become a bad allocation: a block too short to have a footer, + /// and a footer claiming more than a block can hold. + #[rstest] + #[case::level_0_stored(0)] + #[case::level_6_deflate(6)] + fn uncompressed_size_reports_the_footer(#[case] level: u32) { + let payload = one_block_payload(); + let block = first_block_at_level(&payload, level); + assert_eq!(uncompressed_size(&block).expect("valid block"), payload.len()); + } + + /// A claim above the per-block maximum is refused rather than handed back + /// for a caller to allocate from. ISIZE is a `u32`, so the untrusted range + /// runs to 4 GiB. + #[test] + fn uncompressed_size_rejects_a_claim_above_the_block_maximum() { + let mut block = first_block_at_level(&one_block_payload(), 6); + let len = block.len(); + let claimed = u32::try_from(MAX_UNCOMPRESSED_BLOCK_SIZE + 1).expect("fits"); + block[len - 4..].copy_from_slice(&claimed.to_le_bytes()); + let err = uncompressed_size(&block).expect_err("an oversized claim must be rejected"); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(err.to_string().contains("above the"), "got: {err}"); + } + + /// Exactly the per-block maximum is legal — other writers may fill a block + /// to the spec limit even though this crate stops at `BGZF_MAX_BLOCK_SIZE`, + /// so the bound must not be off by one against them. + #[test] + fn uncompressed_size_accepts_the_block_maximum() { + let mut block = first_block_at_level(&one_block_payload(), 6); + let len = block.len(); + let claimed = u32::try_from(MAX_UNCOMPRESSED_BLOCK_SIZE).expect("fits"); + block[len - 4..].copy_from_slice(&claimed.to_le_bytes()); + assert_eq!( + uncompressed_size(&block).expect("the maximum itself is valid"), + MAX_UNCOMPRESSED_BLOCK_SIZE + ); + } + + /// Too short to hold a footer at all — must error rather than read the + /// ISIZE out of whatever bytes happen to be there. + #[test] + fn uncompressed_size_rejects_a_block_without_a_footer() { + let err = uncompressed_size(&[0u8; 10]).expect_err("a short block must be rejected"); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(err.to_string().contains("too short to contain"), "got: {err}"); + } + + // ── `decompress_into_slice` ───────────────────────────────────────────── + + /// Compress `payload` into BGZF blocks at `level` and return the first + /// block's raw bytes. Level 0 emits deflate **stored** blocks (the fast + /// path); any other level emits a real deflate stream. + fn first_block_at_level(payload: &[u8], level: u32) -> Vec { + let mut compressor = crate::writer::InlineBgzfCompressor::new(level); + compressor.write_all(payload).expect("write payload"); + compressor.flush().expect("flush"); + compressor.take_blocks().into_iter().next().expect("payload fits in one block").data + } + + /// A payload that fits comfortably in one BGZF block (< 64 KiB). + fn one_block_payload() -> Vec { + b"the quick brown fox jumps over the lazy dog".repeat(100) + } + + /// Both decompress paths must fill a pre-sized `&mut [u8]` slot with + /// exactly the original bytes: level 0 through the stored-block copy, any + /// other level through libdeflater. `expect_stored` pins which path the + /// case actually takes, so a producer change that stops emitting stored + /// blocks fails here rather than silently reducing coverage to one path. + #[rstest] + #[case::level_0_stored(0, true)] + #[case::level_6_deflate(6, false)] + fn decompress_into_slice_fills_presized_slot(#[case] level: u32, #[case] expect_stored: bool) { + let payload = one_block_payload(); + let block = first_block_at_level(&payload, level); + + let compressed = compressed_data_from_slice(&block); + assert_eq!( + is_stored_block(compressed), + expect_stored, + "level {level} took the unexpected decompress path" + ); + + let mut out = vec![0u8; uncompressed_size_from_slice(&block)]; + let n = decompress_into_slice(&block, &mut Decompressor::new(), &mut out) + .expect("decompress into slot"); + assert_eq!(n, payload.len()); + assert_eq!(&out[..n], payload.as_slice()); + } + + /// The BGZF EOF marker carries no payload, and every stream ends with one, + /// so a caller decompressing each block of a stream in turn hits it. It has + /// to succeed against a zero-length slot, as the two `Vec` siblings do. + /// + /// Its payload is `03 00` — a fixed-Huffman empty frame — so + /// `is_stored_block` is false and it would otherwise reach libdeflater with + /// a zero-length output slice. That returns `Ok(0)` today; this pins the + /// result to the crate rather than to that behaviour. + #[test] + fn decompress_into_slice_accepts_the_eof_marker() { + let compressed = compressed_data_from_slice(&BGZF_EOF); + assert!( + !is_stored_block(compressed), + "the EOF payload is a fixed-Huffman frame, so the stored path must not claim it" + ); + assert_eq!(uncompressed_size(&BGZF_EOF).expect("the EOF marker is a valid block"), 0); + + let n = decompress_into_slice(&BGZF_EOF, &mut Decompressor::new(), &mut []) + .expect("the EOF marker must decompress to nothing"); + assert_eq!(n, 0); + } + + /// A zero-ISIZE block must still reject a wrongly-sized slot. The obvious + /// place to short-circuit on `uncompressed_size == 0` is *before* the + /// slot-size check, which would turn this into a silent `Ok(0)` and hide the + /// caller's sizing bug — the same misreport the up-front check exists to + /// prevent. + #[test] + fn decompress_into_slice_rejects_a_sized_slot_for_an_empty_block() { + let mut out = [0u8; 16]; + let err = decompress_into_slice(&BGZF_EOF, &mut Decompressor::new(), &mut out) + .expect_err("a 16-byte slot for a 0-byte block must be rejected"); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(err.to_string().contains("output slice is"), "got: {err}"); + } + + /// A block at `level` with a flipped footer CRC32 bit: it still produces + /// the right number of bytes, but the checksum no longer matches. Run at + /// both levels because the CRC is verified at two separate call sites — + /// `copy_stored_and_verify_slice` for level 0 and + /// `deflate_into_slice_and_verify` for the rest — and deleting either one + /// used to leave the whole suite green. + fn block_bad_crc(level: u32) -> (Vec, Vec) { + let mut block = first_block_at_level(&one_block_payload(), level); + let len = block.len(); + block[len - 8] ^= 0x01; // footer CRC32 is bytes [len - 8 .. len - 4] + let out = vec![0u8; uncompressed_size_from_slice(&block)]; + (block, out) + } + + /// A level-6 block whose footer ISIZE is 16 too high, with `out` sized to + /// that inflated ISIZE. libdeflater reports the true (shorter) length via + /// `Ok(n)` rather than erroring, so this is what reaches the exact-fill + /// check — the guarantee the docs single out, and the one branch that + /// stops the arena being handed a slot with a stale tail. + fn block_short_fill() -> (Vec, Vec) { + let mut block = first_block_at_level(&one_block_payload(), 6); + let len = block.len(); + let inflated = uncompressed_size_from_slice(&block) + 16; + block[len - 4..].copy_from_slice(&u32::try_from(inflated).expect("fits").to_le_bytes()); + (block, vec![0u8; inflated]) + } + + /// A footer claiming more than a single BGZF block can hold. ISIZE is a + /// `u32`, so an unbounded contract would have a caller allocate up to 4 GiB + /// from a corrupt footer before anything validated it. + fn block_isize_above_max() -> (Vec, Vec) { + let mut block = first_block_at_level(&one_block_payload(), 6); + let len = block.len(); + let claimed = u32::try_from(MAX_UNCOMPRESSED_BLOCK_SIZE + 1).expect("fits"); + block[len - 4..].copy_from_slice(&claimed.to_le_bytes()); + // Deliberately a *correctly* sized slot for the claim: the bound must + // fire on the claim itself, not merely because `out` disagrees with it. + (block, vec![0u8; MAX_UNCOMPRESSED_BLOCK_SIZE + 1]) + } + + /// A caller-oversized `out` (ISIZE + 16). The error must name the slice as + /// the thing that is wrong: before the up-front size check this surfaced as + /// a complaint about the block's own ISIZE, sending a reader after file + /// corruption that isn't there. The level is irrelevant — this check fires + /// before the stored/deflate dispatch — so one case covers it. + fn block_oversized_out() -> (Vec, Vec) { + let block = first_block_at_level(&one_block_payload(), 6); + let out = vec![0u8; uncompressed_size_from_slice(&block) + 16]; + (block, out) + } + + /// Synthesised stored block whose 4-byte payload can't hold the 5-byte + /// deflate stored frame (BFINAL/BTYPE + LEN + NLEN), so the stored fast + /// path errors before copying. `InlineBgzfCompressor` only emits valid + /// frames, so this is built byte by byte: 18-byte header + 4-byte payload + /// + 8-byte footer. + /// + /// The header is a **well-formed** one — `XLEN = 6` and the `BC` subfield's + /// `SLEN = 2` are set, not left zero — so the fixture isolates the stored + /// framing defect it is named for. An invalid header here would be rejected + /// for a different reason by anything that validates one, making the case + /// pass for the wrong cause. + fn block_truncated_stored_framing() -> (Vec, Vec) { + const BLOCK_SIZE: usize = BGZF_HEADER_SIZE + 4 + BGZF_FOOTER_SIZE; + let mut data = vec![0u8; BLOCK_SIZE]; + data[0] = 0x1f; // gzip magic + deflate method + FEXTRA flag + data[1] = 0x8b; + data[2] = 0x08; + data[3] = 0x04; + data[10] = 0x06; // XLEN = 6: the extra field is exactly the BC subfield + data[11] = 0x00; + data[12] = b'B'; // BC subfield ID + data[13] = b'C'; + data[14] = 0x02; // SLEN = 2: BC holds a two-byte BSIZE + data[15] = 0x00; + let bsize_bytes = u16::try_from(BLOCK_SIZE - 1).expect("block fits in u16").to_le_bytes(); + data[16] = bsize_bytes[0]; + data[17] = bsize_bytes[1]; + // Payload bytes 18..22 stay zero → BTYPE bits (1-2) are 00, the stored + // fast path. Footer ISIZE = 1 matches the caller-sized `out`; the CRC + // is never reached. + data[BLOCK_SIZE - 4] = 1; + debug_assert!(crate::header::is_bgzf_header(&data), "fixture header must be well-formed"); + (data, vec![0u8; 1]) + } + + /// A block shorter than the 26-byte header + footer minimum — must error + /// rather than panic on an out-of-bounds slice or a subtract overflow. + fn block_too_short() -> (Vec, Vec) { + (vec![0u8; 10], vec![0u8; 16]) + } + + /// Every malformed input is rejected as `InvalidData` with a message that + /// names what went wrong, so a corrupt block is diagnosable rather than a + /// bare "decompression failed". + #[rstest] + #[case::bad_crc_deflate(block_bad_crc(6), "CRC32")] + #[case::bad_crc_stored(block_bad_crc(0), "CRC32")] + #[case::short_fill(block_short_fill(), "size mismatch")] + #[case::isize_above_max(block_isize_above_max(), "above the")] + #[case::oversized_out(block_oversized_out(), "output slice is")] + #[case::truncated_stored_framing(block_truncated_stored_framing(), "stored block too small")] + #[case::too_short_block(block_too_short(), "too short to contain")] + fn decompress_into_slice_rejects_invalid( + #[case] block_and_out: (Vec, Vec), + #[case] expect_substr: &str, + ) { + let (block, mut out) = block_and_out; + let err = decompress_into_slice(&block, &mut Decompressor::new(), &mut out) + .expect_err("malformed block must be rejected"); + assert_eq!(err.kind(), io::ErrorKind::InvalidData, "expected InvalidData, got {err:?}"); + assert!( + err.to_string().contains(expect_substr), + "error should contain {expect_substr:?}, got: {err}" + ); + } } diff --git a/crates/fgumi-bgzf/src/writer.rs b/crates/fgumi-bgzf/src/writer.rs index d63e9133a..8de60e6c0 100644 --- a/crates/fgumi-bgzf/src/writer.rs +++ b/crates/fgumi-bgzf/src/writer.rs @@ -16,6 +16,28 @@ use std::io; /// Maximum uncompressed size for a BGZF block (64KB - header/footer overhead). pub const BGZF_MAX_BLOCK_SIZE: usize = bgzf::BGZF_BLOCK_SIZE; +/// Cap on the number of compression buffers [`InlineBgzfCompressor`] pools for +/// reuse. +/// +/// The pool only has to cover the *compress* side, which takes one buffer at a +/// time: `compress_current_buffer` pops exactly one per block. The queue on the +/// other side can be arbitrarily long — `write_all` appends a block per 64 KiB +/// written — so a drain of many blocks hands back more buffers than the +/// compressor will ever want at once. The surplus is dropped rather than parked +/// for the compressor's lifetime, which is what keeps the pool's memory a +/// function of this constant instead of of how much the caller buffered before +/// draining. +pub const MAX_POOLED_BUFFERS: usize = 4; + +/// Largest buffer [`InlineBgzfCompressor::recycle_buffer`] will pool. +/// +/// A block buffer is sized by `Compressor::compress` to +/// `header + deflate_compress_bound(input) + footer`. That bound exceeds the +/// input for incompressible data, so a legitimate buffer can run somewhat over +/// [`BGZF_MAX_BLOCK_SIZE`]; double it so no real block buffer is ever refused, +/// while still rejecting a wildly oversized `Vec` a caller hands back. +pub const MAX_POOLED_BUFFER_BYTES: usize = BGZF_MAX_BLOCK_SIZE * 2; + // ============================================================================ // Block types // ============================================================================ @@ -184,6 +206,36 @@ impl InlineBgzfCompressor { std::mem::take(&mut self.completed_blocks) } + /// Return a drained block buffer to the internal pool for reuse by a later + /// block compression. + /// + /// Consumers that drive the compressor with [`write_all`](Self::write_all) + + /// [`flush`](Self::flush) + [`take_blocks`](Self::take_blocks) (rather than + /// [`write_blocks_to`](Self::write_blocks_to), which recycles automatically) + /// otherwise leave `buffer_pool` empty, so every block compression allocates + /// a fresh output `Vec`. Such a consumer can hand back any block `Vec` it is + /// done with via this method to restore the recycling. + /// + /// The buffer is cleared — but not shrunk, since its capacity is the whole + /// point — before being pooled. A buffer is only pooled if the pool has + /// room ([`MAX_POOLED_BUFFERS`]) *and* the buffer is a plausible block + /// buffer ([`MAX_POOLED_BUFFER_BYTES`]); otherwise it is dropped. Bounding + /// the count alone would let one oversized `Vec` handed in by a caller sit + /// in the pool for the compressor's lifetime, so the pool's memory is + /// bounded on both axes. + /// + /// Not to be confused with the pipeline's `WorkerCoreState::recycle_buffer`, + /// which pools the *uncompressed* worker buffers under a different cap. + /// This one pools the compressor's own compressed-block output buffers. + pub fn recycle_buffer(&mut self, mut buffer: Vec) { + if self.buffer_pool.len() < MAX_POOLED_BUFFERS + && buffer.capacity() <= MAX_POOLED_BUFFER_BYTES + { + buffer.clear(); + self.buffer_pool.push(buffer); + } + } + /// Write all completed compressed blocks directly to output and recycle buffers. /// /// This is efficient for single-threaded use as it writes blocks directly @@ -191,14 +243,39 @@ impl InlineBgzfCompressor { /// /// # Errors /// - /// Returns an error if writing to the output fails. + /// Returns an error if writing to the output fails. The block that failed + /// to write, and every block after it, is left in the queue, so they are + /// available to **inspect** rather than being silently discarded. + /// + /// They are *not* safe to replay. [`io::Write::write_all`] loops over + /// [`write`](io::Write::write), advancing past each `Ok(n)`, so it can + /// commit part of the failing block to `output` before an error surfaces — + /// and the whole block is re-queued, not the unwritten remainder. Writing + /// the queue again would repeat those bytes inside a gzip member and + /// produce a stream no BGZF reader can decode. Treat `output` as + /// unrecoverable after this returns an error: the queued blocks are for + /// diagnostics, or for writing somewhere new. pub fn write_blocks_to(&mut self, output: &mut W) -> io::Result<()> { - for block in self.completed_blocks.drain(..) { - output.write_all(&block.data)?; - // Recycle the buffer for reuse - let mut buf = block.data; - buf.clear(); - self.buffer_pool.push(buf); + // Nothing queued is the common call (a flush with no pending blocks). + // Returning here keeps `completed_blocks`'s allocation, which the + // `mem::take` below would otherwise hand away and force a realloc on + // the next compression. + if self.completed_blocks.is_empty() { + return Ok(()); + } + // Drain into a temporary so `recycle_buffer` (which borrows `self` + // mutably) can be called per block without holding a borrow on + // `self.completed_blocks`. + let mut remaining = std::mem::take(&mut self.completed_blocks).into_iter(); + while let Some(block) = remaining.next() { + if let Err(e) = output.write_all(&block.data) { + self.completed_blocks.push(block); + self.completed_blocks.extend(remaining); + return Err(e); + } + // Route the drained buffer through the capped recycle path so the + // pool stays bounded, just like the steady-state recycle path. + self.recycle_buffer(block.data); } Ok(()) } @@ -209,9 +286,12 @@ impl InlineBgzfCompressor { return Ok(()); } - // Get buffer from pool or allocate new + // Get buffer from pool or allocate new. No clear is needed: the only + // thing this buffer is handed to is `Compressor::compress`, which + // resizes it from zero before writing (bgzf 0.4's `resize_uninit` + // begins with `Vec::clear`). Capacity is what the pool is for, and that + // survives. let mut compressed_data = self.buffer_pool.pop().unwrap_or_default(); - compressed_data.clear(); // Compress using bgzf crate's Compressor self.compressor @@ -443,4 +523,169 @@ mod tests { fn test_compress_level_out_of_range_panics() { let _ = InlineBgzfCompressor::new(13); } + + // ── Buffer recycling ──────────────────────────────────────────────────── + + /// Compress `payload` into one block and hand the compressor back. + fn compressor_with_one_block(payload: &[u8]) -> (InlineBgzfCompressor, Vec) { + let mut compressor = InlineBgzfCompressor::new(6); + compressor.write_all(payload).expect("writing data should succeed"); + compressor.flush().expect("flushing compressor should succeed"); + let blocks = compressor.take_blocks(); + (compressor, blocks) + } + + /// A `take_blocks` consumer keeps ownership of every block buffer, so the + /// pool is left empty and each subsequent compression allocates afresh. + /// Handing a drained buffer back must restore the recycling. + #[test] + fn test_recycle_buffer_repopulates_the_pool() { + let (mut compressor, blocks) = compressor_with_one_block(b"payload for one block"); + assert!( + compressor.buffer_pool.is_empty(), + "take_blocks should leave the pool empty, that is the gap recycle_buffer closes" + ); + + let buffer = blocks.into_iter().next().expect("one block").data; + let recycled_capacity = buffer.capacity(); + // The identity of the allocation, not just of the `Vec`. Everything + // below is about this exact heap block being handed back to the next + // compression rather than freed and replaced. + let recycled_ptr = buffer.as_ptr(); + compressor.recycle_buffer(buffer); + assert_eq!(compressor.buffer_pool.len(), 1); + // The allocation is what is being reused, so the capacity has to survive + // pooling -- a `shrink_to_fit` alongside the `clear` would satisfy every + // other assertion here while deleting the point of the method. + assert_eq!( + compressor.buffer_pool[0].capacity(), + recycled_capacity, + "pooling must retain the buffer's allocation, not just the buffer" + ); + + // The next compression must consume the pooled buffer rather than + // allocate. Asserting the pool merely drained would also pass for a + // pop-and-discard implementation, so pin the produced block to the + // recycled allocation instead. + compressor.write_all(b"a second block").expect("writing data should succeed"); + compressor.flush().expect("flushing compressor should succeed"); + assert!(compressor.buffer_pool.is_empty(), "the pooled buffer should have been taken"); + let next = compressor.take_blocks(); + assert_eq!( + next[0].data.as_ptr(), + recycled_ptr, + "the new block should be built in the recycled allocation, not a fresh one" + ); + } + + /// The pool is bounded on capacity as well as count: a caller handing back + /// an outsized `Vec` must not park that allocation for the compressor's + /// lifetime. A real block buffer (~64 KiB) is well under the limit, so this + /// rejects only what it is meant to. + #[test] + fn test_recycle_buffer_refuses_an_oversized_buffer() { + let (mut compressor, blocks) = compressor_with_one_block(b"payload for one block"); + let real = blocks.into_iter().next().expect("one block").data; + assert!( + real.capacity() <= MAX_POOLED_BUFFER_BYTES, + "a real block buffer ({} bytes) must still be poolable", + real.capacity() + ); + + compressor.recycle_buffer(Vec::with_capacity(MAX_POOLED_BUFFER_BYTES + 1)); + assert!(compressor.buffer_pool.is_empty(), "an oversized buffer should not be pooled"); + + compressor.recycle_buffer(real); + assert_eq!(compressor.buffer_pool.len(), 1, "a real block buffer should be pooled"); + } + + /// `write_blocks_to` is the other path into the pool, and the line this + /// change rewrote. It must recycle what it drains, and honour the same cap + /// as the explicit entry point -- a drain of many blocks cannot grow the + /// pool past it. + #[test] + fn test_write_blocks_to_recycles_up_to_the_cap() { + // More blocks than the cap, so both halves of the behaviour are visible. + let payload = vec![b'x'; BGZF_MAX_BLOCK_SIZE * (MAX_POOLED_BUFFERS + 2)]; + let mut compressor = InlineBgzfCompressor::new(6); + compressor.write_all(&payload).expect("writing data should succeed"); + compressor.flush().expect("flushing compressor should succeed"); + let queued = compressor.completed_blocks.len(); + assert!(queued > MAX_POOLED_BUFFERS, "need more than {MAX_POOLED_BUFFERS} blocks"); + assert!(compressor.buffer_pool.is_empty(), "pool starts empty"); + + let mut sink = Vec::new(); + compressor.write_blocks_to(&mut sink).expect("writing blocks should succeed"); + + assert!(compressor.completed_blocks.is_empty(), "every block should have been written"); + assert_eq!( + compressor.buffer_pool.len(), + MAX_POOLED_BUFFERS, + "draining {queued} blocks should fill the pool to the cap and drop the surplus" + ); + } + + /// A consumer that hands back more buffers than the compressor can use + /// must not grow the pool without limit; surplus buffers are dropped. + #[test] + fn test_recycle_buffer_bounds_the_pool() { + let (mut compressor, _blocks) = compressor_with_one_block(b"payload for one block"); + for _ in 0..32 { + compressor.recycle_buffer(vec![0u8; 1024]); + } + assert_eq!( + compressor.buffer_pool.len(), + MAX_POOLED_BUFFERS, + "pool should stop growing at the cap" + ); + } + + /// A sink that accepts `accept` bytes total and then fails, so a + /// multi-block `write_blocks_to` fails partway through. + struct ShortSink { + accept: usize, + } + + impl io::Write for ShortSink { + fn write(&mut self, buf: &[u8]) -> io::Result { + if buf.len() > self.accept { + return Err(io::Error::new(io::ErrorKind::WriteZero, "sink is full")); + } + self.accept -= buf.len(); + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + /// On a write failure, the unwritten blocks must stay queued. Draining them + /// into the void would lose compressed output with no way for the caller to + /// notice, since the error says nothing about how far the write got. + #[test] + fn test_write_blocks_to_keeps_unwritten_blocks_on_error() { + // Two full blocks, then a sink that has room for the first one only. + let payload = vec![b'x'; BGZF_MAX_BLOCK_SIZE * 2]; + let mut compressor = InlineBgzfCompressor::new(6); + compressor.write_all(&payload).expect("writing data should succeed"); + compressor.flush().expect("flushing compressor should succeed"); + let first_block_len = compressor.completed_blocks[0].data.len(); + let queued = compressor.completed_blocks.len(); + assert!(queued >= 2, "expected at least two blocks, got {queued}"); + + let mut sink = ShortSink { accept: first_block_len }; + let err = compressor + .write_blocks_to(&mut sink) + .expect_err("the sink should reject the second block"); + assert_eq!(err.kind(), io::ErrorKind::WriteZero); + // Identity, not just the count: a queue of the right length holding the + // wrong blocks (say, the tail restored but the failing block dropped) + // loses exactly the output this behaviour exists to preserve. Serials + // are contiguous from 0, so the retained queue must start at 1 and run + // to the end. + let retained: Vec = compressor.completed_blocks.iter().map(|b| b.serial).collect(); + let expected: Vec = (1..u64::try_from(queued).expect("block count fits")).collect(); + assert_eq!(retained, expected, "the failing block and its tail should still be queued"); + } }