diff --git a/.github/workflows/leak-guard.yml b/.github/workflows/leak-guard.yml new file mode 100644 index 0000000..ed8efb4 --- /dev/null +++ b/.github/workflows/leak-guard.yml @@ -0,0 +1,35 @@ +name: leak-guard + +# Self-contained public-repo leak gate. Public CI cannot reach the private +# tooling, so this encodes only the generic net: internal-infra references, +# tracked CLAUDE.md/.claude paths, and AI-attribution in commit messages. +# No project-specific reverse-engineering vocabulary lives here. + +on: [push, pull_request] + +jobs: + leak-guard: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + with: + fetch-depth: 0 + - name: Compute commit range + id: range + run: | + if [ "${{ github.event_name }}" = "pull_request" ]; then + base="${{ github.event.pull_request.base.sha }}" + head="${{ github.event.pull_request.head.sha }}" + echo "range=$base..$head" >> "$GITHUB_OUTPUT" + else + before="${{ github.event.before }}" + after="${{ github.sha }}" + # New branch / first push: github.event.before is all-zeros. + if [ -z "$before" ] || [ "$before" = "0000000000000000000000000000000000000000" ]; then + echo "range=$after" >> "$GITHUB_OUTPUT" + else + echo "range=$before..$after" >> "$GITHUB_OUTPUT" + fi + fi + - name: Run leak-guard + run: bash ci/leak-guard.sh "${{ steps.range.outputs.range }}" diff --git a/.gitignore b/.gitignore index 4395118..fbd9c29 100644 --- a/.gitignore +++ b/.gitignore @@ -4,7 +4,13 @@ Cargo.lock *.swo .DS_Store .cargo/ -.claude/worktrees/ # session scratch — never track (may contain RE breadcrumbs) scratch/ + +# stray local build artifact +/rust_out + +# internal agent context — never publish (path AND dir; leak-guard blocks both) +CLAUDE.md +.claude/ diff --git a/CHANGELOG.md b/CHANGELOG.md index f85705d..f51b94e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,31 @@ # Changelog +## 0.31.0 (2026-06-08) + +Hardening and correctness release: a library-wide review-and-fix pass across +the mux pipeline, codec parsers, AACS/CSS decryption, UDF/MPLS/CLPI parsing, +sector prefetch, multi-pass recovery, drive/SCSI, disc labels, and I/O. + +### Fixed + +- AACS: hardened keydb title parsing, redacted secrets from handshake debug + output, and tightened the media-key variant verification gate. A trailing + partial aligned unit is now tolerated when clear and rejected (fail-loud) + when scrambled, rather than passed through. +- Parsers: corrected the MPLS playlist-mark type offset, added per-extent and + cumulative allocation bounds in the UDF reader, and made the drive identity + probe best-effort on a CHECK CONDITION. +- Mux/codec: start-code de-duplication, framing fixes, M2TS packet hardening, + and bounds guards across the HEVC/H.264/VC-1/MPEG-2/TrueHD/DTS/PGS paths. +- Recovery/drive/SCSI: guarded READ CAPACITY short transfers, unified the + error-code / io-kind mapping, and hardened the platform unlock path. +- Robustness: overflow/underflow guards on values derived from untrusted disc + input, deterministic prefetch shutdown, and error display remains code-only. + +### Changed + +- Release profile now builds with thin LTO and a single codegen unit. + ## 0.29.0 (2026-06-06) ### Fixed @@ -90,14 +116,14 @@ ### Changed -- `Drive::is_libredrive_active()` renamed to `Drive::is_raw_read_active()`. - Same semantics; old name removed. Mirrored on the internal - `PlatformDriver::is_libredrive_active()` trait method (now - `is_raw_read_active()`). -- `Error::AacsLibredriveUnsupported` renamed to - `Error::AacsRawReadUnsupported`; the underlying numeric code (E7016) - is unchanged. The `E_AACS_LIBREDRIVE_UNSUPPORTED` constant is - renamed to `E_AACS_RAW_READ_UNSUPPORTED`. +- `Drive::is_raw_read_active()` is the current name for the + raw-read-capability probe (an older internal name was retired in this + pass). Same semantics; the old name was removed. Mirrored on the + internal `PlatformDriver::is_raw_read_active()` trait method. +- `Error::AacsRawReadUnsupported` is the current name for the + raw-read-unsupported variant; the underlying numeric code (E7016) + is unchanged. The corresponding constant is + `E_AACS_RAW_READ_UNSUPPORTED`. No behavioural change — purely a rename pass. @@ -131,26 +157,26 @@ No behavioural change — purely a rename pass. ### Fixed -- **Libredrive raw-read VID shortcut deleted.** v0.25.11 introduced a - `do_handshake` branch that, on libredrive-active drives, skipped the +- **Raw-read VID shortcut deleted.** v0.25.11 introduced a + `do_handshake` branch that, on raw-read-capable drives, skipped the AACS cert handshake and issued `READ_DISC_STRUCTURE` format 0x80 - with AGID=0 directly. The hypothesis was that unlocked + with AGID=0 directly. The hypothesis was that raw-read-capable drives would serve VID without auth. Empirical test (BU40N + UHD disc, 2026-05-21) showed the drive returns `0x05 / 0x6F / 0x02` (`ILLEGAL_REQUEST / Copy protection key exchange failure: KEY NOT ESTABLISHED`) to that CDB regardless of - drive-unlock state. The AACS spec requires a successful + raw-read state. The AACS spec requires a successful `REPORT_KEY` / `SEND_KEY` exchange to establish an AGID before format 0x80 returns VID; that requirement is enforced by the drive - itself and isn't bypassed by libredrive firmware. The shortcut - fired for every libredrive-active drive, so v0.25.11 / v0.25.12 + itself and isn't bypassed by raw-read mode. The shortcut + fired for every raw-read-capable drive, so v0.25.11 / v0.25.12 MOVIE scans were stuck at E7017 instead of progressing to the real wall (no DK walks MKB v77). - `Disc::do_handshake` now always routes through `do_handshake_cert`. - `Drive::is_libredrive_active()` and the Mt1959 MMkv+LbDr marker + `Drive::is_raw_read_active()` and the Mt1959 marker detection are kept as informational signals (logged in the `handshake_entry` warn line) but no longer steer the auth path. -- `read_volume_id_libredrive` deleted (~50 LOC). +- The raw-read VID read helper was deleted (~50 LOC). The corollary: AACS resolution on HRL-burned drives + UHD discs now fails honestly. Either cert auth succeeds (drive unlock may or @@ -170,15 +196,15 @@ v0.25.12 release for details. ### Added -- **Libredrive raw-read VID path.** When the Mt1959 unlock response - confirms both the active-mode (`MMkv`) and mode-ID (`LbDr`) markers, - `Drive::is_libredrive_active()` returns true and `do_handshake` +- **Raw-read VID path.** When the Mt1959 unlock response + confirms both the active-mode and mode-ID markers, + `Drive::is_raw_read_active()` returns true and `do_handshake` skips the AACS cert dance entirely — VID is retrieved via `READ_DISC_STRUCTURE` format 0x80 with AGID=0, and bus encryption is already off. This is what unblocks UHD ripping on drives whose - leaked host cert is on the AACS HRL. + host cert is on the AACS HRL. - New `Error` variants for finer-grained AACS failure reporting: - `AacsHostCertRejected` (E7015), `AacsLibredriveUnsupported` + `AacsHostCertRejected` (E7015), `AacsRawReadUnsupported` (E7016), `AacsVidUnavailable` (E7017), `AacsMkUnavailable` (E7018), `AacsVukNotInKeydb` (E7019). Lets CLIs/UIs render which piece of the AACS chain failed instead of always saying "no keys." @@ -1190,9 +1216,8 @@ LUN/BUS/HOST reset. `READ_RECOVERY_TIMEOUT_MS` (60 s) unchanged. + close-on-timeout in bg thread` to a single synchronous blocking `ioctl(fd, SG_IO, &hdr)`. The old pattern abandoned slow-but-alive commands faster than the drive could drain its internal queue, -deepening the BU40N wedge. Per the audit at -internal SCSI-architecture research, -no reference project (MakeMKV / sg_dd / ddrescue) does what we did — +deepening the BU40N wedge. Per the SCSI-architecture review, +no existing consumer ripper or dd-based tool does what we did — all use sync blocking SG_IO with 8-60 s timeouts and let the kernel's mid-layer (`scsi_eh.rst`) run ABORT TASK / LUN RESET / BUS RESET / HOST RESET escalation internally. @@ -1377,7 +1402,7 @@ at the SCSI + Disc::copy boundaries we can't diagnose where the time goes. This release adds the telemetry. No behavior change; instrumentation only. -- New dep: `tracing = "0.1"`. Per project docs, debug/trace logging is permitted +- New dep: `tracing = "0.1"`. Per CLAUDE.md, debug/trace logging is permitted in libfreemkv (the no-English rule applies to errors, not telemetry). Consumers (autorip) wire a tracing subscriber and pipe events into the JSONL debug log automatically. @@ -1835,7 +1860,7 @@ have recovered tonight's BU40N without operator intervention. ### Zero English in library — typed variants for every error path -Audit pass against the `project docs` rule (no English text in library code). +Audit pass against the `CLAUDE.md` rule (no English text in library code). Found nine call sites that violated the contract by stuffing English into `io::Error::new(kind, "…")` or by abusing `Error::DeviceNotFound { path }` as a free-form description field. Each is now a typed variant with @@ -1916,7 +1941,7 @@ the caller's convenience. Streams, Lower-level surfaces) so `cargo doc` tells callers when to reach for what. - **Dropped `ScanOptions::with_keydb()`**. The `_with_X` constructor - pattern was banned by `project docs` (one method per action). Use the + pattern was banned by `CLAUDE.md` (one method per action). Use the struct literal: `ScanOptions { keydb_path: Some(p.into()) }`. Five external call sites (autorip ×3, freemkv CLI ×3) and three test fixtures migrated. diff --git a/Cargo.toml b/Cargo.toml index 86d10f5..c1e1dea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "libfreemkv" -version = "0.30.7" +version = "0.31.0" edition = "2024" rust-version = "1.86" license = "AGPL-3.0-only" @@ -9,7 +9,11 @@ repository = "https://github.com/freemkv/libfreemkv" keywords = ["bluray", "uhd", "optical", "scsi", "disc"] categories = ["hardware-support", "multimedia"] # Keep internal AI-instruction / private notes out of the published crate. -exclude = [] +exclude = ["CLAUDE.md"] + +[profile.release] +lto = "thin" +codegen-units = 1 [dependencies] serde = { version = "1", features = ["derive"] } @@ -27,7 +31,7 @@ cmac = "0.7" zip = { version = "2", default-features = false, features = ["deflate"] } base64 = "0.22.1" # Trace-level instrumentation for Disc::copy + SgIoTransport::execute. Permitted -# under project docs ("Acceptable strings: debug/trace logging"). Consumers (autorip) +# under CLAUDE.md ("Acceptable strings: debug/trace logging"). Consumers (autorip) # wire a tracing subscriber and pipe events into the JSONL debug log. tracing = "0.1" # Bounded MPSC channel with kernel-wakeup send_timeout. Used by `io::pipeline` diff --git a/ci/leak-guard.sh b/ci/leak-guard.sh new file mode 100755 index 0000000..b44cb91 --- /dev/null +++ b/ci/leak-guard.sh @@ -0,0 +1,120 @@ +#!/usr/bin/env bash +# +# leak-guard.sh — self-contained public-repo leak gate. +# +# This is the LAST line of defense in CI. It is intentionally self-contained: +# public CI cannot reach the private tooling, so this script encodes ONLY the +# generic net — internal infrastructure references, agent-context files, and +# AI-attribution in commit messages. It deliberately contains NO project- +# specific reverse-engineering vocabulary (those words would themselves be a +# leak). The richer private scanner stays private. +# +# Fails (exit 1) if any of the following appear in the repo: +# 1. a tracked CLAUDE.md or .claude/ path (agent context — never public), +# 2. tracked file content matching the internal-infra net, +# 3. a commit message (in the given range) with AI attribution. +# +# Usage: +# leak-guard.sh [] +# optional git rev-list range to scan commit messages +# (e.g. "abc..def"). If omitted, commit-message scan is +# skipped (path + content checks always run). + +set -euo pipefail + +# Absolute path to this script, resolved before any cd, so we can exclude it +# from the content scan (it necessarily contains the detection patterns). +SELF_ABS="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/$(basename "${BASH_SOURCE[0]}")" + +REPO="$(git rev-parse --show-toplevel)" +cd "$REPO" + +fail=0 +note() { printf ' ✗ %s\n' "$1"; fail=1; } + +# Internal-infra net — GENERIC ONLY. This script ships in the public repo, so +# the patterns themselves must not name any org-specific identifier (doing so +# would itself leak the infra they guard). We catch the leak *class*: +# - RFC1918 private IPv4 ranges (10/8, 172.16/12, 192.168/16), +# - private/internal/non-routable TLDs (.internal/.local/.lan/.corp/.invalid), +# - docker.internal. +# The full org-specific net (literal hostnames, service names, repo paths, +# vendor tooling, …) lives ONLY in the private scanner and never ships here. +INFRA_RE='\b10\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}|\b172\.(1[6-9]|2[0-9]|3[01])\.[0-9]{1,3}\.[0-9]{1,3}|\b192\.168\.[0-9]{1,3}\.[0-9]{1,3}|\.internal\b|\.local\b|\.lan\b|\.corp\b|\.invalid\b|docker\.internal' +# Home-path net — GENERIC ONLY. Catches an absolute developer home path +# committed into a tracked file (a macOS /Users//… or Linux /home//… +# path). This names NO specific user — it matches the leak *class* (any home +# path), so the pattern itself reveals nothing org- or person-specific. A real +# leak (e.g. /Users/alice/Developer/x slipping into a public RELEASE.md) trips +# this regardless of whose machine it came from. The username segment is a +# literal-username class ([A-Za-z0-9._-]) so dynamic/templated paths that build +# the user at runtime — shell `/home/$USER/`, doc `/home//`, Rust +# `/home/{user}/` — do NOT false-positive; only a baked-in literal home leaks. +HOMEPATH_RE='/Users/[A-Za-z0-9._-]+/|/home/[A-Za-z0-9._-]+/' +# AI-attribution net (case-insensitive). "claude" matches only as a standalone +# word — NOT preceded by a dot/slash/alnum and NOT followed by .md — so legit +# mentions of CLAUDE.md / .claude/ in a commit message don't false-positive. +ATTR_RE='co-authored-by|generated with|🤖|(?) { if ($l =~ /$rx/) { print "$.: $&\n"; } } + ' "$1" "$2" 2>/dev/null +} + +# This script's own source necessarily contains the detection patterns (e.g. +# the regex tokens in INFRA_RE), so scanning it would always self-flag. Skip it. +SELF="$(git ls-files --full-name -- "$SELF_ABS" 2>/dev/null | head -1)" + +echo "── leak-guard: internal-infra references in tracked files ──" +while IFS= read -r f; do + case "$f" in *.png|*.jpg|*.jpeg|*.ico|*.gif|*.bin|*.crate|*.gz|*.zip|*.pdf) continue ;; esac + [ -n "$SELF" ] && [ "$f" = "$SELF" ] && continue + [ -f "$f" ] || continue + while IFS= read -r hit; do + [ -z "$hit" ] && continue + note "internal-infra reference: $f:$hit" + done < <(pcre_matches "$f" "$INFRA_RE") + while IFS= read -r hit; do + [ -z "$hit" ] && continue + note "[HOME-PATH] absolute home path: $f:$hit (no local home path may be committed to a public repo)" + done < <(pcre_matches "$f" "$HOMEPATH_RE") +done < <(git ls-files) + +RANGE="${1:-}" +if [ -n "$RANGE" ]; then + echo "── leak-guard: AI-attribution in commit messages ($RANGE) ──" + while IFS= read -r sha; do + [ -z "$sha" ] && continue + msg="$(git log -1 --format='%B' "$sha" 2>/dev/null || true)" + # Pass the pattern as an argument (not interpolated into a //) so the + # lookbehind char class and "/" don't break the regex. + hit="$(printf '%s' "$msg" | perl -e ' + my $re = $ARGV[0]; my $rx = qr/$re/i; + while (my $l = ) { if ($l =~ /($rx)/) { print "$1\n"; last; } } + ' "$ATTR_RE" | head -1 || true)" + [ -n "$hit" ] && note "commit ${sha:0:12}: message contains \"$hit\" (owner rule: zero AI attribution, ever)" + done < <(git rev-list "$RANGE" 2>/dev/null || true) +fi + +echo +if [ "$fail" -ne 0 ]; then + echo "✗ leak-guard: blocking finding(s) above — DO NOT MERGE/PUBLISH" + exit 1 +fi +echo "✓ leak-guard: clean" diff --git a/docs/rip-recovery.md b/docs/rip-recovery.md index 0d70c77..a5840a2 100644 --- a/docs/rip-recovery.md +++ b/docs/rip-recovery.md @@ -161,8 +161,8 @@ recoveries to show for it. Recovery responsibility is now layered: layer 1 handles ranges, layer 3 handles request size, neither touches the wedge-prone reset path. -**No `MODE SELECT` to disable drive retries.** Research showed neither ddrescue -nor MakeMKV does this. Drive firmware has access to raw analog signal, laser +**No `MODE SELECT` to disable drive retries.** Neither ddrescue +nor any consumer ripper does this. Drive firmware has access to raw analog signal, laser power control, and drive-specific ECC tuning that userspace can't replicate — disabling it throws away recovery headroom on marginal sectors. We fail fast via short SG_IO timeouts in pass 1 and let the firmware work the long timeout diff --git a/src/aacs/decrypt.rs b/src/aacs/decrypt.rs index 0c0c1cf..9cef653 100644 --- a/src/aacs/decrypt.rs +++ b/src/aacs/decrypt.rs @@ -45,8 +45,15 @@ pub(crate) fn aes_ecb_decrypt(key: &[u8; 16], data: &[u8; 16]) -> [u8; 16] { } /// AES-128-CBC decrypt in-place with the fixed AACS IV. -/// AES-128-CBC decrypt in-place with the fixed AACS IV. +/// +/// Precondition: `data.len()` is a multiple of 16. Any trailing partial +/// block is silently ignored; all callers pass aligned regions (6128 and +/// 2032 bytes), and the assert documents/enforces that contract. pub(crate) fn aes_cbc_decrypt(key: &[u8; 16], data: &mut [u8]) { + debug_assert!( + data.len() % 16 == 0, + "aes_cbc_decrypt requires a block-aligned slice" + ); let cipher = Aes128::new(GenericArray::from_slice(key)); let num_blocks = data.len() / 16; // Process blocks in reverse to avoid clobbering ciphertext needed for XOR @@ -87,14 +94,11 @@ pub fn is_aacs_scrambled(unit: &[u8]) -> bool { unit.len() >= ALIGNED_UNIT_LEN && !ts_syncs_intact(unit) } -/// Most TS packet positions in `unit` carry the `0x47` sync byte — i.e. the -/// unit looks like clear MPEG-TS. Syncs sit at offset 4 and every 192 bytes -/// after (4-byte TP_extra_header + 188-byte TS packet). An encrypted body -/// scrambles all but the first (which lives in the clear 16-byte seed). /// Count the MPEG-TS sync bytes (`0x47`) present at the BD-TS packet stride -/// (offset 4 and every 192 bytes after). A clear or correctly-decrypted m2ts -/// unit shows ~one per packet; an encrypted unit, or a non-content unit -/// decrypted under a key that doesn't apply, shows ~none. +/// (offset 4 and every 192 bytes after — 4-byte TP_extra_header + 188-byte +/// TS packet). A clear or correctly-decrypted m2ts unit shows ~one per +/// packet; an encrypted unit, or a non-content unit decrypted under a key +/// that doesn't apply, shows ~none. pub fn ts_sync_count(unit: &[u8]) -> usize { let mut count = 0; let mut offset = 4; @@ -220,19 +224,39 @@ pub fn unit_key_validates(unit: &[u8], unit_key: &[u8; 16]) -> bool { decrypt_unit(&mut full, unit_key) } -/// Decrypt one aligned unit trying multiple unit keys. Returns the key index that worked. -pub fn decrypt_unit_try_keys(unit: &mut [u8], unit_keys: &[[u8; 16]]) -> Option { +/// Outcome of [`decrypt_unit_try_keys`]. +/// +/// Distinguishes "the unit was already clear, no key was consumed" from "key +/// at index `i` decrypted it" — the bare `Option` form conflated the two +/// (a clear unit reported `Some(0)`, indistinguishable from key index 0, and +/// possibly out of range when `unit_keys` is empty). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum UnitKeyResult { + /// The unit was not scrambled; it was left untouched and no key was used. + AlreadyClear, + /// The unit was decrypted in place by `unit_keys[index]`. + DecryptedWith(usize), +} + +/// Decrypt one aligned unit trying multiple unit keys. +/// +/// Returns [`UnitKeyResult::AlreadyClear`] if the unit was not scrambled (no key +/// consumed), [`UnitKeyResult::DecryptedWith(i)`] if key `i` decrypted it, or +/// `None` if no key worked (the unit is restored to its original bytes). +pub fn decrypt_unit_try_keys(unit: &mut [u8], unit_keys: &[[u8; 16]]) -> Option { if !is_aacs_scrambled(unit) { - return Some(0); + return Some(UnitKeyResult::AlreadyClear); } - // Save original for retry - let original = unit[..ALIGNED_UNIT_LEN].to_vec(); + // Save original for retry. Stack-backed buffer — no heap allocation, and the + // restore-on-failure contract holds uniformly regardless of key count. + let mut original = [0u8; ALIGNED_UNIT_LEN]; + original.copy_from_slice(&unit[..ALIGNED_UNIT_LEN]); for (i, key) in unit_keys.iter().enumerate() { unit[..ALIGNED_UNIT_LEN].copy_from_slice(&original); if decrypt_unit(unit, key) { - return Some(i); + return Some(UnitKeyResult::DecryptedWith(i)); } } @@ -301,6 +325,70 @@ mod tests { assert!(decrypt_unit(&mut unit, &key)); } + #[test] + fn ts_packet_total_no_off_by_one() { + // The maximum sync count is exactly the number of stride + // positions the counting loop visits (offset 4, 196, ...), i.e. + // len / 192, NOT (len - 4) / 192 + 1. For the 6144-byte aligned unit + // the loop checks offsets 4..=5956 → 32 positions. + let unit = vec![0u8; ALIGNED_UNIT_LEN]; + assert_eq!(ts_packet_total(&unit), 32); + // Confirm the loop visits exactly that many stride positions. + let visited = (4..ALIGNED_UNIT_LEN).step_by(TS_PACKET_LEN).count(); + assert_eq!(visited, ts_packet_total(&unit)); + } + + #[test] + fn scramble_detection_at_16_32_boundary() { + // With 32 stride positions the majority threshold is + // total/2 = 16. A unit with EXACTLY half its syncs intact (16) must + // NOT be over-counted into the "scrambled" bucket by an inflated + // total: 16 > 16 is false → not-intact → scrambled. 17 intact → clear. + // The fix is that `total` is 32 (not 33), so the boundary sits cleanly + // at the real midpoint. + let set_syncs = |n: usize| { + let mut unit = vec![0u8; ALIGNED_UNIT_LEN]; + let mut off = 4; + let mut placed = 0; + while off < ALIGNED_UNIT_LEN && placed < n { + unit[off] = TS_SYNC; + off += TS_PACKET_LEN; + placed += 1; + } + unit + }; + + assert_eq!(ts_sync_count(&set_syncs(16)), 16); + assert_eq!(ts_sync_count(&set_syncs(17)), 17); + + // Exactly half intact → classified scrambled (16 > 16 is false). + assert!(is_aacs_scrambled(&set_syncs(16))); + // One past half → classified clear. + assert!(!is_aacs_scrambled(&set_syncs(17))); + } + + #[test] + fn scramble_detection_extremes() { + // Detection semantics for the clear-cut cases must be preserved: + // a fully-clear unit (all 32 syncs) is NOT scrambled; a unit with no + // syncs (fully scrambled body) IS scrambled. + let mut clear = vec![0u8; ALIGNED_UNIT_LEN]; + let mut off = 4; + while off < ALIGNED_UNIT_LEN { + clear[off] = TS_SYNC; + off += TS_PACKET_LEN; + } + assert_eq!(ts_sync_count(&clear), 32); + assert!( + !is_aacs_scrambled(&clear), + "fully-clear unit → not scrambled" + ); + + let scrambled = vec![0u8; ALIGNED_UNIT_LEN]; + assert_eq!(ts_sync_count(&scrambled), 0); + assert!(is_aacs_scrambled(&scrambled), "no syncs → scrambled"); + } + #[test] fn test_aes_cbc_roundtrip() { let key = [ @@ -388,6 +476,8 @@ mod tests { } off += TS_PACKET_LEN; } - assert_eq!(count, (ALIGNED_UNIT_LEN - 4) / TS_PACKET_LEN + 1); + // Assert against the single canonical packet count, not the old + // `(len - 4) / 192 + 1` form that `ts_packet_total` corrected away from. + assert_eq!(count, ts_packet_total(&unit)); } } diff --git a/src/aacs/handshake.rs b/src/aacs/handshake.rs index 83b613d..14e8109 100644 --- a/src/aacs/handshake.rs +++ b/src/aacs/handshake.rs @@ -49,7 +49,6 @@ const EC_A: [u8; 20] = [ 0x9D, 0xC9, 0xD8, 0x13, 0x55, 0xEC, 0xCE, 0xB5, 0x60, 0xBD, 0xB0, 0x9E, 0xF9, 0xEA, 0xE7, 0xC4, 0x79, 0xA7, 0xD7, 0xDC, ]; -#[cfg(test)] const EC_B: [u8; 20] = [ 0x40, 0x2D, 0xAD, 0x3E, 0xC1, 0xCB, 0xCD, 0x16, 0x52, 0x48, 0xD6, 0x8E, 0x12, 0x45, 0xE0, 0xC4, 0xDA, 0xAC, 0xB1, 0xD8, @@ -77,7 +76,6 @@ const P256_A: [u8; 32] = [ 0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFC, ]; -#[cfg(test)] const P256_B: [u8; 32] = [ 0x5A, 0xC6, 0x35, 0xD8, 0xAA, 0x3A, 0x93, 0xE7, 0xB3, 0xEB, 0xBD, 0x55, 0x76, 0x98, 0x86, 0xBC, 0x65, 0x1D, 0x06, 0xB0, 0xCC, 0x53, 0xB0, 0xF6, 0x3B, 0xCE, 0x3C, 0x3E, 0x27, 0xD2, 0x60, 0x4B, @@ -293,6 +291,20 @@ fn ec_double(pt: &EcPoint, a: &BigUint, p: &BigUint) -> EcPoint { } /// Scalar multiplication using double-and-add. +/// +/// NOTE (constant-time tradeoff): this branches on `scalar.bit(0)` and +/// clones BigUints per iteration, so its timing is data-dependent on the +/// secret scalar (the long-term host private key in `ecdsa_sign`, the +/// ephemeral key in ECDH). This is a deliberate tradeoff: the handshake +/// runs once per disc against a local optical drive, so throughput and +/// the narrow local-timing surface do not justify pulling in a vetted +/// constant-time backend. Revisit if this ever signs in a remote/shared +/// context. +/// +/// NOTE (cofactor): both AACS curves used here have cofactor 1, so a +/// point that lies on the curve is automatically in the prime-order +/// subgroup — no small-subgroup defense / `n·P == O` check is required +/// for the inputs this is called with. fn ec_mul(k: &BigUint, pt: &EcPoint, a: &BigUint, p: &BigUint) -> EcPoint { if k.is_zero() { return EcPoint::infinity(); @@ -313,6 +325,20 @@ fn ec_mul(k: &BigUint, pt: &EcPoint, a: &BigUint, p: &BigUint) -> EcPoint { result } +/// True if the point (x, y) satisfies y² ≡ x³ + ax + b (mod p) and lies +/// in the field (x, y < p). Guards the ECDH multiply against the classic +/// invalid-curve attack: a drive that supplies an off-curve key point can +/// otherwise steer the scalar multiply onto a weak curve and leak the host +/// scalar. Caller must reject the point when this returns false. +fn point_on_curve(x: &BigUint, y: &BigUint, a: &BigUint, b: &BigUint, p: &BigUint) -> bool { + if x >= p || y >= p { + return false; + } + let lhs = (y * y) % p; + let rhs = (((x * x) % p) * x + a * x + b) % p; + lhs == rhs +} + /// Convert BigUint to fixed-size big-endian bytes, zero-padded. fn to_bytes_be_padded(n: &BigUint, len: usize) -> Vec { let bytes = n.to_bytes_be(); @@ -341,12 +367,15 @@ fn ecdsa_sign(priv_key: &[u8; 20], data: &[u8]) -> ([u8; 20], [u8; 20]) { let z = BigUint::from_bytes_be(&hash); loop { - // Generate random k + // Generate random k via rejection sampling. Reducing raw RNG bytes + // modulo n would bias k toward small values (n is not a power of + // two); a biased ECDSA nonce is a known key-recovery weakness, so + // we reject and redraw any candidate >= n instead. let mut k_bytes = [0u8; 20]; use rand::RngCore; rand::thread_rng().fill_bytes(&mut k_bytes); - let k = BigUint::from_bytes_be(&k_bytes) % &n; - if k.is_zero() { + let k = BigUint::from_bytes_be(&k_bytes); + if k.is_zero() || k >= n { continue; } @@ -438,11 +467,14 @@ fn ecdsa_sign_p256(priv_key: &[u8; 32], data: &[u8]) -> ([u8; 32], [u8; 32]) { let z = BigUint::from_bytes_be(&hash); loop { + // Rejection sampling for the nonce — see ecdsa_sign for rationale + // (avoid the modulo bias that reducing raw RNG bytes mod n would + // introduce). let mut k_bytes = [0u8; 32]; use rand::RngCore; rand::thread_rng().fill_bytes(&mut k_bytes); - let k = BigUint::from_bytes_be(&k_bytes) % &n; - if k.is_zero() { + let k = BigUint::from_bytes_be(&k_bytes); + if k.is_zero() || k >= n { continue; } @@ -512,28 +544,38 @@ fn ecdsa_verify_p256(pub_x: &[u8], pub_y: &[u8], sig_r: &[u8], sig_s: &[u8], dat &r_point.x % &n == r } -/// Verify an AACS 2.0 certificate (type 0x11, 132 bytes) against AACS 2.0 LA key. +/// Verify an AACS 2.0 certificate (type 0x11) against the AACS 2.0 LA key. +/// +/// Layout: type(1) + flags(1) + padding(2) + serial(6) + pub_x(32) + +/// pub_y(32) + sig_r(32) + sig_s(32) = 138 bytes. The signature covers +/// the first 74 bytes (everything up to and including the public key). +/// +/// The full P-256 certificate is 138 bytes, so the entire 138-byte +/// length must be present before any signature slice is taken — checking +/// `>= 138` up front (rather than the old `>= 132`, which left the +/// `cert[106..138]` slice able to panic on a 132-byte input) keeps this +/// safe against the truncated 132-byte cert the handshake actually +/// passes in (`&response[24..156]`). fn verify_cert_p256(cert: &[u8]) -> bool { - if cert.len() < 132 { + if cert.len() < 138 { return false; } - // AACS 2.0 cert: type(1) + flags(1) + padding(2) + serial(6) + pub_x(32) + pub_y(32) + sig_r(32) + sig_s(32) - // Signature is over the first 74 bytes let sig_r = &cert[74..106]; - let sig_s = &cert[106..138]; // some certs may be padded differently - - // Use what we have — verify over the signed portion - if cert.len() >= 138 { - ecdsa_verify_p256(&AACS2_LA_PUB_X, &AACS2_LA_PUB_Y, sig_r, sig_s, &cert[..74]) - } else { - false - } + let sig_s = &cert[106..138]; + ecdsa_verify_p256(&AACS2_LA_PUB_X, &AACS2_LA_PUB_Y, sig_r, sig_s, &cert[..74]) } /// Extract public key from an AACS 2.0 certificate (32-byte x,y). +/// +/// Returns a zeroed key pair if `cert` is too short to hold the fixed +/// offsets (matches the `>= 138` guard in `verify_cert_p256`), so a +/// short/hostile cert cannot panic on the slice index. fn cert_pub_key_p256(cert: &[u8]) -> ([u8; 32], [u8; 32]) { let mut x = [0u8; 32]; let mut y = [0u8; 32]; + if cert.len() < 74 { + return (x, y); + } x.copy_from_slice(&cert[10..42]); y.copy_from_slice(&cert[42..74]); (x, y) @@ -544,15 +586,20 @@ fn compute_bus_key_p256( host_priv: &[u8; 32], drive_key_point_x: &[u8], drive_key_point_y: &[u8], -) -> [u8; 16] { +) -> Option<[u8; 16]> { let p = BigUint::from_bytes_be(&P256_P); let a = BigUint::from_bytes_be(&P256_A); + let b = BigUint::from_bytes_be(&P256_B); let d = BigUint::from_bytes_be(host_priv); - let dkp = EcPoint::new( - BigUint::from_bytes_be(drive_key_point_x), - BigUint::from_bytes_be(drive_key_point_y), - ); + let dx = BigUint::from_bytes_be(drive_key_point_x); + let dy = BigUint::from_bytes_be(drive_key_point_y); + + // Reject an off-curve drive point before the multiply (invalid-curve attack). + if !point_on_curve(&dx, &dy, &a, &b, &p) { + return None; + } + let dkp = EcPoint::new(dx, dy); let shared = ec_mul(&d, &dkp, &a, &p); @@ -560,7 +607,7 @@ fn compute_bus_key_p256( let x_bytes = to_bytes_be_padded(&shared.x, 32); let mut bus_key = [0u8; 16]; bus_key.copy_from_slice(&x_bytes[16..32]); - bus_key + Some(bus_key) } // ── AACS certificate handling ─────────────────────────────────────────────── @@ -581,9 +628,16 @@ fn verify_cert(cert: &[u8]) -> bool { } /// Extract public key from certificate. +/// +/// Returns a zeroed key pair if `cert` is too short to hold the fixed +/// offsets (matches the `>= 92` guard in `verify_cert`), so a +/// short/hostile cert cannot panic on the slice index. fn cert_pub_key(cert: &[u8]) -> ([u8; 20], [u8; 20]) { let mut x = [0u8; 20]; let mut y = [0u8; 20]; + if cert.len() < 52 { + return (x, y); + } x.copy_from_slice(&cert[12..32]); y.copy_from_slice(&cert[32..52]); (x, y) @@ -596,12 +650,20 @@ fn compute_bus_key( host_priv: &[u8; 20], drive_key_point_x: &[u8; 20], drive_key_point_y: &[u8; 20], -) -> [u8; 16] { +) -> Option<[u8; 16]> { let p = BigUint::from_bytes_be(&EC_P); let a = BigUint::from_bytes_be(&EC_A); + let b = BigUint::from_bytes_be(&EC_B); let d = BigUint::from_bytes_be(host_priv); - let dkp = EcPoint::from_bytes(drive_key_point_x, drive_key_point_y); + let dx = BigUint::from_bytes_be(drive_key_point_x); + let dy = BigUint::from_bytes_be(drive_key_point_y); + + // Reject an off-curve drive point before the multiply (invalid-curve attack). + if !point_on_curve(&dx, &dy, &a, &b, &p) { + return None; + } + let dkp = EcPoint::new(dx, dy); let shared = ec_mul(&d, &dkp, &a, &p); @@ -609,7 +671,7 @@ fn compute_bus_key( let x_bytes = to_bytes_be_padded(&shared.x, 20); let mut bus_key = [0u8; 16]; bus_key.copy_from_slice(&x_bytes[4..20]); // last 16 of 20 - bus_key + Some(bus_key) } /// Generate ephemeral host key pair: (private_key, public_point_x, public_point_y). @@ -620,12 +682,20 @@ fn generate_host_key_pair_p256() -> ([u8; 32], [u8; 32], [u8; 32]) { let n = BigUint::from_bytes_be(&P256_N); let g = EcPoint::from_bytes(&P256_GX, &P256_GY); - let mut priv_bytes = [0u8; 32]; - use rand::RngCore; - rand::thread_rng().fill_bytes(&mut priv_bytes); - let d = BigUint::from_bytes_be(&priv_bytes) % &n; - - let q = ec_mul(&d, &g, &a, &p_mod); + let (d, q) = loop { + let mut priv_bytes = [0u8; 32]; + use rand::RngCore; + rand::thread_rng().fill_bytes(&mut priv_bytes); + // d == 0 (prob ~1/n) would yield the point at infinity / an + // all-zero key and degenerate the bus key — reject and retry, + // matching the AACS 1.0 sibling generate_host_key_pair. + let d = BigUint::from_bytes_be(&priv_bytes) % &n; + if d.is_zero() { + continue; + } + let q = ec_mul(&d, &g, &a, &p_mod); + break (d, q); + }; let mut key = [0u8; 32]; let mut pub_x = [0u8; 32]; @@ -672,7 +742,14 @@ fn generate_host_key_pair() -> ([u8; 20], [u8; 20], [u8; 20]) { // ── AES-CMAC (for MAC verification) ──────────────────────────────────────── -/// AES-128-CMAC over 16 bytes of data. +/// AES-128-CMAC, single-complete-block case ONLY. +/// +/// Implements just the exactly-16-byte message path: it derives subkey +/// K1 and XORs the one full block. It does NOT derive K2 or apply the +/// `0x80` 10*-padding, so it is correct only for a 16-byte input — the +/// `&[u8; 16]` signature enforces that at compile time. Do NOT generalize +/// this to multi-block or short-final-block messages without adding K2 + +/// padding. fn aes_cmac_16(data: &[u8; 16], key: &[u8; 16]) -> [u8; 16] { use aes::Aes128; use aes::cipher::{BlockEncrypt, KeyInit, generic_array::GenericArray}; @@ -746,7 +823,10 @@ fn cdb_report_disc_structure(agid: u8, format: u8, len: u16) -> [u8; 12] { // ── High-level handshake ──────────────────────────────────────────────────── /// Result of a successful AACS authentication handshake. -#[derive(Debug)] +/// +/// `Debug` is implemented manually so the session key material +/// (`bus_key`, `volume_id`, `read_data_key`) is never rendered into logs +/// or `dbg!` output — only its presence is reported. pub struct AacsAuth { /// Bus key (16 bytes) — derived from ECDH pub bus_key: [u8; 16], @@ -756,10 +836,27 @@ pub struct AacsAuth { pub volume_id: Option<[u8; 16]>, /// Read data key (16 bytes) — for AACS 2.0 bus decryption pub read_data_key: Option<[u8; 16]>, - /// Drive certificate (92 bytes) + /// Drive certificate (first 92 bytes of the drive's certificate; + /// an AACS 2.0 P-256 cert is 132 bytes and is truncated to fit this + /// fixed-size field — see [`aacs2_authenticate_p256`]). pub drive_cert: [u8; 92], } +// Manual Debug: bus_key, volume_id, and read_data_key are key material (the +// VID feeds VUK derivation), so they are redacted — a `dbg!`/tracing of +// AacsAuth must never dump them in plaintext. +impl std::fmt::Debug for AacsAuth { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AacsAuth") + .field("bus_key", &"[redacted]") + .field("agid", &self.agid) + .field("volume_id", &self.volume_id.map(|_| "[redacted]")) + .field("read_data_key", &self.read_data_key.map(|_| "[redacted]")) + .field("drive_cert", &self.drive_cert) + .finish() + } +} + /// Perform the full AACS authentication handshake. /// /// Requires a host private key (20 bytes) and host certificate (92 bytes) @@ -873,7 +970,7 @@ pub fn aacs_authenticate( dkp_x.copy_from_slice(&drive_key_point[..20]); dkp_y.copy_from_slice(&drive_key_point[20..40]); - let bus_key = compute_bus_key(&host_key, &dkp_x, &dkp_y); + let bus_key = compute_bus_key(&host_key, &dkp_x, &dkp_y).ok_or(Error::AacsKeyVerify)?; Ok(AacsAuth { bus_key, @@ -904,9 +1001,11 @@ pub fn aacs2_authenticate( } } - // AACS 2.0 native P-256 handshake - let host_priv_v2 = host_priv_key_v2.ok_or(Error::AacsCertShort)?; - let host_cert_v2 = host_cert_v2.ok_or(Error::AacsCertShort)?; + // AACS 2.0 native P-256 handshake. Absent v2 credentials are "no AACS + // 2.0 keys configured" (AacsNoKeys), distinct from a malformed/too-short + // cert (AacsCertShort) — so callers can tell "not provided" from "bad". + let host_priv_v2 = host_priv_key_v2.ok_or(Error::AacsNoKeys)?; + let host_cert_v2 = host_cert_v2.ok_or(Error::AacsNoKeys)?; aacs2_authenticate_p256(session, host_priv_v2, host_cert_v2) } @@ -963,8 +1062,14 @@ fn aacs2_authenticate_p256( // uses certificate formats that differ from the spec, and rejecting them // would break otherwise working drives. The drive is still authenticated // through the ECDH key exchange and P-256 signature verification below. + // The outcome is surfaced as a trace event rather than discarded so the + // trust decision is observable (and so the call is not dead code). if drive_cert[0] == 0x11 && !verify_cert_p256(drive_cert) { - // Certificate verification failed but proceeding for backward compatibility. + tracing::debug!( + target: "freemkv::disc", + phase = "aacs2_cert_verify_skipped", + "drive cert failed P-256 LA verification; proceeding for backward compat" + ); } // Step 6: Read drive key point + signature (P-256: 64+64 = 128 bytes) @@ -1013,7 +1118,8 @@ fn aacs2_authenticate_p256( scsi_write(session, &cdb, &send_buf).map_err(|_| Error::AacsKeyRejected)?; // Step 9: Compute bus key via P-256 ECDH - let bus_key = compute_bus_key_p256(&host_eph_key, drive_key_x, drive_key_y); + let bus_key = compute_bus_key_p256(&host_eph_key, drive_key_x, drive_key_y) + .ok_or(Error::AacsKeyVerify)?; Ok(AacsAuth { bus_key, @@ -1142,9 +1248,11 @@ mod tests { let (priv_b, pub_bx, pub_by) = generate_host_key_pair(); // A computes: priv_a × pub_B - let shared_a = compute_bus_key(&priv_a, &pub_bx, &pub_by); + let shared_a = compute_bus_key(&priv_a, &pub_bx, &pub_by) + .expect("on-curve generated point must be accepted"); // B computes: priv_b × pub_A - let shared_b = compute_bus_key(&priv_b, &pub_ax, &pub_ay); + let shared_b = compute_bus_key(&priv_b, &pub_ax, &pub_ay) + .expect("on-curve generated point must be accepted"); assert_eq!(shared_a, shared_b, "ECDH shared secrets should match"); } @@ -1224,12 +1332,14 @@ mod tests { &priv_a, &to_bytes_be_padded(&pub_b.x, 32), &to_bytes_be_padded(&pub_b.y, 32), - ); + ) + .expect("on-curve generated point must be accepted"); let key_b = compute_bus_key_p256( &priv_b, &to_bytes_be_padded(&pub_a.x, 32), &to_bytes_be_padded(&pub_a.y, 32), - ); + ) + .expect("on-curve generated point must be accepted"); assert_eq!(key_a, key_b, "P-256 ECDH shared secrets should match"); } @@ -1335,6 +1445,62 @@ mod tests { assert_ne!(calc_mac, [0u8; 16], "real CMAC must not be all zeros"); } + #[test] + fn test_verify_cert_p256_short_cert_no_panic() { + // Regression: verify_cert_p256 used to slice cert[106..138] after only + // a `len < 132` guard. The drive cert the handshake passes in is + // exactly 132 bytes (&response[24..156]), so the slice panicked OOB. + // It must now return false (cannot verify) rather than panic. + let cert_132 = [0x11u8; 132]; + assert!( + !verify_cert_p256(&cert_132), + "132-byte cert must be rejected, not panic" + ); + // Boundary lengths around the slice requirement. + for len in [0usize, 73, 74, 105, 106, 131, 137] { + let cert = vec![0x11u8; len]; + assert!(!verify_cert_p256(&cert), "len {len} must not panic"); + } + } + + #[test] + fn test_compute_bus_key_rejects_off_curve_point() { + // An off-curve drive point must be rejected (invalid-curve guard), + // while an on-curve point (here the generator G) is accepted. + let (host_priv, _, _) = generate_host_key_pair(); + + // On-curve: G itself. + assert!( + compute_bus_key(&host_priv, &EC_GX, &EC_GY).is_some(), + "on-curve point must be accepted" + ); + + // Off-curve: G with y flipped by one bit almost never stays on the curve. + let mut bad_y = EC_GY; + bad_y[19] ^= 0x01; + assert!( + compute_bus_key(&host_priv, &EC_GX, &bad_y).is_none(), + "off-curve point must be rejected" + ); + } + + #[test] + fn test_compute_bus_key_p256_rejects_off_curve_point() { + let (host_priv, _, _) = generate_host_key_pair_p256(); + + assert!( + compute_bus_key_p256(&host_priv, &P256_GX, &P256_GY).is_some(), + "on-curve P-256 point must be accepted" + ); + + let mut bad_y = P256_GY; + bad_y[31] ^= 0x01; + assert!( + compute_bus_key_p256(&host_priv, &P256_GX, &bad_y).is_none(), + "off-curve P-256 point must be rejected" + ); + } + #[test] fn test_verify_host_cert_from_keydb() { // Verify the host cert from our KEYDB diff --git a/src/aacs/keydb.rs b/src/aacs/keydb.rs index c8f0785..6c61115 100644 --- a/src/aacs/keydb.rs +++ b/src/aacs/keydb.rs @@ -55,14 +55,22 @@ pub struct DiscEntry { } /// Parse a hex string like "0xABCD..." into bytes. +/// +/// Operates on bytes, not `&str` char boundaries: the keydb is +/// third-party content, so a non-ASCII scalar (e.g. a 4-byte UTF-8 +/// codepoint) must not panic on a mid-codepoint slice. Any non-hex +/// byte yields `None`. pub(crate) fn parse_hex(s: &str) -> Option> { let s = s.trim().trim_start_matches("0x").trim_start_matches("0X"); - if s.len() % 2 != 0 { + let bytes = s.as_bytes(); + if bytes.len() % 2 != 0 { return None; } - let mut out = Vec::with_capacity(s.len() / 2); - for i in (0..s.len()).step_by(2) { - out.push(u8::from_str_radix(&s[i..i + 2], 16).ok()?); + let mut out = Vec::with_capacity(bytes.len() / 2); + for pair in bytes.chunks_exact(2) { + let hi = (pair[0] as char).to_digit(16)?; + let lo = (pair[1] as char).to_digit(16)?; + out.push((hi * 16 + lo) as u8); } Some(out) } @@ -142,12 +150,26 @@ impl KeyDb { continue; } - // Host Certificate (AACS 2.0) + // Host Certificate (AACS 2.0). + // + // An HC2 row normally augments the preceding HC (AACS 1.0) row. + // KEYDB line ordering is third-party, so an HC2 row may appear + // before any HC row; rather than silently dropping the AACS 2.0 + // credentials, carry them on a fresh HostCert with an empty v1 + // cert (the v1 private_key/certificate stay zero/empty and are + // ignored by the v1 handshake, which guards on cert length). if line.starts_with("| HC2") { - if let Some(hc) = db.host_certs.last_mut() { - if let Some((pk, cert)) = Self::parse_host_cert_v2(line) { + if let Some((pk, cert)) = Self::parse_host_cert_v2(line) { + if let Some(hc) = db.host_certs.last_mut() { hc.private_key_v2 = Some(pk); hc.certificate_v2 = Some(cert); + } else { + db.host_certs.push(HostCert { + private_key: [0u8; 20], + certificate: Vec::new(), + private_key_v2: Some(pk), + certificate_v2: Some(cert), + }); } } continue; @@ -173,8 +195,18 @@ impl KeyDb { } /// Load a KEYDB.cfg from disk. - pub fn load(path: &std::path::Path) -> std::io::Result { - let data = std::fs::read_to_string(path)?; + /// + /// A read failure (missing/unreadable file, non-UTF-8 content) surfaces + /// as [`crate::error::Error::KeydbLoad`] carrying the path, per the + /// library contract that a missing/unparseable keydb is a structured + /// error and not a raw `io::Error`. Note that [`Self::parse`] itself is + /// lenient: a syntactically valid but key-less file parses to an empty + /// [`KeyDb`] rather than an error — callers needing a non-empty db must + /// check the parsed contents. + pub fn load(path: &std::path::Path) -> crate::error::Result { + let data = std::fs::read_to_string(path).map_err(|_| crate::error::Error::KeydbLoad { + path: path.display().to_string(), + })?; Ok(Self::parse(&data)) } @@ -234,10 +266,14 @@ impl super::provider::KeyProvider for KeyDb { self.host_certs.clone() } fn lookup_disc_by_hash(&self, disc_hash: &[u8; 20]) -> Option { + use std::fmt::Write; + // Lowercase hex written straight into the pre-sized buffer: find_disc + // lowercases its input anyway, so emitting 'x' here avoids a wasted + // to_lowercase() round-trip, and write! avoids 20 temporary Strings. let mut hex = String::with_capacity(42); hex.push_str("0x"); for b in disc_hash { - hex.push_str(&format!("{b:02X}")); + let _ = write!(hex, "{b:02x}"); } self.find_disc(&hex).cloned() } @@ -324,9 +360,17 @@ impl KeyDb { .next()? .trim(); + let certificate = parse_hex(cert_str)?; + // AACS 1.0 host certs are 92 bytes; drop malformed/short rows at + // parse time so the handshake never attempts junk (mirrors the v2 + // path, which enforces >= 132). + if certificate.len() < 92 { + return None; + } + Some(HostCert { private_key: parse_hex20(priv_str)?, - certificate: parse_hex(cert_str)?, + certificate, private_key_v2: None, certificate_v2: None, }) @@ -371,15 +415,16 @@ impl KeyDb { // Extract title (before first |) let title_part = rest.split(" | ").next().unwrap_or("").trim(); - // Clean title: "TITLE_NAME (Display Title)" → use display title if present - let title = if let Some(start) = title_part.find('(') { - if let Some(end) = title_part.rfind(')') { - title_part[start + 1..end].to_string() - } else { - title_part.to_string() - } - } else { - title_part.to_string() + // Clean title: "TITLE_NAME (Display Title)" → use display title if + // present. keydb.cfg is untrusted third-party content, so a title with + // ')' before '(' (e.g. "FILM) (X") would make start+1 > end; guard the + // slice and fall back to the whole title. + let title = match (title_part.find('('), title_part.rfind(')')) { + (Some(start), Some(end)) => title_part + .get(start + 1..end) + .map(str::to_string) + .unwrap_or_else(|| title_part.to_string()), + _ => title_part.to_string(), }; // Parse fields by tag @@ -536,6 +581,69 @@ mod tests { assert_eq!(hc.certificate.len(), 92); } + #[test] + fn test_parse_hex_rejects_non_ascii_without_panic() { + // A 4-byte UTF-8 scalar has byte-len 4 (passes the even check); the + // old &str-slice path panicked on the mid-codepoint boundary. The + // byte-wise parser must instead return None. + assert!(parse_hex("😀").is_none()); + // Mixed: leading hex then a 2-byte UTF-8 scalar (byte-len even). + assert!(parse_hex("ABé").is_none()); + // Sanity: well-formed hex still parses. + assert_eq!(parse_hex("0x00FF"), Some(vec![0x00, 0xFF])); + // Odd byte length still rejected. + assert!(parse_hex("ABC").is_none()); + } + + #[test] + fn test_hc2_before_hc_is_not_dropped() { + // An HC2 row appearing before any HC row must still land its AACS 2.0 + // credentials on a HostCert rather than being silently discarded. + let cfg = format!( + "| HC2 | HOST_PRIV_KEY 0x{} | HOST_CERT 0x{}\n", + "00".repeat(32), + "00".repeat(132) + ); + let db = KeyDb::parse(&cfg); + assert_eq!( + db.host_certs.len(), + 1, + "HC2-only row must create a HostCert" + ); + assert!(db.host_certs[0].private_key_v2.is_some()); + assert!(db.host_certs[0].certificate_v2.is_some()); + assert!( + db.host_certs[0].certificate.is_empty(), + "v1 cert stays empty for an HC2-only carrier" + ); + } + + #[test] + fn test_hc2_after_hc_augments_existing() { + let cfg = format!( + "| HC | HOST_PRIV_KEY 0x{} | HOST_CERT 0x{}\n| HC2 | HOST_PRIV_KEY 0x{} | HOST_CERT 0x{}\n", + "00".repeat(20), + "00".repeat(92), + "00".repeat(32), + "00".repeat(132) + ); + let db = KeyDb::parse(&cfg); + assert_eq!(db.host_certs.len(), 1, "HC2 augments the preceding HC"); + assert_eq!(db.host_certs[0].certificate.len(), 92); + assert!(db.host_certs[0].certificate_v2.is_some()); + } + + #[test] + fn test_parse_host_cert_rejects_short_v1_cert() { + // A too-short AACS 1.0 cert must be dropped at parse time. + let line = format!( + "| HC | HOST_PRIV_KEY 0x{} | HOST_CERT 0x{}", + "00".repeat(20), + "00".repeat(10) + ); + assert!(KeyDb::parse_host_cert(&line).is_none()); + } + #[test] fn test_parse_full_keydb() { let path = match keydb_path() { diff --git a/src/aacs/keys.rs b/src/aacs/keys.rs index 26e01b7..ddb8b89 100644 --- a/src/aacs/keys.rs +++ b/src/aacs/keys.rs @@ -227,10 +227,19 @@ pub fn derive_media_key_from_pk(mkb: &[u8], processing_keys: &[[u8; 16]]) -> Opt /// Set to 0 to disable walking (entries tried only as terminal PKs). const PK_WALK_MAX_DEPTH: u8 = 3; +/// Hard ceiling on the requested walk depth. The BFS frontier holds `2^depth` +/// 16-byte node keys, so an uncapped `max_depth` (e.g. 26+) would exhaust +/// memory; the walk silently clamps to this. 5 (32-wide frontier) covers every +/// realistic leaked-label case with margin. +const PK_WALK_MAX_DEPTH_CAP: u8 = 5; + /// Same as [`derive_media_key_from_pk`] but with explicit walk depth. /// Each entry is tried as a terminal PK at depth 0, then as a node-key /// whose PK and children are derived via `AES-G3(K, 0|1|2)` for up to /// `max_depth` additional levels. +/// +/// The BFS frontier grows as `2^max_depth`; `max_depth` is clamped to +/// [`PK_WALK_MAX_DEPTH_CAP`] so a large value cannot exhaust memory. pub fn derive_media_key_from_pk_walked( mkb: &[u8], processing_keys: &[[u8; 16]], @@ -252,6 +261,9 @@ fn walk_pk_against_tables_impl( mk_dv: &[u8; 16], max_depth: u8, ) -> Option<[u8; 16]> { + // Clamp the frontier depth (2^depth node keys) so a caller-supplied value + // cannot OOM the process. + let max_depth = max_depth.min(PK_WALK_MAX_DEPTH_CAP); let num_uvs = uvs .chunks(5) .take_while(|c| c.len() == 5 && (c[0] & 0xC0) == 0) @@ -415,7 +427,8 @@ pub mod probe { /// (16-byte entries); `mk_dv` is from the verify record. Each entry in /// `keys` is tried as a terminal PK and as an SD node-key descending via /// `AES-G3(K, 0|1|2)` for `max_depth` levels — identical logic to the - /// production walk. Returns the verified Media Key, if any. + /// production walk (`max_depth` is clamped to the same internal cap to + /// bound the `2^depth` frontier). Returns the verified Media Key, if any. pub fn walk_pk_against_tables( keys: &[[u8; 16]], subdiff: &[u8], @@ -446,7 +459,7 @@ fn mkb_find_mk_dv(mkb: &[u8]) -> Option<[u8; 16]> { // mk_dv is at offset 4 of the record (after the 4-byte header) let mut dv = [0u8; 16]; dv.copy_from_slice(&mkb[pos + 4..pos + 20]); - tracing::warn!( + tracing::debug!( target: "freemkv::disc", phase = "mkb_mk_dv_found", rec_type, @@ -524,9 +537,6 @@ fn find_record_body(mkb: &[u8], rec_type_wanted: u8) -> Option> { if rec_type == rec_type_wanted && rec_len > 4 { return Some(mkb[pos + 4..pos + rec_len].to_vec()); } - if rec_len == 0 { - break; - } pos += rec_len; } None @@ -613,7 +623,17 @@ fn calc_pk_from_dk(dk: &[u8; 16], uv: u32, v_mask: u32, dev_key_v_mask: u32) -> let mut right_child = aesg3(dk, 2); let mut current_v_mask = dev_key_v_mask; + // The subset-difference tree is at most 32 levels deep (u32 mask), so the + // walk must converge in <= 32 steps. The arithmetic `>> 1` sign-extends + // current_v_mask, so a v_mask coarser than dev_key_v_mask (reachable from + // a crafted/corrupt MKB) would otherwise saturate at 0xFFFF_FFFF and spin + // forever — bound the loop to keep a bad disc from hanging the rip thread. + let mut steps = 0u32; while current_v_mask != v_mask { + if steps >= 32 { + break; + } + steps += 1; // Find the highest unset bit in current_v_mask let mut bit_pos: i32 = -1; for i in (0..32).rev() { @@ -662,6 +682,13 @@ pub fn derive_media_key_from_dk(mkb: &[u8], device_keys: &[DeviceKey]) -> Option if u_mask_shift & 0xC0 != 0 { break; // device revoked } + // Shifts of 32..=63 (0x20..=0x3F pass the 0xC0 mask above) would + // panic in debug / wrap to a wrong mask in release. The MKB byte + // is disc-controlled, so a crafted/corrupt MKB must not crash the + // ripper: skip an out-of-range slot rather than `<<` it. + if u_mask_shift >= 32 { + continue; + } let uv = u32::from_be_bytes([p_uv[0], p_uv[1], p_uv[2], p_uv[3]]); if uv == 0 { @@ -674,7 +701,12 @@ pub fn derive_media_key_from_dk(mkb: &[u8], device_keys: &[DeviceKey]) -> Option if ((device_number & u_mask) == (uv & u_mask)) && ((device_number & v_mask) != (uv & v_mask)) { - // Found matching subset-difference — find the right device key + // Found matching subset-difference — find the right device key. + // dk.u_mask_shift is a u8 from keydb with no range check; + // guard the shift the same way as the MKB byte above. + if dk.u_mask_shift >= 32 { + continue; + } let dev_key_v_mask = calc_v_mask(dk.uv); let dev_key_u_mask: u32 = 0xFFFF_FFFF << dk.u_mask_shift; @@ -928,7 +960,7 @@ pub fn resolve_keys_v21(ctx: &ResolveContext<'_>) -> Option { } }; - tracing::warn!( + tracing::info!( target: "freemkv::disc", phase = "resolve_keys_v21_start", bus_encryption, @@ -953,7 +985,7 @@ pub fn resolve_keys_v21(ctx: &ResolveContext<'_>) -> Option { ctx.volume_id, ) { Ok((_km, kvu)) => { - tracing::warn!( + tracing::debug!( target: "freemkv::disc", phase = "resolve_keys_v21_path1_hit", "Variant chain produced Km + Kvu" @@ -961,7 +993,7 @@ pub fn resolve_keys_v21(ctx: &ResolveContext<'_>) -> Option { return Some(build(Some(kvu), derive_uks(&kvu), 1)); } Err(e) => { - tracing::warn!( + tracing::debug!( target: "freemkv::disc", phase = "resolve_keys_v21_path1_miss", error_code = %e, @@ -974,14 +1006,18 @@ pub fn resolve_keys_v21(ctx: &ResolveContext<'_>) -> Option { // Path 3: pre-computed MK + matching VID → derived VUK. // Short-circuit: first provider with a matching VID wins. if let Some(entry) = providers.lookup_disc_by_vid(ctx.volume_id) { - if let (Some(mk), Some(_)) = (entry.media_key, entry.disc_id) { + // The entry already matched by VID and derive_vuk needs only mk + + // ctx.volume_id, so a provider that matches by VID without + // populating disc_id (e.g. a webservice) must not have its MK + // dropped — gate on the MK alone. + if let Some(mk) = entry.media_key { let vuk = derive_vuk(&mk, ctx.volume_id); - tracing::warn!(target: "freemkv::disc", phase = "resolve_keys_v21_path3_hit", "MK+VID entry matched volume_id"); + tracing::debug!(target: "freemkv::disc", phase = "resolve_keys_v21_path3_hit", "MK+VID entry matched volume_id"); return Some(build(Some(vuk), derive_uks(&vuk), 3)); } } } else { - tracing::warn!( + tracing::debug!( target: "freemkv::disc", phase = "resolve_keys_v21_no_vid", "VID unavailable; paths 1/3 skipped" @@ -991,10 +1027,10 @@ pub fn resolve_keys_v21(ctx: &ResolveContext<'_>) -> Option { // Paths 4 and 5: hash lookup, prefer V over U on the same entry. if let Some(entry) = providers.lookup_disc_by_hash(&uk_file.disc_hash) { if let Some(vuk) = entry.vuk { - tracing::warn!(target: "freemkv::disc", phase = "resolve_keys_v21_path4_hit", "VUK from KEYDB"); + tracing::debug!(target: "freemkv::disc", phase = "resolve_keys_v21_path4_hit", "VUK from KEYDB"); return Some(build(Some(vuk), derive_uks(&vuk), 4)); } else if let Some(unit_keys) = match_keydb_unit_keys(&uk_file, &entry.unit_keys) { - tracing::warn!( + tracing::debug!( target: "freemkv::disc", phase = "resolve_keys_v21_path5_hit", uk_count = unit_keys.len(), @@ -1055,7 +1091,7 @@ fn resolve_keys_classical(ctx: &ResolveContext<'_>, version: AacsVersion) -> Opt } }; - tracing::warn!( + tracing::info!( target: "freemkv::disc", phase = "resolve_keys_start", version = ?version, @@ -1075,7 +1111,7 @@ fn resolve_keys_classical(ctx: &ResolveContext<'_>, version: AacsVersion) -> Opt let mk_dv = mkb_find_mk_dv(mkb); let subdiff = mkb_find_subdiff_records(mkb); let cvalues = mkb_find_cvalues(mkb); - tracing::warn!( + tracing::debug!( target: "freemkv::disc", phase = "resolve_keys_mkb_records", mk_dv_found = mk_dv.is_some(), @@ -1090,59 +1126,68 @@ fn resolve_keys_classical(ctx: &ResolveContext<'_>, version: AacsVersion) -> Opt let all_dks = providers.device_keys(); if let Some(mk) = derive_media_key_from_dk(mkb, &all_dks) { let vuk = derive_vuk(&mk, ctx.volume_id); - tracing::warn!(target: "freemkv::disc", phase = "resolve_keys_path1_hit", "media key derived from device key"); + tracing::debug!(target: "freemkv::disc", phase = "resolve_keys_path1_hit", "media key derived from device key"); return Some(build(Some(vuk), derive_uks(&vuk), 1)); } - tracing::warn!(target: "freemkv::disc", phase = "resolve_keys_path1_miss", dk_count = all_dks.len(), "DK derivation failed"); + tracing::debug!(target: "freemkv::disc", phase = "resolve_keys_path1_miss", dk_count = all_dks.len(), "DK derivation failed"); // Path 2: MKB + processing keys → media key → VUK let all_pks = providers.processing_keys(); if let Some(mk) = derive_media_key_from_pk(mkb, &all_pks) { let vuk = derive_vuk(&mk, ctx.volume_id); - tracing::warn!(target: "freemkv::disc", phase = "resolve_keys_path2_hit", "media key derived from processing key"); + tracing::debug!(target: "freemkv::disc", phase = "resolve_keys_path2_hit", "media key derived from processing key"); return Some(build(Some(vuk), derive_uks(&vuk), 2)); } - tracing::warn!(target: "freemkv::disc", phase = "resolve_keys_path2_miss", pk_count = all_pks.len(), "PK derivation failed"); + tracing::debug!(target: "freemkv::disc", phase = "resolve_keys_path2_miss", pk_count = all_pks.len(), "PK derivation failed"); // Path 2.5: MK-pool brute. keydb stores Media Keys per-disc, but an // MK is MKB-scoped (shared across a pressing/MKB-family). A disc // whose own hash/VID isn't keyed can still resolve if ANY stored MK // verifies against its MKB. Try every distinct MK via km_verifies; // a UNIQUE pass is this disc's Km → derive VUK (needs VID) → UK. - // km_verifies is one AES-D + magic check per candidate (cheap). + // One AES-D + magic check per candidate (cheap). mk_dv is hoisted + // out of the loop so the MKB is not re-walked per candidate. let mks = providers.media_keys(); let mut mk_hits: Vec<[u8; 16]> = Vec::new(); - for mk in &mks { - if probe::km_verifies(mkb, mk) && !mk_hits.contains(mk) { - mk_hits.push(*mk); - if mk_hits.len() > 1 { - break; // ambiguous — bail to avoid a wrong key + if let Some(mk_dv) = mkb_find_mk_dv(mkb) { + for mk in &mks { + let verifies = aes_ecb_decrypt(mk, &mk_dv)[..8] + == [0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF]; + if verifies && !mk_hits.contains(mk) { + mk_hits.push(*mk); + if mk_hits.len() > 1 { + break; // ambiguous — bail to avoid a wrong key + } } } } if mk_hits.len() == 1 { let vuk = derive_vuk(&mk_hits[0], ctx.volume_id); - tracing::warn!(target: "freemkv::disc", phase = "resolve_keys_path2_5_hit", mk_pool = mks.len(), "media key from keydb MK-pool brute (km_verifies)"); + tracing::debug!(target: "freemkv::disc", phase = "resolve_keys_path2_5_hit", mk_pool = mks.len(), "media key from keydb MK-pool brute (km_verifies)"); // Same class as path 3 (KEYDB MK → derived VUK). return Some(build(Some(vuk), derive_uks(&vuk), 3)); } - tracing::warn!(target: "freemkv::disc", phase = "resolve_keys_path2_5_miss", mk_pool = mks.len(), mk_hits = mk_hits.len(), "MK-pool brute: no unique verifying MK"); + tracing::debug!(target: "freemkv::disc", phase = "resolve_keys_path2_5_miss", mk_pool = mks.len(), mk_hits = mk_hits.len(), "MK-pool brute: no unique verifying MK"); } else { - tracing::warn!(target: "freemkv::disc", phase = "resolve_keys_no_mkb", "no MKB; paths 1/2 skipped"); + tracing::debug!(target: "freemkv::disc", phase = "resolve_keys_no_mkb", "no MKB; paths 1/2 skipped"); } // Path 3: pre-computed MK + matching VID → derived VUK. // Short-circuit: first provider with a matching VID wins. if let Some(entry) = providers.lookup_disc_by_vid(ctx.volume_id) { - if let (Some(mk), Some(_)) = (entry.media_key, entry.disc_id) { + // The entry already matched by VID and derive_vuk needs only mk + + // ctx.volume_id, so a provider that matches by VID without + // populating disc_id (e.g. a webservice) must not have its MK + // dropped — gate on the MK alone. + if let Some(mk) = entry.media_key { let vuk = derive_vuk(&mk, ctx.volume_id); - tracing::warn!(target: "freemkv::disc", phase = "resolve_keys_path3_hit", "MK+VID entry matched volume_id"); + tracing::debug!(target: "freemkv::disc", phase = "resolve_keys_path3_hit", "MK+VID entry matched volume_id"); return Some(build(Some(vuk), derive_uks(&vuk), 3)); } } - tracing::warn!(target: "freemkv::disc", phase = "resolve_keys_path3_miss", "no MK+VID entry matched volume_id"); + tracing::debug!(target: "freemkv::disc", phase = "resolve_keys_path3_miss", "no MK+VID entry matched volume_id"); } else { - tracing::warn!( + tracing::debug!( target: "freemkv::disc", phase = "resolve_keys_no_vid", "VID unavailable; paths 1/2/3 require VID and are skipped" @@ -1153,12 +1198,12 @@ fn resolve_keys_classical(ctx: &ResolveContext<'_>, version: AacsVersion) -> Opt // U (path 5). They are not independent checks — path 5 only fires // because path 4 had no VUK on the same entry. if let Some(entry) = providers.lookup_disc_by_hash(&uk_file.disc_hash) { - tracing::warn!(target: "freemkv::disc", phase = "resolve_keys_keydb_hit_entry", "disc hash found in provider"); + tracing::debug!(target: "freemkv::disc", phase = "resolve_keys_keydb_hit_entry", "disc hash found in provider"); if let Some(vuk) = entry.vuk { - tracing::warn!(target: "freemkv::disc", phase = "resolve_keys_path4_hit", "VUK from provider"); + tracing::debug!(target: "freemkv::disc", phase = "resolve_keys_path4_hit", "VUK from provider"); return Some(build(Some(vuk), derive_uks(&vuk), 4)); } else if let Some(unit_keys) = match_keydb_unit_keys(&uk_file, &entry.unit_keys) { - tracing::warn!( + tracing::debug!( target: "freemkv::disc", phase = "resolve_keys_path5_hit", uk_count = unit_keys.len(), @@ -1206,6 +1251,36 @@ mod tests { if path.exists() { Some(path) } else { None } } + #[test] + fn derive_media_key_from_dk_survives_out_of_range_u_mask_shift() { + // Regression: a crafted/corrupt MKB with a Subset-Difference + // u_mask_shift of 32..=63 (passes the 0xC0 revoked-marker check but + // overflows `0xFFFF_FFFF << shift`) used to panic in debug / compute a + // wrong mask in release. The walk must now skip the bad slot and + // return cleanly (no panic) on disc-controlled bytes. + let mut mkb: Vec = Vec::new(); + // 0x81 record: 4-byte header + 16-byte mk_dv body (rec_len = 20). + mkb.extend_from_slice(&[0x81, 0x00, 0x00, 0x14]); + mkb.extend_from_slice(&[0xAB; 16]); + // 0x04 Subset-Difference: one 5-byte entry with u_mask_shift = 0x30 + // (48 — out of range, but 0x30 & 0xC0 == 0 so the revoke check passes). + mkb.extend_from_slice(&[0x04, 0x00, 0x00, 0x09]); + mkb.extend_from_slice(&[0x30, 0x00, 0x00, 0x00, 0x01]); + // 0x05 cvalues: one 16-byte entry (rec_len = 20). + mkb.extend_from_slice(&[0x05, 0x00, 0x00, 0x14]); + mkb.extend_from_slice(&[0xCD; 16]); + + let dk = DeviceKey { + key: [0x11; 16], + node: 1, + uv: 1, + u_mask_shift: 0x30, // also out of range on the device-key side + }; + + // Must not panic; no valid derivation is expected from this junk. + let _ = derive_media_key_from_dk(&mkb, &[dk]); + } + #[test] fn test_vuk_derivation() { // Pick any UHD entry with a known MK, VID, and VUK from KEYDB. @@ -1284,12 +1359,17 @@ mod tests { // Try decrypting a real encrypted aligned unit from a UHD sample. // This disc is AACS 2.0 (BEE) so unit key alone won't work — // we need bus decryption first. But this verifies the pipeline. - let unit_path = std::path::Path::new("/tmp/encrypted_unit.bin"); + // Path comes from ENCRYPTED_UNIT_PATH (same env-driven pattern as the + // KEYDB_PATH / MKB_SAMPLE_DIR fixtures); no-ops in CI when unset. + let unit_path = match std::env::var("ENCRYPTED_UNIT_PATH").ok() { + Some(p) => std::path::PathBuf::from(p), + None => return, + }; if !unit_path.exists() { return; } - let original = std::fs::read(unit_path).unwrap(); + let original = std::fs::read(&unit_path).unwrap(); assert_eq!(original.len(), ALIGNED_UNIT_LEN); assert!( super::super::decrypt::is_aacs_scrambled(&original), @@ -1316,10 +1396,10 @@ mod tests { let keys: Vec<[u8; 16]> = entry.unit_keys.iter().map(|(_, k)| *k).collect(); let mut unit = original.clone(); - if let Some(idx) = super::super::decrypt::decrypt_unit_try_keys(&mut unit, &keys) { + if let Some(res) = super::super::decrypt::decrypt_unit_try_keys(&mut unit, &keys) { eprintln!( - "SUCCESS: Decrypted with entry {} key {}", - entry.disc_hash, idx + "SUCCESS: Decrypted with entry {} ({res:?})", + entry.disc_hash ); // Count TS sync bytes let ts = (0..32).filter(|&i| unit[4 + i * 192] == 0x47).count(); diff --git a/src/aacs/mod.rs b/src/aacs/mod.rs index 251c669..6d61c3f 100644 --- a/src/aacs/mod.rs +++ b/src/aacs/mod.rs @@ -8,6 +8,7 @@ //! | DK | DEVICE_KEY 0x... | DEVICE_NODE 0x... | KEY_UV 0x... | KEY_U_MASK_SHIFT 0x... //! | PK | 0x... //! | HC | HOST_PRIV_KEY 0x... | HOST_CERT 0x... +//! | HC2 | HOST_PRIV_KEY 0x... | HOST_CERT 0x... //! 0x = | D | <date> | M | 0x<media_key> | I | 0x<disc_id> | V | 0x<vuk> | U | <unit_keys> //! //! The VUK decrypts title keys from AACS/Unit_Key_RO.inf on disc. @@ -23,20 +24,19 @@ pub mod variants; // Explicit re-exports — only items needed by external consumers and sibling crate modules. // AES primitives (aes_ecb_encrypt, aes_ecb_decrypt, aes_cbc_decrypt) are pub(crate) in decrypt.rs. pub use decrypt::{ - ALIGNED_UNIT_LEN, decrypt_bus, decrypt_unit, decrypt_unit_full, decrypt_unit_try_keys, - is_aacs_scrambled, unit_key_validates, + ALIGNED_UNIT_LEN, UnitKeyResult, decrypt_bus, decrypt_unit, decrypt_unit_full, + decrypt_unit_try_keys, is_aacs_scrambled, ts_packet_total, ts_sync_count, unit_key_validates, }; pub use keydb::{DeviceKey, DiscEntry, HostCert, KeyDb}; pub use keys::probe; pub use keys::{ AacsVersion, ContentCert, ResolveContext, ResolvedKeys, UnitKeyFile, decrypt_unit_key, - derive_media_key_from_dk, derive_media_key_from_pk, derive_vuk, disc_hash, disc_hash_hex, - mkb_content_len, mkb_version, parse_content_cert, parse_unit_key_ro, read_mkb_from_drive, - resolve_keys_v1, resolve_keys_v2, resolve_keys_v21, + derive_media_key_from_dk, derive_media_key_from_pk, derive_media_key_from_pk_walked, + derive_vuk, disc_hash, disc_hash_hex, mkb_content_len, mkb_version, parse_content_cert, + parse_unit_key_ro, read_mkb_from_drive, resolve_keys_v1, resolve_keys_v2, resolve_keys_v21, }; pub use provider::KeyProvider; pub use variants::{ KEY_CORRECTION_DATA_PLACEHOLDER, MediaKeyVariantError, MkbRecord, ProcessingKeyMatch, - derive_media_key_variant, is_variant_mkb, variant_data_record, variant_key_data, variant_nonce, - walk_mkb, walk_processing_key, + derive_media_key_variant, is_variant_mkb, variant_nonce, walk_mkb, walk_processing_key, }; diff --git a/src/aacs/provider.rs b/src/aacs/provider.rs index 251dc94..6ee2e68 100644 --- a/src/aacs/provider.rs +++ b/src/aacs/provider.rs @@ -7,17 +7,24 @@ //! Methods come in two flavors: //! //! - **Bulk material** ([`device_keys`], [`processing_keys`], -//! [`host_certs`]) — the resolver unions results across all -//! providers and tries each candidate. +//! [`media_keys`]) — the resolver unions (and dedups) results +//! across all providers and tries each candidate. //! - **Disc-keyed lookup** ([`lookup_disc_by_hash`], //! [`lookup_disc_by_vid`]) — the resolver short-circuits on the //! first hit, so providers are queried in array order with //! fastest/closest first. //! +//! [`host_certs`] is a sixth method but is NOT consumed by the +//! resolver chain: the SCSI handshake reads host certs directly from +//! the caller-supplied credentials, not from the provider array. A +//! provider that overrides `host_certs` today has no effect on the +//! handshake; the method is retained as a forward-looking extension +//! point only. +//! //! Default impls return empty / `None` so backends only override //! the methods they actually support — an external key service might //! implement only `lookup_disc_by_hash`, while a local file might -//! implement all five. +//! implement all six. //! //! Calls may block (disk I/O, network round-trips). The resolver //! invokes each method at most a handful of times per scan; for @@ -25,6 +32,7 @@ //! //! [`device_keys`]: KeyProvider::device_keys //! [`processing_keys`]: KeyProvider::processing_keys +//! [`media_keys`]: KeyProvider::media_keys //! [`host_certs`]: KeyProvider::host_certs //! [`lookup_disc_by_hash`]: KeyProvider::lookup_disc_by_hash //! [`lookup_disc_by_vid`]: KeyProvider::lookup_disc_by_vid @@ -58,6 +66,11 @@ pub trait KeyProvider: Send + Sync { /// AACS host certificates (with their private keys) for drive /// authentication. Multiple in case some are revoked. + /// + /// NOTE: not consumed by the resolver chain — the handshake reads + /// host certs from the caller-supplied credentials directly, so + /// overriding this method has no effect on drive authentication + /// today. Retained as a forward-looking extension point. fn host_certs(&self) -> Vec<HostCert> { Vec::new() } @@ -78,19 +91,28 @@ pub trait KeyProvider: Send + Sync { /// Resolver-side helpers that aggregate across a provider array. /// -/// The resolver consumes `&[&dyn KeyProvider]` directly; these -/// helpers wrap the union-vs-short-circuit policy per method. +/// The resolver wraps `ctx.providers` (`&[&dyn KeyProvider]`) in this +/// struct; these helpers apply the union-vs-short-circuit policy per +/// method. The bulk unions dedup so overlapping providers don't make +/// the resolver re-walk/re-validate identical material. pub(crate) struct Providers<'a>(pub &'a [&'a dyn KeyProvider]); impl Providers<'_> { - /// Union — gather DKs from every provider. + /// Union (deduped) — gather DKs from every provider. pub fn device_keys(&self) -> Vec<DeviceKey> { - self.0.iter().flat_map(|p| p.device_keys()).collect() + let mut v: Vec<DeviceKey> = self.0.iter().flat_map(|p| p.device_keys()).collect(); + // DeviceKey has no Ord/Hash; dedup on the value-defining tuple. + v.sort_unstable_by_key(|d| (d.key, d.node, d.uv, d.u_mask_shift)); + v.dedup_by_key(|d| (d.key, d.node, d.uv, d.u_mask_shift)); + v } - /// Union — gather PKs from every provider. + /// Union (deduped) — gather PKs from every provider. pub fn processing_keys(&self) -> Vec<[u8; 16]> { - self.0.iter().flat_map(|p| p.processing_keys()).collect() + let mut v: Vec<[u8; 16]> = self.0.iter().flat_map(|p| p.processing_keys()).collect(); + v.sort_unstable(); + v.dedup(); + v } /// Union of distinct Media Keys across every provider, for the MK-pool @@ -102,9 +124,9 @@ impl Providers<'_> { v } - /// Union — gather host certs from every provider. Not yet wired into - /// the SCSI handshake (which still reads `KeyDb.host_certs` directly); - /// kept here so a provider-aware handshake refactor is a drop-in. + /// Union — gather host certs from every provider. The SCSI handshake + /// reads host certs from the caller-supplied credentials directly and + /// does not call this, so it is currently unused by the resolver chain. #[allow(dead_code)] pub fn host_certs(&self) -> Vec<HostCert> { self.0.iter().flat_map(|p| p.host_certs()).collect() diff --git a/src/aacs/variants.rs b/src/aacs/variants.rs index f0a0a47..b762c41 100644 --- a/src/aacs/variants.rs +++ b/src/aacs/variants.rs @@ -14,6 +14,19 @@ //! a disc carries neither, callers should fall back to the classical //! single-stage derivation in [`super::keys`]. //! +//! **Status: the chain cannot yet produce a key on a real disc.** Two +//! sub-fields are unfinished: +//! - [`variants_for_uv`] (the `VARIANTS[uv]` lookup in the `0x83` +//! record) is a stub that always returns `None`, so the chain +//! short-circuits with [`MediaKeyVariantError::VariantsTableUnavailable`]. +//! - The Encrypted Media Key Variant Data (C) and the Variant Key +//! Data (VKD) table are *distinct* sub-fields of the `0x82` record +//! per AACS 2.1, but [`variant_data_record`] (C) and +//! [`variant_key_data`] (VKD) both currently return the *whole* +//! first `0x82` body — so on a single-`0x82` disc they alias. The +//! `0x82` sub-field offsets must be fixed against a real Variant +//! disc before this chain is wired into `resolve_keys`. +//! //! The chain follows the published spec: //! //! ```text @@ -28,6 +41,19 @@ //! Two condition bits on `Kmp[15]` route off the hardcoded-KCD path //! (Soft Correction and Online Challenge). The chain refuses to run in //! either case — callers must handle those modes out of band. +//! +//! # Status: Kp verification +//! +//! On the classical path [`walk_processing_key`] gates each match on +//! the VERIFY_MAGIC relation, which authenticates the Processing Key. +//! On a variant MKB that magic check does NOT hold (the walk yields a +//! Media Key *Precursor*, not the Media Key), so the walk accepts a +//! variant match without it. The replacement gate lives at the END of +//! [`derive_media_key_variant`]: the derived final `Km` is verified +//! against the MKB's Verify-Media-Key record before any `(Km, Kvu)` is +//! returned. A future implementer wiring [`variants_for_uv`] must keep +//! that final gate — the per-match magic check no longer protects the +//! variant path. use super::decrypt::aes_ecb_decrypt; use super::keydb::DeviceKey; @@ -94,7 +120,13 @@ pub fn is_variant_mkb(records: &[MkbRecord]) -> bool { } /// Body of the Encrypted Media Key Variant Data record (type `0x82`). -pub fn variant_data_record(records: &[MkbRecord]) -> Option<&[u8]> { +/// +/// Returns the whole first `0x82` body; the internal C / VKD sub-field +/// split is not yet decoded, so this aliases [`variant_key_data`] on a +/// single-`0x82` disc. `pub(crate)` until the sub-field offsets are fixed +/// against a real variant disc — it is not part of the public surface +/// because it knowingly returns an undecoded composite. +pub(crate) fn variant_data_record(records: &[MkbRecord]) -> Option<&[u8]> { records .iter() .find(|r| r.rec_type == 0x82) @@ -115,7 +147,11 @@ pub fn variant_nonce(records: &[MkbRecord]) -> Option<[u8; 16]> { /// Body of the Variant Key Data record. Returns the first `0x82` body /// that is a non-empty multiple of 16 bytes. -pub fn variant_key_data(records: &[MkbRecord]) -> Option<&[u8]> { +/// +/// Like [`variant_data_record`], this returns the whole `0x82` body and +/// aliases it on a single-`0x82` disc; the C / VKD sub-field split is +/// undecoded. `pub(crate)` until fixed against a real variant disc. +pub(crate) fn variant_key_data(records: &[MkbRecord]) -> Option<&[u8]> { records .iter() .find(|r| r.rec_type == 0x82 && !r.body.is_empty() && r.body.len() % 16 == 0) @@ -167,7 +203,17 @@ fn calc_pk_from_dk(dk: &[u8; 16], uv: u32, v_mask: u32, dev_key_v_mask: u32) -> let mut right_child = aesg3_step(dk, 2); let mut current_v_mask = dev_key_v_mask; + // Bound the walk to the 32-level depth of a u32 subset-difference tree. + // `current_v_mask` advances via an arithmetic `>> 1` which sign-extends, so + // a disc-supplied v_mask coarser than dev_key_v_mask would otherwise drive + // current_v_mask up to 0xFFFF_FFFF and spin forever — a crafted MKB must + // not hang the rip thread (this runs before the KCD placeholder gate). + let mut steps = 0u32; while current_v_mask != v_mask { + if steps >= 32 { + break; + } + steps += 1; let mut bit_pos: i32 = -1; for i in (0..32).rev() { if (current_v_mask & (1u32 << i)) == 0 { @@ -243,11 +289,23 @@ pub fn walk_processing_key( for uvs_idx in 0..num_uvs { let p_uv = &uvs[1 + 5 * uvs_idx..]; + // `num_uvs` was computed by `take_while(.. (c[0] & 0xC0) == 0)`, so + // every chunk in `0..num_uvs` already has its revoked-marker bits + // clear — that `take_while` is the single authoritative place the + // parse stops, no inner re-check needed. let u_mask_shift = uvs[5 * uvs_idx]; if u_mask_shift & 0xC0 != 0 { break; } + // 0x20..=0x3F (32..=63) pass the 0xC0 revoked-marker check but are + // out of range for a u32 shift. `wrapping_shl` would silently + // compute shift % 32 (e.g. 32 → no shift → 0xFFFF_FFFF), matching a + // wrong uv slot and deriving a wrong key. Disc-controlled byte: + // skip the slot instead. + if u_mask_shift >= 32 { + continue; + } let uv = u32::from_be_bytes([p_uv[0], p_uv[1], p_uv[2], p_uv[3]]); if uv == 0 { @@ -260,6 +318,11 @@ pub fn walk_processing_key( if ((device_number & u_mask) == (uv & u_mask)) && ((device_number & v_mask) != (uv & v_mask)) { + // dk.u_mask_shift is a u8 from keydb with no range check; guard + // it the same way before the wrapping_shl below. + if dk.u_mask_shift >= 32 { + continue; + } let dev_key_v_mask = calc_v_mask(dk.uv); let dev_key_u_mask: u32 = 0xFFFF_FFFFu32.wrapping_shl(dk.u_mask_shift as u32); @@ -334,6 +397,10 @@ pub enum MediaKeyVariantError { VariantsTableUnavailable, /// VKD index resolved out of the supplied `vkd_table`. VkdIndexOutOfRange, + /// The derived Media Key failed the MKB's Verify-Media-Key relation. + /// On the variant path this final gate replaces the per-match magic + /// check (which does not hold for a Precursor). + MediaKeyVerifyFailed, } impl std::fmt::Display for MediaKeyVariantError { @@ -347,6 +414,7 @@ impl std::fmt::Display for MediaKeyVariantError { MediaKeyVariantError::KcdNotProvided => 7105, MediaKeyVariantError::VariantsTableUnavailable => 7106, MediaKeyVariantError::VkdIndexOutOfRange => 7107, + MediaKeyVariantError::MediaKeyVerifyFailed => 7108, }; write!(f, "E{code}") } @@ -356,11 +424,15 @@ impl std::error::Error for MediaKeyVariantError {} // ── Chain ───────────────────────────────────────────────────────────────── -/// Look up `VARIANTS[uv]` for the matched uv. The byte layout of the -/// per-uv slot in the Variant Number record is undocumented and is -/// disc-specific; this helper returns `None` until a Variant disc is -/// available to fix the layout against. -fn variants_for_uv(_records: &[MkbRecord], _uv_index: usize) -> Option<u16> { +/// Look up the per-slot `VARIANTS` value for the matched subset-difference +/// slot. AACS 2.1 keys the VARIANTS table by the matched SD slot (the same +/// index that selected the cvalue), so the caller passes +/// [`ProcessingKeyMatch::cvalue_index`]. The byte layout of the per-slot entry +/// in the Variant Number record is undocumented and disc-specific; this helper +/// returns `None` until a Variant disc is available to fix the layout against. +/// +/// `sd_slot_index` is the matched subset-difference slot (== cvalue index). +fn variants_for_uv(_records: &[MkbRecord], _sd_slot_index: usize) -> Option<u16> { None } @@ -377,6 +449,12 @@ fn variants_for_uv(_records: &[MkbRecord], _uv_index: usize) -> Option<u16> { /// the final VUK alongside the Media Key. /// /// Returns `(Km, Kvu)` on success. +/// +/// NOTE: the `VARIANTS[uv]` lookup ([`variants_for_uv`]) is not yet +/// implemented, so on a real Variant disc this always returns +/// `Err(`[`MediaKeyVariantError::VariantsTableUnavailable`]`)` before a +/// key is produced. The chain can only succeed against synthetic test +/// fixtures today. pub fn derive_media_key_variant( mkb_records: &[MkbRecord], device_keys: &[DeviceKey], @@ -446,6 +524,17 @@ pub fn derive_media_key_variant( km[12 + i] ^= uv_bytes[i]; } + // Gate: verify the derived Media Key against the MKB's Verify-Media-Key + // record. On the variant path the per-match magic check in + // `walk_processing_key` does NOT hold (it only saw the Precursor), so this + // is the authoritative Kp/Km verification — it MUST run before returning a + // real key. + let mk_dv = mkb_find_mk_dv(mkb_records).ok_or(MediaKeyVariantError::MkbIncomplete)?; + const VERIFY_MAGIC: [u8; 8] = [0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF]; + if aes_ecb_decrypt(&km, &mk_dv)[..8] != VERIFY_MAGIC { + return Err(MediaKeyVariantError::MediaKeyVerifyFailed); + } + // Step: Kvu = AES-G(Km, VID). let kvu = aes_g(&km, vid); @@ -456,6 +545,19 @@ pub fn derive_media_key_variant( mod tests { use super::*; + #[test] + fn calc_pk_from_dk_terminates_on_nonconvergent_mask() { + // Regression for the unbounded-loop hang: pick a (dev_key_v_mask, + // v_mask) pair the arithmetic `>> 1` walk can never reconcile. + // dev_key_v_mask has the MSB set, so `>> 1` sign-extends and the + // mask saturates at 0xFFFF_FFFF, never reaching a coarser v_mask. + // The 32-step bound must let this return rather than spin forever. + let dk = [0x11u8; 16]; + let pk = calc_pk_from_dk(&dk, 0x0000_0002, 0x0000_0000, 0xFFFF_FFFE); + // Bounded exit yields *some* key; we only assert it terminated. + let _ = pk; + } + // ── Helpers ── fn synthetic_mkb_classical() -> Vec<u8> { @@ -575,6 +677,7 @@ mod tests { MediaKeyVariantError::KcdNotProvided, MediaKeyVariantError::VariantsTableUnavailable, MediaKeyVariantError::VkdIndexOutOfRange, + MediaKeyVariantError::MediaKeyVerifyFailed, ]; for e in cases { let s = e.to_string(); @@ -626,7 +729,7 @@ mod tests { mkb.extend_from_slice(&[0x03, 0x00, 0x00, 0x00, 0x02]); // Pick a known DK; with dk.uv == MKB.uv (==2) and - // dk.u_mask_shift == MKB.u_mask_shift (==1), dev_key_v_mask + // dk.u_mask_shift == MKB.u_mask_shift (==3), dev_key_v_mask // equals the MKB's v_mask and the calc_pk_from_dk loop is a // no-op — Kp = aesg3_step(dk, 1). let dk_bytes: [u8; 16] = [ diff --git a/src/clpi.rs b/src/clpi.rs index 498b056..d5df4ac 100644 --- a/src/clpi.rs +++ b/src/clpi.rs @@ -11,17 +11,22 @@ use crate::error::{Error, Result}; /// Parsed CLPI clip info. #[derive(Debug)] -#[allow(dead_code)] -pub struct ClipInfo { +pub(crate) struct ClipInfo { + /// CLPI version string. Parsed for completeness; not yet consumed. + #[allow(dead_code)] pub version: String, /// Total source packets in the m2ts (each 192 bytes) pub source_packet_count: u32, - /// Coarse EP entries for the primary video stream + /// Coarse EP entries for the primary video stream. Populated for the + /// EP-map → sector-extent lookup (`get_extents`), which is exercised by + /// tests and reserved for the timestamp-range read path. + #[allow(dead_code)] pub ep_coarse: Vec<EpCoarse>, - /// Fine EP entries for the primary video stream + /// Fine EP entries for the primary video stream (see `ep_coarse`). + #[allow(dead_code)] pub ep_fine: Vec<EpFine>, /// Per-stream metadata from the ProgramInfo section (BD spec). - /// Cross-validates the MPLS STN view — see `labels/clpi.rs`. + /// Cross-validates the MPLS STN view — see `labels/clpi_audit.rs`. /// Empty when program_info is missing or malformed. pub streams: Vec<ClpiStream>, } @@ -30,57 +35,82 @@ pub struct ClipInfo { /// table. Mirrors the same fields the MPLS STN table carries — see /// `mpls::StreamEntry` for the playlist-side equivalent. #[derive(Debug, Clone)] -#[allow(dead_code)] -pub struct ClpiStream { +pub(crate) struct ClpiStream { /// PID of the stream in the MPEG-TS (matches MPLS). pub pid: u16, - /// SCSI/BD coding type byte (0x80 LPCM, 0x83 TrueHD, 0x86 DTS-HD MA, + /// BD stream coding type byte (0x80 LPCM, 0x83 TrueHD, 0x86 DTS-HD MA, /// 0x90 PG, etc.). See `labels::mpls_universal::coding_type_to_codec_hint`. pub coding_type: u8, /// ISO 639-2 3-char language code. Empty for video streams. pub language: String, + // The CLPI cross-validation consumer (labels/clpi_audit.rs) reads only + // pid/coding_type/language. The codec sub-fields below are parsed from + // the BD stream_coding_info for completeness but have no reader yet. /// Audio format byte (1=mono, 3=stereo, 6=5.1, 12=7.1). /// Zero for non-audio streams. + #[allow(dead_code)] pub audio_format: u8, /// Audio sample rate (1=48kHz, 4=96kHz, 5=192kHz). Zero for non-audio. + #[allow(dead_code)] pub audio_rate: u8, /// Video format byte (1=480i, 4=1080i, 5=720p, 6=1080p, 8=2160p). /// Zero for non-video. + #[allow(dead_code)] pub video_format: u8, /// Video rate (1=23.976, 2=24, 3=25, 4=29.97, 6=50, 7=59.94). + #[allow(dead_code)] pub video_rate: u8, } +/// Coarse EP-map entry. Fields feed the EP-map resolution used by +/// `get_extents` (test-exercised; reserved for the timestamp-range path). #[derive(Debug, Clone)] #[allow(dead_code)] -pub struct EpCoarse { +pub(crate) struct EpCoarse { pub ref_to_fine_id: u32, pub pts_coarse: u32, pub spn_coarse: u32, } +/// Fine EP-map entry (see `EpCoarse`). #[derive(Debug, Clone)] #[allow(dead_code)] -pub struct EpFine { +pub(crate) struct EpFine { pub pts_fine: u32, pub spn_fine: u32, } +// EP-map → sector-extent resolution. Exercised by the unit tests and +// reserved for the timestamp-range read path; no production caller yet. #[allow(dead_code)] impl ClipInfo { /// Reconstruct full PTS from coarse + fine entry. - pub fn full_pts(coarse: &EpCoarse, fine: &EpFine) -> u32 { - (coarse.pts_coarse << 19) + (fine.pts_fine << 8) + /// + /// The BD spec PTS is 33-bit: `pts_coarse` is 14 bits (max 16383) and + /// `16383 << 19` exceeds `u32::MAX`, so the result must be `u64` to + /// avoid overflow (panic in debug, silent wrap in release). + pub fn full_pts(coarse: &EpCoarse, fine: &EpFine) -> u64 { + ((coarse.pts_coarse as u64) << 19) + ((fine.pts_fine as u64) << 8) } /// Reconstruct full SPN from coarse + fine entry. pub fn full_spn(coarse: &EpCoarse, fine: &EpFine) -> u32 { - (coarse.spn_coarse & 0xFFFE_0000) + fine.spn_fine + // The two operands occupy non-overlapping bit ranges (coarse holds + // the high bits, fine the low 17), so OR expresses intent and is + // robust to a hand-constructed EpFine. + debug_assert!(fine.spn_fine <= 0x1_FFFF); + (coarse.spn_coarse & 0xFFFE_0000) | fine.spn_fine } /// Get all EP entries as (PTS, SPN) pairs, fully resolved. - pub fn resolved_ep_map(&self) -> Vec<(u32, u32)> { - let mut entries = Vec::new(); + /// + /// PTS resets at each coarse-group boundary on disc, so the raw + /// concatenation is not globally monotonic. The returned vector is + /// sorted by PTS so callers (e.g. [`get_extents`]) can binary-search it. + /// + /// [`get_extents`]: ClipInfo::get_extents + pub fn resolved_ep_map(&self) -> Vec<(u64, u32)> { + let mut entries = Vec::with_capacity(self.ep_fine.len()); for (ci, coarse) in self.ep_coarse.iter().enumerate() { let fine_start = coarse.ref_to_fine_id as usize; @@ -98,6 +128,12 @@ impl ClipInfo { } } + // get_extents binary-searches by PTS, so the map must be ordered. + // Real discs have globally increasing PTS in coarse order; sort by + // (pts, spn) so a cross-group PTS collision can't leave the search + // landing on the wrong group's SPN. + entries.sort_by_key(|&(pts, spn)| (pts, spn)); + entries } @@ -105,7 +141,9 @@ impl ClipInfo { /// /// Converts PTS timestamps to SPN ranges, then SPN to LBA /// using the file's starting LBA on disc. - pub fn get_extents(&self, in_time: u32, out_time: u32) -> Vec<Extent> { + pub fn get_extents(&self, in_time: u64, out_time: u64) -> Vec<Extent> { + // resolved_ep_map() returns entries sorted by PTS, so binary search + // is valid here. let ep_map = self.resolved_ep_map(); if ep_map.is_empty() { return Vec::new(); @@ -122,7 +160,7 @@ impl ClipInfo { let end_spn = match ep_map.binary_search_by_key(&out_time, |(pts, _)| *pts) { Ok(i) => ep_map[i].1, Err(i) if i < ep_map.len() => ep_map[i].1, - _ => ep_map.last().unwrap().1 + 1, + _ => ep_map.last().unwrap().1.saturating_add(1), }; if end_spn <= start_spn { @@ -162,7 +200,7 @@ pub fn parse(data: &[u8]) -> Result<ClipInfo> { // ClipInfo section at offset 40 // source_packet_count at offset 40 + 4(len) + 2(reserved) + 1(stream_type) + 1(app_type) + 4(reserved) + 4(ts_rate) - let source_packet_count = if data.len() > 56 { + let source_packet_count = if data.len() >= 60 { u32::from_be_bytes([data[56], data[57], data[58], data[59]]) } else { 0 @@ -322,8 +360,17 @@ fn parse_cpi(data: &[u8]) -> Result<(Vec<EpCoarse>, Vec<EpFine>)> { return Ok((Vec::new(), Vec::new())); } + // Bound all EP-map reads to this CPI section. The length field counts + // bytes after itself, so the section spans data[..cpi_length + 4]. A + // bogus ep_map_offset within data.len() but past the CPI section would + // otherwise read into an adjacent CLPI section; clamp first. + let data = &data[..(cpi_length + 4).min(data.len())]; + // CPI type at bits 44-47 (byte 5, lower 4 bits) // Skip to EP map: offset 4 (after length) + 2 (reserved/type) + if data.len() < 6 { + return Ok((Vec::new(), Vec::new())); + } let ep_map = &data[6..]; if ep_map.len() < 4 { return Ok((Vec::new(), Vec::new())); @@ -351,9 +398,6 @@ fn parse_cpi(data: &[u8]) -> Result<(Vec<EpCoarse>, Vec<EpFine>)> { // num_EP_coarse: 16 bits │ (10+4+16+18+32 = 80) // num_EP_fine: 18 bits │ // EP_map_start_address: 32 bits ┘ - if ep_map.len() < 16 { - return Ok((Vec::new(), Vec::new())); - } let _stream_pid = u16::from_be_bytes([ep_map[2], ep_map[3]]); // Read 10 bytes (80 bits) from ep_map[4..14] for bit extraction @@ -389,7 +433,10 @@ fn parse_cpi(data: &[u8]) -> Result<(Vec<EpCoarse>, Vec<EpFine>)> { // Coarse entries start at offset 4, 8 bytes each let coarse_data = &stream_ep[4..]; - let mut ep_coarse = Vec::with_capacity(num_coarse); + // Cap the pre-reservation by what the slice can actually hold: + // num_coarse is a 16-bit disc field, so a hostile value would + // otherwise reserve up to ~0.5 MB for an entry table that doesn't exist. + let mut ep_coarse = Vec::with_capacity(num_coarse.min(coarse_data.len() / 8)); for i in 0..num_coarse { let off = i * 8; if off + 8 > coarse_data.len() { @@ -419,7 +466,13 @@ fn parse_cpi(data: &[u8]) -> Result<(Vec<EpCoarse>, Vec<EpFine>)> { } // Fine entries at fine_start, 4 bytes each - let mut ep_fine = Vec::with_capacity(num_fine); + // Cap the pre-reservation: num_fine is an 18-bit disc field (max + // 262143), so reserve only what the slice can actually hold. + let mut ep_fine = if fine_start < stream_ep.len() { + Vec::with_capacity(num_fine.min((stream_ep.len() - fine_start) / 4)) + } else { + Vec::new() + }; if fine_start < stream_ep.len() { let fine_data = &stream_ep[fine_start..]; for i in 0..num_fine { @@ -643,10 +696,49 @@ mod tests { }; // full_pts = (100 << 19) + (50 << 8) = 52_428_800 + 12_800 = 52_441_600 let pts = ClipInfo::full_pts(&coarse, &fine); - assert_eq!(pts, (100 << 19) + (50 << 8)); + assert_eq!(pts, (100u64 << 19) + (50u64 << 8)); assert_eq!(pts, 52_441_600); } + #[test] + fn full_pts_no_u32_overflow() { + // pts_coarse is a 14-bit field (max 0x3FFF = 16383); 16383 << 19 + // overflows u32, so full_pts must use u64. + let coarse = EpCoarse { + ref_to_fine_id: 0, + pts_coarse: 0x3FFF, + spn_coarse: 0, + }; + let fine = EpFine { + pts_fine: 0x7FF, + spn_fine: 0, + }; + let pts = ClipInfo::full_pts(&coarse, &fine); + assert_eq!(pts, (0x3FFFu64 << 19) + (0x7FFu64 << 8)); + assert!(pts > u32::MAX as u64); + } + + #[test] + fn resolved_ep_map_sorted_for_binary_search() { + // Two coarse groups whose fine PTS reset across the boundary + // (50,100 then 25,75) produce a non-monotonic raw concatenation. + // resolved_ep_map must sort so get_extents' binary search is valid. + let cpi = build_cpi( + 0x1011, + &[(0, 0, 0x00020000), (2, 0, 0x00040000)], + &[(50, 1024), (100, 2048), (25, 512), (75, 1536)], + ); + let data = build_clpi(1_000_000, Some(&cpi)); + let clip = parse(&data).expect("should parse"); + + let resolved = clip.resolved_ep_map(); + assert_eq!(resolved.len(), 4); + // Strictly sorted by PTS. + for w in resolved.windows(2) { + assert!(w[0].0 <= w[1].0, "ep_map not sorted: {resolved:?}"); + } + } + #[test] fn full_spn_calculation() { let coarse = EpCoarse { @@ -674,6 +766,22 @@ mod tests { assert_eq!(spn2, 0x00FE0000 + 0x1234); } + #[test] + fn parse_truncated_clipinfo_no_panic() { + // 57/58/59-byte CLPI with valid magic: passes the data.len() < 40 + // guard but data[56..60] needs 60 bytes. Must not panic. + for len in 40..60usize { + let mut data = vec![0u8; len]; + data[0..4].copy_from_slice(b"HDMV"); + if len >= 8 { + data[4..8].copy_from_slice(b"0200"); + } + let clip = parse(&data).expect("short CLPI should parse, not panic"); + // source_packet_count is unreadable below 60 bytes → 0. + assert_eq!(clip.source_packet_count, 0); + } + } + #[test] fn parse_invalid_magic() { let mut data = build_clpi(1000, None); diff --git a/src/css/auth.rs b/src/css/auth.rs index d0a8ba5..0231548 100644 --- a/src/css/auth.rs +++ b/src/css/auth.rs @@ -106,7 +106,7 @@ const CRYPT_TAB2: [u8; 256] = [ 0x45, 0x78, 0xA9, 0xA8, 0xEA, 0xC9, 0x6A, 0xF7, 0x29, 0x91, 0xF0, 0x02, 0x18, 0x3A, 0x4E, 0x7C, ]; -const CRYPT_TAB3: [u8; 288] = [ +const CRYPT_TAB3: [u8; 256] = [ 0x73, 0x51, 0x95, 0xE1, 0x12, 0xE4, 0xC0, 0x58, 0xEE, 0xF2, 0x08, 0x1B, 0xA9, 0xFA, 0x98, 0x4C, 0xA7, 0x33, 0xE2, 0x1B, 0xA7, 0x6D, 0xF5, 0x30, 0x97, 0x1D, 0xF3, 0x02, 0x60, 0x5A, 0x82, 0x0F, 0x91, 0xD0, 0x9C, 0x10, 0x39, 0x7A, 0x83, 0x85, 0x3B, 0xB2, 0xB8, 0xAE, 0x0C, 0x09, 0x52, 0xEA, @@ -123,8 +123,6 @@ const CRYPT_TAB3: [u8; 288] = [ 0xBD, 0xC1, 0x0E, 0x56, 0x54, 0x3E, 0x14, 0x5F, 0x8C, 0x8F, 0x6E, 0x75, 0x1C, 0x07, 0x39, 0x7B, 0x4B, 0xDB, 0xD3, 0x4B, 0x1E, 0xC8, 0x7E, 0xFE, 0x3E, 0x72, 0x16, 0x83, 0x7D, 0xEE, 0xF5, 0xCA, 0xC5, 0x18, 0xF9, 0xD8, 0x68, 0xAB, 0x38, 0x85, 0xA8, 0xF0, 0xA1, 0x73, 0x9F, 0x5D, 0x19, 0x0B, - 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x33, 0x72, 0x39, 0x25, 0x67, 0x26, 0x6D, 0x71, - 0x36, 0x77, 0x3C, 0x20, 0x62, 0x23, 0x68, 0x74, 0xC3, 0x82, 0xC9, 0x15, 0x57, 0x16, 0x5D, 0x81, ]; const VARIANTS: [u8; 32] = [ @@ -153,10 +151,6 @@ const PERM_VARIANT: [[u8; 32]; 2] = [ ], ]; -// ── SCSI constants ──────────────────────────────────────────────────────── - -const SCSI_READ_DVD_STRUCTURE: u8 = 0xAD; - // ── Public API ──────────────────────────────────────────────────────────── /// Perform CSS bus authentication only. @@ -307,7 +301,7 @@ fn read_disc_key(drive: &mut Drive, agid: u8, bus_key: &[u8; 5]) -> Result<[u8; // READ DVD STRUCTURE, format 0x02 (disc key), 2048+4 bytes let alloc_len: u16 = 2048 + 4; let mut cdb = [0u8; 12]; - cdb[0] = SCSI_READ_DVD_STRUCTURE; + cdb[0] = crate::scsi::SCSI_READ_DISC_STRUCTURE; // bytes 2-5: address = 0 cdb[6] = 0; // layer cdb[7] = 0x02; // format = disc key @@ -333,12 +327,23 @@ fn read_disc_key(drive: &mut Drive, agid: u8, bus_key: &[u8; 5]) -> Result<[u8; } // Try each player key against each of 408 disc key entries. - // Each entry in the block is the disc key encrypted with a specific player key. - // We try all known player keys and verify by checking that two different - // entries produce the same disc key. - let mut candidates: Vec<([u8; 5], usize, usize)> = Vec::new(); // (disc_key, pk_idx, pos) + // Each entry in the block is the disc key encrypted with a specific player + // key. We collect every decryption and accept the disc key as soon as two + // independent decryptions agree on the same 5-byte value (the agreement may + // come from two different player keys or from one player key decrypting two + // different entries to the same value). + // + // NOTE: this is a collision heuristic, not the canonical CSS disc-key + // self-verification (which decrypts the verification entry with the + // candidate and checks the result equals the candidate). A coincidental + // collision among the ~12,648 candidate decryptions could in principle + // accept a wrong disc key; in practice a chance collision on 5 bytes is + // improbable enough to serve as the validity check, and this path is the + // production DVD disc-key recovery. Left as-is to avoid regressing it + // without a real disc-key-block test vector to validate against. + let mut candidates: Vec<[u8; 5]> = Vec::new(); - for (pk_idx, player_key) in PLAYER_KEYS.iter().enumerate() { + for player_key in PLAYER_KEYS.iter() { for pos in 0..408 { let offset = pos * 5; if offset + 5 > disc_key_block.len() { @@ -348,13 +353,11 @@ fn read_disc_key(drive: &mut Drive, agid: u8, bus_key: &[u8; 5]) -> Result<[u8; enc.copy_from_slice(&disc_key_block[offset..offset + 5]); let candidate = super::lfsr::decrypt_key(0x00, player_key, &enc); - // Check if any previous candidate matches (same disc key from different entry/pk) - for (prev, _, _) in &candidates { - if *prev == candidate { - return Ok(candidate); - } + // Accept on the first agreement between two independent decryptions. + if candidates.contains(&candidate) { + return Ok(candidate); } - candidates.push((candidate, pk_idx, pos)); + candidates.push(candidate); } } @@ -392,65 +395,16 @@ fn read_raw_title_key(drive: &mut Drive, agid: u8, lba: u32) -> Result<[u8; 5]> Ok(key) } -#[allow(dead_code)] -fn read_title_key( - drive: &mut Drive, - agid: u8, - lba: u32, - bus_key: &[u8; 5], - disc_key: &[u8; 5], -) -> Result<[u8; 5]> { - let scsi = drive.scsi_mut(); - - let mut cdb = [0u8; 12]; - cdb[0] = crate::scsi::SCSI_REPORT_KEY; - cdb[2] = (lba >> 24) as u8; - cdb[3] = (lba >> 16) as u8; - cdb[4] = (lba >> 8) as u8; - cdb[5] = lba as u8; - cdb[8] = 0x00; - cdb[9] = 0x0C; - cdb[10] = (agid << 6) | 0x04; - - let mut buf = [0u8; 12]; - let tk_result = scsi.execute( - &cdb, - crate::scsi::DataDirection::FromDevice, - &mut buf, - 5_000, - ); - tk_result.map_err(|_| Error::CssAuthFailed)?; - - // Title key at bytes 5..10, byte-reversed - let mut title_key = [0u8; 5]; - for i in 0..5 { - title_key[i] = buf[5 + (4 - i)]; - } - - // XOR with reversed bus key (same pattern as disc key block) - for i in 0..5 { - title_key[i] ^= bus_key[4 - i]; - } - - // Check for null key (title not encrypted) - if title_key == [0u8; 5] { - return Ok(title_key); - } - - // Decrypt with disc key (invert=0xFF for title keys) - let title_key = super::lfsr::decrypt_key(0xFF, disc_key, &title_key); - - Ok(title_key) -} - // ── CSSCryptKey ─────────────────────────────────────────────────────────── -/// Exposed for testing only. -pub fn test_crypt_key(key_type: usize, variant: u8, challenge: &[u8; 10]) -> [u8; 5] { - crypt_key(key_type, variant, challenge) -} - fn crypt_key(key_type: usize, variant: u8, challenge: &[u8; 10]) -> [u8; 5] { + // key_type indexes PERM_CHALLENGE ([_;3]); variant indexes + // VARIANTS/PERM_VARIANT ([_;32]). All internal callers pass key_type in + // 0..3 and variant in 0..32; the asserts document the contract for the + // pub(crate) test entry point test_crypt_key and turn a would-be + // out-of-bounds panic into an explicit precondition violation. + debug_assert!(key_type < 3, "crypt_key: key_type out of range"); + debug_assert!((variant as usize) < 32, "crypt_key: variant out of range"); let perm = &PERM_CHALLENGE[key_type]; let mut scratch = [0u8; 10]; for i in 0..10 { diff --git a/src/css/crack.rs b/src/css/crack.rs index 7727f29..e920886 100644 --- a/src/css/crack.rs +++ b/src/css/crack.rs @@ -1,13 +1,22 @@ //! CSS title key recovery — Stevenson's divide-and-conquer attack (1999). //! //! Given a scrambled DVD sector with known plaintext (MPEG-2 PES headers), -//! recovers the 5-byte title key by: +//! this would recover the 5-byte title key by: //! -//! 1. XORing ciphertext with TAB1[ciphertext] to cancel the mangling +//! 1. Computing `TAB1[ciphertext] ^ plaintext` to cancel the TAB1 output +//! mangling and expose the raw LFSR-combination keystream //! 2. Iterating all 2^16 LFSR1 states //! 3. For each: deducing what LFSR0 must produce, then verifying //! -//! Total work: ~65536 iterations with 10-byte validation = instant. +//! NOTE: this recovery path is currently non-functional. It models the +//! textbook direct-seed CSS cipher, whereas the in-repo descrambler +//! ([`super::lfsr::descramble_sector`]) seeds its LFSRs from a key that +//! has been run through an additional `decrypt_key` mangling step. The two +//! are therefore inconsistent and [`recover_title_key`] never returns a key +//! for a sector scrambled by this crate's own descrambler. The production +//! DVD path does NOT use this fallback — it derives the title key over SCSI +//! ([`super::auth::authenticate_and_read_title_key`]). See the ignored +//! regression test below. //! //! Algorithm: Frank A. Stevenson, "Divide and conquer attack" (1999). @@ -25,7 +34,12 @@ const FLAG_BYTE: usize = 0x14; /// region (bytes 0x80+). For MPEG-2 sectors, the first bytes are typically /// a PES header: `00 00 01 [stream_id] ...` /// -/// Returns the recovered 5-byte title key, or None if recovery fails. +/// Returns the recovered 5-byte title key, or `None` if recovery fails. +/// +/// NOTE: see the module docs — this attack models the textbook direct-seed +/// CSS cipher and is inconsistent with this crate's descrambler, so it +/// currently returns `None` even for an exact known plaintext. It is not on +/// the production DVD decrypt path. pub fn recover_title_key(sector: &[u8], plain: &[u8]) -> Option<[u8; 5]> { if sector.len() < SECTOR_SIZE || plain.len() < 10 { return None; @@ -39,10 +53,11 @@ pub fn recover_title_key(sector: &[u8], plain: &[u8]) -> Option<[u8; 5]> { let crypted = §or[ENCRYPTED_START..]; let seed = §or[SEED_OFFSET..SEED_OFFSET + 5]; - // Phase 1: Cancel the TAB1 mangling layer - // The CSS cipher applies TAB1 as an output permutation. - // XORing ciphertext with TAB1[ciphertext] and plaintext removes it, - // leaving the raw LFSR combination output. + // Phase 1: Cancel the TAB1 mangling layer and subtract the known plaintext. + // The CSS cipher applies TAB1 as an output permutation. Computing + // `buf[i] = TAB1[crypted[i]] ^ plain[i]` both undoes that permutation and + // XORs out the known plaintext, leaving the raw LFSR-combination keystream + // bytes for the attack to match against. let mut buf = [0u8; 10]; for i in 0..10 { if i >= crypted.len() || i >= plain.len() { @@ -82,8 +97,12 @@ pub fn recover_title_key(sector: &[u8], plain: &[u8]) -> Option<[u8; 5]> { t5 += t6 + t4_perm as u32; let t6_inv = TAB4[t6 as usize & 0xFF]; - // Build LFSR0 candidate from deduced output bytes - t3 = (t3 << 8) | t6_inv as u32; + // Build LFSR0 candidate from deduced output bytes. + // wrapping_shl: the accumulator is a rolling 32-bit window; + // the top byte is intentionally shifted out. Matches the + // release-mode wrap (no behaviour change) without a debug + // overflow panic. + t3 = t3.wrapping_shl(8) | t6_inv as u32; t5 >>= 8; } @@ -97,9 +116,11 @@ pub fn recover_title_key(sector: &[u8], plain: &[u8]) -> Option<[u8; 5]> { t1 = ((t1 & 1) << 8) ^ t4 as u32; let t4_perm = TAB5[t4 as usize]; - // Clock LFSR0 forward + // Clock LFSR0 forward. wrapping_shl keeps the rolling 32-bit + // window semantics (top byte shifted out) identical to the + // release build while avoiding a debug overflow panic. let t6 = ((((((t3 >> 8) ^ t3) >> 1) ^ t3) >> 3) ^ t3) >> 7; - t3 = (t3 << 8) | (t6 & 0xFF); + t3 = t3.wrapping_shl(8) | (t6 & 0xFF); let t6_perm = TAB4[(t6 & 0xFF) as usize]; t5 += t6_perm as u32 + t4_perm as u32; @@ -143,7 +164,11 @@ pub fn recover_title_key(sector: &[u8], plain: &[u8]) -> Option<[u8; 5]> { let t4 = (t3 >> 1).wrapping_sub(4); for t5_off in 0u32..8 { let val = t4.wrapping_add(t5_off); - if (val * 2 + 8 - (val & 7)) == t3 { + // Reconstruction probe: val can sit near u32::MAX, so the + // (val*2 + 8 - (val & 7)) expression must wrap rather than + // panic in debug. wrapping_* reproduces the release result + // exactly (the comparison against t3 is unaffected). + if val.wrapping_mul(2).wrapping_add(8).wrapping_sub(val & 7) == t3 { result_key[0] = (i_try >> 8) as u8; result_key[1] = (i_try & 0xFF) as u8; result_key[2] = (val & 0xFF) as u8; @@ -172,9 +197,15 @@ pub fn recover_title_key(sector: &[u8], plain: &[u8]) -> Option<[u8; 5]> { Some(result_key) } -/// Crack the CSS title key from an encrypted sector using MPEG-2 pattern attack. +/// Crack the CSS title key from an encrypted sector using an MPEG-2 +/// pattern attack. /// -/// Detects the PES header pattern at byte 0x80 and uses it as known plaintext. +/// Detects the PES header pattern at byte 0x80 and uses it as known +/// plaintext. This is a best-effort fallback for the SCSI auth path +/// (see [`super::resolve`]): it only succeeds on a sector whose +/// encrypted region begins with one of the tried PES header patterns, +/// and returns `None` otherwise. The production DVD path obtains the +/// title key via drive authentication, not cracking. pub fn crack_title_key(sector: &[u8]) -> Option<[u8; 5]> { if sector.len() < SECTOR_SIZE { return None; @@ -199,7 +230,8 @@ pub fn crack_title_key(sector: &[u8]) -> Option<[u8; 5]> { // Try many PES header patterns at byte 0x80. // Structure: 00 00 01 [stream_id] [len_hi] [len_lo] [flags1] [flags2] [hdr_len] [data] - let mut patterns: Vec<[u8; 10]> = Vec::with_capacity(128); + // 24 padding-stream + 144 video/audio + 1 navigation = 169 patterns. + let mut patterns: Vec<[u8; 10]> = Vec::with_capacity(169); // Padding stream (0xBE): payload is 0xFF bytes, various lengths for len_hi in 0u8..8 { @@ -242,23 +274,6 @@ pub fn crack_title_key(sector: &[u8]) -> Option<[u8; 5]> { None } -/// Crack CSS key from multiple sectors. -pub fn crack_from_sectors(sectors: &[Vec<u8>]) -> Option<[u8; 5]> { - for sector in sectors { - if sector.len() < SECTOR_SIZE { - continue; - } - let flags = (sector[FLAG_BYTE] >> 4) & 0x03; - if flags == 0 { - continue; - } - if let Some(key) = crack_title_key(sector) { - return Some(key); - } - } - None -} - #[cfg(test)] mod tests { use super::*; @@ -269,6 +284,30 @@ mod tests { assert!(crack_title_key(§or).is_none()); } + /// Regression: the LFSR0 reconstruction arithmetic must not overflow + /// (panic) in a debug build for scrambled sector content. Exercises + /// the full 2^16 Stevenson search via recover_title_key directly (the + /// overflow site), with the assertion simply that it does not panic. + /// recover_title_key is used rather than crack_title_key to avoid + /// re-running the search for all 169 PES patterns. + #[test] + fn crack_scrambled_sectors_never_overflow() { + for seed in 0u32..4 { + let mut sector = vec![0u8; SECTOR_SIZE]; + sector[FLAG_BYTE] = 0x30; // scramble flag set + let mut x = seed.wrapping_mul(2_654_435_761).wrapping_add(1); + for b in sector.iter_mut().skip(0x80) { + x = x.wrapping_mul(1_103_515_245).wrapping_add(12_345); + *b = (x >> 16) as u8; + } + for (i, b) in sector[SEED_OFFSET..SEED_OFFSET + 5].iter_mut().enumerate() { + *b = seed.wrapping_add(i as u32) as u8; + } + let plain = [0x00u8, 0x00, 0x01, 0xE0, 0x00, 0x00, 0x80, 0x80, 0x05, 0x21]; + let _ = recover_title_key(§or, &plain); + } + } + #[test] fn crack_too_short_returns_none() { let sector = vec![0u8; 100]; @@ -282,13 +321,22 @@ mod tests { assert!(recover_title_key(§or, &short_plain).is_none()); } - /// Test 3: css_crack_recovers_key_from_scrambled_sector + /// Build a scrambled sector with known plaintext (both an MPEG PES header + /// at 0x80 and an exact-plaintext probe), then assert that the Stevenson + /// recovery actually recovers a key whose descramble round-trips the body. /// - /// Build a plaintext sector with known MPEG-2 PES headers, scramble it - /// with a known title key, then run crack_title_key() on the scrambled - /// sector. If the Stevenson attack succeeds, verify that descrambling - /// with the recovered key produces the original plaintext at bytes 128..132. + /// This is the regression gate for the CSS crack/recover path. It is + /// `#[ignore]`d because that path is currently non-functional: the + /// recovery models the textbook direct-seed cipher, whereas this crate's + /// [`descramble_sector`] seeds from a `decrypt_key`-mangled key, so the + /// two are inconsistent and recovery returns `None`. When the crack + /// algorithm is re-derived against this crate's actual descrambler, this + /// test must pass with `--ignored` removed. The production DVD path does + /// not use crack/recover (it authenticates over SCSI), so the broken + /// fallback does not affect shipped behavior. #[test] + #[ignore = "CSS crack/recover path is non-functional vs this crate's descrambler; \ + see module docs. Regression gate for a future fix."] fn css_crack_recovers_key_from_scrambled_sector() { use super::super::lfsr::descramble_sector; @@ -309,90 +357,47 @@ mod tests { // Sector seed at bytes 0x54-0x58 plaintext[SEED_OFFSET..SEED_OFFSET + 5].copy_from_slice(&[0x11, 0x22, 0x33, 0x44, 0x55]); - // PES header at byte 0x80: 00 00 01 E0 (video stream) - // Then typical PES header bytes for a stream with PTS - plaintext[0x80] = 0x00; - plaintext[0x81] = 0x00; - plaintext[0x82] = 0x01; - plaintext[0x83] = 0xE0; - plaintext[0x84] = 0x00; // PES length hi - plaintext[0x85] = 0x00; // PES length lo - plaintext[0x86] = 0x80; // flags: data_alignment, copyright - plaintext[0x87] = 0x80; // PTS flag - plaintext[0x88] = 0x05; // PES header data length - plaintext[0x89] = 0x21; // PTS byte 1 + // PES header at byte 0x80: 00 00 01 E0 (video stream) with PTS. + let exact_plain: [u8; 10] = [0x00, 0x00, 0x01, 0xE0, 0x00, 0x00, 0x80, 0x80, 0x05, 0x21]; + plaintext[0x80..0x80 + 10].copy_from_slice(&exact_plain); let original_plaintext = plaintext.clone(); - // "Scramble" the sector by calling descramble (which XORs the keystream) - // on the plaintext. This produces a scrambled sector. + // "Scramble" the sector by XORing the keystream over the plaintext. descramble_sector(&title_key, &mut plaintext); - // The scramble flag was cleared by descramble_sector. Restore it so - // the cracker sees it as encrypted. + // descramble_sector cleared the flag; restore it so the cracker sees + // the sector as encrypted. plaintext[FLAG_BYTE] = 0x30; - // Now we have a scrambled sector. Try to crack the title key. - let cracked_key = crack_title_key(&plaintext); + // 1) Pattern-guessing entry point must recover a key. + let cracked = crack_title_key(&plaintext); + assert!( + cracked.is_some(), + "crack_title_key returned None for a sector scrambled with a known key" + ); + let cracked = cracked.unwrap(); + let mut body = plaintext.clone(); + descramble_sector(&cracked, &mut body); + assert_eq!( + &body[0x80..SECTOR_SIZE], + &original_plaintext[0x80..SECTOR_SIZE], + "crack_title_key key did not round-trip the body" + ); - match cracked_key { - Some(key) => { - // Verify: descramble with the cracked key should recover plaintext - let mut test = plaintext.clone(); - descramble_sector(&key, &mut test); - - // Check that the PES header is recovered - assert_eq!(test[0x80], 0x00, "PES byte 0 mismatch"); - assert_eq!(test[0x81], 0x00, "PES byte 1 mismatch"); - assert_eq!(test[0x82], 0x01, "PES byte 2 mismatch"); - assert_eq!(test[0x83], 0xE0, "PES byte 3 mismatch"); - - // Also verify the rest of the encrypted region matches original - assert_eq!( - &test[0x80..SECTOR_SIZE], - &original_plaintext[0x80..SECTOR_SIZE], - "Decrypted content does not match original plaintext" - ); - - eprintln!( - "Stevenson attack succeeded: cracked key = {:02X?}, original = {:02X?}", - key, title_key - ); - } - None => { - // The Stevenson attack may not always find a key for all title keys - // and sector seeds. This is expected for some combinations where the - // known plaintext pattern doesn't match what crack_title_key tries. - eprintln!( - "Stevenson attack did not find key for title_key={:02X?} seed={:02X?}. \ - This can happen when the cipher output doesn't match the tried patterns. \ - Testing with recover_title_key directly with exact plaintext.", - title_key, - &[0x11u8, 0x22, 0x33, 0x44, 0x55], - ); - - // Try with exact known plaintext instead of guessing - let exact_plain: [u8; 10] = - [0x00, 0x00, 0x01, 0xE0, 0x00, 0x00, 0x80, 0x80, 0x05, 0x21]; - let recovered = recover_title_key(&plaintext, &exact_plain); - if let Some(key) = recovered { - let mut test = plaintext.clone(); - descramble_sector(&key, &mut test); - assert_eq!(test[0x80], 0x00); - assert_eq!(test[0x81], 0x00); - assert_eq!(test[0x82], 0x01); - eprintln!( - "recover_title_key with exact plaintext succeeded: {:02X?}", - key - ); - } else { - eprintln!( - "recover_title_key also returned None. The attack may not converge \ - for this particular key/seed combination. This is a known limitation \ - of the brute-force LFSR0 recovery phase." - ); - } - } - } + // 2) Exact known plaintext must also recover a round-tripping key. + let recovered = recover_title_key(&plaintext, &exact_plain); + assert!( + recovered.is_some(), + "recover_title_key returned None for exact known plaintext" + ); + let recovered = recovered.unwrap(); + let mut body2 = plaintext.clone(); + descramble_sector(&recovered, &mut body2); + assert_eq!( + &body2[0x80..SECTOR_SIZE], + &original_plaintext[0x80..SECTOR_SIZE], + "recover_title_key key did not round-trip the body" + ); } } diff --git a/src/css/lfsr.rs b/src/css/lfsr.rs index d2fd525..e85c82d 100644 --- a/src/css/lfsr.rs +++ b/src/css/lfsr.rs @@ -1,7 +1,8 @@ //! CSS cipher implementation based on the Stevenson 1999 analysis. //! //! The CSS cipher uses two table-driven feedback circuits: -//! - LFSR1: 9-bit state (two halves), driven by TAB2/TAB3 +//! - LFSR1: 17-bit state (9-bit lo + 8-bit hi register, seeded from +//! key[0..2]), driven by TAB2/TAB3 //! - LFSR0: 32-bit state, driven by a feedback polynomial through TAB4 //! //! The keystream is the bytewise sum (with carry) of both LFSR outputs. @@ -12,6 +13,45 @@ use super::tables::{TAB1, TAB2, TAB3, TAB4, TAB5}; +/// Seed the 32-bit LFSR0 register from the 5-byte working key, applying +/// the per-byte TAB4 bit-reversal. Shared by [`descramble_sector`] and +/// [`decrypt_key`] so the seeding lives in one place. +#[inline] +fn seed_lfsr0(key: &[u8; 5]) -> u32 { + let lfsr0: u32 = ((key[4] as u32) << 17) + | ((key[3] as u32) << 9) + | (((key[2] as u32) << 1) + 8 - (key[2] as u32 & 7)); + (TAB4[(lfsr0 & 0xFF) as usize] as u32) << 24 + | (TAB4[((lfsr0 >> 8) & 0xFF) as usize] as u32) << 16 + | (TAB4[((lfsr0 >> 16) & 0xFF) as usize] as u32) << 8 + | TAB4[((lfsr0 >> 24) & 0xFF) as usize] as u32 +} + +/// One CSS keystream step. Advances both LFSRs, folds their permuted +/// outputs into `combined` (carry kept across calls), and returns the +/// low keystream byte. `invert` XORs the LFSR0 output index (0x00 on the +/// descramble path, 0xFF on the key-decrypt path). +#[inline] +fn css_step( + lfsr1_lo: &mut u32, + lfsr1_hi: &mut u32, + lfsr0: &mut u32, + combined: &mut u32, + invert: u8, +) -> u8 { + let o_lfsr1 = TAB2[*lfsr1_hi as usize] ^ TAB3[*lfsr1_lo as usize]; + *lfsr1_hi = *lfsr1_lo >> 1; + *lfsr1_lo = ((*lfsr1_lo & 1) << 8) ^ o_lfsr1 as u32; + + let o_lfsr0 = (((((((*lfsr0 >> 8) ^ *lfsr0) >> 1) ^ *lfsr0) >> 3) ^ *lfsr0) >> 7) as u8; + *lfsr0 = (*lfsr0 >> 8) | ((o_lfsr0 as u32) << 24); + + *combined += TAB5[o_lfsr1 as usize] as u32 + TAB4[(o_lfsr0 ^ invert) as usize] as u32; + let out = (*combined & 0xFF) as u8; + *combined >>= 8; + out +} + /// Descramble a CSS-encrypted DVD sector in place. /// /// The sector seed (bytes 0x54-0x58) is XORed with the title key to produce @@ -20,7 +60,18 @@ use super::tables::{TAB1, TAB2, TAB3, TAB4, TAB5}; /// /// The scramble flag at byte 0x14 (bits 4-5) indicates encryption. /// After descrambling, the flag is cleared. +/// +/// No-op (returns without modifying `sector`) in two cases: +/// - `sector.len() < 2048`: the encrypted region (0x80..0x800) is not +/// fully present. Callers chunk by 2048, so a trailing partial chunk is +/// left untouched. The `debug_assert!` flags this misuse in debug/test +/// builds; a DVD sector is always exactly 2048 bytes. +/// - scramble flags are zero: the sector is not CSS-encrypted. pub fn descramble_sector(title_key: &[u8; 5], sector: &mut [u8]) { + debug_assert!( + sector.len() >= 2048, + "descramble_sector: buffer shorter than one 2048-byte sector" + ); if sector.len() < 2048 { return; } @@ -39,36 +90,40 @@ pub fn descramble_sector(title_key: &[u8; 5], sector: &mut [u8]) { title_key[4] ^ sector[0x58], ]; - // Decrypt the key through the CSS mangling function to get the working key - let working_key = decrypt_key(0xFF, &key, §or[0x54..0x59]); + // Decrypt the key through the CSS mangling function to get the working key. + // The sector seed is bytes 0x54..0x59 (5 bytes). + let seed: [u8; 5] = [ + sector[0x54], + sector[0x55], + sector[0x56], + sector[0x57], + sector[0x58], + ]; + let working_key = decrypt_key(0xFF, &key, &seed); // Generate keystream and XOR with encrypted region let mut lfsr1_lo: u32 = working_key[0] as u32 | 0x100; let mut lfsr1_hi: u32 = working_key[1] as u32; - - let mut lfsr0: u32 = ((working_key[4] as u32) << 17) - | ((working_key[3] as u32) << 9) - | (((working_key[2] as u32) << 1) + 8 - (working_key[2] as u32 & 7)); - lfsr0 = (TAB4[(lfsr0 & 0xFF) as usize] as u32) << 24 - | (TAB4[((lfsr0 >> 8) & 0xFF) as usize] as u32) << 16 - | (TAB4[((lfsr0 >> 16) & 0xFF) as usize] as u32) << 8 - | TAB4[((lfsr0 >> 24) & 0xFF) as usize] as u32; + let mut lfsr0: u32 = seed_lfsr0(&working_key); let mut combined: u32 = 0; - // Generate 1920 keystream bytes (for sector bytes 128..2048) - // Per libdvdcss css_unscramble: TAB1 permutation on ciphertext, no invert on LFSR0 + // Generate 1920 keystream bytes (for sector bytes 128..2048) and XOR them + // into the encrypted region. Each keystream byte is the carrying sum of the + // TAB5-permuted LFSR1 output and the TAB4-permuted LFSR0 output. No TAB1 + // permutation is applied to the ciphertext here (TAB1 is only used inside + // decrypt_key); the working key was already produced by decrypt_key above, + // so this keystream is paired with that mangling step, not a plain + // direct-seed unscramble. No invert is applied on the LFSR0 output. for byte in sector.iter_mut().take(2048).skip(128) { - let o_lfsr1 = TAB2[lfsr1_hi as usize] ^ TAB3[lfsr1_lo as usize]; - lfsr1_hi = lfsr1_lo >> 1; - lfsr1_lo = ((lfsr1_lo & 1) << 8) ^ o_lfsr1 as u32; - - let o_lfsr0 = (((((((lfsr0 >> 8) ^ lfsr0) >> 1) ^ lfsr0) >> 3) ^ lfsr0) >> 7) as u8; - lfsr0 = (lfsr0 >> 8) | ((o_lfsr0 as u32) << 24); - - combined += TAB5[o_lfsr1 as usize] as u32 + TAB4[o_lfsr0 as usize] as u32; - *byte ^= (combined & 0xFF) as u8; - combined >>= 8; + let ks = css_step( + &mut lfsr1_lo, + &mut lfsr1_hi, + &mut lfsr0, + &mut combined, + 0x00, + ); + *byte ^= ks; } // Clear scramble flags @@ -80,37 +135,23 @@ pub fn descramble_sector(title_key: &[u8; 5], sector: &mut [u8]) { /// Decrypts `p_crypted` using `p_key` with the CSS two-LFSR cipher. /// The `invert` parameter controls the XOR applied to LFSR0 output /// (0x00 for disc key decryption, 0xFF for title key / sector key). -pub(crate) fn decrypt_key(invert: u8, p_key: &[u8; 5], p_crypted: &[u8]) -> [u8; 5] { - if p_crypted.len() < 5 { - return *p_key; - } - +pub(crate) fn decrypt_key(invert: u8, p_key: &[u8; 5], p_crypted: &[u8; 5]) -> [u8; 5] { let mut lfsr1_lo: u32 = p_key[0] as u32 | 0x100; let mut lfsr1_hi: u32 = p_key[1] as u32; - - let mut lfsr0: u32 = ((p_key[4] as u32) << 17) - | ((p_key[3] as u32) << 9) - | (((p_key[2] as u32) << 1) + 8 - (p_key[2] as u32 & 7)); - lfsr0 = (TAB4[(lfsr0 & 0xFF) as usize] as u32) << 24 - | (TAB4[((lfsr0 >> 8) & 0xFF) as usize] as u32) << 16 - | (TAB4[((lfsr0 >> 16) & 0xFF) as usize] as u32) << 8 - | TAB4[((lfsr0 >> 24) & 0xFF) as usize] as u32; + let mut lfsr0: u32 = seed_lfsr0(p_key); let mut combined: u32 = 0; let mut k = [0u8; 5]; + // TAB5 for LFSR1 output, TAB4 for LFSR0^invert (per libdvdcss css_DecryptKey). for byte in &mut k { - let o_lfsr1 = TAB2[lfsr1_hi as usize] ^ TAB3[lfsr1_lo as usize]; - lfsr1_hi = lfsr1_lo >> 1; - lfsr1_lo = ((lfsr1_lo & 1) << 8) ^ o_lfsr1 as u32; - - let o_lfsr0 = (((((((lfsr0 >> 8) ^ lfsr0) >> 1) ^ lfsr0) >> 3) ^ lfsr0) >> 7) as u8; - lfsr0 = (lfsr0 >> 8) | ((o_lfsr0 as u32) << 24); - - // TAB5 for LFSR1 output, TAB4 for LFSR0^invert (per libdvdcss css_DecryptKey) - combined += TAB5[o_lfsr1 as usize] as u32 + TAB4[(o_lfsr0 ^ invert) as usize] as u32; - *byte = (combined & 0xFF) as u8; - combined >>= 8; + *byte = css_step( + &mut lfsr1_lo, + &mut lfsr1_hi, + &mut lfsr0, + &mut combined, + invert, + ); } // Two rounds of chained XOR through TAB1 @@ -184,7 +225,7 @@ mod tests { assert_ne!(result, [0u8; 5]); } - /// Test 1: css_decrypt_key_roundtrip + /// css_decrypt_key_roundtrip /// /// decrypt_key is not a simple encrypt/decrypt pair — it is a one-way mangling /// function. However, we can verify consistency: calling it twice with the same @@ -228,11 +269,13 @@ mod tests { } } - /// Test 2: css_descramble_produces_valid_mpeg2 + /// Test 2: descramble_modifies_encrypted_region /// - /// descramble_sector XORs a keystream into bytes 128..2048. Calling it - /// twice with the same key and restored scramble flag should roundtrip, - /// since XOR is its own inverse. + /// descramble_sector XORs a keystream into bytes 128..2048. The keystream + /// depends only on (title_key, sector_seed), so applying descramble twice + /// with the scramble flag restored between calls re-XORs the same keystream + /// and restores the original encrypted region — the keystream XOR is its + /// own inverse. This pins the cipher's involution property over the body. #[test] fn css_descramble_modifies_encrypted_region() { let title_key = [0x42, 0x13, 0x37, 0xBE, 0xEF]; @@ -255,9 +298,20 @@ mod tests { } // Encrypted region modified assert_ne!(§or[128..256], &original[128..256]); + + // Round-trip: restore the scramble flag and descramble again. The same + // keystream is regenerated (it depends only on title_key + seed, both + // unchanged), so the body is restored to its original bytes. + sector[0x14] = 0x30; + descramble_sector(&title_key, &mut sector); + assert_eq!( + §or[128..2048], + &original[128..2048], + "double descramble did not restore the encrypted region" + ); } - /// Test 4: css_tab1_relationship + /// css_tab1_relationship /// /// Verify the structure of TAB1: it is a substitution table used in /// key mangling. Check that no two inputs map to the same output @@ -281,7 +335,7 @@ mod tests { } } - /// Test 5: css_tab4_is_bit_reversal + /// css_tab4_is_bit_reversal /// /// TAB4 reverses the bits of each byte: TAB4[0x01] = 0x80, TAB4[0x80] = 0x01, etc. #[test] diff --git a/src/css/mod.rs b/src/css/mod.rs index eae97bd..463b0fa 100644 --- a/src/css/mod.rs +++ b/src/css/mod.rs @@ -1,13 +1,21 @@ //! CSS (Content Scramble System) — DVD disc encryption. //! //! CSS uses a weak 40-bit LFSR stream cipher (broken since 1999). -//! No keys needed — the title key is cracked from encrypted content -//! using a known-plaintext attack on MPEG-2 PES headers. +//! +//! The production entry point is [`resolve`]. Two title-key acquisition +//! paths exist behind it: +//! - The SCSI auth path drives bus authentication with the compiled-in CSS +//! player keys and reads the title key from the drive (the production DVD +//! path on a live drive). +//! - The crack fallback ([`crack_key`]) needs no keys — it attempts the +//! Stevenson known-plaintext attack on MPEG-2 PES headers. (Currently +//! non-functional; see the `crack` module docs.) //! //! Usage: //! ```rust,ignore -//! let key = css::crack_key(reader, &extents)?; -//! css::descramble_sector(&key, &mut sector); +//! if let Some(state) = css::resolve(&mut ctx) { +//! css::descramble_sector(&state, &mut sector); +//! } //! ``` pub mod auth; @@ -22,7 +30,7 @@ use crate::sector::SectorSource; /// CSS decryption state for a DVD title. #[derive(Debug, Clone)] pub struct CssState { - /// Cracked 5-byte title key + /// 5-byte CSS title key (from SCSI auth or the crack fallback). pub title_key: [u8; 5], } @@ -37,7 +45,7 @@ pub struct CssState { /// headers; works on disc images and on drives whose CSS auth path /// is unavailable). /// -/// `live_drive` always wins when both modes are populated. +/// The `drive` (auth) path always wins when both modes are populated. pub struct CssContext<'a> { /// Live SCSI drive — when present, [`resolve`] tries the auth path. pub drive: Option<&'a mut Drive>, @@ -70,23 +78,32 @@ pub fn resolve(ctx: &mut CssContext<'_>) -> Option<CssState> { None } -/// Crack the CSS title key by reading encrypted sectors and applying -/// a known-plaintext attack on MPEG-2 headers. -/// -/// Crack the CSS title key by scanning scrambled sectors across extents. +/// Crack the CSS title key by scanning scrambled sectors across extents and +/// applying a known-plaintext attack on MPEG-2 PES headers. /// /// The Stevenson attack needs a sector where a PES header starts at byte /// 0x80 (start of the encrypted region). This only happens when a new PES /// packet begins at exactly sector offset 128. We scan up to 50000 /// scrambled sectors sequentially across all extents. +/// +/// NOTE: the underlying recovery ([`crack::recover_title_key`]) is currently +/// non-functional against this crate's descrambler (see `crack` module +/// docs), so this scan returns `None`. The production DVD path uses the SCSI +/// auth path, not this crack fallback. pub fn crack_key(reader: &mut dyn SectorSource, extents: &[Extent]) -> Option<CssState> { let mut tried = 0u32; let max_tries = 50_000; + // Reused across every scanned sector; read_sectors overwrites all 2048 + // bytes on success, so no re-zeroing is needed between iterations. + let mut buf = vec![0u8; 2048]; + for ext in extents { let mut i = 0; while i < ext.sector_count && tried < max_tries { - let mut buf = vec![0u8; 2048]; + // Every scanned sector counts toward the cap, so a long run + // of unscrambled sectors can't read past the budget. + tried += 1; if reader .read_sectors(ext.start_lba + i, 1, &mut buf, true) .is_ok() @@ -95,7 +112,6 @@ pub fn crack_key(reader: &mut dyn SectorSource, extents: &[Extent]) -> Option<Cs if let Some(key) = crack::crack_title_key(&buf) { return Some(CssState { title_key: key }); } - tried += 1; } i += 1; } diff --git a/src/css/tables.rs b/src/css/tables.rs index d966b75..4fa4186 100644 --- a/src/css/tables.rs +++ b/src/css/tables.rs @@ -44,7 +44,10 @@ pub const TAB2: [u8; 256] = [ 0xff, 0xfe, 0xfd, 0xfc, 0xfb, 0xfa, 0xf9, 0xf8, 0xf6, 0xf7, 0xf4, 0xf5, 0xf2, 0xf3, 0xf0, 0xf1, ]; -/// Table 3: LFSR1 low-byte feedback permutation. +/// Table 3: LFSR1 9-bit low-word feedback table (512 entries). +/// +/// Indexed by the 9-bit LFSR1 low word (the upper feedback bit makes the +/// index 9-bit, hence 512 entries, not 256). pub const TAB3: [u8; 512] = [ 0x00, 0x24, 0x49, 0x6d, 0x92, 0xb6, 0xdb, 0xff, 0x00, 0x24, 0x49, 0x6d, 0x92, 0xb6, 0xdb, 0xff, 0x00, 0x24, 0x49, 0x6d, 0x92, 0xb6, 0xdb, 0xff, 0x00, 0x24, 0x49, 0x6d, 0x92, 0xb6, 0xdb, 0xff, @@ -100,8 +103,10 @@ pub const TAB4: [u8; 256] = [ 0x0f, 0x8f, 0x4f, 0xcf, 0x2f, 0xaf, 0x6f, 0xef, 0x1f, 0x9f, 0x5f, 0xdf, 0x3f, 0xbf, 0x7f, 0xff, ]; -/// Table 5: LFSR1 output permutation for the Stevenson attack. -/// This is the inverse byte-reversal of TAB4. +/// Table 5: LFSR1 output permutation used in the keystream combiner. +/// `TAB5[i] == TAB4[i] ^ 0xFF` (bitwise complement of the TAB4 bit-reversal +/// table). Applied on the normal descramble/recrypt path (lfsr.rs) as well as +/// in the key-recovery fallback (crack.rs). pub const TAB5: [u8; 256] = [ 0xff, 0x7f, 0xbf, 0x3f, 0xdf, 0x5f, 0x9f, 0x1f, 0xef, 0x6f, 0xaf, 0x2f, 0xcf, 0x4f, 0x8f, 0x0f, 0xf7, 0x77, 0xb7, 0x37, 0xd7, 0x57, 0x97, 0x17, 0xe7, 0x67, 0xa7, 0x27, 0xc7, 0x47, 0x87, 0x07, @@ -120,3 +125,21 @@ pub const TAB5: [u8; 256] = [ 0xf8, 0x78, 0xb8, 0x38, 0xd8, 0x58, 0x98, 0x18, 0xe8, 0x68, 0xa8, 0x28, 0xc8, 0x48, 0x88, 0x08, 0xf0, 0x70, 0xb0, 0x30, 0xd0, 0x50, 0x90, 0x10, 0xe0, 0x60, 0xa0, 0x20, 0xc0, 0x40, 0x80, 0x00, ]; + +#[cfg(test)] +mod tests { + use super::*; + + /// Pins the documented relationship `TAB5[i] == TAB4[i] ^ 0xFF` so the + /// table doc cannot drift from the data. + #[test] + fn tab5_is_complement_of_tab4() { + for i in 0..256 { + assert_eq!( + TAB5[i], + TAB4[i] ^ 0xFF, + "TAB5[{i:#04x}] != TAB4[{i:#04x}] ^ 0xFF" + ); + } + } +} diff --git a/src/decrypt.rs b/src/decrypt.rs index b40f255..bd16cc0 100644 --- a/src/decrypt.rs +++ b/src/decrypt.rs @@ -60,8 +60,8 @@ static DECRYPT_POOL: RwLock<Option<Arc<rayon::ThreadPool>>> = RwLock::new(None); /// Configure how many threads to use for AACS unit decryption. A value /// of `0` resets to the env / default resolution. `1` forces serial. -/// `N > 1` builds a new rayon pool of size N and atomically replaces -/// the live pool. +/// `N > 1` builds a new rayon pool of size N (capped at [`MAX_THREADS`]) +/// and atomically replaces the live pool. /// /// Thread-safe. Live decrypt calls keep their previously-acquired /// pool reference for the rest of the call — no mid-call pool @@ -82,29 +82,36 @@ pub fn set_decrypt_threads(n: usize) { /// Get (or lazily build) the active rayon thread pool. Returns an /// `Arc` so in-flight work survives a concurrent /// [`set_decrypt_threads`] swap. -fn decrypt_pool() -> Arc<rayon::ThreadPool> { - // Fast path: pool already built. - if let Ok(guard) = DECRYPT_POOL.read() { +/// +/// Returns `None` if the pool cannot be built (e.g. the OS refuses the +/// worker threads under a pid/thread limit). The caller falls back to +/// the serial decrypt path — library code never panics here. +fn decrypt_pool() -> Option<Arc<rayon::ThreadPool>> { + // Fast path: pool already built. A poisoned read lock still yields a + // usable guard (the pool Arc is immutable once stored). + { + let guard = DECRYPT_POOL.read().unwrap_or_else(|e| e.into_inner()); if let Some(pool) = guard.as_ref() { - return Arc::clone(pool); + return Some(Arc::clone(pool)); } } - // Slow path: build a new one under the write lock. Double-check - // after acquiring in case another caller built it first. - let mut guard = DECRYPT_POOL.write().expect("DECRYPT_POOL RwLock poisoned"); + // Slow path: build a new one under the write lock. Recover the guard + // on poisoning (a prior panic) rather than propagating a secondary + // panic — we simply rebuild. Double-check after acquiring in case + // another caller built it first. + let mut guard = DECRYPT_POOL.write().unwrap_or_else(|e| e.into_inner()); if let Some(pool) = guard.as_ref() { - return Arc::clone(pool); + return Some(Arc::clone(pool)); } let n = decrypt_threads(); - let pool = Arc::new( - rayon::ThreadPoolBuilder::new() - .num_threads(n) - .thread_name(|i| format!("freemkv-decrypt-{i}")) - .build() - .expect("rayon decrypt pool build failed"), - ); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(n) + .thread_name(|i| format!("freemkv-decrypt-{i}")) + .build() + .ok() + .map(Arc::new)?; *guard = Some(Arc::clone(&pool)); - pool + Some(pool) } /// Current effective decrypt thread count. Resolution order: @@ -181,9 +188,46 @@ pub fn decrypt_sectors( }; let rdk: Option<[u8; 16]> = *read_data_key; let unit_len = aacs::ALIGNED_UNIT_LEN; + // AACS decrypts whole 6144-byte aligned units. The live mux path + // (mux/disc.rs::fill_extents) issues 1- or 2-sector reads at every + // extent tail, so a buffer is commonly NOT a multiple of the unit + // length. We process the whole leading units exactly as a fully + // aligned buffer would be, then make a deliberate decision about any + // trailing partial unit. + // + // Trailing-partial contract: + // * A clear partial (incomplete final unit / clear nav-TS tail) is + // what AACS legitimately leaves in the clear on disc, so we leave + // it untouched and return Ok. This is the proven, shipped + // behavior every production UHD MKV was made with — no regression + // on conformant discs. + // * A *scrambled* partial can only arise from a structurally + // malformed UDF layout that splits an encrypted unit across an + // extent boundary. Those bytes are encrypted content that cannot + // be decrypted standalone; passing them through as clear would be + // silent corruption. We fail loud (Error::DecryptFailed), matching + // the highway path's Error::ExtentNotUnitAligned policy. + // + // Detection: is_aacs_scrambled() short-circuits to false for any + // buffer shorter than a full unit, so it cannot judge a partial. We + // instead apply the same TS-sync-intactness test it uses internally + // (ts_sync_count vs ts_packet_total) directly to the available + // partial bytes. A clear TS tail carries 0x47 syncs at the 192-byte + // stride (> half the packets) → intact → not scrambled → tolerate. An + // encrypted tail has those syncs destroyed (≤ half) → scrambled → + // reject. If the partial is too short to hold even one TS packet + // (< 192 bytes, ts_packet_total == 0) we cannot judge confidently and + // tolerate rather than risk a false positive on conformant tails. + let partial_len = buf.len() % unit_len; + if partial_len != 0 { + let partial = &buf[buf.len() - partial_len..]; + let packets = aacs::ts_packet_total(partial); + if packets > 0 && aacs::ts_sync_count(partial) <= packets / 2 { + return Err(crate::error::Error::DecryptFailed); + } + } let nthreads = decrypt_threads(); - let chunks: Vec<&mut [u8]> = buf.chunks_mut(unit_len).collect(); - let nunits = chunks.len(); + let nunits = buf.len() / unit_len; // Per-unit decrypt closure. The is_aacs_scrambled check reads the // raw TS syncs; a non-m2ts unit (e.g. MPLS/CLPI nav file) can look @@ -202,22 +246,34 @@ pub fn decrypt_sectors( if nthreads <= 1 || nunits < PARALLEL_MIN_UNITS { // Serial path: avoids thread-pool overhead for tiny // buffers; also the only path when caller pinned - // single-threaded via FREEMKV_THREADS=1. - for chunk in chunks { + // single-threaded via FREEMKV_THREADS=1. Iterate the + // chunks directly — no Vec of slice pointers needed. + for chunk in buf.chunks_mut(unit_len) { decrypt_one(chunk); } } else { - // Parallel path via rayon's persistent global pool. - // The pool is built once on first use (lazy_static-style) - // and reused across every decrypt_sectors call — no - // per-call OS thread spawn, no thread-creation latency - // amortised per batch. Each unit decrypts independently - // (own key derivation), so par_iter is sound. - decrypt_pool().install(|| { - chunks.into_par_iter().for_each(|chunk| { - decrypt_one(chunk); - }); - }); + // Parallel path via rayon's persistent thread pool. + // The pool is built once on first use and reused across + // every decrypt_sectors call — no per-call OS thread + // spawn. Each unit decrypts independently (own key + // derivation), so par_iter is sound. On a pool-build + // failure (e.g. thread/pid-limit exhaustion) we fall + // back to the serial path rather than panic. + match decrypt_pool() { + Some(pool) => { + let chunks: Vec<&mut [u8]> = buf.chunks_mut(unit_len).collect(); + pool.install(|| { + chunks.into_par_iter().for_each(|chunk| { + decrypt_one(chunk); + }); + }); + } + None => { + for chunk in buf.chunks_mut(unit_len) { + decrypt_one(chunk); + } + } + } } } DecryptKeys::Css { title_key } => { @@ -260,4 +316,111 @@ mod tests { "non-m2ts unit must be restored after failed decrypt" ); } + + /// Build a clear-TS region: a 0x47 sync byte at offset 4 of every 192-byte + /// BD-TS packet (matching `ts_sync_count`'s probe stride), filler elsewhere. + /// Reads as NOT scrambled. + fn clear_ts_region(len: usize) -> Vec<u8> { + let mut v: Vec<u8> = (0..len).map(|i| (i as u8).wrapping_mul(31)).collect(); + let mut off = 4; + while off < len { + v[off] = 0x47; + off += 192; + } + v + } + + /// Build a scrambled region: the 192-byte-stride sync positions are NOT + /// 0x47 (encrypted content destroys them), so it reads as scrambled. + fn scrambled_region(len: usize) -> Vec<u8> { + let mut v: Vec<u8> = (0..len).map(|i| (i as u8).wrapping_mul(31)).collect(); + let mut off = 4; + while off < len { + // Force a non-sync byte at every probe position. + v[off] = 0xA5; + off += 192; + } + v + } + + /// Whole leading units plus a CLEAR trailing partial (the benign, + /// conformant case): AACS leaves an incomplete final unit / clear nav-TS + /// tail in the clear on disc. We must return `Ok` and leave the partial + /// bytes byte-for-byte unchanged — no regression on real discs. + #[test] + fn aacs_clear_trailing_partial_is_tolerated_unchanged() { + let keys = DecryptKeys::Aacs { + unit_keys: vec![(0, [0xAB; 16])], + read_data_key: None, + }; + // One full scrambled unit + a 2048-byte (single-sector) CLEAR tail. + let unit = scrambled_region(aacs::ALIGNED_UNIT_LEN); + let tail = clear_ts_region(2048); + let mut buf = unit; + buf.extend_from_slice(&tail); + + decrypt_sectors(&mut buf, &keys, 0).expect("clear trailing partial is Ok"); + + assert_eq!( + &buf[aacs::ALIGNED_UNIT_LEN..], + &tail[..], + "clear trailing partial unit must be left unchanged" + ); + } + + /// Whole leading units plus a SCRAMBLED trailing partial (the malformed + /// danger case): an encrypted unit split across an extent boundary cannot be + /// decrypted standalone. Passing it through as clear would be silent + /// corruption, so we must fail loud with `DecryptFailed`. + #[test] + fn aacs_scrambled_trailing_partial_is_rejected() { + let keys = DecryptKeys::Aacs { + unit_keys: vec![(0, [0xAB; 16])], + read_data_key: None, + }; + // One full unit + a 4096-byte (two-sector) SCRAMBLED tail. + let unit = clear_ts_region(aacs::ALIGNED_UNIT_LEN); + let tail = scrambled_region(4096); + let mut buf = unit; + buf.extend_from_slice(&tail); + + let err = decrypt_sectors(&mut buf, &keys, 0) + .expect_err("scrambled trailing partial must be rejected"); + assert_eq!( + err.code(), + crate::error::Error::DecryptFailed.code(), + "scrambled trailing partial must fail with DecryptFailed" + ); + } + + /// An empty buffer is a valid no-op (zero units), not an error. + #[test] + fn aacs_empty_buffer_is_ok() { + let keys = DecryptKeys::Aacs { + unit_keys: vec![(0, [0xAB; 16])], + read_data_key: None, + }; + let mut buf: Vec<u8> = Vec::new(); + assert!(decrypt_sectors(&mut buf, &keys, 0).is_ok()); + } + + /// An exact multiple of the unit length has no trailing partial: behavior + /// is unchanged — clear units stay clear, scrambled units are decrypt- + /// attempted. Two clear units must round-trip untouched and return `Ok`. + #[test] + fn aacs_exact_multiple_unchanged() { + let keys = DecryptKeys::Aacs { + unit_keys: vec![(0, [0xAB; 16])], + read_data_key: None, + }; + let mut buf = clear_ts_region(aacs::ALIGNED_UNIT_LEN * 2); + let snapshot = buf.clone(); + + decrypt_sectors(&mut buf, &keys, 0).expect("exact-multiple buffer is Ok"); + + assert_eq!( + buf, snapshot, + "clear exact-multiple buffer must be left unchanged" + ); + } } diff --git a/src/disc/bluray.rs b/src/disc/bluray.rs index a4ade26..f32d8e2 100644 --- a/src/disc/bluray.rs +++ b/src/disc/bluray.rs @@ -30,6 +30,12 @@ impl Disc { titles } + /// Parse one MPLS playlist into a [`DiscTitle`]. + /// + /// Sums PlayItem durations; returns `None` if the playlist is under + /// 30 seconds (skips menu / clip-info stub playlists) or fails to + /// parse. Physical sector extents are pulled from the UDF allocation + /// descriptors of each referenced `.m2ts` (deduplicated by clip_id). pub(super) fn parse_playlist( reader: &mut dyn SectorSource, udf_fs: &udf::UdfFs, @@ -55,27 +61,41 @@ impl Disc { let mut extents = Vec::new(); let mut total_size: u64 = 0; let mut clips = Vec::with_capacity(parsed.play_items.len()); + // BD playlists legally reference the same .m2ts clip_id from + // multiple PlayItems (multi-angle, seamless splits, looped + // segments). The physical extents and packet count must be + // counted ONCE per unique clip — mux reads extents in order, so + // a duplicate would mux the A/V twice and inflate size_bytes. + // Per-PlayItem Clip entries (differing in/out times) still get + // recorded. + let mut seen_clips: std::collections::HashSet<String> = std::collections::HashSet::new(); for play_item in &parsed.play_items { let clip_dur = play_item.out_time.saturating_sub(play_item.in_time) as f64 / 45000.0; let mut pkt_count: u32 = 0; + let first_ref = seen_clips.insert(play_item.clip_id.clone()); let clpi_path = format!("/BDMV/CLIPINF/{}.clpi", play_item.clip_id); if let Ok(clpi_data) = udf_fs.read_file(reader, &clpi_path) { if let Ok(clip_info) = clpi::parse(&clpi_data) { pkt_count = clip_info.source_packet_count; - total_size += pkt_count as u64 * 192; - // Get m2ts file extents from UDF allocation descriptors. - // Dual-layer discs split files across layers — UDF knows the real layout. - let m2ts_path = format!("/BDMV/STREAM/{}.m2ts", play_item.clip_id); - if let Ok(file_exts) = udf_fs.file_extents(reader, &m2ts_path) { - for (lba, sectors) in file_exts { - if sectors > 0 && lba > 0 { - extents.push(Extent { - start_lba: lba, - sector_count: sectors, - }); + // Only fetch/push the physical extents and add to the + // total size the first time this clip_id is seen. + if first_ref { + total_size += pkt_count as u64 * 192; + + // Get m2ts file extents from UDF allocation descriptors. + // Dual-layer discs split files across layers — UDF knows the real layout. + let m2ts_path = format!("/BDMV/STREAM/{}.m2ts", play_item.clip_id); + if let Ok(file_exts) = udf_fs.file_extents(reader, &m2ts_path) { + for (lba, sectors) in file_exts { + if sectors > 0 && lba > 0 { + extents.push(Extent { + start_lba: lba, + sector_count: sectors, + }); + } } } } @@ -170,23 +190,49 @@ impl Disc { }) .collect(); - // Convert marks to chapters (mark_type 0 or 1 = chapter entry, 2 = link) - let first_in_time = parsed.play_items.first().map(|pi| pi.in_time).unwrap_or(0); + // Convert marks to chapters. mark_type == 1 is an entry-mark + // (chapter); type 2 is a link point and type 0 is reserved, so + // neither is a chapter. + // + // Each mark's timestamp is in the timebase of the PlayItem it + // references (play_item_ref). The chapter's position on the + // muxed timeline is the summed duration of every preceding + // PlayItem plus the mark's offset within its own PlayItem. Using + // play_items[0].in_time for every mark would misplace chapters in + // multi-PlayItem playlists. let chapters: Vec<Chapter> = parsed .marks .iter() - .filter(|m| m.mark_type <= 1) - .enumerate() - .map(|(i, m)| { - let time_secs = (m.timestamp as f64 - first_in_time as f64) / 45000.0; - Chapter { + .filter(|m| m.mark_type == 1) + .filter_map(|m| { + let pi_idx = m.play_item_ref as usize; + let pi = parsed.play_items.get(pi_idx)?; + let preceding: f64 = parsed.play_items[..pi_idx] + .iter() + .map(|p| p.out_time.saturating_sub(p.in_time) as f64 / 45000.0) + .sum(); + let within = (m.timestamp as f64 - pi.in_time as f64) / 45000.0; + let time_secs = preceding + within; + Some(Chapter { time_secs: if time_secs < 0.0 { 0.0 } else { time_secs }, - name: format!("Chapter {}", i + 1), - } + name: String::new(), // filled with the ordinal below + }) + }) + .enumerate() + .map(|(i, mut ch)| { + ch.name = super::chapter_name(i); + ch }) .collect(); - let playlist_num = filename.trim_end_matches(".mpls").trim_end_matches(".MPLS"); + // Strip the .mpls suffix case-insensitively before parsing the + // numeric playlist id (the dir scan accepts any-case .mpls). + let playlist_num = filename + .get(..filename.len().saturating_sub(5)) + .filter(|_| { + filename.len() >= 5 && filename[filename.len() - 5..].eq_ignore_ascii_case(".mpls") + }) + .unwrap_or(filename); let playlist_id = playlist_num.parse::<u16>().unwrap_or(0); Some(DiscTitle { diff --git a/src/disc/dvd.rs b/src/disc/dvd.rs index c340878..faed991 100644 --- a/src/disc/dvd.rs +++ b/src/disc/dvd.rs @@ -34,15 +34,27 @@ impl Disc { label: String::new(), }); - // Map DvdAudioAttr to Stream::Audio + // Map DvdAudioAttr to Stream::Audio. The PID is derived from the + // stream's REAL on-wire private_stream_1 sub-stream id (assigned + // by per-codec ordinal in the IFO scan) via the same + // `dvd_audio_pid` table the demuxer's `PsPacket::dvd_pid` uses, + // so a mixed-codec title (AC-3 + DTS + LPCM) routes correctly + // instead of colliding on 0xBD00. Streams carried as a regular + // MPEG-audio PES (MP1/MP2, no sub-id) fall back to a distinct + // 0xBD00+ordinal PID — disjoint from the 0xBD80+ canonical audio + // space — though they are not routed via `dvd_pid` today. let audio_streams: Vec<Stream> = ts .audio_streams .iter() .enumerate() .map(|(i, a)| { let codec = a.codec; + let pid = a + .sub_stream_id + .and_then(crate::mux::ps::dvd_audio_pid) + .unwrap_or(0xBD00 + i as u16); Stream::Audio(AudioStream { - pid: 0xBD00 + i as u16, // DVD private stream 1 sub-IDs + pid, codec, channels: AudioChannels::from_count(a.channels), language: a.language.clone(), @@ -88,8 +100,13 @@ impl Disc { .iter() .enumerate() .map(|(i, s)| { + // VobSub sub-stream ids run 0x20..=0x3F; PID = sub-id + // (identity), shared with the demuxer via + // `dvd_subtitle_pid`. + let sub_id = 0x20u8.saturating_add(i.min(0x1F) as u8); + let pid = crate::mux::ps::dvd_subtitle_pid(sub_id).unwrap_or(sub_id as u16); Stream::Subtitle(SubtitleStream { - pid: 0x20 + i as u16, // DVD sub-stream IDs 0x20-0x3F + pid, codec: Codec::DvdSub, language: s.language.clone(), forced: false, @@ -109,7 +126,7 @@ impl Disc { .enumerate() .map(|(i, &t)| Chapter { time_secs: t, - name: format!("Chapter {}", i + 1), + name: chapter_name(i), }) .collect(); diff --git a/src/disc/encrypt.rs b/src/disc/encrypt.rs index c5f6912..91ab396 100644 --- a/src/disc/encrypt.rs +++ b/src/disc/encrypt.rs @@ -88,7 +88,7 @@ impl Disc { } let mut vid = [0u8; 16]; vid.copy_from_slice(&buf[4..20]); - tracing::warn!( + tracing::debug!( target: "freemkv::disc", phase = "oem_vid_ok", "OEM VID retrieved" @@ -119,10 +119,10 @@ impl Disc { /// /// Returns `(handshake, error)`: /// * `(Some(_), None)` — VID acquired - /// * `(None, Some(_))` — specific failure mode (see - /// `AacsHostCertRejected` / `AacsRawReadUnsupported` / - /// `AacsVidUnavailable` / `DriveProfileMissing` / - /// `VidCdbUnavailable` variants in `error.rs`) + /// * `(None, Some(_))` — specific failure mode; only + /// `AacsHostCertRejected` and `AacsVidUnavailable` are returned + /// here (the OEM-path `DriveProfileMissing` / `VidCdbUnavailable` + /// errors are caught internally and fall through to cert auth) /// * `(None, None)` — handshake not attempted (no keydb; /// resolution will proceed with VID=zero and rely on path 1 /// disc-hash → VUK lookup) @@ -131,7 +131,7 @@ impl Disc { opts: &ScanOptions, ) -> (Option<HandshakeResult>, Option<Error>) { let unlocked = session.is_unlocked(); - tracing::warn!( + tracing::debug!( target: "freemkv::disc", phase = "handshake_entry", unlocked, @@ -199,25 +199,24 @@ impl Disc { }; let host_cert_count = host_certs.len(); - tracing::warn!( + tracing::debug!( target: "freemkv::disc", phase = "handshake_start", host_cert_count, "handshake starting" ); - // v0.25.7 wedge fix. Pre-0.25.7 this loop fired up to 16 AACS - // authenticate attempts back-to-back with no pause. Each attempt - // is 5-10 SCSI REPORT_KEY/SEND_KEY exchanges. On a disc whose - // host cert isn't in our KEYDB (or one the drive rejects), - // that's 80-160 SCSI commands hammered at the drive in a - // few hundred milliseconds — and the BU40N (and most consumer - // optical drives) responds by entering a fast-fail firmware - // wedge state where every subsequent CDB returns - // ILLEGAL_REQUEST/INVALID_FIELD_IN_CDB (sense 05/24) until - // power-cycled. Hit live on rip1 2026-05-20 during a MOVIE - // UHD scan: KEYDB miss → 16 cert attempts in a tight loop → - // wedge → forced host reboot + drive disconnect to recover. + // Cert-attempt wedge guard. An earlier version fired up to 16 + // AACS authenticate attempts back-to-back with no pause. Each + // attempt is 5-10 SCSI REPORT_KEY/SEND_KEY exchanges. On a disc + // whose host cert isn't in the KEYDB (or one the drive rejects), + // that's 80-160 SCSI commands hammered at the drive in a few + // hundred milliseconds — and consumer optical drives can respond + // by entering a fast-fail firmware wedge state where every + // subsequent CDB returns ILLEGAL_REQUEST/INVALID_FIELD_IN_CDB + // (sense 05/24) until power-cycled. Observed live on a UHD scan: + // KEYDB miss → many cert attempts in a tight loop → wedge → + // forced power cycle to recover. // // Defense-in-depth: cap attempts, sleep between, and bail // early on the drive's wedge sense so any later regression @@ -262,20 +261,31 @@ impl Disc { ); } Err(e) => { - let code = e.code(); - last_err_code = Some(code); - // Drive wedge senses (any with high byte 0x05 = - // ILLEGAL_REQUEST). The drive isn't merely - // rejecting our cert — it's saying "I won't talk - // to you anymore." Trying more certs makes the - // wedge worse. Bail out immediately. - let sense_key = ((code >> 8) & 0xFF) as u8; - if sense_key == 0x05 { + last_err_code = Some(e.code()); + // Log the real SCSI sense triple, not `e.code()` — + // `code()` collapses every ScsiError to the flat + // E_SCSI_ERROR constant and carries no sense key, + // so it has no diagnostic value for auth-failure + // routing. + let sense = e.scsi_sense(); + // Drive wedge senses (ILLEGAL_REQUEST, sense key + // 0x05). The drive isn't merely rejecting our + // cert — it's signalling it won't talk to us + // anymore. Trying more certs makes the wedge worse, + // so bail out immediately. NOTE: this must read the + // sense key off the structured ScsiSense, NOT off + // `e.code()`; `code()` is a flat constant for every + // ScsiError so the old `(code >> 8) & 0xFF` guard + // never matched and was dead code (the very wedge + // this defense exists to prevent could recur). + if sense.map(|s| s.is_illegal_request()).unwrap_or(false) { tracing::warn!( target: "freemkv::disc", phase = "handshake_wedge_detected", cert_index = idx, - error_code = code, + sense_key = sense.map(|s| s.sense_key), + asc = sense.map(|s| s.asc), + ascq = sense.map(|s| s.ascq), "drive returned ILLEGAL_REQUEST during auth; bailing out to avoid wedge" ); return (None, Some(Error::AacsHostCertRejected)); @@ -339,13 +349,15 @@ impl Disc { .or_else(|_| udf_fs.read_file(reader, "/AACS/MKB_RW.inf")) .ok() .unwrap_or_default(); + // Trim to the real record length. truncate is a no-op when n >= + // len and correctly empties the vec when n == 0 (zeroed/corrupt + // MKB), so it never leaves the full ~128 MiB zero-pad on + // AacsState.mkb. let n = aacs::mkb_content_len(&mkb_bytes); - if n > 0 && n < mkb_bytes.len() { - mkb_bytes.truncate(n); - } + mkb_bytes.truncate(n); let mkb_ver = aacs::mkb_version(&mkb_bytes); - tracing::warn!( + tracing::debug!( target: "freemkv::disc", phase = "scan_aacs_vid_only", disc_hash = %aacs::disc_hash_hex(&dh), diff --git a/src/disc/mapfile.rs b/src/disc/mapfile.rs index 4c4974f..3f8ab7a 100644 --- a/src/disc/mapfile.rs +++ b/src/disc/mapfile.rs @@ -6,7 +6,7 @@ //! //! Format: //! ```text -//! # Rescue Logfile. Created by libfreemkv v0.11.21 +//! # Rescue Logfile. Created by libfreemkv vX.Y.Z //! # Current pos / status / pass / pass_time (ddrescue state machine — we only populate pos) //! 0x000000000 ? 1 0 //! # pos size status @@ -50,6 +50,8 @@ pub enum SectorStatus { } impl SectorStatus { + /// The single ddrescue status character for this status + /// (`?`/`*`/`/`/`-`/`+`). pub fn to_char(self) -> char { match self { Self::NonTried => '?', @@ -59,6 +61,8 @@ impl SectorStatus { Self::Finished => '+', } } + /// Parse a ddrescue status character into a `SectorStatus`. Returns + /// `None` for any character that is not one of `?*/-+`. pub fn from_char(c: char) -> Option<Self> { Some(match c { '?' => Self::NonTried, @@ -102,8 +106,8 @@ pub struct MapStats { /// retry" UI bucket; `bytes_pending` over-counts because it folds /// in `bytes_nontried`. pub bytes_retryable: u64, - /// Number of unreadable ranges (for UI display). Computed from - /// `ranges_with(&[Unreadable])`. + /// Number of distinct `Unreadable` ranges (for UI display). + /// Computed by `compute_stats` (counts coalesced `-` entries). pub num_bad_ranges: u32, /// Largest gap among unreadable ranges in milliseconds. Computed as /// largest range size / bytes_per_sec * 1000. Set by caller (autorip) @@ -231,6 +235,16 @@ impl Mapfile { } let pos = parse_hex(fields[0])?; let size = parse_hex(fields[1])?; + // Reject an entry whose pos+size overflows u64 up front. The + // downstream overlap/coalesce/next_with code adds pos+size + // freely; a crafted/corrupt line like + // `0xfffffffffffffff0 0x20 +` would otherwise panic (debug) + // or wrap to a tiny range (release), corrupting stats and + // resume logic. + if pos.checked_add(size).is_none() { + let e: io::Error = crate::error::Error::MapfileInvalid { kind: "range" }.into(); + return Err(e); + } let status = fields[2] .chars() .next() @@ -247,7 +261,31 @@ impl Mapfile { entries.push(MapEntry { pos, size, status }); } entries.sort_by_key(|e| e.pos); - let total_size = entries.last().map(|e| e.pos + e.size).unwrap_or(0); + // Reject overlapping ranges. A well-formed ddrescue mapfile is a + // disjoint partition; overlaps (from a corrupt/hand-edited file) + // would make compute_stats double-count, so bytes_good / + // bytes_unreadable / bytes_pending could exceed bytes_total and + // inflate resume / abort-on-loss decisions and >100% progress. + for pair in entries.windows(2) { + let prev_end = pair[0].pos.saturating_add(pair[0].size); + if prev_end > pair[1].pos { + let e: io::Error = crate::error::Error::MapfileInvalid { kind: "overlap" }.into(); + return Err(e); + } + } + let total_size = entries + .last() + .map(|e| e.pos.saturating_add(e.size)) + .unwrap_or(0); + // Enforce the keys-XOR-vid invariant that set_unit_keys() + // guarantees: a corrupt/hand-edited file carrying both comment + // types would otherwise load with vid=Some AND non-empty + // unit_keys, violating the invariant downstream code relies on. + // Unit keys win, matching the setter (it clears vid when keys + // are present). + if !unit_keys.is_empty() { + vid = None; + } let stats = Self::compute_stats(&entries, total_size); Ok(Self { path: path.to_path_buf(), @@ -265,7 +303,26 @@ impl Mapfile { /// Load if the file exists, otherwise create a fresh mapfile. pub fn open_or_create(path: &Path, total_size: u64, version: &str) -> io::Result<Self> { match Self::load(path) { - Ok(mf) => Ok(mf), + Ok(mf) => { + // load() derives total_size from the last entry's + // pos+size; if that disagrees with the caller's + // expected disc size (different disc, edited/partial + // file, trimmed trailing region) the downstream + // resume/progress math keys off the wrong basis. Surface + // it so an operator can spot a mismatched mapfile rather + // than failing the resume outright. + if mf.total_size != total_size { + tracing::warn!( + target: "freemkv::disc", + phase = "mapfile_total_size_mismatch", + loaded_total = mf.total_size, + supplied_total = total_size, + path = %path.display(), + "loaded mapfile coverage differs from supplied disc size" + ); + } + Ok(mf) + } Err(e) if e.kind() == io::ErrorKind::NotFound => { Self::create(path, total_size, version) } @@ -280,11 +337,18 @@ impl Mapfile { if size == 0 { return Ok(()); } - let end = pos.saturating_add(size); + // Mirror load()'s overflow contract: reject a range that would + // wrap u64 rather than storing a saturated entry narrower than + // its size, which load() would then reject on the next resume + // (making the mapfile unreadable). + let Some(end) = pos.checked_add(size) else { + let e: io::Error = crate::error::Error::MapfileInvalid { kind: "range" }.into(); + return Err(e); + }; let mut new_entries = Vec::with_capacity(self.entries.len() + 2); for e in self.entries.drain(..) { - let e_end = e.pos + e.size; + let e_end = e.pos.saturating_add(e.size); if e_end <= pos || e.pos >= end { // entirely before or after — keep new_entries.push(e); @@ -313,8 +377,8 @@ impl Mapfile { let mut merged: Vec<MapEntry> = Vec::with_capacity(new_entries.len()); for e in new_entries { if let Some(last) = merged.last_mut() { - if last.pos + last.size == e.pos && last.status == e.status { - last.size += e.size; + if last.pos.saturating_add(last.size) == e.pos && last.status == e.status { + last.size = last.size.saturating_add(e.size); continue; } } @@ -383,10 +447,13 @@ impl Mapfile { &self.unit_keys } + /// All map entries, sorted ascending by `pos` and (after load) + /// guaranteed disjoint and non-overflowing. pub fn entries(&self) -> &[MapEntry] { &self.entries } + /// Total image size in bytes, i.e. the end byte of the last entry. pub fn total_size(&self) -> u64 { self.total_size } @@ -397,7 +464,7 @@ impl Mapfile { if e.status != status { continue; } - let e_end = e.pos + e.size; + let e_end = e.pos.saturating_add(e.size); if e_end <= from { continue; } @@ -416,6 +483,8 @@ impl Mapfile { .collect() } + /// Snapshot of the incrementally-maintained summary statistics. + /// O(1) — returns the cached `MapStats`. pub fn stats(&self) -> MapStats { self.stats } @@ -428,7 +497,10 @@ impl Mapfile { for e in entries { match e.status { SectorStatus::Finished => s.bytes_good += e.size, - SectorStatus::Unreadable => s.bytes_unreadable += e.size, + SectorStatus::Unreadable => { + s.bytes_unreadable += e.size; + s.num_bad_ranges += 1; + } SectorStatus::NonTried => { s.bytes_pending += e.size; s.bytes_nontried += e.size; @@ -514,16 +586,36 @@ impl Drop for Mapfile { /// error, so a corrupt header never fails a mapfile load. fn parse_vid_hex(s: &str) -> Option<[u8; 16]> { let s = s.strip_prefix("0x").unwrap_or(s); - if s.len() != 32 { + // Parse on bytes, not on the &str: slicing a &str by byte index + // (`&s[i*2..i*2+2]`) panics when the cut lands inside a multi-byte + // UTF-8 char. A hand-edited/corrupt `# freemkv-vid:` comment of + // exactly 32 bytes containing a multi-byte char would otherwise + // kill the whole load. ASCII hex is one byte per char, so anything + // non-ASCII is simply rejected here as malformed. + let bytes = s.as_bytes(); + if bytes.len() != 32 { return None; } let mut out = [0u8; 16]; for (i, b) in out.iter_mut().enumerate() { - *b = u8::from_str_radix(&s[i * 2..i * 2 + 2], 16).ok()?; + let hi = hex_nibble(bytes[i * 2])?; + let lo = hex_nibble(bytes[i * 2 + 1])?; + *b = (hi << 4) | lo; } Some(out) } +/// Map a single ASCII hex digit byte to its 0-15 value. Returns `None` +/// for any non-hex byte (including any non-ASCII / multi-byte lead byte). +fn hex_nibble(c: u8) -> Option<u8> { + match c { + b'0'..=b'9' => Some(c - b'0'), + b'a'..=b'f' => Some(c - b'a' + 10), + b'A'..=b'F' => Some(c - b'A' + 10), + _ => None, + } +} + /// Parse a `# freemkv-uk:` value `<cps>:<32hex>` into `(cps_unit, key)`. Returns /// `None` on any malformation so a corrupt line is ignored, never fatal. fn parse_uk_line(s: &str) -> Option<(u32, [u8; 16])> { @@ -557,7 +649,9 @@ mod tests { tag, n ); - std::env::temp_dir().join(name) + let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("target/test-scratch"); + let _ = std::fs::create_dir_all(&dir); + dir.join(name) } #[test] @@ -758,6 +852,54 @@ mod tests { let _ = std::fs::remove_file(&p2); } + #[test] + fn load_rejects_entry_whose_range_overflows_u64() { + let p = tmpfile("load_overflow"); + let _ = std::fs::remove_file(&p); + // pos near u64::MAX with a nonzero size overflows pos+size. + let body = format!("0x{:x} 0x10 +\n", u64::MAX - 4); + std::fs::write(&p, body).unwrap(); + let kind = match Mapfile::load(&p) { + Ok(_) => panic!("overflowing entry must be rejected"), + Err(e) => e.kind(), + }; + assert_eq!(kind, io::ErrorKind::InvalidData); + let _ = std::fs::remove_file(&p); + } + + #[test] + fn record_rejects_range_overflowing_u64() { + let p = tmpfile("record_overflow"); + let _ = std::fs::remove_file(&p); + let mut mf = Mapfile::create(&p, 1000, "test").unwrap(); + let err = mf + .record(u64::MAX - 4, 16, SectorStatus::Finished) + .expect_err("overflowing record must be rejected"); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + let _ = std::fs::remove_file(&p); + } + + #[test] + fn load_enforces_keys_xor_vid_on_malformed_file() { + let p = tmpfile("load_keys_xor_vid"); + let _ = std::fs::remove_file(&p); + // Hand-craft a file carrying BOTH a vid comment and a uk comment + // (which write_to_disk would never emit together). load() must + // resolve to keys-only, matching set_unit_keys()'s invariant. + let body = "# freemkv-vid:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n\ + # freemkv-uk: 0:11111111111111111111111111111111\n\ + 0x0 0x200 +\n"; + std::fs::write(&p, body).unwrap(); + let loaded = Mapfile::load(&p).unwrap(); + assert_eq!( + loaded.vid(), + None, + "load() must clear vid when unit keys are present" + ); + assert_eq!(loaded.unit_keys(), &[(0u32, [0x11u8; 16])]); + let _ = std::fs::remove_file(&p); + } + #[test] fn vid_round_trips_and_data_lines_unaffected() { let p = tmpfile("vid_round_trips"); @@ -831,6 +973,72 @@ mod tests { let _ = std::fs::remove_file(&resaved); } + #[test] + fn parse_vid_hex_does_not_panic_on_multibyte_32_byte_input() { + // A 32-BYTE comment containing a multi-byte char would make the + // old `&s[i*2..i*2+2]` slice fall inside a char boundary and + // panic. Must return None instead. + let s = "中".to_string() + &"a".repeat(29); // 3 + 29 = 32 bytes + assert_eq!(s.len(), 32); + assert_eq!(parse_vid_hex(&s), None); + // A valid 32-char ASCII hex string still parses. + assert_eq!( + parse_vid_hex("00112233445566778899aabbccddeeff"), + Some([ + 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, + 0xee, 0xff, + ]) + ); + } + + #[test] + fn load_rejects_overflowing_pos_plus_size() { + let p = tmpfile("load_rejects_overflow"); + let _ = std::fs::remove_file(&p); + std::fs::write( + &p, + "# Rescue Logfile. Created by test\n\ + 0x000000000 ? 1 0\n\ + 0xfffffffffffffff0 0x20 +\n", + ) + .unwrap(); + assert!( + Mapfile::load(&p).is_err(), + "a pos+size that overflows u64 must be rejected, not wrap" + ); + let _ = std::fs::remove_file(&p); + } + + #[test] + fn load_rejects_overlapping_ranges() { + let p = tmpfile("load_rejects_overlap"); + let _ = std::fs::remove_file(&p); + std::fs::write( + &p, + "# Rescue Logfile. Created by test\n\ + 0x000000000 ? 1 0\n\ + 0x000000000 0x00000100 +\n\ + 0x000000080 0x00000100 -\n", + ) + .unwrap(); + assert!( + Mapfile::load(&p).is_err(), + "overlapping ranges must be rejected so stats can't double-count" + ); + let _ = std::fs::remove_file(&p); + } + + #[test] + fn num_bad_ranges_counts_unreadable_entries() { + let p = tmpfile("num_bad_ranges"); + let _ = std::fs::remove_file(&p); + let mut mf = Mapfile::create(&p, 1000, "test").unwrap(); + mf.record(100, 50, SectorStatus::Unreadable).unwrap(); + mf.record(300, 50, SectorStatus::Unreadable).unwrap(); + assert_eq!(mf.stats().num_bad_ranges, 2); + let _ = std::fs::remove_file(&p); + } + #[test] fn stats_consistent_after_split_record() { let p = tmpfile("stats_consistent_after_split"); diff --git a/src/disc/mod.rs b/src/disc/mod.rs index a9e298a..db84dc3 100644 --- a/src/disc/mod.rs +++ b/src/disc/mod.rs @@ -364,10 +364,19 @@ pub enum ColorSpace { pub struct Chapter { /// Chapter start time in seconds pub time_secs: f64, - /// Chapter name (e.g. "Chapter 1", "Chapter 2") + /// Chapter name — a bare 1-based index ("1", "2", …). The library + /// emits no localized prose; consuming apps prepend any "Chapter " + /// prefix in the user's language. pub name: String, } +/// Default chapter name for the 0-based chapter index `i`: the bare +/// 1-based ordinal as a string. Keeps chapter labelling language-neutral +/// (apps localize) and gives BD and DVD a single source of truth. +pub(crate) fn chapter_name(i: usize) -> String { + (i + 1).to_string() +} + /// A contiguous range of sectors on disc. #[derive(Debug, Clone, Copy)] pub struct Extent { @@ -454,24 +463,24 @@ pub fn bytes_bad_in_title(title: &DiscTitle, bad_ranges: &[(u64, u64)]) -> u64 { if bad_ranges.is_empty() || title.extents.is_empty() { return 0; } - let t_start = title.extents.first().map(|e| (e.start_lba as u64) * 2048); - let t_end = title - .extents - .last() - .map(|e| ((e.start_lba as u64) + (e.sector_count as u64)) * 2048); - let (Some(ts), Some(te)) = (t_start, t_end) else { - return 0; - }; - bad_ranges - .iter() - .map(|(pos, size)| { + // Overlap each bad range against every extent individually. A single + // bounding box (first extent start → last extent end) would count + // bad sectors in inter-extent gaps (other titles' data, BDMV + // metadata) as bad bytes in this title, over-counting lost_ms for + // titles with non-contiguous clips. + let mut total: u64 = 0; + for ext in &title.extents { + let es = (ext.start_lba as u64) * 2048; + let ee = ((ext.start_lba as u64) + (ext.sector_count as u64)) * 2048; + for (pos, size) in bad_ranges { let r_start = *pos; - let r_end = *pos + *size; - let overlap_start = r_start.max(ts); - let overlap_end = r_end.min(te); - overlap_end.saturating_sub(overlap_start) - }) - .sum() + let r_end = pos.saturating_add(*size); + let overlap_start = r_start.max(es); + let overlap_end = r_end.min(ee); + total = total.saturating_add(overlap_end.saturating_sub(overlap_start)); + } + } + total } // ─── Display helpers ──────────────────────────────────────────────────────── @@ -535,7 +544,10 @@ impl Codec { 0x81 => Codec::Ac3, 0x84 | 0xA1 => Codec::Ac3Plus, 0x80 => Codec::Lpcm, - 0xA2 => Codec::DtsHdHr, + // 0x86 (primary) / 0xA2 (secondary) are the DTS-HD MA + // lossless pair, parallel to 0x81/0xA1 for AC-3. 0xA2 is + // lossless MA, not lossy HR. + 0xA2 => Codec::DtsHdMa, 0x90 | 0x91 => Codec::Pgs, ct => Codec::Unknown(ct), } @@ -693,7 +705,6 @@ impl AudioChannels { 3 => AudioChannels::Stereo, 6 => AudioChannels::Surround51, 12 => AudioChannels::Surround71, - _ if af > 0 => AudioChannels::Unknown, _ => AudioChannels::Unknown, } } @@ -1176,21 +1187,19 @@ impl Disc { Ok((capacity, buffered, udf_fs)) } - /// Scan a disc -- parse filesystem, playlists, streams, and set up AACS decryption. + /// Scan a disc — parse filesystem, playlists, streams, and set up + /// AACS decryption. This is the main entry point; after `scan()` the + /// Disc is ready (titles populated with streams, AACS inputs + /// captured, content readable and decryptable transparently). /// - /// This is the main entry point. After scan(), the Disc is ready: - /// - titles are populated with streams - /// - AACS keys are derived (if KEYDB available) - /// - content can be read and decrypted transparently - /// - /// Scan a disc. One pipeline, one order: + /// One pipeline, one order: /// 1. Read capacity + UDF filesystem /// 2. AACS handshake + key resolution /// 3. Parse playlists + streams /// 4. Apply labels /// - /// The session must be open and unlocked (Drive::open handles this). - /// All disc reads use standard READ(10) via UDF -- no vendor SCSI commands. + /// The session must be open and unlocked (`Drive::open` handles this). + /// All disc reads use standard READ(10) via UDF — no vendor SCSI commands. pub fn scan(session: &mut Drive, opts: &ScanOptions) -> Result<Self> { // AACS handshake (Blu-ray/UHD). Routes through Disc::read_vid, // which prefers the per-drive OEM CDB path when the drive is @@ -1280,58 +1289,52 @@ impl Disc { Self::scan_with(reader, capacity, None, None, opts, udf_fs) } + /// Read a disc's AACS key-input files from a sector source: returns + /// `(Unit_Key_RO.inf, MKB)` raw bytes. Shared body for + /// [`Disc::read_aacs_inputs`] (ISO) and + /// [`Disc::read_aacs_inputs_from_drive`] (live drive). + /// + /// Prefers MKB_RO, falls back to MKB_RW, then TRIMS to the real + /// record length. Both files are allocated to a fixed ~128 MiB and + /// zero-padded, so reading either ships up to ~124 MiB of nothing — + /// trim to the record stream so callers send/store a few MB, not + /// 128 MiB. + fn read_aacs_inputs_from_reader( + reader: &mut dyn SectorSource, + udf_fs: &udf::UdfFs, + ) -> Result<(Vec<u8>, Vec<u8>)> { + let inf = udf_fs + .read_file(reader, "/AACS/Unit_Key_RO.inf") + .or_else(|_| udf_fs.read_file(reader, "/AACS/DUPLICATE/Unit_Key_RO.inf")) + .map_err(|_| Error::AacsNoKeys)?; + let mut mkb = udf_fs + .read_file(reader, "/AACS/MKB_RO.inf") + .or_else(|_| udf_fs.read_file(reader, "/AACS/MKB_RW.inf")) + .map_err(|_| Error::AacsNoKeys)?; + let n = crate::aacs::mkb_content_len(&mkb); + mkb.truncate(n); + Ok((inf, mkb)) + } + /// Read a disc's AACS key-input files from an ISO image: returns /// `(Unit_Key_RO.inf, MKB)` raw bytes. For callers that resolve a Unit Key - /// out-of-band: obtain the key however you like, then scan with - /// `ScanOptions { unit_key: Some(uk), .. }`. libfreemkv never makes a - /// network call. + /// out-of-band: obtain the key however you like, then apply it via + /// [`Disc::decrypt_with`]. libfreemkv never makes a network call. pub fn read_aacs_inputs(iso_path: &std::path::Path) -> Result<(Vec<u8>, Vec<u8>)> { let mut reader = crate::io::file_sector_source::FileSectorSource::open(iso_path) .map_err(|_| Error::AacsNoKeys)?; let udf_fs = udf::read_filesystem(&mut reader)?; - let inf = udf_fs - .read_file(&mut reader, "/AACS/Unit_Key_RO.inf") - .or_else(|_| udf_fs.read_file(&mut reader, "/AACS/DUPLICATE/Unit_Key_RO.inf")) - .map_err(|_| Error::AacsNoKeys)?; - // Prefer MKB_RO, fall back to MKB_RW, then TRIM to the real record - // length. Both files are allocated to a fixed ~128 MiB and zero-padded, - // so reading either ships up to ~124 MiB of nothing — trim to the - // record stream so callers send/store a few MB, not 128 MiB. - let mut mkb = udf_fs - .read_file(&mut reader, "/AACS/MKB_RO.inf") - .or_else(|_| udf_fs.read_file(&mut reader, "/AACS/MKB_RW.inf")) - .map_err(|_| Error::AacsNoKeys)?; - let n = crate::aacs::mkb_content_len(&mkb); - if n > 0 && n < mkb.len() { - mkb.truncate(n); - } - Ok((inf, mkb)) + Self::read_aacs_inputs_from_reader(&mut reader, &udf_fs) } /// Same as [`Disc::read_aacs_inputs`] but reads from a live drive. The /// out-of-band Unit Key path fetches the disc's key files from the drive, - /// resolves a key from them however it likes, then scans with - /// `ScanOptions { unit_key: Some(uk), .. }`. These files are plaintext UDF - /// metadata — no AACS handshake or keys are required to read them. + /// resolves a key from them however it likes, then applies it via + /// [`Disc::decrypt_with`]. These files are plaintext UDF metadata — no + /// AACS handshake or keys are required to read them. pub fn read_aacs_inputs_from_drive(drive: &mut Drive) -> Result<(Vec<u8>, Vec<u8>)> { let (_, mut reader, udf_fs) = Self::read_udf(drive)?; - let inf = udf_fs - .read_file(&mut reader, "/AACS/Unit_Key_RO.inf") - .or_else(|_| udf_fs.read_file(&mut reader, "/AACS/DUPLICATE/Unit_Key_RO.inf")) - .map_err(|_| Error::AacsNoKeys)?; - // Prefer MKB_RO, fall back to MKB_RW, then TRIM to the real record - // length. Both files are allocated to a fixed ~128 MiB and zero-padded, - // so reading either ships up to ~124 MiB of nothing — trim to the - // record stream so callers send/store a few MB, not 128 MiB. - let mut mkb = udf_fs - .read_file(&mut reader, "/AACS/MKB_RO.inf") - .or_else(|_| udf_fs.read_file(&mut reader, "/AACS/MKB_RW.inf")) - .map_err(|_| Error::AacsNoKeys)?; - let n = crate::aacs::mkb_content_len(&mkb); - if n > 0 && n < mkb.len() { - mkb.truncate(n); - } - Ok((inf, mkb)) + Self::read_aacs_inputs_from_reader(&mut reader, &udf_fs) } /// Core scan pipeline — works with any SectorSource. @@ -2119,13 +2122,11 @@ impl Disc { let pipe: Pipeline<WorkItem, sweep::ConsumerSummary> = Pipeline::spawn_named("freemkv-sweep-consumer", DEFAULT_PIPELINE_DEPTH, sink)?; - // Translate `Pipeline::send` failure (consumer gone) into the - // same `Error` shape the 0.17.x `send_or_abort` produced, so - // the producer-error semantics are unchanged. + // Translate `Pipeline::send` failure (consumer gone) into a + // numeric library error so the producer-error semantics are + // unchanged but no English leaks into an io::Error. fn consumer_gone() -> Error { - Error::IoError { - source: std::io::Error::other("sweep consumer terminated unexpectedly"), - } + Error::PipelineConsumerGone } let mut buf = vec![0u8; batch as usize * 2048]; @@ -2637,6 +2638,10 @@ pub(crate) fn sleep_secs_or_halt( } } +/// Mapfile path for a regular output file: appends `.mapfile` to the +/// output path. For `/dev/null` (benchmark) output use +/// [`Disc::mapfile_for`], which special-cases it to a temp-dir path +/// derived from the disc title. pub fn mapfile_path_for(iso_path: &std::path::Path) -> std::path::PathBuf { let mut s = iso_path.as_os_str().to_os_string(); s.push(".mapfile"); @@ -2646,8 +2651,10 @@ pub fn mapfile_path_for(iso_path: &std::path::Path) -> std::path::PathBuf { impl Disc { /// Path to the mapfile for a given output path. /// - /// For `/dev/null` output, returns `/tmp/{volume_id_or_title}.mapfile`. - /// For regular files, returns `{path}.mapfile`. + /// For `/dev/null` output, returns + /// `{temp_dir}/{volume_id_or_title}.mapfile` (temp dir is + /// `TMPDIR`-aware and cross-platform). For regular files, returns + /// `{path}.mapfile`. pub fn mapfile_for(&self, path: &std::path::Path) -> std::path::PathBuf { if path.as_os_str() == "/dev/null" { let name: String = self @@ -2663,7 +2670,7 @@ impl Disc { } }) .collect(); - std::path::PathBuf::from(format!("/tmp/{name}.mapfile")) + std::env::temp_dir().join(format!("{name}.mapfile")) } else { mapfile_path_for(path) } @@ -2755,25 +2762,25 @@ pub fn detect_max_batch_sectors(device_path: &str) -> u16 { return DEFAULT_BATCH_SECTORS_OPTICAL; } - // Check if optical drive (0x05 = CD/DVD) - let is_optical = (|| -> bool { - use std::path::Path; - let scsi_device_dir = "/sys/class/scsi_device/".to_string(); - if let Ok(entries) = std::fs::read_dir(&scsi_device_dir) { - for entry in entries.flatten() { - let device_type_path = entry.path().join("device/type"); - if Path::new(&device_type_path).exists() { - if let Ok(content) = std::fs::read_to_string(&device_type_path) { - // Type 0x05 (decimal 5) = CD/DVD drive - if content.trim().parse::<u32>() == Ok(5) { - return true; - } - } - } - } - } - false - })(); + // Check whether THIS device (not any device on the host) is an + // optical drive: read the SCSI peripheral type of the target node + // only. Type 0x05 (decimal 5) = CD/DVD. A previous version scanned + // every /sys/class/scsi_device entry and returned true if any was + // optical, misclassifying a block device as optical on a host that + // also has an optical drive. + let is_optical = { + // For an sg node the type lives at scsi_generic/<sg>/device/type; + // for a block node (sr0/sdX) at /sys/block/<name>/device/type. + let type_path = if dev_name.starts_with("sg") { + format!("/sys/class/scsi_generic/{dev_name}/device/type") + } else { + format!("/sys/block/{dev_name}/device/type") + }; + std::fs::read_to_string(&type_path) + .ok() + .map(|c| c.trim().parse::<u32>() == Ok(5)) + .unwrap_or(false) + }; if is_optical { // For sg devices, find the corresponding block device name @@ -3417,7 +3424,6 @@ mod tests { #[test] fn sweep_to_dev_null_real() { - let _cleanup = CleanupGuard(std::path::PathBuf::from("/tmp/T2.mapfile")); let sectors: u32 = 1000; let bad: std::collections::HashSet<u32> = [500u32, 501, 502].into_iter().collect(); let mut reader = MockReader { @@ -3425,6 +3431,7 @@ mod tests { bad_sectors: bad, }; let disc = make_test_disc(sectors, "T2"); + let _cleanup = CleanupGuard(disc.mapfile_for(std::path::Path::new("/dev/null"))); let opts = CopyOptions { decrypt: false, multipass: true, @@ -3450,13 +3457,13 @@ mod tests { #[test] fn sweep_dev_null_full_good() { - let _cleanup = CleanupGuard(std::path::PathBuf::from("/tmp/T3.mapfile")); let sectors: u32 = 2000; let mut reader = MockReader { total_sectors: sectors, bad_sectors: std::collections::HashSet::new(), }; let disc = make_test_disc(sectors, "T3"); + let _cleanup = CleanupGuard(disc.mapfile_for(std::path::Path::new("/dev/null"))); let opts = CopyOptions { decrypt: false, multipass: false, @@ -3615,4 +3622,51 @@ mod tests { let meta = std::fs::metadata(&iso_path).unwrap(); assert_eq!(meta.len(), sectors as u64 * 2048); } + + /// bytes_bad_in_title must overlap per-extent, not against a single + /// bounding box: a bad range in the gap between two extents of the + /// same title must NOT be counted. + #[test] + fn bytes_bad_in_title_ignores_inter_extent_gap() { + let mut title = title_with_video(Codec::Hevc, Resolution::R2160p); + // Two extents: sectors [0,10) and [100,110). Gap = [10,100). + title.extents = vec![ + Extent { + start_lba: 0, + sector_count: 10, + }, + Extent { + start_lba: 100, + sector_count: 10, + }, + ]; + // A bad range entirely inside the gap (sector 50 == byte 50*2048). + let gap = vec![(50 * 2048, 2048)]; + assert_eq!( + bytes_bad_in_title(&title, &gap), + 0, + "bad bytes in the inter-extent gap must not be counted" + ); + // A bad range overlapping the first extent counts. + let in_first = vec![(0, 4096)]; + assert_eq!(bytes_bad_in_title(&title, &in_first), 4096); + // A bad range spanning both extents plus the gap counts only the + // bytes that fall inside the two extents (10 + 10 sectors). + let spanning = vec![(0, 110 * 2048)]; + assert_eq!(bytes_bad_in_title(&title, &spanning), 20 * 2048); + } + + /// 0xA2 is secondary DTS-HD MA (lossless), not lossy HR. + #[test] + fn coding_type_a2_is_dts_hd_ma() { + assert_eq!(Codec::from_coding_type(0xA2), Codec::DtsHdMa); + assert_eq!(Codec::from_coding_type(0x86), Codec::DtsHdMa); + } + + /// chapter_name emits a bare 1-based ordinal (no localized prose). + #[test] + fn chapter_name_is_bare_ordinal() { + assert_eq!(chapter_name(0), "1"); + assert_eq!(chapter_name(41), "42"); + } } diff --git a/src/disc/patch.rs b/src/disc/patch.rs index c60b8fe..f3cf563 100644 --- a/src/disc/patch.rs +++ b/src/disc/patch.rs @@ -319,14 +319,14 @@ const CACHE_PRIME_SECTORS: u32 = 3; /// scatter its sample LBAs across the failing region rather than /// hammering the same neighborhood. pub(super) fn skip_sectors_for_probe(idx: usize) -> u64 { - let base = PASSN_SKIP_SECTORS_BASE as i64; - let escalation = (idx * 3) as i64; - let shifted = if escalation < 64 { - base << escalation - } else { - base - }; - shifted.min(PASSN_SKIP_SECTORS_CAP as i64) as u64 + let escalation = (idx.saturating_mul(3)).min(u32::MAX as usize) as u32; + // Saturating shift: a large `idx` would overflow a fixed-width shift + // (32 << 60 = 2^65), so fall back to the cap instead of panicking + // (debug) or wrapping to 0 (release). + PASSN_SKIP_SECTORS_BASE + .checked_shl(escalation) + .unwrap_or(PASSN_SKIP_SECTORS_CAP) + .min(PASSN_SKIP_SECTORS_CAP) } /// Send a `PatchItem` and translate a `SendError` (consumer thread died @@ -336,9 +336,7 @@ pub(super) fn send_or_abort( pipe: &Pipeline<PatchItem, PatchSummary>, item: PatchItem, ) -> Result<()> { - pipe.send(item).map_err(|_| Error::IoError { - source: std::io::Error::other("patch consumer terminated unexpectedly"), - }) + pipe.send(item).map_err(|_| Error::PipelineConsumerGone) } /// Phase A pre-snapshot. Loads the mapfile, captures the fields the @@ -667,21 +665,21 @@ pub(super) fn handle_read_success<R: SectorSource + ?Sized>( state.damage_window.push(true); if state.damage_window.len() > PASSN_DAMAGE_WINDOW { state.damage_window.remove(0); - - tracing::info!( - target: "freemkv::disc", - phase = "patch_read_ok", - lba, - count, - bytes, - blocks_read_ok = state.blocks_read_ok, - consecutive_failures = state.consecutive_failures, - read_duration_ms, - range_idx = frame.range_idx, - pos, - "Read succeeded" - ); } + + tracing::info!( + target: "freemkv::disc", + phase = "patch_read_ok", + lba, + count, + bytes, + blocks_read_ok = state.blocks_read_ok, + consecutive_failures = state.consecutive_failures, + read_duration_ms, + range_idx = frame.range_idx, + pos, + "Read succeeded" + ); // Plaintext: DecryptingSectorSource applied AACS / CSS in-place // during the read_sectors call above. The pre-0.18 inline // decrypt_sectors call lived here. @@ -909,8 +907,12 @@ pub(super) fn handle_read_failure<R: SectorSource + ?Sized>( // Check if this is a NOT_READY error that should be retried let sense = err.scsi_sense(); - // ASC values indicating temporary drive unresponsiveness: - // 0x02 = medium not present, 0x03 = becoming ready, 0x04 = initialization required + // ASC values (under NOT READY, sense_key 0x02) indicating temporary + // drive unresponsiveness worth retrying: + // 0x02 = LUN not ready, no reference position (mechanism still seeking) + // 0x03 = LUN not ready, manual intervention required + // 0x04 = LUN not ready, in process of becoming ready / initializing + // (Medium-not-present is ASC 0x3A, not handled here — nothing to retry.) let is_not_ready_retryable = sense .map(|s| s.sense_key == 0x02 && (s.asc == 0x02 || s.asc == 0x03 || s.asc == 0x04)) .unwrap_or(false); @@ -923,7 +925,7 @@ pub(super) fn handle_read_failure<R: SectorSource + ?Sized>( lba, consecutive_failures = state.consecutive_failures, err_asc = sense.map(|s| s.asc as u32).unwrap_or(0), - "NOT_READY with ASC=0x03/0x04; pausing for drive recovery before retry" + "NOT_READY with ASC in 0x02/0x03/0x04; pausing for drive recovery before retry" ); // Extended pause for NOT_READY - let drive complete internal mechanical recovery @@ -1020,17 +1022,30 @@ pub(super) fn handle_read_failure<R: SectorSource + ?Sized>( ); } - // Probe good sectors to differentiate wedge vs bad sector + // Probe good sectors to differentiate wedge vs bad sector. + // `skip_sectors_for_probe` returns a SECTOR distance; scale to bytes + // before adding to `pos` (a byte offset). The previous code compared + // a sector count against `block_bytes` and added a sector count to a + // byte offset, so the only probe that ran landed back on the failing + // LBA — the responsive-vs-wedged heuristic never scattered. if state.consecutive_failures >= 3 && state.consecutive_failures % 5 == 0 { - let probe_offsets: [u64; 3] = [0, skip_sectors_for_probe(1), skip_sectors_for_probe(2)]; + let probe_offsets_sectors: [u64; 3] = + [0, skip_sectors_for_probe(1), skip_sectors_for_probe(2)]; let mut probes_ok = 0; - for (probe_idx, &offset) in probe_offsets.iter().enumerate() { - if offset >= block_bytes || (offset == 0 && state.consecutive_failures < 5) { + for (probe_idx, &offset_sectors) in probe_offsets_sectors.iter().enumerate() { + let offset = offset_sectors.saturating_mul(2048); + let probe_pos = pos.saturating_add(offset); + // Skip the zero-distance re-read until failures are well + // established (it just re-confirms the current LBA), and + // never probe past the end of the current bad range (the + // probe scatters sample LBAs across the failing region — + // `block_bytes`, one block, was the wrong bound and in the + // wrong units). + if probe_pos >= frame.end || (offset == 0 && state.consecutive_failures < 5) { continue; } - let probe_pos = pos + offset; let probe_lba = (probe_pos / 2048) as u32; let probe_count = 1u16; let mut probe_buf = [0u8; 2048]; @@ -1229,20 +1244,11 @@ pub(super) fn check_range_watchdog( frame: &RangeFrame, shared: &Mutex<SharedPatchState>, ) -> bool { - if state.range_start.elapsed().as_secs() > frame.range_budget_secs { - tracing::warn!( - target: "freemkv::disc", - phase = "patch_range_timeout", - range_lba = frame.range_pos / 2048, - range_sectors = frame.range_sectors, - elapsed_secs = state.range_start.elapsed().as_secs(), - budget_secs = frame.range_budget_secs, - bytes_recovered = state.range_bytes_good.saturating_sub(state.bytes_good_before), - "Range timeout - moving to next range" - ); - return true; - } - + // Refresh the forward-progress baseline FIRST, then do a single + // elapsed-vs-budget check. Reading bytes_good before the budget + // test means a range that committed a recovered sector since the + // previous tick resets its clock instead of being abandoned in the + // budget-boundary window. let bytes_good_now = { let g = shared .lock() @@ -1292,30 +1298,42 @@ pub(super) fn handle_skip_limit( // them on a later pass when state has evolved (cache, mechanical // settle). 2026-05-07 dd-as-oracle test confirmed ~36% of patch- // marked Unreadable sectors are actually readable. - let unmarked_bytes = frame.block_end.saturating_sub(frame.range_pos); - if opts.reverse { - send_or_abort( - pipe, - PatchItem::NonTrimmed { - pos: frame.range_pos, - len: unmarked_bytes, - }, - )?; - } else { - let remaining_start = frame.range_pos + (frame.end - frame.block_end); - if remaining_start < frame.end { - send_or_abort( - pipe, - PatchItem::NonTrimmed { - pos: remaining_start, - len: frame.end - remaining_start, - }, - )?; - } + if let Some((pos, len)) = + skip_limit_remainder(opts.reverse, frame.range_pos, frame.end, frame.block_end) + { + send_or_abort(pipe, PatchItem::NonTrimmed { pos, len })?; } Ok(()) } +/// The never-attempted remainder of a range when the skip limit is +/// reached, as `Some((pos, len))` or `None` if nothing is left. +/// +/// `block_end` is the per-iteration cursor. In reverse mode it moved +/// DOWN from `end` toward `range_pos`, so the attempted region is +/// `[block_end, end)` and the remainder is `[range_pos, block_end)`. In +/// forward mode it moved UP from `range_pos` toward `end`, so the +/// attempted region is `[range_pos, block_end)` and the remainder is +/// `[block_end, end)`. The pre-fix forward formula +/// `range_pos + (end - block_end)` was a mirror reflection that, once +/// `block_end` passed the midpoint, produced a start BELOW `block_end` +/// and overlapped the already-recovered region — downgrading Finished +/// sectors to NonTrimmed. +fn skip_limit_remainder( + reverse: bool, + range_pos: u64, + end: u64, + block_end: u64, +) -> Option<(u64, u64)> { + if reverse { + let len = block_end.saturating_sub(range_pos); + (len > 0).then_some((range_pos, len)) + } else { + let len = end.saturating_sub(block_end); + (len > 0).then_some((block_end, len)) + } +} + /// Damage-cluster size-aware skip decision. Inspects `state.damage_window` /// against the `PASSN_DAMAGE_THRESHOLD_PCT` threshold; if crossed, /// advances `frame.block_end` by an escalating skip (capped at 1/4 of @@ -1348,7 +1366,9 @@ pub(super) fn compute_damage_skip( }; let range_remaining_sectors = range_remaining_bytes / 2048; let range_quarter = (range_remaining_sectors / 4).max(1); - let escalated = (PASSN_SKIP_SECTORS_BASE << state.consecutive_skips_without_recovery) + let escalated = PASSN_SKIP_SECTORS_BASE + .checked_shl(state.consecutive_skips_without_recovery) + .unwrap_or(PASSN_SKIP_SECTORS_CAP) .min(PASSN_SKIP_SECTORS_CAP); let skip_sectors = escalated.min(range_quarter); let skip_bytes = skip_sectors * 2048; @@ -1539,7 +1559,11 @@ impl Disc { // "split decisions", not recorded failures // - drop-to-1 retries the SAME starting position, so every // sector in the failed batch is individually probed - let initial_batch = opts.block_sectors.unwrap_or(1); + // Clamp to at least 1 sector. block_sectors is public + // (Option<u16>); Some(0) would compute a zero-length read per + // iteration, never advance block_end, and busy-spin the range + // until its watchdog fired. + let initial_batch = opts.block_sectors.unwrap_or(1).max(1); let recovery = opts.full_recovery; let mut state = PatchLoopState::new( bytes_good_before, @@ -1785,3 +1809,51 @@ impl Disc { )) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn skip_sectors_for_probe_does_not_overflow_for_large_idx() { + // idx=20 (escalation 60) and idx=21 (63) previously overflowed + // i64 via `32i64 << escalation`. Must saturate to the cap. + for idx in [0usize, 1, 2, 20, 21, 100, usize::MAX] { + let v = skip_sectors_for_probe(idx); + assert!( + v <= PASSN_SKIP_SECTORS_CAP, + "idx {idx}: {v} exceeds cap {PASSN_SKIP_SECTORS_CAP}" + ); + } + // Small indices still escalate as before. + assert_eq!(skip_sectors_for_probe(0), PASSN_SKIP_SECTORS_BASE); + assert_eq!(skip_sectors_for_probe(1), PASSN_SKIP_SECTORS_BASE << 3); + } + + #[test] + fn skip_limit_remainder_forward_does_not_overlap_recovered_region() { + // Forward mode: range [1000, 2000), cursor advanced past the + // midpoint to block_end=1700. The recovered region is + // [1000, 1700); the never-attempted remainder must be exactly + // [1700, 2000) — NOT a mirror start below block_end. + let r = skip_limit_remainder(false, 1000, 2000, 1700); + assert_eq!(r, Some((1700, 300))); + // The pre-fix mirror formula would have produced start = + // 1000 + (2000 - 1700) = 1300, which overlaps [1000, 1700). + assert!(r.unwrap().0 >= 1700, "must not overlap recovered region"); + } + + #[test] + fn skip_limit_remainder_forward_none_when_fully_attempted() { + assert_eq!(skip_limit_remainder(false, 1000, 2000, 2000), None); + } + + #[test] + fn skip_limit_remainder_reverse_marks_low_unattempted_region() { + // Reverse mode: cursor moved down to block_end=1300, so + // [1300, 2000) was attempted and [1000, 1300) is the remainder. + let r = skip_limit_remainder(true, 1000, 2000, 1300); + assert_eq!(r, Some((1000, 300))); + assert_eq!(skip_limit_remainder(true, 1000, 2000, 1000), None); + } +} diff --git a/src/disc/read_error.rs b/src/disc/read_error.rs index a708add..0f6ee55 100644 --- a/src/disc/read_error.rs +++ b/src/disc/read_error.rs @@ -34,14 +34,20 @@ pub struct ReadCtx { /// Sliding window of recent read outcomes (true=ok, false=fail). /// Capped at `damage_window_max`. Drives damage-jump decisions. pub damage_window: Vec<bool>, + /// Maximum number of outcome entries kept in `damage_window`; the + /// oldest is evicted once this is exceeded. A whole count (e.g. 16). pub damage_window_max: usize, + /// Fraction of `damage_window` entries that must be failures before + /// the window-based damage-jump fires, as a whole-number percentage + /// (e.g. `12` = 12%). pub damage_threshold_pct: usize, /// Trigger a damage-jump after this many consecutive outer-batch /// failures, even when the damage_window isn't full yet. Pass 1 - /// uses a small value (4) so we don't spend ~40 minutes grinding - /// to fill a 16-block window before the first jump on a damage - /// zone we entered cleanly. Pass N uses a larger value (or - /// disables this — see `bisect_on_marginal`) because Pass N's + /// uses a small value (1 — jump on the first outer failure; see + /// the 2026-05-11 rewrite in `for_sweep`) so we don't spend ~40 + /// minutes grinding to fill a 16-block window before the first jump + /// on a damage zone we entered cleanly. Pass N uses a larger value + /// (or disables this — see `bisect_on_marginal`) because Pass N's /// whole job IS to grind on the bad ranges. pub fast_jump_threshold: u64, /// Multiplier applied to damage-jump distance. Doubles each jump, @@ -240,6 +246,13 @@ impl ReadCtx { // drive recovered, so further wedges should reset the skip // budget instead of accumulating toward a real abort. self.wedge_count = 0; + // A successful read also means the bridge recovered, so the + // 15s-cooldown retry budget should be available again for the + // next bridge-degradation event. Without this reset the budget + // saturates permanently after 5 cumulative events across the + // whole pass and later degradations skip the cooldown retry, + // needlessly losing data. + self.bridge_degradation_count = 0; // Outer-success only: a good single-sector read inside a // bisect doesn't mean we've left the damaged batch. Only an // outer-batch success resets the outer-failure counter. @@ -259,6 +272,13 @@ impl ReadCtx { if self.in_damage_zone && self.consecutive_good >= self.damage_window_max as u64 { self.in_damage_zone = false; self.last_error_family = None; + // Reset the damage-jump multiplier so the NEXT zone starts + // from the base jump distance. Without this the multiplier + // stays at whatever the prior zone inflated it to (up to + // MAX_JUMP_MULTIPLIER=64), so the next zone's first jump is + // 64x oversized and skips recoverable data. The field doc + // promises this reset. + self.jump_multiplier = 1; } } @@ -375,9 +395,9 @@ const JUMP_BASE_SECTORS: u64 = 1024; // When the BU40N (or similar drives) hits a physical-damage cluster, // its firmware can transition into a "wedge" state where it returns // HARDWARE_ERROR or ILLEGAL_REQUEST for every subsequent read — -// often for many LBAs after the actual bad sector. Per project docs -// "Bad-sector handling" rule #2: "Recovery requires eject+reload OR -// significant cool-down." +// often for many LBAs after the actual bad sector. Once wedged, +// recovery requires either a physical eject + reload or a significant +// cool-down period; hammering the same LBA only deepens the state. // // Pass 1's pre-fix behavior was to immediately AbortPass on the // first HARDWARE_ERROR / ILLEGAL_REQUEST, killing the rip at @@ -396,10 +416,10 @@ const JUMP_BASE_SECTORS: u64 = 1024; /// One-gigabyte jump (1024 MiB) on each wedge. Big enough to clear /// almost any single-cluster damage zone we've seen. const WEDGE_JUMP_SECTORS: u64 = 524_288; -/// Cooldown pause after each wedge. Per project docs the drive needs -/// "significant cool-down"; 30 s strikes a balance between giving -/// the drive a chance to recover and not stalling the rip if the -/// drive is permanently stuck. +/// Cooldown pause after each wedge. A wedged drive needs a +/// significant cool-down to leave fast-fail; 30 s strikes a balance +/// between giving the drive a chance to recover and not stalling the +/// rip if the drive is permanently stuck. const WEDGE_PAUSE_SECS: u64 = 30; /// Bail after this many consecutive wedges with no good read in /// between. At 1 GB jumps this lets us scan ~16 GB worth of fully @@ -463,8 +483,13 @@ pub fn handle_read_error(err: &Error, ctx: &mut ReadCtx) -> ReadAction { .unwrap_or(SenseFamily::Other); // Zone-entry tracking: this is the first error after a clean run - // (or the first error of the sweep). - if !ctx.in_damage_zone && !ctx.bisecting { + // (or the first error of the sweep). Capture the genuine + // clean->damaged transition here, BEFORE mutating in_damage_zone, + // so the 30s zone-entry cooldown below keys off the real + // transition rather than re-deriving it from a counter that the + // fast-jump path resets after every jump. + let is_zone_entry_transition = !ctx.in_damage_zone && !ctx.bisecting; + if is_zone_entry_transition { ctx.in_damage_zone = true; ctx.zones_entered += 1; } @@ -569,9 +594,13 @@ pub fn handle_read_error(err: &Error, ctx: &mut ReadCtx) -> ReadAction { // AbortPass after N consecutive wedges with no successful // read in between. if sense_key == scsi::SENSE_KEY_HARDWARE_ERROR || sense_key == scsi::SENSE_KEY_ILLEGAL_REQUEST { - if !ctx.bisecting { - ctx.wedge_count += 1; - } + // Count every wedge, including bisect-inner ones. A wedge is a + // firmware fast-fail state regardless of whether we're inside a + // bisect; if we did NOT count bisect-inner wedges, a drive that + // wedges mid-bisect would burn a 30s WEDGE_PAUSE cooldown per + // inner sector and never reach WEDGE_ABORT_THRESHOLD from inside + // the bisect — ~16 min of cooldown sleeping on a batch=32 bisect. + ctx.wedge_count += 1; if ctx.wedge_count >= WEDGE_ABORT_THRESHOLD { tracing::warn!( target: "freemkv::disc", @@ -665,8 +694,7 @@ pub fn handle_read_error(err: &Error, ctx: &mut ReadCtx) -> ReadAction { // branch for future tuning. Pass N (bisect_on_marginal=true) // uses the standard pauses — it's running single-sector retries // on already-known-bad LBAs by design. - let is_zone_entry = - ctx.consecutive_outer_failures == 1 && !ctx.bisecting && !ctx.bisect_on_marginal; + let is_zone_entry = is_zone_entry_transition && !ctx.bisecting && !ctx.bisect_on_marginal; let pause_secs = if is_zone_entry { ZONE_ENTRY_COOLDOWN_SECS } else if ctx.consecutive_failures >= CONSECUTIVE_FAIL_LONG_PAUSE_THRESHOLD { @@ -692,7 +720,7 @@ pub fn handle_read_error(err: &Error, ctx: &mut ReadCtx) -> ReadAction { // Two triggers, evaluated in order: // // a. **Fast-entry** — `consecutive_outer_failures >= fast_jump_threshold`. - // Fires on Pass 1 (threshold=4) so we don't spend ~40 min + // Fires on Pass 1 (threshold=1) so we don't spend ~40 min // grinding to fill a 16-block damage window before the // first jump on a damage zone we entered cleanly. Doesn't // fire on Pass N (threshold=u64::MAX). @@ -1029,6 +1057,43 @@ mod tests { } } + #[test] + fn pass_1_subsequent_in_zone_errors_skip_long_cooldown() { + // Regression: the fast-jump path resets consecutive_outer_failures + // to 0 after each jump, so the next in-zone error re-increments it + // to 1. Zone-entry must key off the genuine clean->damaged + // transition (in_damage_zone), not the counter, otherwise every + // error in a damaged region pays the 30 s cooldown. + let mut ctx = ReadCtx::for_sweep(32); + // First error: genuine zone entry, gets the long cooldown. + let first = handle_read_error(&medium_err(), &mut ctx); + match first { + ReadAction::JumpAhead { pause_secs, .. } => assert_eq!( + pause_secs, + ZONE_ENTRY_COOLDOWN_SECS + POST_JUMP_EXTRA_PAUSE_SECS + ), + other => panic!("expected JumpAhead on first error, got {other:?}"), + } + // We are now still in the damage zone; the jump reset the outer + // counter. A second error must NOT re-arm the 30 s cooldown. + assert!(ctx.in_damage_zone); + let second = handle_read_error(&medium_err(), &mut ctx); + let pause = match second { + ReadAction::JumpAhead { pause_secs, .. } => pause_secs, + ReadAction::SkipBlock { pause_secs } => pause_secs, + other => panic!("expected pausing action, got {other:?}"), + }; + assert_ne!( + pause, + ZONE_ENTRY_COOLDOWN_SECS + POST_JUMP_EXTRA_PAUSE_SECS, + "subsequent in-zone error must not pay the 30 s zone-entry cooldown" + ); + assert!( + pause <= FAIL_PAUSE_SECS + POST_JUMP_EXTRA_PAUSE_SECS, + "subsequent in-zone pause should be the standard fail pause, got {pause}" + ); + } + #[test] fn pass_n_pauses_uniformly_on_failed_read() { // Pass N (bisect_on_marginal=true) is exempt from the @@ -1067,6 +1132,66 @@ mod tests { ); } + #[test] + fn jump_multiplier_resets_after_damage_zone_exit() { + // A zone that doubles the multiplier must not carry the inflated + // value into the next zone — otherwise the next zone's first + // jump is up to 64x oversized and skips recoverable data. + let mut ctx = ReadCtx::for_sweep(32); + // First zone: a few errors push jumps and double the multiplier. + for _ in 0..4 { + handle_read_error(&medium_err(), &mut ctx); + } + assert!( + ctx.jump_multiplier > 1, + "expected the multiplier to inflate inside a damage zone" + ); + // Exit the zone: damage_window_max consecutive good reads. + ctx.bisecting = false; + for _ in 0..ctx.damage_window_max { + ctx.on_success(); + } + assert!(!ctx.in_damage_zone, "zone should have exited"); + assert_eq!( + ctx.jump_multiplier, 1, + "jump_multiplier must reset to 1 on zone exit" + ); + } + + #[test] + fn bridge_degradation_count_resets_on_success() { + // After a good read the bridge recovered; the 15s-cooldown retry + // budget must be available again instead of staying saturated + // for the whole pass. + let mut ctx = ReadCtx::for_patch(1); + ctx.bridge_degradation_count = BRIDGE_DEGRADATION_MAX_RETRIES; + ctx.on_success(); + assert_eq!(ctx.bridge_degradation_count, 0); + } + + #[test] + fn wedge_abort_reachable_during_bisect() { + // A drive that wedges mid-bisect must still reach the abort + // threshold rather than burning a WEDGE_PAUSE cooldown per inner + // sector forever. + let mut ctx = ReadCtx::for_patch(32); + ctx.bisecting = true; + let mut aborted = false; + for _ in 0..WEDGE_ABORT_THRESHOLD { + if matches!( + handle_read_error(&hardware_err(), &mut ctx), + ReadAction::AbortPass + ) { + aborted = true; + break; + } + } + assert!( + aborted, + "wedge abort threshold must be reachable from inside a bisect" + ); + } + #[test] fn on_success_resets_failure_counters_and_pushes_window() { let mut ctx = ReadCtx::for_sweep(32); diff --git a/src/disc/sweep.rs b/src/disc/sweep.rs index d4a8c91..2c86838 100644 --- a/src/disc/sweep.rs +++ b/src/disc/sweep.rs @@ -8,15 +8,13 @@ //! during the post-read work; throughput tops out at the *sum* of //! both costs. //! -//! 0.17.11 introduced a bespoke producer/consumer split (the now- -//! removed `disc/sweep_pipeline.rs`) to overlap the two stages. 0.18 -//! collapses that split — together with the analogous splits patch -//! and mux need — onto the generic [`crate::io::Pipeline`] + -//! [`crate::io::Sink`] primitive. This module is the sweep-specific -//! `Sink` impl; the producer-side state machine (read_error context, -//! decrypt, set_speed, halt) stays in `Disc::sweep` in `disc/mod.rs`. +//! A producer/consumer split overlaps the two stages on the generic +//! [`crate::io::Pipeline`] + [`crate::io::Sink`] primitive. This module +//! is the sweep-specific `Sink` impl; the producer-side state machine +//! (read_error context, decrypt, set_speed, halt) stays in +//! `Disc::sweep` in `disc/mod.rs`. //! -//! Correctness invariants preserved (same as 0.17.11): +//! Correctness invariants preserved: //! - Mapfile is single-writer (consumer-only). No locking. //! - All `read_error::ReadCtx` state stays on the producer thread. //! - `set_speed` calls happen on the producer thread (same thread that @@ -25,9 +23,8 @@ //! intact in the consumer (write before record), so the on-disk //! invariant "mapfile only marks Finished what the file has //! received" survives a crash mid-pass. -//! - The BU40N+Initio bridge wedge concern is unchanged: only one -//! SCSI command in flight at a time, error-path timing identical, -//! no new retry logic. +//! - Only one SCSI command is in flight at a time; error-path timing +//! is identical and no new retry logic is introduced. use std::io::{Seek, SeekFrom, Write}; use std::sync::mpsc::{Receiver, SyncSender, sync_channel}; @@ -40,7 +37,7 @@ use super::mapfile::{MapStats, Mapfile, SectorStatus}; /// Reusable zero buffer for SkipFill / GapFill / BisectBad. 64 KB /// matches the existing zero_gap chunk size used by the pre-split /// sweep loop. -const ZERO_CHUNK: usize = 65 * 1024; +const ZERO_CHUNK: usize = 64 * 1024; /// Producer → Consumer messages. The consumer applies these in FIFO /// order; ordering of file writes and mapfile records across items is @@ -151,55 +148,31 @@ impl Sink<WorkItem> for SweepSink { WorkItem::Good { pos, buf } => { // Decrypt is on the producer; consumer assumes plaintext. let len = buf.len() as u64; - self.file - .seek(SeekFrom::Start(pos)) - .map_err(|e| Error::IoError { source: e })?; - self.file - .write_all(&buf) - .map_err(|e| Error::IoError { source: e })?; - self.map - .record(pos, len, SectorStatus::Finished) - .map_err(|e| Error::IoError { source: e })?; + self.file.seek(SeekFrom::Start(pos))?; + self.file.write_all(&buf)?; + self.map.record(pos, len, SectorStatus::Finished)?; } WorkItem::BisectGood { pos, buf } => { - self.file - .seek(SeekFrom::Start(pos)) - .map_err(|e| Error::IoError { source: e })?; - self.file - .write_all(&buf[..]) - .map_err(|e| Error::IoError { source: e })?; - self.map - .record(pos, 2048, SectorStatus::Finished) - .map_err(|e| Error::IoError { source: e })?; + self.file.seek(SeekFrom::Start(pos))?; + self.file.write_all(&buf[..])?; + self.map.record(pos, 2048, SectorStatus::Finished)?; } WorkItem::BisectBad { pos } => { - self.file - .seek(SeekFrom::Start(pos)) - .map_err(|e| Error::IoError { source: e })?; - self.file - .write_all(&self.zero[..2048]) - .map_err(|e| Error::IoError { source: e })?; - self.map - .record(pos, 2048, SectorStatus::NonTrimmed) - .map_err(|e| Error::IoError { source: e })?; + self.file.seek(SeekFrom::Start(pos))?; + self.file.write_all(&self.zero[..2048])?; + self.map.record(pos, 2048, SectorStatus::NonTrimmed)?; } WorkItem::SkipFill { pos, len } | WorkItem::GapFill { pos, len } => { - self.file - .seek(SeekFrom::Start(pos)) - .map_err(|e| Error::IoError { source: e })?; + self.file.seek(SeekFrom::Start(pos))?; // Subsequent writes are sequential; `WritebackFile`'s // seek-elision keeps them on the writeback pipeline path. let mut filled = 0u64; while filled < len { let chunk = (len - filled).min(self.zero.len() as u64) as usize; - self.file - .write_all(&self.zero[..chunk]) - .map_err(|e| Error::IoError { source: e })?; + self.file.write_all(&self.zero[..chunk])?; filled += chunk as u64; } - self.map - .record(pos, len, SectorStatus::NonTrimmed) - .map_err(|e| Error::IoError { source: e })?; + self.map.record(pos, len, SectorStatus::NonTrimmed)?; } WorkItem::StatsRequest => { let stats = self.map.stats(); @@ -230,7 +203,7 @@ impl Sink<WorkItem> for SweepSink { // Non-regular outputs (/dev/null, pipes) always fail // sync_all; that's not a real error. } - self.map.flush().map_err(|e| Error::IoError { source: e })?; + self.map.flush()?; Ok(ConsumerSummary { stats: self.map.stats(), diff --git a/src/drive/capture.rs b/src/drive/capture.rs index 6c35185..15db7f0 100644 --- a/src/drive/capture.rs +++ b/src/drive/capture.rs @@ -25,8 +25,14 @@ pub struct DriveCapture { /// A single GET CONFIGURATION feature response from the drive. #[derive(Debug, Clone)] pub struct CapturedFeature { + /// MMC-6 GET CONFIGURATION feature code (e.g. `0x010D` = AACS). pub code: u16, + /// Static human-readable label from the internal `FEATURES` table — + /// not a device-reported string. pub name: &'static str, + /// Raw feature-descriptor payload bytes, with the 8-byte GET + /// CONFIGURATION header stripped (i.e. `buf[8..]`). Unlike + /// [`DriveCapture::gc_010c`], which retains the full header. pub data: Vec<u8>, } diff --git a/src/drive/linux.rs b/src/drive/linux.rs index 7e1fe11..03fdcbb 100644 --- a/src/drive/linux.rs +++ b/src/drive/linux.rs @@ -1,18 +1,33 @@ //! Linux drive discovery and device resolution. +use crate::drive::DeviceResolution; use crate::error::{Error, Result}; use crate::identity::DriveId; +/// SCSI peripheral device type 5 = MMC / optical (CD/DVD/BD), held in the +/// low 5 bits of INQUIRY byte 0 (the high 3 bits are the peripheral +/// qualifier, masked off here). +const SCSI_PERIPHERAL_TYPE_OPTICAL: u8 = 0x05; + +/// Discover optical drives by enumerating `/dev/sg*` SCSI-generic nodes, +/// opening each, running INQUIRY, and keeping only devices whose +/// peripheral device type is optical (MMC, type 0x05). +/// +/// Devices where `scsi::open` or `DriveId::from_drive` fail are silently +/// skipped — that is intentional for enumeration (a busy or wedged node +/// shouldn't abort discovery of the others). pub fn find_drives() -> Vec<(String, DriveId)> { let mut drives = Vec::new(); - for i in 0..16 { - let path = format!("/dev/sg{i}"); + for name in enumerate_sg_names() { + let path = format!("/dev/{name}"); if !std::path::Path::new(&path).exists() { continue; } if let Ok(mut transport) = crate::scsi::open(std::path::Path::new(&path)) { if let Ok(id) = DriveId::from_drive(transport.as_mut()) { - if !id.raw_inquiry.is_empty() && (id.raw_inquiry[0] & 0x1F) == 0x05 { + if !id.raw_inquiry.is_empty() + && (id.raw_inquiry[0] & 0x1F) == SCSI_PERIPHERAL_TYPE_OPTICAL + { drives.push((path, id)); } } @@ -21,41 +36,78 @@ pub fn find_drives() -> Vec<(String, DriveId)> { drives } +/// Enumerate `sg*` device names. Linux assigns `/dev/sgN` sequentially +/// across *all* SCSI-generic devices (disks, tape, HBAs, optical), so a +/// fixed `sg0..15` range can miss an optical drive on a host with many +/// targets. Prefer the exact present-device list from +/// `/sys/class/scsi_generic/`; fall back to a bounded `sg0..15` probe +/// only when sysfs is unreadable (minimal containers). +fn enumerate_sg_names() -> Vec<String> { + let mut names = Vec::new(); + if let Ok(entries) = std::fs::read_dir("/sys/class/scsi_generic") { + for entry in entries.flatten() { + let name = entry.file_name().to_string_lossy().to_string(); + if name.starts_with("sg") { + names.push(name); + } + } + } else { + for i in 0..16 { + let name = format!("sg{i}"); + if std::path::Path::new(&format!("/dev/{name}")).exists() { + names.push(name); + } + } + } + names.sort(); + names +} + +/// Resolve a device path to its raw `/dev/sg*` SCSI-generic node. +/// +/// - `/dev/sg*` paths pass through unchanged ([`DeviceResolution::Direct`]). +/// - `/dev/sr*` block paths are matched (by vendor/product/serial) to the +/// corresponding `/dev/sg*` node ([`DeviceResolution::SrToSg`]); if no +/// match is found the original path is returned with +/// [`DeviceResolution::SrNoSgMatch`]. +/// - Any other existing path passes through as [`DeviceResolution::Direct`]. #[allow(dead_code)] -pub fn resolve_device(path: &str) -> Result<(String, Option<String>)> { +pub fn resolve_device(path: &str) -> Result<(String, DeviceResolution)> { if path.contains("/sg") { if !std::path::Path::new(path).exists() { return Err(Error::DeviceNotFound { path: path.to_string(), }); } - return Ok((path.to_string(), None)); + return Ok((path.to_string(), DeviceResolution::Direct)); } if path.contains("/sr") { let mut sr_transport = crate::scsi::open(std::path::Path::new(path))?; let sr_id = DriveId::from_drive(sr_transport.as_mut())?; drop(sr_transport); for (sg_path, sg_id) in find_drives() { - if sg_id.vendor_id == sr_id.vendor_id + // Require a non-empty serial before treating vendor/product/ + // serial as a unique match. serial_number falls back to an + // empty string when GET CONFIGURATION 0108h is unavailable + // (common on OEM drives); two same-model drives would then + // both compare equal and the first in enumeration order would + // win silently, resolving sr1 to sr0's sg node. An empty + // serial can't disambiguate, so fall through to the no-match + // path instead. + if !sr_id.serial_number.is_empty() + && sg_id.vendor_id == sr_id.vendor_id && sg_id.product_id == sr_id.product_id && sg_id.serial_number == sr_id.serial_number { - let warning = - format!("{path} is a block device (sr) — using {sg_path} (sg) for raw access"); - return Ok((sg_path, Some(warning))); + return Ok((sg_path, DeviceResolution::SrToSg)); } } - return Ok(( - path.to_string(), - Some(format!( - "{path} is a block device (sr) — no matching sg device found" - )), - )); + return Ok((path.to_string(), DeviceResolution::SrNoSgMatch)); } if !std::path::Path::new(path).exists() { return Err(Error::DeviceNotFound { path: path.to_string(), }); } - Ok((path.to_string(), None)) + Ok((path.to_string(), DeviceResolution::Direct)) } diff --git a/src/drive/macos.rs b/src/drive/macos.rs index 409a699..75eda85 100644 --- a/src/drive/macos.rs +++ b/src/drive/macos.rs @@ -4,9 +4,20 @@ //! to discover optical drives without exclusive access or unmounts. Only //! the returned paths are then opened for INQUIRY to build full `DriveId`. +use crate::drive::DeviceResolution; use crate::error::{Error, Result}; use crate::identity::DriveId; +/// SCSI peripheral device type 5 = MMC / optical, in the low 5 bits of +/// INQUIRY byte 0. +const SCSI_PERIPHERAL_TYPE_OPTICAL: u8 = 0x05; + +/// Discover optical drives via the IOKit registry (`scsi::list_drives`), +/// then open each candidate for INQUIRY to build a full `DriveId`. +/// +/// Any drive where `scsi::open` or `DriveId::from_drive` fails, or whose +/// peripheral device type is not optical (MMC, type 0x05), is silently +/// skipped — the same MMC filter the Linux and Windows backends apply. pub fn find_drives() -> Vec<(String, DriveId)> { let mut drives = Vec::new(); let discovered = crate::scsi::list_drives(); @@ -15,7 +26,11 @@ pub fn find_drives() -> Vec<(String, DriveId)> { match crate::scsi::open(path) { Ok(mut transport) => { if let Ok(id) = DriveId::from_drive(transport.as_mut()) { - drives.push((info.path.clone(), id)); + if !id.raw_inquiry.is_empty() + && (id.raw_inquiry[0] & 0x1F) == SCSI_PERIPHERAL_TYPE_OPTICAL + { + drives.push((info.path.clone(), id)); + } } } Err(_) => { @@ -26,20 +41,15 @@ pub fn find_drives() -> Vec<(String, DriveId)> { drives } -pub fn resolve_device(path: &str) -> Result<(String, Option<String>)> { - // Accept /dev/diskN or /dev/rdiskN paths as-is - if path.contains("/disk") || path.contains("/rdisk") { - if !std::path::Path::new(path).exists() { - return Err(Error::DeviceNotFound { - path: path.to_string(), - }); - } - return Ok((path.to_string(), None)); - } +/// Resolve a device path on macOS. There is no `sr`→`sg` style +/// substitution here (that is a Linux concern), so any existing path is +/// returned unchanged as [`DeviceResolution::Direct`]; the +/// [`DeviceResolution`] return exists for cross-platform signature parity. +pub fn resolve_device(path: &str) -> Result<(String, DeviceResolution)> { if !std::path::Path::new(path).exists() { return Err(Error::DeviceNotFound { path: path.to_string(), }); } - Ok((path.to_string(), None)) + Ok((path.to_string(), DeviceResolution::Direct)) } diff --git a/src/drive/mod.rs b/src/drive/mod.rs index 44ecf15..fbf24c8 100644 --- a/src/drive/mod.rs +++ b/src/drive/mod.rs @@ -1,6 +1,8 @@ //! Drive session — open, identify, and read from optical drives. //! -//! 4. `probe_disc()` — probe disc surface. Drive learns optimal speeds. +//! A `Drive` is opened from a device path, identifies itself via INQUIRY, +//! optionally unlocks/initializes via a platform driver, and reads sectors. +//! `probe_disc()` primes the firmware's per-region speed table. pub(crate) fn extract_scsi_context(e: &Error) -> (u8, Option<crate::scsi::ScsiSense>) { match e { @@ -23,7 +25,7 @@ pub(crate) mod macos; pub(crate) mod windows; use crate::error::{Error, Result}; -use crate::event::{Event, EventKind}; +use crate::event::Event; use crate::identity::DriveId; use crate::platform::PlatformDriver; use crate::platform::mt1959::Mt1959; @@ -114,6 +116,35 @@ impl Drive { }) } + /// Test-only constructor: build a `Drive` over an arbitrary + /// [`ScsiTransport`] (no profile, no platform driver, no block-device + /// fallback) so command-builder/response-parser logic can be exercised + /// against a scripted mock transport. + #[cfg(test)] + fn from_transport_for_test(scsi: Box<dyn ScsiTransport>) -> Self { + Drive { + scsi, + driver: None, + profile: None, + platform: None, + drive_id: DriveId { + vendor_id: String::new(), + product_id: String::new(), + product_revision: String::new(), + vendor_specific: String::new(), + firmware_date: String::new(), + serial_number: String::new(), + raw_inquiry: Vec::new(), + raw_gc_010c: Vec::new(), + }, + device_path: "test".to_string(), + halt: Arc::new(AtomicBool::new(false)), + event_fn: None, + #[cfg(target_os = "linux")] + block_dev_fd: None, + } + } + /// Get a clone of the halt flag. Set to true to interrupt Drive::read(). pub fn halt_flag(&self) -> Arc<AtomicBool> { self.halt.clone() @@ -134,15 +165,6 @@ impl Drive { self.event_fn = Some(Box::new(f)); } - #[allow(dead_code)] // public on_event registration kept; Drive currently - // has no internal emission sites after the 0.13.6 recovery strip. - // DiscStream is the BytesRead source. Plan to drop on_event in 0.14. - fn emit(&self, kind: EventKind) { - if let Some(ref f) = self.event_fn { - f(Event { kind }); - } - } - fn is_halted(&self) -> bool { self.halt.load(Ordering::Relaxed) } @@ -168,7 +190,7 @@ impl Drive { Ok(r) } - /// Close the drive cleanly. Unlocks tray, flushes SCSI state, closes fd. + /// Close the drive cleanly. Unlocks the tray and closes the fd. /// Also runs automatically on Drop as a safety net. pub fn close(self) { // cleanup() runs here via Drop @@ -245,9 +267,13 @@ impl Drive { // Bit 1: media present, Bit 0: tray open match media_status & 0x03 { 0x00 => DriveStatus::NoDisc, // tray closed, no disc - 0x01 => DriveStatus::TrayOpen, // tray open + 0x01 => DriveStatus::TrayOpen, // tray open, no media 0x02 => DriveStatus::DiscPresent, // tray closed, disc present - 0x03 => DriveStatus::DiscPresent, // tray closed, disc present + // 0x03 = tray-open bit AND media-present bit both set: + // a contradictory/transient state. Don't report it as + // ready — autorip must not start a rip on a drive that + // is still settling. Treat as tray-open. + 0x03 => DriveStatus::TrayOpen, _ => DriveStatus::Unknown, } } @@ -338,8 +364,12 @@ impl Drive { 5_000, ) .ok()?; - if r.bytes_transferred > 8 { - Some(buf[8..r.bytes_transferred].to_vec()) + // Clamp the transport-reported count to the buffer length: a + // misbehaving driver/bridge could report more bytes than the + // buffer holds, which would panic the slice. + let end = r.bytes_transferred.min(buf.len()); + if end > 8 { + Some(buf[8..end].to_vec()) } else { None } @@ -372,8 +402,9 @@ impl Drive { 5_000, ) .ok()?; - if r.bytes_transferred > 0 { - Some(buf[..r.bytes_transferred].to_vec()) + let end = r.bytes_transferred.min(buf.len()); + if end > 0 { + Some(buf[..end].to_vec()) } else { None } @@ -404,8 +435,9 @@ impl Drive { 5_000, ) .ok()?; - if r.bytes_transferred > 0 { - Some(buf[..r.bytes_transferred].to_vec()) + let end = r.bytes_transferred.min(buf.len()); + if end > 0 { + Some(buf[..end].to_vec()) } else { None } @@ -425,8 +457,9 @@ impl Drive { 5_000, ) .ok()?; - if r.bytes_transferred > 0 { - Some(buf[..r.bytes_transferred].to_vec()) + let end = r.bytes_transferred.min(buf.len()); + if end > 0 { + Some(buf[..end].to_vec()) } else { None } @@ -465,8 +498,7 @@ impl Drive { /// /// `recovery=true` uses [`crate::scsi::READ_RECOVERY_TIMEOUT_MS`] (60 s, /// matches sg_dd) for the `Disc::patch` pass; `recovery=false` uses - /// [`crate::scsi::READ_TIMEOUT_MS`] (30 s, matches the kernel's - /// `/sys/block/sr*/device/timeout` default) for `Disc::copy`'s fast + /// [`crate::scsi::READ_TIMEOUT_MS`] (10 s) for `Disc::copy`'s fast /// skip-forward sweep. Both budgets are generous enough that the drive /// can finish ECC recovery on a marginal sector — pre-0.13.21 this was /// 1.5 s on the fast path which forced the kernel mid-layer to time @@ -475,12 +507,10 @@ impl Drive { /// `DiscStream` adaptive batch halving) handles retry policy. /// /// Inline retry phases (5× gentle + reset+reopen + 5× more) were - /// removed in 0.13.6. Per - /// the stop-wedge postmortem (2026-04-25), - /// the inline reset on the LG BU40N (Initio bridge) wedged drive - /// firmware without ever recovering a sector. The remaining recovery - /// layers (Disc::patch multi-pass, DiscStream batch halving) do not - /// touch the wedge-prone reset path. + /// removed in 0.13.6: on some USB-SATA bridges the inline reset wedged + /// drive firmware without ever recovering a sector. The remaining + /// recovery layers (Disc::patch multi-pass, DiscStream batch halving) + /// do not touch the wedge-prone reset path. pub fn read(&mut self, lba: u32, count: u16, buf: &mut [u8], recovery: bool) -> Result<usize> { let timeout_ms = if recovery { crate::scsi::READ_RECOVERY_TIMEOUT_MS @@ -598,14 +628,13 @@ impl Drive { 0x00, ]; let mut buf = [0u8; 8]; - self.scsi.as_mut().execute( + let result = self.scsi.as_mut().execute( &cdb, crate::scsi::DataDirection::FromDevice, &mut buf, 5_000, )?; - let last_lba = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]); - Ok(last_lba + 1) + decode_read_capacity(&buf, result.bytes_transferred) } pub fn set_speed(&mut self, speed_kbs: u16) { @@ -762,6 +791,22 @@ pub fn find_drive() -> Option<Drive> { .find_map(|(path, _)| Drive::open(std::path::Path::new(&path)).ok()) } +/// Decode a READ CAPACITY (10) response into a sector count. +/// +/// A short transfer (`bytes_transferred < 4`, which would leave the high +/// bytes zero-initialised and decode to a bogus 1-sector disc) is rejected +/// as [`Error::DiscCapacityMalformed`]. The `0xFFFF_FFFF` "capacity exceeds +/// 32-bit" sentinel, whose `last_lba + 1` overflows `u32`, is reported as the +/// distinct [`Error::DiscCapacityOverflow`] so callers can tell an unusable +/// response apart from an over-large disc. +fn decode_read_capacity(buf: &[u8; 8], bytes_transferred: usize) -> Result<u32> { + if bytes_transferred < 4 { + return Err(Error::DiscCapacityMalformed); + } + let last_lba = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]); + last_lba.checked_add(1).ok_or(Error::DiscCapacityOverflow) +} + /// Halt-aware sleep primitive — wakes within ~100 ms of `halt` flipping /// to true. Kept for the unit tests that cover the slicing behaviour; /// production code paths no longer sleep on the recovery hot path @@ -799,9 +844,25 @@ fn discover_drives() -> Vec<(String, DriveId)> { } } -/// Resolve a device path to its raw SCSI device, with optional warning message. +/// Structured outcome of [`resolve_device`] — a machine-readable signal +/// (no English prose) the application layer can render however it likes. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeviceResolution { + /// Path resolved directly to a SCSI-generic device; no substitution. + Direct, + /// A `/dev/sr*` block path was substituted with the matching + /// `/dev/sg*` SCSI-generic device for raw access (Linux only). + SrToSg, + /// A `/dev/sr*` block path was given but no matching `/dev/sg*` + /// device could be found; the original path is returned (Linux only). + SrNoSgMatch, +} + +/// Resolve a device path to its raw SCSI device. Returns the resolved +/// path plus a structured [`DeviceResolution`] signal describing whether +/// any substitution happened; the application layer maps that to UX text. #[allow(dead_code)] -pub(crate) fn resolve_device(path: &str) -> Result<(String, Option<String>)> { +pub(crate) fn resolve_device(path: &str) -> Result<(String, DeviceResolution)> { #[cfg(target_os = "linux")] { linux::resolve_device(path) @@ -877,4 +938,110 @@ mod halt_tests { let r = sleep_until_halted(&flag, Duration::ZERO); assert!(r.is_ok()); } + + #[test] + fn read_capacity_short_transfer_is_rejected() { + // bytes_transferred < 4 must NOT decode to capacity=1 from + // zero-init bytes. + let buf = [0u8; 8]; + assert!(matches!( + decode_read_capacity(&buf, 0), + Err(Error::DiscCapacityMalformed) + )); + assert!(matches!( + decode_read_capacity(&buf, 3), + Err(Error::DiscCapacityMalformed) + )); + } + + #[test] + fn read_capacity_full_transfer_decodes_last_lba_plus_one() { + // last_lba = 0x00012344 -> capacity 0x00012345. + let buf = [0x00, 0x01, 0x23, 0x44, 0, 0, 0, 0]; + assert_eq!(decode_read_capacity(&buf, 8).unwrap(), 0x0001_2345); + } + + #[test] + fn read_capacity_overflow_is_rejected() { + // last_lba = u32::MAX (the "capacity exceeds 32-bit" sentinel) -> +1 + // overflows; reported as the distinct DiscCapacityOverflow, not the + // short-transfer DiscCapacityMalformed. + let buf = [0xFF, 0xFF, 0xFF, 0xFF, 0, 0, 0, 0]; + assert!(matches!( + decode_read_capacity(&buf, 8), + Err(Error::DiscCapacityOverflow) + )); + } +} + +#[cfg(test)] +mod command_tests { + use super::*; + use crate::scsi::{DataDirection, ScsiResult, ScsiTransport}; + + /// Mock transport: returns a fixed data payload (copied into the + /// caller's buffer, truncated to fit) on every `execute()`. + struct FixedTransport { + payload: Vec<u8>, + } + + impl ScsiTransport for FixedTransport { + fn execute( + &mut self, + _cdb: &[u8], + _direction: DataDirection, + data: &mut [u8], + _timeout_ms: u32, + ) -> Result<ScsiResult> { + let n = self.payload.len().min(data.len()); + data[..n].copy_from_slice(&self.payload[..n]); + Ok(ScsiResult { + status: 0, + bytes_transferred: n, + sense: [0u8; 32], + }) + } + } + + fn drive_with(payload: Vec<u8>) -> Drive { + Drive::from_transport_for_test(Box::new(FixedTransport { payload })) + } + + #[test] + fn read_capacity_normal_adds_one() { + // last_lba = 0x0000_0063 (99) → capacity 100 sectors. + let mut d = drive_with(vec![0x00, 0x00, 0x00, 0x63, 0x00, 0x00, 0x08, 0x00]); + assert_eq!(d.read_capacity().unwrap(), 100); + } + + #[test] + fn read_capacity_sentinel_does_not_overflow() { + // last_lba = 0xFFFF_FFFF is the "capacity exceeds 32-bit" sentinel; + // +1 would overflow. Must surface DiscCapacityOverflow, not panic + // (debug) or wrap to 0 (release). + let mut d = drive_with(vec![0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x08, 0x00]); + assert!(matches!( + d.read_capacity(), + Err(Error::DiscCapacityOverflow) + )); + } + + #[test] + fn drive_status_tray_open_and_media_present_is_not_ready_to_rip() { + // GET EVENT STATUS reply: byte 5 (media_status) low bits = 0b11 + // (tray-open AND media-present, contradictory). Must NOT report + // DiscPresent. Buffer is 8 bytes; bytes_transferred >= 6. + let mut buf = vec![0u8; 8]; + buf[5] = 0x03; + let mut d = drive_with(buf); + assert_eq!(d.drive_status(), DriveStatus::TrayOpen); + } + + #[test] + fn drive_status_disc_present_maps_correctly() { + let mut buf = vec![0u8; 8]; + buf[5] = 0x02; // media present, tray closed + let mut d = drive_with(buf); + assert_eq!(d.drive_status(), DriveStatus::DiscPresent); + } } diff --git a/src/drive/windows.rs b/src/drive/windows.rs index 8c077f6..249da9b 100644 --- a/src/drive/windows.rs +++ b/src/drive/windows.rs @@ -1,9 +1,18 @@ //! Windows drive discovery and device resolution. +use crate::drive::DeviceResolution; use crate::error::Result; use crate::identity::DriveId; use std::path::Path; +/// SCSI peripheral device type 5 = MMC / optical, in the low 5 bits of +/// INQUIRY byte 0. +const SCSI_PERIPHERAL_TYPE_OPTICAL: u8 = 0x05; + +/// Discover optical drives. Probes `\\.\CdRom0..15` first; only if none +/// are found does it fall back to scanning drive letters `D..Z`. Each +/// candidate is opened, INQUIRY'd, and kept only if its peripheral device +/// type is optical (MMC, type 0x05). Returns normalized `\\.\` paths. pub fn find_drives() -> Vec<(String, DriveId)> { let mut drives = Vec::new(); @@ -12,7 +21,9 @@ pub fn find_drives() -> Vec<(String, DriveId)> { let path = format!("\\\\.\\CdRom{}", i); if let Ok(mut transport) = crate::scsi::open(Path::new(&path)) { if let Ok(id) = DriveId::from_drive(transport.as_mut()) { - if !id.raw_inquiry.is_empty() && (id.raw_inquiry[0] & 0x1F) == 0x05 { + if !id.raw_inquiry.is_empty() + && (id.raw_inquiry[0] & 0x1F) == SCSI_PERIPHERAL_TYPE_OPTICAL + { drives.push((path, id)); } } @@ -25,8 +36,12 @@ pub fn find_drives() -> Vec<(String, DriveId)> { let path = format!("{}:", letter as char); if let Ok(mut transport) = crate::scsi::open(Path::new(&path)) { if let Ok(id) = DriveId::from_drive(transport.as_mut()) { - if !id.raw_inquiry.is_empty() && (id.raw_inquiry[0] & 0x1F) == 0x05 { - drives.push((path, id)); + if !id.raw_inquiry.is_empty() + && (id.raw_inquiry[0] & 0x1F) == SCSI_PERIPHERAL_TYPE_OPTICAL + { + // Normalize so returned paths are consistently in + // \\.\ form regardless of which loop matched. + drives.push((normalize_path(&path), id)); } } } @@ -36,8 +51,11 @@ pub fn find_drives() -> Vec<(String, DriveId)> { drives } -pub fn resolve_device(path: &str) -> Result<(String, Option<String>)> { - Ok((normalize_path(path), None)) +/// Resolve a device path to its normalized Windows `\\.\` form. Windows +/// has no `sr`→`sg` symlink-target indirection, so resolution is purely a +/// path normalization and always reports [`DeviceResolution::Direct`]. +pub fn resolve_device(path: &str) -> Result<(String, DeviceResolution)> { + Ok((normalize_path(path), DeviceResolution::Direct)) } /// Normalize a device path to Windows \\.\X: format. @@ -55,9 +73,6 @@ fn normalize_path(path: &str) -> String { if trimmed.len() == 2 && trimmed.as_bytes()[1] == b':' { return format!("\\\\.\\{}", trimmed); } - if path.to_lowercase().starts_with("cdrom") { - return format!("\\\\.\\{}", path); - } format!("\\\\.\\{}", path) } diff --git a/src/drm/mod.rs b/src/drm/mod.rs index a3138fd..513de73 100644 --- a/src/drm/mod.rs +++ b/src/drm/mod.rs @@ -8,17 +8,25 @@ //! | [`DrmScheme::Css`] | DVD probe sector flagged scrambled | //! | [`DrmScheme::Aacs10`] | Content cert type byte `0x00` | //! | [`DrmScheme::Aacs20`] | Content cert type byte `!= 0x00`, no Variant | -//! | [`DrmScheme::Aacs21`] | Content cert + MKB records `0x82` / `0x83` | +//! | [`DrmScheme::Aacs21`] | as Aacs20 + MKB Variant records `0x82`/`0x83` | +//! +//! The content cert type byte only ever decodes to V10 (`0x00`) or V20 +//! (`!= 0x00`); the V21 promotion is decided solely by the MKB Variant +//! walk in [`DrmScheme::detect`], never by the cert byte. //! //! Detection happens from a [`DrmProbe`] (raw inputs the caller has //! already extracted from the disc); resolution runs through a //! [`DrmContext`] (the full set of inputs the loaders need). //! -//! The AACS 2.1 arm is wired but disabled. The dispatcher leaves -//! [`crate::aacs::resolve_keys_v21`] reachable as a library entry point -//! for fixture-driven validation, but production consumers go through -//! [`DrmScheme::load`], which short-circuits V21 to `None` until the -//! Variant chain has a real Variant-scheme disc to validate against. +//! AACS 2.1 discs that are fully keyed in `keydb.cfg` decrypt through +//! the same classical Media Key chain as AACS 2.0, so [`DrmScheme::load`] +//! routes the `Aacs21` arm to [`crate::aacs::resolve_keys_v2`] — the +//! KEYDB lookup paths (MK+VID, disc-hash VUK, pre-decrypted unit keys) +//! succeed for any disc present in the keydb. The dedicated Variant/KCD +//! chain ([`crate::aacs::resolve_keys_v21`]) stays reachable as a +//! library entry point for fixture-driven validation but is not yet on +//! the dispatch path; it is enabled once KCD validation against a real +//! Variant-scheme disc lands. use crate::aacs; use crate::css; @@ -101,9 +109,15 @@ impl DrmScheme { /// Run key resolution for this scheme against `ctx`. /// /// Returns `None` when the scheme's resolver could not produce keys - /// (missing context, KEYDB miss, failed crypto walk, etc.) or when - /// the scheme itself is gated off (see the inline comment on the - /// `Aacs21` arm). + /// (missing context, KEYDB miss, failed crypto walk, etc.). + /// + /// The `Aacs21` arm resolves through the classical [`aacs::resolve_keys_v2`] + /// chain: AACS 2.1 discs already keyed in `keydb.cfg` decrypt identically + /// to AACS 2.0 (the KEYDB MK+VID / disc-hash VUK / pre-decrypted unit-key + /// paths succeed for any disc in the keydb), and `resolve_keys_v2` promotes + /// the resolved version to V21 when Variant MKB records are present. The + /// dedicated Variant/KCD chain ([`aacs::resolve_keys_v21`]) is kept opt-in + /// until KCD validation against a real Variant-scheme disc lands. pub fn load(self, ctx: &mut DrmContext<'_>) -> Option<ResolvedScheme> { match self { DrmScheme::Css => ctx @@ -116,20 +130,14 @@ impl DrmScheme { .as_ref() .and_then(aacs::resolve_keys_v1) .map(ResolvedScheme::Aacs), - DrmScheme::Aacs20 => ctx + // Both Aacs20 and Aacs21 route through the classical V2 chain. + // The Variant/KCD chain (resolve_keys_v21) is wired but gated; + // a keyed-in V21 disc resolves via the KEYDB paths here. + DrmScheme::Aacs20 | DrmScheme::Aacs21 => ctx .aacs .as_ref() .and_then(aacs::resolve_keys_v2) .map(ResolvedScheme::Aacs), - // AACS 2.1 derivation is wired but disabled. KCD validation - // against a Variant-scheme disc is pending. To enable, - // uncomment the line below. - // DrmScheme::Aacs21 => ctx - // .aacs - // .as_ref() - // .and_then(aacs::resolve_keys_v21) - // .map(ResolvedScheme::Aacs), - DrmScheme::Aacs21 => None, } } } @@ -240,25 +248,34 @@ mod tests { } #[test] - fn load_aacs21_returns_none() { - // The Aacs21 dispatch arm is commented out; load() must - // return None until KCD validation lands. + fn load_aacs21_routes_through_v2_resolver() { + // The Aacs21 arm shares the classical V2 resolver with Aacs20, so a + // V21 disc keyed in keydb.cfg resolves instead of being short-circuited + // to None at the dispatcher. With an EMPTY keydb both schemes fail key + // resolution identically — proving Aacs21 takes the resolver path + // rather than an unconditional `None` gate. let uk_ro = vec![0u8; 256]; let vid = [0u8; 16]; let keydb = aacs::KeyDb::empty(); let providers: &[&dyn aacs::KeyProvider] = &[&keydb]; - let ctx_aacs = aacs::ResolveContext { - unit_key_ro: &uk_ro, - content_cert: None, - volume_id: &vid, - providers, - mkb: None, - }; - let mut ctx = DrmContext { - aacs: Some(ctx_aacs), + + let make_ctx = || DrmContext { + aacs: Some(aacs::ResolveContext { + unit_key_ro: &uk_ro, + content_cert: None, + volume_id: &vid, + providers, + mkb: None, + }), css: None, }; - assert!(DrmScheme::Aacs21.load(&mut ctx).is_none()); + + let v20 = DrmScheme::Aacs20.load(&mut make_ctx()); + let v21 = DrmScheme::Aacs21.load(&mut make_ctx()); + // Same resolver, same (empty-keydb) outcome. + assert_eq!(v20.is_none(), v21.is_none()); + // Empty keydb -> no keys for either. + assert!(v21.is_none()); } /// Exercises the V21 helper directly. Gated `#[ignore]` because diff --git a/src/error.rs b/src/error.rs index 37eead9..6bd01ef 100644 --- a/src/error.rs +++ b/src/error.rs @@ -30,6 +30,7 @@ pub const E_IOKIT_PLUGIN_FAILED: u16 = 1006; // Profile (2xxx) pub const E_UNSUPPORTED_DRIVE: u16 = 2000; +// 2001: burned/retired — do not reuse. pub const E_PROFILE_PARSE: u16 = 2002; pub const E_UNSUPPORTED_PLATFORM: u16 = 2003; pub const E_PLATFORM_NOT_IMPLEMENTED: u16 = 2004; @@ -46,15 +47,18 @@ pub const E_IO_ERROR: u16 = 5000; // Disc format (6xxx) pub const E_DISC_READ: u16 = 6000; -pub const E_HALTED: u16 = 6010; pub const E_MPLS_PARSE: u16 = 6001; pub const E_CLPI_PARSE: u16 = 6002; pub const E_UDF_NOT_FOUND: u16 = 6003; +// 6004: burned/retired — do not reuse. pub const E_DISC_TITLE_RANGE: u16 = 6005; +// 6006: burned/retired — do not reuse. pub const E_IFO_PARSE: u16 = 6007; pub const E_MKV_INVALID: u16 = 6008; pub const E_NO_STREAMS: u16 = 6009; +pub const E_HALTED: u16 = 6010; pub const E_MAPFILE_INVALID: u16 = 6011; +pub const E_UDF_BUFFER_TOO_SMALL: u16 = 6012; // AACS (7xxx) pub const E_AACS_NO_KEYS: u16 = 7000; @@ -69,6 +73,7 @@ pub const E_AACS_KEY_VERIFY: u16 = 7008; pub const E_AACS_VID_READ: u16 = 7009; pub const E_AACS_VID_MAC: u16 = 7010; pub const E_AACS_DATA_KEY: u16 = 7011; +// 7012: burned/retired — do not reuse. pub const E_DECRYPT_FAILED: u16 = 7013; pub const E_CSS_AUTH_FAILED: u16 = 7014; pub const E_AACS_HOST_CERT_REJECTED: u16 = 7015; @@ -87,6 +92,8 @@ pub const E_KEYDB_INVALID: u16 = 8002; pub const E_KEYDB_WRITE: u16 = 8003; pub const E_KEYDB_PARSE: u16 = 8004; pub const E_KEYDB_LOAD: u16 = 8005; +pub const E_KEYDB_UNSUPPORTED_SCHEME: u16 = 8006; +pub const E_KEYDB_TOO_MANY_REDIRECTS: u16 = 8007; // Stream/mux (9xxx) pub const E_STREAM_READ_ONLY: u16 = 9000; @@ -99,11 +106,29 @@ pub const E_PES_INVALID_MAGIC: u16 = 9006; pub const E_ISO_TOO_LARGE: u16 = 9007; pub const E_NO_METADATA: u16 = 9008; pub const E_DISC_URL_NOT_DIRECT: u16 = 9009; +pub const E_HEVC_PARAM_PARSE: u16 = 9010; +pub const E_MUX_TRACK_RANGE: u16 = 9011; +pub const E_FMP4_UNIMPLEMENTED: u16 = 9012; +pub const E_DEMUX_THREAD_PANICKED: u16 = 9013; +pub const E_PIPELINE_JOIN_TIMEOUT: u16 = 9014; +pub const E_PIPELINE_CONSUMER_PANICKED: u16 = 9015; +pub const E_SWEEP_CONSUMER_GONE: u16 = 9016; +pub const E_PES_TRACK_TOO_LARGE: u16 = 9017; +pub const E_PIPELINE_CONSUMER_GONE: u16 = 9018; +pub const E_DISC_CAPACITY_OVERFLOW: u16 = 9020; +pub const E_M2TS_PACKET_MALFORMED: u16 = 9021; +pub const E_EXTENT_NOT_UNIT_ALIGNED: u16 = 9030; +/// READ CAPACITY returned a short or overflowing transfer. +pub const E_DISC_CAPACITY_MALFORMED: u16 = 9047; // ── Error enum ────────────────────────────────────────────────────────────── /// Structured error with numeric code and context data. No English text. +/// +/// Marked `#[non_exhaustive]`: downstream crates must not match it +/// exhaustively, so new variants can be added without a semver break. #[derive(Debug)] +#[non_exhaustive] pub enum Error { // Device (1xxx) DeviceNotFound { @@ -201,6 +226,10 @@ pub enum Error { UdfNotFound { path: String, }, + /// A `SectorSource` caller passed a destination buffer smaller than one + /// 2048-byte sector. A contract violation on the public reader API — + /// returned instead of panicking on the slice. + UdfBufferTooSmall, DiscTitleRange { index: usize, count: usize, @@ -280,6 +309,14 @@ pub enum Error { KeydbLoad { path: String, }, + /// A redirect (or the configured URL) targets a scheme this + /// dependency-light HTTP client cannot fetch (e.g. `https://`). + /// Carries the offending scheme for diagnostics. + KeydbUnsupportedScheme { + scheme: String, + }, + /// The redirect chain exceeded the follow limit. + KeydbTooManyRedirects, // Stream/mux (9xxx) StreamReadOnly, @@ -297,6 +334,12 @@ pub enum Error { size: usize, }, PesInvalidMagic, + /// PES frame track index exceeds the 1-byte on-wire field (> 255). + /// Carries the offending index. Distinct from [`Error::PesInvalidMagic`], + /// which signals corrupt input on the read side. + PesTrackTooLarge { + track: usize, + }, IsoTooLarge { path: String, }, @@ -305,6 +348,65 @@ pub enum Error { /// `Drive::open() + Disc::scan() + DiscStream::new()` directly. This /// is a structural API constraint, not a parse failure. DiscUrlNotDirect, + /// A non-empty `HEVCDecoderConfigurationRecord` (hvcC) was supplied to + /// a muxer but failed to parse into any VPS/SPS/PPS NAL — emitting the + /// stream without parameter sets would yield an undecodable result. + HevcParamParse, + /// A muxer `write_frame` / `set_codec_private` was given a track index + /// beyond the configured PID/track count. + MuxTrackRange { + track: usize, + tracks: usize, + }, + /// The fragmented-MP4 sink cannot emit media — `moof`/`mdat` framing is + /// not implemented. Surfaced instead of silently discarding samples. + Fmp4Unimplemented, + /// A worker thread in the threaded mux pipeline terminated without + /// sending its terminal sentinel — i.e. it panicked or was dropped + /// mid-stream. Surfaced so a parser/demux panic is never silently + /// reported to the caller as a clean end-of-stream (which would + /// truncate output without any error). + DemuxThreadPanicked, + /// A pipeline `join()` exceeded its deadline while waiting for the + /// consumer thread to drain. The consumer is intentionally leaked; + /// the caller should fall back to a degraded path. + PipelineJoinTimeout, + /// The pipeline consumer thread panicked. The original panic + /// payload is not preserved (no English text in the library); it is + /// logged at the panic site instead. + PipelineConsumerPanicked, + /// A pipeline producer's `send` failed because the consumer thread + /// has already terminated (the receiver end is gone). + SweepConsumerGone, + /// A producer thread tried to hand work to its pipeline consumer + /// (sweep / patch sink) but the consumer thread had already + /// terminated (panicked or dropped the receiver). The producer + /// surfaces this so the outer pass can abort cleanly instead of + /// blocking on a dead channel. + PipelineConsumerGone, + /// READ CAPACITY(10) reported a last-LBA of `0xFFFFFFFF` — the SPC + /// sentinel meaning "capacity exceeds 32-bit addressing". Adding 1 to + /// derive the sector count would overflow `u32`. Reachable from + /// disc-reported bytes and synthetic [`crate::sector::SectorSource`] + /// fixtures. + DiscCapacityOverflow, + /// An extent fed to the prefetch producer has a `sector_count` + /// whose trailing 1-2 sectors cannot form a complete AACS aligned + /// unit (3 sectors / 6144 bytes). Emitting that tail as a + /// standalone batch would hand the decrypt step a sub-unit chunk + /// it silently leaves encrypted. The producer surfaces this rather + /// than emit still-encrypted bytes. + ExtentNotUnitAligned, + /// An MPEG-TS packet under construction violated the 188-byte fixed + /// size (over-long adaptation field, overflowing payload, or a + /// short/mis-assembled packet). Indicates a muxer invariant break, + /// not untrusted input — surfaced instead of writing a corrupt + /// transport stream. + M2tsPacketMalformed, + /// READ CAPACITY transferred fewer than 4 bytes, or the decoded + /// last-LBA + 1 overflowed `u32`. Either case means the capacity + /// response is unusable; no English commentary. + DiscCapacityMalformed, } impl Error { @@ -330,6 +432,7 @@ impl Error { Error::MplsParse => E_MPLS_PARSE, Error::ClpiParse => E_CLPI_PARSE, Error::UdfNotFound { .. } => E_UDF_NOT_FOUND, + Error::UdfBufferTooSmall => E_UDF_BUFFER_TOO_SMALL, Error::DiscTitleRange { .. } => E_DISC_TITLE_RANGE, Error::IfoParse => E_IFO_PARSE, Error::MkvInvalid => E_MKV_INVALID, @@ -363,6 +466,8 @@ impl Error { Error::KeydbWrite { .. } => E_KEYDB_WRITE, Error::KeydbParse => E_KEYDB_PARSE, Error::KeydbLoad { .. } => E_KEYDB_LOAD, + Error::KeydbUnsupportedScheme { .. } => E_KEYDB_UNSUPPORTED_SCHEME, + Error::KeydbTooManyRedirects => E_KEYDB_TOO_MANY_REDIRECTS, Error::StreamReadOnly => E_STREAM_READ_ONLY, Error::StreamWriteOnly => E_STREAM_WRITE_ONLY, Error::StreamUrlInvalid { .. } => E_STREAM_URL_INVALID, @@ -370,9 +475,22 @@ impl Error { Error::StreamUrlMissingPort { .. } => E_STREAM_URL_MISSING_PORT, Error::PesFrameTooLarge { .. } => E_PES_FRAME_TOO_LARGE, Error::PesInvalidMagic => E_PES_INVALID_MAGIC, + Error::PesTrackTooLarge { .. } => E_PES_TRACK_TOO_LARGE, Error::IsoTooLarge { .. } => E_ISO_TOO_LARGE, Error::NoMetadata => E_NO_METADATA, Error::DiscUrlNotDirect => E_DISC_URL_NOT_DIRECT, + Error::HevcParamParse => E_HEVC_PARAM_PARSE, + Error::MuxTrackRange { .. } => E_MUX_TRACK_RANGE, + Error::Fmp4Unimplemented => E_FMP4_UNIMPLEMENTED, + Error::DemuxThreadPanicked => E_DEMUX_THREAD_PANICKED, + Error::PipelineJoinTimeout => E_PIPELINE_JOIN_TIMEOUT, + Error::PipelineConsumerPanicked => E_PIPELINE_CONSUMER_PANICKED, + Error::SweepConsumerGone => E_SWEEP_CONSUMER_GONE, + Error::PipelineConsumerGone => E_PIPELINE_CONSUMER_GONE, + Error::DiscCapacityOverflow => E_DISC_CAPACITY_OVERFLOW, + Error::ExtentNotUnitAligned => E_EXTENT_NOT_UNIT_ALIGNED, + Error::M2tsPacketMalformed => E_M2TS_PACKET_MALFORMED, + Error::DiscCapacityMalformed => E_DISC_CAPACITY_MALFORMED, } } } @@ -443,7 +561,13 @@ impl std::fmt::Display for Error { ), None => write!(f, "E{}: 0x{:02x}/0x{:02x}", self.code(), opcode, status,), }, - Error::IoError { source } => write!(f, "E{}: {}", self.code(), source), + // Language-neutral: std::io::Error's Display is English + // ("permission denied"); emit the raw OS errno when present, + // else the ErrorKind debug name (an identifier, not prose). + Error::IoError { source } => match source.raw_os_error() { + Some(errno) => write!(f, "E{}: {}", self.code(), errno), + None => write!(f, "E{}: {:?}", self.code(), source.kind()), + }, Error::DiscRead { sector, status, @@ -451,21 +575,23 @@ impl std::fmt::Display for Error { } => match (status, sense) { (Some(st), Some(s)) => write!( f, - "E{}: {} 0x{:02x}/0x{:02x}/0x{:02x}", + "E{}: {} 0x{:02x}/0x{:02x}/0x{:02x}/0x{:02x}", self.code(), sector, st, s.sense_key, s.asc, + s.ascq, ), (Some(st), None) => write!(f, "E{}: {} 0x{:02x}", self.code(), sector, st,), (None, Some(s)) => write!( f, - "E{}: {} 0x{:02x}/0x{:02x}", + "E{}: {} 0x{:02x}/0x{:02x}/0x{:02x}", self.code(), sector, s.sense_key, s.asc, + s.ascq, ), (None, None) => write!(f, "E{}: {}", self.code(), sector), }, @@ -478,12 +604,25 @@ impl std::fmt::Display for Error { Error::KeydbHttp { status } => write!(f, "E{}: {}", self.code(), status), Error::KeydbWrite { path } => write!(f, "E{}: {}", self.code(), path), Error::KeydbLoad { path } => write!(f, "E{}: {}", self.code(), path), + Error::KeydbUnsupportedScheme { scheme } => { + write!(f, "E{}: {}", self.code(), scheme) + } Error::StreamUrlInvalid { url } => write!(f, "E{}: {}", self.code(), url), Error::StreamUrlMissingPath { scheme } => write!(f, "E{}: {}", self.code(), scheme), Error::StreamUrlMissingPort { addr } => write!(f, "E{}: {}", self.code(), addr), Error::PesFrameTooLarge { size } => write!(f, "E{}: {}", self.code(), size), + Error::PesTrackTooLarge { track } => write!(f, "E{}: {}", self.code(), track), Error::IsoTooLarge { path } => write!(f, "E{}: {}", self.code(), path), - Error::NoDiscKey { disc_hash } => write!(f, "E{}: {}", self.code(), disc_hash), + Error::NoDiscKey { disc_hash } => { + if disc_hash.is_empty() { + write!(f, "E{}", self.code()) + } else { + write!(f, "E{}: {}", self.code(), disc_hash) + } + } + Error::MuxTrackRange { track, tracks } => { + write!(f, "E{}: {}/{}", self.code(), track, tracks) + } _ => write!(f, "E{}", self.code()), } } @@ -506,10 +645,21 @@ impl From<std::io::Error> for Error { impl From<Error> for std::io::Error { fn from(e: Error) -> Self { + // An `Error::IoError` is just a wrapper around an underlying + // `io::Error` that entered via `From<io::Error> for Error`. + // Round-trip it back unchanged so the original `ErrorKind` and + // raw OS error code survive instead of being flattened to + // `Other` with a stringified message. + if let Error::IoError { source } = e { + return source; + } let code = e.code(); let msg = e.to_string(); // Map our error categories to io::ErrorKind let kind = match code { + // Device access-denied semantics map to PermissionDenied; + // the rest of the 1xxx block is "device absent" -> NotFound. + E_DEVICE_PERMISSION | E_DEVICE_LOCKED => std::io::ErrorKind::PermissionDenied, 1000..=1999 => std::io::ErrorKind::NotFound, 2000..=2999 => std::io::ErrorKind::Unsupported, 3000..=3999 => std::io::ErrorKind::PermissionDenied, @@ -523,6 +673,28 @@ impl From<Error> for std::io::Error { // 9009 DiscUrlNotDirect: structurally unsupported entry point, // not a parse failure — caller used the wrong API. 9009 => std::io::ErrorKind::Unsupported, + // 9010 HevcParamParse: malformed hvcC payload. + 9010 => std::io::ErrorKind::InvalidData, + // 9011 MuxTrackRange: caller passed a bad track index. + 9011 => std::io::ErrorKind::InvalidInput, + // 9012 Fmp4Unimplemented: sink can't emit media yet. + 9012 => std::io::ErrorKind::Unsupported, + // 9014 PipelineJoinTimeout: consumer drain exceeded deadline. + E_PIPELINE_JOIN_TIMEOUT => std::io::ErrorKind::TimedOut, + // 9017 PesTrackTooLarge: out-of-range track index on serialize. + 9017 => std::io::ErrorKind::InvalidInput, + // 9020 DiscCapacityOverflow: disc reported a capacity sentinel + // we can't represent — treat as bad/invalid device data. + 9020 => std::io::ErrorKind::InvalidData, + // 9021 M2tsPacketMalformed: a muxer invariant break produced + // a non-188-byte packet — treat as invalid data. + 9021 => std::io::ErrorKind::InvalidData, + // 9030 ExtentNotUnitAligned: a malformed/non-AACS-aligned + // extent was handed to the prefetch producer. + 9030 => std::io::ErrorKind::InvalidInput, + // 9047 DiscCapacityMalformed: the drive returned an unusable + // READ CAPACITY response (short transfer / overflow). + 9047 => std::io::ErrorKind::InvalidData, _ => std::io::ErrorKind::Other, }; std::io::Error::new(kind, msg) @@ -590,6 +762,8 @@ impl Error { /// /// - MEDIUM ERROR (sense key 3) — canonical bad-sector signal /// - ABORTED COMMAND (sense key B) — transient; retry usually works + /// - NOT READY (sense key 2) — the dominant bad-sector response on + /// the BU40N (ASC 0x04/ASCQ 0x3E); a pause + retry often recovers /// - RECOVERED ERROR (sense key 1) / NO SENSE (sense key 0) — not /// classified as fatal; treat as recoverable /// @@ -636,6 +810,9 @@ mod tests { .code(), Error::MapfileInvalid { kind: "hex" }.code(), Error::DiscUrlNotDirect.code(), + Error::ExtentNotUnitAligned.code(), + Error::M2tsPacketMalformed.code(), + Error::DiscCapacityMalformed.code(), ]; let mut sorted = codes.to_vec(); sorted.sort(); @@ -680,6 +857,7 @@ mod tests { ), (Error::MapfileInvalid { kind: "hex" }, E_MAPFILE_INVALID), (Error::DiscUrlNotDirect, E_DISC_URL_NOT_DIRECT), + (Error::ExtentNotUnitAligned, E_EXTENT_NOT_UNIT_ALIGNED), ]; for (e, want_code) in cases { let s = e.to_string(); @@ -695,9 +873,7 @@ mod tests { // not). for word in s.split(|c: char| !c.is_ascii_alphabetic()) { assert!( - word.len() <= 8 - || word.eq_ignore_ascii_case("renesas") - || word.eq_ignore_ascii_case("freebsd"), + word.len() <= 8, "Display contains suspicious English-looking word `{word}` in `{s}`" ); } @@ -711,17 +887,26 @@ mod tests { let io: std::io::Error = e.into(); io.kind() }; - // 1xxx range → NotFound + // 1xxx "device absent" → NotFound assert_eq!( mapped(Error::ScsiInterfaceUnavailable { path: "p".into() }), ErrorKind::NotFound ); + assert_eq!( + mapped(Error::DeviceNotFound { path: "p".into() }), + ErrorKind::NotFound + ); + // 1xxx access-denied semantics → PermissionDenied (not NotFound) + assert_eq!( + mapped(Error::DevicePermission { path: "p".into() }), + ErrorKind::PermissionDenied + ); assert_eq!( mapped(Error::DeviceLocked { path: "p".into(), kr: 0 }), - ErrorKind::NotFound + ErrorKind::PermissionDenied ); // 2xxx range → Unsupported assert_eq!( @@ -741,5 +926,62 @@ mod tests { ); // 9009 special-cased to Unsupported assert_eq!(mapped(Error::DiscUrlNotDirect), ErrorKind::Unsupported); + // 9021 special-cased to InvalidData + assert_eq!(mapped(Error::M2tsPacketMalformed), ErrorKind::InvalidData); + // 9047 DiscCapacityMalformed → InvalidData + assert_eq!(mapped(Error::DiscCapacityMalformed), ErrorKind::InvalidData); + } + + /// `Error::IoError` must round-trip back to the *original* + /// `io::Error` — preserving its `ErrorKind` and raw OS error — + /// rather than being flattened to `Other` with a stringified + /// message. + #[test] + fn ioerror_roundtrips_preserving_kind_and_oscode() { + use std::io::ErrorKind; + let original = std::io::Error::from_raw_os_error(13); // EACCES + let original_kind = original.kind(); + let wrapped: Error = original.into(); // From<io::Error> for Error + let back: std::io::Error = wrapped.into(); // From<Error> for io::Error + assert_eq!(back.kind(), original_kind); + assert_eq!(back.raw_os_error(), Some(13)); + + // A synthesized kind (no OS code) must also survive. + let timeout: Error = std::io::Error::from(ErrorKind::TimedOut).into(); + let back2: std::io::Error = timeout.into(); + assert_eq!(back2.kind(), ErrorKind::TimedOut); + } + + /// `DiscRead` Display must include the ASCQ byte (the 5th field) so + /// NOT_READY substates (0x04/0x3E vs 0x04/0x01) are distinguishable + /// in logs and bug reports. + #[test] + fn discread_display_includes_ascq() { + let e = Error::DiscRead { + sector: 42, + status: Some(0x02), + sense: Some(crate::scsi::ScsiSense { + sense_key: 0x02, + asc: 0x04, + ascq: 0x3e, + }), + }; + let s = e.to_string(); + // sense_key/asc/ascq triple all present. + assert!(s.contains("0x02/0x04/0x3e"), "ascq missing from `{s}`"); + } + + /// `NoDiscKey` with an empty hash must not emit a dangling + /// "colon space" suffix. + #[test] + fn nodisckey_empty_hash_has_no_trailing_colon() { + let e = Error::NoDiscKey { + disc_hash: String::new(), + }; + assert_eq!(e.to_string(), format!("E{}", E_NO_DISC_KEY)); + let e2 = Error::NoDiscKey { + disc_hash: "abc".into(), + }; + assert_eq!(e2.to_string(), format!("E{}: abc", E_NO_DISC_KEY)); } } diff --git a/src/event.rs b/src/event.rs index ebb2cf4..b90a063 100644 --- a/src/event.rs +++ b/src/event.rs @@ -8,11 +8,17 @@ //! disc.rip(&mut session, 0, output, |event| { //! match event.kind { //! EventKind::BytesRead { bytes, total } => update_progress(bytes, total), -//! EventKind::ReadError { sector, .. } => log_error(sector), +//! EventKind::SectorSkipped { sector } => log_skip(sector), +//! EventKind::BatchSizeChanged { new_size, .. } => note_recovery(new_size), //! _ => {} //! } //! }); //! ``` +//! +//! Note: the library currently emits only `BytesRead`, `SectorSkipped`, +//! and `BatchSizeChanged`. The other [`EventKind`] variants are part of +//! the stable event vocabulary for consumers (and future emit sites) but +//! are not produced by the library today. use crate::error::Error; diff --git a/src/halt.rs b/src/halt.rs index 06bf932..08f6a11 100644 --- a/src/halt.rs +++ b/src/halt.rs @@ -34,11 +34,9 @@ impl Halt { Self(Arc::new(AtomicBool::new(false))) } - /// Wrap an existing `Arc<AtomicBool>` as a `Halt`. Useful as a - /// bridge during the 0.18 deprecation window: callers that already - /// hold an `Arc<AtomicBool>` (e.g. `Drive::halt_flag()`, the - /// deprecated `DiscStream::set_halt`) can adopt the new token API - /// without changing the underlying flag. + /// Wrap an existing `Arc<AtomicBool>` as a `Halt`. A bridge for + /// callers that already hold an `Arc<AtomicBool>` cancellation flag + /// and want to adopt the token API without allocating a new flag. /// /// Cancelling either side flips the same bit — the wrapping `Halt` /// and the original `Arc` are two views over one shared flag. @@ -46,10 +44,9 @@ impl Halt { Self(flag) } - /// Borrow the underlying `Arc<AtomicBool>`. Used at boundaries with - /// pre-`Halt` APIs that still take an `Arc<AtomicBool>` directly - /// (`CopyOptions::halt`, the deprecated `DiscStream::set_halt`). - /// Round 3 deletes those boundaries and this accessor with them. + /// Borrow the underlying `Arc<AtomicBool>`. The inverse of + /// [`from_arc`](Self::from_arc): hand the shared flag to an API that + /// still takes a raw `Arc<AtomicBool>` rather than a `Halt`. pub fn as_arc(&self) -> &Arc<AtomicBool> { &self.0 } diff --git a/src/identity.rs b/src/identity.rs index 317227f..c5d6cef 100644 --- a/src/identity.rs +++ b/src/identity.rs @@ -57,27 +57,48 @@ impl DriveId { let cdb_inq = [0x12, 0x00, 0x00, 0x00, 0x60, 0x00]; transport.execute(&cdb_inq, DataDirection::FromDevice, &mut inquiry, 5000)?; - // GET CONFIGURATION Feature 010Ch — MMC-6 §6.6 + // GET CONFIGURATION Feature 010Ch — MMC-6 §6.6. + // Best-effort: 010Ch (Firmware Information) is an optional feature. + // A drive that lacks it may CHECK CONDITION rather than return an + // empty descriptor, so a failure here is treated as feature-absent + // (empty firmware date + empty raw bytes) instead of aborting the + // whole identity probe. let mut gc = vec![0u8; 256]; let cdb_gc = [0x46, 0x02, 0x01, 0x0C, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00]; - let result = transport.execute(&cdb_gc, DataDirection::FromDevice, &mut gc, 5000)?; + // `bytes_transferred` is device-reported and untrusted; clamp every + // slice end to the actual buffer length before indexing. + let (firmware_date, raw_gc_010c) = + match transport.execute(&cdb_gc, DataDirection::FromDevice, &mut gc, 5000) { + Ok(result) => { + let end = result.bytes_transferred.min(gc.len()); + let date = if end > 12 { + String::from_utf8_lossy(&gc[12..24.min(end)]) + .trim() + .to_string() + } else { + String::new() + }; + (date, gc[..end].to_vec()) + } + Err(_) => (String::new(), Vec::new()), + }; - let firmware_date = if result.bytes_transferred > 12 { - String::from_utf8_lossy(&gc[12..24.min(result.bytes_transferred)]) - .trim() - .to_string() - } else { - String::new() - }; - - // GET CONFIGURATION Feature 0108h — Serial Number + // GET CONFIGURATION Feature 0108h — Serial Number. + // Best-effort, like 010Ch above: the serial-number feature is + // optional, so a drive that lacks it (CHECK CONDITION) or reports + // too few bytes deliberately yields an empty serial rather than + // failing the identity probe. let mut gc_serial = vec![0u8; 256]; let cdb_serial = [0x46, 0x02, 0x01, 0x08, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00]; let serial_number = if let Ok(r) = transport.execute(&cdb_serial, DataDirection::FromDevice, &mut gc_serial, 5000) { if r.bytes_transferred > 12 { - String::from_utf8_lossy(&gc_serial[12..r.bytes_transferred]) + // `bytes_transferred` is device-reported and untrusted; clamp + // the slice end to the buffer length to avoid an out-of-range + // panic on an oversized reported count. + let end = r.bytes_transferred.min(gc_serial.len()); + String::from_utf8_lossy(&gc_serial[12..end]) .trim() .to_string() } else { @@ -94,8 +115,8 @@ impl DriveId { vendor_specific: ascii_field(&inquiry, 36, 43), firmware_date, serial_number, - raw_inquiry: inquiry.to_vec(), - raw_gc_010c: gc[..result.bytes_transferred].to_vec(), + raw_inquiry: inquiry, + raw_gc_010c, }) } @@ -155,6 +176,49 @@ fn ascii_field(data: &[u8], start: usize, end: usize) -> String { #[cfg(test)] mod tests { use super::*; + use crate::scsi::{ScsiResult, ScsiTransport}; + + /// Transport that returns the requested data length but reports a + /// bytes_transferred larger than the caller's buffer — models a drive + /// that lies about its transfer count. The old slicing code panicked + /// on this; the clamps must keep it from indexing out of range. + struct OversizedCountTransport; + + impl ScsiTransport for OversizedCountTransport { + fn execute( + &mut self, + cdb: &[u8], + _dir: DataDirection, + buf: &mut [u8], + _timeout_ms: u32, + ) -> Result<ScsiResult> { + // Fill plausible ASCII so the from_utf8_lossy paths run. + for b in buf.iter_mut() { + *b = b'A'; + } + // INQUIRY (0x12): honest count. GET CONFIGURATION (0x46): lie. + let bytes_transferred = if cdb.first() == Some(&0x12) { + buf.len() + } else { + buf.len() + 4096 + }; + Ok(ScsiResult { + status: 0, + bytes_transferred, + sense: [0u8; 32], + }) + } + } + + #[test] + fn from_drive_clamps_oversized_bytes_transferred() { + // Must not panic despite the transport reporting a transfer count + // far beyond the 256-byte GET CONFIGURATION buffers. + let mut t = OversizedCountTransport; + let id = DriveId::from_drive(&mut t).expect("from_drive must not error"); + // raw_gc_010c is clamped to the 256-byte buffer, never the lie. + assert_eq!(id.raw_gc_010c.len(), 256); + } #[test] fn test_bu40n_identity() { diff --git a/src/ifo.rs b/src/ifo.rs index 4a07348..5d40ef6 100644 --- a/src/ifo.rs +++ b/src/ifo.rs @@ -77,6 +77,14 @@ pub struct DvdAudioAttr { pub channels: u8, pub sample_rate: u32, pub language: String, + /// The PES `private_stream_1` sub-stream id this audio stream carries + /// on the wire (AC-3: `0x80..=0x87`, DTS: `0x88..=0x8F`, LPCM: + /// `0xA0..=0xA7`), assigned by per-codec ordinal during the scan. + /// `None` for codecs carried as a regular MPEG-audio PES (MP1/MP2, + /// stream_id `0xC0..`) which don't use a private-stream-1 sub-id. + /// This is the single routing key shared with the muxer's `dvd_pid()` + /// so the two never disagree on a mixed-codec title. + pub sub_stream_id: Option<u8>, } /// DVD subtitle stream attributes. @@ -301,6 +309,12 @@ fn parse_vts( } audio_streams.push(parse_audio_attr(&vts_data, aoff)?); } + // Assign each audio stream its on-wire private_stream_1 sub-stream id + // by per-codec ordinal — the same convention DVD authoring uses (AC-3 + // 0x80+, DTS 0x88+, LPCM 0xA0+). This is the routing key shared with + // the muxer; per-codec ordinals (not the positional index) are what + // keep mixed-codec titles from colliding. + assign_audio_sub_stream_ids(&mut audio_streams); // Subtitle streams: count at 0x254 (u16 BE), then 6 bytes each starting at 0x256 let num_subs = if vts_data.len() >= 0x256 { @@ -415,9 +429,49 @@ fn parse_audio_attr(data: &[u8], offset: usize) -> Result<DvdAudioAttr> { channels, sample_rate, language, + // Assigned by `assign_audio_sub_stream_ids` once all streams in the + // title set are known (the sub-id is a per-codec ordinal). + sub_stream_id: None, }) } +/// Assign the on-wire `private_stream_1` sub-stream id to each audio +/// stream by per-codec ordinal, matching DVD authoring convention and the +/// muxer's `dvd_pid()` routing: +/// - AC-3 → `0x80 + n` (n = 0-based index among AC-3 streams) +/// - DTS → `0x88 + n` +/// - LPCM → `0xA0 + n` +/// - MP1/MP2 and anything else → `None` (regular MPEG-audio PES, not a +/// private-stream-1 sub-id). +/// +/// Indices saturate at the codec range ceiling (8 AC-3/DTS, 8 LPCM) so a +/// malformed over-count never produces an out-of-range sub-id. +fn assign_audio_sub_stream_ids(streams: &mut [DvdAudioAttr]) { + let mut n_ac3 = 0u8; + let mut n_dts = 0u8; + let mut n_lpcm = 0u8; + for s in streams.iter_mut() { + s.sub_stream_id = match s.codec { + Codec::Ac3 => { + let id = 0x80 + n_ac3.min(7); + n_ac3 = n_ac3.saturating_add(1); + Some(id) + } + Codec::Dts => { + let id = 0x88 + n_dts.min(7); + n_dts = n_dts.saturating_add(1); + Some(id) + } + Codec::Lpcm => { + let id = 0xA0 + n_lpcm.min(7); + n_lpcm = n_lpcm.saturating_add(1); + Some(id) + } + _ => None, + }; + } +} + /// Parse one subtitle stream attribute block (6 bytes at `offset`). fn parse_subtitle_attr(data: &[u8], offset: usize) -> Result<DvdSubtitleAttr> { // Language code: bytes 2-3 as ISO 639 @@ -481,7 +535,10 @@ fn parse_pgcit( match parse_pgc(data, pgc_abs, chapter_count) { Ok(title) => titles.push(title), - Err(_) => continue, // skip malformed PGCs + // By design: a single unparseable PGC (truncated/corrupt entry, + // authoring-tool quirk) must not lose the whole title list. + // Skip it and keep collecting the titles that do parse. + Err(_) => continue, } } @@ -707,6 +764,7 @@ mod tests { channels: 6, sample_rate: 48000, language: "en".to_string(), + sub_stream_id: Some(0x80), }; assert_eq!(audio.channels, 6); @@ -826,6 +884,54 @@ mod tests { assert_eq!(attr.language, "en"); } + #[test] + fn mixed_codec_sub_stream_ids_are_distinct() { + // A title mixing AC-3, DTS and LPCM must get per-codec ordinal + // sub-ids (0x80, 0x88, 0xA0...), all distinct — this is the + // routing key that keeps mixed-codec audio from colliding. + let mut streams = vec![ + DvdAudioAttr { + codec: Codec::Ac3, + channels: 6, + sample_rate: 48000, + language: "en".into(), + sub_stream_id: None, + }, + DvdAudioAttr { + codec: Codec::Dts, + channels: 6, + sample_rate: 48000, + language: "en".into(), + sub_stream_id: None, + }, + DvdAudioAttr { + codec: Codec::Lpcm, + channels: 2, + sample_rate: 48000, + language: "fr".into(), + sub_stream_id: None, + }, + DvdAudioAttr { + codec: Codec::Ac3, + channels: 2, + sample_rate: 48000, + language: "es".into(), + sub_stream_id: None, + }, + ]; + assign_audio_sub_stream_ids(&mut streams); + assert_eq!(streams[0].sub_stream_id, Some(0x80)); // AC-3 #0 + assert_eq!(streams[1].sub_stream_id, Some(0x88)); // DTS #0 + assert_eq!(streams[2].sub_stream_id, Some(0xA0)); // LPCM #0 + assert_eq!(streams[3].sub_stream_id, Some(0x81)); // AC-3 #1 + // All sub-ids unique. + let ids: Vec<u8> = streams.iter().filter_map(|s| s.sub_stream_id).collect(); + let mut sorted = ids.clone(); + sorted.sort_unstable(); + sorted.dedup(); + assert_eq!(ids.len(), sorted.len(), "sub-stream ids must be unique"); + } + #[test] fn audio_attr_dts() { let mut data = vec![0u8; 16]; diff --git a/src/io/bounded.rs b/src/io/bounded.rs index 45b83f9..b96248a 100644 --- a/src/io/bounded.rs +++ b/src/io/bounded.rs @@ -67,10 +67,12 @@ pub(crate) enum BoundedError { /// The deadline elapsed before the syscall returned. Same leak /// semantics as `Halted`. Timeout, - /// The worker thread panicked, or its sender disconnected before - /// sending a result. Treat as a benign no-op (callers usually - /// log and continue) rather than a hard error — by definition no - /// syscall observably ran to completion in this case. + /// The worker thread panicked, the OS rejected the thread spawn, + /// or its sender disconnected before sending a result. Treat as a + /// benign no-op (callers usually log and continue) rather than a + /// hard error — by definition no syscall observably ran to + /// completion in this case. In the spawn-failure case no thread is + /// leaked. WorkerLost, } @@ -101,6 +103,12 @@ where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { + // If the caller already requested halt, don't spawn (and leak) a + // worker that would run `op` to completion in the background. + if halt.is_some_and(|h| h.is_cancelled()) { + return Err(BoundedError::Halted); + } + // Rendezvous channel: the worker sends exactly one value (the // op's return) and then exits. Capacity-0 means the send blocks // until we receive — fine on the happy path; on the timeout / diff --git a/src/io/byte_channel.rs b/src/io/byte_channel.rs deleted file mode 100644 index 80a3693..0000000 --- a/src/io/byte_channel.rs +++ /dev/null @@ -1,327 +0,0 @@ -//! Byte-sized bounded producer/consumer channel. -//! -//! Wraps `std::sync::mpsc::sync_channel` with a byte-accounting -//! `Mutex<usize> + Condvar` cap. Sender blocks (cooperatively) when -//! `used_bytes + item.byte_size() > capacity_bytes`. Receiver -//! decrements `used_bytes` when it takes the item. -//! -//! Why: the existing producer→consumer channel between `DiscStream` -//! (PES producer) and `MuxSink` (PES consumer) is bounded by frame -//! count. Frame sizes vary 100× between metadata and keyframes, so a -//! count-based cap either starves on small frames or buffers far too -//! much memory on big ones. Byte-sized accounting sizes the buffer for -//! the worst-case input stall (NFS read p99 ≈ 1–2 s × ~15 MB/s peak -//! compressed bitrate ≈ ~30 MB) directly. -//! -//! The underlying mpsc channel is created with a very large slot count -//! so the byte cap (not the slot count) is the real backpressure. Slot -//! count is only there to give the kernel a small chunk to wake on. - -use std::sync::mpsc::{Receiver as MpscReceiver, RecvError, SendError, SyncSender, sync_channel}; -use std::sync::{Arc, Condvar, Mutex}; - -/// Default byte cap for the muxer's input channel. Sized to hide a -/// worst-case ~2 s NFS read refill at UHD peak compressed bitrate -/// (~15 MB/s); 64 MiB gives headroom. Tweakable; not magic. -pub const BYTE_CHANNEL_DEFAULT_CAPACITY: usize = 64 * 1024 * 1024; - -/// Slot capacity of the inner `sync_channel`. Large so the byte cap is -/// the real backpressure mechanism — the mpsc slot count only exists -/// to give the kernel a chunk to wake on. PES frames are typically -/// ~700 B each, so 64 MiB ≈ 90 k frames; 200 k is comfortable headroom. -const INNER_SLOT_CAPACITY: usize = 200_000; - -/// Anything whose in-memory cost can be accounted by a single -/// `usize`. Implement on the item type sent through [`Sender`]. -pub trait HasByteSize { - /// Bytes this item contributes to the channel's used budget. - /// Must be > 0 to make progress (a 0-byte item would never - /// block the sender no matter the cap; see send_blocks_at_capacity - /// test). - fn byte_size(&self) -> usize; -} - -impl HasByteSize for crate::pes::PesFrame { - fn byte_size(&self) -> usize { - // Frame data + the fixed header overhead the serializer - // writes (track + pts + keyframe + len). The `Vec<u8>` heap - // allocation also has alloc-header overhead but that's - // <0.1 % at typical frame sizes — folding it in would just - // add noise to the budget. - self.data.len() + 14 - } -} - -/// Shared book-keeping between [`Sender`] and [`Receiver`]. Wrapped in -/// an `Arc` because both halves hold it independently. -struct Accounting { - used: Mutex<usize>, - cv: Condvar, - capacity: usize, -} - -/// Send half of the byte-bounded channel. -/// -/// `send` blocks (on a `Condvar`) when adding the item would push -/// `used_bytes` past `capacity_bytes`. Unblocks when the receiver -/// `recv`s items out and notifies. Returns `Err(item)` if the -/// receiver has been dropped — mirrors `mpsc::SyncSender::send`. -pub struct Sender<T: HasByteSize> { - tx: SyncSender<T>, - acct: Arc<Accounting>, -} - -impl<T: HasByteSize> Clone for Sender<T> { - fn clone(&self) -> Self { - Sender { - tx: self.tx.clone(), - acct: self.acct.clone(), - } - } -} - -impl<T: HasByteSize> Sender<T> { - /// Push one item. Blocks until adding it would not exceed the - /// capacity, then sends through the inner mpsc channel. - pub fn send(&self, item: T) -> Result<(), SendError<T>> { - let sz = item.byte_size(); - // Reserve capacity first. The reservation is observable to - // other senders via `used`; only after we win the slot do we - // hand the item to the inner mpsc channel. That ordering means - // `used` is always a conservative upper bound on what's in the - // mpsc queue + about-to-be-sent. - { - let mut used = self.acct.used.lock().expect("byte_channel poisoned"); - // An item bigger than the whole capacity will never fit; let - // it through anyway as a one-shot reservation, otherwise the - // sender deadlocks forever waiting for `used == 0` AND - // nothing in flight. The receiver will drain it on the - // other side. Same behaviour as `std::sync::mpsc` for - // arbitrarily large messages. - while *used + sz > self.acct.capacity && *used > 0 { - used = self.acct.cv.wait(used).expect("byte_channel cv poisoned"); - } - *used += sz; - } - match self.tx.send(item) { - Ok(()) => Ok(()), - Err(SendError(returned)) => { - // Receiver dropped — refund the reservation so a later - // sender on a clone doesn't observe phantom used bytes - // (the receiver is gone so nobody will decrement). - let mut used = self.acct.used.lock().expect("byte_channel poisoned"); - *used = used.saturating_sub(sz); - self.acct.cv.notify_all(); - Err(SendError(returned)) - } - } - } -} - -/// Receive half of the byte-bounded channel. -/// -/// `recv` blocks on the inner mpsc until an item is available, then -/// decrements the byte-accounting and wakes any sender waiting on -/// capacity. -pub struct Receiver<T: HasByteSize> { - rx: MpscReceiver<T>, - acct: Arc<Accounting>, -} - -impl<T: HasByteSize> Receiver<T> { - /// Take the next item. Returns `Err(RecvError)` when all senders - /// have been dropped and the channel is empty. - pub fn recv(&self) -> Result<T, RecvError> { - let item = self.rx.recv()?; - let sz = item.byte_size(); - let mut used = self.acct.used.lock().expect("byte_channel poisoned"); - *used = used.saturating_sub(sz); - // Notify all so multi-sender setups wake every blocked sender, - // not just one. Wasted wakeups are cheap; missed wakeups would - // be a deadlock. - self.acct.cv.notify_all(); - Ok(item) - } -} - -/// Create a byte-bounded channel with the given capacity in bytes. -/// Returns a `(Sender, Receiver)` pair; clone the `Sender` for -/// multi-producer setups. -pub fn channel<T: HasByteSize>(capacity_bytes: usize) -> (Sender<T>, Receiver<T>) { - let (tx, rx) = sync_channel::<T>(INNER_SLOT_CAPACITY); - let acct = Arc::new(Accounting { - used: Mutex::new(0), - cv: Condvar::new(), - capacity: capacity_bytes, - }); - ( - Sender { - tx, - acct: acct.clone(), - }, - Receiver { rx, acct }, - ) -} - -#[cfg(test)] -mod tests { - use super::*; - use std::sync::Arc; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::thread; - use std::time::{Duration, Instant}; - - /// Test payload — its `byte_size` returns whatever we passed at - /// construction so capacity math is exact and predictable. - #[derive(Clone, Debug, PartialEq, Eq)] - struct Item { - sz: usize, - tag: u32, - } - - impl HasByteSize for Item { - fn byte_size(&self) -> usize { - self.sz - } - } - - #[test] - fn send_recv_round_trip() { - let (tx, rx) = channel::<Item>(1024); - for i in 0..5 { - tx.send(Item { sz: 100, tag: i }).unwrap(); - } - for i in 0..5 { - let got = rx.recv().unwrap(); - assert_eq!(got, Item { sz: 100, tag: i }); - } - } - - #[test] - fn byte_accounting_decrements_on_recv() { - // Internal book-keeping check via observable side-effect: after - // sending K items totalling N bytes and receiving them all, a - // subsequent send of an N-byte item must NOT block (no items - // in flight, all capacity refunded). - let (tx, rx) = channel::<Item>(1024); - for _ in 0..4 { - tx.send(Item { sz: 256, tag: 0 }).unwrap(); - } - for _ in 0..4 { - rx.recv().unwrap(); - } - // Cap is now fully available again. Send a 1024-byte item; the - // `used > 0` guard means it goes through alone (no wait). - let start = Instant::now(); - tx.send(Item { sz: 1024, tag: 99 }).unwrap(); - assert!(start.elapsed() < Duration::from_millis(100)); - let got = rx.recv().unwrap(); - assert_eq!(got.tag, 99); - } - - #[test] - fn send_blocks_at_capacity_unblocks_on_recv() { - // Cap = 200 bytes, item = 100 bytes. First two sends fit - // exactly; the third must block until a recv frees capacity. - let (tx, rx) = channel::<Item>(200); - tx.send(Item { sz: 100, tag: 0 }).unwrap(); - tx.send(Item { sz: 100, tag: 1 }).unwrap(); - - let tx2 = tx.clone(); - let sent_at = Arc::new(Mutex::new(None::<Instant>)); - let sent_at2 = sent_at.clone(); - let h = thread::spawn(move || { - tx2.send(Item { sz: 100, tag: 2 }).unwrap(); - *sent_at2.lock().unwrap() = Some(Instant::now()); - }); - - // Give the sender thread a head start; it should be parked in - // `cv.wait` because used (200) + 100 > capacity (200). - thread::sleep(Duration::from_millis(100)); - assert!( - sent_at.lock().unwrap().is_none(), - "third send should be blocked at capacity" - ); - - // Drain one. Sender wakes and completes. - let recv_at = Instant::now(); - let got = rx.recv().unwrap(); - assert_eq!(got.tag, 0); - h.join().unwrap(); - - let sent_when = sent_at.lock().unwrap().unwrap(); - assert!( - sent_when >= recv_at, - "sender must complete AFTER receiver freed capacity" - ); - - // Drain the remaining two. - assert_eq!(rx.recv().unwrap().tag, 1); - assert_eq!(rx.recv().unwrap().tag, 2); - } - - #[test] - fn item_larger_than_capacity_still_goes_through() { - // Pathological case: a single item bigger than the capacity. - // The guard `*used > 0` lets it through when the channel is - // empty (otherwise the sender deadlocks forever). Matches - // `mpsc::SyncSender` semantics for oversize messages. - let (tx, rx) = channel::<Item>(100); - tx.send(Item { sz: 1000, tag: 7 }).unwrap(); - let got = rx.recv().unwrap(); - assert_eq!(got, Item { sz: 1000, tag: 7 }); - } - - #[test] - fn concurrent_send_recv_stress() { - // 4 sender threads × 1k items each, 1 receiver. Verify byte - // accounting stays sane (channel never deadlocks, every item - // arrives exactly once) under contention. - const SENDERS: u32 = 4; - const PER_SENDER: u32 = 1000; - const TOTAL: u32 = SENDERS * PER_SENDER; - - let (tx, rx) = channel::<Item>(8 * 1024); - let sent = Arc::new(AtomicUsize::new(0)); - let mut handles = Vec::new(); - for s in 0..SENDERS { - let tx = tx.clone(); - let sent = sent.clone(); - handles.push(thread::spawn(move || { - for i in 0..PER_SENDER { - // Vary item size so accounting actually has to - // multiplex differently-sized blockers. 1B → 256B. - let sz = 1 + ((i as usize) % 256); - tx.send(Item { - sz, - tag: s * PER_SENDER + i, - }) - .unwrap(); - sent.fetch_add(1, Ordering::SeqCst); - } - })); - } - // Drop our local sender so the receiver can eventually see - // RecvError once all sender clones are done. Cloning the - // sender into each producer means each clone Drop'd separately. - drop(tx); - - let mut received = 0u32; - while let Ok(_item) = rx.recv() { - received += 1; - } - for h in handles { - h.join().unwrap(); - } - assert_eq!(received, TOTAL); - assert_eq!(sent.load(Ordering::SeqCst) as u32, TOTAL); - } - - #[test] - fn send_after_recv_dropped_returns_err() { - let (tx, rx) = channel::<Item>(1024); - drop(rx); - let r = tx.send(Item { sz: 10, tag: 0 }); - assert!(r.is_err()); - } -} diff --git a/src/io/byte_prefetcher.rs b/src/io/byte_prefetcher.rs index bd8aa8e..a139e55 100644 --- a/src/io/byte_prefetcher.rs +++ b/src/io/byte_prefetcher.rs @@ -10,11 +10,12 @@ //! //! This is the byte-stream half of the freemkv mux highway — //! `BytePrefetcher` feeds [`crate::mux::demux_thread::DemuxThread`] -//! for `m2ts://`, `network://`, `stdio://`, and any other stream -//! whose source is an `io::Read` rather than a `SectorSource`. +//! for `m2ts://` (the only in-tree caller today, via +//! [`crate::mux::resolve`]), and works for any stream whose source is +//! an `io::Read` rather than a `SectorSource`. -use crate::halt::Halt; -use crossbeam_channel::{Receiver, Sender, bounded}; +use crate::halt::{Halt, POLL_INTERVAL}; +use crossbeam_channel::{Receiver, RecvTimeoutError, SendTimeoutError, Sender, bounded}; use std::io::Read; use std::thread::JoinHandle; @@ -38,6 +39,13 @@ pub const DEFAULT_CHUNK_BYTES: usize = 16 * 1024 * 1024; /// Returned from [`BytePrefetcher::into_channels`]. Owns the /// producer-thread join handle so dropping the shell joins the /// producer. +/// +/// Drop blocks the calling thread until the producer exits. To +/// guarantee a prompt exit, drop the forward receiver and the recycle +/// sender first so the producer observes channel disconnection (or +/// cancel the [`Halt`] passed to [`BytePrefetcher::new`], which the +/// producer polls at [`POLL_INTERVAL`] granularity even while parked +/// on a channel op). pub struct PrefetchShell { producer: Option<JoinHandle<()>>, } @@ -66,7 +74,13 @@ impl BytePrefetcher { mut reader: R, chunk_bytes: usize, halt: Option<Halt>, - ) -> Self { + ) -> std::io::Result<Self> { + // A zero-length chunk makes every recycled buffer an empty + // slice; `reader.read(&mut [])` returns Ok(0), which the loop + // below treats as EOF — the consumer would see a clean, + // silent zero-byte stream. Callers pass the downstream + // demuxer's batch size, which is always > 0. + debug_assert!(chunk_bytes > 0, "BytePrefetcher chunk_bytes must be > 0"); let (tx, rx) = bounded::<Batch>(FORWARD_DEPTH); let (recycle_tx, recycle_rx) = bounded::<Vec<u8>>(RECYCLE_DEPTH); @@ -80,16 +94,32 @@ impl BytePrefetcher { let producer = std::thread::Builder::new() .name("freemkv-byte-prefetch".into()) .spawn(move || { + let cancelled = || halt.as_ref().map(|h| h.is_cancelled()).unwrap_or(false); loop { - if halt.as_ref().map(|h| h.is_cancelled()).unwrap_or(false) { + if cancelled() { return; } - let mut buf = match recycle_rx.recv() { - Ok(b) => b, - Err(_) => return, // consumer dropped both channels + // Park on the recycle channel, but re-poll halt + // every POLL_INTERVAL: a pure-AtomicBool Halt does + // not disconnect the channel, so a blocking recv() + // would never re-reach the cancel check. + let mut buf = loop { + match recycle_rx.recv_timeout(POLL_INTERVAL) { + Ok(b) => break b, + Err(RecvTimeoutError::Timeout) => { + if cancelled() { + return; + } + } + // Consumer dropped both channels. + Err(RecvTimeoutError::Disconnected) => return, + } }; - // Re-expose the full extent (previous iteration - // may have truncated after a short read). + // Re-expose the full extent. After a short read the + // prior iteration truncated to n < chunk_bytes, so + // this regrows the length back to chunk_bytes + // without reallocating (capacity was fixed at + // construction and never shrinks). if buf.len() < chunk_bytes { buf.resize(chunk_bytes, 0); } else { @@ -109,18 +139,31 @@ impl BytePrefetcher { } }; buf.truncate(n); - if tx.send(Ok(buf)).is_err() { - return; // consumer dropped + // Hand off the filled buffer, re-polling halt on + // each timeout slice so a cancel can interrupt a + // producer parked on a saturated forward channel. + let mut pending = Ok(buf); + loop { + match tx.send_timeout(pending, POLL_INTERVAL) { + Ok(()) => break, + Err(SendTimeoutError::Timeout(returned)) => { + if cancelled() { + return; + } + pending = returned; + } + // Consumer dropped. + Err(SendTimeoutError::Disconnected(_)) => return, + } } } - }) - .expect("freemkv-byte-prefetch thread spawn failed"); + })?; - Self { + Ok(Self { rx, recycle_tx, producer: Some(producer), - } + }) } /// Peel off the channels for zero-copy pipeline consumption. The @@ -128,11 +171,30 @@ impl BytePrefetcher { /// drains `rx`, runs the demuxer in place on each filled buffer, /// and recycles back through `recycle_tx`. pub fn into_channels(self) -> (Receiver<Batch>, Sender<Vec<u8>>, PrefetchShell) { - let mut me = self; - let producer = me.producer.take(); - let rx = me.rx.clone(); - let recycle = me.recycle_tx.clone(); - std::mem::forget(me); + // MOVE the three fields out cleanly — never clone. Each of + // `rx` and `recycle_tx` ends up with exactly ONE live copy: + // the one in the returned tuple. The pre-1.0.0 implementation + // cloned both and then `mem::forget`-ed `self`, leaking the + // originals so an extra live receiver + sender survived + // forever. That defeated the channel-disconnection shutdown: + // when the demux consumer exited early (halt, or a `tx.send` + // error in `demux_thread`), the producer's `recycle_rx.recv()` + // and `tx.send()` never saw all-peers-dropped, so the producer + // never returned and `PrefetchShell::drop`'s `join()` hung. + // + // `ManuallyDrop` + `ptr::read` reads each field out by value + // and suppresses `self`'s own `Drop` (which would otherwise + // double-`join`), leaving NO extra live endpoint behind. This + // is the panic-free equivalent of the `Option::take` approach + // and mirrors `sector::prefetched::into_channels`. + let me = std::mem::ManuallyDrop::new(self); + // SAFETY: `me` is `ManuallyDrop`, so none of these fields will + // be dropped by `me`. Each `ptr::read` performs exactly one + // bitwise move out; every field is read exactly once and never + // touched again, so there are no double-frees and no aliasing. + let producer = unsafe { std::ptr::read(&me.producer) }; + let rx = unsafe { std::ptr::read(&me.rx) }; + let recycle = unsafe { std::ptr::read(&me.recycle_tx) }; (rx, recycle, PrefetchShell { producer }) } } @@ -144,3 +206,72 @@ impl Drop for BytePrefetcher { } } } + +#[cfg(test)] +mod tests { + use super::*; + + /// Endless reader: every `read` fills the whole buffer and never + /// hits EOF, so the producer keeps trying to push batches forward + /// until the forward channel disconnects. Exactly the shape that + /// wedged the pre-1.0.0 `clone + mem::forget` `into_channels`. + struct EndlessReader; + impl Read for EndlessReader { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> { + buf.fill(0); + Ok(buf.len()) + } + } + + /// Run `f` on a helper thread and fail if it does not finish within + /// `secs`. Turns a join-deadlock into a test failure instead of a + /// hung CI run. + fn within<F: FnOnce() + Send + 'static>(secs: u64, f: F) { + let (done_tx, done_rx) = bounded::<()>(1); + std::thread::spawn(move || { + f(); + let _ = done_tx.send(()); + }); + assert!( + done_rx + .recv_timeout(std::time::Duration::from_secs(secs)) + .is_ok(), + "operation did not complete within {secs}s (deadlock)" + ); + } + + /// The CRITICAL regression: after `into_channels`, dropping the + /// returned forward receiver + recycle sender must let the producer + /// observe disconnection and exit, so dropping the `PrefetchShell` + /// (which joins the producer) returns promptly. With the old + /// clone+forget the leaked endpoints kept the producer blocked and + /// this join hung forever. + #[test] + fn into_channels_drop_releases_producer() { + within(10, || { + // Small chunk so the producer cycles quickly and fills the + // forward channel without allocating much. + let pf = BytePrefetcher::new(EndlessReader, 4096, None).expect("spawn"); + let (rx, recycle_tx, shell) = pf.into_channels(); + // Consumer goes away early (halt / abort analogue): drop + // both channel endpoints without draining to EOF. + drop(rx); + drop(recycle_tx); + // Joining the producer must not hang. + drop(shell); + }); + } + + /// Same property via the halt path: cancel the token, then the + /// producer must exit and the shell join must complete. + #[test] + fn halt_releases_producer() { + within(10, || { + let halt = Halt::new(); + let pf = BytePrefetcher::new(EndlessReader, 4096, Some(halt.clone())).expect("spawn"); + let (_rx, _recycle_tx, shell) = pf.into_channels(); + halt.cancel(); + drop(shell); + }); + } +} diff --git a/src/io/file_sector_source/macos.rs b/src/io/file_sector_source/macos.rs index d3519cd..850cacf 100644 --- a/src/io/file_sector_source/macos.rs +++ b/src/io/file_sector_source/macos.rs @@ -8,32 +8,22 @@ use std::fs::File; use std::os::unix::io::AsRawFd; -/// `F_RDADVISE` opcode — not in libc's named constants on all SDKs. -const F_RDADVISE: libc::c_int = 44; - /// Cap on the byte length we pass to `F_RDADVISE`. Asking for a /// multi-GB readahead window is counterproductive — the OS doesn't /// have that much cache to throw at one fd. 64 MiB is generous for -/// our use case (sweep, mux) and matches the byte-channel cap so the -/// kernel's prefetch ≥ our app-level pipeline depth. +/// our use case (sweep, mux) so the kernel's prefetch ≥ our app-level +/// pipeline depth. const RDADVISE_MAX_BYTES: i64 = 64 * 1024 * 1024; -/// `radvisory` per `<sys/fcntl.h>`. repr(C) layout is stable. -#[repr(C)] -struct RadAdvisory { - ra_offset: libc::off_t, - ra_count: libc::c_int, -} - pub(super) fn hint_sequential(file: &File, len_bytes: u64) { let bytes = (len_bytes as i64).min(RDADVISE_MAX_BYTES); - let mut ra = RadAdvisory { + let mut ra = libc::radvisory { ra_offset: 0, ra_count: bytes as libc::c_int, }; // Best-effort. unsafe { - libc::fcntl(file.as_raw_fd(), F_RDADVISE, &mut ra); + libc::fcntl(file.as_raw_fd(), libc::F_RDADVISE, &mut ra); } } @@ -52,11 +42,12 @@ pub(super) fn drop_window(_file: &File, _start: u64, _len: u64) {} /// returns immediately. pub(super) fn prefetch(file: &File, offset: u64, len: u64) { let bytes = (len as i64).min(RDADVISE_MAX_BYTES); - let mut ra = RadAdvisory { + let mut ra = libc::radvisory { ra_offset: offset as libc::off_t, ra_count: bytes as libc::c_int, }; + // Best-effort — kernel hint only. unsafe { - libc::fcntl(file.as_raw_fd(), F_RDADVISE, &mut ra); + libc::fcntl(file.as_raw_fd(), libc::F_RDADVISE, &mut ra); } } diff --git a/src/io/file_sector_source/mod.rs b/src/io/file_sector_source/mod.rs index 3d6624c..d24f174 100644 --- a/src/io/file_sector_source/mod.rs +++ b/src/io/file_sector_source/mod.rs @@ -17,9 +17,17 @@ //! Without page-cache eviction an 85 GB streaming ISO read pins the //! entire file in memory, starves the concurrent writer, and collapses //! mux throughput (observed: 2.7 MB/s mux on 0.21.5 vs. 70 MB/s -//! isolated NFS reads). Every [`READ_DROP_CHUNK_BYTES`] of consumed -//! bytes we call `posix_fadvise(DONTNEED)` over that window, mirroring -//! the write-side [`crate::io::writeback::WritebackPipeline`] policy. +//! isolated NFS reads). Every [`READ_DROP_CHUNK_BYTES_DEFAULT`] of +//! consumed bytes we call `posix_fadvise(DONTNEED)` over that window, +//! mirroring the write-side [`crate::io::writeback::WritebackPipeline`] +//! policy. +//! +//! The drop window is accounted by a monotonic forward byte counter, +//! which matches the sequential streaming pattern the mux highway +//! drives. Under random or backward access the dropped range no longer +//! lines up with the bytes actually read — but `DONTNEED` is purely an +//! advisory cache hint with no correctness impact, so this degrades to +//! a slightly imprecise hint rather than a bug. //! //! ## Platform open hint //! @@ -102,7 +110,10 @@ pub struct FileSectorSource { bytes_read_since_drop: u64, /// File offset at which the current drop window starts. The next /// DONTNEED drops from `drop_window_start` for - /// `bytes_read_since_drop` bytes. + /// `bytes_read_since_drop` bytes. This advances monotonically with + /// the byte count, so it tracks the actual reads only under the + /// forward-sequential access the mux highway uses; under random + /// access it degrades to a harmless, imprecise advisory hint. drop_window_start: u64, /// Cached drop chunk size (resolved from env once at open). drop_chunk_bytes: u64, @@ -116,16 +127,18 @@ impl FileSectorSource { /// /// Issues the platform's "sequential access expected" hint on the /// fd (Linux `posix_fadvise(SEQUENTIAL)`, macOS `fcntl(F_RDADVISE)`, - /// Windows TODO stub) so the kernel's readahead widens. - pub fn open(path: &Path) -> std::io::Result<Self> { - let file = File::open(path)?; - let len = file.metadata()?.len(); + /// Windows no-op) so the kernel's readahead widens. + pub fn open(path: &Path) -> Result<Self> { + let file = File::open(path).map_err(|e| Error::IoError { source: e })?; + let len = file + .metadata() + .map_err(|e| Error::IoError { source: e })? + .len(); let sectors = len / SECTOR_SIZE as u64; if sectors > u32::MAX as u64 { return Err(Error::IsoTooLarge { path: path.to_string_lossy().into_owned(), - } - .into()); + }); } let capacity = sectors as u32; diff --git a/src/io/file_sector_source/windows.rs b/src/io/file_sector_source/windows.rs index 81fbf94..ab0c495 100644 --- a/src/io/file_sector_source/windows.rs +++ b/src/io/file_sector_source/windows.rs @@ -1,20 +1,18 @@ //! Windows: the canonical sequential-access hint is -//! `FILE_FLAG_SEQUENTIAL_SCAN` passed to `CreateFile` at open time — -//! it cannot be set after the fact via `SetFileInformationByHandle`. -//! Routing the open call through this module would mean a custom -//! `File::from_raw_handle` plumb for every `FileSectorSource::open` -//! caller, which is more invasive than the Phase 1 scope. -//! -//! TODO: replumb `FileSectorSource::open` to take an -//! `OpenOptions`-style builder so the Windows path can flip the flag -//! at open time. For now this is a no-op stub. +//! `FILE_FLAG_SEQUENTIAL_SCAN`, which must be passed to `CreateFile` +//! at open time and cannot be set afterward via +//! `SetFileInformationByHandle`. Since `FileSectorSource::open` uses a +//! plain `File::open`, the hints in this module are no-op stubs. use std::fs::File; +/// No-op stub. `FILE_FLAG_SEQUENTIAL_SCAN` can only be set at +/// `CreateFile` open time, which the plain `File::open` path does not +/// do, so there is no post-open hint to issue here. pub(super) fn hint_sequential(_file: &File, _len_bytes: u64) { tracing::debug!( target: "mux", - "FileSectorSource hint_sequential: windows stub (TODO: FILE_FLAG_SEQUENTIAL_SCAN at open)" + "FileSectorSource hint_sequential: windows no-op stub" ); } diff --git a/src/io/mod.rs b/src/io/mod.rs index 8c10a4e..1d4701f 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -18,15 +18,16 @@ //! consumed window so an 85 GB streaming ISO read doesn't fill the //! page cache and starve the concurrent MKV write. //! -//! `Pipeline` + `Sink` (0.18) is the generic producer/consumer primitive +//! `Pipeline` + `Sink` is the generic producer/consumer primitive //! used by sweep, patch, and mux to overlap reads with writes via a //! bounded channel + dedicated consumer thread. //! -//! `byte_channel` is a byte-sized producer/consumer channel for the -//! mux pipeline, sized to absorb worst-case input read stalls. +//! `byte_prefetcher` is the read-ahead producer feeding the mux +//! pipeline for `io::Read`-backed sources: a worker thread fills a +//! recycled pool of buffers and ships them through a channel, exposing +//! `BytePrefetcher` / `PrefetchShell`. pub(crate) mod bounded; -pub mod byte_channel; pub mod byte_prefetcher; pub mod file_sector_source; pub mod sink; @@ -40,8 +41,6 @@ pub mod pipeline; pub(crate) use writeback_file::WritebackFile; -// Re-exports for the 0.18 redesign. Sweep, patch, and mux are all -// wired up (disc/sweep.rs, disc/patch.rs, autorip's ripper/mux.rs). pub use pipeline::{ DEFAULT_PIPELINE_DEPTH, Flow, Pipeline, READ_PIPELINE_DEPTH, Sink, WRITE_PIPELINE_DEPTH, WRITE_THROUGH_DEPTH, diff --git a/src/io/pipeline.rs b/src/io/pipeline.rs index 180c5a4..ef85f1a 100644 --- a/src/io/pipeline.rs +++ b/src/io/pipeline.rs @@ -5,10 +5,9 @@ //! The consumer's behaviour is supplied by a [`Sink`] implementation: //! `apply` is called once per item, `close` is called once at the end. //! -//! Three call sites in libfreemkv want a producer/consumer split — -//! sweep (migrated to `disc/sweep.rs::SweepSink`), patch, and mux. -//! 0.18 collapses all three onto this primitive; sweep is in, -//! patch and mux migrate in later 0.18 slices. +//! Call sites in libfreemkv that want a producer/consumer split — +//! sweep (`disc/sweep.rs::SweepSink`) and the file-backed mux highway +//! — are built on this primitive. //! //! ## Cancellation and error semantics //! @@ -24,7 +23,8 @@ //! blocks on a dead receiver, and the first error is propagated as //! the `JoinHandle` result. //! - Consumer panic is converted into -//! `Error::IoError { source: io::Error::other(...) }`. +//! [`Error::PipelineConsumerPanicked`] (the panic message is logged, +//! not embedded in the error value). //! //! ## Debug logging //! @@ -32,7 +32,6 @@ //! logging throughout the pipeline (channel sends/receives, backpressure, //! consumer lag detection). This is critical for diagnosing stalls. -use std::io; use std::thread::{self, JoinHandle}; use std::time::{Duration, Instant}; @@ -41,15 +40,14 @@ use crossbeam_channel::{Sender, TrySendError, bounded}; use crate::error::Error; use crate::halt::Halt; -/// Deadline for [`Pipeline::finish_with_halt`]'s polling join. Chosen -/// to be comfortably longer than the autorip hard watchdog -/// (`HARD_WATCHDOG_STALL_SECS = 300s`) so the watchdog's `exit(1)` -/// fires first when both are racing on the same wedged consumer. +/// Deadline for [`Pipeline::finish_with_halt`]'s polling join. /// /// 10 minutes is a backstop, not a normal timeout — the consumer is /// expected to drain in seconds. If we hit this, something is wedged -/// inside a kernel call the consumer thread can't unwind from, and the -/// caller has already lost the rip. +/// inside a kernel call the consumer thread can't unwind from. It is +/// deliberately long so a consuming application's own (shorter) stall +/// watchdog gets the first chance to escalate; this join only fires +/// when no such watchdog intervenes. pub const JOIN_TIMEOUT_SECS: u64 = 600; /// Halt-check cadence for the send loop. Producer blocks on @@ -72,11 +70,40 @@ use crate::halt::POLL_INTERVAL; const SEND_HALT_CHECK_INTERVAL: Duration = POLL_INTERVAL; /// Check if verbose debug logging is enabled via FREEMKV_DEBUG env var. +/// +/// The value cannot change mid-run, and this is called multiple times +/// per item on the mux highway hot loop, so the env lookup (a String +/// allocation behind the global env lock) is cached after the first +/// call. pub fn debug_enabled() -> bool { - std::env::var("FREEMKV_DEBUG") - .ok() - .map(|v| v == "1" || v == "true" || v == "yes") - .unwrap_or(false) + static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new(); + *ENABLED.get_or_init(|| { + std::env::var("FREEMKV_DEBUG") + .ok() + .map(|v| v == "1" || v == "true" || v == "yes") + .unwrap_or(false) + }) +} + +/// Turn a consumer-thread panic payload into the numeric +/// [`Error::PipelineConsumerPanicked`] variant. The original panic +/// message (the two stdlib formats `panic!` produces: `&str` / +/// `String`) is logged at the join site for diagnostics — it is NOT +/// baked into the error value, since the library carries no English +/// text in its errors. Callers discriminate on the variant. +fn consumer_panicked(payload: Box<dyn std::any::Any + Send>) -> Error { + let msg = payload + .downcast_ref::<&'static str>() + .copied() + .or_else(|| payload.downcast_ref::<String>().map(|s| s.as_str())) + .unwrap_or("(no message)"); + tracing::error!( + target: "freemkv::pipeline", + phase = "consumer_panicked", + panic_message = msg, + "pipeline consumer thread panicked" + ); + Error::PipelineConsumerPanicked } /// Default channel depth for callers without a specific reason to @@ -104,11 +131,12 @@ pub const WRITE_THROUGH_DEPTH: usize = 1; /// ([`Flow::Continue`]), or stop the pipeline early and run `close()` /// ([`Flow::Stop`]). /// -/// `Stop` has no in-tree caller in this slice — sweep never returns -/// it (it always processes the producer's full work-list before the -/// channel is dropped). Patch and mux are the intended consumers and -/// migrate in later 0.18 slices. The variant ships now so the contract -/// is fixed; the targeted `#[allow]` is removed when patch lands. +/// `Stop` currently has no in-tree caller — sweep never returns it (it +/// always processes the producer's full work-list before the channel +/// is dropped), and the mux highway drains to EOF. The variant is part +/// of the fixed `Sink` contract for early-stop consumers, so the +/// `#[allow(dead_code)]` is intentional and permanent until such a +/// consumer lands. pub enum Flow { Continue, #[allow(dead_code)] @@ -184,47 +212,54 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> { let mut stopped = false; while let Ok(item) = rx.recv() { - if debug_enabled() { + let debug = debug_enabled(); + if debug { tracing::debug!("Pipeline receive: item={}", std::any::type_name::<I>()); } - let apply_start = std::time::Instant::now(); - if first_err.is_some() || stopped { // Drain remaining items so the producer never // blocks on a dead receiver. `apply` is not // called once we've decided to stop. continue; } + + // Only pay for the timestamp when debug tracing is + // on — this runs per item on the mux highway hot + // path. + let apply_start = debug.then(Instant::now); + match sink.apply(item) { Ok(Flow::Continue) => {} Ok(Flow::Stop) => { stopped = true; - if debug_enabled() { + if debug { tracing::debug!("Pipeline: consumer returned Flow::Stop"); } } Err(e) => { - if debug_enabled() { + if debug { tracing::debug!("Pipeline: apply error, stopping, err={:?}", e); } first_err = Some(e); } } - let apply_elapsed = apply_start.elapsed(); - if debug_enabled() && apply_elapsed > std::time::Duration::from_millis(100) { - tracing::debug!( - "Pipeline apply: took {:.2}s, item={}", - apply_elapsed.as_secs_f64(), - std::any::type_name::<I>() - ); - } else if debug_enabled() { - tracing::debug!( - "Pipeline apply: OK in {:.3}ms, item={}", - apply_elapsed.as_micros(), - std::any::type_name::<I>() - ); + if let Some(start) = apply_start { + let apply_elapsed = start.elapsed(); + if apply_elapsed > Duration::from_millis(100) { + tracing::debug!( + "Pipeline apply: took {:.2}s, item={}", + apply_elapsed.as_secs_f64(), + std::any::type_name::<I>() + ); + } else { + tracing::debug!( + "Pipeline apply: OK in {:.3}ms, item={}", + apply_elapsed.as_micros(), + std::any::type_name::<I>() + ); + } } } @@ -250,31 +285,37 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> { /// independent signal (e.g. `Halt`) — `send` alone is not the /// notification edge. pub fn send(&self, item: I) -> Result<(), I> { - let start = std::time::Instant::now(); + // Only timestamp when debug tracing is on — `send` runs per + // item on the mux highway hot path. + let start = debug_enabled().then(Instant::now); match self.tx.send(item) { Ok(()) => { - let elapsed = start.elapsed(); - if debug_enabled() && elapsed > std::time::Duration::from_millis(10) { - tracing::debug!( - "Pipeline send: blocked {:.2}s, item={}", - elapsed.as_secs_f64(), - std::any::type_name::<I>() - ); - } else if debug_enabled() { - tracing::debug!("Pipeline send: OK in {:.3}ms", elapsed.as_micros()); + if let Some(start) = start { + let elapsed = start.elapsed(); + if elapsed > Duration::from_millis(10) { + tracing::debug!( + "Pipeline send: blocked {:.2}s, item={}", + elapsed.as_secs_f64(), + std::any::type_name::<I>() + ); + } else { + tracing::debug!("Pipeline send: OK in {:.3}ms", elapsed.as_micros()); + } } Ok(()) } Err(e) => { - let elapsed = start.elapsed(); - if debug_enabled() && elapsed > std::time::Duration::from_millis(10) { - tracing::debug!( - "Pipeline send: blocked {:.2}s before channel closed, item={}", - elapsed.as_secs_f64(), - std::any::type_name::<I>() - ); - } else if debug_enabled() { - tracing::debug!("Pipeline send: failed after {:.3}ms", elapsed.as_micros()); + if let Some(start) = start { + let elapsed = start.elapsed(); + if elapsed > Duration::from_millis(10) { + tracing::debug!( + "Pipeline send: blocked {:.2}s before channel closed, item={}", + elapsed.as_secs_f64(), + std::any::type_name::<I>() + ); + } else { + tracing::debug!("Pipeline send: failed after {:.3}ms", elapsed.as_micros()); + } } Err(e.0) } @@ -367,9 +408,9 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> { /// Drop the producer-side channel and wait for the consumer /// thread to finish. Returns whatever the consumer's `close()` /// produced, or the first `apply` error, or — on consumer panic — - /// an `Error::IoError` whose source is `io::Error::other(...)` - /// with a "pipeline consumer panicked: <payload>" message - /// (callers can match on the constant prefix). + /// [`Error::PipelineConsumerPanicked`]. The panic payload is + /// logged at the join site (the library carries no English in its + /// error values), so callers discriminate on the variant. pub fn finish(self) -> Result<R, Error> { let Pipeline { tx, handle } = self; // Explicit drop, although the destructure already drops `tx` @@ -377,20 +418,7 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> { drop(tx); match handle.join() { Ok(result) => result, - Err(payload) => { - // Preserve the original panic message when the - // consumer's panic payload was a `&str` or `String` - // (the two stdlib formats that `panic!` produces). - // Anything else falls back to "(no message)". - let msg = payload - .downcast_ref::<&'static str>() - .copied() - .or_else(|| payload.downcast_ref::<String>().map(|s| s.as_str())) - .unwrap_or("(no message)"); - Err(Error::IoError { - source: io::Error::other(format!("pipeline consumer panicked: {msg}")), - }) - } + Err(payload) => Err(consumer_panicked(payload)), } } @@ -402,19 +430,18 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> { /// [`JOIN_TIMEOUT_SECS`] deadline. Returns: /// /// - `Ok(R)` on a clean consumer exit. - /// - `Err(Error::IoError)` with one of three message prefixes for - /// wedge cases: - /// - `"pipeline join halted"` — halt fired while waiting. - /// - `"pipeline join timed out"` — `JOIN_TIMEOUT_SECS` elapsed. - /// - `"pipeline consumer panicked"` — same as `finish()`. + /// - One of three numeric error variants for the wedge cases: + /// - [`Error::Halted`] — halt fired while waiting. + /// - [`Error::PipelineJoinTimeout`] — `JOIN_TIMEOUT_SECS` elapsed. + /// - [`Error::PipelineConsumerPanicked`] — same as `finish()`. /// /// In the `halted` and `timed out` branches the consumer thread is /// intentionally leaked — exactly the same trade-off the /// `bounded_syscall` primitive makes. The wedged kernel call /// inside the consumer will unwind whenever it does, or at /// process exit. The caller is free to fall back to a degraded - /// path (in autorip's case: `exit(1)` after the hard watchdog - /// escalation, letting Docker restart the container). + /// path (e.g. abort the session and let a supervisor restart the + /// process). /// /// Plain [`Pipeline::finish`] is preserved for callers without a /// halt-token plumbed through; that path still blocks indefinitely @@ -427,31 +454,18 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> { if handle.is_finished() { return match handle.join() { Ok(result) => result, - Err(payload) => { - let msg = payload - .downcast_ref::<&'static str>() - .copied() - .or_else(|| payload.downcast_ref::<String>().map(|s| s.as_str())) - .unwrap_or("(no message)"); - Err(Error::IoError { - source: io::Error::other(format!("pipeline consumer panicked: {msg}")), - }) - } + Err(payload) => Err(consumer_panicked(payload)), }; } if let Some(h) = halt { if h.is_cancelled() { // Consumer thread is intentionally leaked. - return Err(Error::IoError { - source: io::Error::other("pipeline join halted"), - }); + return Err(Error::Halted); } } if Instant::now() >= deadline { // Consumer thread is intentionally leaked. - return Err(Error::IoError { - source: io::Error::other("pipeline join timed out"), - }); + return Err(Error::PipelineJoinTimeout); } thread::sleep(POLL_INTERVAL); } @@ -700,25 +714,13 @@ mod tests { std::panic::set_hook(prev); - match res { - Err(Error::IoError { source }) => { - let msg = source.to_string(); - // Constant prefix lets callers match without parsing - // the variable payload tail. - assert!( - msg.contains("pipeline consumer panicked"), - "expected constant panic prefix, got: {msg}" - ); - // The original `panic!` payload (a `&'static str`) must - // be preserved — without the downcast the message - // would just be the prefix. - assert!( - msg.contains("synthetic test panic"), - "expected original panic payload, got: {msg}" - ); - } - other => panic!("expected Err(IoError), got {other:?}"), - } + // A consumer panic surfaces as the numeric variant, not an + // English-carrying io::Error. The original panic payload is + // logged at the join site, not embedded in the error value. + assert!( + matches!(res, Err(Error::PipelineConsumerPanicked)), + "expected Err(PipelineConsumerPanicked), got {res:?}" + ); } /// Never-completing sink — `apply` blocks until cancelled. Signals @@ -810,8 +812,8 @@ mod tests { #[test] fn send_with_halt_returns_item_on_halt() { // Same setup, but the halt fires before the deadline elapses. - // The send loop must observe the halt within ~50 ms (the - // SEND_POLL_INTERVAL) and return the item. + // The send loop must observe the halt within ~250 ms (the + // SEND_HALT_CHECK_INTERVAL) and return the item. let cancel = Arc::new(std::sync::atomic::AtomicBool::new(false)); let started = Arc::new(std::sync::atomic::AtomicBool::new(false)); let pipe = Pipeline::spawn( @@ -850,8 +852,7 @@ mod tests { #[test] fn finish_with_halt_returns_halted_when_consumer_wedged() { // Consumer wedges on the first apply; halt fires; finish - // returns the documented "pipeline join halted" error rather - // than blocking forever. + // returns Error::Halted rather than blocking forever. let cancel = Arc::new(std::sync::atomic::AtomicBool::new(false)); let started = Arc::new(std::sync::atomic::AtomicBool::new(false)); let pipe = Pipeline::spawn( @@ -879,15 +880,10 @@ mod tests { // Release the leaked consumer so the test process exits cleanly. cancel.store(true, Ordering::SeqCst); - match res { - Err(Error::IoError { source }) => { - assert!( - source.to_string().contains("pipeline join halted"), - "expected halt-prefix error, got: {source}" - ); - } - other => panic!("expected Err(IoError) halted, got {other:?}"), - } + assert!( + matches!(res, Err(Error::Halted)), + "expected Err(Halted), got {res:?}" + ); // Bailed out within ~1 second of the halt firing (worst case // one POLL_INTERVAL = 250 ms of slack). assert!( diff --git a/src/io/sink/local_file.rs b/src/io/sink/local_file.rs index 5eef032..58c2fc0 100644 --- a/src/io/sink/local_file.rs +++ b/src/io/sink/local_file.rs @@ -12,14 +12,18 @@ //! size patch, Cues index, segment header backpatch) to land on the //! right offset. //! -//! `RandomAccessSink` is satisfied via the blanket impl in -//! [`super::mod`]; no explicit impl needed here. +//! [`SequentialSink`](super::SequentialSink) is implemented explicitly +//! (not via a blanket impl) so its `finish()` flushes the `BufWriter` +//! and `fsync`s the file even when called through a `dyn` trait object; +//! [`RandomAccessSink`](super::RandomAccessSink) is implemented over the +//! `Seek` impl below. use std::fs::{File, OpenOptions}; use std::io::{self, BufWriter, Seek, SeekFrom, Write}; use std::path::Path; use super::preallocate; +use super::{RandomAccessSink, SequentialSink}; const BUFFER_BYTES: usize = 4 * 1024 * 1024; @@ -80,13 +84,26 @@ impl LocalFileSink { /// Drain the internal buffer and `fsync` the underlying file. /// Idempotent with `Drop` (the `BufWriter` also flushes on drop; /// this call additionally surfaces fsync errors to the caller). - #[allow(dead_code)] // exposed for parity with WritebackFile::sync_all + /// [`SequentialSink::finish`](super::SequentialSink::finish) + /// delegates here so the durable flush happens through a trait + /// object too. pub fn sync_all(&mut self) -> io::Result<()> { self.inner.flush()?; self.inner.get_ref().sync_all() } } +impl SequentialSink for LocalFileSink { + /// Flush the 4 MiB `BufWriter` and `fsync` the file. Overriding the + /// trait default is what makes a `dyn SequentialSink` / `dyn + /// RandomAccessSink` `finish()` actually durable instead of a no-op. + fn finish(&mut self) -> io::Result<()> { + self.sync_all() + } +} + +impl RandomAccessSink for LocalFileSink {} + impl Write for LocalFileSink { fn write(&mut self, buf: &[u8]) -> io::Result<usize> { self.inner.write(buf) diff --git a/src/io/sink/mod.rs b/src/io/sink/mod.rs index 8299c9c..828533e 100644 --- a/src/io/sink/mod.rs +++ b/src/io/sink/mod.rs @@ -35,14 +35,22 @@ pub use socket::{SocketSink, UdpSocketSink}; /// trait does not impose or hide any buffering of its own. /// /// `finish` drains any internal buffering and signals end-of-stream to -/// the underlying transport (close-write on a socket, flush on a -/// buffered writer, etc.). The default impl is a no-op; concrete -/// implementations that need explicit shutdown can override it but the -/// blanket impl below keeps it optional for adapter types like -/// `&mut File`. +/// the underlying transport (close-write on a socket, flush + fsync on +/// a buffered file, etc.). The default impl flushes via [`Write::flush`] +/// — correct for an unbuffered destination — but every concrete sink in +/// this module overrides it to drain its own buffer and run its +/// transport-specific finalisation (socket `shutdown(Write)`, file +/// `fsync`). There is deliberately NO blanket `impl SequentialSink for +/// T`: a blanket impl would force the no-op-style default on every +/// concrete sink (a blanket impl cannot be overridden per-type without a +/// coherence conflict), so a `Box<dyn SequentialSink>` / `&mut dyn +/// SequentialSink` `finish()` call would silently skip the flush and +/// transport shutdown. With explicit per-type impls the vtable dispatches +/// `finish` to the real implementation, so flush + durable-finish +/// actually happen through a trait object. pub trait SequentialSink: Write + Send { fn finish(&mut self) -> std::io::Result<()> { - Ok(()) + self.flush() } } @@ -51,15 +59,6 @@ pub trait SequentialSink: Write + Send { /// random-access sink is always usable as a sequential sink. pub trait RandomAccessSink: SequentialSink + Seek {} -// Blanket impls so any `Write + Send` type acts as a `SequentialSink` -// (with default `finish`), and any sink that also impls `Seek` is -// automatically a `RandomAccessSink`. Keeps call-site ergonomics simple -// — `&mut File`, `LocalFileSink`, `WritebackFile`, `BufWriter<File>`, -// and `Cursor<Vec<u8>>` all satisfy the right trait without per-type -// boilerplate. -impl<T: Write + Send + ?Sized> SequentialSink for T {} -impl<T: SequentialSink + Seek + ?Sized> RandomAccessSink for T {} - /// Pick the right `RandomAccessSink` impl for `dest` based on its /// filesystem type. /// @@ -82,13 +81,9 @@ pub fn open_for_mkv( dest: &std::path::Path, size_hint: Option<u64>, ) -> std::io::Result<Box<dyn RandomAccessSink>> { - #[cfg(not(target_os = "linux"))] - use crate::platform::fs_type::detect; - #[cfg(target_os = "linux")] - use crate::platform::fs_type::{FsType, detect}; - #[cfg(target_os = "linux")] { + use crate::platform::fs_type::{FsType, detect}; if detect(dest) == FsType::Nfs { let wf = match size_hint { Some(n) => crate::io::WritebackFile::create_with_size_hint(dest, n)?, @@ -97,12 +92,13 @@ pub fn open_for_mkv( return Ok(Box::new(wf)); } } - // Silence the unused-binding warning on non-Linux where the only - // branch above is cfg-gated out. + // Only Linux differentiates the sink by filesystem type (NFS gets + // the WritebackFile machinery); every other OS always uses + // `LocalFileSink`. Reference `detect` as a value (no call, no + // `statfs` syscall) so it isn't flagged dead on non-Linux while + // still avoiding the wasted probe whose result we'd discard. #[cfg(not(target_os = "linux"))] - { - let _ = detect(dest); - } + let _ = crate::platform::fs_type::detect; let sink = match size_hint { Some(n) => LocalFileSink::with_size_hint(dest, n)?, @@ -114,32 +110,25 @@ pub fn open_for_mkv( #[cfg(test)] mod tests { use super::*; - use std::fs::File; - // Type-level assertion: the blanket impls cover the shapes we care - // about. These functions never run; they just have to type-check. - fn _assert_file_is_sequential(_: &mut dyn SequentialSink) {} - fn _assert_file_is_random_access(_: &mut dyn RandomAccessSink) {} + // Type-level assertion: the concrete sinks satisfy the trait + // objects. These functions never run; they just have to type-check. + fn _assert_is_sequential(_: &mut dyn SequentialSink) {} + fn _assert_is_random_access(_: &mut dyn RandomAccessSink) {} #[test] - fn blanket_impls_cover_file_and_localfilesink() { - // `File` directly via blanket impls. + fn concrete_sinks_satisfy_traits() { let dir = tempfile::tempdir().unwrap(); - let mut f = File::create(dir.path().join("a.bin")).unwrap(); - _assert_file_is_sequential(&mut f); - _assert_file_is_random_access(&mut f); - // `LocalFileSink` ditto. + // `LocalFileSink` is a random-access (and thus sequential) sink. let mut s = LocalFileSink::create(&dir.path().join("b.bin")).unwrap(); - _assert_file_is_sequential(&mut s); - _assert_file_is_random_access(&mut s); + _assert_is_sequential(&mut s); + _assert_is_random_access(&mut s); - // `WritebackFile` — confirms the Phase 1 type still satisfies - // the trait via the blanket impl without needing an explicit - // `impl RandomAccessSink for WritebackFile {}`. + // `WritebackFile` ditto, via its explicit per-type impls. let mut wf = crate::io::WritebackFile::create(&dir.path().join("c.bin")).unwrap(); - _assert_file_is_sequential(&mut wf); - _assert_file_is_random_access(&mut wf); + _assert_is_sequential(&mut wf); + _assert_is_random_access(&mut wf); } #[test] @@ -155,4 +144,24 @@ mod tests { let bytes = std::fs::read(&p).unwrap(); assert_eq!(&bytes[..5], b"hello"); } + + /// finish() through a `dyn SequentialSink` trait object must + /// dispatch to the concrete sink's override (flush + fsync), not a + /// no-op default. This is the regression test for the silent-no-op + /// finish() bug. + #[test] + fn finish_through_trait_object_flushes_local_file() { + let dir = tempfile::tempdir().unwrap(); + let p = dir.path().join("trait-finish.bin"); + let sink = LocalFileSink::create(&p).unwrap(); + // Box as the trait object the production path uses. + let mut boxed: Box<dyn SequentialSink> = Box::new(sink); + boxed.write_all(b"buffered-tail").unwrap(); + // finish() through the vtable must drain the 4 MiB BufWriter and + // fsync; the bytes must be visible to a separate reader BEFORE + // we drop the sink (drop-flush must not be what saves us). + boxed.finish().unwrap(); + let bytes = std::fs::read(&p).unwrap(); + assert_eq!(&bytes[..], b"buffered-tail"); + } } diff --git a/src/io/sink/preallocate/linux.rs b/src/io/sink/preallocate/linux.rs index 5995020..8d423be 100644 --- a/src/io/sink/preallocate/linux.rs +++ b/src/io/sink/preallocate/linux.rs @@ -5,13 +5,16 @@ //! the file naturally. use std::fs::File; -#[cfg(unix)] use std::os::unix::io::AsRawFd; pub(super) fn preallocate_impl(file: &File, size_bytes: u64) { let fd = file.as_raw_fd(); + // Clamp to the signed `off_t` range fallocate expects; an unchecked + // `as i64` cast would wrap a >= 2^63 size to a negative length that + // fallocate rejects with EINVAL (silent no-op). + let len = i64::try_from(size_bytes).unwrap_or(i64::MAX); // FALLOC_FL_KEEP_SIZE = 0x01. - let rc = unsafe { libc::fallocate(fd, libc::FALLOC_FL_KEEP_SIZE, 0, size_bytes as i64) }; + let rc = unsafe { libc::fallocate(fd, libc::FALLOC_FL_KEEP_SIZE, 0, len) }; tracing::debug!( target: "mux", "LocalFileSink fallocate size_hint={size_bytes} rc={rc} ok={}", diff --git a/src/io/sink/preallocate/macos.rs b/src/io/sink/preallocate/macos.rs index f4c8360..ecfdd4c 100644 --- a/src/io/sink/preallocate/macos.rs +++ b/src/io/sink/preallocate/macos.rs @@ -1,7 +1,8 @@ //! macOS `F_PREALLOCATE` extent reservation. //! -//! `fcntl(F_PREALLOCATE)` with `F_ALLOCATECONTIG` first (try for a -//! contiguous run) and fall back to `F_ALLOCATEALL` (non-contig OK). +//! `fcntl(F_PREALLOCATE)` with `F_ALLOCATECONTIG | F_ALLOCATEALL` first +//! (prefer a contiguous run but accept scattered extents to satisfy the +//! full length) and fall back to `F_ALLOCATEALL` alone on failure. //! Reported file size is unchanged — the muxer's writes still grow it. use std::fs::File; @@ -13,16 +14,24 @@ use crate::io::platform_macos::{ pub(super) fn preallocate_impl(file: &File, size_bytes: u64) { let fd = file.as_raw_fd(); + // Clamp to the signed `off_t` range; an unchecked `as off_t` cast + // would wrap a >= 2^63 size to a negative length. + let len = i64::try_from(size_bytes).unwrap_or(i64::MAX) as libc::off_t; let mut store = Fstore { - fst_flags: F_ALLOCATECONTIG, + // Prefer a contiguous run but accept scattered extents to + // satisfy the full length. Without F_ALLOCATEALL the first + // attempt is best-effort and can return rc=0 with a partial + // allocation, so the fallback below would never fire. Matches + // writeback_file/macos.rs. + fst_flags: F_ALLOCATECONTIG | F_ALLOCATEALL, fst_posmode: F_PEOFPOSMODE, fst_offset: 0, - fst_length: size_bytes as libc::off_t, + fst_length: len, fst_bytesalloc: 0, }; let mut rc = unsafe { libc::fcntl(fd, F_PREALLOCATE, &mut store as *mut Fstore) }; if rc == -1 { - // Fall back to non-contiguous. + // Fall back to non-contiguous only. store.fst_flags = F_ALLOCATEALL; rc = unsafe { libc::fcntl(fd, F_PREALLOCATE, &mut store as *mut Fstore) }; } diff --git a/src/io/sink/socket.rs b/src/io/sink/socket.rs index 312a6bf..fbaa969 100644 --- a/src/io/sink/socket.rs +++ b/src/io/sink/socket.rs @@ -11,16 +11,21 @@ //! conventional choice). `finish()` is a no-op; UDP has no end-of-stream //! marker. //! -//! Both types satisfy [`SequentialSink`] via the blanket impl in -//! `super::mod`. Neither implements `Seek`, so neither satisfies -//! [`RandomAccessSink`] — using one with `MkvMux` is a compile error, -//! which is the design intent. +//! Both types implement [`SequentialSink`] explicitly so their +//! `finish()` dispatches correctly through a `dyn SequentialSink` trait +//! object (the `SocketSink` override drains the buffer and +//! `shutdown(Write)`s; the `UdpSocketSink` override flushes only). +//! Neither implements `Seek`, so neither satisfies [`RandomAccessSink`] +//! — using one with `MkvMux` is a compile error, which is the design +//! intent. //! //! [`SequentialSink`]: super::SequentialSink //! [`RandomAccessSink`]: super::RandomAccessSink use std::io::{self, BufWriter, Write}; -use std::net::{Shutdown, TcpStream, ToSocketAddrs, UdpSocket}; +use std::net::{Shutdown, SocketAddr, TcpStream, ToSocketAddrs, UdpSocket}; + +use super::SequentialSink; /// `BufWriter` capacity for [`SocketSink`]. 1 MiB matches the typical /// kernel send-buffer ceiling and keeps small-write amplification from @@ -54,8 +59,10 @@ impl SocketSink { // `set_nodelay(true)` keeps small writes (TS packet trains, fMP4 // moof headers) from sitting in Nagle's algorithm until the buffer // fills. The BufWriter already absorbs syscall overhead; Nagle - // would just add latency without coalescing more. - stream.set_nodelay(true)?; + // would just add latency without coalescing more. It is a latency + // hint, not a correctness requirement, so a platform that rejects + // TCP_NODELAY must not fail the connect — demote the error. + let _ = stream.set_nodelay(true); if let Some(n) = sndbuf_bytes { set_send_buffer(&stream, n)?; } @@ -76,17 +83,12 @@ impl Write for SocketSink { } } -impl SocketSink { +impl SequentialSink for SocketSink { /// Drain the BufWriter and `shutdown(Write)` the underlying socket - /// so the peer sees a clean EOF. - /// - /// Note: [`SequentialSink::finish`](super::SequentialSink::finish)'s - /// blanket-impl default is a no-op. Trait-object call sites that - /// need socket shutdown should call this inherent method directly - /// before dropping the sink, or hold the concrete `SocketSink` type - /// (typical pattern: each muxer's `finish()` calls the appropriate - /// inherent close method on its captured concrete sink). - pub fn finish(&mut self) -> io::Result<()> { + /// so the peer sees a clean EOF. Overriding the trait default is + /// what makes a `dyn SequentialSink` `finish()` send the buffered + /// tail and the EOF instead of silently dropping them. + fn finish(&mut self) -> io::Result<()> { self.buf.flush()?; // `shutdown(Write)` signals clean EOF to the peer. Errors here // are non-fatal — the connection may have already been torn down @@ -117,10 +119,22 @@ impl UdpSocketSink { /// /// `sndbuf_bytes`, when present, is a hint to `SO_SNDBUF`. pub fn connect<A: ToSocketAddrs>(peer: A, sndbuf_bytes: Option<usize>) -> io::Result<Self> { - // Bind to all-zeros / any port. The kernel picks an ephemeral - // source port and the source IP at first send. - let socket = UdpSocket::bind("0.0.0.0:0")?; - socket.connect(peer)?; + // Resolve the peer first so the local bind matches its address + // family. Binding `0.0.0.0:0` (IPv4) and then connecting to an + // IPv6 peer fails with EAFNOSUPPORT, so pick the wildcard that + // matches the resolved family. + let peer_addr = peer + .to_socket_addrs()? + .next() + .ok_or_else(|| io::Error::from(io::ErrorKind::AddrNotAvailable))?; + let bind_addr = match peer_addr { + SocketAddr::V4(_) => "0.0.0.0:0", + SocketAddr::V6(_) => "[::]:0", + }; + // Bind to the matching wildcard / any port. The kernel picks an + // ephemeral source port and the source IP at first send. + let socket = UdpSocket::bind(bind_addr)?; + socket.connect(peer_addr)?; if let Some(n) = sndbuf_bytes { set_udp_send_buffer(&socket, n)?; } @@ -140,10 +154,12 @@ impl Write for UdpSocketSink { } } -impl UdpSocketSink { - /// No-op — UDP has no end-of-stream marker. Provided for parity - /// with [`SocketSink::finish`] so call sites can treat them uniformly. - pub fn finish(&mut self) -> io::Result<()> { +impl SequentialSink for UdpSocketSink { + /// UDP has no end-of-stream marker, so there is nothing to shut + /// down; `write` already sent each datagram unbuffered. Flush is a + /// no-op but kept explicit so the trait-object `finish()` matches + /// the concrete behaviour. + fn finish(&mut self) -> io::Result<()> { Ok(()) } } diff --git a/src/io/writeback/linux.rs b/src/io/writeback/linux.rs index 2aee7f0..3431f49 100644 --- a/src/io/writeback/linux.rs +++ b/src/io/writeback/linux.rs @@ -54,7 +54,6 @@ use std::collections::VecDeque; use std::fs::File; use std::os::unix::io::{AsRawFd, RawFd}; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::{Duration, Instant}; @@ -94,13 +93,11 @@ pub(crate) struct WritebackPipeline { is_nfs: bool, /// Set the first time WAIT_AFTER exceeds [`WAIT_AFTER_TIMEOUT`]. /// Once set, behaviour matches the NFS path for the rest of the - /// pipeline's life. Wrapped in `Arc` only because both this - /// struct and the spawned worker thread (which itself doesn't - /// touch the flag) share-via-fd patterns might one day need it; - /// today it's effectively a single-owner cell — the `Arc` shape - /// keeps the door open for moving the read side into a worker - /// without re-plumbing types. - degraded: Arc<AtomicBool>, + /// pipeline's life. A plain `AtomicBool`: the flag is only ever + /// touched on the owning thread (the spawned WAIT_AFTER worker never + /// reads or writes it). `AtomicBool` over `bool` only because the + /// load/store sites read cleanly; no sharing is needed today. + degraded: AtomicBool, } impl WritebackPipeline { @@ -124,7 +121,7 @@ impl WritebackPipeline { wait_after_window: VecDeque::with_capacity(ADAPTIVE_WINDOW), chunk_count: 0, is_nfs, - degraded: Arc::new(AtomicBool::new(false)), + degraded: AtomicBool::new(false), } } @@ -143,8 +140,12 @@ impl WritebackPipeline { if pos < self.last_flush_pos.saturating_add(self.chunk_bytes) { return; } - let chunk_off = self.last_flush_pos as i64; - let chunk_len = (pos - self.last_flush_pos) as i64; + // Byte offsets are unsigned throughout; the signed cast happens + // only at the libc call boundary where the kernel ABI requires + // `i64`. `saturating_sub` documents and hardens the line-above + // guard that `pos >= last_flush_pos`. + let chunk_off: u64 = self.last_flush_pos; + let chunk_len: u64 = pos.saturating_sub(self.last_flush_pos); let mut wait_ms: u64 = 0; let mut fadvise_ms: u64 = 0; // Async kickoff for the just-completed chunk runs on every @@ -152,7 +153,12 @@ impl WritebackPipeline { // by spec and gives the kernel an early hint that this range // is ready to flush. unsafe { - libc::sync_file_range(self.fd, chunk_off, chunk_len, libc::SYNC_FILE_RANGE_WRITE); + libc::sync_file_range( + self.fd, + chunk_off as i64, + chunk_len as i64, + libc::SYNC_FILE_RANGE_WRITE, + ); } if let Some((prev_off, prev_len)) = self.pending.take() { if self.skip_wait() { @@ -198,12 +204,12 @@ impl WritebackPipeline { } } } - self.pending = Some((chunk_off as u64, chunk_len as u64)); + self.pending = Some((chunk_off, chunk_len)); self.last_flush_pos = pos; self.chunk_count += 1; tracing::trace!( target: "mux", - "WritebackPipeline chunk off={} len={} sync_file_range_ms={wait_ms} fadvise_ms={fadvise_ms} chunk_bytes={} skip_wait={}", + "WritebackPipeline chunk off={} len={} wait_after_ms={wait_ms} fadvise_ms={fadvise_ms} chunk_bytes={} skip_wait={}", chunk_off, chunk_len, self.chunk_bytes, @@ -231,10 +237,14 @@ impl WritebackPipeline { if self.wait_after_window.len() < ADAPTIVE_WINDOW { return; } - // p95 of 16 samples ≈ sorted[14] (5 % of 16 = 0.8 ≈ 1 above). + // p95 index, derived from the window size so it stays valid if + // ADAPTIVE_WINDOW changes (a hard-coded `[14]` would panic OOB + // for a window <= 14). For the default 16 this is index 15 + // (ceil(16 * 95 / 100) - 1 = 15), i.e. the top sample. let mut sorted: Vec<u64> = self.wait_after_window.iter().copied().collect(); sorted.sort_unstable(); - let p95 = sorted[14]; + let p95_idx = (ADAPTIVE_WINDOW * 95).div_ceil(100).min(ADAPTIVE_WINDOW) - 1; + let p95 = sorted[p95_idx]; let old = self.chunk_bytes; let new = if p95 > ADAPTIVE_GROW_MS && self.chunk_bytes < CHUNK_BYTES_MAX { (self.chunk_bytes * 2).min(CHUNK_BYTES_MAX) @@ -321,11 +331,14 @@ fn detect_nfs(fd: RawFd) -> bool { /// Run `sync_file_range(WAIT_AFTER)` on a worker thread and wait up /// to [`WAIT_AFTER_TIMEOUT`] for it to return. `Some(elapsed_ms)` on /// success; `None` on timeout. On timeout the worker thread is -/// 0.20.6 generalizes the worker-thread + recv_timeout pattern into -/// [`crate::io::bounded::bounded_syscall`]; this helper now just adapts -/// the generic primitive to the WAIT_AFTER call shape (returns elapsed_ms -/// instead of the syscall's `()` return, treats `WorkerLost` as a benign -/// no-op to match the original semantics). +/// intentionally leaked — it unwinds whenever the syscall eventually +/// returns or the process exits. +/// +/// This delegates to [`crate::io::bounded::bounded_syscall`], the +/// generic worker-thread + `recv_timeout` primitive, and just adapts it +/// to the WAIT_AFTER call shape: it returns `elapsed_ms` instead of the +/// syscall's `()`, and treats `WorkerLost` as a benign no-op to match +/// the original semantics. fn wait_after_with_timeout(fd: RawFd, off: u64, len: u64) -> Option<u64> { let started = Instant::now(); match crate::io::bounded::bounded_syscall(None, WAIT_AFTER_TIMEOUT, move || unsafe { diff --git a/src/io/writeback_file/linux.rs b/src/io/writeback_file/linux.rs index f7204a1..3747a9b 100644 --- a/src/io/writeback_file/linux.rs +++ b/src/io/writeback_file/linux.rs @@ -19,14 +19,10 @@ use std::time::Duration; pub(super) fn preallocate(file: &File, size_bytes: u64) { // FALLOC_FL_KEEP_SIZE = 0x01 — keep the reported file size at 0 // (writes grow it normally) while still pre-reserving the extents. - let rc = unsafe { - libc::fallocate( - file.as_raw_fd(), - libc::FALLOC_FL_KEEP_SIZE, - 0, - size_bytes as i64, - ) - }; + // Clamp to the signed `off_t` range; an unchecked `as i64` cast + // would wrap a >= 2^63 size to a negative length (EINVAL no-op). + let len = i64::try_from(size_bytes).unwrap_or(i64::MAX); + let rc = unsafe { libc::fallocate(file.as_raw_fd(), libc::FALLOC_FL_KEEP_SIZE, 0, len) }; tracing::debug!( target: "mux", "WritebackFile fallocate size_hint={size_bytes} rc={rc} ok={}", @@ -34,10 +30,14 @@ pub(super) fn preallocate(file: &File, size_bytes: u64) { ); } -/// Run `fsync` on `file` with a 60 s deadline. On timeout we log loudly -/// and return `Ok(())` — the kernel will still flush on close, so the -/// data is best-effort durable; the alternative (trap the thread for -/// the rest of the rip) defeats `/api/stop`. +/// Run `fsync` on `file` with a 60 s deadline. On timeout — and +/// likewise on halt or a lost worker — we log and return `Ok(())`: the +/// kernel will still flush on close, so the data is best-effort durable. +/// The alternative (trap the thread for the rest of the rip, or return +/// an error that aborts an otherwise-complete mux) is worse, so all +/// three fallbacks return `Ok(())`. `Ok(())` from these paths is NOT a +/// durability barrier — the durable flush did not complete; only the +/// hang is bounded. pub(super) fn durable_sync(file: &File) -> io::Result<()> { let fd = file.as_raw_fd(); match crate::io::bounded::bounded_syscall( @@ -60,7 +60,19 @@ pub(super) fn durable_sync(file: &File) -> io::Result<()> { ); Ok(()) } - Err(crate::io::bounded::BoundedError::Halted) => Ok(()), - Err(crate::io::bounded::BoundedError::WorkerLost) => Ok(()), + Err(crate::io::bounded::BoundedError::Halted) => { + tracing::warn!( + target: "mux", + "WritebackFile::sync_all fsync skipped (halt requested); data not durably flushed, kernel will flush on close" + ); + Ok(()) + } + Err(crate::io::bounded::BoundedError::WorkerLost) => { + tracing::error!( + target: "mux", + "WritebackFile::sync_all fsync worker lost before completion; data not durably flushed, kernel will flush on close" + ); + Ok(()) + } } } diff --git a/src/io/writeback_file/macos.rs b/src/io/writeback_file/macos.rs index ed9aefb..8c4f2bd 100644 --- a/src/io/writeback_file/macos.rs +++ b/src/io/writeback_file/macos.rs @@ -1,10 +1,12 @@ //! macOS platform impl for [`super::WritebackFile`]. //! //! - `preallocate`: `fcntl(F_PREALLOCATE)` — macOS's fallocate-equiv. -//! Reserves a contiguous extent when possible, falling back to a -//! non-contiguous reservation if the FS can't satisfy it. Reported -//! file size is unchanged (`F_ALLOCATEALL` is not set, so allocation -//! is "best effort up to length"; growth happens via writes). +//! First attempt requests `F_ALLOCATECONTIG | F_ALLOCATEALL` (prefer a +//! contiguous run but accept scattered extents to satisfy the full +//! length), falling back to `F_ALLOCATEALL` alone on failure. +//! `F_PREALLOCATE` never advances EOF regardless of the flags — only +//! `ftruncate`/writes grow the file — so the reported file size is +//! unchanged; `F_ALLOCATEALL` governs the contiguity fallback, not size. //! - `durable_sync`: `fcntl(F_FULLFSYNC)` wrapped in //! [`crate::io::bounded::bounded_syscall`] with a 60 s deadline. //! F_FULLFSYNC is HFS+/APFS's true-fsync (flushes the disk's own @@ -25,11 +27,14 @@ use crate::io::platform_macos::{ const F_FULLFSYNC: libc::c_int = 51; pub(super) fn preallocate(file: &File, size_bytes: u64) { + // Clamp to the signed `off_t` range; an unchecked `as off_t` cast + // would wrap a >= 2^63 size to a negative length. + let len = i64::try_from(size_bytes).unwrap_or(i64::MAX) as libc::off_t; let mut fst = Fstore { fst_flags: F_ALLOCATECONTIG | F_ALLOCATEALL, fst_posmode: F_PEOFPOSMODE, fst_offset: 0, - fst_length: size_bytes as libc::off_t, + fst_length: len, fst_bytesalloc: 0, }; // First attempt: contiguous. diff --git a/src/io/writeback_file/mod.rs b/src/io/writeback_file/mod.rs index 89b00b7..5d90e7a 100644 --- a/src/io/writeback_file/mod.rs +++ b/src/io/writeback_file/mod.rs @@ -24,29 +24,32 @@ //! ## Platform split //! //! The platform-specific pieces of this wrapper — extent preallocation -//! (Linux `fallocate(KEEP_SIZE)`, macOS `F_PREALLOCATE`, Windows -//! `SetFileValidData`) and the durable-flush primitive (Linux/macOS -//! `fsync`/`F_FULLFSYNC` wrapped in a bounded syscall, Windows -//! `FlushFileBuffers`) — live in per-OS sibling modules. The dispatch -//! happens once at the bottom of this file via cfg-gated `mod` decls. -//! No inline `#[cfg(target_os = "...")]` in the business-logic above. +//! (Linux `fallocate(KEEP_SIZE)`, macOS `F_PREALLOCATE`, Windows no-op +//! today) and the durable-flush primitive (Linux/macOS +//! `fsync`/`F_FULLFSYNC` wrapped in a bounded syscall; Windows plain +//! `FlushFileBuffers`, unbounded) — live in per-OS sibling modules. The +//! dispatch happens once at the bottom of this file via cfg-gated `mod` +//! decls. No inline `#[cfg(target_os = "...")]` in the business-logic +//! above. //! //! ## Write path //! //! Writes are direct passthrough to the underlying `File` (no writer -//! thread, no ring, no batching). Empirically the Phase-2.5 -//! writer-thread architecture introduced a ~60% mux throughput -//! regression on NFS bidirectional workloads; reverting the write path -//! to direct passthrough restores the 0.20.7 baseline. The writeback -//! pipeline still runs (it's called inline from `write` / `write_all` / -//! `seek`) so the bounded-cache invariant on Linux is preserved. +//! thread, no ring, no batching). Empirically a writer-thread +//! architecture introduced a ~60% mux throughput regression on NFS +//! bidirectional workloads; the direct-passthrough write path is faster. +//! The writeback pipeline still runs (it's called inline from `write` / +//! `write_all` / `seek`) so the bounded-cache invariant on Linux is +//! preserved. //! //! ## Halt-safety //! -//! `sync_all` runs the per-OS durable-flush primitive, which on -//! Linux/macOS is wrapped in [`crate::io::bounded::bounded_syscall`] -//! with a 60 s deadline. A wedged NFS server cannot trap the muxer -//! indefinitely on the final fsync. +//! `sync_all` runs the per-OS durable-flush primitive. On Linux/macOS +//! it is wrapped in [`crate::io::bounded::bounded_syscall`] with a 60 s +//! deadline, so a wedged NFS server cannot trap the muxer indefinitely +//! on the final fsync. Windows is a known deviation: its `durable_sync` +//! calls `File::sync_all` (`FlushFileBuffers`) directly and is NOT +//! bounded — a wedged UNC/SMB share can block the final flush there. #[cfg(target_os = "linux")] mod linux; @@ -73,18 +76,23 @@ use std::path::Path; use super::writeback::WritebackPipeline; /// Granularity at which the Linux writeback pipeline issues -/// `sync_file_range` pairs. 32 MiB is the empirically best value on -/// the rip1 test bed (NFS to unraid-1 over 1 GbE, single-disk SAS): -/// 8 MiB / 64 MiB / 128 MiB all measured worse in the 0.21.x mux -/// iteration runs. Override via `FREEMKV_WRITEBACK_CHUNK_MIB` — -/// faster backends (NVMe, RAID) may tolerate larger windows. +/// `sync_file_range` pairs. 32 MiB is the empirically best value on a +/// 1 GbE NFS mount backed by a single spinning disk: 8 MiB / 64 MiB / +/// 128 MiB all measured worse. Override via `FREEMKV_WRITEBACK_CHUNK_MIB` +/// — faster backends (NVMe, RAID) may tolerate larger windows. const WRITEBACK_CHUNK_BYTES_DEFAULT: u64 = 32 * 1024 * 1024; +/// Upper bound (in MiB) accepted from `FREEMKV_WRITEBACK_CHUNK_MIB`. +/// 64 GiB — far above `CHUNK_BYTES_MAX` (256 MiB), generous for any +/// real backend, and small enough that `n * 1024 * 1024` cannot wrap +/// `u64`. Out-of-range values fall back to the default. +const WRITEBACK_CHUNK_MIB_MAX: u64 = 64 * 1024; + fn writeback_chunk_bytes() -> u64 { std::env::var("FREEMKV_WRITEBACK_CHUNK_MIB") .ok() .and_then(|v| v.parse::<u64>().ok()) - .filter(|&n| n > 0) + .filter(|&n| n > 0 && n <= WRITEBACK_CHUNK_MIB_MAX) .map(|n| n * 1024 * 1024) .unwrap_or(WRITEBACK_CHUNK_BYTES_DEFAULT) } @@ -160,6 +168,13 @@ impl WritebackFile { /// trap the calling thread indefinitely. On timeout the page cache /// is left to the kernel's normal flush-on-close path — best /// effort, but bounded. + /// + /// IMPORTANT: on Linux/macOS a successful `Ok(())` does NOT + /// guarantee the data is durable if the bounded fsync timed out or + /// was halted — only the hang is bounded, the fsync may not have + /// completed. Callers needing crash-consistency (e.g. mux-finish + /// then external commit/DB update) must not treat `Ok(())` as a + /// durability barrier. pub(crate) fn sync_all(&mut self) -> io::Result<()> { self.pipeline.finalize(); platform::durable_sync(&self.file) @@ -215,6 +230,21 @@ impl Seek for WritebackFile { } } +impl super::sink::SequentialSink for WritebackFile { + /// Drain the writeback pipeline and run the bounded durable flush — + /// the same work [`Self::sync_all`] does. Implemented explicitly (no + /// blanket impl) so a `dyn SequentialSink` / `dyn RandomAccessSink` + /// `finish()` actually finalises + fsyncs instead of hitting a no-op + /// default. Note the bounded-fsync caveat from [`Self::sync_all`] + /// applies: `Ok(())` is not a durability barrier if the fsync timed + /// out or was halted. + fn finish(&mut self) -> io::Result<()> { + self.sync_all() + } +} + +impl super::sink::RandomAccessSink for WritebackFile {} + impl Drop for WritebackFile { fn drop(&mut self) { // Run the pipeline's tail finalize so the last in-flight chunk @@ -312,4 +342,20 @@ mod tests { drop(w); assert_eq!(read_back(&p), b"onetwothree"); } + + /// finish() through a `dyn RandomAccessSink` trait object must + /// dispatch to WritebackFile's override (finalize + durable_sync), + /// not a no-op default. Bytes must be visible to a separate reader + /// before drop. + #[test] + fn finish_through_trait_object_persists() { + use crate::io::sink::RandomAccessSink; + let dir = tempfile::tempdir().unwrap(); + let p = dir.path().join("finish-dyn.bin"); + let w = WritebackFile::create(&p).unwrap(); + let mut boxed: Box<dyn RandomAccessSink> = Box::new(w); + boxed.write_all(b"durable-tail").unwrap(); + boxed.finish().unwrap(); + assert_eq!(read_back(&p), b"durable-tail"); + } } diff --git a/src/io/writeback_file/windows.rs b/src/io/writeback_file/windows.rs index f1de892..3097d94 100644 --- a/src/io/writeback_file/windows.rs +++ b/src/io/writeback_file/windows.rs @@ -1,16 +1,18 @@ //! Windows platform impl for [`super::WritebackFile`]. //! -//! TODO: this stub matches the design's "validate without a Windows -//! build env, leave a stub" carve-out. The real impl should use: +//! Current behaviour: //! -//! - `SetEndOfFile` + `SetFileValidData` for extent preallocation -//! (caller needs `SE_MANAGE_VOLUME_NAME` privilege; if unavailable -//! fall back to a write-zero path or just skip). -//! - `FlushFileBuffers` for fsync-equivalent durable flush. -//! -//! Until then: preallocate is a debug-logged no-op; durable_sync calls -//! the std `File::sync_all` (which on Windows maps to -//! `FlushFileBuffers` internally). +//! - `preallocate` is a debug-logged no-op. Windows has no +//! `fallocate`-equivalent that keeps the reported size, so extent +//! reservation is not wired up. +//! - `durable_sync` delegates to the std `File::sync_all`, which on +//! Windows maps to `FlushFileBuffers`. Unlike the Linux/macOS impls +//! this is NOT wrapped in the bounded-syscall primitive (that would +//! need an `unsafe impl Send` for `RawHandle`, which cannot be +//! validated without a Windows test env), so a wedged UNC/SMB share +//! can block the final flush. This deviation is documented on +//! [`super::WritebackFile::sync_all`] and the parent module's +//! Halt-safety section. use std::fs::File; use std::io; @@ -18,15 +20,12 @@ use std::io; pub(super) fn preallocate(_file: &File, size_bytes: u64) { tracing::debug!( target: "mux", - "WritebackFile preallocate size_hint={size_bytes} skipped (windows stub; TODO: SetFileValidData)" + "WritebackFile preallocate size_hint={size_bytes} skipped (no-op on windows)" ); } pub(super) fn durable_sync(file: &File) -> io::Result<()> { - // `File::sync_all` on Windows is `FlushFileBuffers`. Acceptable - // for now; the bounded-syscall wrapper is not used here because - // the stub also skips the worker-thread + leak machinery (the - // wrapper would need an `unsafe impl Send` for `RawHandle`, and - // designing that without a Windows test env is asking for it). + // `File::sync_all` on Windows is `FlushFileBuffers`. Not wrapped in + // the bounded-syscall primitive (see the module doc) — unbounded. file.sync_all() } diff --git a/src/keydb.rs b/src/keydb.rs index b1fc039..c996e59 100644 --- a/src/keydb.rs +++ b/src/keydb.rs @@ -5,8 +5,43 @@ use crate::error::{Error, Result}; use std::io::{Read, Write}; -use std::net::TcpStream; +use std::net::{TcpStream, ToSocketAddrs}; use std::path::PathBuf; +use std::time::Duration; + +/// Network operation timeout (connect / read / write). Keeps the daily +/// refresh thread from blocking indefinitely on an unresponsive mirror. +const NET_TIMEOUT: Duration = Duration::from_secs(10); + +/// Read timeout — longer than connect/write since the keydb body can be +/// several MiB over a slow link. +const READ_TIMEOUT: Duration = Duration::from_secs(30); + +/// Maximum redirects to follow before giving up. +const MAX_REDIRECTS: usize = 5; + +/// Upper bound on decompressed keydb size. The published keydb is a few +/// MiB; 64 MiB is a generous ceiling that still caps a decompression +/// bomb (a tiny zip/gz can otherwise inflate to GiB and OOM the daily +/// refresh thread). +const MAX_KEYDB_BYTES: u64 = 64 * 1024 * 1024; + +/// Read a decompressed stream into a String with a hard size ceiling. +/// Returns `Error::KeydbInvalid` if the input exceeds the cap or is not +/// valid UTF-8. +fn read_capped_to_string<R: Read>(reader: R) -> Result<String> { + let mut buf = Vec::new(); + // Read one byte past the cap so an exactly-at-cap stream is accepted + // but anything larger is rejected. + reader + .take(MAX_KEYDB_BYTES + 1) + .read_to_end(&mut buf) + .map_err(|_| Error::KeydbParse)?; + if buf.len() as u64 > MAX_KEYDB_BYTES { + return Err(Error::KeydbInvalid); + } + String::from_utf8(buf).map_err(|_| Error::KeydbParse) +} /// Standard keydb storage path. pub fn default_path() -> Result<PathBuf> { @@ -30,13 +65,11 @@ pub fn save(data: &[u8]) -> Result<UpdateResult> { let text = if data.starts_with(b"PK\x03\x04") { extract_zip(data)? } else if data.starts_with(&[0x1f, 0x8b]) { - let mut dec = flate2::read::GzDecoder::new(data); - let mut out = String::new(); - dec.read_to_string(&mut out) - .map_err(|_| Error::KeydbParse)?; - out + read_capped_to_string(flate2::read::GzDecoder::new(data))? } else { - String::from_utf8(data.to_vec()).map_err(|_| Error::KeydbParse)? + std::str::from_utf8(data) + .map(str::to_string) + .map_err(|_| Error::KeydbParse)? }; let entries = text @@ -82,16 +115,25 @@ pub struct UpdateResult { fn http_get(url: &str) -> Result<Vec<u8>> { let (mut host, mut port, mut path) = parse_url(url)?; - for _ in 0..5 { - let addr = format!("{host}:{port}"); - let mut stream = - TcpStream::connect(&addr).map_err(|_| Error::KeydbConnect { host: host.clone() })?; - stream - .set_read_timeout(Some(std::time::Duration::from_secs(30))) - .ok(); + for _ in 0..MAX_REDIRECTS { + // Resolve to a concrete socket address so we can bound the connect + // with connect_timeout (plain connect() uses the OS default, which + // can be minutes). + let addr = (host.as_str(), port) + .to_socket_addrs() + .ok() + .and_then(|mut it| it.next()) + .ok_or_else(|| Error::KeydbConnect { host: host.clone() })?; + let mut stream = TcpStream::connect_timeout(&addr, NET_TIMEOUT) + .map_err(|_| Error::KeydbConnect { host: host.clone() })?; + stream.set_read_timeout(Some(READ_TIMEOUT)).ok(); + stream.set_write_timeout(Some(NET_TIMEOUT)).ok(); + // HTTP/1.0 forces close-delimited framing: the server cannot reply + // with Transfer-Encoding: chunked, so the raw body is the keydb + // bytes with no chunk-size lines to de-frame. let request = format!( - "GET {path} HTTP/1.1\r\nHost: {host}\r\nConnection: close\r\nAccept-Encoding: identity\r\n\r\n" + "GET {path} HTTP/1.0\r\nHost: {host}\r\nConnection: close\r\nAccept-Encoding: identity\r\n\r\n" ); stream .write_all(request.as_bytes()) @@ -104,18 +146,27 @@ fn http_get(url: &str) -> Result<Vec<u8>> { .map_err(|_| Error::KeydbConnect { host: host.clone() })?; let header_end = find_header_end(&response).ok_or(Error::KeydbParse)?; - let headers = std::str::from_utf8(&response[..header_end]).unwrap_or(""); + // Lossy: a stray non-UTF-8 byte in the header block must not blank + // out the whole status line / Location header (which would surface + // as an undiagnosable KeydbHttp{status:0}). + let headers = String::from_utf8_lossy(&response[..header_end]); let body = &response[header_end + 4..]; - if let Some(location) = extract_header(headers, "Location") { - let parsed = parse_url(&location)?; - host = parsed.0; - port = parsed.1; - path = parsed.2; + let status = parse_status(&headers); + + // Only treat a Location header as a redirect when the status is + // actually 3xx; a 200 carrying a stray Location (some proxies) is + // not a redirect, and a 3xx without Location is a malformed redirect. + if (300..=399).contains(&status) { + let location = + extract_header(&headers, "Location").ok_or(Error::KeydbHttp { status })?; + let (next_host, next_port, next_path) = resolve_redirect(&location, &host, port)?; + host = next_host; + port = next_port; + path = next_path; continue; } - let status = parse_status(headers); if status != 200 { return Err(Error::KeydbHttp { status }); } @@ -123,17 +174,66 @@ fn http_get(url: &str) -> Result<Vec<u8>> { return Ok(body.to_vec()); } - Err(Error::KeydbHttp { status: 302 }) + Err(Error::KeydbTooManyRedirects) +} + +/// Resolve a `Location` value against the current request target. +/// Handles absolute `http://` URLs, scheme-relative `//host/path`, +/// absolute paths `/path`, and rejects unsupported schemes (e.g. +/// `https://`, which this dependency-light client cannot fetch) with a +/// diagnosable error rather than a generic parse failure. +fn resolve_redirect( + location: &str, + cur_host: &str, + cur_port: u16, +) -> Result<(String, u16, String)> { + let loc = location.trim(); + + if let Some(rest) = loc.strip_prefix("//") { + // Scheme-relative: //host[:port]/path — inherit http. + return parse_url(&format!("http://{rest}")); + } + if loc.starts_with('/') { + // Absolute path on the same host/port. + return Ok((cur_host.to_string(), cur_port, loc.to_string())); + } + if let Some(scheme) = loc.split("://").next() { + if loc.contains("://") && !scheme.eq_ignore_ascii_case("http") { + return Err(Error::KeydbUnsupportedScheme { + scheme: scheme.to_string(), + }); + } + } + parse_url(loc) } fn parse_url(url: &str) -> Result<(String, u16, String)> { + // Reject non-http(s) up front so the caller gets a scheme diagnostic + // rather than an opaque parse error. + if let Some(scheme) = url.split("://").next() { + if url.contains("://") && !scheme.eq_ignore_ascii_case("http") { + return Err(Error::KeydbUnsupportedScheme { + scheme: scheme.to_string(), + }); + } + } let url = url.strip_prefix("http://").ok_or(Error::KeydbParse)?; let (host_port, path) = match url.find('/') { Some(i) => (&url[..i], &url[i..]), None => (url, "/"), }; let (host, port) = match host_port.find(':') { - Some(i) => (&host_port[..i], host_port[i + 1..].parse().unwrap_or(80)), + Some(i) => { + let port_str = &host_port[i + 1..]; + // A non-empty-but-unparseable port is a malformed URL; only an + // omitted port defaults to 80. + let port = if port_str.is_empty() { + 80 + } else { + port_str.parse().map_err(|_| Error::KeydbParse)? + }; + (&host_port[..i], port) + } None => (host_port, 80u16), }; Ok((host.to_string(), port, path.to_string())) @@ -153,12 +253,15 @@ fn find_header_end(data: &[u8]) -> Option<usize> { } fn extract_header(headers: &str, name: &str) -> Option<String> { + // Split on the first ':' rather than byte-indexing at name.len(), + // which would panic on a multibyte UTF-8 codepoint straddling that + // offset (headers are decoded from untrusted network bytes). Also + // accepts single-character values (e.g. "Location:x"). for line in headers.lines() { - if line.len() > name.len() + 2 - && line[..name.len()].eq_ignore_ascii_case(name) - && line.as_bytes()[name.len()] == b':' - { - return Some(line[name.len() + 1..].trim().to_string()); + if let Some((key, value)) = line.split_once(':') { + if key.trim().eq_ignore_ascii_case(name) { + return Some(value.trim().to_string()); + } } } None @@ -169,14 +272,88 @@ fn extract_zip(data: &[u8]) -> Result<String> { let mut archive = zip::ZipArchive::new(cursor).map_err(|_| Error::KeydbParse)?; for i in 0..archive.len() { - let mut file = archive.by_index(i).map_err(|_| Error::KeydbParse)?; + let file = archive.by_index(i).map_err(|_| Error::KeydbParse)?; if file.name().ends_with(".cfg") || file.name().ends_with(".CFG") { - let mut text = String::new(); - file.read_to_string(&mut text) - .map_err(|_| Error::KeydbParse)?; - return Ok(text); + return read_capped_to_string(file); } } Err(Error::KeydbInvalid) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_url_defaults_and_paths() { + let (h, p, path) = parse_url("http://example.com/keydb.zip").unwrap(); + assert_eq!( + (h.as_str(), p, path.as_str()), + ("example.com", 80, "/keydb.zip") + ); + + let (h, p, path) = parse_url("http://example.com:8080").unwrap(); + assert_eq!((h.as_str(), p, path.as_str()), ("example.com", 8080, "/")); + } + + #[test] + fn parse_url_rejects_https_scheme() { + // TLS is unsupported by this client; surface a scheme diagnostic + // rather than a generic parse error. + assert!(matches!( + parse_url("https://example.com/k.zip"), + Err(Error::KeydbUnsupportedScheme { .. }) + )); + } + + #[test] + fn parse_url_rejects_malformed_port() { + // Non-empty-but-unparseable port must error, not silently fall to 80. + assert!(matches!( + parse_url("http://example.com:abc/path"), + Err(Error::KeydbParse) + )); + // An empty port still defaults to 80. + let (_, p, _) = parse_url("http://example.com:/path").unwrap(); + assert_eq!(p, 80); + } + + #[test] + fn redirect_to_https_is_unsupported_scheme_not_parse_error() { + // The bplaced-style mirror enabling TLS on a redirect must produce a + // diagnosable scheme error, not KeydbParse. + assert!(matches!( + resolve_redirect("https://mirror.example/keydb.zip", "old.host", 80), + Err(Error::KeydbUnsupportedScheme { .. }) + )); + } + + #[test] + fn redirect_scheme_relative_and_absolute_path() { + // Scheme-relative //host/path inherits http. + let (h, p, path) = resolve_redirect("//mirror.example/a.zip", "old.host", 80).unwrap(); + assert_eq!( + (h.as_str(), p, path.as_str()), + ("mirror.example", 80, "/a.zip") + ); + + // Absolute path stays on the current host/port. + let (h, p, path) = resolve_redirect("/new/path.zip", "cur.host", 8080).unwrap(); + assert_eq!( + (h.as_str(), p, path.as_str()), + ("cur.host", 8080, "/new/path.zip") + ); + + // Absolute http URL is followed normally. + let (h, _, path) = resolve_redirect("http://other.host/x.zip", "cur.host", 80).unwrap(); + assert_eq!((h.as_str(), path.as_str()), ("other.host", "/x.zip")); + } + + #[test] + fn parse_status_extracts_code() { + assert_eq!(parse_status("HTTP/1.0 200 OK\r\nFoo: bar"), 200); + assert_eq!(parse_status("HTTP/1.1 301 Moved Permanently"), 301); + assert_eq!(parse_status("garbage"), 0); + } +} diff --git a/src/keysource.rs b/src/keysource.rs index 20f7923..970b288 100644 --- a/src/keysource.rs +++ b/src/keysource.rs @@ -60,14 +60,20 @@ pub struct DiscInputs { /// cache hold exactly one. The caller drives the loop: `next_key` → /// `Disc::decrypt_with` → on `Err`, ask again → until a key decrypts or the /// source returns `None` (a genuine "no key for this disc"). Compose several -/// sources, in the caller's chosen order, with [`crate`]'s `MultiSource`. +/// sources, in the caller's chosen order, with the companion +/// `freemkv-keysources` crate's `MultiSource`. pub trait KeySource { /// Hand the NEXT candidate key for this disc, or `None` once this source is /// exhausted. Stateful: the source tracks what it already handed out this /// session, so asking again after a rejected key yields the next candidate /// (or `None`) — it never re-offers a key or re-hits a one-shot backend (an - /// online service is asked at most once). A source failure (I/O, network, - /// parse) surfaces as `None` — there is simply nothing more to try. + /// online service is asked at most once). + /// + /// `None` means only "no more candidates from this source"; it does NOT by + /// itself distinguish a genuine "no key for this disc" from a source + /// failure (I/O, network, parse). After exhaustion the caller must consult + /// [`KeySource::errored`] to tell the two apart — a failed source records + /// the failure there and still returns `None` here. fn next_key(&mut self, inputs: &DiscInputs) -> Option<Key>; /// Whether this source needs [`DiscInputs::samples`] populated (encrypted diff --git a/src/labels/bdmt.rs b/src/labels/bdmt.rs index d462f40..3b771e3 100644 --- a/src/labels/bdmt.rs +++ b/src/labels/bdmt.rs @@ -14,8 +14,9 @@ //! //! This module is intentionally separate from the BD-J `StreamLabel` //! parsers under `labels/*.rs`. The XML here is disc-level (title, -//! description, set position), not per-stream — wiring into the main -//! parser registry happens elsewhere. +//! description, set position), not per-stream. It is invoked from the +//! disc-scan path in [`labels`](super) ([`detect`] then [`parse`]), +//! and [`DiscMetadata`] is re-exported there. //! //! Real-world XML is irregular: missing description elements, multiple //! title elements (first one wins), and occasional malformed content. @@ -23,13 +24,17 @@ //! metadata" (returns `None` from the helper), and the caller can //! still get metadata from sibling-language XML files. -// The module wiring (registry hook + public re-export) is added -// separately. Until then the parse/detect entry points have no use super::xml; use crate::sector::SectorSource; use crate::udf::UdfFs; use std::collections::BTreeMap; +/// Upper bound on the size of a single `bdmt_<lang>.xml` we will read. +/// The size comes from attacker-controlled UDF metadata; real files are +/// a few KB, so 1 MiB is generous while preventing a crafted huge-size +/// entry from triggering an oversized allocation in `read_file`. +const MAX_BDMT_BYTES: u64 = 1024 * 1024; + /// Disc-level metadata extracted from `/BDMV/META/DL/bdmt_*.xml`. /// /// All maps are keyed by 3-char ISO 639-2 language code (e.g. @@ -71,6 +76,13 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<DiscMetadata> let Some(lang) = lang_code_from_filename(&entry.name) else { continue; }; + // entry.size is attacker-controlled UDF metadata and flows into + // a Vec::with_capacity in read_file. A real BDMV bdmt XML is a + // few KB; cap well above that so a crafted multi-GB size can't + // trigger a huge allocation before any parsing. + if !bdmt_size_acceptable(entry.size) { + continue; + } let path = format!("/BDMV/META/DL/{}", entry.name); let Ok(bytes) = udf.read_file(reader, &path) else { continue; @@ -78,7 +90,7 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<DiscMetadata> let Ok(text) = std::str::from_utf8(&bytes) else { continue; }; - let Some((title, description, disc_set)) = parse_bdmt_xml(&lang, text) else { + let Some((title, description, disc_set)) = parse_bdmt_xml(text) else { continue; }; out.titles.insert(lang.clone(), title); @@ -102,6 +114,13 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<DiscMetadata> } } +/// Gate a `bdmt_<lang>.xml` file by its declared (untrusted) UDF size +/// before reading it. Anything over [`MAX_BDMT_BYTES`] is skipped to +/// avoid an oversized allocation in `read_file`. +fn bdmt_size_acceptable(size: u64) -> bool { + size <= MAX_BDMT_BYTES +} + /// True if `name` matches the `bdmt_<lang>.xml` convention with a /// 3-character ISO 639-2 lang code segment. Case-insensitive. fn is_bdmt_filename(name: &str) -> bool { @@ -134,13 +153,13 @@ pub(crate) type BdmtFields = (String, Option<String>, Option<(u32, u32)>); /// Title-element preference: `<di:name>` → `<di:title>` → /// `<di:tableOfContents>/<di:titleName>` (first match wins, per the /// authoring-tool conventions documented at the module level). -pub(crate) fn parse_bdmt_xml(_lang_code: &str, xml_text: &str) -> Option<BdmtFields> { +pub(crate) fn parse_bdmt_xml(xml_text: &str) -> Option<BdmtFields> { let title = extract_title(xml_text)?; + // xml::text already returns a trimmed string (see xml::text), so the + // description is only filtered for emptiness and XML-fragment noise. let description = xml::text(xml_text, "description") .filter(|s| !s.is_empty()) - .filter(|s| !looks_like_xml(s)) - .map(|s| s.trim().to_string()) - .filter(|s| !s.is_empty()); + .filter(|s| !looks_like_xml(s)); let disc_set = extract_disc_set(xml_text); Some((title, description, disc_set)) } @@ -162,11 +181,12 @@ fn extract_title(xml_text: &str) -> Option<String> { // Order matches the module-level convention: <di:name> first // (Paramount-style), then <di:title>, then the nested // tableOfContents/titleName form. + // xml::text already trims its result, so an empty string after + // extraction means a genuinely empty element. for tag in ["name", "title"] { if let Some(s) = xml::text(xml_text, tag) { - let trimmed = s.trim(); - if !trimmed.is_empty() { - return Some(trimmed.to_string()); + if !s.is_empty() { + return Some(s); } } } @@ -175,9 +195,8 @@ fn extract_title(xml_text: &str) -> Option<String> { if let Some((s, e)) = xml::find_element(xml_text, "tableOfContents", 0) { let block = &xml_text[s..e]; if let Some(t) = xml::text(block, "titleName") { - let trimmed = t.trim(); - if !trimmed.is_empty() { - return Some(trimmed.to_string()); + if !t.is_empty() { + return Some(t); } } } @@ -197,6 +216,12 @@ fn extract_disc_set(xml_text: &str) -> Option<(u32, u32)> { .trim() .parse::<u32>() .ok()?; + // Reject nonsensical "Disc N of M" values: (0,0), (0,5), (5,2)... + // These serialize to JSON and reach downstream consumers as + // meaningless metadata. + if n < 1 || total < 1 || n > total { + return None; + } Some((n, total)) } @@ -212,7 +237,7 @@ mod tests { <discInfo xmlns:di="urn:BDA:bdmv;disclibmeta"> <di:name>Aurora Drift</di:name> </discInfo>"#; - let (title, desc, set) = parse_bdmt_xml("eng", xml).expect("title should parse"); + let (title, desc, set) = parse_bdmt_xml(xml).expect("title should parse"); assert_eq!(title, "Aurora Drift"); assert_eq!(desc, None); assert_eq!(set, None); @@ -226,7 +251,7 @@ mod tests { <di:title>Echo Chamber</di:title> <di:description>A film about machines.</di:description> </discInfo>"#; - let (title, desc, _) = parse_bdmt_xml("eng", xml).unwrap(); + let (title, desc, _) = parse_bdmt_xml(xml).unwrap(); assert_eq!(title, "Echo Chamber"); assert_eq!(desc.as_deref(), Some("A film about machines.")); } @@ -241,10 +266,54 @@ mod tests { <di:titleName>Feelings Two</di:titleName> </di:tableOfContents> </discInfo>"#; - let (title, _, _) = parse_bdmt_xml("eng", xml).unwrap(); + let (title, _, _) = parse_bdmt_xml(xml).unwrap(); assert_eq!(title, "Feelings Two"); } + #[test] + fn bdmt_size_gate_rejects_oversized_entries() { + assert!(bdmt_size_acceptable(0)); + assert!(bdmt_size_acceptable(4096)); + assert!(bdmt_size_acceptable(MAX_BDMT_BYTES)); + assert!(!bdmt_size_acceptable(MAX_BDMT_BYTES + 1)); + // A crafted multi-GB size is rejected before any allocation. + assert!(!bdmt_size_acceptable(8 * 1024 * 1024 * 1024)); + assert!(!bdmt_size_acceptable(u64::MAX)); + } + + #[test] + fn disc_set_rejects_nonsensical_pairs() { + // n > total, zero numerator, zero denominator → all None. + let over = r#"<discInfo xmlns:di="urn:BDA:bdmv;disclibmeta"> + <di:name>X</di:name> + <di:discNumber>5</di:discNumber> + <di:numSets>2</di:numSets> +</discInfo>"#; + assert_eq!(parse_bdmt_xml(over).unwrap().2, None); + + let zero_n = r#"<discInfo xmlns:di="urn:BDA:bdmv;disclibmeta"> + <di:name>X</di:name> + <di:discNumber>0</di:discNumber> + <di:numSets>5</di:numSets> +</discInfo>"#; + assert_eq!(parse_bdmt_xml(zero_n).unwrap().2, None); + + let zero_total = r#"<discInfo xmlns:di="urn:BDA:bdmv;disclibmeta"> + <di:name>X</di:name> + <di:discNumber>1</di:discNumber> + <di:numSets>0</di:numSets> +</discInfo>"#; + assert_eq!(parse_bdmt_xml(zero_total).unwrap().2, None); + + // A valid pair still passes. + let ok = r#"<discInfo xmlns:di="urn:BDA:bdmv;disclibmeta"> + <di:name>X</di:name> + <di:discNumber>2</di:discNumber> + <di:numSets>3</di:numSets> +</discInfo>"#; + assert_eq!(parse_bdmt_xml(ok).unwrap().2, Some((2, 3))); + } + #[test] fn extract_box_set_position() { let xml = r#"<discInfo xmlns:di="urn:BDA:bdmv;disclibmeta"> @@ -252,7 +321,7 @@ mod tests { <di:discNumber>2</di:discNumber> <di:numSets>5</di:numSets> </discInfo>"#; - let (_, _, set) = parse_bdmt_xml("eng", xml).unwrap(); + let (_, _, set) = parse_bdmt_xml(xml).unwrap(); assert_eq!(set, Some((2, 5))); } @@ -264,7 +333,7 @@ mod tests { <di:discNumber>3</di:discNumber> <di:numberOfSets>6</di:numberOfSets> </discInfo>"#; - let (_, _, set) = parse_bdmt_xml("eng", xml).unwrap(); + let (_, _, set) = parse_bdmt_xml(xml).unwrap(); assert_eq!(set, Some((3, 6))); } @@ -276,7 +345,7 @@ mod tests { <di:name>X</di:name> <di:discNumber>1</di:discNumber> </discInfo>"#; - let (_, _, set) = parse_bdmt_xml("eng", xml).unwrap(); + let (_, _, set) = parse_bdmt_xml(xml).unwrap(); assert_eq!(set, None); } @@ -296,7 +365,7 @@ mod tests { let mut meta = DiscMetadata::default(); for (lang, blob) in [("eng", eng_xml), ("fra", fra_xml)] { - let (title, desc, ds) = parse_bdmt_xml(lang, blob).unwrap(); + let (title, desc, ds) = parse_bdmt_xml(blob).unwrap(); meta.titles.insert(lang.to_string(), title); if let Some(d) = desc { meta.descriptions.insert(lang.to_string(), d); @@ -333,20 +402,20 @@ mod tests { // surfaces as None to its caller. Either is documented as // acceptable per the module spec. let bad = "this is not xml &&& <<< nope"; - assert!(parse_bdmt_xml("eng", bad).is_none()); + assert!(parse_bdmt_xml(bad).is_none()); // Half-open tag, no body, no close: also yields no title. let truncated = "<discInfo><di:name>"; - assert!(parse_bdmt_xml("eng", truncated).is_none()); + assert!(parse_bdmt_xml(truncated).is_none()); } #[test] fn description_with_only_child_xml_is_dropped() { - // Real-world bug from a captured disc (2026-05-11 - // capture): <di:description> contained only <di:thumbnail/> - // child elements with no actual prose. The previous parser - // surfaced the raw XML fragment as the description string. - // Now we reject candidates that begin with `<`. + // Real-world bug: <di:description> contained only + // <di:thumbnail/> child elements with no actual prose. The + // previous parser surfaced the raw XML fragment as the + // description string. Now we reject candidates that begin + // with `<`. let xml = r#"<discInfo> <di:name>Skyline Run</di:name> <di:description> @@ -355,7 +424,7 @@ mod tests { </di:description> </discInfo>"#; let (title, description, _) = - parse_bdmt_xml("eng", xml).expect("title is present so parse must succeed"); + parse_bdmt_xml(xml).expect("title is present so parse must succeed"); assert_eq!(title, "Skyline Run"); assert!( description.is_none(), @@ -371,7 +440,7 @@ mod tests { <di:name>Some Movie</di:name> <di:description>An epic tale of one man's quest for tea.</di:description> </discInfo>"#; - let (_, description, _) = parse_bdmt_xml("eng", xml).expect("must parse"); + let (_, description, _) = parse_bdmt_xml(xml).expect("must parse"); assert_eq!( description.as_deref(), Some("An epic tale of one man's quest for tea.") @@ -383,7 +452,7 @@ mod tests { let xml = r#"<discInfo><di:name> Aurora Drift </di:name></discInfo>"#; - let (title, _, _) = parse_bdmt_xml("eng", xml).unwrap(); + let (title, _, _) = parse_bdmt_xml(xml).unwrap(); assert_eq!(title, "Aurora Drift"); } diff --git a/src/labels/class_reader.rs b/src/labels/class_reader.rs index 7b56d21..f508051 100644 --- a/src/labels/class_reader.rs +++ b/src/labels/class_reader.rs @@ -16,8 +16,6 @@ // callers land. Tests below cover the API in isolation. #![allow(dead_code)] -use std::fmt; - const CLASS_MAGIC: u32 = 0xCAFEBABE; // --------------------------------------------------------------------------- @@ -34,24 +32,10 @@ pub enum Error { BadInstruction { pc: usize, opcode: u8 }, } -impl fmt::Display for Error { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Error::UnexpectedEof { needed } => write!(f, "unexpected EOF reading {}", needed), - Error::BadMagic(m) => write!(f, "bad class file magic: 0x{:08X}", m), - Error::BadCpTag { index, tag } => { - write!(f, "unknown constant pool tag {} at index {}", tag, index) - } - Error::BadUtf8 { index } => write!(f, "invalid modified-UTF-8 at cp index {}", index), - Error::BadCodeAttribute => write!(f, "malformed Code attribute"), - Error::BadInstruction { pc, opcode } => { - write!(f, "unrecognized opcode 0x{:02X} at pc={}", opcode, pc) - } - } - } -} - -impl std::error::Error for Error {} +// No Display/std::error::Error impl: this is a crate-internal, typed error +// used only for `match`/`?` within the label parsers (callers discard it via +// `let Ok(_) = ... else continue`). Per the library's zero-English rule there +// is no user-facing text; the variant fields carry the structured detail. pub type Result<T> = std::result::Result<T, Error>; @@ -673,8 +657,14 @@ fn instruction_size(code: &[u8], pc: usize) -> Option<usize> { if high < low { return None; } - let entries = (high - low + 1) as usize; - Some(padded_start - pc + 12 + entries * 4) + // `high - low + 1` can overflow i32 for adversarial bytecode + // (e.g. low=i32::MIN/high=0, or low=0/high=i32::MAX), so widen + // to i64 before adding. The product and final sum are saturating + // so they cannot overflow usize on a 32-bit target either. + let entries = (high as i64 - low as i64 + 1) as u64; + let table_bytes = entries.saturating_mul(4); + let base = (padded_start - pc + 12) as u64; + usize::try_from(base.saturating_add(table_bytes)).ok() } LOOKUPSWITCH => { let padded_start = (pc + 1 + 3) & !3; @@ -686,7 +676,11 @@ fn instruction_size(code: &[u8], pc: usize) -> Option<usize> { if npairs < 0 { return None; } - Some(padded_start - pc + 8 + (npairs as usize) * 8) + // Saturating product/sum so an attacker-supplied npairs cannot + // overflow usize on a 32-bit target. + let pair_bytes = (npairs as u64).saturating_mul(8); + let base = (padded_start - pc + 8) as u64; + usize::try_from(base.saturating_add(pair_bytes)).ok() } WIDE => { // `wide` prefixes one of: iload/lload/fload/dload/aload/ @@ -958,7 +952,12 @@ impl<'a> Reader<'a> { if self.pos + 4 > self.data.len() { return Err(Error::UnexpectedEof { needed }); } - let v = u32::from_be_bytes(self.data[self.pos..self.pos + 4].try_into().unwrap()); + let v = u32::from_be_bytes([ + self.data[self.pos], + self.data[self.pos + 1], + self.data[self.pos + 2], + self.data[self.pos + 3], + ]); self.pos += 4; Ok(v) } @@ -971,7 +970,16 @@ impl<'a> Reader<'a> { if self.pos + 8 > self.data.len() { return Err(Error::UnexpectedEof { needed }); } - let v = u64::from_be_bytes(self.data[self.pos..self.pos + 8].try_into().unwrap()); + let v = u64::from_be_bytes([ + self.data[self.pos], + self.data[self.pos + 1], + self.data[self.pos + 2], + self.data[self.pos + 3], + self.data[self.pos + 4], + self.data[self.pos + 5], + self.data[self.pos + 6], + self.data[self.pos + 7], + ]); self.pos += 8; Ok(v) } @@ -1026,9 +1034,9 @@ mod tests { } #[test] - fn modified_utf8_three_byte_bmp() { - // U+00E9 'é' as 3-byte BMP form is unusual but legal; the 2-byte - // form is normative. Test the 2-byte form (0xC3 0xA9). + fn modified_utf8_two_byte() { + // U+00E9 'é' in the standard 2-byte modified-UTF-8 encoding + // (0xC3 0xA9), exercising the decoder's 2-byte branch. let s = decode_modified_utf8(&[0xC3, 0xA9]).unwrap(); assert_eq!(s, "é"); } @@ -1069,6 +1077,39 @@ mod tests { assert_eq!(instruction_size(&code, 0), Some(28)); } + #[test] + fn instruction_size_tableswitch_overflow_does_not_panic() { + // Adversarial low/high spanning the full i32 range. `high - low + 1` + // overflows i32; the widened i64 count then saturates the byte + // products. Must return a value (possibly None on a 32-bit usize) + // without panicking. + for (low, high) in [ + (i32::MIN, 0i32), + (0i32, i32::MAX), + (i32::MIN, i32::MAX), + (-1i32, i32::MAX), + ] { + let mut code = vec![TABLESWITCH]; + code.extend_from_slice(&[0, 0, 0]); // padding + code.extend_from_slice(&[0, 0, 0, 0]); // default offset + code.extend_from_slice(&low.to_be_bytes()); + code.extend_from_slice(&high.to_be_bytes()); + // No need to supply the (enormous) jump table; size computation + // must not read it. + let _ = instruction_size(&code, 0); + } + } + + #[test] + fn instruction_size_lookupswitch_overflow_does_not_panic() { + // Maximal npairs; `npairs * 8` must saturate rather than overflow. + let mut code = vec![LOOKUPSWITCH]; + code.extend_from_slice(&[0, 0, 0]); // padding + code.extend_from_slice(&[0, 0, 0, 0]); // default + code.extend_from_slice(&i32::MAX.to_be_bytes()); // npairs = i32::MAX + let _ = instruction_size(&code, 0); + } + #[test] fn instruction_size_wide() { // wide iload: 4 bytes. wide iinc: 6 bytes. diff --git a/src/labels/clpi_audit.rs b/src/labels/clpi_audit.rs index af0b25b..5332132 100644 --- a/src/labels/clpi_audit.rs +++ b/src/labels/clpi_audit.rs @@ -1,33 +1,20 @@ //! CLPI vs MPLS cross-validation diagnostic. //! -//! Empirical question (raised 2026-05-11): is CLPI's per-stream -//! language and codec data truly redundant with MPLS's STN-table data -//! on real-world Blu-rays? +//! Walks both per-clip CLPI program info and per-playlist MPLS STN +//! tables, normalizes their stream lists by `(PID, language, +//! coding_type)`, and classifies each PID into one of four buckets: //! -//! Build a quick audit that walks both sources, normalizes their stream -//! lists by (PID, language, coding_type), and flags any disagreement. +//! 1. **CLPI only** — a stream present in a `.clpi` ProgramInfo that no +//! playlist STN table references (orphan on disc). +//! 2. **MPLS only** — a stream a playlist references that no `.clpi` +//! ProgramInfo lists (indicates a parser disagreement). +//! 3. **Match** — both sources agree on coding_type and language. +//! 4. **Divergent** — both sources see the PID but disagree on +//! coding_type or language. //! -//! Three classes of mismatch we want to detect: -//! -//! 1. **CLPI has streams MPLS doesn't reference.** Orphan streams in -//! the .m2ts that no playlist's STN table includes. Means the user -//! can't reach them through the menu but they're physically on the -//! disc. -//! 2. **MPLS has streams CLPI doesn't list.** Should never happen if -//! both parsers are correct — playlists reference clips which -//! reference streams. If it happens, one of our parsers has a bug. -//! 3. **Same PID, different language / coding_type.** The playlist re- -//! tagged a stream's metadata. Rare but spec-permitted. Means CLPI -//! and MPLS disagree about the same physical stream's properties. -//! -//! If audits across the corpus show zero mismatches of any class, CLPI -//! program_info extraction is **empirically redundant** for labels and -//! we can leave it out of the registry. If even one mismatch surfaces, -//! we add a CLPI label parser to the registry as belt-and-suspenders. -//! -//! This module exposes `audit(reader, udf)` returning a structured -//! report. Surfaced via the labels-analyze tool — not part of the -//! `analyze()` pipeline (no impact on the label output). +//! [`audit`] returns a structured [`ClpiVsMplsAudit`] report. This is a +//! diagnostic surface only; it does not feed the label-selection +//! pipeline. use crate::sector::SectorSource; use crate::udf::UdfFs; @@ -45,11 +32,14 @@ pub struct ClpiVsMplsRow { } impl ClpiVsMplsRow { - /// Three rules for classification: - /// - both sources missing (impossible — caller wouldn't insert) - /// - one source missing → class A or B (orphan-on-disc / playlist-only) - /// - both present but fields differ → class C (metadata divergence) - /// - both present and identical → no mismatch + /// Classification rules: + /// - one coding_type present, the other missing → `ClpiOnly` / + /// `MplsOnly` + /// - both coding_types present, fields differ → `Divergent` + /// - both coding_types present and identical → `Match` + /// - both coding_types missing (`audit` never builds this, but a + /// caller can construct such a row) → compare the language fields: + /// `Divergent` if they differ, else `Match` pub fn class(&self) -> ClpiVsMplsClass { match ( self.clpi_coding_type.is_some(), @@ -66,7 +56,13 @@ impl ClpiVsMplsRow { ClpiVsMplsClass::Divergent } } - (false, false) => ClpiVsMplsClass::Match, + (false, false) => { + if self.clpi_language == self.mpls_language { + ClpiVsMplsClass::Match + } else { + ClpiVsMplsClass::Divergent + } + } } } } @@ -95,6 +91,9 @@ pub struct ClpiVsMplsAudit { } impl ClpiVsMplsAudit { + /// Count rows by class, returned in the fixed order + /// `(clpi_only, mpls_only, matches, divergent)` matching the + /// [`ClpiVsMplsClass`] variants. pub fn class_counts(&self) -> (usize, usize, usize, usize) { let mut clpi_only = 0; let mut mpls_only = 0; @@ -138,6 +137,11 @@ pub fn audit(reader: &mut dyn SectorSource, udf: &UdfFs) -> ClpiVsMplsAudit { continue; }; for s in clip.streams { + if s.pid == 0 { + // PID 0 means "no PID in stream entry" — skip rather + // than collide, mirroring the MPLS side below. + continue; + } clpi_by_pid .entry(s.pid) .or_insert((s.coding_type, s.language)); @@ -248,6 +252,32 @@ mod tests { assert_eq!(r.class(), ClpiVsMplsClass::Divergent); } + #[test] + fn class_both_coding_missing_divergent_on_lang() { + // Caller-built row with neither coding_type but disagreeing + // languages must classify Divergent, not Match. + let r = ClpiVsMplsRow { + pid: 0x1100, + clpi_coding_type: None, + clpi_language: Some("eng".into()), + mpls_coding_type: None, + mpls_language: Some("fra".into()), + }; + assert_eq!(r.class(), ClpiVsMplsClass::Divergent); + } + + #[test] + fn class_both_coding_missing_match_on_equal_lang() { + let r = ClpiVsMplsRow { + pid: 0x1100, + clpi_coding_type: None, + clpi_language: Some("eng".into()), + mpls_coding_type: None, + mpls_language: Some("eng".into()), + }; + assert_eq!(r.class(), ClpiVsMplsClass::Match); + } + #[test] fn class_counts_sum_rows() { let audit = ClpiVsMplsAudit { diff --git a/src/labels/criterion.rs b/src/labels/criterion.rs index 780c4a6..b6be237 100644 --- a/src/labels/criterion.rs +++ b/src/labels/criterion.rs @@ -2,16 +2,29 @@ //! //! Clean structured XML with Content/Qualifier per stream and //! stream number mapping via playbackconfig. +//! +//! When `playbackconfig.xml` is absent or maps only some streams, +//! unmapped streams get 1-based-per-type stream numbers synthesized in +//! `streamproperties.xml` order, skipping any number already claimed by +//! the map so synthesized and mapped numbers never collide. See +//! [`assign_stream_numbers`]. use super::{LabelPurpose, LabelQualifier, ParseResult, StreamLabel, StreamLabelType, xml}; use crate::sector::SectorSource; use crate::udf::UdfFs; use std::collections::HashMap; +/// Cheap signature check: a Criterion disc ships `streamproperties.xml` +/// inside a `/BDMV/JAR/*` archive. pub fn detect(udf: &UdfFs) -> bool { super::jar_file_exists(udf, "streamproperties.xml") } +/// Parse `streamproperties.xml` (+ optional `playbackconfig.xml`) into +/// per-stream labels. Returns `None` if `streamproperties.xml` is +/// absent/unparseable or yields no streams. Stream numbering follows +/// the contract documented at module level (see +/// [`assign_stream_numbers`]). pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> { let sp_data = super::read_jar_file(reader, udf, "streamproperties.xml")?; let sp_text = std::str::from_utf8(&sp_data).ok()?; @@ -29,28 +42,10 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> } } + let stream_nums = assign_stream_numbers(&stream_infos, &stream_map); + let mut labels = Vec::new(); - let mut audio_idx: u16 = 1; - let mut sub_idx: u16 = 1; - - for info in &stream_infos { - let stream_num = - stream_map - .get(&info.id) - .copied() - .unwrap_or_else(|| match info.stream_type { - StreamLabelType::Audio => { - let n = audio_idx; - audio_idx += 1; - n - } - StreamLabelType::Subtitle => { - let n = sub_idx; - sub_idx += 1; - n - } - }); - + for (info, &stream_num) in stream_infos.iter().zip(stream_nums.iter()) { labels.push(StreamLabel { stream_number: stream_num, stream_type: info.stream_type, @@ -70,6 +65,58 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> Some(ParseResult::high(labels)) } +/// Assign a 1-based stream number per `StreamInfo`, parallel to +/// `infos`. +/// +/// A stream mapped in `playbackconfig.xml` (`stream_map`) keeps its +/// mapped number. Streams with no mapping (absent or incomplete +/// `playbackconfig.xml`, or an unmatched `StreamInfo_ID`) are numbered +/// 1-based per type — but the fallback counter SKIPS any number already +/// claimed via the map, so a synthesized number can never collide with a +/// map-assigned one. (Both numbering domains are 1-based per type, and +/// `apply_labels` matches on `(type, stream_number)`, so a collision +/// would mislabel tracks.) +fn assign_stream_numbers(infos: &[StreamInfo], stream_map: &HashMap<String, u16>) -> Vec<u16> { + // Numbers already claimed by the map, per type. + let mut taken_audio: Vec<u16> = Vec::new(); + let mut taken_sub: Vec<u16> = Vec::new(); + for info in infos { + if let Some(&n) = stream_map.get(&info.id) { + match info.stream_type { + StreamLabelType::Audio => taken_audio.push(n), + StreamLabelType::Subtitle => taken_sub.push(n), + } + } + } + + let mut audio_idx: u16 = 1; + let mut sub_idx: u16 = 1; + let mut out = Vec::with_capacity(infos.len()); + for info in infos { + let n = match stream_map.get(&info.id).copied() { + Some(n) => n, + None => { + let (idx, taken) = match info.stream_type { + StreamLabelType::Audio => (&mut audio_idx, &taken_audio), + StreamLabelType::Subtitle => (&mut sub_idx, &taken_sub), + }; + // Advance past any number already claimed via the map. + // saturating: a crafted XML with >65k stream entries must + // not overflow (panic in debug, wrap-to-0 in release) on + // untrusted disc bytes. + while taken.contains(idx) { + *idx = idx.saturating_add(1); + } + let n = *idx; + *idx = idx.saturating_add(1); + n + } + }; + out.push(n); + } + out +} + struct StreamInfo { id: String, stream_type: StreamLabelType, @@ -137,10 +184,75 @@ fn parse_playback_config(text: &str, map: &mut HashMap<String, u16>) { xml::text(block, "StreamInfo_ID"), ) { if let Ok(stream_num) = stream_id_str.parse::<u16>() { - map.insert(info_id, stream_num); + // Stream numbers are 1-based per the apply_labels + // contract; a mapped 0 is unmatchable and silently + // drops the label. Skip it rather than store it. + if stream_num != 0 { + map.insert(info_id, stream_num); + } } } from = end; } } } + +#[cfg(test)] +mod tests { + use super::*; + + fn info(id: &str, t: StreamLabelType) -> StreamInfo { + StreamInfo { + id: id.into(), + stream_type: t, + language: "eng".into(), + variant: String::new(), + purpose: LabelPurpose::Normal, + qualifier: LabelQualifier::None, + } + } + + #[test] + fn fallback_numbers_dense_when_map_empty() { + let infos = vec![ + info("a0", StreamLabelType::Audio), + info("a1", StreamLabelType::Audio), + info("s0", StreamLabelType::Subtitle), + ]; + let nums = assign_stream_numbers(&infos, &HashMap::new()); + // Per-type 1-based: audio 1,2 ; subtitle 1. + assert_eq!(nums, vec![1, 2, 1]); + } + + #[test] + fn fallback_does_not_collide_with_partial_map() { + // Map claims audio "a1" -> 1. The unmapped audio "a0" must NOT + // also get 1 (the pre-fix bug); it must skip to 2. + let mut map = HashMap::new(); + map.insert("a1".to_string(), 1u16); + let infos = vec![ + info("a0", StreamLabelType::Audio), // unmapped → fallback + info("a1", StreamLabelType::Audio), // mapped → 1 + info("a2", StreamLabelType::Audio), // unmapped → fallback + ]; + let nums = assign_stream_numbers(&infos, &map); + // a0 skips the taken 1 → 2; a1 keeps 1; a2 → 3. All distinct. + assert_eq!(nums, vec![2, 1, 3]); + let mut sorted = nums.clone(); + sorted.sort_unstable(); + sorted.dedup(); + assert_eq!(sorted.len(), 3, "stream numbers must be unique"); + } + + #[test] + fn map_fully_drives_numbers_when_complete() { + let mut map = HashMap::new(); + map.insert("a0".to_string(), 5u16); + map.insert("a1".to_string(), 9u16); + let infos = vec![ + info("a0", StreamLabelType::Audio), + info("a1", StreamLabelType::Audio), + ]; + assert_eq!(assign_stream_numbers(&infos, &map), vec![5, 9]); + } +} diff --git a/src/labels/ctrm.rs b/src/labels/ctrm.rs index 6d257fc..e460f14 100644 --- a/src/labels/ctrm.rs +++ b/src/labels/ctrm.rs @@ -9,11 +9,17 @@ use crate::sector::SectorSource; use crate::udf::UdfFs; use std::collections::HashMap; +/// Cheap signature check: a CTRM disc ships `menu_base.prop` and/or +/// `language_streams.txt` inside a `/BDMV/JAR/*` archive. pub fn detect(udf: &UdfFs) -> bool { super::jar_file_exists(udf, "menu_base.prop") || super::jar_file_exists(udf, "language_streams.txt") } +/// Full extraction: parses `language_streams.txt` (structured types) and +/// `menu_base.prop` (stream numbers + button names), merging when both +/// are present. Returns `None` when neither file is present/parseable or +/// no labels result. pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> { // Try language_streams.txt first (richer structured data) let ls_labels = parse_language_streams(reader, udf); @@ -50,9 +56,32 @@ fn merge(ls: Vec<StreamLabel>, mb: Vec<StreamLabel>) -> Vec<StreamLabel> { } } } + // Append any menu_base-only stream (present in mb but not in ls by + // (stream_type, stream_number)). Without this the both-files path + // silently drops streams the menu_base-only path would have emitted: + // language_streams is authoritative for type/purpose but is not + // necessarily a superset of menu_base. + for mb_label in mb { + let already = result.iter().any(|l| { + l.stream_type == mb_label.stream_type && l.stream_number == mb_label.stream_number + }); + if !already { + result.push(mb_label); + } + } result } +/// True if a property-key prefix denotes a commentary stream group. +/// Tightened from a bare `prefix.contains("comm")` substring scan, which +/// over-matched unrelated prefixes like `common_*` / `community_*`. We +/// split on `_` and require a `commentary` (or `comm`) segment. +fn prefix_is_commentary(prefix: &str) -> bool { + prefix + .split('_') + .any(|seg| seg.eq_ignore_ascii_case("commentary") || seg.eq_ignore_ascii_case("comm")) +} + // ── language_streams.txt parser ──────────────────────────────────────────── fn parse_language_streams(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<Vec<StreamLabel>> { @@ -73,9 +102,12 @@ fn parse_language_streams(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option< } let type_str = parts[1]; + // STN indices are 1-based; apply_labels pre-increments from 0 and + // never matches a 0, so a 0 here would emit a dead label. Skip it + // (matching the `n > 0` guard in parse_menu_base). let stream_num: u16 = match parts[2].parse() { - Ok(n) => n, - Err(_) => continue, + Ok(n) if n > 0 => n, + _ => continue, }; let language = parts[3].to_string(); let variant = if parts.len() > 4 { @@ -145,18 +177,18 @@ fn parse_language_streams(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option< if !variant.is_empty() { match variant.as_str() { - // Codec variants — use shared label vocab - "atmos" | "MLP" | "AC3" | "DTS" | "DDL" => { - codec_hint = vocab::codec(&variant).to_string(); - } // Purpose variants "eda" => final_purpose = LabelPurpose::Descriptive, // Dialect variants — pass through raw code from disc "csp" | "cs" | "lsp" | "ls" | "cf" | "pf" | "bp" | "pp" => { variant_code = variant.clone(); } - // Unknown — store as-is in codec_hint - _ => codec_hint = variant.clone(), + // Everything else: defer to vocab::codec as the single + // source of codec-name truth. If it recognizes the token + // (returns something other than the input) it's a known + // codec — store the canonical name. Otherwise it's an + // unknown token, stored as-is. + _ => codec_hint = vocab::codec(&variant).to_string(), } } @@ -182,85 +214,10 @@ fn parse_language_streams(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option< mod tests { use super::*; - /// Build a minimal menu_base.prop text and run `parse_menu_base`'s - /// inner logic via a temporary closure. This isolates the prop - /// parsing without needing a SectorSource. + /// Run the real shipping parser ([`parse_menu_base_text`]) on a + /// menu_base.prop body so tests exercise production code directly. fn parse_props(text: &str) -> Vec<StreamLabel> { - // Mirror the inner loop of parse_menu_base exactly. Kept - // separate so the test doesn't need disc fixtures. - use std::collections::HashMap; - let mut entries: HashMap<String, HashMap<String, String>> = HashMap::new(); - for line in text.lines() { - let line = line.trim(); - if line.is_empty() || line.starts_with('#') { - continue; - } - let Some(eq_pos) = line.find('=') else { - continue; - }; - let full_key = &line[..eq_pos]; - let value = &line[eq_pos + 1..]; - if let Some(dot_pos) = full_key.rfind('.') { - entries - .entry(full_key[..dot_pos].to_string()) - .or_default() - .insert(full_key[dot_pos + 1..].to_string(), value.to_string()); - } - } - let mut labels = Vec::new(); - for (prefix, props) in &entries { - let is_audio = props - .get("class") - .is_some_and(|c| c.contains("AudioButton")) - || prefix.starts_with("audio_"); - let is_subtitle = props - .get("class") - .is_some_and(|c| c.contains("SubtitleButton")) - || prefix.starts_with("subtitle_"); - let stream_num_str = props - .get("streamNumber") - .or_else(|| props.get("audioStream")) - .or_else(|| props.get("subtitleStream")); - let stream_num: u16 = match stream_num_str.and_then(|s| s.parse().ok()) { - Some(n) if n > 0 => n, - _ => continue, - }; - if !is_audio && !is_subtitle { - continue; - } - let name = props.get("name").cloned().unwrap_or_default(); - let purpose = match vocab::purpose(&name) { - LabelPurpose::Normal if prefix.contains("comm") => LabelPurpose::Commentary, - p => p, - }; - let qualifier = if is_subtitle { - vocab::qualifier(&name) - } else { - LabelQualifier::None - }; - let stream_type = if is_audio { - StreamLabelType::Audio - } else { - StreamLabelType::Subtitle - }; - let language = props - .get("audioLanguage") - .or_else(|| props.get("subtitleLanguage")) - .cloned() - .unwrap_or_default(); - labels.push(StreamLabel { - stream_number: stream_num, - stream_type, - language, - name, - purpose, - qualifier, - codec_hint: String::new(), - variant: String::new(), - }); - } - labels.sort_by_key(|l| (l.stream_type as u8, l.stream_number)); - labels + parse_menu_base_text(text) } #[test] @@ -334,6 +291,74 @@ mod tests { ); assert_eq!(labels[0].qualifier, LabelQualifier::None); } + + #[test] + fn dual_flag_entry_resolves_to_audio_with_no_subtitle_qualifier() { + // An entry tripping BOTH flags (audio_ prefix sets is_audio, + // class "SubtitleButton" sets is_subtitle). Audio wins the type, + // and the subtitle qualifier (SDH) must NOT be carried onto the + // resulting Audio label. Regression for the type/qualifier split. + let labels = parse_props( + "audio_1.class=SubtitleButton\n\ + audio_1.streamNumber=6\n\ + audio_1.name=English SDH\n", + ); + assert_eq!(labels.len(), 1); + assert_eq!(labels[0].stream_type, StreamLabelType::Audio); + assert_eq!(labels[0].qualifier, LabelQualifier::None); + } + + #[test] + fn prefix_commentary_segment_match_not_substring() { + // Genuine commentary group segments match. + assert!(prefix_is_commentary("audio_commentary")); + assert!(prefix_is_commentary("audio_commentary_1")); + assert!(prefix_is_commentary("comm")); + // Substring-only prefixes must NOT match (the over-match bug). + assert!(!prefix_is_commentary("common")); + assert!(!prefix_is_commentary("audio_common_1")); + assert!(!prefix_is_commentary("community")); + assert!(!prefix_is_commentary("audio_1")); + } + + fn lbl(t: StreamLabelType, n: u16, name: &str) -> StreamLabel { + StreamLabel { + stream_number: n, + stream_type: t, + language: String::new(), + name: name.to_string(), + purpose: LabelPurpose::Normal, + qualifier: LabelQualifier::None, + codec_hint: String::new(), + variant: String::new(), + } + } + + #[test] + fn merge_preserves_menu_base_only_streams() { + // language_streams covers audio 1; menu_base has audio 1 (name) + // AND a menu_base-only audio 2. The merge must keep audio 2 — + // the both-files path previously dropped it. + let ls = vec![lbl(StreamLabelType::Audio, 1, "")]; + let mb = vec![ + lbl(StreamLabelType::Audio, 1, "Main"), + lbl(StreamLabelType::Audio, 2, "Commentary"), + ]; + let merged = merge(ls, mb); + assert_eq!(merged.len(), 2, "menu_base-only stream must survive"); + // ls audio 1 takes its name from mb. + let a1 = merged + .iter() + .find(|l| l.stream_type == StreamLabelType::Audio && l.stream_number == 1) + .unwrap(); + assert_eq!(a1.name, "Main"); + // mb-only audio 2 is appended. + assert!( + merged + .iter() + .any(|l| l.stream_number == 2 && l.name == "Commentary") + ); + } } // ── menu_base.prop parser ────────────────────────────────────────────────── @@ -341,7 +366,18 @@ mod tests { fn parse_menu_base(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<Vec<StreamLabel>> { let data = super::read_jar_file(reader, udf, "menu_base.prop")?; let text = std::str::from_utf8(&data).ok()?; + let labels = parse_menu_base_text(text); + if labels.is_empty() { + return None; + } + Some(labels) +} +/// Parse the body of a `menu_base.prop` file into stream labels. Split +/// out from [`parse_menu_base`] (which only handles file I/O + UTF-8 +/// decode) so unit tests exercise the real parsing logic instead of a +/// hand-copied duplicate. Returns the labels sorted by (type, number). +fn parse_menu_base_text(text: &str) -> Vec<StreamLabel> { // Parse key=value, group by prefix let mut entries: HashMap<String, HashMap<String, String>> = HashMap::new(); @@ -394,6 +430,15 @@ fn parse_menu_base(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<Vec<Str continue; } + // Resolve the stream type FIRST: when an entry trips both flags + // (e.g. an `audio_` prefix with a class containing + // "SubtitleButton"), audio wins the type. + let stream_type = if is_audio { + StreamLabelType::Audio + } else { + StreamLabelType::Subtitle + }; + let name = props.get("name").cloned().unwrap_or_default(); // Purpose: ask vocab first (word-boundary matched — avoids the @@ -402,23 +447,19 @@ fn parse_menu_base(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<Vec<Str // (`audio_commentary.foo`-style keys group commentary streams // regardless of display name). let purpose = match vocab::purpose(&name) { - LabelPurpose::Normal if prefix.contains("comm") => LabelPurpose::Commentary, + LabelPurpose::Normal if prefix_is_commentary(prefix) => LabelPurpose::Commentary, p => p, }; - // Qualifier: only apply to subtitles (SDH is a subtitle concept). - let qualifier = if is_subtitle { + // Qualifier (SDH/Forced) is a subtitle-only concept. Gate on the + // RESOLVED type, not the raw is_subtitle flag, so an entry that + // resolved to Audio never carries a subtitle qualifier. + let qualifier = if stream_type == StreamLabelType::Subtitle { vocab::qualifier(&name) } else { LabelQualifier::None }; - let stream_type = if is_audio { - StreamLabelType::Audio - } else { - StreamLabelType::Subtitle - }; - // Try to extract language from audioLanguage/subtitleLanguage prop let language = props .get("audioLanguage") @@ -438,9 +479,6 @@ fn parse_menu_base(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<Vec<Str }); } - if labels.is_empty() { - return None; - } labels.sort_by_key(|l| (l.stream_type as u8, l.stream_number)); - Some(labels) + labels } diff --git a/src/labels/dbp.rs b/src/labels/dbp.rs index ac5f1bd..8d4eb10 100644 --- a/src/labels/dbp.rs +++ b/src/labels/dbp.rs @@ -1,8 +1,6 @@ -//! "dbp" framework — Magnolia Pictures BD-J authoring shop (per -//! `bd-live.magpictures.com` referenced in the disc's -//! `com/dbp/bluray.MenuXlet.perm`). Detected on UHD discs whose -//! `/BDMV/JAR/<x>.jar` (top-level, not in a subdir) contains -//! `com/dbp/` package paths. +//! "dbp" framework — a BD-J authoring framework identified by +//! `com/dbp/` package paths in a top-level `/BDMV/JAR/<x>.jar` (not in +//! a subdir). Seen on UHD discs. //! //! Stream labels live as plain ASCII strings inside compiled `.class` //! files in the jar — a quirk of the menu-rendering layer encoding @@ -16,21 +14,20 @@ //! ATextField,Subtitle0,None,Fontstrip_Composite,... //! ``` //! -//! The single uppercase letter before `TextField` is string-pool -//! prefix noise — the parser anchors on `TextField,` regardless of -//! what precedes it. `Subtitle0` is the disable-subtitles menu -//! button and is skipped (not a real subtitle stream). +//! The parser ignores any prefix before the first `TextField,` +//! occurrence — whatever string-pool ordering placed ahead of it is +//! irrelevant. `Subtitle0` is the disable-subtitles menu button and is +//! skipped (not a real subtitle stream). //! //! ## Implementation //! -//! v2 (2026-05-10): rewritten on top of [`super::class_reader`] — -//! iterates `CpInfo::Utf8` constant-pool entries instead of raw byte -//! scanning each class file. Equivalent label coverage (the literal -//! `TextField,...` strings live in the CP as Utf8 entries), but -//! structurally cleaner: no false-positive risk from method bytecode -//! or attribute names happening to contain `TextField,`. Language / -//! purpose / qualifier classification moved to [`super::vocab`] so all -//! Java-parser families share one source of truth. +//! Iterates `CpInfo::Utf8` constant-pool entries rather than raw +//! byte-scanning each class file. Equivalent label coverage (the literal +//! `TextField,...` strings live in the CP as Utf8 entries) with no +//! false-positive risk from method bytecode or attribute names that +//! happen to contain `TextField,`. Language / purpose / qualifier +//! classification lives in [`super::vocab`] so all Java-parser families +//! share one source of truth. use super::class_reader::CpInfo; use super::{ParseResult, StreamLabel, StreamLabelType, jar, vocab}; @@ -48,6 +45,9 @@ pub fn detect(udf: &UdfFs) -> bool { jar::has_any_top_level_jar(udf) } +/// Scan every top-level `/BDMV/JAR/*.jar` for the dbp framework and +/// extract its stream labels. Returns `None` if no jar carries a +/// `com/dbp/` package path or none yields any labels. pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> { jar::for_each_jar(reader, udf, |_entry_name, archive| { if !jar::has_path_prefix(archive, "com/dbp/") { diff --git a/src/labels/deluxe.rs b/src/labels/deluxe.rs index 3cd2f46..f880587 100644 --- a/src/labels/deluxe.rs +++ b/src/labels/deluxe.rs @@ -16,11 +16,13 @@ //! | Purpose | 8 ldcs starting `Normal, Commentary, PiP, Trivia, ...` | //! | VideoFormat | 7 ldcs starting `HD, HDR10 Plus, HD Dolby, ...` | //! | Region | 22 ldcs starting `USA_D1, LIC1, LIC2, LIC3, ...` | -//! | Studio | 6 ldcs starting `Disney, Marvel, Pixar, ...` | -//! | Codec | many `new` instructions, 0 ldcs in `<clinit>` (codec strings live in subclasses) | +//! | Studio | 6 ldcs in `<clinit>` | //! //! Matching on the shape rather than the class name keeps the parser -//! working across obfuscation variants. +//! working across obfuscation variants. Codec strings come from the +//! standard BD-J `org/bluray/ti/CodingType` enum referenced directly by +//! the binding constructors (see [`StackVal::CodingType`]), not from a +//! Deluxe-internal enum. //! //! ## Implementation phases //! @@ -29,14 +31,6 @@ //! the framework-stable fingerprints. Output: `Vec<(label, MasterEnum)>` //! with full ordinal → string-value tables. //! -//! - **Phase B** — codec enum subclass walk (`decode_codec_enum`). -//! The codec enum's `<clinit>` has many `new` instructions and zero -//! string ldcs — codec name strings live in the subclasses each -//! `new` constructs. Walks every referenced subclass's constant -//! pool, extracts the codec name string, following the standard Java -//! enum compilation convention (each enum value's `<init>` is called -//! with its name string as the first arg). -//! //! - **Phase C** — binding-class identification (`find_binding_classes`). //! The per-stream table is built by some class via repeated //! `getstatic` references to the master enums identified in A. @@ -105,21 +99,6 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> // Build a fast-lookup table for Phase D's bytecode decoder. let master_table = MasterEnumTable::from(&enums); - // Phase B — codec enum (structural + subclass walk). - let codec_shape = find_codec_enum(archive); - let codec_table = match codec_shape.as_ref() { - Some(shape) => decode_codec_enum(archive, shape), - None => CodecTable::default(), - }; - if let Some(shape) = &codec_shape { - tracing::info!( - jar = %entry_name, - class = %shape.class_name, - count = codec_table.codecs.len(), - "deluxe codec enum decoded", - ); - } - // Phase C — find ALL binding-class candidates (audio + subtitle // are often split across two classes on Deluxe). Each gets its // own `<clinit>` walk; constructions union into a single @@ -323,203 +302,6 @@ fn ldcs_match_prefix(ldcs: &[String], prefix: &[&str]) -> bool { .all(|(got, want)| got == want) } -/// Phase B (structural): identify the codec enum class. The codec -/// enum's `<clinit>` has many `new` instructions (one per codec value) -/// and zero string ldcs — codec name strings live in the subclasses -/// each `new` constructs, not in the enum class itself. This function -/// returns the candidate enum's class name + the ordered list of -/// subclass class names; [`decode_codec_enum`] walks those subclasses -/// to extract the codec strings. -pub(crate) fn find_codec_enum(archive: &mut jar::Jar) -> Option<CodecEnumShape> { - let mut best: Option<(String, Vec<String>)> = None; - jar::for_each_class(archive, |class_name, class| { - let Some((news, ldcs)) = clinit_news_and_ldcs(class) else { - return; - }; - // Codec enum's <clinit> has many `new` ops, 0 string ldcs. - if news.len() < 20 || !ldcs.is_empty() { - return; - } - match &best { - None => best = Some((class_name.to_string(), news)), - Some((_, prev)) => { - if news.len() > prev.len() { - best = Some((class_name.to_string(), news)); - } - } - } - }); - best.map(|(class_name, subclass_news)| CodecEnumShape { - class_name, - subclass_news, - }) -} - -#[derive(Debug)] -pub(crate) struct CodecEnumShape { - pub class_name: String, - /// Ordered list of class names referenced by `new` in <clinit>. - /// One entry per codec enum value; subclass walking resolves - /// each to a codec string. - pub subclass_news: Vec<String>, -} - -/// Phase B (subclass walk): given the codec enum's structural shape, -/// walk each referenced subclass's constant pool to extract its -/// codec name string. Output is ordinal-indexed: `codecs[i]` is the -/// codec name for the i-th `new` instruction in the enum's `<clinit>`. -/// -/// The codec name extraction heuristic: each subclass's constant -/// pool typically contains a small number of Utf8 entries; the -/// codec-name-shaped one is uppercase, ≥4 chars, optionally with -/// underscores or digits. We pick the first matching Utf8 entry that -/// isn't a method-descriptor sigil, class-name fragment, or attribute -/// name. Empty string when no candidate is found — the parser can -/// surface "unknown codec at ordinal N" via tracing. -pub(crate) fn decode_codec_enum(archive: &mut jar::Jar, shape: &CodecEnumShape) -> CodecTable { - // Two-pass: first pass extracts the codec-name candidate from - // every class in the jar (cheap to do all at once, cache for the - // ordinal-ordered second pass). - let mut name_by_class: HashMap<String, String> = HashMap::new(); - let wanted: HashSet<&str> = shape.subclass_news.iter().map(String::as_str).collect(); - jar::for_each_class(archive, |class_name, class| { - if !wanted.contains(class_name) { - return; - } - if let Some(name) = extract_codec_name(class) { - name_by_class.insert(class_name.to_string(), name); - } - }); - - let codecs: Vec<String> = shape - .subclass_news - .iter() - .map(|c| name_by_class.get(c).cloned().unwrap_or_default()) - .collect(); - CodecTable { codecs } -} - -/// Per-codec name table — `codecs[ordinal]` is the codec string for -/// that enum value. Empty string for ordinals where Phase B couldn't -/// extract a name (rare; logged via tracing). -#[derive(Debug, Default, Clone)] -pub(crate) struct CodecTable { - pub codecs: Vec<String>, -} - -impl CodecTable { - /// Resolve a codec enum ordinal to its name string. Returns None - /// for out-of-range ordinals or for entries Phase B couldn't - /// extract (those slots are stored as empty strings, which this - /// helper normalizes to None). - #[allow(dead_code)] // surface for callers; interpret_streams uses - // binding_type substring match for now (codec-ordinal wiring - // deferred until corpus bytecode confirms the codec arg position). - pub fn get(&self, ordinal: u16) -> Option<&str> { - let s = self.codecs.get(ordinal as usize)?; - if s.is_empty() { None } else { Some(s.as_str()) } - } -} - -/// Heuristic: extract the codec-name string from a codec-enum -/// subclass's constant pool. Codec names are uppercase tokens with -/// optional underscores/digits, ≥4 chars (e.g. "ATMOS_HD_AUDIO", -/// "DOLBY_AC3_AUDIO", "DTS_HD_MA", "PCM_5_1"). We scan the pool's -/// Utf8 entries and pick the first that: -/// - is ≥4 chars -/// - contains only A-Z, 0-9, and _ -/// - contains at least one underscore OR is a known codec token -/// (the underscore signal is what separates "ATMOS_HD_AUDIO" -/// from "Utf8" / "Code" / "Object" attribute names). -/// -/// Returns `None` when no candidate matches — the caller's `codecs[i]` -/// will be empty for that ordinal. -fn extract_codec_name(class: &ClassFile) -> Option<String> { - for (_, entry) in class.constant_pool.iter() { - let CpInfo::Utf8(s) = entry else { - continue; - }; - if s.len() < 4 { - continue; - } - if !s - .chars() - .all(|c| c.is_ascii_uppercase() || c.is_ascii_digit() || c == '_') - { - continue; - } - if !s.contains('_') { - // Single-token all-caps strings might still be valid - // (e.g. "ATMOS", "DTS"). Require at least one of the - // known codec token roots to avoid false positives like - // attribute names that happen to be uppercase. For now - // we only accept these as a fallback. - let is_known_root = [ - "ATMOS", "DOLBY", "DTS", "TRUEHD", "MLP", "AC3", "EAC3", "PCM", - ] - .iter() - .any(|root| s == *root); - if !is_known_root { - continue; - } - } - return Some(s.clone()); - } - None -} - -/// Walk `<clinit>` and return `(new_class_names, ldc_strings)`. Used -/// for the codec-enum shape match where we care about both counts. -#[allow(dead_code)] -fn clinit_news_and_ldcs( - class: &super::class_reader::ClassFile, -) -> Option<(Vec<String>, Vec<String>)> { - let mut news = Vec::new(); - let mut ldcs = Vec::new(); - let mut found = false; - let mut _aastore = 0u32; - for m in &class.methods { - let Some(name) = class.member_name(m) else { - continue; - }; - if name != "<clinit>" { - continue; - } - found = true; - let Some(code) = m.code(&class.constant_pool) else { - continue; - }; - for insn in code.instructions() { - match insn.opcode { - NEW => { - if let Some(idx) = insn.cp_index() { - if let Some(n) = class.constant_pool.class_name(idx) { - news.push(n.to_string()); - } - } - } - LDC | LDC_W => { - if let Some(idx) = insn.cp_index() { - let s = match class.constant_pool.get(idx) { - Some(CpInfo::String { string_index }) => { - class.constant_pool.utf8(*string_index).map(str::to_string) - } - Some(CpInfo::Utf8(s)) => Some(s.clone()), - _ => None, - }; - if let Some(s) = s { - ldcs.push(s); - } - } - } - AASTORE => _aastore += 1, - _ => {} - } - } - } - if found { Some((news, ldcs)) } else { None } -} - // ── Phase C: find the binding class ───────────────────────────────────────── /// Phase C: identify the class that builds the per-stream label table. @@ -647,15 +429,18 @@ pub(crate) fn decode_binding( binding_class_name: &str, master: &MasterEnumTable, ) -> Vec<Construction> { - let mut out: Vec<Construction> = Vec::new(); let target_name = binding_class_name.to_string(); - jar::for_each_class(archive, |class_name, class| { + // Short-circuit on the name match: try_each_class stops iterating + // (and stops decompressing/parsing remaining .class entries) as soon + // as the closure returns Some, instead of walking the whole jar past + // the target. + jar::try_each_class(archive, |class_name, class| { if class_name != target_name { - return; + return None; } - out = decode_binding_class(class, master); - }); - out + Some(decode_binding_class(class, master)) + }) + .unwrap_or_default() } /// Walk every method named `<clinit>` (typically only one) on this @@ -1565,25 +1350,6 @@ mod tests { ); } - #[test] - fn extract_codec_name_picks_uppercase_with_underscore() { - // Synthetic class file built via ClassFile::parse would be - // overkill; here we directly invoke extract_codec_name via a - // minimal hand-built ClassFile. Skip — covered indirectly by - // the end-to-end Phase B tests at corpus runtime. Tested - // signal: the matcher logic itself. - // (Helper inlined for clarity rather than spinning up a fake - // class.) - let candidate_strings = ["Code", "Utf8", "ATMOS_HD_AUDIO", "MyVar"]; - let result = candidate_strings.iter().find(|s| { - s.len() >= 4 - && s.chars() - .all(|c| c.is_ascii_uppercase() || c.is_ascii_digit() || c == '_') - && s.contains('_') - }); - assert_eq!(result, Some(&"ATMOS_HD_AUDIO")); - } - #[test] fn master_enum_table_resolves_field_to_ordinal() { let table = lang_enum_master(); diff --git a/src/labels/jar.rs b/src/labels/jar.rs index a475898..4fe1411 100644 --- a/src/labels/jar.rs +++ b/src/labels/jar.rs @@ -7,17 +7,20 @@ //! "open every top-level jar, look at every .class inside" without //! repeating the zip-archive boilerplate. -// `try_each_class` is staged for `labels::deluxe`, which needs the -// early-return form to short-circuit class iteration on a match. -// dead-code allow comes off when deluxe lands. -#![allow(dead_code)] - use super::class_reader::ClassFile; use crate::sector::SectorSource; use crate::udf::UdfFs; -use std::io::Cursor; +use std::io::{Cursor, Read}; use zip::ZipArchive; +/// Upper bound on bytes read out of a single `.class` entry. The jar's +/// uncompressed-size field is attacker-controlled disc metadata, so the +/// buffer is grown incrementally and the read is capped here rather than +/// pre-sized from the declared size. A real BD-J `.class` is far under +/// this ceiling (64 MiB); a lying header simply gets truncated and the +/// class fails to parse, which is skipped like any other bad entry. +const MAX_CLASS_BYTES: u64 = 64 * 1024 * 1024; + /// In-memory zip archive: backed by a `Vec<u8>` read from UDF. Owns /// the buffer; callers pass it to [`has_path_prefix`], [`for_each_class`], /// etc. @@ -81,15 +84,8 @@ where /// Used by parsers as a cheap "is this MY framework's jar?" check /// (e.g. `has_path_prefix(archive, "com/dbp/")` for dbp, /// `has_path_prefix(archive, "com/bydeluxe/")` for Deluxe). -pub fn has_path_prefix(archive: &mut Jar, prefix: &str) -> bool { - for i in 0..archive.len() { - if let Ok(f) = archive.by_index(i) { - if f.name().starts_with(prefix) { - return true; - } - } - } - false +pub fn has_path_prefix(archive: &Jar, prefix: &str) -> bool { + archive.file_names().any(|n| n.starts_with(prefix)) } /// Iterate every `.class` entry in the jar, parse it with @@ -103,23 +99,12 @@ pub fn for_each_class<F>(archive: &mut Jar, mut f: F) where F: FnMut(&str, &ClassFile), { - for i in 0..archive.len() { - let Ok(mut entry) = archive.by_index(i) else { - continue; - }; - if !entry.name().ends_with(".class") { - continue; - } - let name = entry.name().to_string(); - let mut bytes = Vec::with_capacity(entry.size() as usize); - if std::io::Read::read_to_end(&mut entry, &mut bytes).is_err() { - continue; - } - let Ok(class) = ClassFile::parse(&bytes) else { - continue; - }; - f(&name, &class); - } + // Defer to try_each_class; the callback always yields None so + // iteration never short-circuits. + try_each_class(archive, |name, class| { + f(name, class); + None::<()> + }); } /// Like [`for_each_class`] but allows the callback to short-circuit @@ -129,15 +114,18 @@ where F: FnMut(&str, &ClassFile) -> Option<R>, { for i in 0..archive.len() { - let Ok(mut entry) = archive.by_index(i) else { + let Ok(entry) = archive.by_index(i) else { continue; }; if !entry.name().ends_with(".class") { continue; } let name = entry.name().to_string(); - let mut bytes = Vec::with_capacity(entry.size() as usize); - if std::io::Read::read_to_end(&mut entry, &mut bytes).is_err() { + // The declared uncompressed size is attacker-controlled, so the + // buffer grows incrementally and the read is capped at + // MAX_CLASS_BYTES rather than pre-sized from entry.size(). + let mut bytes = Vec::new(); + if entry.take(MAX_CLASS_BYTES).read_to_end(&mut bytes).is_err() { continue; } let Ok(class) = ClassFile::parse(&bytes) else { @@ -149,3 +137,166 @@ where } None } + +#[cfg(test)] +mod tests { + use super::*; + + /// Smallest constant-pool-empty `.class`: magic, versions, cp_count=1 + /// (zero real entries), then empty access/this/super/interfaces/ + /// fields/methods/attributes. + const MINIMAL_CLASS: &[u8] = &[ + 0xCA, 0xFE, 0xBA, 0xBE, // magic + 0x00, 0x00, // minor + 0x00, 0x00, // major + 0x00, 0x01, // constant_pool_count = 1 -> no entries + 0x00, 0x00, // access_flags + 0x00, 0x00, // this_class + 0x00, 0x00, // super_class + 0x00, 0x00, // interfaces_count + 0x00, 0x00, // fields_count + 0x00, 0x00, // methods_count + 0x00, 0x00, // attributes_count + ]; + + /// Build a raw, single-entry, Stored (uncompressed) ZIP whose local + /// header and central directory both declare `declared_size` as the + /// uncompressed size, while the actual stored payload is `payload`. + /// This lets a test forge an attacker-controlled size field that does + /// not match the real data length. + fn build_stored_zip(name: &str, payload: &[u8], declared_size: u32) -> Vec<u8> { + let name_bytes = name.as_bytes(); + let crc: u32 = { + // CRC-32 (IEEE) over payload. + let mut crc = 0xFFFF_FFFFu32; + for &b in payload { + crc ^= b as u32; + for _ in 0..8 { + let mask = (crc & 1).wrapping_neg(); + crc = (crc >> 1) ^ (0xEDB8_8320 & mask); + } + } + !crc + }; + let mut out = Vec::new(); + // ----- Local file header ----- + let lfh_offset = out.len() as u32; + out.extend_from_slice(&0x0403_4b50u32.to_le_bytes()); // signature + out.extend_from_slice(&20u16.to_le_bytes()); // version needed + out.extend_from_slice(&0u16.to_le_bytes()); // flags + out.extend_from_slice(&0u16.to_le_bytes()); // method = Stored + out.extend_from_slice(&0u16.to_le_bytes()); // mod time + out.extend_from_slice(&0u16.to_le_bytes()); // mod date + out.extend_from_slice(&crc.to_le_bytes()); // crc-32 + out.extend_from_slice(&(payload.len() as u32).to_le_bytes()); // compressed size + out.extend_from_slice(&declared_size.to_le_bytes()); // uncompressed size (forged) + out.extend_from_slice(&(name_bytes.len() as u16).to_le_bytes()); + out.extend_from_slice(&0u16.to_le_bytes()); // extra len + out.extend_from_slice(name_bytes); + out.extend_from_slice(payload); + // ----- Central directory header ----- + let cd_offset = out.len() as u32; + out.extend_from_slice(&0x0201_4b50u32.to_le_bytes()); // signature + out.extend_from_slice(&20u16.to_le_bytes()); // version made by + out.extend_from_slice(&20u16.to_le_bytes()); // version needed + out.extend_from_slice(&0u16.to_le_bytes()); // flags + out.extend_from_slice(&0u16.to_le_bytes()); // method = Stored + out.extend_from_slice(&0u16.to_le_bytes()); // mod time + out.extend_from_slice(&0u16.to_le_bytes()); // mod date + out.extend_from_slice(&crc.to_le_bytes()); + out.extend_from_slice(&(payload.len() as u32).to_le_bytes()); // compressed size + out.extend_from_slice(&declared_size.to_le_bytes()); // uncompressed size (forged) + out.extend_from_slice(&(name_bytes.len() as u16).to_le_bytes()); + out.extend_from_slice(&0u16.to_le_bytes()); // extra len + out.extend_from_slice(&0u16.to_le_bytes()); // comment len + out.extend_from_slice(&0u16.to_le_bytes()); // disk number start + out.extend_from_slice(&0u16.to_le_bytes()); // internal attrs + out.extend_from_slice(&0u32.to_le_bytes()); // external attrs + out.extend_from_slice(&lfh_offset.to_le_bytes()); // local header offset + out.extend_from_slice(name_bytes); + let cd_size = out.len() as u32 - cd_offset; + // ----- End of central directory ----- + out.extend_from_slice(&0x0605_4b50u32.to_le_bytes()); // signature + out.extend_from_slice(&0u16.to_le_bytes()); // disk number + out.extend_from_slice(&0u16.to_le_bytes()); // cd start disk + out.extend_from_slice(&1u16.to_le_bytes()); // entries on this disk + out.extend_from_slice(&1u16.to_le_bytes()); // total entries + out.extend_from_slice(&cd_size.to_le_bytes()); + out.extend_from_slice(&cd_offset.to_le_bytes()); + out.extend_from_slice(&0u16.to_le_bytes()); // comment len + out + } + + fn open(bytes: Vec<u8>) -> Jar { + ZipArchive::new(Cursor::new(bytes)).expect("valid zip") + } + + #[test] + fn try_each_class_reads_minimal_class() { + let mut jar = open(build_stored_zip( + "Foo.class", + MINIMAL_CLASS, + MINIMAL_CLASS.len() as u32, + )); + let mut seen = Vec::new(); + let r: Option<()> = try_each_class(&mut jar, |name, _class| { + seen.push(name.to_string()); + None + }); + assert!(r.is_none()); + assert_eq!(seen, vec!["Foo.class".to_string()]); + } + + #[test] + fn for_each_class_visits_every_class() { + let mut jar = open(build_stored_zip( + "Bar.class", + MINIMAL_CLASS, + MINIMAL_CLASS.len() as u32, + )); + let mut count = 0usize; + for_each_class(&mut jar, |_, _| count += 1); + assert_eq!(count, 1); + } + + /// The uncompressed-size field is attacker-controlled. A tiny stored + /// entry that declares 0xFFFF_FFFF (≈4 GiB) must NOT trigger a 4 GiB + /// pre-allocation; with the incremental read the call completes and + /// the real (small) payload parses fine. + #[test] + fn forged_huge_uncompressed_size_does_not_preallocate() { + let mut jar = open(build_stored_zip("Evil.class", MINIMAL_CLASS, 0xFFFF_FFFF)); + let mut parsed = false; + for_each_class(&mut jar, |name, _class| { + assert_eq!(name, "Evil.class"); + parsed = true; + }); + // Reached here without OOM/abort, and the real bytes parsed. + assert!(parsed); + } + + /// The read is bounded by MAX_CLASS_BYTES: a stored entry whose real + /// payload exceeds the cap yields only the first MAX_CLASS_BYTES + /// bytes to the parser, never the full (unbounded) entry. Verified + /// here on a small cap via the entry-count path: the truncated bytes + /// still parse a valid class prefix, but no read beyond the cap + /// occurs. We assert the entry is still surfaced exactly once (the + /// cap does not drop legitimate entries) and the call returns. + #[test] + fn read_is_bounded_by_cap() { + // Padding past MINIMAL_CLASS is harmless trailing data the parser + // ignores; the point is that read_to_end stops at the cap rather + // than following a (potentially huge) declared size. + let mut payload = MINIMAL_CLASS.to_vec(); + payload.extend(std::iter::repeat(0u8).take(4096)); + let mut jar = open(build_stored_zip( + "Padded.class", + &payload, + // Forge a size far larger than the real payload. + 0xFFFF_FFFF, + )); + let mut visited = 0usize; + for_each_class(&mut jar, |_, _| visited += 1); + assert_eq!(visited, 1); + } +} diff --git a/src/labels/mod.rs b/src/labels/mod.rs index 6ccdde1..6b60b42 100644 --- a/src/labels/mod.rs +++ b/src/labels/mod.rs @@ -4,7 +4,9 @@ //! To add a new format: //! 1. Create `src/labels/myformat.rs` //! 2. Implement `pub fn detect(udf: &UdfFs) -> bool` -//! 3. Implement `pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<Vec<StreamLabel>>` +//! 3. Implement `pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult>` +//! (set [`ParseResult::confidence`]; it drives parser selection on +//! a tie) //! 4. Add `mod myformat;` below and one line to `PARSERS` array mod bdmt; @@ -69,6 +71,9 @@ pub enum LabelPurpose { Commentary, Descriptive, Score, + /// Alternate music track (e.g. an alternate end-credits / closing- + /// theme music stream), tagged by the `ime` token some BD-J + /// authoring tools emit on the secondary music audio. Ime, } @@ -377,12 +382,19 @@ fn codec_hint_consistent(hint: &str, codec: &crate::disc::Codec) -> bool { let says_dts = !says_dts_ma && !says_dts_hr && h.contains("dts"); let says_lpcm = h.contains("lpcm") || h.contains("pcm"); let says_atmos = h.contains("atmos"); + // DTS:X is an object-audio extension carried on a DTS-HD MA (or HR) + // core, exactly as Atmos rides TrueHD / DD+. The spec Codec enum has + // no DtsX variant, so a correctly-authored DTS:X hint must be judged + // consistent with its DtsHdMa/DtsHdHr carrier rather than discarded. + let says_dtsx = h.contains("dts:x") || h.contains("dts-x") || h.contains("dtsx"); let names_family = says_truehd || says_ddp || says_ac3 || says_dts_ma || says_dts_hr || says_dts || says_lpcm; // Pure-editorial hint (no codec family named) isn't asserting a codec → // consistent. "Atmos" alone implies a lossless carrier (TrueHD or DD+). + // ("DTS:X" always also matches the "dts" family above, so it never + // reaches this branch — it is handled in the DtsHdMa/DtsHdHr arms.) if !names_family { return if says_atmos { matches!(codec, Codec::TrueHd | Codec::Ac3Plus) @@ -395,8 +407,8 @@ fn codec_hint_consistent(hint: &str, codec: &crate::disc::Codec) -> bool { Codec::TrueHd => says_truehd || says_atmos, Codec::Ac3Plus => says_ddp || says_atmos, Codec::Ac3 => says_ac3, - Codec::DtsHdMa => says_dts_ma, - Codec::DtsHdHr => says_dts_hr, + Codec::DtsHdMa => says_dts_ma || says_dtsx, + Codec::DtsHdHr => says_dts_hr || says_dtsx, Codec::Dts => says_dts, Codec::Lpcm => says_lpcm, // Unknown / other stream codec — don't second-guess the parser's hint. @@ -512,9 +524,9 @@ fn extract(reader: &mut dyn SectorSource, udf: &UdfFs) -> Vec<StreamLabel> { } // CLPI orphan streams: PIDs in /BDMV/CLIPINF/*.clpi ProgramInfo - // that no MPLS playlist references. Empirical (2026-05-11): ~5% - // of streams across the 11-disc corpus are CLPI-only — physically - // on disc, not menu-reachable. Append them as Low-confidence + // that no MPLS playlist references. Empirically a small fraction of + // streams are CLPI-only — physically on disc, not menu-reachable. + // Append them as Low-confidence // labels at the tail of each stream_type (next slot after the // highest existing stream_number). let _orphans_added = append_clpi_orphans(&mut labels, reader, udf); @@ -616,10 +628,13 @@ fn append_clpi_orphans( continue; } // Translate CLPI coding_type → label stream_type. + // 0x90 = Presentation Graphics (PG subtitle). 0x91 = + // Interactive Graphics (BD-J menu overlay), NOT a user-facing + // subtitle — skip it, matching the MPLS path which drops IG. let stype = match s.coding_type { 0x80..=0x86 | 0xA1 | 0xA2 => StreamLabelType::Audio, - 0x90 | 0x91 => StreamLabelType::Subtitle, - _ => continue, // video / unknown — skip + 0x90 => StreamLabelType::Subtitle, + _ => continue, // 0x91 IG / video / unknown — skip }; // Same dedup logic as MPLS: normalize language, build codec // hint, check against existing label set. @@ -689,6 +704,26 @@ fn append_clpi_orphans( added } +/// Pick the winning parser result from `results` (built in PARSERS +/// order): highest [`Confidence`] among non-empty results, with the +/// earliest array position winning on a tie — matching `extract()`'s +/// strict-`>` first-wins scan. +/// +/// `Iterator::max_by_key` returns the LAST maximal element, so the key +/// is `(confidence, Reverse(index))`: among equal-confidence entries the +/// one with the smallest index has the largest `Reverse(index)` and is +/// selected, i.e. first wins. +fn select_result<'a>( + results: &'a [(&'static str, ParseResult)], +) -> Option<&'a (&'static str, ParseResult)> { + results + .iter() + .enumerate() + .filter(|(_, (_, r))| !r.labels.is_empty()) + .max_by_key(|(idx, (_, r))| (r.confidence, std::cmp::Reverse(*idx))) + .map(|(_, entry)| entry) +} + /// Diagnostic introspection — returns the parser that matched, the /// labels it emitted, and the inventory of files under `/BDMV/JAR/*/` /// that the discriminators looked at. Intended for `freemkv-tools @@ -714,18 +749,8 @@ pub fn analyze(reader: &mut dyn SectorSource, udf: &UdfFs) -> LabelAnalysis { } // Selection logic mirrors `extract`: highest confidence + non-empty, - // array order tiebreaker. - let chosen = all_results - .iter() - .filter(|(_, r)| !r.labels.is_empty()) - .max_by(|(_, a), (_, b)| { - // Cmp first by confidence (higher first), then position - // (earlier first). max_by yields the maximum, so we - // invert the index comparison. - a.confidence - .cmp(&b.confidence) - .then(std::cmp::Ordering::Equal) - }); + // with first-in-array-order winning on a confidence tie. + let chosen = select_result(&all_results); let (parser, confidence, mut labels) = match chosen { Some((name, r)) => (Some(*name), Some(r.confidence), r.labels.clone()), @@ -976,6 +1001,65 @@ mod registry_tests { ); } + fn one_label() -> StreamLabel { + StreamLabel { + stream_number: 1, + stream_type: StreamLabelType::Audio, + language: "eng".into(), + name: String::new(), + purpose: LabelPurpose::Normal, + qualifier: LabelQualifier::None, + codec_hint: String::new(), + variant: String::new(), + } + } + + fn result(conf: Confidence) -> ParseResult { + ParseResult { + labels: vec![one_label()], + confidence: conf, + } + } + + /// `select_result` must pick the highest-confidence non-empty result + /// and, on a confidence tie, the FIRST in array order — matching + /// `extract()`'s strict-`>` first-wins scan (regression for the old + /// `analyze()` `max_by(...then(Equal))` no-op that picked the LAST). + #[test] + fn select_result_first_wins_on_tie() { + // Two parsers, equal (Medium) confidence: the first must win. + let results = vec![ + ("alpha", result(Confidence::Medium)), + ("beta", result(Confidence::Medium)), + ]; + assert_eq!(select_result(&results).map(|(n, _)| *n), Some("alpha")); + } + + #[test] + fn select_result_highest_confidence_wins() { + let results = vec![ + ("low", result(Confidence::Low)), + ("high", result(Confidence::High)), + ("medium", result(Confidence::Medium)), + ]; + assert_eq!(select_result(&results).map(|(n, _)| *n), Some("high")); + } + + #[test] + fn select_result_skips_empty_and_handles_none() { + let empty = ParseResult { + labels: Vec::new(), + confidence: Confidence::High, + }; + // High-confidence but empty must be skipped in favour of a + // non-empty lower-confidence result. + let results = vec![("empty", empty), ("real", result(Confidence::Low))]; + assert_eq!(select_result(&results).map(|(n, _)| *n), Some("real")); + // No non-empty results → None. + let none: Vec<(&'static str, ParseResult)> = Vec::new(); + assert!(select_result(&none).is_none()); + } + /// Per-parser sanity: every parser has both detect and parse /// hooked up. Catches accidental nullification (e.g. someone /// stubbing `parse` to always-None during a refactor). @@ -1339,6 +1423,44 @@ mod apply_tests { } } + #[test] + fn apply_keeps_consistent_dtsx_hint_on_dts_hd_ma() { + // DTS:X rides a DTS-HD MA core just as Atmos rides TrueHD. A + // correctly-authored "DTS:X" hint on a DtsHdMa stream is richer + // than the spec codec yet consistent, so it's kept verbatim — + // not discarded and regenerated to "DTS-HD Master Audio". + let mut titles = vec![title_with(vec![audio( + 0x1100, + Codec::DtsHdMa, + AudioChannels::Surround71, + "eng", + )])]; + let labels = vec![audio_label(1, "eng", "DTS:X", "")]; + apply_labels(&labels, &mut titles); + if let Stream::Audio(a) = &titles[0].streams[0] { + assert_eq!(a.label, "DTS:X"); + } else { + panic!("expected audio stream"); + } + } + + #[test] + fn dtsx_hint_consistent_with_dts_hd_carriers() { + use crate::disc::Codec; + // The MED fix: a DTS:X hint must now be judged consistent with + // its DTS-HD lossless carriers (previously it was rejected, + // because says_dts_ma/says_dts_hr were both false for "DTS:X"). + assert!(codec_hint_consistent("DTS:X", &Codec::DtsHdMa)); + assert!(codec_hint_consistent("DTS-X 7.1", &Codec::DtsHdHr)); + assert!(codec_hint_consistent("dtsx", &Codec::DtsHdMa)); + // It still names the DTS family, so plain-DTS streams remain + // consistent (family match) — never discarded. + assert!(codec_hint_consistent("DTS:X", &Codec::Dts)); + // But a DTS:X hint on a non-DTS stream is a genuine mismatch. + assert!(!codec_hint_consistent("DTS:X", &Codec::TrueHd)); + assert!(!codec_hint_consistent("DTS:X", &Codec::Ac3Plus)); + } + #[test] fn apply_normalizes_plain_consistent_hint_to_marketing() { // Wicked's French track: a DD+ stream whose hint "AC-3+ 5.1" is correct diff --git a/src/labels/mpls_universal.rs b/src/labels/mpls_universal.rs index 8ea506b..83ea7d5 100644 --- a/src/labels/mpls_universal.rs +++ b/src/labels/mpls_universal.rs @@ -66,7 +66,7 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> // canonical "same physical stream" key; type+lang+codec round // out the rare case where two distinct logical streams happen // to share a PID across playlists with different metadata. - let mut seen: Vec<(u8, String, String, u16)> = Vec::new(); + let mut seen: Vec<(StreamLabelType, String, String, u16)> = Vec::new(); // Global 1-based counters keyed by StreamLabelType. Incremented // only when an entry survives dedup, so stream_numbers are dense @@ -101,8 +101,7 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> let name = language_display_name(&language); let codec_hint = build_codec_hint(label_type, entry); - let type_tag = type_tag(label_type); - let key = (type_tag, language.clone(), codec_hint.clone(), entry.pid); + let key = (label_type, language.clone(), codec_hint.clone(), entry.pid); if seen.contains(&key) { continue; } @@ -147,11 +146,12 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> fn has_mpls_extension(name: &str) -> bool { // Case-insensitive ".mpls" suffix. Some discs use uppercase, // some lowercase; UDF filenames preserve case but we don't. - let n = name.len(); - if n < 5 { - return false; - } - name[n - 5..].eq_ignore_ascii_case(".mpls") + // + // UDF names are decoded via from_utf8_lossy, so a multi-byte + // replacement char (EF BF BD) can straddle byte index n-5; a raw + // byte slice there panics on a non-char-boundary. `ends_with` on a + // lowercased copy is char-boundary-safe and still case-insensitive. + name.len() >= 5 && name.to_ascii_lowercase().ends_with(".mpls") } /// Lowercase + trim the raw 3-char ISO 639-2 code. If the lowered @@ -236,7 +236,7 @@ pub(crate) fn codec_name(coding_type: u8) -> &'static str { 0x82 => "DTS", 0x83 => "TrueHD", 0x84 => "AC-3+", - 0x85 => "DTS-HD", + 0x85 => "DTS-HD HR", // BD-ROM Part 3-1: 0x85 = DTS-HD High Resolution 0x86 => "DTS-HD MA", 0x90 => "PG", 0x91 => "IG", @@ -287,17 +287,6 @@ fn build_codec_hint(label_type: StreamLabelType, entry: &crate::mpls::StreamEntr out } -/// Dedup-tag for the label type. `u8` instead of `StreamLabelType` -/// itself because the enum does not derive `Hash` / `Eq`-by-discriminant -/// in a way that we want to couple to (and `==` works fine for the -/// linear `Vec::contains` lookup we do). -fn type_tag(t: StreamLabelType) -> u8 { - match t { - StreamLabelType::Audio => 1, - StreamLabelType::Subtitle => 2, - } -} - // ── Tests ──────────────────────────────────────────────────────────── #[cfg(test)] @@ -351,18 +340,32 @@ mod tests { /// don't have to synthesize valid MPLS bytes. fn labels_from_playlists(playlists: &[Playlist]) -> Vec<StreamLabel> { let mut labels: Vec<StreamLabel> = Vec::new(); - let mut seen: Vec<(u8, String, String, u16)> = Vec::new(); + let mut seen: Vec<(StreamLabelType, String, String, u16)> = Vec::new(); + + // Global counters hoisted OUT of the playlist loop to match + // production `parse()` (lines 77-78): stream_numbers are dense + // per type across the whole disc, not reset per playlist. + let mut audio_idx: u16 = 0; + let mut sub_idx: u16 = 0; for playlist in playlists { - let mut audio_idx: u16 = 0; - let mut sub_idx: u16 = 0; - for entry in &playlist.streams { let label_type = match entry.stream_type { 2 | 5 => StreamLabelType::Audio, 3 => StreamLabelType::Subtitle, _ => continue, }; + // Dedup BEFORE consuming a counter value, matching prod + // parse() ordering so a deduped duplicate does not burn a + // stream number. + let language = normalize_language(&entry.language); + let name = language_display_name(&language); + let codec_hint = build_codec_hint(label_type, entry); + let key = (label_type, language.clone(), codec_hint.clone(), entry.pid); + if seen.contains(&key) { + continue; + } + seen.push(key); let stream_number = match label_type { StreamLabelType::Audio => { audio_idx += 1; @@ -373,19 +376,6 @@ mod tests { sub_idx } }; - let language = normalize_language(&entry.language); - let name = language_display_name(&language); - let codec_hint = build_codec_hint(label_type, entry); - let key = ( - type_tag(label_type), - language.clone(), - codec_hint.clone(), - entry.pid, - ); - if seen.contains(&key) { - continue; - } - seen.push(key); labels.push(StreamLabel { stream_number, stream_type: label_type, @@ -474,6 +464,43 @@ mod tests { let mut langs: Vec<String> = labels.iter().map(|l| l.language.clone()).collect(); langs.sort(); assert_eq!(langs, vec!["deu", "eng", "fra"]); + + // Stream numbers must be DENSE and GLOBAL across playlists, not + // reset per playlist. eng (pl1) = 1, fra (pl1) = 2, the duplicate + // eng in pl2 is deduped (no number consumed), and deu (pl2) = 3. + // Regression guard for the per-playlist counter-reset divergence. + let num = |lang: &str| { + labels + .iter() + .find(|l| l.language == lang) + .map(|l| l.stream_number) + }; + assert_eq!(num("eng"), Some(1)); + assert_eq!(num("fra"), Some(2)); + assert_eq!(num("deu"), Some(3)); + } + + #[test] + fn has_mpls_extension_handles_short_and_non_ascii_names() { + // Short names: no panic, just false. + assert!(!has_mpls_extension("")); + assert!(!has_mpls_extension("a")); + assert!(!has_mpls_extension(".mpl")); + // Exact-length and longer valid suffixes, case-insensitive. + assert!(has_mpls_extension("0.mpls")); + assert!(has_mpls_extension("00000.MPLS")); + assert!(has_mpls_extension("Movie.MpLs")); + // Non-matching suffix. + assert!(!has_mpls_extension("file.clpi")); + // Multi-byte char near the tail must NOT panic on a byte-slice + // boundary (from_utf8_lossy U+FFFD = EF BF BD is the real-disc + // case). A name ending in such a char is simply not ".mpls". + assert!(!has_mpls_extension("na\u{FFFD}me")); + // And a name where a multi-byte char sits exactly at the n-5 + // boundary used by the old slice index. + assert!(!has_mpls_extension("ab\u{FFFD}cd")); + // A genuine .mpls preceded by a multi-byte char still matches. + assert!(has_mpls_extension("f\u{FFFD}.mpls")); } #[test] @@ -490,7 +517,7 @@ mod tests { (0x82, "DTS"), (0x83, "TrueHD"), (0x84, "AC-3+"), - (0x85, "DTS-HD"), + (0x85, "DTS-HD HR"), (0x86, "DTS-HD MA"), (0x90, "PG"), (0x91, "IG"), diff --git a/src/labels/paramount.rs b/src/labels/paramount.rs index 590a00a..b1a7bd7 100644 --- a/src/labels/paramount.rs +++ b/src/labels/paramount.rs @@ -27,24 +27,51 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> // Find the feature playlist — longest duration or name="Feature" let feature = find_feature_playlist(text)?; + let labels = labels_from_feature(&feature); + + if labels.is_empty() { + return None; + } + // High confidence: paramount's playlists.xml is fully structured + // and we extract every documented field. + Some(ParseResult::high(labels)) +} + +/// Build the stream labels from a single `<playlist .../>` feature +/// element. Split out from `parse` so the per-type numbering and +/// commentary/forced-index logic is unit-testable without a +/// `SectorSource`/`UdfFs`. +fn labels_from_feature(feature: &str) -> Vec<StreamLabel> { let mut labels = Vec::new(); // Parse audio streams - if let Some(aud) = xml::attr(&feature, "aud") { - let com_idx = xml::attr(&feature, "aud_com1_idx").and_then(|s| s.parse::<usize>().ok()); + if let Some(aud) = xml::attr(feature, "aud") { + // aud_com1_idx is a trimmed, comma-separated list of CSV positions + // (some authoring tools emit whitespace, and multiple commentary + // tracks are possible) — symmetric with sub_com1_idx below. + let com_indices: Vec<usize> = xml::attr(feature, "aud_com1_idx") + .map(|s| s.split(',').filter_map(|i| i.trim().parse().ok()).collect()) + .unwrap_or_default(); + // stream_number must match apply_labels' monotonic 1-based + // per-type counter, which increments once per *real* stream — so + // it counts only non-empty slots, not the raw CSV index. The + // commentary index comparison stays on the raw CSV index `i`, + // since aud_com1_idx is positional against the original CSV. + let mut audio_num: u16 = 0; for (i, lang) in aud.split(',').enumerate() { let lang = lang.trim(); if lang.is_empty() { continue; } - let purpose = if com_idx == Some(i) { + let purpose = if com_indices.contains(&i) { LabelPurpose::Commentary } else { LabelPurpose::Normal }; + audio_num = audio_num.saturating_add(1); labels.push(StreamLabel { - stream_number: (i + 1) as u16, + stream_number: audio_num, stream_type: StreamLabelType::Audio, language: lang.to_string(), name: String::new(), @@ -57,15 +84,18 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> } // Parse subtitle streams - if let Some(sub) = xml::attr(&feature, "sub") { - let forced: Vec<bool> = xml::attr(&feature, "forced_sub") + if let Some(sub) = xml::attr(feature, "sub") { + let forced: Vec<bool> = xml::attr(feature, "forced_sub") .map(|s| s.split(',').map(|f| f.trim() == "1").collect()) .unwrap_or_default(); - let com_indices: Vec<usize> = xml::attr(&feature, "sub_com1_idx") + let com_indices: Vec<usize> = xml::attr(feature, "sub_com1_idx") .map(|s| s.split(',').filter_map(|i| i.trim().parse().ok()).collect()) .unwrap_or_default(); + // As with audio: count only non-empty slots for stream_number, + // but keep com/forced lookups on the raw CSV index `i`. + let mut sub_num: u16 = 0; for (i, lang) in sub.split(',').enumerate() { let lang = lang.trim(); if lang.is_empty() { @@ -84,8 +114,9 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> LabelQualifier::None }; + sub_num = sub_num.saturating_add(1); labels.push(StreamLabel { - stream_number: (i + 1) as u16, + stream_number: sub_num, stream_type: StreamLabelType::Subtitle, language: lang.to_string(), name: String::new(), @@ -97,15 +128,11 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> } } - if labels.is_empty() { - return None; - } - // High confidence: paramount's playlists.xml is fully structured - // and we extract every documented field. - Some(ParseResult::high(labels)) + labels } -/// Find the feature playlist element (the one with the most audio tracks). +/// Find the feature playlist element (the one with the most non-empty +/// audio slots). fn find_feature_playlist(text: &str) -> Option<String> { let mut best: Option<String> = None; let mut best_aud_count = 0; @@ -121,9 +148,11 @@ fn find_feature_playlist(text: &str) -> Option<String> { } } - // Otherwise pick the one with the most audio streams. + // Otherwise pick the one with the most audio streams. Count only + // non-empty slots so a malformed `aud=",,,,,"` can't outscore a + // legitimate feature. if let Some(aud) = xml::attr(element, "aud") { - let count = aud.split(',').count(); + let count = aud.split(',').filter(|s| !s.trim().is_empty()).count(); if count > best_aud_count { best_aud_count = count; best = Some(element.to_string()); @@ -134,3 +163,78 @@ fn find_feature_playlist(text: &str) -> Option<String> { } best } + +#[cfg(test)] +mod tests { + use super::*; + + fn audio(labels: &[StreamLabel]) -> Vec<&StreamLabel> { + labels + .iter() + .filter(|l| l.stream_type == StreamLabelType::Audio) + .collect() + } + + fn subs(labels: &[StreamLabel]) -> Vec<&StreamLabel> { + labels + .iter() + .filter(|l| l.stream_type == StreamLabelType::Subtitle) + .collect() + } + + #[test] + fn empty_middle_slot_does_not_inflate_stream_number() { + // aud="eng,,fra": the empty middle slot is skipped, and the + // second real stream (fra) must be numbered 2, matching + // apply_labels' monotonic counter — not 3 (its raw CSV index). + let feature = r#"<playlist name="Feature" aud="eng,,fra" />"#; + let labels = labels_from_feature(feature); + let a = audio(&labels); + assert_eq!(a.len(), 2); + assert_eq!(a[0].language, "eng"); + assert_eq!(a[0].stream_number, 1); + assert_eq!(a[1].language, "fra"); + assert_eq!(a[1].stream_number, 2); + } + + #[test] + fn aud_com1_idx_trimmed_and_multivalue() { + // Whitespace around the index, and a multi-value list, must both + // resolve. com index is positional against the raw CSV, so with + // an empty slot at position 1, " 2 " marks the 'fra' track + // (CSV index 2) as commentary. + let feature = r#"<playlist aud="eng,,fra" aud_com1_idx=" 2 " />"#; + let labels = labels_from_feature(feature); + let a = audio(&labels); + assert_eq!(a.len(), 2); + assert_eq!(a[1].language, "fra"); + assert_eq!(a[1].purpose, LabelPurpose::Commentary); + assert_eq!(a[0].purpose, LabelPurpose::Normal); + } + + #[test] + fn forced_sub_aligns_with_raw_csv_index() { + // sub="eng,eng,zho,ces" forced_sub="0,0,0,1": the forced flag is + // positional on the raw CSV, so 'ces' (index 3) is forced; its + // stream_number is its non-empty position (4 here, no gaps). + let feature = r#"<playlist sub="eng,eng,zho,ces" forced_sub="0,0,0,1" />"#; + let labels = labels_from_feature(feature); + let s = subs(&labels); + assert_eq!(s.len(), 4); + assert_eq!(s[3].language, "ces"); + assert_eq!(s[3].qualifier, LabelQualifier::Forced); + assert_eq!(s[3].stream_number, 4); + } + + #[test] + fn find_feature_skips_empty_audio_slot_playlist() { + // A playlist of all-empty audio slots must not outscore a real + // two-language feature. + let xml = r#" + <playlist name="Junk" aud=",,,,," /> + <playlist name="Movie" aud="eng,fra" /> + "#; + let feature = find_feature_playlist(xml).expect("a feature is found"); + assert!(feature.contains(r#"name="Movie""#)); + } +} diff --git a/src/labels/pixelogic.rs b/src/labels/pixelogic.rs index 4073435..ece15aa 100644 --- a/src/labels/pixelogic.rs +++ b/src/labels/pixelogic.rs @@ -1,7 +1,7 @@ //! Pixelogic — `bluray_project.bin` //! //! Binary file with embedded UTF-8 token strings in STN order per -//! playlist section. Most common format (5/10 test discs). +//! playlist section. A common Pixelogic layout. //! //! Token format: `{lang}_{codec?}_{purpose?}_{region?}_` @@ -11,10 +11,15 @@ use super::{ }; use crate::sector::SectorSource; use crate::udf::UdfFs; -use std::sync::atomic::{AtomicBool, Ordering}; /// Known audio codec tokens const AUDIO_CODECS: &[&str] = &["MLP", "AC3", "DTS", "DDL", "WAV", "AC"]; +/// Sane upper bound on streams of one type within a single feature +/// section. The BD STN table caps audio at 32; this generous ceiling +/// stops a crafted blob with tens of thousands of stream tokens from +/// overflowing the u16 STN counters (panic in debug, wrap-to-0 in +/// release, which would misnumber subsequent labels). +const MAX_STREAMS_PER_TYPE: u16 = 512; /// Known region tokens const REGIONS: &[&str] = &[ "US", "UK", "CF", "PF", "CS", "LS", "BP", "PP", "SM", "TM", "CAN", "DUM", "FLE", @@ -34,15 +39,16 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> // Tracked across all parse_token calls in this run: did any stream // hit an unrecognized token component (skip-unknown path)? If yes // we downgrade confidence to Medium — the labels are still valid - // but the corpus surfaced something we don't catalogue. - let saw_unknown = AtomicBool::new(false); + // but the corpus surfaced something we don't catalogue. Parsing is + // single-threaded and sequential, so a plain bool suffices. + let mut saw_unknown = false; - let labels = assign_labels(&strings, &saw_unknown); + let labels = assign_labels(&strings, &mut saw_unknown); if labels.is_empty() { return None; } - let confidence = if saw_unknown.load(Ordering::Relaxed) { + let confidence = if saw_unknown { Confidence::Medium } else { Confidence::High @@ -54,7 +60,7 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult> /// `StreamLabel` per editorial token, numbered in STN order. Split out /// from `parse` so the section/numbering logic is unit-testable without /// a `SectorSource`/`UdfFs`. -fn assign_labels(strings: &[String], saw_unknown: &AtomicBool) -> Vec<StreamLabel> { +fn assign_labels(strings: &[String], saw_unknown: &mut bool) -> Vec<StreamLabel> { // The authoritative per-feature stream list lives in the `FPL_` // (FeaturePLaylist) section, in STN order. `SEG_*` entries are menu // segments (intros, logos, disclaimers, previews) that can also carry @@ -109,14 +115,25 @@ fn assign_labels(strings: &[String], saw_unknown: &AtomicBool) -> Vec<StreamLabe // left exactly as-is — the corpus snapshots show forced/commentary // subtitle tokens already align with STN without counting the // placeholders, and counting them regresses several discs. + // Stop accumulating once both counters reach the sane cap — a + // crafted blob can't drive them to u16 overflow. + if audio_num >= MAX_STREAMS_PER_TYPE && sub_num >= MAX_STREAMS_PER_TYPE { + break; + } + if s.starts_with("Audio Stream") { - audio_num += 1; + if audio_num < MAX_STREAMS_PER_TYPE { + audio_num += 1; + } continue; } - if let Some(label) = parse_token_inner(s, Some(saw_unknown)) { + if let Some(label) = parse_token_inner(s, Some(&mut *saw_unknown)) { match label.stream_type { StreamLabelType::Audio => { + if audio_num >= MAX_STREAMS_PER_TYPE { + continue; + } audio_num += 1; labels.push(StreamLabel { stream_number: audio_num, @@ -124,6 +141,9 @@ fn assign_labels(strings: &[String], saw_unknown: &AtomicBool) -> Vec<StreamLabe }); } StreamLabelType::Subtitle => { + if sub_num >= MAX_STREAMS_PER_TYPE { + continue; + } sub_num += 1; labels.push(StreamLabel { stream_number: sub_num, @@ -137,7 +157,7 @@ fn assign_labels(strings: &[String], saw_unknown: &AtomicBool) -> Vec<StreamLabe labels } -fn parse_token_inner(s: &str, saw_unknown: Option<&AtomicBool>) -> Option<StreamLabel> { +fn parse_token_inner(s: &str, mut saw_unknown: Option<&mut bool>) -> Option<StreamLabel> { let clean = s.trim().trim_start_matches('\t').trim_end_matches('_'); let parts: Vec<&str> = clean.split('_').collect(); if parts.len() < 2 { @@ -156,10 +176,18 @@ fn parse_token_inner(s: &str, saw_unknown: Option<&AtomicBool>) -> Option<Stream let mut is_subtitle = false; let mut is_audio = false; - for &part in &parts[1..] { - if part.is_empty() { + for &raw_part in &parts[1..] { + if raw_part.is_empty() { continue; } + // Token components are spec-uppercase (codec IDs, ADES/ACOM/SDH, + // region codes). vocab elsewhere is deliberately case-insensitive, + // so normalize each component to uppercase before the gate to + // avoid silently dropping a lowercase-authored token (which would + // fall through to the unknown branch and, with no is_audio/ + // is_subtitle set, get the whole stream discarded below). + let part_up = raw_part.to_ascii_uppercase(); + let part = part_up.as_str(); if AUDIO_CODECS.contains(&part) { codec = vocab::codec(part).to_string(); is_audio = true; @@ -182,10 +210,16 @@ fn parse_token_inner(s: &str, saw_unknown: Option<&AtomicBool>) -> Option<Stream } else if part == "STRI" || part == "TXT" { is_subtitle = true; } else if part == "FOR" { + // `FOR` (forced) is a subtitle-domain qualifier. A token whose + // only non-language component is FOR (e.g. `eng_FOR_`) would + // otherwise classify as neither audio nor subtitle and be + // dropped at the `!is_audio && !is_subtitle` guard below. Treat + // a forced marker as a subtitle signal so the stream survives. qualifier = LabelQualifier::Forced; + is_subtitle = true; } else if REGIONS.contains(&part) { variant = part.to_string(); - } else if part.starts_with("PGStream") { + } else if part.starts_with("PGSTREAM") { is_subtitle = true; } else { // Unknown token component — skip this single part rather @@ -197,8 +231,8 @@ fn parse_token_inner(s: &str, saw_unknown: Option<&AtomicBool>) -> Option<Stream // but flag the parse as Medium-confidence so callers know // some data was elided. tracing::debug!(part = %part, "pixelogic: unrecognized token component, skipping"); - if let Some(flag) = saw_unknown { - flag.store(true, Ordering::Relaxed); + if let Some(flag) = saw_unknown.as_deref_mut() { + *flag = true; } } } @@ -207,7 +241,14 @@ fn parse_token_inner(s: &str, saw_unknown: Option<&AtomicBool>) -> Option<Stream return None; } - let stream_type = if is_subtitle { + // Tie-break for tokens that signal both domains (e.g. `eng_MLP_SDH_` + // sets is_audio via the codec and is_subtitle via SDH). An audio + // codec hint is the stronger, audio-domain signal, so prefer Audio + // when one is present (keeps the parsed codec_hint instead of + // discarding it); otherwise file as Subtitle. Pure-subtitle and + // pure-audio tokens are unaffected. + let has_audio_codec = is_audio && !codec.is_empty(); + let stream_type = if is_subtitle && !has_audio_codec { StreamLabelType::Subtitle } else { StreamLabelType::Audio @@ -294,6 +335,52 @@ mod tests { assert!(parse_token_inner("ENG_MLP_", None).is_none()); // uppercase not accepted as ISO 639-2 } + #[test] + fn parse_token_dual_type_with_codec_prefers_audio() { + // `eng_MLP_SDH_` sets the audio codec (MLP) and the subtitle SDH + // qualifier. Policy: a codec hint wins -> Audio, and codec_hint is + // preserved rather than discarded. + let l = parse_token_inner("eng_MLP_SDH_", None).unwrap(); + assert_eq!(l.stream_type, StreamLabelType::Audio); + assert_eq!(l.codec_hint, "TrueHD"); + assert_eq!(l.qualifier, LabelQualifier::Sdh); + } + + #[test] + fn parse_token_solo_forced_is_subtitle() { + // A token whose only non-language component is FOR must survive as + // a forced subtitle rather than being dropped. + let l = parse_token_inner("eng_FOR_", None).unwrap(); + assert_eq!(l.stream_type, StreamLabelType::Subtitle); + assert_eq!(l.language, "eng"); + assert_eq!(l.qualifier, LabelQualifier::Forced); + } + + #[test] + fn parse_token_components_are_case_insensitive() { + // Regression for the case-sensitive gate: a lowercase codec/ + // qualifier component must classify identically to uppercase + // rather than falling through to the unknown branch and getting + // the whole stream dropped. The ISO 639-2 lang prefix is still + // required lowercase. + let l = parse_token_inner("eng_mlp_", None).unwrap(); + assert_eq!(l.stream_type, StreamLabelType::Audio); + assert_eq!(l.codec_hint, "TrueHD"); + + let l = parse_token_inner("eng_ac3_acom_", None).unwrap(); + assert_eq!(l.stream_type, StreamLabelType::Audio); + assert_eq!(l.purpose, LabelPurpose::Commentary); + assert_eq!(l.codec_hint, "Dolby Digital"); + + let l = parse_token_inner("eng_sdh_", None).unwrap(); + assert_eq!(l.stream_type, StreamLabelType::Subtitle); + assert_eq!(l.qualifier, LabelQualifier::Sdh); + + // Mixed-case region token still recognized as a variant. + let l = parse_token_inner("eng_MLP_us_", None).unwrap(); + assert_eq!(l.variant, "US"); + } + fn strs(v: &[&str]) -> Vec<String> { v.iter().map(|s| s.to_string()).collect() } @@ -305,7 +392,7 @@ mod tests { // `eng_ACOM_` commentary at STN slot 4. The commentary must land on // audio #4, not collapse onto #1 (which would tag the main feature // track as commentary). - let flag = AtomicBool::new(false); + let mut flag = false; let tokens = strs(&[ "FPL_MainFeature", "Audio Stream 1", @@ -313,7 +400,7 @@ mod tests { "Audio Stream 3", "eng_ACOM_", ]); - let labels = assign_labels(&tokens, &flag); + let labels = assign_labels(&tokens, &mut flag); let audio: Vec<_> = labels .iter() .filter(|l| l.stream_type == StreamLabelType::Audio) @@ -330,7 +417,7 @@ mod tests { // token, but the real playlist is `FPL_MainFeature`. When an FPL_ // section exists, the SEG_ one must be ignored as an anchor — so we // number from the FPL playlist, putting the commentary at slot 2. - let flag = AtomicBool::new(false); + let mut flag = false; let tokens = strs(&[ "SEG_MainFeature", "eng_ACOM_", // stray token in the menu segment — must be ignored @@ -338,7 +425,7 @@ mod tests { "Audio Stream 1", "eng_ACOM_", ]); - let labels = assign_labels(&tokens, &flag); + let labels = assign_labels(&tokens, &mut flag); let audio: Vec<_> = labels .iter() .filter(|l| l.stream_type == StreamLabelType::Audio) @@ -351,9 +438,9 @@ mod tests { #[test] fn assign_labels_falls_back_to_seg_without_fpl() { // Discs with no FPL_ playlist still anchor on SEG_MainFeature. - let flag = AtomicBool::new(false); + let mut flag = false; let tokens = strs(&["SEG_MainFeature", "eng_MLP_", "spa_AC3_"]); - let labels = assign_labels(&tokens, &flag); + let labels = assign_labels(&tokens, &mut flag); let audio: Vec<_> = labels .iter() .filter(|l| l.stream_type == StreamLabelType::Audio) diff --git a/src/labels/text.rs b/src/labels/text.rs index 40cd949..4e00661 100644 --- a/src/labels/text.rs +++ b/src/labels/text.rs @@ -1,13 +1,14 @@ //! Text-extraction helpers used by parsers that scan binary blobs for //! embedded label strings. //! -//! Promoted from two near-duplicate implementations: -//! - `pixelogic::extract_strings` (`bluray_project.bin`, min_len=4) -//! - `dbp::extract_printable` (`.class` files in jars, min_len=5) +//! Promoted from a byte-scanning helper (`bluray_project.bin`, +//! min_len=4). Single implementation, threshold passed in. //! -//! Single implementation, threshold passed in. Callers that have a -//! more structured parse path (e.g. `class_reader` for .class) should -//! prefer that — this helper is for genuinely unstructured input. +//! `dbp` no longer uses a byte-scanning helper — it iterates +//! `class_reader::CpInfo::Utf8` constant-pool entries directly. Callers +//! that have a more structured parse path (e.g. `class_reader` for +//! `.class`) should prefer that; this helper is for genuinely +//! unstructured input. /// Walk `data`, emit every maximal run of printable-ASCII bytes /// (`0x20..=0x7E`) whose length is at least `min_len`. @@ -21,13 +22,13 @@ pub fn extract_ascii_strings(data: &[u8], min_len: usize) -> Vec<String> { for &b in data { if (0x20..=0x7E).contains(&b) { current.push(b as char); - } else if current.len() >= min_len { + } else if !current.is_empty() && current.len() >= min_len { out.push(std::mem::take(&mut current)); } else { current.clear(); } } - if current.len() >= min_len { + if !current.is_empty() && current.len() >= min_len { out.push(current); } out @@ -86,4 +87,14 @@ mod tests { let got = extract_ascii_strings(b"a\0b", 0); assert_eq!(got, vec!["a", "b"]); } + + #[test] + fn min_len_zero_skips_empty_runs_on_consecutive_separators() { + // Consecutive separators must NOT emit empty strings even at + // min_len=0 — an empty string is not a "run of printable bytes". + let got = extract_ascii_strings(b"\0\0abc", 0); + assert_eq!(got, vec!["abc"]); + let got = extract_ascii_strings(b"ab\0\0\0cd\0\0", 0); + assert_eq!(got, vec!["ab", "cd"]); + } } diff --git a/src/labels/vocab.rs b/src/labels/vocab.rs index e6ab56d..14bf4cf 100644 --- a/src/labels/vocab.rs +++ b/src/labels/vocab.rs @@ -32,16 +32,19 @@ use super::{LabelPurpose, LabelQualifier}; /// Map a codec identifier found in label data to its display name. /// /// These are well-known codec identifiers used across multiple BD-J -/// authoring tools. Unknown codes pass through unchanged so callers -/// can still surface vendor-specific tokens we haven't catalogued. +/// authoring tools. Matching is case-insensitive (on-disc tokens vary: +/// `ATMOS`, `Atmos`, `atmos`). Unknown codes pass through unchanged (in +/// their original casing) so callers can still surface vendor-specific +/// tokens we haven't catalogued. pub fn codec(code: &str) -> &str { - match code { + match code.to_ascii_uppercase().as_str() { "MLP" => "TrueHD", "AC3" | "AC" => "Dolby Digital", - "DTS" => "DTS", "DDL" => "Dolby Digital Plus", "WAV" => "PCM", - "atmos" => "Dolby Atmos", + "ATMOS" => "Dolby Atmos", + // "DTS" is recognized but has no distinct display alias — return + // the original token rather than a re-cased copy. _ => code, } } @@ -54,10 +57,9 @@ pub fn codec(code: &str) -> &str { /// `variant` is the regional dialect as a human-readable English word /// (`"Brazilian"`, `"Castilian"`, `"Canadian"`, `"Simplified"`, ...) /// or `""` when the input names just a bare language without -/// dialect ("Spanish" → variant=""). The variant matches the -/// convention pixelogic / ctrm / criterion already use for their -/// `StreamLabel::variant` field: a short display token the UI can -/// surface verbatim. +/// dialect ("Spanish" → variant=""). It is a short display token +/// suitable for the [`StreamLabel::variant`](super::StreamLabel) field, +/// to be surfaced verbatim by the UI. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct LangInfo { pub code: &'static str, @@ -70,10 +72,14 @@ pub struct LangInfo { /// Handles both bare English names ("English", "Spanish") and the /// multi-word vendor variants we've seen in the corpus ("Brazilian /// Portuguese", "Castilian Spanish", "Canadian French"). Match is -/// case-insensitive; longer compound phrases win over their bare -/// counterparts (so "Brazilian Portuguese" returns -/// `LangInfo { code: "por", variant: "Brazilian" }`, not consumed by -/// the bare "Portuguese" entry). +/// case-insensitive. Compound phrases are scanned BEFORE bare names, so +/// "Brazilian Portuguese" returns +/// `LangInfo { code: "por", variant: "Brazilian" }` rather than being +/// consumed by the bare "Portuguese" entry. Within `COMPOUND_LANGS` the +/// scan is positional (first `contains` hit wins), so that table MUST be +/// maintained longest-first — a longer phrase must precede any shorter +/// phrase it contains (e.g. "latin american spanish" before +/// "latin spanish"). /// /// Bare-name matches return `variant: ""`. /// @@ -81,15 +87,16 @@ pub struct LangInfo { /// fall back to MPLS spec codes, pass through raw, or drop the stream. /// Never guesses. /// -/// Why the variant: the prior `lang() -> Option<&str>` shape silently -/// dropped regional dialect info. "Brazilian Portuguese 5.1" became -/// `language="por", variant=""` — UI displayed plain "Portuguese" -/// even though the disc had explicitly labeled this stream Brazilian. -/// Capturing the variant here parallels how pixelogic and ctrm -/// populate `StreamLabel::variant` from their own region tables. +/// Why the variant: returning only the ISO code would silently drop +/// regional dialect info — "Brazilian Portuguese 5.1" would become +/// `language="por", variant=""` and the UI would display plain +/// "Portuguese" even though the disc explicitly labeled the stream +/// Brazilian. Returning the variant lets callers populate +/// [`StreamLabel::variant`](super::StreamLabel) with the dialect. pub fn lang(text: &str) -> Option<LangInfo> { let lower = text.to_lowercase(); - // Multi-word compounds first — longest-match wins. + // Multi-word compounds first. Scan is positional (first hit wins), + // so COMPOUND_LANGS MUST stay ordered longest-first. for (needle, code, variant) in COMPOUND_LANGS { if lower.contains(needle) { return Some(LangInfo { code, variant }); @@ -250,22 +257,26 @@ fn has_word(haystack: &str, needle: &str) -> bool { if needle.is_empty() { return false; } - let bytes = haystack.as_bytes(); - let nb = needle.as_bytes(); - let mut i = 0; - while i + nb.len() <= bytes.len() { - if &bytes[i..i + nb.len()] == nb { - let before = if i == 0 { None } else { Some(bytes[i - 1]) }; - let after = bytes.get(i + nb.len()).copied(); - let bound = |c: Option<u8>| match c { - None => true, - Some(b) => !b.is_ascii_alphanumeric(), - }; - if bound(before) && bound(after) { - return true; - } + // Boundary check is char-aware (not byte-level): a non-ASCII letter + // adjacent to the match (e.g. an accented or CJK char, which is + // multiple UTF-8 bytes) is alphanumeric and so is NOT a boundary, + // preventing false positives like "sdh" inside "cafésch". Needles + // are ASCII tokens, so a byte-offset match aligns with char + // boundaries in `haystack`. + for (idx, _) in haystack.match_indices(needle) { + // Char immediately before the match. + let before_is_alnum = haystack[..idx] + .chars() + .next_back() + .is_some_and(char::is_alphanumeric); + // Char immediately after the match. + let after_is_alnum = haystack[idx + needle.len()..] + .chars() + .next() + .is_some_and(char::is_alphanumeric); + if !before_is_alnum && !after_is_alnum { + return true; } - i += 1; } false } @@ -283,12 +294,26 @@ mod tests { assert_eq!(codec("AC"), "Dolby Digital"); assert_eq!(codec("DDL"), "Dolby Digital Plus"); assert_eq!(codec("atmos"), "Dolby Atmos"); + assert_eq!(codec("WAV"), "PCM"); + assert_eq!(codec("DTS"), "DTS"); + } + + #[test] + fn codec_case_insensitive() { + // On-disc casing varies; all forms must canonicalize. + assert_eq!(codec("ATMOS"), "Dolby Atmos"); + assert_eq!(codec("Atmos"), "Dolby Atmos"); + assert_eq!(codec("atmos"), "Dolby Atmos"); + assert_eq!(codec("mlp"), "TrueHD"); + assert_eq!(codec("ac3"), "Dolby Digital"); } #[test] fn codec_unknown_passes_through() { assert_eq!(codec("FX9"), "FX9"); assert_eq!(codec(""), ""); + // Unknown tokens keep their original casing. + assert_eq!(codec("Vendor_X"), "Vendor_X"); } fn li(code: &'static str, variant: &'static str) -> LangInfo { @@ -390,6 +415,26 @@ mod tests { assert_eq!(purpose(""), LabelPurpose::Normal); } + #[test] + fn purpose_recognizes_ime() { + assert_eq!(purpose("IME"), LabelPurpose::Ime); + assert_eq!(purpose("English ime"), LabelPurpose::Ime); + // Word-boundary: "ime" inside "time" must not match. + assert_eq!(purpose("Showtime audio"), LabelPurpose::Normal); + } + + #[test] + fn has_word_treats_non_ascii_letter_as_a_letter_boundary() { + // A non-ASCII (multi-byte) letter glued to the needle is NOT a + // word boundary, so the needle must not match there. + assert!(!has_word("cafésdh", "sdh")); // 'é' precedes "sdh" + assert!(!has_word("日本sdh", "sdh")); + // But a real boundary (space / punctuation / non-letter) matches. + assert!(has_word("café sdh", "sdh")); + assert!(has_word("日本 sdh", "sdh")); + assert!(has_word("sdh", "sdh")); + } + #[test] fn qualifier_recognizes_sdh() { assert_eq!(qualifier("English SDH"), LabelQualifier::Sdh); diff --git a/src/labels/xml.rs b/src/labels/xml.rs index 1b324f3..9d4cf5d 100644 --- a/src/labels/xml.rs +++ b/src/labels/xml.rs @@ -37,6 +37,19 @@ pub fn attr(element: &str, name: &str) -> Option<String> { let name_bytes = name_lower.as_bytes(); let mut i = 0; while i + name_bytes.len() < bytes.len() { + // Skip over a quoted attribute value entirely so a name token + // embedded inside another attribute's value (e.g. + // `y="name='inner'"`) is never matched as a real attribute. + if bytes[i] == b'"' || bytes[i] == b'\'' { + let q = bytes[i]; + i += 1; + while i < bytes.len() && bytes[i] != q { + i += 1; + } + // Step past the closing quote (or to EOF). + i += 1; + continue; + } // Find the next position where `name=` could start. We need // a word boundary before the name (whitespace or `<` or `:`). if i > 0 && is_name_char(bytes[i - 1]) { @@ -96,17 +109,13 @@ pub fn attr(element: &str, name: &str) -> Option<String> { /// handled — the first close encountered wins (this matches the /// prior behavior in criterion.rs). pub fn text(xml: &str, tag: &str) -> Option<String> { - let (open_end, body_start) = find_open_tag(xml, tag, 0)?; - // Self-closing — already consumed in find_open_tag if `/>`. - if open_end == body_start { - // Means find_open_tag returned the same offset twice for - // self-closing form. (Not currently the case in our impl, - // but defensive.) - return Some(String::new()); - } - // For self-closing tags, body_start is past `/>` and we have no - // content. Detect by checking the char at body_start - 1 was `/`. - if body_start >= 2 && &xml[body_start - 2..body_start] == "/>" { + let (_open_end, body_start) = find_open_tag(xml, tag, 0)?; + // For self-closing tags, body_start is past `/>` and there is no + // content. Detect with a *byte* comparison: slicing `&xml[..]` two + // bytes back can land inside a multi-byte UTF-8 char and panic + // (untrusted on-disc XML), but indexing the byte slice never does. + let b = xml.as_bytes(); + if body_start >= 2 && b[body_start - 2] == b'/' && b[body_start - 1] == b'>' { return Some(String::new()); } // Find the matching close tag. Case-insensitive + namespace-aware. @@ -121,9 +130,9 @@ pub fn text(xml: &str, tag: &str) -> Option<String> { /// iterating over repeated elements like `<playlist>` blocks in /// `paramount`. /// -/// Self-closing elements return the same offset for body_end as the -/// element_end (i.e. `element_end - element_start` includes only the -/// `<tag .../>` text). +/// For self-closing elements, `element_end` points just past `/>` and +/// there is no separate body range (`element_end - element_start` +/// spans only the `<tag .../>` text). pub fn find_element(xml: &str, tag: &str, from: usize) -> Option<(usize, usize)> { let bytes = xml.as_bytes(); let mut i = from; @@ -190,8 +199,9 @@ pub fn find_element(xml: &str, tag: &str, from: usize) -> Option<(usize, usize)> /// The character after the tag name must not be a name-continuation /// (so `<player>` doesn't match `<play>`). fn matches_tag_name_at(bytes: &[u8], start: usize, tag: &str) -> bool { - let tag_lower = tag.to_ascii_lowercase(); - let tag_bytes = tag_lower.as_bytes(); + // Compare case-insensitively without allocating a lowercased copy + // of `tag` on every call (hot path: once per `<`/`</`). + let tag_bytes = tag.as_bytes(); // Skip optional `prefix:` (one or more name chars + `:`). let mut name_start = start; let mut scan = start; @@ -204,7 +214,7 @@ fn matches_tag_name_at(bytes: &[u8], start: usize, tag: &str) -> bool { if name_start + tag_bytes.len() > bytes.len() { return false; } - if !slice_eq_ignore_case(&bytes[name_start..name_start + tag_bytes.len()], tag_bytes) { + if !bytes[name_start..name_start + tag_bytes.len()].eq_ignore_ascii_case(tag_bytes) { return false; } // Boundary: char after the tag name must be `>`, `/`, whitespace. @@ -449,4 +459,28 @@ mod tests { let (s, e) = find_element(xml, "item", 0).unwrap(); assert_eq!(&xml[s..e], r#"<ns:item id="1" />"#); } + + #[test] + fn text_multibyte_before_self_close_does_not_panic() { + // A multi-byte UTF-8 char ending right before the `/>` used to + // panic on a non-char-boundary str slice in `text()`. The + // byte-level self-closing check must handle it cleanly. + // 'é' (0xC3 0xA9) directly precedes the `/>`. + assert_eq!(text("<x>é</x>", "x"), Some("é".into())); + // Self-closing form with a multi-byte char in an attr value. + assert_eq!(text(r#"<x a="é"/>"#, "x"), Some("".into())); + assert_eq!(text("<x>日本語</x>", "x"), Some("日本語".into())); + } + + #[test] + fn attr_not_matched_inside_quoted_value() { + // `name` appears only inside another attribute's quoted value; + // it must NOT be returned as a real attribute. + assert_eq!(attr(r#"<x y="name='inner'"/>"#, "name"), None); + // A real `name` attribute after a decoy value still resolves. + assert_eq!( + attr(r#"<x y="name='inner'" name="real"/>"#, "name"), + Some("real".into()) + ); + } } diff --git a/src/lib.rs b/src/lib.rs index 1254359..138409e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,16 +16,27 @@ //! for title in &disc.titles { //! println!("{} -- {} streams", title.duration_display(), title.streams.len()); //! } +//! ``` //! -//! // Stream via PES pipeline +//! Muxing to an output container runs through the PES pipeline. A live +//! `disc://` cannot be opened via [`input`] — it returns +//! [`Error::DiscUrlNotDirect`] by design (use `Drive` + `Disc::scan` + +//! `DiscStream::new` directly for a live drive). Any file-backed source +//! (`iso://`, `m2ts://`) opens through [`input`]: +//! +//! ```no_run +//! # fn run() -> std::io::Result<()> { //! let opts = libfreemkv::InputOptions::default(); -//! let mut input = libfreemkv::input("disc://", &opts).unwrap(); +//! let mut input = libfreemkv::input("iso://disc.iso", &opts)?; //! let title = input.info().clone(); -//! let mut output = libfreemkv::output("mkv://Movie.mkv", &title).unwrap(); -//! while let Ok(Some(frame)) = input.read() { -//! output.write(&frame).unwrap(); +//! let mut output = libfreemkv::output("mkv://Movie.mkv", &title)?; +//! // Propagate read errors instead of silently stopping on the first one. +//! while let Some(frame) = input.read()? { +//! output.write(&frame)?; //! } -//! output.finish().unwrap(); +//! output.finish()?; +//! # Ok(()) +//! # } //! ``` //! //! # Architecture @@ -70,6 +81,8 @@ //! | E5xxx | I/O errors | //! | E6xxx | Disc format errors | //! | E7xxx | AACS errors | +//! | E8xxx | Keydb errors (fetch, parse, load) | +//! | E9xxx | Stream / mux errors (URL, PES, pipeline) | pub mod aacs; pub(crate) mod clpi; @@ -145,7 +158,7 @@ pub use io::pipeline::{ }; // ─── Drive events (low-level callbacks) ───────────────────────────────────── -pub use event::{Event, EventKind}; +pub use event::{BatchSizeReason, Event, EventKind}; pub use identity::DriveId; pub use profile::DriveProfile; // Platform trait is pub(crate) — callers use Drive, not Platform directly. @@ -184,7 +197,7 @@ pub use keysource::{DiscInputs, KeySource}; // // - `DiscStream` — physical drive or ISO (any `SectorSource`). Read-only. // - `MkvStream` — Matroska container. Read on `open()`, write on `create()`. -// - `M2tsStream` — Blu-ray Transport Stream. Read on `open()`, write on `create()`. +// - `M2tsStream` — Blu-ray Transport Stream. Write-only sink (`create()`). // - `NetworkStream` — TCP. Read on `listen()`, write on `connect()`. // - `NullStream` — write-only black-hole sink. Useful for benchmarks. // - `StdioStream` — pipe to/from stdin/stdout. Read or write. @@ -204,6 +217,7 @@ pub use mux::MkvStream; pub use mux::NetworkStream; pub use mux::NullStream; pub use mux::StdioStream; +pub use mux::WriteSeek; pub use mux::{InputOptions, StreamUrl, input, output, parse_url}; // ─── Lower-level surfaces ─────────────────────────────────────────────────── diff --git a/src/mpls.rs b/src/mpls.rs index d47881f..fa28e64 100644 --- a/src/mpls.rs +++ b/src/mpls.rs @@ -10,9 +10,10 @@ use crate::error::{Error, Result}; /// Parsed MPLS playlist. #[derive(Debug)] -#[allow(dead_code)] -pub struct Playlist { - /// MPLS version (e.g. "0200" or "0300") +pub(crate) struct Playlist { + /// MPLS version (e.g. "0200" or "0300"). Parsed for completeness; + /// no production reader yet. + #[allow(dead_code)] pub version: String, /// Play items in playback order pub play_items: Vec<PlayItem>, @@ -24,11 +25,15 @@ pub struct Playlist { /// A playlist mark entry from the PlayListMark section. #[derive(Debug, Clone)] -#[allow(dead_code)] -pub struct PlaylistMark { - /// Mark type: 1 = chapter entry mark +pub(crate) struct PlaylistMark { + /// PlayListMark mark_type (BD-ROM PlayListMark spec): + /// 0 = reserved, 1 = entry mark (chapter), 2 = link point. + /// Chapter filters should test `== 1`, not `<= 1`. pub mark_type: u8, - /// Which play item this mark belongs to + /// Which play item this mark belongs to. Carries the per-PlayItem + /// timebase needed to place a mark in a multi-PlayItem playlist; + /// the chapter builder does not consume it yet. + #[allow(dead_code)] pub play_item_ref: u16, /// Timestamp in 45kHz PTS ticks pub timestamp: u32, @@ -36,22 +41,25 @@ pub struct PlaylistMark { /// A play item — one clip reference with in/out times. #[derive(Debug)] -#[allow(dead_code)] -pub struct PlayItem { +pub(crate) struct PlayItem { /// Clip filename without extension (e.g. "00001") pub clip_id: String, /// In-time in 45kHz ticks pub in_time: u32, /// Out-time in 45kHz ticks pub out_time: u32, - /// Connection condition (1=seamless, 5/6=non-seamless) + /// Connection condition (1=seamless, 5/6=non-seamless). Parsed for + /// completeness; no production reader yet. + #[allow(dead_code)] pub connection_condition: u8, } /// A stream entry from the STN table. #[derive(Debug, Clone)] pub struct StreamEntry { - /// Stream category: 1=video, 2=audio, 3=PG subtitle, 4=IG, 5=secondary audio, 6=secondary video, 7=DV EL + /// Stream category: 1=video, 2=audio, 3=PG subtitle, 5=secondary audio, + /// 6=secondary video, 7=DV EL. IG (4) is consumed during parsing to keep + /// the STN cursor aligned but is never retained as a StreamEntry. pub stream_type: u8, /// MPEG-TS PID pub pid: u16, @@ -76,6 +84,14 @@ pub struct StreamEntry { } /// Parse an MPLS file from raw bytes. +/// +/// `data` is the raw contents of a `BDMV/PLAYLIST/*.mpls` file. Returns +/// [`Error::MplsParse`] on malformed or truncated input. +/// +/// Note: [`Playlist::streams`] is extracted ONLY from the first play +/// item's STN table. Multi-item playlists whose later items carry a +/// different codec/track set are not fully represented by `streams`; +/// callers selecting tracks for mux should account for this. pub fn parse(data: &[u8]) -> Result<Playlist> { if data.len() < 40 { return Err(Error::MplsParse); @@ -95,7 +111,10 @@ pub fn parse(data: &[u8]) -> Result<Playlist> { let pl = &data[playlist_start..]; let num_play_items = u16::from_be_bytes([pl[6], pl[7]]) as usize; - let mut play_items = Vec::with_capacity(num_play_items); + // num_play_items is an untrusted u16 (max 65535); cap the pre-allocation + // so a truncated/fuzz input can't force a large reservation that the + // bounds-checked loop never fills. 256 covers any realistic playlist. + let mut play_items = Vec::with_capacity(num_play_items.min(256)); let mut streams = Vec::new(); let mut pos = 10; @@ -198,8 +217,12 @@ pub fn parse(data: &[u8]) -> Result<Playlist> { entry.stream_type = 6; entry.secondary = true; streams.push(entry); - // Skip extra ref bytes (audio refs + PG refs) - if next + 2 < item.len() { + // Skip extra ref bytes (audio refs + PG refs). + // Use `next < item.len()` to match the sibling secondary + // blocks; the inner `after_arefs < item.len()` re-guards + // the second read, so the stricter `+2` only mis-aligned + // spos when the aref count sits in the last 1-2 bytes. + if next < item.len() { let n_arefs = item[next] as usize; let after_arefs = next + 2 + n_arefs + (n_arefs % 2); if after_arefs < item.len() { @@ -256,16 +279,21 @@ pub fn parse(data: &[u8]) -> Result<Playlist> { // Parse PlayListMark section let mut marks = Vec::new(); - if mark_start > 0 && mark_start + 4 <= data.len() { + // The first real read is num_marks at ms[4..6], so the section needs + // at least 6 bytes (length(4) + num_marks(2)). + if mark_start > 0 && mark_start + 6 <= data.len() { let ms = &data[mark_start..]; - if ms.len() >= 6 { + { let num_marks = u16::from_be_bytes([ms[4], ms[5]]) as usize; let mut mpos = 6; for _ in 0..num_marks { if mpos + 14 > ms.len() { break; } - let mark_type = ms[mpos]; + // PlayListMark entry: reserved(1) + mark_type(1) + + // ref_to_PlayItem_id(2) + mark_time_stamp(4) + + // entry_ES_PID(2) + duration(4). mark_type is at +1, not +0. + let mark_type = ms[mpos + 1]; let play_item_ref = u16::from_be_bytes([ms[mpos + 2], ms[mpos + 3]]); let timestamp = u32::from_be_bytes([ms[mpos + 4], ms[mpos + 5], ms[mpos + 6], ms[mpos + 7]]); @@ -314,7 +342,10 @@ fn parse_stream_entry(item: &[u8], pos: usize, stream_type: u8) -> Option<(Strea 0x03 | 0x04 => 3, _ => 0, }; - let pid = if pid_off != 0 && pos + pid_off + 2 <= item.len() { + // Bound the PID read by the entry's declared end (se_end), not just by + // item.len(): a short se_len must not let us read PID bytes out of the + // following stream_attributes region. + let pid = if pid_off != 0 && pos + pid_off + 2 <= se_end { u16::from_be_bytes([item[pos + pid_off], item[pos + pid_off + 1]]) } else { 0 @@ -370,8 +401,10 @@ fn parse_stream_entry(item: &[u8], pos: usize, stream_type: u8) -> Option<(Strea } } } - 3 | 4 => { - // PG/IG: coding_type(1) + language(3) + 3 => { + // PG: coding_type(1) + language(3). + // IG (type 4) is parsed only to advance spos and is then + // discarded by the caller, so it deliberately has no arm here. if sa.len() >= 4 { language = String::from_utf8_lossy(&sa[1..4]).to_string(); } @@ -552,11 +585,11 @@ mod tests { buf.extend_from_slice(&(mark_section_len as u32).to_be_bytes()); buf.extend_from_slice(&(marks.len() as u16).to_be_bytes()); for m in marks { - buf.push(m.mark_type); // [0] mark_type - buf.push(0); // [1] reserved + buf.push(0); // [0] reserved + buf.push(m.mark_type); // [1] mark_type buf.extend_from_slice(&m.play_item_ref.to_be_bytes()); // [2-3] play_item_ref buf.extend_from_slice(&m.timestamp.to_be_bytes()); // [4-7] timestamp - buf.extend_from_slice(&[0u8; 6]); // [8-13] padding (entry_ES_PID + duration + mark_data) + buf.extend_from_slice(&[0u8; 6]); // [8-13] entry_ES_PID(2) + duration(4) } buf @@ -783,6 +816,40 @@ mod tests { assert_eq!(playlist.streams[2].pid, 0x1B00); } + #[test] + fn parse_secondary_video_then_dv_alignment() { + // Regression: the secondary-video ref-skip must use the same + // `next < item.len()` guard as the sibling secondary blocks so spos + // stays aligned for a following stream (here a Dolby Vision EL). + let video = build_stream_entry_video(0x1011, 0x24, 8, 1, Some(0x12)); + + // Secondary video with audio-ref + PG-ref blocks present. + let mut sec_video_with_refs = build_stream_entry_video(0x1B00, 0x1B, 4, 1, None); + sec_video_with_refs.push(0); // n_arefs = 0 + sec_video_with_refs.push(0); // reserved + sec_video_with_refs.push(0); // n_prefs = 0 + sec_video_with_refs.push(0); // reserved + + // Dolby Vision enhancement layer immediately after. + let dv_el = build_stream_entry_video(0x1015, 0x24, 8, 1, Some(0x12)); + + let data = build_mpls( + &[(b"00001", 1, 0, 9000000)], + (1, 0, 0, 0, 0, 1, 0, 1), // 1 video, 1 sec_video, 1 dv + &[video, sec_video_with_refs, dv_el], + ); + + let playlist = parse(&data).expect("should parse"); + assert_eq!(playlist.streams.len(), 3); + // Secondary video + assert_eq!(playlist.streams[1].stream_type, 6); + assert_eq!(playlist.streams[1].pid, 0x1B00); + // DV EL parsed at the correct offset → correct PID and type 7. + assert_eq!(playlist.streams[2].stream_type, 7); + assert_eq!(playlist.streams[2].pid, 0x1015); + assert!(playlist.streams[2].secondary); + } + #[test] fn parse_marks_chapter_entries() { let video = build_stream_entry_video(0x1011, 0x1B, 6, 1, None); @@ -874,6 +941,52 @@ mod tests { assert!((ch2_secs - 200.0).abs() < 0.001); } + #[test] + fn mark_type_read_from_correct_offset() { + // Regression for the mark_type off-by-one: each PlayListMark entry is + // reserved(1) + mark_type(1) + .... The parser must read byte[1], not + // byte[0]. build_mpls_with_marks writes reserved=0 at byte[0] and the + // mark_type at byte[1], so a parser that read byte[0] would see 0 for + // every mark. Use distinct non-zero, non-1 types to make the offset + // error unmistakable. + let video = build_stream_entry_video(0x1011, 0x1B, 6, 1, None); + let marks = vec![ + TestMark { + mark_type: 1, // entry mark (chapter) + play_item_ref: 0, + timestamp: 90000, + }, + TestMark { + mark_type: 2, // link point (not a chapter) + play_item_ref: 0, + timestamp: 180000, + }, + TestMark { + mark_type: 3, // arbitrary other type + play_item_ref: 0, + timestamp: 270000, + }, + ]; + + let data = build_mpls_with_marks( + &[(b"00001", 1, 0, 9000000)], + (1, 0, 0, 0, 0, 0, 0, 0), + &[video], + &marks, + ); + + let playlist = parse(&data).expect("should parse marks"); + assert_eq!(playlist.marks.len(), 3); + // If the parser read the reserved byte (byte[0] == 0) these would all + // be 0; reading byte[1] yields the real types. + assert_eq!(playlist.marks[0].mark_type, 1); + assert_eq!(playlist.marks[1].mark_type, 2); + assert_eq!(playlist.marks[2].mark_type, 3); + // Only the type-1 mark is a chapter under the corrected convention. + let chapters = playlist.marks.iter().filter(|m| m.mark_type == 1).count(); + assert_eq!(chapters, 1); + } + #[test] fn parse_no_marks_section() { // When mark_start is 0, no marks should be returned diff --git a/src/mux/codec/ac3.rs b/src/mux/codec/ac3.rs index 0952f35..bd86c2a 100644 --- a/src/mux/codec/ac3.rs +++ b/src/mux/codec/ac3.rs @@ -7,10 +7,22 @@ use super::{CodecParser, Frame, PesPacket, pts_to_ns}; /// Sample rates indexed by fscod (0=48kHz, 1=44.1kHz, 2=32kHz). fscod=3 is -/// reserved in AC-3 and signals "fscod2" (reduced rates) in E-AC-3; we treat -/// the base rate as 48 kHz in that case for duration purposes. +/// reserved in AC-3; in E-AC-3 it signals "fscod2" (reduced rates: 24/22.05/16 +/// kHz, selected by byte-4 bits [5:4]). `frame_sample_rate` decodes fscod2 in +/// the E-AC-3 case; this table's index-3 entry (48 kHz) is only the fallback +/// when the header is too short to read fscod2. const SAMPLE_RATES: [u32; 4] = [48_000, 44_100, 32_000, 48_000]; +/// E-AC-3 reduced sample rates indexed by fscod2 (byte-4 bits [5:4]), used when +/// fscod==3. Index 3 is reserved; we fall back to 48 kHz for it. +const EAC3_REDUCED_RATES: [u32; 4] = [24_000, 22_050, 16_000, 48_000]; + +/// Minimum byte length of a valid (E-)AC-3 frame. A real E-AC-3 frame must carry +/// at least the syncword (2) + BSI header (~4) before any audio. `eac3_frame_size` +/// returns `(frmsiz + 1) * 2`, so frmsiz=0/1 yield 2/4-byte "frames" that are +/// sub-header junk; rejecting anything below this guards against emitting them. +const MIN_FRAME_BYTES: usize = 6; + /// AC-3 (legacy) always carries 6 audio blocks × 256 samples = 1536 samples. const AC3_SAMPLES_PER_FRAME: u32 = 1536; @@ -89,8 +101,9 @@ impl CodecParser for Ac3Parser { ac3_frame_size(remaining) }; - if frame_size == 0 || frame_size > 8192 { - // Invalid frame size — skip this sync word + if !(MIN_FRAME_BYTES..=8192).contains(&frame_size) { + // Invalid/sub-header frame size (e.g. an E-AC-3 frmsiz of 0/1 + // sizing to a 2/4-byte fragment) — skip this sync word. pos = start + 2; continue; } @@ -112,13 +125,15 @@ impl CodecParser for Ac3Parser { } // Keep unconsumed data for the next call. `pos` is the start of the - // unconsumed region: either a partial frame that straddles this PES - // boundary — which, by construction, begins at a syncword (every byte - // before `pos` was emitted as a frame or skipped as pre-sync junk) — or - // trailing bytes too short to size/complete a frame. Carry from `pos`, - // NOT from the next syncword: discarding bytes between `pos` and the - // next sync would drop the partial frame we are deliberately keeping - // across the boundary. + // last unprocessed search region. On the `start + frame_size > len` + // break it sits exactly at the straddling frame's syncword; on the + // `remaining.len() < 6` break it is the value from the top of that + // iteration, with the syncword possibly sitting after some pre-sync + // junk — so the re-scan below (from `pos`, NOT a recomputed sync) is + // required to locate the carry-over syncword. Carry from `pos`, NOT + // from the next syncword: discarding bytes between `pos` and the next + // sync would drop the partial frame we are deliberately keeping across + // the boundary. let keep_from = if pos < data.len() { // A syncword at/after `pos` marks the carry-over start (anything // before it is junk with no sync). With no full sync, retain the @@ -139,6 +154,11 @@ impl CodecParser for Ac3Parser { // No frame could be parsed out of a buffer this large — this is // not valid AC-3 here. Drop it and resync on the next PES rather // than grow without bound on pathological input. + tracing::debug!( + target: "mux", + "ac3: carry-over buffer exceeded {} bytes without a frame; dropping and resyncing", + MAX_AC3_BUF + ); self.buf.clear(); } else { self.buf = tail.to_vec(); @@ -174,7 +194,7 @@ impl CodecParser for Ac3Parser { } else { ac3_frame_size(frame) }; - if frame_size == 0 || frame_size > 8192 || off + frame_size > buf.len() { + if !(MIN_FRAME_BYTES..=8192).contains(&frame_size) || off + frame_size > buf.len() { return Vec::new(); } let duration_ns = frame_duration_ns(frame, bsid); @@ -212,12 +232,20 @@ fn eac3_samples_per_frame(data: &[u8]) -> u32 { numblks * 256 } -/// Sample rate (Hz) of an AC-3/E-AC-3 frame from its fscod field (byte 4 bits 7-6). -fn frame_sample_rate(data: &[u8]) -> u32 { +/// Sample rate (Hz) of an AC-3/E-AC-3 frame from its fscod field (byte 4 bits +/// 7-6). For E-AC-3 (`bsid >= 11`) an fscod of 3 selects a reduced rate via +/// fscod2 (byte 4 bits [5:4]); decoding it keeps the frame duration correct +/// instead of mistiming reduced-rate frames at 48 kHz (A/V drift). +fn frame_sample_rate(data: &[u8], bsid: u8) -> u32 { if data.len() < 5 { return SAMPLE_RATES[0]; } - SAMPLE_RATES[((data[4] >> 6) & 0x03) as usize] + let fscod = (data[4] >> 6) & 0x03; + if fscod == 0x03 && bsid >= 11 { + let fscod2 = (data[4] >> 4) & 0x03; + return EAC3_REDUCED_RATES[fscod2 as usize]; + } + SAMPLE_RATES[fscod as usize] } /// Duration of one AC-3/E-AC-3 frame in nanoseconds: samples_per_frame / @@ -228,7 +256,7 @@ fn frame_duration_ns(data: &[u8], bsid: u8) -> u64 { } else { AC3_SAMPLES_PER_FRAME } as u64; - let rate = frame_sample_rate(data) as u64; + let rate = frame_sample_rate(data, bsid) as u64; // samples / rate seconds → ns, rounded to nearest. (samples * 1_000_000_000 + rate / 2) / rate } @@ -240,7 +268,7 @@ fn find_ac3_sync(data: &[u8]) -> Option<usize> { /// Extract bsid from an AC-3/E-AC-3 frame starting at the syncword. /// bsid is at byte 5, bits 7..3. -pub fn get_bsid(data: &[u8]) -> u8 { +fn get_bsid(data: &[u8]) -> u8 { if data.len() < 6 { return 0; } @@ -256,8 +284,11 @@ fn eac3_frame_size(data: &[u8]) -> usize { (frmsiz + 1) * 2 } -/// Calculate AC-3 frame size in bytes from fscod and frmsizecod. -fn ac3_frame_size(data: &[u8]) -> usize { +/// Calculate AC-3 frame size in bytes from fscod and frmsizecod. Returns 0 for +/// an unmappable header (reserved fscod==3, or frmsizecod out of table range). +/// `pub(crate)` so the TrueHD parser can reuse it when skipping interleaved AC-3 +/// frames instead of duplicating the size table. +pub(crate) fn ac3_frame_size(data: &[u8]) -> usize { if data.len() < 5 { return 0; } @@ -439,7 +470,7 @@ mod tests { #[test] fn buffer_stays_bounded_across_many_garbage_pes() { - // Finding 14: the carry-over buffer must never grow without bound. Feed + // The carry-over buffer must never grow without bound. Feed // many large PES packets that contain no usable frame and assert the // retained buffer stays tiny — carry-from-`pos` drops all pre-sync junk, // and a never-completing frame is bounded by the 8192-byte frame cap and @@ -573,6 +604,44 @@ mod tests { assert_eq!(frame_duration_ns(&frame, bsid), 32_000_000); } + #[test] + fn eac3_subheader_sized_frame_is_rejected() { + // An E-AC-3 sync with frmsiz=0 sizes to a 2-byte "frame"; frmsiz=1 to + // 4 bytes. Both are sub-header junk that must NOT be emitted as audio. + // bsid must be >= 11 for the E-AC-3 sizing path. Byte 5 bits 7..3 = bsid. + let mut parser = Ac3Parser::new(); + // Build an E-AC-3 sync: 0x0B 0x77, frmsiz=0 (bytes 2-3 low bits = 0), + // bsid=16 (>=11) at byte 5. Pad to a few bytes so find_ac3_sync + sizing + // run. eac3_frame_size = (0 + 1) * 2 = 2 < MIN_FRAME_BYTES. + let mut data = vec![0x0B, 0x77, 0x00, 0x00, 0x00, 16 << 3, 0x00, 0x00]; + // Append a real AC-3 frame after the junk so we can confirm the parser + // resyncs past the junk and still emits the valid frame. + let good = make_ac3_frame(0, 2); + data.extend_from_slice(&good); + let pes = PesPacket { + pid: 0, + pts: Some(90000), + dts: None, + data, + }; + let frames = parser.parse(&pes); + assert_eq!(frames.len(), 1, "only the real AC-3 frame is emitted"); + assert_eq!(frames[0].data.len(), 160); + } + + #[test] + fn eac3_fscod2_reduced_rate_duration() { + // E-AC-3 with fscod==3 (reduced rate) and fscod2==0 → 24 kHz, not 48. + // bsid>=11 selects the E-AC-3 path. When fscod==3 the block count is + // fixed at 6 → 1536 samples. Byte 4 layout: fscod(2)|fscod2(2)|... + // fscod=3 (0b11), fscod2=0 (0b00) → byte4 = 0b1100_0000 = 0xC0. + let data = [0x0B, 0x77, 0x00, 0x00, 0xC0, 16 << 3]; + let bsid = get_bsid(&data); + assert!(bsid >= 11, "test frame is E-AC-3"); + // 1536 samples / 24000 Hz = 64 ms. + assert_eq!(frame_duration_ns(&data, bsid), 64_000_000); + } + #[test] fn ac3_frame_size_table() { // fscod=0 (48kHz), frmsizecod=0: 64 words = 128 bytes diff --git a/src/mux/codec/dts.rs b/src/mux/codec/dts.rs index c1322f3..a6ed62c 100644 --- a/src/mux/codec/dts.rs +++ b/src/mux/codec/dts.rs @@ -15,6 +15,11 @@ const DTS_CORE_SYNC: [u8; 4] = [0x7F, 0xFE, 0x80, 0x01]; #[cfg(test)] const DTS_HD_EXT_SYNC: [u8; 4] = [0x64, 0x58, 0x20, 0x25]; +/// DTS / DTS-HD elementary-stream parser. Buffers DTS across PES boundaries so +/// a core frame plus all of its trailing DTS-HD extension substreams are +/// emitted together as one access unit, delimited by the next valid core sync. +/// This preserves the lossless extension data instead of downgrading to lossy +/// core (the Dunkirk / Fight Club lossy-core bug). pub struct DtsParser { buf: Vec<u8>, /// PTS of the access unit currently being assembled in `buf` (the unit @@ -92,11 +97,22 @@ impl DtsParser { /// this without a clean boundary we resync rather than stall or balloon. const MAX_AU_BYTES: usize = 65536; -/// Minimum plausible DTS core frame size. The core header alone is ~10-14 -/// bytes; a decoded `core_size` below this means we matched a false/corrupt -/// core sync (the 14-bit `fsize` field decoded to a tiny value) rather than a -/// real frame, so we resync instead of emitting a junk access unit. -const MIN_CORE_FRAME_BYTES: usize = 10; +/// Number of leading bytes that must be buffered before the core `fsize` field +/// (bytes 5-7) can be decoded. This is a HEADER-LAYOUT minimum — "enough bytes +/// to read the size field" — and is deliberately distinct from +/// `MIN_CORE_FRAME_BYTES` (the decoded-frame-size validity floor). They must not +/// be conflated: this one gates buffer reads of the header, the other rejects +/// implausible decoded sizes. +const CORE_HEADER_MIN_BYTES: usize = 10; + +/// Minimum plausible decoded DTS core frame size, per ETSI TS 102 114: the +/// on-wire FSIZE floor is 95, so a real core frame is at least 96 bytes. A +/// decoded `core_size` below this means we matched a false/corrupt core sync +/// (a lucky 0x7FFE8001 in extension-substream payload whose 14-bit `fsize` +/// decoded to a tiny value) rather than a real frame, so we resync instead of +/// closing an access unit at a junk boundary and dropping the DTS-HD extension +/// tail. +const MIN_CORE_FRAME_BYTES: usize = 96; /// Sentinel for "no valid PTS base captured yet". Real PTS-in-ns values are /// non-negative (derived from the unsigned 90 kHz PES timestamp), so a negative @@ -156,25 +172,30 @@ impl CodecParser for DtsParser { }; if start > 0 { self.drain_front(start); - if find_sync(&self.buf, &DTS_CORE_SYNC) != Some(0) { - // Shouldn't happen, but never loop forever. - break; - } + // The sync `find_sync` located at offset `start` is now at + // offset 0 by construction, so a re-scan would be a redundant + // O(buf_len) walk per iteration; assert the invariant instead. + debug_assert_eq!( + find_sync(&self.buf, &DTS_CORE_SYNC), + Some(0), + "drain_front(start) must leave the core sync at offset 0" + ); } // Need the core header to size the core frame. - if self.buf.len() < 10 { + if self.buf.len() < CORE_HEADER_MIN_BYTES { break; } let core_size = dts_core_frame_size(&self.buf); // `dts_core_frame_size` returns a 14-bit `fsize + 1`, so it is // always in [1, 16384]; the bare `== 0` / `> MAX_AU_BYTES` checks - // can never fire. A real DTS core header is at least ~10-14 bytes, - // so any decoded size below that came from a false/corrupt sync. - // Reject it (drain the 4 syncword bytes and resync) instead of - // letting a tiny bogus size close the current access unit at a junk - // boundary and drop the trailing extension substreams. The - // `> MAX_AU_BYTES` upper bound is kept as a harmless guard. + // can never fire. A real DTS core frame is at least + // MIN_CORE_FRAME_BYTES (96, the ETSI spec floor), so any decoded + // size below that came from a false/corrupt sync. Reject it (drain + // the 4 syncword bytes and resync) instead of letting a tiny bogus + // size close the current access unit at a junk boundary and drop the + // trailing extension substreams. The `> MAX_AU_BYTES` upper bound is + // kept as a harmless guard. if !(MIN_CORE_FRAME_BYTES..=MAX_AU_BYTES).contains(&core_size) { // Bogus core sync — skip past it and resync. self.drain_front(4); @@ -257,7 +278,8 @@ impl CodecParser for DtsParser { // core + its extension substreams, which had no following core sync to // close it during streaming). Require a complete core frame; drop a // bare partial sync tail. - if find_sync(&self.buf, &DTS_CORE_SYNC) != Some(0) || self.buf.len() < 10 { + if find_sync(&self.buf, &DTS_CORE_SYNC) != Some(0) || self.buf.len() < CORE_HEADER_MIN_BYTES + { self.buf.clear(); return Vec::new(); } @@ -315,7 +337,7 @@ fn next_core_boundary(buf: &[u8], core_size: usize) -> NextCore { while let Some(rel) = find_sync(&buf[from..], &DTS_CORE_SYNC) { let pos = from + rel; // Need the candidate's core header to judge it. - if buf.len() - pos < 10 { + if buf.len() - pos < CORE_HEADER_MIN_BYTES { return NextCore::NeedMore; } let sz = dts_core_frame_size(&buf[pos..]); @@ -328,10 +350,17 @@ fn next_core_boundary(buf: &[u8], core_size: usize) -> NextCore { NextCore::None } -/// DTS core frame size from header bits. -/// fsize is at bits 46-59 (14 bits) of the header: bytes 5-7. +/// DTS core frame size from header bits. `fsize` is the 14-bit field at bits +/// 46-59 of the header (bytes 5-7). On the wire `fsize` is the frame length +/// minus one, so this returns `fsize + 1`, i.e. the core frame length in bytes +/// (range 1..=16384). Callers treat the result as the actual byte length and +/// the MIN..=MAX range checks assume so. +/// +/// Returns `0` when `data` is shorter than `CORE_HEADER_MIN_BYTES` — every call +/// site rejects that via the minimum-frame lower bound, so a `0` is never +/// mistaken for a valid tiny frame. fn dts_core_frame_size(data: &[u8]) -> usize { - if data.len() < 10 { + if data.len() < CORE_HEADER_MIN_BYTES { return 0; } // fsize field: 14 bits starting at bit 46 @@ -629,6 +658,43 @@ mod tests { assert_eq!(tail[0].data.len(), 640); } + #[test] + fn sub_spec_core_size_is_rejected_as_false_sync() { + // A core sync whose decoded fsize+1 lands in [CORE_HEADER_MIN_BYTES, + // MIN_CORE_FRAME_BYTES) — i.e. a "frame" smaller than the 96-byte ETSI + // spec minimum — is a false sync inside extension payload and must NOT + // close an access unit. Pick a decoded size of 64 (well inside the old + // 10..96 false-positive window the raised floor now rejects). + let false_size = 64usize; + assert!( + (CORE_HEADER_MIN_BYTES..MIN_CORE_FRAME_BYTES).contains(&false_size), + "test fixture must sit in the widened reject window" + ); + let mut parser = DtsParser::new(); + + // Frame 1: real core(512) + extension that embeds a sub-spec "core sync" + // whose fsize decodes to 64 bytes. + let mut ext = make_dts_ext(256); + let bogus = make_dts_core(false_size); // valid-looking sync, size 64 + ext[64..64 + bogus.len()].copy_from_slice(&bogus); + let mut frame1 = make_dts_core(512); + frame1.extend_from_slice(&ext); + + assert!( + parser.parse(&make_pes(frame1, Some(90000))).is_empty(), + "sub-spec core size must not close the AU" + ); + + // Real next core closes frame 1 as core + full extension. + let f = parser.parse(&make_pes(make_dts_core(640), Some(93000))); + assert_eq!(f.len(), 1); + assert_eq!( + f[0].data.len(), + 512 + 256, + "AU must not be split at the sub-spec false sync" + ); + } + #[test] fn forced_emit_does_not_corrupt_next_au_pts() { // When the buffer exceeds MAX_AU_BYTES with no next core sync, the diff --git a/src/mux/codec/dvdsub.rs b/src/mux/codec/dvdsub.rs index 34e487b..de3428f 100644 --- a/src/mux/codec/dvdsub.rs +++ b/src/mux/codec/dvdsub.rs @@ -2,10 +2,14 @@ //! //! DVD subtitles are carried in PS private stream 1 with sub-stream IDs 0x20-0x3F. //! A single subpicture unit (SPU — one displayed bitmap) may span multiple PES -//! packets: only the first PES carries a PTS, continuations carry PTS=0. The SPU -//! begins with a 2-byte big-endian `SPU_size` giving the total byte length of the -//! whole unit. We reassemble across PES boundaries into one Frame so large -//! subtitles aren't split/garbled, inheriting the head PES's PTS. +//! packets: only the first PES carries a PTS; continuation PES packets have no +//! PTS field (the PS demuxer leaves `pts` as `None`). The SPU begins with a +//! 2-byte big-endian `SPU_size` giving the total byte length of the whole unit. +//! We reassemble across PES boundaries into one Frame so large subtitles aren't +//! split/garbled, inheriting the head PES's PTS. The presence of a PTS — not +//! merely an open `pending` — is the authoritative SPU-boundary signal, so a +//! lost continuation or a corrupt SPU_size can't merge the next subtitle into +//! the stuck unit. //! //! For MKV: codec ID "S_VOBSUB". //! All frames are keyframes (each is a complete bitmap). @@ -35,7 +39,7 @@ impl DvdSubParser { /// Emit `pending` as a Frame if it is complete (or `force` at EOF), /// returning it and clearing the buffer. Returns None if nothing to emit. fn take_if_complete(&mut self, force: bool) -> Option<Frame> { - let (pts_ns, size, buf) = self.pending.as_ref()?; + let (_, size, buf) = self.pending.as_ref()?; if force || buf.len() >= *size { let (pts_ns, _, data) = self.pending.take().unwrap(); return Some(Frame { @@ -45,7 +49,6 @@ impl DvdSubParser { duration_ns: None, }); } - let _ = pts_ns; None } } @@ -58,32 +61,63 @@ impl CodecParser for DvdSubParser { let mut out = Vec::new(); - if self.pending.is_some() { - // Continuation of an in-progress SPU (PTS=0 on these). Append, - // bounded by MAX_SPU_BYTES. - if let Some((_, _, buf)) = self.pending.as_mut() { - let room = MAX_SPU_BYTES.saturating_sub(buf.len()); - let take = room.min(pes.data.len()); - buf.extend_from_slice(&pes.data[..take]); + // A PES carrying a real PTS is the START of a new SPU; continuations of + // an in-progress SPU carry no PTS (the PS demuxer leaves `pts` None when + // the PES has no PTS field — see the module doc). PTS is therefore the + // authoritative SPU-boundary signal, NOT merely `pending.is_some()`. + // + // Append-as-continuation ONLY when this PES has no PTS. When it has a + // PTS but a stale `pending` is still open (a lost continuation, or a + // corrupt/oversized declared SPU_size that real data never reaches), + // force-emit the stuck unit truncated and fall through to start a fresh + // SPU from this PES. Without this, one bad SPU_size would swallow every + // later subtitle until EOF — exactly the damaged-disc case we target. + if pes.pts.is_none() { + if self.pending.is_some() { + // Continuation: append, bounded by MAX_SPU_BYTES. + if let Some((_, _, buf)) = self.pending.as_mut() { + let room = MAX_SPU_BYTES.saturating_sub(buf.len()); + let take = room.min(pes.data.len()); + buf.extend_from_slice(&pes.data[..take]); + } + if let Some(frame) = self.take_if_complete(false) { + out.push(frame); + } + return out; } - if let Some(frame) = self.take_if_complete(false) { - out.push(frame); - } - return out; + // No pending and no PTS: nothing to attach this to. Pass it through + // as a lone frame (PTS unknown → 0) rather than drop it. + } else if let Some(frame) = self.take_if_complete(true) { + // New SPU starting while a previous one is still open → flush stale. + out.push(frame); } // Start of a new SPU. The first 2 bytes are the big-endian total size. let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0); let declared = if pes.data.len() >= 2 { - ((pes.data[0] as usize) << 8) | pes.data[1] as usize + // SPU_size includes the 2-byte header, so a declared size < 2 is + // always malformed; treat it like the too-short path (lone frame) + // rather than emit an immediate oversized unit. + let d = ((pes.data[0] as usize) << 8) | pes.data[1] as usize; + if d < 2 { + out.push(Frame { + pts_ns, + keyframe: true, + data: pes.data.clone(), + duration_ns: None, + }); + return out; + } + d } else { // Too short to carry SPU_size — pass through as a lone frame. - return vec![Frame { + out.push(Frame { pts_ns, keyframe: true, data: pes.data.clone(), duration_ns: None, - }]; + }); + return out; }; let mut buf = pes.data.clone(); @@ -114,6 +148,18 @@ impl CodecParser for DvdSubParser { /// /// Input: `[padding, Y, Cb, Cr]` (as stored in DVD IFO PGC data). /// Returns `[R, G, B]`. +/// +/// Range convention (deliberate): this uses the **full-range (JFIF) BT.601** +/// coefficients with no 16/235 luma scaling. DVD IFO palette YCbCr is nominally +/// studio-swing BT.601, so studio-swing math would be more colorimetrically +/// "correct" in isolation. But the output here is a VobSub `.idx` `palette:` +/// line, and the entire VobSub ecosystem (the original tooling, mkvtoolnix, +/// players that read the .idx palette) is built around this full-range formula — +/// it is the de-facto on-disk convention. Emitting studio-swing-scaled RGB here +/// would make freemkv's palettes inconsistent with every other tool and wrong in +/// players that assume the VobSub convention. We therefore intentionally keep +/// full-range; do NOT "fix" this to studio-swing without changing the consuming +/// side in lockstep. pub fn ycbcr_to_rgb(color: &[u8; 4]) -> [u8; 3] { let y = color[1] as f64; let cb = color[2] as f64; @@ -240,9 +286,10 @@ mod tests { let f = parser.parse(&make_pes(head.clone(), Some(90000))); assert!(f.is_empty(), "incomplete SPU should not emit yet"); - let f = parser.parse(&make_pes(cont1.clone(), Some(0))); + // Continuations carry NO PTS (None), per the PS demuxer. + let f = parser.parse(&make_pes(cont1.clone(), None)); assert!(f.is_empty(), "still incomplete"); - let frames = parser.parse(&make_pes(cont2.clone(), Some(0))); + let frames = parser.parse(&make_pes(cont2.clone(), None)); assert_eq!(frames.len(), 1, "completed SPU emits exactly one frame"); // Reassembled bytes = head + cont1 + cont2, in order. @@ -268,6 +315,67 @@ mod tests { assert_eq!(frames[0].pts_ns, 1_000_000_000); } + #[test] + fn real_pts_pes_force_emits_stale_pending_and_starts_new_spu() { + // A lost continuation leaves an incomplete pending SPU. The NEXT real + // subtitle arrives with its own PTS — it must force-emit the stuck unit + // (truncated) and begin a fresh SPU, not be appended as a continuation. + let mut parser = DvdSubParser::new(None); + + // SPU 1 declares 100 bytes but only 6 arrive; the continuation is lost. + let head1 = vec![0x00, 0x64, 0xDE, 0xAD, 0xBE, 0xEF]; + assert!( + parser + .parse(&make_pes(head1.clone(), Some(90000))) + .is_empty(), + "SPU 1 incomplete, held pending" + ); + + // SPU 2 arrives with a real PTS — declares 4 bytes, fully present. + let head2 = vec![0x00, 0x04, 0x11, 0x22]; + let frames = parser.parse(&make_pes(head2.clone(), Some(180000))); + // First the truncated stale SPU 1, then complete SPU 2. + assert_eq!(frames.len(), 2, "stale flushed + new emitted"); + assert_eq!(frames[0].data, head1, "stale SPU 1 emitted truncated"); + assert_eq!(frames[0].pts_ns, 1_000_000_000, "SPU 1 keeps its PTS"); + assert_eq!(frames[1].data, head2, "SPU 2 emitted fresh"); + assert_eq!(frames[1].pts_ns, 2_000_000_000, "SPU 2 keeps its own PTS"); + } + + #[test] + fn corrupt_oversized_size_recovers_on_next_real_pts() { + // A corrupt SPU_size that real data never reaches must not swallow every + // later subtitle. The next real-PTS PES resets pending and recovers the + // track. + let mut parser = DvdSubParser::new(None); + + // Declares 0xFFFF but only a few bytes ever arrive (corrupt size). + let bad = vec![0xFF, 0xFF, 0x01, 0x02, 0x03]; + assert!(parser.parse(&make_pes(bad.clone(), Some(90000))).is_empty()); + // A no-PTS stray continuation appends (still stuck under the bad size). + assert!(parser.parse(&make_pes(vec![0x04, 0x05], None)).is_empty()); + + // Next real subtitle (PTS present) recovers: stale flushed + new SPU. + let good = vec![0x00, 0x04, 0xAA, 0xBB]; + let frames = parser.parse(&make_pes(good.clone(), Some(270000))); + assert_eq!(frames.len(), 2, "track recovers, not swallowed to EOF"); + assert_eq!(frames[1].data, good); + assert_eq!(frames[1].pts_ns, 3_000_000_000); + } + + #[test] + fn declared_size_below_two_passes_through_as_lone_frame() { + // SPU_size includes its own 2-byte header, so a declared size < 2 is + // malformed. It must pass through as a lone frame, not emit an oversized + // unit or get stuck pending. + let mut parser = DvdSubParser::new(None); + let data = vec![0x00, 0x00, 0xAB, 0xCD]; // declared = 0 + let frames = parser.parse(&make_pes(data.clone(), Some(90000))); + assert_eq!(frames.len(), 1); + assert_eq!(frames[0].data, data, "passed through whole"); + assert!(parser.pending.is_none(), "no pending left open"); + } + // ── YCbCr → RGB conversion tests ────────────────────────────────────── #[test] diff --git a/src/mux/codec/h264.rs b/src/mux/codec/h264.rs index ef3a1ab..ea5621c 100644 --- a/src/mux/codec/h264.rs +++ b/src/mux/codec/h264.rs @@ -4,6 +4,7 @@ //! Detects keyframes (IDR slices). //! Each PES packet = one access unit = one frame. +use super::startcode::{find_start_code, skip_start_code}; use super::{CodecParser, Frame, PesPacket, pts_to_ns}; /// H.264 NAL unit types we care about. @@ -12,6 +13,9 @@ const NAL_SPS: u8 = 7; const NAL_PPS: u8 = 8; const NAL_AUD: u8 = 9; +/// H.264 (AVC) Annex B → MKV codec parser: extracts SPS/PPS for the avcC +/// codecPrivate, detects IDR keyframes, and converts each PES access unit into +/// length-prefixed NAL units. Implements [`CodecParser`]. pub struct H264Parser { // First-seen SPS/PPS seed the MKV codecPrivate (avcC) — the only out-of-band // copy the player gets. BD H.264 repeats the parameter sets at every IDR; @@ -32,6 +36,7 @@ impl Default for H264Parser { } impl H264Parser { + /// Create a fresh H.264 parser with no parameter sets captured yet. pub fn new() -> Self { Self { sps: None, @@ -56,7 +61,14 @@ fn handle_param_set(first: &mut Option<Vec<u8>>, nal: &[u8], frame_data: &mut Ve Some(f) if f.as_slice() == nal => {} // == codecPrivate → player has it Some(_) => { // Differs from codecPrivate → emit in-band so it wins at this AU. - frame_data.extend_from_slice(&(nal.len() as u32).to_be_bytes()); + // A NAL longer than u32::MAX cannot be length-prefixed in the + // 4-byte field; skip it rather than emit a truncated length over + // the full body (mis-framed NALU). Unreachable in practice — no + // real access unit is >4 GiB. + let Ok(len) = u32::try_from(nal.len()) else { + return; + }; + frame_data.extend_from_slice(&len.to_be_bytes()); frame_data.extend_from_slice(nal); } } @@ -78,7 +90,10 @@ impl CodecParser for H264Parser { // Annex B (start-code prefixed) NALUs to length-prefixed NALUs (MKV with // AVCDecoderConfigurationRecord expects a 4-byte length prefix per NAL). let mut keyframe = false; - let mut frame_data = Vec::new(); + // Pre-size: output is ~input bytes plus a few 4-byte NAL length prefixes. + // The unsized Vec growth chain otherwise reallocs several times per + // frame in the mux hot path (mirrors the HEVC parser). + let mut frame_data = Vec::with_capacity(pes.data.len() + 64); for nal in NalIterator::new(&pes.data) { let nal_type = nal[0] & 0x1F; @@ -88,13 +103,21 @@ impl CodecParser for H264Parser { // mid-title redefinition differs from the avcC copy. NAL_SPS => handle_param_set(&mut self.sps, nal, &mut frame_data), NAL_PPS => handle_param_set(&mut self.pps, nal, &mut frame_data), - // Access unit delimiters: drop. + // Access unit delimiters: drop. Intentional and spec-correct — + // Matroska H.264 frame data omits AUDs (the container delimits + // access units), so keeping them in-band is redundant. Mirrors + // the HEVC parser. NAL_AUD => {} _ => { if nal_type == NAL_SLICE_IDR { keyframe = true; } - let len = nal.len() as u32; + // A NAL longer than u32::MAX can't be length-prefixed in the + // 4-byte field; skip it rather than mis-frame the output. + // Unreachable in practice (no real AU is >4 GiB). + let Ok(len) = u32::try_from(nal.len()) else { + continue; + }; frame_data.extend_from_slice(&len.to_be_bytes()); frame_data.extend_from_slice(nal); } @@ -182,61 +205,42 @@ impl<'a> Iterator for NalIterator<'a> { type Item = &'a [u8]; fn next(&mut self) -> Option<&'a [u8]> { - if self.pos >= self.data.len() { - return None; - } + // Loop (not tail-recursion) over empty NALs: a crafted/garbled Annex B + // stream with many adjacent start codes (e.g. 00 00 01 00 00 01 ...) + // yields empty NALs back-to-back; recursing once per empty NAL would + // overflow the stack. `self.pos` advances to `nal_end` each iteration, + // so the loop always terminates. Mirrors the HEVC parser's while-scan. + loop { + if self.pos >= self.data.len() { + return None; + } - // Skip the start code at current position - let nal_start = skip_start_code(self.data, self.pos)?; + // Skip the start code at current position + let nal_start = skip_start_code(self.data, self.pos)?; - // Find next start code (or end of data) - let nal_end = find_start_code(self.data, nal_start).unwrap_or(self.data.len()); + // Find next start code (or end of data) + let nal_end = find_start_code(self.data, nal_start).unwrap_or(self.data.len()); - // Remove trailing zeros (part of next start code's zero prefix) - let mut end = nal_end; - while end > nal_start && self.data[end - 1] == 0x00 { - end -= 1; - } + // Strip the leading zeros of the following start code. For a + // conforming bitstream this is lossless: rbsp_trailing_bits() sets a + // stop-one bit, so the final byte of any RBSP is never 0x00 — the only + // trailing zeros here belong to the next 00 00 (00) 01 prefix, never to + // the NAL's RBSP payload. (Mirrors the HEVC parser.) + let mut end = nal_end; + while end > nal_start && self.data[end - 1] == 0x00 { + end -= 1; + } - self.pos = nal_end; + self.pos = nal_end; - if end > nal_start { - Some(&self.data[nal_start..end]) - } else { - self.next() + if end > nal_start { + return Some(&self.data[nal_start..end]); + } + // Empty NAL — continue scanning instead of recursing. } } } -/// Find the position of the next start code (00 00 01) at or after `from`. -/// -/// Backed by `memchr::memmem::find` for SIMD-accelerated bytestring -/// search. On AVX2-capable x86_64 this runs ~5–10× the byte-by-byte -/// scan that preceded it; on a 200 KB UHD HEVC frame the saving is -/// in the hundreds of microseconds per call. -pub fn find_start_code(data: &[u8], from: usize) -> Option<usize> { - if data.len() < from + 3 { - return None; - } - memchr::memmem::find(&data[from..], b"\x00\x00\x01").map(|rel| from + rel) -} - -/// Skip past the start code at position `pos`, returning the first byte after it. -pub fn skip_start_code(data: &[u8], pos: usize) -> Option<usize> { - if pos + 2 >= data.len() { - return None; - } - if data[pos] == 0x00 && data[pos + 1] == 0x00 { - if pos + 3 < data.len() && data[pos + 2] == 0x00 && data[pos + 3] == 0x01 { - return Some(pos + 4); // 4-byte start code - } - if data[pos + 2] == 0x01 { - return Some(pos + 3); // 3-byte start code - } - } - None -} - #[cfg(test)] mod tests { use super::*; @@ -251,39 +255,6 @@ mod tests { } } - // --- find_start_code tests --- - - #[test] - fn find_start_code_3byte() { - let data = [0x00, 0x00, 0x01, 0x65]; - assert_eq!(find_start_code(&data, 0), Some(0)); - } - - #[test] - fn find_start_code_4byte() { - let data = [0x00, 0x00, 0x00, 0x01, 0x65]; - // find_start_code looks for 00 00 01 pattern, which starts at offset 1 in a 4-byte start code - assert_eq!(find_start_code(&data, 0), Some(1)); - } - - #[test] - fn find_start_code_offset() { - let data = [0xFF, 0xFF, 0x00, 0x00, 0x01, 0x09]; - assert_eq!(find_start_code(&data, 0), Some(2)); - } - - #[test] - fn find_start_code_none() { - let data = [0x00, 0x00, 0x00, 0x00]; - assert_eq!(find_start_code(&data, 0), None); - } - - #[test] - fn find_start_code_too_short() { - let data = [0x00, 0x00]; - assert_eq!(find_start_code(&data, 0), None); - } - // --- parse SPS+PPS → codec_private --- #[test] @@ -590,6 +561,31 @@ mod tests { ); } + #[test] + fn many_empty_nals_do_not_overflow_stack() { + // Regression: NalIterator::next must iterate, not recurse, over empty + // NALs. A crafted Annex B stream of tens of thousands of adjacent start + // codes (each producing an empty NAL) would blow the stack under the old + // tail-recursive implementation. Iterating handles it in bounded stack. + let mut data = Vec::new(); + // 50_000 back-to-back 3-byte start codes → 50_000 empty NALs. + for _ in 0..50_000 { + data.extend_from_slice(&[0x00, 0x00, 0x01]); + } + // One real NAL at the end so the iterator yields something. + data.extend_from_slice(&[0x41, 0xAA, 0xBB]); + + let mut parser = H264Parser::new(); + let frames = parser.parse(&make_pes(data, Some(0))); + // Exactly one populated frame; the empty NALs are skipped without + // overflowing. + assert_eq!(frames.len(), 1); + let fd = &frames[0].data; + let len = u32::from_be_bytes([fd[0], fd[1], fd[2], fd[3]]) as usize; + assert_eq!(len, 3, "the single real NAL is length-prefixed"); + assert_eq!(fd[4], 0x41); + } + #[test] fn avcc_oversized_param_set_returns_none() { // A param set > 65535 bytes can't be length-encoded in avcC's 16-bit diff --git a/src/mux/codec/hevc.rs b/src/mux/codec/hevc.rs index 2785e6a..6f3d25e 100644 --- a/src/mux/codec/hevc.rs +++ b/src/mux/codec/hevc.rs @@ -4,7 +4,7 @@ //! Detects keyframes (IRAP pictures: IDR, CRA, BLA). //! Each PES packet = one access unit = one frame. -use super::h264::{find_start_code, skip_start_code}; +use super::startcode::{find_start_code, skip_start_code}; use super::{CodecParser, Frame, PesPacket, pts_to_ns}; // HEVC NAL unit types @@ -20,6 +20,9 @@ const _NAL_UNSPEC62_DV_RPU: u8 = 62; const NAL_BLA_W_LP: u8 = 16; const NAL_RSV_IRAP_VCL23: u8 = 23; +/// HEVC (H.265) Annex B → MKV codec parser: extracts VPS/SPS/PPS for the hvcC +/// codecPrivate, detects IRAP keyframes, and converts each PES access unit into +/// length-prefixed NAL units. Implements [`CodecParser`]. pub struct HevcParser { // First-seen parameter set of each type → seeds the MKV codecPrivate (hvcC). // This is the ONLY copy the player gets out-of-band, and a player re-applies @@ -42,6 +45,7 @@ impl Default for HevcParser { } impl HevcParser { + /// Create a fresh HEVC parser with no parameter sets captured yet. pub fn new() -> Self { Self { vps: None, @@ -71,12 +75,30 @@ fn handle_param_set(first: &mut Option<Vec<u8>>, nal: &[u8], frame_data: &mut Ve Some(f) if f.as_slice() == nal => {} // == codecPrivate → player has it Some(_) => { // Differs from codecPrivate → emit in-band so it wins at this AU. - frame_data.extend_from_slice(&(nal.len() as u32).to_be_bytes()); + // A NAL longer than u32::MAX can't be length-prefixed in the 4-byte + // field; skip it rather than mis-frame the output. Unreachable in + // practice (no real access unit is >4 GiB). + let Ok(len) = u32::try_from(nal.len()) else { + return; + }; + frame_data.extend_from_slice(&len.to_be_bytes()); frame_data.extend_from_slice(nal); } } } +/// Append `nal` to `out` as a 4-byte big-endian length prefix followed by the +/// NAL body. A NAL longer than `u32::MAX` can't be length-prefixed in the +/// 4-byte field, so it is skipped rather than mis-framed. Unreachable in +/// practice (no real access unit is >4 GiB). +fn push_length_prefixed(out: &mut Vec<u8>, nal: &[u8]) { + let Ok(len) = u32::try_from(nal.len()) else { + return; + }; + out.extend_from_slice(&len.to_be_bytes()); + out.extend_from_slice(nal); +} + impl CodecParser for HevcParser { fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> { if pes.data.is_empty() { @@ -112,7 +134,12 @@ impl CodecParser for HevcParser { end -= 1; } - if nal_start < data.len() { + // Skip empty NALs entirely. When the trailing-zero strip reduces + // `end` back to `nal_start` (e.g. `00 00 01 00 00 01`, or a + // zero-filled bad sector between two start codes), the slice is + // empty; emitting a 4-byte 0x00000000 length prefix with no NAL + // body produces a structurally invalid NALU a decoder rejects. + if nal_start < data.len() && end > nal_start { // HEVC NAL header: 2 bytes. Type is bits 1-6 of first byte. let nal_type = (data[nal_start] >> 1) & 0x3F; @@ -126,18 +153,18 @@ impl CodecParser for HevcParser { NAL_PPS => { handle_param_set(&mut self.pps, &data[nal_start..end], &mut frame_data) } - NAL_AUD => {} // Skip access unit delimiters + // Drop Access Unit Delimiters. This is intentional and + // spec-correct: Matroska HEVC frame data omits AUDs + // (the container delimits access units), so carrying + // them in-band is redundant. H.264 does the same below. + NAL_AUD => {} t if (NAL_BLA_W_LP..=NAL_RSV_IRAP_VCL23).contains(&t) => { keyframe = true; - let nal = &data[nal_start..end]; - frame_data.extend_from_slice(&(nal.len() as u32).to_be_bytes()); - frame_data.extend_from_slice(nal); + push_length_prefixed(&mut frame_data, &data[nal_start..end]); } _ => { // All other NAL types (slices, SEI, DV RPU, etc.) pass through - let nal = &data[nal_start..end]; - frame_data.extend_from_slice(&(nal.len() as u32).to_be_bytes()); - frame_data.extend_from_slice(nal); + push_length_prefixed(&mut frame_data, &data[nal_start..end]); } } } @@ -173,8 +200,9 @@ impl CodecParser for HevcParser { return None; } - // Simplified: store as arrays in Annex B format - // Full HEVCDecoderConfigurationRecord is complex — for now, concatenate + // Build a conforming HEVCDecoderConfigurationRecord: fixed header + // (configurationVersion, profile_tier_level fields, parallelism, parsed + // chroma/bit depths) followed by numOfArrays length-prefixed NAL arrays. let mut record = Vec::new(); // Minimal HEVCDecoderConfigurationRecord header. @@ -198,17 +226,17 @@ impl CodecParser for HevcParser { if sps.len() > 7 { record.extend_from_slice(&sps[4..8]); } else { - let avail = sps.len().saturating_sub(4).min(4); + let target = record.len() + 4; record.extend_from_slice(&sps[sps.len().min(4)..sps.len().min(8)]); - record.extend_from_slice(&vec![0u8; 4 - avail]); + record.resize(target, 0u8); // zero-pad the missing bytes in place } // general_constraint_indicator_flags (6 bytes) — SPS bytes 8..14 if sps.len() > 13 { record.extend_from_slice(&sps[8..14]); } else { - let avail = sps.len().saturating_sub(8).min(6); + let target = record.len() + 6; record.extend_from_slice(&sps[sps.len().min(8)..sps.len().min(14)]); - record.extend_from_slice(&vec![0u8; 6 - avail]); + record.resize(target, 0u8); // zero-pad the missing bytes in place } // general_level_idc — SPS byte 14 record.push(if sps.len() > 14 { sps[14] } else { 0 }); @@ -224,6 +252,8 @@ impl CodecParser for HevcParser { chroma_format_idc: 1, bit_depth_luma_minus8: 0, bit_depth_chroma_minus8: 0, + max_sub_layers_minus1: 0, + temporal_id_nesting_flag: 0, }); // chromaFormat (6 reserved bits set + 2-bit chroma_format_idc) record.push(0xFC | (chroma.chroma_format_idc & 0x03)); @@ -233,8 +263,14 @@ impl CodecParser for HevcParser { record.push(0xF8 | (chroma.bit_depth_chroma_minus8 & 0x07)); // avgFrameRate record.extend_from_slice(&[0, 0]); - // constantFrameRate + numTemporalLayers + temporalIdNested + lengthSizeMinusOne - record.push(0x03); // lengthSizeMinusOne = 3 (4 bytes) + // Byte 21 packs four fields (ISO/IEC 14496-15): + // constantFrameRate u(2) = 0 (unknown / not constant) + // numTemporalLayers u(3) = sps_max_sub_layers_minus1 + 1 + // temporalIdNested u(1) = sps_temporal_id_nesting_flag + // lengthSizeMinusOne u(2) = 3 (4-byte length prefix) + let num_temporal_layers = (chroma.max_sub_layers_minus1 + 1) & 0x07; + let temporal_id_nested = chroma.temporal_id_nesting_flag & 0x01; + record.push((num_temporal_layers << 3) | (temporal_id_nested << 2) | 0x03); // numOfArrays record.push(3); // VPS, SPS, PPS @@ -271,6 +307,10 @@ struct SpsChroma { chroma_format_idc: u8, bit_depth_luma_minus8: u8, bit_depth_chroma_minus8: u8, + /// sps_max_sub_layers_minus1 (u3): numTemporalLayers = this + 1 for hvcC. + max_sub_layers_minus1: u8, + /// sps_temporal_id_nesting_flag (u1) for hvcC temporalIdNested. + temporal_id_nesting_flag: u8, } /// Minimal MSB-first bit reader over a byte slice. @@ -365,7 +405,7 @@ fn parse_sps_chroma(sps: &[u8]) -> Option<SpsChroma> { // sps_max_sub_layers_minus1 u(3) let max_sub_layers_minus1 = r.read_bits(3)?; // sps_temporal_id_nesting_flag u(1) - r.skip_bits(1)?; + let temporal_id_nesting_flag = r.read_bit()?; // profile_tier_level( 1, sps_max_sub_layers_minus1 ) parse_profile_tier_level(&mut r, max_sub_layers_minus1)?; @@ -396,17 +436,18 @@ fn parse_sps_chroma(sps: &[u8]) -> Option<SpsChroma> { chroma_format_idc, bit_depth_luma_minus8, bit_depth_chroma_minus8, + max_sub_layers_minus1: max_sub_layers_minus1 as u8, + temporal_id_nesting_flag: temporal_id_nesting_flag as u8, }) } /// Consume a profile_tier_level(profilePresentFlag=1, maxNumSubLayersMinus1) /// structure from the bit reader (HEVC 7.3.3). fn parse_profile_tier_level(r: &mut BitReader, max_sub_layers_minus1: u32) -> Option<()> { - // general: profile_space u(2) + tier u(1) + profile_idc u(5) = 8 bits, - // profile_compatibility_flags u(32), 4× constraint/flags + reserved = 44 - // bits, general_inbld/reserved = 1 bit (total constraint area 48 bits), - // general_level_idc u(8). 88 bits = 11 bytes... but the spec packs the - // general PTL as 8 + 32 + 48 + 8 = 96 bits = 12 bytes. Skip 96 bits. + // general PTL fixed layout (HEVC 7.3.3): profile_space u(2) + tier u(1) + + // profile_idc u(5) = 8, general_profile_compatibility_flags u(32), + // constraint-flags/reserved area = 48, general_level_idc u(8). + // Total = 8 + 32 + 48 + 8 = 96 bits = 12 bytes. Skip 96 bits. r.skip_bits(96)?; if max_sub_layers_minus1 > 0 { @@ -859,6 +900,30 @@ mod tests { ); } + // --- empty NAL between adjacent start codes is skipped --- + + #[test] + fn empty_nal_between_start_codes_emits_no_bare_prefix() { + // `00 00 01 00 00 01 <real NAL>`: the first start code is immediately + // followed by another, so the in-between NAL is empty after the + // trailing-zero strip. It must be skipped, NOT written as a bare + // 0x00000000 length prefix (which a decoder treats as malformed). + let mut parser = HevcParser::new(); + let mut data = Vec::new(); + data.extend_from_slice(&[0x00, 0x00, 0x01]); // start code, empty NAL + data.extend_from_slice(&[0x00, 0x00, 0x01]); // next start code + data.extend_from_slice(&hevc_nal_header(1)); // TRAIL_R + data.extend_from_slice(&[0x10, 0x20]); + + let frames = parser.parse(&make_pes(data, Some(0))); + assert_eq!(frames.len(), 1); + let fd = &frames[0].data; + // Exactly one length-prefixed NAL — no zero-length entry. + let len = u32::from_be_bytes([fd[0], fd[1], fd[2], fd[3]]) as usize; + assert!(len > 0, "no bare zero-length prefix emitted"); + assert_eq!(len + 4, fd.len(), "exactly one NAL in frame data"); + } + // --- empty PES --- #[test] @@ -1129,6 +1194,20 @@ mod tests { assert_eq!(cp[18], 0xF8 | 4); } + #[test] + fn hvcc_byte21_from_sps_temporal_layers() { + // make_sps_with_chroma sets sps_max_sub_layers_minus1 = 0 and + // sps_temporal_id_nesting_flag = 1, so byte 21 must encode + // numTemporalLayers = 1, temporalIdNested = 1, lengthSizeMinusOne = 3: + // (1 << 3) | (1 << 2) | 3 = 0x0F. + let sps = make_sps_with_chroma(1, 2, 2); + let cp = codec_private_from_sps(&sps); + assert_eq!( + cp[21], 0x0F, + "byte 21: numTemporalLayers=1, temporalIdNested=1, lengthSizeMinusOne=3" + ); + } + #[test] fn hvcc_handles_emulation_prevention_in_sps() { // Insert an emulation-prevention byte (00 00 03) into the SPS RBSP and diff --git a/src/mux/codec/lpcm.rs b/src/mux/codec/lpcm.rs index 7954b07..4893ead 100644 --- a/src/mux/codec/lpcm.rs +++ b/src/mux/codec/lpcm.rs @@ -18,7 +18,10 @@ //! DVD = leave intact. The raw PCM data is otherwise one complete audio frame //! per PES; no framing is needed. //! -//! For MKV: codec ID "A_PCM/INT/BIG" (BD) or "A_PCM/INT/LIT" (DVD). +//! For MKV: both BD and DVD LPCM map to codec ID "A_PCM/INT/BIG" (big-endian). +//! DVD-Video LPCM is big-endian per the DVD-Video spec, and `mkv.rs` emits +//! "A_PCM/INT/BIG" unconditionally for `Codec::Lpcm` — there is no DVD/BD branch +//! and no "A_PCM/INT/LIT" path, so no byte-swap or alternate codec ID applies. //! All frames are keyframes (uncompressed audio). use super::{CodecParser, Frame, PesPacket, pts_to_ns}; @@ -131,8 +134,9 @@ mod tests { fn dvd_lpcm_preserves_all_pcm_bytes() { // DVD-PS LPCM: PsDemuxer already removed the 7-byte private sub-header, // so the payload handed to this parser is raw PCM. The DVD parser must - // NOT strip any further bytes (the round-2 audit Finding 3 bug: the BD - // 4-byte strip dropped one sample pair per PES, drifting the audio). + // NOT strip any further bytes — applying the BD 4-byte strip to a DVD + // payload would drop one sample pair per PES and progressively drift + // the audio. let mut parser = LpcmParser::new_dvd(); let pcm = vec![0xDE, 0xAD, 0xBE, 0xEF, 0xCA, 0xFE, 0x01, 0x02]; let frames = parser.parse(&make_pes(pcm.clone(), Some(90000))); diff --git a/src/mux/codec/mod.rs b/src/mux/codec/mod.rs index 8d77aaa..1b26bea 100644 --- a/src/mux/codec/mod.rs +++ b/src/mux/codec/mod.rs @@ -7,15 +7,27 @@ //! - Determine keyframe status //! - Convert PTS from 90kHz to nanoseconds +/// AC-3 / E-AC-3 (Dolby Digital / Digital Plus) elementary-stream parser. pub mod ac3; +/// DTS / DTS-HD elementary-stream parser. pub mod dts; +/// DVD bitmap subtitle (VobSub) parser. pub mod dvdsub; +/// H.264 (AVC) Annex-B elementary-stream parser. pub mod h264; +/// HEVC (H.265) Annex-B elementary-stream parser. pub mod hevc; +/// BD/DVD LPCM (Linear PCM) audio parser. pub mod lpcm; +/// MPEG-2 Video elementary-stream parser. pub mod mpeg2; +/// HDMV PGS (Presentation Graphics Stream) subtitle parser. pub mod pgs; +/// Shared MPEG/Annex-B start-code scanning helpers. +pub(crate) mod startcode; +/// Dolby TrueHD / Atmos elementary-stream parser. pub mod truehd; +/// VC-1 (SMPTE 421M) elementary-stream parser. pub mod vc1; use super::ts::PesPacket; @@ -70,12 +82,20 @@ pub trait CodecParser: Send { } /// Passthrough parser — treats each PES as one frame, no parsing. -/// Used for codecs where PES = frame (AC3, DTS, PGS). +/// +/// Used for the audio codecs that have no dedicated parser and whose PES +/// boundaries already line up with frame boundaries (Aac, Mp2, Mp3, Flac, +/// Opus). AC3/DTS/TrueHD have their own parsers; PGS/DvdSub have their own +/// subtitle parsers. Video codecs must NOT use the all-keyframe form of this +/// parser — see `parser_for_codec`. pub struct PassthroughParser { keyframe: bool, } impl PassthroughParser { + /// Create a passthrough parser. Pass `true` for codecs where every PES is + /// independently decodable (audio / subtitle keyframes), `false` for the + /// video fallback where no frame-boundary or keyframe detection occurs. pub fn new(always_keyframe: bool) -> Self { Self { keyframe: always_keyframe, @@ -124,6 +144,69 @@ pub fn parser_for_codec( Codec::Lpcm if is_dvd_ps => Box::new(lpcm::LpcmParser::new_dvd()), Codec::Lpcm => Box::new(lpcm::LpcmParser::new()), Codec::DvdSub => Box::new(dvdsub::DvdSubParser::new(codec_data)), - _ => Box::new(PassthroughParser::new(true)), + // Video codecs with no dedicated parser. There is no frame-boundary + // detection here, so a PES carrying multiple access units is emitted as + // one oversized block — but marking every frame a keyframe (as the + // audio passthrough does) would explode Cues density and mislead + // seeking. Use the non-keyframe passthrough and warn that framing is + // approximate. Mpeg1/Av1 are real Codec variants without a parser yet. + Codec::Mpeg1 | Codec::Av1 => { + tracing::warn!( + target: "mux", + "no dedicated parser for video codec {:?}; using non-keyframe passthrough (frame boundaries/keyframes not detected)", + codec + ); + Box::new(PassthroughParser::new(false)) + } + // Remaining audio-only codecs (Aac, Mp2, Mp3, Flac, Opus) where PES = + // frame: all-keyframe passthrough is correct. Subtitle/Unknown also land + // here; keyframe flag is irrelevant for them. + Codec::Aac | Codec::Mp2 | Codec::Mp3 | Codec::Flac | Codec::Opus => { + Box::new(PassthroughParser::new(true)) + } + Codec::Srt | Codec::Ssa | Codec::Unknown(_) => Box::new(PassthroughParser::new(true)), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn pes(pts: Option<i64>, data: Vec<u8>) -> PesPacket { + PesPacket { + pid: 0x1011, + pts, + dts: None, + data, + } + } + + #[test] + fn unhandled_video_codecs_use_non_keyframe_passthrough() { + // Mpeg1/Av1 have no dedicated parser. They must NOT be marked + // all-keyframe (that would explode Cues density and mislead seeking); + // the non-keyframe passthrough is the safe fallback. + for codec in [Codec::Mpeg1, Codec::Av1] { + let mut parser = parser_for_codec(codec, None, false); + let frames = parser.parse(&pes(Some(9000), vec![0xDE, 0xAD, 0xBE, 0xEF])); + assert_eq!(frames.len(), 1, "{codec:?}"); + assert!( + !frames[0].keyframe, + "{codec:?} must not be flagged keyframe by the fallback parser" + ); + assert_eq!(frames[0].data, vec![0xDE, 0xAD, 0xBE, 0xEF]); + } + } + + #[test] + fn unhandled_audio_codecs_use_keyframe_passthrough() { + // PES = frame audio codecs: every frame is independently decodable, so + // all-keyframe passthrough is correct. + for codec in [Codec::Aac, Codec::Mp2, Codec::Mp3, Codec::Flac, Codec::Opus] { + let mut parser = parser_for_codec(codec, None, false); + let frames = parser.parse(&pes(Some(0), vec![0x01, 0x02])); + assert_eq!(frames.len(), 1, "{codec:?}"); + assert!(frames[0].keyframe, "{codec:?} should be keyframe"); + } } } diff --git a/src/mux/codec/mpeg2.rs b/src/mux/codec/mpeg2.rs index f3783fe..d009017 100644 --- a/src/mux/codec/mpeg2.rs +++ b/src/mux/codec/mpeg2.rs @@ -9,6 +9,7 @@ //! - Sequence extension: 00 00 01 B5 //! - Picture header: 00 00 01 00 +use super::startcode::find_start_code; use super::{CodecParser, Frame, pts_to_ns}; use crate::mux::ts::PesPacket; @@ -61,6 +62,7 @@ impl Default for Mpeg2Parser { } impl Mpeg2Parser { + /// Create a new MPEG-2 parser with no captured sequence-header state. pub fn new() -> Self { Self { seq_header: None, @@ -102,8 +104,13 @@ impl CodecParser for Mpeg2Parser { // PTS-based seeking. Fall back to DTS only if PTS is absent. let pts_ns = pes.pts.or(pes.dts).map(pts_to_ns).unwrap_or(0); let data = &pes.data; - let mut keyframe = false; + // Keyframe-ness is a property of the coded PICTURE, not of a sequence + // header. A PES may carry a sequence header followed by a P/B-frame + // (open-GOP / re-encoded MPEG-2); the picture, not the seq header, + // decides the cue point. Set this only from the PICTURE_CODE arm. + let mut picture_is_keyframe = false; let mut has_picture = false; + let mut saw_seq_header = false; // Scan for start codes in the elementary stream data. let mut pos = 0; @@ -165,7 +172,19 @@ impl CodecParser for Mpeg2Parser { }; self.seq_header = Some(data[hdr_start..hdr_end].to_vec()); - keyframe = true; + // A NEW sequence header replaces the stored one, so its B5 + // sequence extension must be re-captured. Reset the flag the + // SEQ_EXT_CODE arm guards on; otherwise, once the first + // header's B3+B5 pair was seen, every later header (channel + // change, title boundary, parser reuse) would be stored + // without its extension bytes — corrupting codecPrivate + // (interlace, chroma format, progressive-sequence flags). + self.has_extension = false; + // NOTE: a sequence header does NOT make the access unit a + // keyframe — that is decided solely by the PICTURE_CODE arm + // (picture_is_keyframe). Setting it here would mis-cue a + // seq-header-followed-by-P/B-frame PES. + saw_seq_header = true; pos = if next_sc.is_some() { hdr_end } else { sc + 4 }; } SEQ_EXT_CODE if self.seq_header.is_some() && !self.has_extension => { @@ -185,7 +204,7 @@ impl CodecParser for Mpeg2Parser { if sc + 5 < data.len() { let picture_coding_type = (data[sc + 5] >> 3) & 0x07; if picture_coding_type == PICTURE_TYPE_I { - keyframe = true; + picture_is_keyframe = true; } } pos = sc + 4; @@ -209,13 +228,15 @@ impl CodecParser for Mpeg2Parser { // header and no picture. A PES with neither (e.g. a slice // continuation) still passes through unchanged, preserving real // keyframe detection. - if !has_picture && contains_seq_header(data) { + // `saw_seq_header` is set by the scan loop's SEQ_HEADER_CODE arm above, + // so this reuses that single pass instead of re-scanning the PES bytes. + if !has_picture && saw_seq_header { return Vec::new(); } vec![Frame { pts_ns, - keyframe, + keyframe: picture_is_keyframe, data: pes.data.clone(), duration_ns: None, }] @@ -264,29 +285,6 @@ fn parse_aspect_ratio(hdr: &[u8]) -> Option<(u8, u8)> { Some(ASPECT_RATIOS[ar_code]) } -/// Returns true if `data` contains a sequence-header start code (00 00 01 B3). -fn contains_seq_header(data: &[u8]) -> bool { - let mut pos = 0; - while let Some(sc) = find_start_code(data, pos) { - if sc + 3 >= data.len() { - break; - } - if data[sc + 3] == SEQ_HEADER_CODE { - return true; - } - pos = sc + 4; - } - false -} - -/// Find the position of the next start code (00 00 01) at or after `from`. -fn find_start_code(data: &[u8], from: usize) -> Option<usize> { - if data.len() < from + 3 { - return None; - } - (from..data.len() - 2).find(|&i| data[i] == 0x00 && data[i + 1] == 0x00 && data[i + 2] == 0x01) -} - #[cfg(test)] mod tests { use super::*; @@ -487,6 +485,33 @@ mod tests { assert!(parser.codec_private().is_some()); } + // --- seq-header keyframe flag must not leak into a P/B-frame --- + + #[test] + fn seq_header_then_p_frame_is_not_keyframe() { + // A PES carrying a sequence header followed by a P-frame (open-GOP / + // re-encoded MPEG-2) must NOT be flagged a keyframe — the keyframe-ness + // belongs to the coded picture, not the sequence header. A spurious + // keyframe here produces a bad MKV cue point. + let mut parser = Mpeg2Parser::new(); + + let mut data = Vec::new(); + data.extend_from_slice(&make_seq_header(720, 480, 3, 4)); + data.extend_from_slice(&make_picture_header(2)); // P-frame + data.extend_from_slice(&[0xFF; 16]); + + let pes = make_pes(data, Some(0)); + let frames = parser.parse(&pes); + + assert_eq!(frames.len(), 1); + assert!( + !frames[0].keyframe, + "seq-header + P-frame must not be a keyframe" + ); + // The sequence header is still captured for codecPrivate. + assert!(parser.codec_private().is_some()); + } + // --- parameter-set-only PES (seq header, no picture) emits no frame --- #[test] @@ -522,6 +547,49 @@ mod tests { assert!(frames2[0].keyframe); } + // --- a SECOND sequence header re-captures its extension --- + + #[test] + fn new_sequence_header_recaptures_extension() { + // Regression: has_extension was never reset when a new sequence header + // replaced the stored one, so a second header (channel change / title + // boundary) was stored WITHOUT its B5 sequence extension. To exercise + // the SEQ_EXT_CODE arm (which the has_extension flag guards), each + // header and its extension arrive in SEPARATE PES packets. + let mut parser = Mpeg2Parser::new(); + + // Header A (no trailing start code → captured alone), then its B5 + // extension in the next PES. + let _ = parser.parse(&make_pes(make_seq_header(1920, 1080, 3, 4), Some(0))); + let mut ext_a = vec![0x00, 0x00, 0x01, SEQ_EXT_CODE]; + ext_a.extend_from_slice(&[0x11, 0x22, 0x33, 0x44, 0x55, 0x66]); + let _ = parser.parse(&make_pes(ext_a, Some(0))); + assert!( + parser + .codec_private() + .unwrap() + .windows(6) + .any(|w| w == [0x11, 0x22, 0x33, 0x44, 0x55, 0x66]), + "first header's extension captured (has_extension now true)" + ); + + // A NEW header B, then ITS extension in a separate PES. With the bug, + // has_extension stayed true and this extension would be dropped. + let _ = parser.parse(&make_pes(make_seq_header(720, 480, 2, 4), Some(3600))); + let mut ext_b = vec![0x00, 0x00, 0x01, SEQ_EXT_CODE]; + ext_b.extend_from_slice(&[0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF]); + let _ = parser.parse(&make_pes(ext_b, Some(3600))); + + let cp2 = parser.codec_private().unwrap(); + assert!( + cp2.windows(6) + .any(|w| w == [0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF]), + "second header's extension must be re-captured, not dropped" + ); + // It is header B (720x480), not stale header A. + assert_eq!(parser.resolution(), Some((720, 480))); + } + // --- PTS conversion --- #[test] diff --git a/src/mux/codec/pgs.rs b/src/mux/codec/pgs.rs index df87b3f..d61038c 100644 --- a/src/mux/codec/pgs.rs +++ b/src/mux/codec/pgs.rs @@ -31,6 +31,8 @@ const MAX_PGS_PENDING_BYTES: usize = 4 * 1024 * 1024; // palette_id_ref) = 13. const PCS_NUM_OBJECTS_OFFSET: usize = 13; +/// Stateful parser that collapses PGS display/clear PCS pairs into +/// duration-bearing Matroska frames. Implements [`CodecParser`]. pub struct PgsParser { pending: Option<(i64, Vec<u8>)>, } @@ -42,9 +44,25 @@ impl Default for PgsParser { } impl PgsParser { + /// Create a fresh PGS parser with no pending display set. pub fn new() -> Self { Self { pending: None } } + + /// Take the pending display set (if any) and emit it as a Frame whose + /// duration runs from its start PTS to `end_pts_ns` (the PTS of the PCS that + /// closes or replaces it), clamped to >= 0. Shared by the clear-PCS and + /// replace-PCS arms so the Frame shape stays in one place. + fn emit_pending(&mut self, end_pts_ns: i64) -> Option<Frame> { + let (start_pts, data) = self.pending.take()?; + let duration = end_pts_ns.saturating_sub(start_pts).max(0) as u64; + Some(Frame { + pts_ns: start_pts, + keyframe: true, + data, + duration_ns: Some(duration), + }) + } } impl CodecParser for PgsParser { @@ -52,10 +70,36 @@ impl CodecParser for PgsParser { if pes.data.is_empty() { return Vec::new(); } - let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0); + // Keep PTS as Option: a PCS with no PTS has an UNKNOWN start/clear time. + // Collapsing it to a 0 sentinel produces a frame with a wrong start time + // and an absurd duration (the full elapsed time of the disc). PGS PCS + // packets carry a PTS on well-formed BD streams, so a missing PTS is a + // malformed-stream path that we skip cleanly rather than corrupt. + let pts = pes.pts.map(pts_to_ns); let is_pcs = pes.data[0] == SEGMENT_PCS; - let pcs_num_objects = if is_pcs && pes.data.len() > PCS_NUM_OBJECTS_OFFSET { + + // A PCS too short to carry number_of_composition_objects is malformed. + // Don't let it fall through to the non-PCS arm (where it would pollute + // the pending display set or pass through as a lone frame): close any + // pending set undurated (mirroring the no-PTS display path) and drop + // the truncated header so the parser resyncs on the next PCS. + if is_pcs && pes.data.len() <= PCS_NUM_OBJECTS_OFFSET { + return self + .pending + .take() + .map(|(start_pts, data)| { + vec![Frame { + pts_ns: start_pts, + keyframe: true, + data, + duration_ns: None, + }] + }) + .unwrap_or_default(); + } + + let pcs_num_objects = if is_pcs { Some(pes.data[PCS_NUM_OBJECTS_OFFSET]) } else { None @@ -65,33 +109,40 @@ impl CodecParser for PgsParser { match pcs_num_objects { // Clear/empty PCS — closes any pending display. Drop the // clear segment itself; BlockDuration covers the screen - // wipe. + // wipe. A clear PCS with no PTS can't time the duration, so + // emit the pending set with no duration (it lingers to EOF). Some(0) => { - if let Some((start_pts, data)) = self.pending.take() { - let duration = pts_ns.saturating_sub(start_pts).max(0) as u64; - out.push(Frame { + let frame = match pts { + Some(end) => self.emit_pending(end), + None => self.pending.take().map(|(start_pts, data)| Frame { pts_ns: start_pts, keyframe: true, data, - duration_ns: Some(duration), - }); - } + duration_ns: None, + }), + }; + out.extend(frame); } // Display PCS — start a new pending. If a prior display // was never explicitly cleared (replace-without-clear), // emit it with the new PCS's PTS as its end. - Some(_) => { - if let Some((start_pts, data)) = self.pending.take() { - let duration = pts_ns.saturating_sub(start_pts).max(0) as u64; - out.push(Frame { + Some(_) => match pts { + Some(start) => { + out.extend(self.emit_pending(start)); + self.pending = Some((start, pes.data.clone())); + } + // A display PCS with no PTS has an unknown start time. Don't + // store it with a 0 sentinel (wrong start, absurd duration). + // Flush any prior pending undurated and skip storing this one. + None => { + out.extend(self.pending.take().map(|(start_pts, data)| Frame { pts_ns: start_pts, keyframe: true, data, - duration_ns: Some(duration), - }); + duration_ns: None, + })); } - self.pending = Some((pts_ns, pes.data.clone())); - } + }, // Non-PCS first segment — either a continuation of the // current display set, or non-standard layout. If we have // a pending display, append; otherwise emit as-is. @@ -103,14 +154,21 @@ impl CodecParser for PgsParser { if buf.len() + pes.data.len() <= MAX_PGS_PENDING_BYTES { buf.extend_from_slice(&pes.data); } - } else { + } else if pes.pts.is_some() { + // A lone non-PCS segment with a real PTS — pass it through. + // (A missing PTS falls through to the drop path below: a + // bitmap with no timing reference would land at 00:00:00.) out.push(Frame { - pts_ns, + pts_ns: pts.unwrap_or(0), keyframe: true, data: pes.data.clone(), duration_ns: None, }); } + // No pending set AND no PTS: drop it. Emitting at pts_ns=0 would + // place a stray bitmap at 00:00:00.000 with no timing reference; + // the no-PTS PCS arms above avoid the 0 sentinel for the same + // reason. } } @@ -255,6 +313,82 @@ mod tests { assert_eq!(frames[0].duration_ns, None); } + #[test] + fn display_pcs_without_pts_is_not_stored_with_zero_start() { + // A display PCS with no PTS has an unknown start time. It must NOT be + // stored with a 0 sentinel — otherwise a later clear PCS at real PTS T + // would emit a frame with pts_ns=0 and duration_ns=T (hours of ns for a + // mid-disc subtitle). The malformed display PCS is skipped instead. + let mut parser = PgsParser::new(); + let frames = parser.parse(&make_pes(pcs_bytes(1), None)); + assert!(frames.is_empty(), "no-PTS display PCS emits nothing"); + assert!( + parser.pending.is_none(), + "no-PTS display PCS must not be stored as pending" + ); + + // A subsequent well-formed display + clear pair must time correctly, + // unpolluted by the skipped no-PTS PCS. + let _ = parser.parse(&make_pes(pcs_bytes(1), Some(90000))); + let f = parser.parse(&make_pes(pcs_bytes(0), Some(270000))); + assert_eq!(f.len(), 1); + assert_eq!(f[0].pts_ns, 1_000_000_000); + assert_eq!(f[0].duration_ns, Some(2_000_000_000)); + } + + #[test] + fn clear_pcs_without_pts_emits_pending_undurated() { + // A clear PCS that lacks a PTS can't compute a duration; the pending + // display is still emitted, but with no duration (lingers to EOF) + // instead of a bogus absurd one. + let mut parser = PgsParser::new(); + let _ = parser.parse(&make_pes(pcs_bytes(1), Some(90000))); + let f = parser.parse(&make_pes(pcs_bytes(0), None)); + assert_eq!(f.len(), 1); + assert_eq!(f[0].pts_ns, 1_000_000_000, "pending keeps its real start"); + assert_eq!(f[0].duration_ns, None, "no duration without a clear PTS"); + } + + #[test] + fn truncated_pcs_flushes_pending_and_resyncs() { + // A PCS too short to carry number_of_composition_objects arriving with a + // pending display must close that display (undurated) and drop the + // truncated header, not append its bytes into the pending bitmap. + let mut parser = PgsParser::new(); + let display = pcs_bytes(1); + assert!( + parser + .parse(&make_pes(display.clone(), Some(90000))) + .is_empty() + ); + + // A 13-byte (<= PCS_NUM_OBJECTS_OFFSET) PCS: truncated. + let truncated = vec![SEGMENT_PCS; PCS_NUM_OBJECTS_OFFSET]; + let frames = parser.parse(&make_pes(truncated, Some(180000))); + assert_eq!(frames.len(), 1, "pending display flushed on truncated PCS"); + assert_eq!(frames[0].data, display, "pending bitmap not polluted"); + assert_eq!(frames[0].duration_ns, None, "flushed undurated"); + assert!(parser.pending.is_none(), "parser resynced"); + } + + #[test] + fn lone_non_pcs_without_pts_is_dropped() { + // A non-PCS segment with no pending set and no PTS must be dropped, not + // emitted at pts_ns = 0 (which would land a stray bitmap at time zero). + let mut parser = PgsParser::new(); + let frames = parser.parse(&make_pes(vec![0x15, 0x00, 0x02, 0xAA], None)); + assert!(frames.is_empty(), "no pending + no PTS → dropped"); + } + + #[test] + fn lone_non_pcs_with_pts_passes_through() { + // A lone non-PCS segment WITH a PTS still passes through. + let mut parser = PgsParser::new(); + let frames = parser.parse(&make_pes(vec![0x15, 0x00, 0x02, 0xAA], Some(90000))); + assert_eq!(frames.len(), 1); + assert_eq!(frames[0].pts_ns, 1_000_000_000); + } + #[test] fn flush_with_nothing_pending_is_empty() { let mut parser = PgsParser::new(); diff --git a/src/mux/codec/startcode.rs b/src/mux/codec/startcode.rs new file mode 100644 index 0000000..9d156bf --- /dev/null +++ b/src/mux/codec/startcode.rs @@ -0,0 +1,93 @@ +//! Shared MPEG/Annex-B start-code scanning helpers. +//! +//! H.264, HEVC, MPEG-2 and the MPEG-2 Program Stream demuxer all locate the +//! 3-byte `00 00 01` start-code prefix to delimit NAL units / PES units. A +//! single memchr-backed implementation lives here so every caller gets the +//! same SIMD-accelerated scan instead of a hand-rolled byte-by-byte loop. + +/// Find the position of the next start code (`00 00 01`) at or after `from`. +/// +/// Backed by `memchr::memmem::find` for SIMD-accelerated bytestring search. On +/// AVX2-capable x86_64 this runs several times faster than a byte-by-byte scan; +/// on a 200 KB UHD HEVC frame the saving is in the hundreds of microseconds per +/// call. The reported offset is the start of the `00 00 01` triple, so a 4-byte +/// `00 00 00 01` start code is reported at the second `00`. +pub fn find_start_code(data: &[u8], from: usize) -> Option<usize> { + if data.len() < from + 3 { + return None; + } + memchr::memmem::find(&data[from..], b"\x00\x00\x01").map(|rel| from + rel) +} + +/// Skip past the start code at position `pos`, returning the first byte after +/// it. Handles both the 3-byte (`00 00 01`) and 4-byte (`00 00 00 01`) forms. +/// Returns `None` if `pos` does not begin a start code or the buffer is too +/// short to contain one. +pub fn skip_start_code(data: &[u8], pos: usize) -> Option<usize> { + if pos + 2 >= data.len() { + return None; + } + if data[pos] == 0x00 && data[pos + 1] == 0x00 { + if pos + 3 < data.len() && data[pos + 2] == 0x00 && data[pos + 3] == 0x01 { + return Some(pos + 4); // 4-byte start code + } + if data[pos + 2] == 0x01 { + return Some(pos + 3); // 3-byte start code + } + } + None +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn find_start_code_3byte() { + let data = [0x00, 0x00, 0x01, 0x65]; + assert_eq!(find_start_code(&data, 0), Some(0)); + } + + #[test] + fn find_start_code_4byte() { + let data = [0x00, 0x00, 0x00, 0x01, 0x65]; + // The 00 00 01 triple starts at offset 1 in a 4-byte start code. + assert_eq!(find_start_code(&data, 0), Some(1)); + } + + #[test] + fn find_start_code_offset() { + let data = [0xFF, 0xFF, 0x00, 0x00, 0x01, 0x09]; + assert_eq!(find_start_code(&data, 0), Some(2)); + } + + #[test] + fn find_start_code_none() { + let data = [0x00, 0x00, 0x00, 0x00]; + assert_eq!(find_start_code(&data, 0), None); + } + + #[test] + fn find_start_code_too_short() { + let data = [0x00, 0x00]; + assert_eq!(find_start_code(&data, 0), None); + } + + #[test] + fn skip_3byte() { + let data = [0x00, 0x00, 0x01, 0x65]; + assert_eq!(skip_start_code(&data, 0), Some(3)); + } + + #[test] + fn skip_4byte() { + let data = [0x00, 0x00, 0x00, 0x01, 0x65]; + assert_eq!(skip_start_code(&data, 0), Some(4)); + } + + #[test] + fn skip_not_a_start_code() { + let data = [0xFF, 0x00, 0x01, 0x65]; + assert_eq!(skip_start_code(&data, 0), None); + } +} diff --git a/src/mux/codec/truehd.rs b/src/mux/codec/truehd.rs index 02bf28b..9c79ede 100644 --- a/src/mux/codec/truehd.rs +++ b/src/mux/codec/truehd.rs @@ -4,9 +4,10 @@ //! Access units span PES boundaries — must buffer and reassemble. //! //! TrueHD access unit header (4 bytes): -//! [0..1] upper 4 bits = parity, lower 12 bits = length in 2-byte words -//! [2..3] timing value -//! [4..] substream data (major sync 0xF8726FBA may appear at offset 4) +//! bytes 0-1: top nibble = MLP check/access-unit nibble, lower 12 bits = +//! access-unit length in 2-byte words +//! bytes 2-3: timing value +//! bytes 4..: substream data (major sync 0xF8726FBA may appear at offset 4) //! //! AC-3 frames (interleaved, same PID): start with sync word 0x0B77. //! We skip AC-3 frames and only emit TrueHD access units. @@ -41,55 +42,70 @@ impl TrueHdParser { } } - /// Skip an AC-3 frame starting at the current buffer position. - /// Returns number of bytes consumed, or 0 if not enough data. - fn skip_ac3_frame(&self) -> usize { + /// Size (bytes) of the AC-3 frame at the buffer head. + /// + /// Distinguishes three cases the caller must treat differently: + /// - `Unmappable`: the header's fscod/frmsizecod don't map to a real frame + /// size (reserved fscod==3, or frmsizecod >= 38). The caller must drain + /// and resync, NOT wait for more data — waiting would stall forever. + /// - `NeedMore`: a valid size, but the frame isn't fully buffered yet. + /// - `Frame(n)`: a complete `n`-byte AC-3 frame is buffered. + /// + /// Frame sizing reuses `ac3::ac3_frame_size` so the AC-3 size table has a + /// single source of truth shared with the AC-3 parser; a returned `0` there + /// (reserved fscod or out-of-range frmsizecod) is the unmappable case. + fn ac3_frame_at_head(&self) -> Ac3Size { if self.buf.len() < 6 { - return 0; + return Ac3Size::NeedMore; } - // AC-3 frame size from frmsizcod + fscod - // Byte 4: [fscod:2][frmsizecod:6] - let fscod = (self.buf[4] >> 6) & 0x03; - let frmsizecod = (self.buf[4] & 0x3F) as usize; - // Frame size in 16-bit words per fscod (simplified table for common rates) - let frame_words = match fscod { - 0 => { - // 48 kHz - static SIZES: [usize; 38] = [ - 64, 64, 80, 80, 96, 96, 112, 112, 128, 128, 160, 160, 192, 192, 224, 224, 256, - 256, 320, 320, 384, 384, 448, 448, 512, 512, 640, 640, 768, 768, 896, 896, - 1024, 1024, 1152, 1152, 1280, 1280, - ]; - SIZES.get(frmsizecod).copied().unwrap_or(0) - } - 1 => { - // 44.1 kHz - static SIZES: [usize; 38] = [ - 69, 70, 87, 88, 104, 105, 121, 122, 139, 140, 174, 175, 208, 209, 243, 244, - 278, 279, 348, 349, 417, 418, 487, 488, 557, 558, 696, 697, 835, 836, 975, 976, - 1114, 1115, 1253, 1254, 1393, 1394, - ]; - SIZES.get(frmsizecod).copied().unwrap_or(0) - } - 2 => { - // 32 kHz - static SIZES: [usize; 38] = [ - 96, 96, 120, 120, 144, 144, 168, 168, 192, 192, 240, 240, 288, 288, 336, 336, - 384, 384, 480, 480, 576, 576, 672, 672, 768, 768, 960, 960, 1152, 1152, 1344, - 1344, 1536, 1536, 1728, 1728, 1920, 1920, - ]; - SIZES.get(frmsizecod).copied().unwrap_or(0) - } - _ => 0, - }; - let frame_bytes = frame_words * 2; - if frame_bytes == 0 || self.buf.len() < frame_bytes { - return 0; + let frame_bytes = super::ac3::ac3_frame_size(&self.buf); + if frame_bytes == 0 { + // Reserved fscod or out-of-range frmsizecod → unmappable header. + return Ac3Size::Unmappable; } - frame_bytes + if self.buf.len() < frame_bytes { + return Ac3Size::NeedMore; + } + Ac3Size::Frame(frame_bytes) } } +/// Secondary validation for an AC-3 frame of `frame_bytes` at the buffer head: +/// is its computed end a plausible boundary? Accept when the frame fills the +/// rest of the buffer, or the bytes that follow start another AC-3 sync +/// (0x0B77) or a plausible TrueHD access unit (non-zero 12-bit length within +/// the 32 KiB cap). If none holds, the leading 0x0B77 is more likely a TrueHD +/// AU header that happens to look like AC-3, so the AC-3 reading is rejected. +fn ac3_boundary_corroborated(buf: &[u8], frame_bytes: usize) -> bool { + if frame_bytes >= buf.len() { + // The AC-3 frame is fully buffered and ends the data — consistent. + return true; + } + let tail = &buf[frame_bytes..]; + if tail.len() < 2 { + // Not enough following bytes to judge; accept (the next call will see + // the continuation). + return true; + } + // Another AC-3 sync immediately after? + if tail[0] == 0x0B && tail[1] == 0x77 { + return true; + } + // A plausible TrueHD AU header after? (non-zero 12-bit length, <= 32 KiB) + let next_words = (((tail[0] as usize) << 8) | tail[1] as usize) & 0xFFF; + next_words != 0 && next_words * 2 <= 32768 +} + +/// Outcome of sizing the AC-3 frame at the TrueHD buffer head. +enum Ac3Size { + /// fscod/frmsizecod don't map to a real frame size — resync, don't wait. + Unmappable, + /// A valid size, but the frame is not fully buffered yet. + NeedMore, + /// A complete `n`-byte AC-3 frame is buffered. + Frame(usize), +} + impl CodecParser for TrueHdParser { fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> { if pes.data.is_empty() { @@ -118,29 +134,48 @@ impl CodecParser for TrueHdParser { break; } - // AC-3 frame (interleaved): starts with sync word 0x0B77 + // AC-3 frame (interleaved): starts with sync word 0x0B77. + // + // 0x0B 0x77 is also a legal TrueHD AU header (check-nibble 0, + // length-high-bits 0xB → length 0xB77 words). To avoid an AC-3 + // misread stealing a real TrueHD AU, an AC-3 frame is only accepted + // when its computed end is corroborated by what follows: end of + // buffer (frame fills the rest), another AC-3 sync, or a plausible + // TrueHD AU header. If none holds, this is treated as a TrueHD AU. if self.buf[0] == 0x0B && self.buf[1] == 0x77 { - let skip = self.skip_ac3_frame(); - if skip == 0 { - break; // incomplete AC-3 frame, wait for more data + match self.ac3_frame_at_head() { + Ac3Size::Unmappable => { + // Permanently unmappable header at the head would stall + // the parser forever; resync by dropping 2 bytes so one + // bad frame costs one frame, not the whole buffer. + self.buf.drain(..2); + continue; + } + Ac3Size::NeedMore => break, // wait for the rest of the frame + Ac3Size::Frame(skip) => { + if ac3_boundary_corroborated(&self.buf, skip) { + self.buf.drain(..skip); + continue; + } + // Not corroborated — fall through and interpret the + // 0x0B77 bytes as a TrueHD access unit instead. + } } - self.buf.drain(..skip); - continue; } // TrueHD access unit: lower 12 bits of first 2 bytes = length in words let unit_words = (((self.buf[0] as usize) << 8) | self.buf[1] as usize) & 0xFFF; if unit_words == 0 { - self.buf.drain(..2); + // A zero-length AU is malformed/padding. The AU header is 4 bytes + // (length + timing); drain the whole header, not just the length + // word, otherwise the timing bytes get misread as the next + // length word and produce a spurious parse on the next iteration. + self.buf.drain(..4); continue; } + // unit_words is masked to 12 bits, so unit_bytes <= 4095 * 2 = 8190; + // no separate oversize-resync guard is reachable. let unit_bytes = unit_words * 2; - if unit_bytes > 32768 { - // Likely misaligned — try to resync by scanning for AC-3 sync or - // a valid TrueHD length - self.buf.drain(..2); - continue; - } if self.buf.len() < unit_bytes { break; // incomplete access unit, wait for more data } @@ -352,6 +387,67 @@ mod tests { ); } + #[test] + fn zero_length_au_drains_full_header() { + // A zero-length AU header (4 bytes: length=0 + timing) must be skipped + // whole. If only 2 bytes were drained the timing bytes would be misread + // as a bogus length word. Here the timing bytes are 0x01 0x90 (= 0x190 = + // 400 words = 800 bytes) which, if misread, would stall the parser + // waiting for 800 bytes that never come. Draining 4 lets the following + // real unit parse. + let mut parser = TrueHdParser::new(); + let mut data = vec![0x00, 0x00, 0x01, 0x90]; // length=0, timing=0x0190 + data.extend_from_slice(&make_truehd_unit(200)); + let frames = parser.parse(&make_pes(data, Some(90000))); + assert_eq!(frames.len(), 1, "real unit parses after zero-length header"); + assert_eq!(frames[0].data.len(), 200); + } + + #[test] + fn unmappable_ac3_header_resyncs_not_stalls() { + // A permanently unmappable 0x0B77 header at the buffer head (reserved + // fscod==3) must NOT stall the parser. It used to be treated as + // "incomplete, wait" and break forever, dropping every following AU. + // Now it resyncs (drains 2 bytes) so a clean TrueHD unit behind it is + // eventually emitted. + let mut parser = TrueHdParser::new(); + // Unmappable AC-3-looking head: 0x0B77, byte4 fscod=3 (0xC0). + let mut data = vec![0x0B, 0x77, 0x00, 0x00, 0xC0, 0x00]; + // A clean TrueHD AU follows. + data.extend_from_slice(&make_truehd_unit(200)); + let frames = parser.parse(&make_pes(data, Some(90000))); + assert_eq!( + frames.len(), + 1, + "TrueHD AU behind a bad header is recovered" + ); + assert_eq!(frames[0].data.len(), 200); + assert!(parser.buf.is_empty(), "buffer fully consumed, no stall"); + } + + #[test] + fn truehd_au_with_0b77_head_not_stolen_by_ac3() { + // A TrueHD AU whose first two bytes are 0x0B 0x77 (length 0xB77 = 2935 + // words = 5870 bytes) must NOT be misrouted to the AC-3 path. The AC-3 + // size for this header (fscod from byte4) would close the boundary in + // the wrong place; the secondary corroboration rejects it because the + // computed AC-3 end is not followed by another AC-3 sync / TrueHD AU. + let mut parser = TrueHdParser::new(); + // 5870-byte AU starting with 0x0B 0x77. Byte 4 = 0x00 → AC-3 would + // size it as fscod=0, frmsizecod=0 → 128 bytes. The bytes at offset 128 + // are zeros (next_words==0) → not corroborated → kept as TrueHD. + let mut unit = vec![0u8; 5870]; + unit[0] = 0x0B; // 0xB high nibble of the 12-bit length, check nibble 0 + unit[1] = 0x77; // low byte of length 0xB77 + let frames = parser.parse(&make_pes(unit, Some(90000))); + assert_eq!(frames.len(), 1, "0x0B77-headed TrueHD AU kept whole"); + assert_eq!( + frames[0].data.len(), + 5870, + "AU sized by TrueHD length, not AC-3 frame size" + ); + } + #[test] fn codec_private_none() { let parser = TrueHdParser::new(); diff --git a/src/mux/codec/vc1.rs b/src/mux/codec/vc1.rs index d1b9e88..6cf89d1 100644 --- a/src/mux/codec/vc1.rs +++ b/src/mux/codec/vc1.rs @@ -3,7 +3,8 @@ //! VC-1 uses start codes similar to MPEG-2. //! Sequence header (0x0F) contains codec initialization data. //! Frame start = Frame header start code (0x0D). -//! I-frames (keyframes) are identified from the frame header. +//! I-frames (keyframes) are signalled by the presence of a Sequence Header +//! (0x0F) in the PES, per the BD VC-1 convention (see `parse`). use super::{CodecParser, Frame, PesPacket, pts_to_ns}; @@ -164,25 +165,54 @@ fn parse_vc1_resolution(sh: &[u8]) -> Option<(u32, u32)> { // Simple/Main profile: resolution not in sequence header return None; } - // Advanced profile layout (bit-level starting from sh[4]): - // profile(2) + level(3) + chroma_format(2) + quantizer_spec(3) + - // postproc_flag(1) + max_coded_width(12) + max_coded_height(12) ... - // Total bits before width: 2+3+2+3+1 = 11 bits - // We need at least 11+12+12 = 35 bits = 5 bytes from sh[4..] + // Advanced profile sequence-header layout (SMPTE 421M, bit-level from sh[4]): + // PROFILE(2) + LEVEL(3) + COLORDIFF_FORMAT(2) + FRMRTQ_POSTPROC(3) + + // BITRTQ_POSTPROC(5) + POSTPROCFLAG(1) + MAX_CODED_WIDTH(12) + + // MAX_CODED_HEIGHT(12) ... + // Total bits before MAX_CODED_WIDTH: 2+3+2+3+5+1 = 16 bits. + // We need 16+12+12 = 40 bits = 5 de-escaped bytes from sh[4..]. if sh.len() < 9 { return None; } - // Build a u64 from bytes 4..9 for easy bit extraction - let mut bits: u64 = 0; - for j in 0..5 { - bits = (bits << 8) | sh[4 + j] as u64; + // VC-1 Annex-B EBDU payload may carry emulation-prevention bytes (an + // inserted 0x03 after a 00 00 run). De-escape the payload before bit + // extraction so an EP byte landing within the first few bytes can't shift + // every subsequent bit and corrupt MAX_CODED_WIDTH/HEIGHT. Collect just the + // 5 de-escaped bytes the bit fields need. + let payload = &sh[4..]; + let mut deesc = Vec::with_capacity(5); + let mut zeros = 0u8; + for &b in payload { + if zeros >= 2 && b == 0x03 { + zeros = 0; // drop the emulation-prevention byte + continue; + } + deesc.push(b); + if deesc.len() == 5 { + break; + } + zeros = if b == 0x00 { zeros + 1 } else { 0 }; } - // bits has 40 bits. Skip first 11 bits, then read 12+12. - let coded_width = ((bits >> (40 - 11 - 12)) & 0xFFF) as u32 + 1; - let coded_height = ((bits >> (40 - 11 - 24)) & 0xFFF) as u32 + 1; + if deesc.len() < 5 { + return None; + } + // Build a u64 from the 5 de-escaped bytes for easy bit extraction. + let mut bits: u64 = 0; + for &b in &deesc { + bits = (bits << 8) | b as u64; + } + // bits holds 40 significant bits laid out as: + // [16 leading bits][MAX_CODED_WIDTH:12][MAX_CODED_HEIGHT:12] + // so MAX_CODED_WIDTH starts 12 bits from the LSB end and MAX_CODED_HEIGHT + // occupies the low 12 bits (shift 0). + const WIDTH_SHIFT: u64 = 12; // 40 - 16 - 12 + let coded_width = ((bits >> WIDTH_SHIFT) & 0xFFF) as u32 + 1; + let coded_height = (bits & 0xFFF) as u32 + 1; + // coded_width/height are `(bits & 0xFFF) + 1`, so always >= 1; after the + // ×2 both are always >= 2. Only the upper bound can fail. let w = coded_width * 2; let h = coded_height * 2; - if w > 0 && h > 0 && w <= 8192 && h <= 8192 { + if w <= 8192 && h <= 8192 { Some((w, h)) } else { None @@ -444,6 +474,59 @@ mod tests { assert_eq!(frames[0].pts_ns, 2_000_000_000); } + // --- advanced-profile resolution parsing (bit-offset regression) --- + + /// Build an advanced-profile VC-1 sequence header encoding the given + /// width/height. Layout from sh[4]: PROFILE(2)=3, LEVEL(3), COLORDIFF(2), + /// FRMRTQ(3), BITRTQ(5), POSTPROCFLAG(1) = 16 bits, then + /// MAX_CODED_WIDTH(12) = width/2 - 1, MAX_CODED_HEIGHT(12) = height/2 - 1. + fn make_ap_seq_header(width: u32, height: u32) -> Vec<u8> { + let coded_w = (width / 2) - 1; + let coded_h = (height / 2) - 1; + // Accumulate 40 bits MSB-first: 16 leading bits then 12+12. + let mut acc: u64 = 0; + let mut nbits = 0u32; + let put = |val: u64, n: u32, acc: &mut u64, nbits: &mut u32| { + *acc = (*acc << n) | (val & ((1u64 << n) - 1)); + *nbits += n; + }; + // PROFILE = 3 (advanced), then 14 more leading bits (all zero here). + put(0b11, 2, &mut acc, &mut nbits); + put(0, 14, &mut acc, &mut nbits); // level+colordiff+frmrtq+bitrtq+postproc + put(coded_w as u64, 12, &mut acc, &mut nbits); + put(coded_h as u64, 12, &mut acc, &mut nbits); + // 40 bits → 5 bytes, MSB-first. + let mut payload = Vec::with_capacity(5); + for i in (0..5).rev() { + payload.push(((acc >> (i * 8)) & 0xFF) as u8); + } + let mut sh = vec![0x00, 0x00, 0x01, SC_SEQUENCE_HEADER]; + sh.extend_from_slice(&payload); + sh + } + + #[test] + fn advanced_profile_resolution_uses_16bit_offset() { + // Regression: the parser skipped 11 bits (omitting BITRTQ_POSTPROC's 5 + // bits) instead of 16, reading width/height 5 bits too early. Encode a + // non-default 1280x720 and confirm it round-trips, proving the 16-bit + // pre-width offset. + let mut parser = Vc1Parser::new(); + let mut data = make_ap_seq_header(1280, 720); + // A frame so the parser emits and stores the header. + data.extend_from_slice(&[0x00, 0x00, 0x01, SC_FRAME]); + data.extend_from_slice(&[0x55, 0x66]); + parser.parse(&make_pes(data, Some(0))); + + let cp = parser.codec_private(); + // codec_private needs an entry point too; resolution is in width/height + // fields regardless. Read them off the parser via codec_private when + // available, else assert the internal fields directly. + assert_eq!(parser.width, 1280, "width parsed at the 16-bit offset"); + assert_eq!(parser.height, 720, "height parsed at the 16-bit offset"); + let _ = cp; + } + // --- find_next_sc utility --- #[test] diff --git a/src/mux/demux_thread.rs b/src/mux/demux_thread.rs index d4f9816..708a147 100644 --- a/src/mux/demux_thread.rs +++ b/src/mux/demux_thread.rs @@ -18,13 +18,14 @@ //! //! ## Lifecycle //! -//! [`DemuxThread::spawn`] takes ownership of the inner reader and the -//! demuxer state, returns a handle plus a `Receiver<DemuxBatch>`. -//! Dropping the handle closes the channel which signals the thread -//! to exit; the join in `Drop::drop` is bounded. +//! [`DemuxThread::spawn_zero_copy`] consumes the prefetch channels and +//! the demuxer state, returning a handle plus a `Receiver<DemuxBatch>`. +//! Dropping the handle closes the channel which signals the worker to +//! exit; the join in `Drop::drop` blocks until the worker observes +//! channel closure and returns (no timeout — a wedged downstream would +//! block the drop until it releases the channel). use crate::halt::Halt; -use crate::sector::SectorSource; use crossbeam_channel::{Receiver, Sender, bounded}; use std::thread::JoinHandle; @@ -39,6 +40,13 @@ pub enum DemuxBatch { Ps(Vec<super::ps::PsPacket>), /// Underlying reader returned an error. Terminal. Err(std::io::Error), + /// Explicit clean-completion sentinel. The worker sends this as its + /// LAST message on every non-error exit (input exhausted, or halt + /// cancelled) so the consumer can distinguish a normal end-of-stream + /// from a bare channel disconnection. A worker that panics mid-stream + /// drops `tx` without sending this, so the consumer sees `RecvError` + /// and reports the panic rather than silently truncating output. + Eof, } /// Spawned demux thread. Drop joins. @@ -56,158 +64,7 @@ pub struct DemuxThread { } impl DemuxThread { - /// Spawn the demux thread. Returns the thread handle and a - /// receiver for [`DemuxBatch`] items. - /// - /// `reader` is the fully-composed read+decrypt stack (e.g. - /// [`PrefetchedSectorSource`](crate::sector::PrefetchedSectorSource) - /// wrapping - /// [`DecryptingSectorSource`](crate::sector::DecryptingSectorSource)). - /// `extents` is what the thread walks; it issues one - /// `read_sectors` per batch of `batch_sectors` sectors (aligned - /// to 3-sector AACS units when possible). - pub fn spawn<S: SectorSource + Send + 'static>( - mut reader: S, - extents: Vec<crate::disc::Extent>, - batch_sectors: u16, - halt: Option<Halt>, - ts: Option<super::ts::TsDemuxer>, - ps: Option<super::ps::PsDemuxer>, - ) -> (Self, Receiver<DemuxBatch>) { - let (tx, rx) = bounded::<DemuxBatch>(DEMUX_CHANNEL_DEPTH); - let mut ts = ts; - let mut ps = ps; - - let handle = std::thread::Builder::new() - .name("freemkv-demux".into()) - .spawn(move || { - let mut buf = vec![0u8; batch_sectors as usize * 2048]; - let mut ext_idx = 0usize; - let mut offset: u32 = 0; - let prof = std::env::var_os("FREEMKV_PROFILE").is_some(); - let mut prof_started = std::time::Instant::now(); - let mut prof_last_dump = prof_started; - let mut prof_read_ns: u128 = 0; - let mut prof_feed_ns: u128 = 0; - let mut prof_send_ns: u128 = 0; - let mut prof_bytes: u64 = 0; - while ext_idx < extents.len() { - if halt.as_ref().map(|h| h.is_cancelled()).unwrap_or(false) { - return; - } - let ext = &extents[ext_idx]; - let remaining = ext.sector_count.saturating_sub(offset); - if remaining == 0 { - ext_idx += 1; - offset = 0; - continue; - } - let mut sectors = remaining.min(batch_sectors as u32) as u16; - if sectors >= 3 { - sectors -= sectors % 3; - } - let bytes = sectors as usize * 2048; - if buf.len() < bytes { - buf.resize(bytes, 0); - } - let lba = ext.start_lba + offset; - let t0 = if prof { - Some(std::time::Instant::now()) - } else { - None - }; - let n = match reader.read_sectors(lba, sectors, &mut buf[..bytes], false) { - Ok(n) => n, - Err(e) => { - let _ = tx.send(DemuxBatch::Err(e.into())); - return; - } - }; - let t1 = if prof { - Some(std::time::Instant::now()) - } else { - None - }; - offset += sectors as u32; - - // Demux this batch immediately so the channel - // carries already-parsed PesPackets, not raw - // sector bytes. - if let Some(ref mut d) = ts { - let pkts = d.feed(&buf[..n]); - let t2 = if prof { - Some(std::time::Instant::now()) - } else { - None - }; - if !pkts.is_empty() && tx.send(DemuxBatch::Ts(pkts)).is_err() { - return; // consumer dropped - } - let t3 = if prof { - Some(std::time::Instant::now()) - } else { - None - }; - if prof { - prof_read_ns += t1.unwrap().duration_since(t0.unwrap()).as_nanos(); - prof_feed_ns += t2.unwrap().duration_since(t1.unwrap()).as_nanos(); - prof_send_ns += t3.unwrap().duration_since(t2.unwrap()).as_nanos(); - prof_bytes += n as u64; - let now = t3.unwrap(); - if now.duration_since(prof_last_dump) - >= std::time::Duration::from_secs(5) - { - let el = now.duration_since(prof_started).as_millis().max(1); - let mbps = prof_bytes as u128 * 1000 / 1_000_000 / el; - eprintln!( - "[demux] elapsed={}ms in={}MB/s read={}% feed={}% send={}%", - el, - mbps, - prof_read_ns / 10_000 / el, - prof_feed_ns / 10_000 / el, - prof_send_ns / 10_000 / el, - ); - prof_last_dump = now; - prof_started = now; - prof_read_ns = 0; - prof_feed_ns = 0; - prof_send_ns = 0; - prof_bytes = 0; - } - } - } else if let Some(ref mut d) = ps { - let pkts = d.feed(&buf[..n]); - if !pkts.is_empty() && tx.send(DemuxBatch::Ps(pkts)).is_err() { - return; - } - } - } - // EOF — emit any flushed packets too. - if let Some(ref mut d) = ts { - let tail = d.flush(); - if !tail.is_empty() { - let _ = tx.send(DemuxBatch::Ts(tail)); - } - } else if let Some(ref mut d) = ps { - let tail = d.flush(); - if !tail.is_empty() { - let _ = tx.send(DemuxBatch::Ps(tail)); - } - } - // Sender drops here -> consumer sees RecvError → EOF. - }) - .expect("freemkv-demux thread spawn failed"); - - ( - Self { - handle: Some(handle), - producer_shell: None, - }, - rx, - ) - } - - /// Zero-copy variant. Instead of taking a `SectorSource` and + /// Spawn the demux thread. Instead of taking a `SectorSource` and /// memcpy-ing through its `read_sectors` API, this constructor /// consumes the prefetch channels directly: filled buffers come /// in via `prefetch_rx`, the demux thread feeds them, then @@ -230,7 +87,7 @@ impl DemuxThread { halt: Option<Halt>, ts: Option<super::ts::TsDemuxer>, ps: Option<super::ps::PsDemuxer>, - ) -> (Self, Receiver<DemuxBatch>) { + ) -> crate::error::Result<(Self, Receiver<DemuxBatch>)> { let (tx, rx) = bounded::<DemuxBatch>(DEMUX_CHANNEL_DEPTH); let mut ts = ts; let mut ps = ps; @@ -246,6 +103,10 @@ impl DemuxThread { let mut prof_bytes: u64 = 0; loop { if halt.as_ref().map(|h| h.is_cancelled()).unwrap_or(false) { + // Caller-initiated stop is a clean termination — + // send the Eof sentinel so the consumer doesn't + // mistake it for a worker panic. + let _ = tx.send(DemuxBatch::Eof); return; } let t0 = if prof { @@ -328,16 +189,21 @@ impl DemuxThread { let _ = tx.send(DemuxBatch::Ps(tail)); } } + // Clean end-of-stream sentinel. Reaching here means no + // panic occurred; a panic during `feed`/`flush` skips + // this and drops `tx`, which the consumer reads as an + // error rather than a clean EOF. + let _ = tx.send(DemuxBatch::Eof); }) - .expect("freemkv-demux thread spawn failed"); + .map_err(|e| crate::error::Error::IoError { source: e })?; - ( + Ok(( Self { handle: Some(handle), producer_shell: Some(Box::new(producer_shell)), }, rx, - ) + )) } } diff --git a/src/mux/disc.rs b/src/mux/disc.rs index ec3968a..846c254 100644 --- a/src/mux/disc.rs +++ b/src/mux/disc.rs @@ -5,7 +5,7 @@ //! //! Read-only. For disc→ISO (raw sector copy), use `Disc::copy()`. -use crate::disc::{Disc, DiscTitle, Extent}; +use crate::disc::{DiscTitle, Extent}; use crate::drive::extract_scsi_context; use crate::event::{BatchSizeReason, Event, EventKind}; use crate::halt::Halt; @@ -107,7 +107,6 @@ pub struct DiscStream { /// (raw / unencrypted disc) makes the decorator a pass-through. reader: DecryptingSectorSource<Box<dyn SectorSource>>, title: DiscTitle, - disc: Option<Disc>, /// Mirror of the keys handed in at construction. The decorator /// owns the cryptographic state; this field is kept for /// metadata-side callers (`info()` and friends) that want to @@ -159,6 +158,17 @@ pub struct DiscStream { parsers: Vec<(u16, Box<dyn super::codec::CodecParser>)>, pending_frames: std::collections::VecDeque<crate::pes::PesFrame>, pid_to_track: Vec<(u16, usize)>, + /// Cached `FREEMKV_SKIP_PARSE` profiling flag. The env var cannot + /// change at runtime, and `std::env::var_os` takes a process-wide + /// lock; reading it once at construction keeps it out of the + /// per-batch read() hot loop. + skip_parse: bool, + /// Cached `FREEMKV_PROFILE` presence, read once at construction. When + /// false, the read() loop skips the four `Instant::now()` captures and the + /// `prof_tick` calls entirely, so profiling-off runs pay no per-iteration + /// timestamp cost or `prof_active()` env-var lookup (which takes a + /// process-wide lock). + profiling: bool, } impl DiscStream { @@ -177,11 +187,15 @@ impl DiscStream { let extents = title.extents.clone(); let bytes_total_extents: u64 = extents.iter().map(|e| e.sector_count as u64 * 2048).sum(); - // Debug log reader type at construction — critical for diagnosing mux reading from drive instead of ISO + // Debug log reader type at construction — critical for diagnosing mux + // reading from drive instead of ISO. `type_name_of_val(&*reader)` + // resolves the CONCRETE type behind the box (Drive / FileSectorSource), + // unlike `type_name::<dyn SectorSource>()` which always prints the + // trait-object name regardless of the underlying source. tracing::debug!( target: "mux", "DiscStream constructed with reader type: {}", - std::any::type_name::<dyn SectorSource>() + std::any::type_name_of_val(&*reader) ); let mut pids = Vec::new(); @@ -220,7 +234,6 @@ impl DiscStream { // the decorator is a pass-through. reader: DecryptingSectorSource::new(reader, decrypt_keys.clone()), title, - disc: None, decrypt_keys, extents, current_extent: 0, @@ -240,6 +253,8 @@ impl DiscStream { parsers, pending_frames: std::collections::VecDeque::new(), pid_to_track, + skip_parse: std::env::var_os("FREEMKV_SKIP_PARSE").is_some(), + profiling: std::env::var_os("FREEMKV_PROFILE").is_some(), } } @@ -298,11 +313,6 @@ impl DiscStream { self.reader.set_keys(crate::decrypt::DecryptKeys::None); } - /// Get the scanned Disc (for listing all titles). - pub fn disc(&self) -> Option<&Disc> { - self.disc.as_ref() - } - fn fill_extents(&mut self) -> io::Result<bool> { if self.current_extent >= self.extents.len() { return Ok(false); @@ -317,7 +327,10 @@ impl DiscStream { return self.fill_extents(); } - let lba = ext_start + self.current_offset; + // start_lba comes from UDF/MPLS extents; a malformed extent near + // u32::MAX would overflow (debug panic / release wrap to a wrong LBA). + // Saturate for consistency with the rest of the file's arithmetic. + let lba = ext_start.saturating_add(self.current_offset); // Adaptive sizer: start at current (preferred until a failure), shrink // on failure, advance on success. One 5s read attempt per try — no @@ -349,15 +362,20 @@ impl DiscStream { let bytes = sectors as usize * 2048; self.read_buf.resize(bytes, 0); - let ok = self + let res = self .reader - .read_sectors(lba, sectors, &mut self.read_buf[..bytes], false) - .is_ok(); + .read_sectors(lba, sectors, &mut self.read_buf[..bytes], false); - if ok { + if let Ok(&got) = res.as_ref() { + // SectorSource::read_sectors returns the number of bytes + // written into buf. All in-tree sources return full-or-error, + // but a short count would leave the stale/zeroed tail of + // read_buf in place; trust the returned count, not `bytes`. + debug_assert!(got <= bytes, "read_sectors over-reported byte count"); if let Some(ev) = self.adaptive.on_success(sectors) { self.emit(ev); } + let bytes = got.min(bytes); self.buf_valid = bytes; self.current_offset += sectors as u32; self.bytes_read_total = self.bytes_read_total.saturating_add(bytes as u64); @@ -379,10 +397,14 @@ impl DiscStream { self.current_offset += 1; break; } else { - let err = self - .reader - .read_sectors(lba, sectors, &mut self.read_buf[..2048], false) - .err(); + // Build the error from the failure we ALREADY hold. + // Re-reading the same known-bad LBA here doubled drive + // abuse (hard rule #2: repeated failed reads on the + // same LBA push the BU40N into fast-fail) and, if the + // retry transiently succeeded, dropped the good data + // and returned a bogus status=0/sense=None error for a + // readable sector. + let err = res.err(); let (status, sense) = err.as_ref().map(extract_scsi_context).unwrap_or((0, None)); return Err(crate::error::Error::DiscRead { @@ -409,9 +431,9 @@ impl DiscStream { } /// Per-stage profiling state — populated only when `FREEMKV_PROFILE` -/// is set. Dumps a percentage breakdown to stderr every -/// [`PROFILE_INTERVAL`]. Zero overhead in normal runs (Option check -/// is the only added cost). +/// is set. Logs a percentage breakdown via `tracing` (target "mux") +/// every [`PROFILE_INTERVAL`]. Zero overhead in normal runs (the +/// `DiscStream::profiling` check is the only added cost). struct StageProf { started: std::time::Instant, last_dump: std::time::Instant, @@ -465,7 +487,8 @@ fn prof_tick(stage: &str, ns: u128, bytes: u64) { let feed_pct = p.feed_ns / 10_000 / elapsed_ms; let consume_pct = p.consume_ns / 10_000 / elapsed_ms; let mbps = p.bytes_in as u128 * 1000 / 1_000_000 / elapsed_ms; - eprintln!( + tracing::debug!( + target: "mux", "[profile] elapsed={}ms in={}MB/s fill={}% feed={}% consume={}%", elapsed_ms, mbps, fill_pct, feed_pct, consume_pct, ); @@ -484,7 +507,9 @@ impl crate::pes::Stream for DiscStream { } loop { - let t0 = std::time::Instant::now(); + // Profiling timestamps only when FREEMKV_PROFILE is set; otherwise + // these stay None and no Instant::now() is taken in the hot loop. + let t0 = self.profiling.then(std::time::Instant::now); if !self.fill_extents()? { self.eof = true; // Flush demuxer — last PES packet may still be in the assembler @@ -565,8 +590,10 @@ impl crate::pes::Stream for DiscStream { } let bytes = self.buf_valid; - let t1 = std::time::Instant::now(); - prof_tick("fill", t1.duration_since(t0).as_nanos(), bytes as u64); + let t1 = self.profiling.then(std::time::Instant::now); + if let (Some(t0), Some(t1)) = (t0, t1) { + prof_tick("fill", t1.duration_since(t0).as_nanos(), bytes as u64); + } // Plaintext: the wrapped reader (DecryptingSectorSource) // applied AACS / CSS in-place during fill_extents' // read_sectors call. The pre-0.18 inline decrypt step @@ -574,9 +601,11 @@ impl crate::pes::Stream for DiscStream { if let Some(ref mut demuxer) = self.ts_demuxer { let packets = demuxer.feed(&self.read_buf[..bytes]); - let t2 = std::time::Instant::now(); - prof_tick("feed", t2.duration_since(t1).as_nanos(), 0); - let skip_parse = std::env::var_os("FREEMKV_SKIP_PARSE").is_some(); + let t2 = self.profiling.then(std::time::Instant::now); + if let (Some(t1), Some(t2)) = (t1, t2) { + prof_tick("feed", t2.duration_since(t1).as_nanos(), 0); + } + let skip_parse = self.skip_parse; for pes in packets { if let Some((_, track)) = self .pid_to_track @@ -608,8 +637,10 @@ impl crate::pes::Stream for DiscStream { } } } - let t3 = std::time::Instant::now(); - prof_tick("consume", t3.duration_since(t2).as_nanos(), 0); + let t3 = self.profiling.then(std::time::Instant::now); + if let (Some(t2), Some(t3)) = (t2, t3) { + prof_tick("consume", t3.duration_since(t2).as_nanos(), 0); + } } else if let Some(ref mut demuxer) = self.ps_demuxer { let packets = demuxer.feed(&self.read_buf[..bytes]); for ps in &packets { @@ -691,7 +722,7 @@ impl crate::pes::Stream for DiscStream { // bottleneck profiling, so codec_private is never populated. // Pretend headers are ready immediately in that mode so the // CLI loop doesn't hang waiting for them. - if std::env::var_os("FREEMKV_SKIP_PARSE").is_some() { + if self.skip_parse { return true; } for (idx, s) in self.title.streams.iter().enumerate() { diff --git a/src/mux/ebml.rs b/src/mux/ebml.rs index ce240a7..16ef93e 100644 --- a/src/mux/ebml.rs +++ b/src/mux/ebml.rs @@ -40,6 +40,13 @@ pub fn write_size(w: &mut impl Write, size: u64) -> io::Result<()> { (size >> 8) as u8, size as u8, ]) + } else if size >= 0x00FF_FFFF_FFFF_FFFF { + // 0x00FF_FFFF_FFFF_FFFF (max 56-bit) encodes byte-for-byte + // identical to write_unknown_size (the EBML all-ones + // "unknown/open-ended" sentinel), and anything larger doesn't fit + // the 7-byte payload. Reject so a finite size can never be emitted + // as the unknown-size marker. + Err(crate::error::Error::MkvInvalid.into()) } else { // 8-byte size for large elements w.write_all(&[ @@ -119,9 +126,23 @@ pub fn start_master<W: Write + Seek>(w: &mut W, id: u32) -> io::Result<u64> { } /// End a master element: seek back and write the actual size. +/// +/// `size_pos` must be the offset returned by [`start_master`], which always +/// writes the 8-byte size placeholder before any body bytes. Therefore +/// `end_pos >= size_pos + 8` always holds, and the resulting `data_size` +/// fits the 7-byte VINT payload (a single MKV element exceeding 2^56 bytes +/// is not representable and never produced here). pub fn end_master<W: Write + Seek>(w: &mut W, size_pos: u64) -> io::Result<()> { let end_pos = w.stream_position()?; + debug_assert!( + end_pos >= size_pos + 8, + "end_master: end_pos {end_pos} < size_pos {size_pos} + 8 (placeholder not written?)" + ); let data_size = end_pos - size_pos - 8; // subtract the 8-byte size field itself + debug_assert!( + data_size < 0x0100_0000_0000_0000, + "end_master: data_size {data_size} exceeds the 7-byte VINT payload" + ); w.seek(SeekFrom::Start(size_pos))?; // Write as 8-byte VINT: 0x01 followed by 7 bytes of size w.write_all(&[ @@ -216,6 +237,9 @@ pub fn read_size(r: &mut impl Read) -> io::Result<(u64, usize)> { | (b[1] as u64) << 16 | (b[2] as u64) << 8 | b[3] as u64; + if val == 0x07_FFFF_FFFF { + return Ok((u64::MAX, 5)); + } Ok((val, 5)) } else if b0 & 0x04 != 0 { let mut b = [0u8; 5]; @@ -226,6 +250,9 @@ pub fn read_size(r: &mut impl Read) -> io::Result<(u64, usize)> { | (b[2] as u64) << 16 | (b[3] as u64) << 8 | b[4] as u64; + if val == 0x3FF_FFFF_FFFF { + return Ok((u64::MAX, 6)); + } Ok((val, 6)) } else if b0 & 0x02 != 0 { let mut b = [0u8; 6]; @@ -237,8 +264,11 @@ pub fn read_size(r: &mut impl Read) -> io::Result<(u64, usize)> { | (b[3] as u64) << 16 | (b[4] as u64) << 8 | b[5] as u64; + if val == 0x01_FFFF_FFFF_FFFF { + return Ok((u64::MAX, 7)); + } Ok((val, 7)) - } else { + } else if b0 & 0x01 != 0 { let mut b = [0u8; 7]; r.read_exact(&mut b)?; let val = (b[0] as u64) << 48 @@ -252,6 +282,13 @@ pub fn read_size(r: &mut impl Read) -> io::Result<(u64, usize)> { return Ok((u64::MAX, 8)); } Ok((val, 8)) + } else { + // b0 == 0x00: no length marker in the first byte. A VINT wider than + // 8 bytes is not representable by Matroska's size encoding, so this + // is a malformed/over-long size field rather than a valid 8-byte + // length. Reject it instead of silently building a size from the + // following 7 bytes (which would desync the parse). + Err(crate::error::Error::MkvInvalid.into()) } } @@ -280,13 +317,10 @@ pub fn read_uint_val(r: &mut impl Read, len: usize) -> io::Result<u64> { Ok(val) } -/// Read a float value. EBML floats are exactly 0, 4, or 8 bytes. -/// -/// The previous `else` branch read a fixed 8 bytes for ANY non-4 length, -/// so a malformed element with `len > 8` left `len - 8` unconsumed bytes -/// (mis-read as the next EBML header → desync of the rest of the parent -/// element) and `len < 4` over-read. Consume exactly `len` bytes and -/// reject anything that isn't a valid float width. +/// Read a float value. EBML floats are exactly 0, 4, or 8 bytes; any other +/// length is rejected as [`Error::MkvInvalid`] and exactly the float width is +/// consumed (so a malformed element never under- or over-reads and desyncs the +/// rest of the parent element). pub fn read_float_val(r: &mut impl Read, len: usize) -> io::Result<f64> { match len { 0 => Ok(0.0), @@ -311,7 +345,9 @@ pub fn read_string_val(r: &mut impl Read, len: usize) -> io::Result<String> { while buf.last() == Some(&0) { buf.pop(); } - String::from_utf8(buf).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e)) + // Library rule: errors are numeric variants, never English strings. + // A non-UTF-8 string element is malformed input → MkvInvalid. + String::from_utf8(buf).map_err(|_| crate::error::Error::MkvInvalid.into()) } /// Read binary data of `len` bytes. @@ -330,7 +366,10 @@ fn read_exact_bounded(r: &mut impl Read, len: usize) -> io::Result<Vec<u8>> { let mut buf = Vec::new(); let got = r.take(len as u64).read_to_end(&mut buf)?; if got != len { - return Err(io::ErrorKind::UnexpectedEof.into()); + // A truncated element is malformed input. Use the typed crate error + // so callers matching on Error::MkvInvalid catch short reads rather + // than a bare io::ErrorKind that bypasses the numeric-code identity. + return Err(crate::error::Error::MkvInvalid.into()); } Ok(buf) } @@ -456,6 +495,25 @@ mod tests { assert_eq!(buf, [126 | 0x80]); // 126 < 0x7F, uses 1 byte } + #[test] + fn write_size_rejects_unknown_size_sentinel() { + // 0x00FF_FFFF_FFFF_FFFF would encode byte-for-byte identical to the + // EBML unknown-size marker; it must be rejected, not silently emitted. + let mut buf = Vec::new(); + let e = write_size(&mut buf, 0x00FF_FFFF_FFFF_FFFF).unwrap_err(); + assert_eq!(e.kind(), io::ErrorKind::InvalidData); + assert!(buf.is_empty(), "no bytes should be written on rejection"); + + // One below the boundary still encodes as a normal 8-byte size whose + // payload is NOT all-ones, so read_size yields the finite value back. + buf.clear(); + let v = 0x00FF_FFFF_FFFF_FFFE; + write_size(&mut buf, v).unwrap(); + let (back, consumed) = read_size(&mut Cursor::new(&buf)).unwrap(); + assert_eq!(consumed, 8); + assert_eq!(back, v); + } + #[test] fn test_write_uint() { let mut buf = Vec::new(); @@ -632,6 +690,79 @@ mod tests { } } + #[test] + fn read_size_unknown_sentinel_all_widths() { + // The all-ones VINT of each width is the EBML "unknown size" marker + // and must read back as u64::MAX. write_size never emits the 5/6/7-byte + // widths, so these are hand-crafted. Each entry is (bytes, expected_len). + let cases: &[(&[u8], usize)] = &[ + // 1-byte: 0x80 | 0x7F + (&[0xFF], 1), + // 2-byte: 0x40 marker, value bits all 1 + (&[0x7F, 0xFF], 2), + // 3-byte + (&[0x3F, 0xFF, 0xFF], 3), + // 4-byte + (&[0x1F, 0xFF, 0xFF, 0xFF], 4), + // 5-byte (0x08 marker) + (&[0x0F, 0xFF, 0xFF, 0xFF, 0xFF], 5), + // 6-byte (0x04 marker) + (&[0x07, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF], 6), + // 7-byte (0x02 marker) + (&[0x03, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF], 7), + // 8-byte (0x01 marker) + (&[0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF], 8), + ]; + for (bytes, expected_len) in cases { + let mut cursor = Cursor::new(*bytes); + let (size, consumed) = read_size(&mut cursor).unwrap(); + assert_eq!( + size, + u64::MAX, + "all-ones {}-byte VINT should be unknown-size", + expected_len + ); + assert_eq!(consumed, *expected_len); + } + } + + #[test] + fn read_size_concrete_5_6_7_byte_values() { + // A non-sentinel 5/6/7-byte size must read back as its concrete value, + // not be mistaken for unknown-size. + // 5-byte: marker 0x08, value 0x01 (0x0800000001 with width bit only). + let mut c = Cursor::new(&[0x08u8, 0x00, 0x00, 0x00, 0x01]); + assert_eq!(read_size(&mut c).unwrap(), (1, 5)); + // 6-byte + let mut c = Cursor::new(&[0x04u8, 0x00, 0x00, 0x00, 0x00, 0x05]); + assert_eq!(read_size(&mut c).unwrap(), (5, 6)); + // 7-byte + let mut c = Cursor::new(&[0x02u8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x09]); + assert_eq!(read_size(&mut c).unwrap(), (9, 7)); + } + + #[test] + fn read_size_rejects_zero_first_byte() { + // b0 == 0x00 has no width marker — an over-long/invalid VINT. It must + // be rejected, not silently treated as an 8-byte size. + let mut c = Cursor::new(&[0x00u8, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF]); + let e = read_size(&mut c).unwrap_err(); + assert_eq!(e.kind(), io::ErrorKind::InvalidData); + } + + #[test] + fn write_size_rejects_at_or_above_2_56() { + // 2^56 cannot be encoded in the 7-payload-byte 8-byte VINT and must + // error rather than silently truncate. + let mut buf = Vec::new(); + let e = write_size(&mut buf, 0x0100_0000_0000_0000).unwrap_err(); + assert_eq!(e.kind(), io::ErrorKind::InvalidData); + // The largest encodable size still succeeds. + let mut buf = Vec::new(); + write_size(&mut buf, 0x00FF_FFFF_FFFF_FFFE).unwrap(); + assert_eq!(buf.len(), 8); + } + #[test] fn unknown_size() { let mut buf = Vec::new(); diff --git a/src/mux/fmp4/mod.rs b/src/mux/fmp4/mod.rs index 4aac3fa..a6dfdd9 100644 --- a/src/mux/fmp4/mod.rs +++ b/src/mux/fmp4/mod.rs @@ -1,18 +1,20 @@ -//! Fragmented MP4 muxer — **stub** for Phase 3. +//! Fragmented MP4 muxer — **STUB**: fragment emission is not implemented. //! //! Goal: ISO/IEC 14496-12 fragmented MP4 (`ftyp` + `moov` init segment, //! then a sequence of `moof+mdat` media fragments) targeting a //! [`SequentialSink`](crate::io::sink::SequentialSink). DASH-friendly, //! no Cues backpatch. //! -//! Status (v0.21.0 Phase 3): **STUB**. We ship the init segment -//! (`ftyp` + a minimal HEVC `moov` skeleton with one video track) so -//! the muxer's shape and call site are validated, but media fragments -//! are NOT yet emitted — calls to [`Fmp4Mux::write_video`] currently -//! accumulate frames into an internal buffer and discard them on -//! [`Fmp4Mux::finish`]. +//! Status: **STUB**. The muxer can emit the init segment (`ftyp` + a +//! minimal HEVC `moov` skeleton with one video track) via +//! [`Fmp4Mux::write_init_segment`], so the shape and call site are +//! validated, but media fragments (`moof`/`mdat`) are NOT emitted. +//! [`Fmp4Mux::write_video`] therefore returns +//! [`Error::Fmp4Unimplemented`](crate::error::Error::Fmp4Unimplemented) +//! rather than silently accepting and discarding frames. It buffers +//! nothing, so it cannot accumulate memory. //! -//! ## What's TODO (tracked in Phase 4 / v0.22.0 scope) +//! ## Not yet implemented //! //! - `moof` box: `mfhd` (sequence_number) + `traf` (`tfhd` + `tfdt` //! + `trun` with sample sizes, durations, flags, composition offsets). @@ -20,7 +22,7 @@ //! - Fragment cadence: one fragment per GOP or every N seconds, //! whichever comes first. //! - HEVC `hvcC` box inside `moov.trak.mdia.minf.stbl.stsd` so the -//! init segment is self-describing. +//! init segment is self-describing (`stsd` currently has zero entries). //! - Sample-flags computation (sync vs. delta, depends_on, etc.). //! - Edit lists / fragment_duration for accurate seeking. //! @@ -64,83 +66,66 @@ const VIDEO_TRACK_ID: u32 = 1; /// Fragmented MP4 muxer — stub. /// /// See the module-level doc comment for what is and isn't shipped in -/// this stub. +/// this stub. Fragment emission is not implemented: +/// [`write_video`](Self::write_video) returns +/// [`Error::Fmp4Unimplemented`](crate::error::Error::Fmp4Unimplemented) +/// rather than discarding media. pub struct Fmp4Mux<W: Write> { writer: W, header_written: bool, - /// Pending frames — held for the future fragment-emit path. The - /// stub drops these on `finish` but keeping them around lets the - /// post-stub work re-attach without changing the public API. - pending: Vec<PendingSample>, /// hvcC bytes, if provided. Embedded in the `moov.…stsd.hvc1.hvcC` - /// box once that path lands. + /// box once the emission path lands. #[allow(dead_code)] codec_private: Option<Vec<u8>>, } -struct PendingSample { - #[allow(dead_code)] - pts_ns: i64, - #[allow(dead_code)] - keyframe: bool, - #[allow(dead_code)] - data: Vec<u8>, -} - impl<W: Write> Fmp4Mux<W> { pub fn new(writer: W) -> Self { Self { writer, header_written: false, - pending: Vec::new(), codec_private: None, } } /// Provide the `HEVCDecoderConfigurationRecord` for the video track. /// The stub stores it but doesn't yet embed it in `moov` — that's - /// part of the post-stub work. + /// part of the unimplemented emission path. pub fn set_video_codec_private(&mut self, hvcc: Vec<u8>) { self.codec_private = Some(hvcc); } - /// Write one video PES frame. - /// - /// **Stub behaviour:** the first call emits the init segment - /// (`ftyp` + `moov`) so any consumer that just wants the shape can - /// receive it. Subsequent calls accumulate frames in memory for - /// the future fragmenting path; **no media bytes are written yet**. - pub fn write_video(&mut self, pts_ns: i64, keyframe: bool, data: &[u8]) -> io::Result<()> { - if !self.header_written { - self.write_init_segment()?; - self.header_written = true; + /// Emit the init segment (`ftyp` + `moov`) once. Idempotent — a + /// second call is a no-op. Lets a consumer that just wants the + /// container shape obtain a valid (if sample-less) init segment. + pub fn write_init_segment(&mut self) -> io::Result<()> { + if self.header_written { + return Ok(()); } - // TODO(0.22.0): emit one `moof+mdat` per GOP. For now stash the - // frame so the future patch can hot-wire emission without API - // churn. - self.pending.push(PendingSample { - pts_ns, - keyframe, - data: data.to_vec(), - }); - Ok(()) - } - - /// Flush. The stub additionally drops accumulated `pending` frames. - pub fn finish(&mut self) -> io::Result<()> { - // TODO(0.22.0): emit final fragment from pending; today the - // stub just clears the buffer to release memory. - self.pending.clear(); - self.writer.flush() - } - - fn write_init_segment(&mut self) -> io::Result<()> { let ftyp = build_ftyp(); let moov = build_moov(); self.writer.write_all(&ftyp)?; self.writer.write_all(&moov)?; + self.header_written = true; Ok(()) } + + /// Write one video PES frame. + /// + /// **Stub:** `moof`/`mdat` emission is not implemented. To avoid + /// silently dropping media (and avoid unbounded buffering), this + /// emits the init segment on the first call and then returns + /// [`Error::Fmp4Unimplemented`](crate::error::Error::Fmp4Unimplemented). + /// No frame bytes are buffered or written. + pub fn write_video(&mut self, _pts_ns: i64, _keyframe: bool, _data: &[u8]) -> io::Result<()> { + self.write_init_segment()?; + Err(crate::error::Error::Fmp4Unimplemented.into()) + } + + /// Flush the underlying writer. + pub fn finish(&mut self) -> io::Result<()> { + self.writer.flush() + } } /// Build the `ftyp` box. `major_brand = "iso6"`, `minor_version = 1`, @@ -218,7 +203,9 @@ fn build_tkhd() -> Vec<u8> { for v in [0x1_0000u32, 0, 0, 0, 0x1_0000, 0, 0, 0, 0x4000_0000] { body.extend_from_slice(&v.to_be_bytes()); } - // width / height in 16.16 fixed point — placeholder 1920x1080. + // width / height in 16.16 fixed point — placeholder; replace with + // SPS-derived dimensions (and matching stsd visual width/height) when + // fragment emission lands. body.extend_from_slice(&(1920u32 << 16).to_be_bytes()); body.extend_from_slice(&(1080u32 << 16).to_be_bytes()); wrap_box(&TKHD, &body) @@ -291,7 +278,9 @@ fn build_dinf() -> Vec<u8> { fn build_stbl() -> Vec<u8> { // Stub stsd: empty sample description (zero entries). Replace with // hvc1+hvcC once the fragmenting path lands so the init segment is - // actually decodable. + // actually decodable. Must be populated together with build_mvex: + // when the hvc1+hvcC sample entry lands and entry_count becomes 1, + // the trex default_sample_description_index=1 becomes valid. let mut stsd_body = Vec::new(); stsd_body.extend_from_slice(&[0, 0, 0, 0]); stsd_body.extend_from_slice(&0u32.to_be_bytes()); // entry_count @@ -317,6 +306,9 @@ fn build_stbl() -> Vec<u8> { fn build_mvex() -> Vec<u8> { // trex: track_ID=1, default_sample_description_index=1, others=0. + // dsdi=1 only becomes valid once build_stbl's stsd carries the + // matching hvc1 sample entry (entry_count=1) — keep the two in sync + // when fragment emission lands. let mut trex_body = Vec::new(); trex_body.extend_from_slice(&[0, 0, 0, 0]); // version + flags trex_body.extend_from_slice(&VIDEO_TRACK_ID.to_be_bytes()); @@ -331,9 +323,19 @@ fn build_mvex() -> Vec<u8> { /// Wrap a box body in `[size:u32-BE][type:4]`. Suitable for any body /// that fits in u32; oversized boxes (size > 4 GiB) need the 64-bit /// large-size extension which we don't generate in the stub. +/// +/// All callers build tiny init-segment boxes (kilobytes at most), so the +/// `u32` size never overflows; the saturating cast plus the debug assert +/// documents and guards that invariant rather than silently emitting a +/// truncated, structurally corrupt size field. `body` is always internally +/// constructed here, never untrusted input — a future caller feeding a +/// multi-gigabyte body trips the debug assert instead of writing a malformed +/// box. fn wrap_box(box_type: &[u8; 4], body: &[u8]) -> Vec<u8> { - let size = (body.len() + 8) as u32; - let mut out = Vec::with_capacity(body.len() + 8); + let total = body.len() + 8; + debug_assert!(total <= u32::MAX as usize, "fMP4 box exceeds u32 size"); + let size = u32::try_from(total).unwrap_or(u32::MAX); + let mut out = Vec::with_capacity(total); out.extend_from_slice(&size.to_be_bytes()); out.extend_from_slice(box_type); out.extend_from_slice(body); @@ -355,9 +357,7 @@ mod tests { fn init_segment_starts_with_ftyp_then_moov() { let mut sink: Vec<u8> = Vec::new(); let mut mux = Fmp4Mux::new(&mut sink); - // Trigger init emission via a single (stubbed) write. - mux.write_video(0, true, &[0x00, 0x00, 0x00, 0x01, 0x40]) - .unwrap(); + mux.write_init_segment().unwrap(); mux.finish().unwrap(); drop(mux); @@ -375,17 +375,26 @@ mod tests { } #[test] - fn moov_contains_trak_mvex() { + fn write_video_reports_unimplemented_and_buffers_nothing() { + // write_video must NOT silently accept-and-drop media: it emits the + // init segment, then signals that fragment emission is unimplemented. let mut sink: Vec<u8> = Vec::new(); let mut mux = Fmp4Mux::new(&mut sink); - mux.write_video(0, true, &[]).unwrap(); + let err = mux.write_video(0, true, &[0xDE; 4096]).unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::Unsupported); mux.finish().unwrap(); - drop(sink); + drop(mux); + // Only the init segment (ftyp + moov) was written — no media bytes. + let (ftyp_size, _) = read_box_header(&sink); + let (moov_size, _) = read_box_header(&sink[ftyp_size as usize..]); + assert_eq!(sink.len(), ftyp_size as usize + moov_size as usize); + } - // Re-emit into a fresh buffer for parsing. + #[test] + fn moov_contains_trak_mvex() { let mut buf: Vec<u8> = Vec::new(); let mut mux2 = Fmp4Mux::new(&mut buf); - mux2.write_video(0, true, &[]).unwrap(); + mux2.write_init_segment().unwrap(); mux2.finish().unwrap(); drop(mux2); diff --git a/src/mux/hevc/mod.rs b/src/mux/hevc/mod.rs index 41aba01..30d7bd9 100644 --- a/src/mux/hevc/mod.rs +++ b/src/mux/hevc/mod.rs @@ -69,19 +69,32 @@ impl<W: Write> HevcMux<W> { /// - Length-prefixed: `[u32-BE len][NAL bytes]` repeated. This is /// the form emitted by libfreemkv's HEVC parser (the MKV-native /// layout). Converted to Annex B. - /// - Already Annex B: bytes containing `00 00 00 01` start codes - /// anywhere in the buffer. Passed through unchanged. + /// - Already Annex B: a buffer beginning with a `00 00 00 01` or + /// `00 00 01` start code. Passed through unchanged. /// /// `_pts_ns` is accepted for symmetry with other muxers but ignored /// — Annex B has no timing layer. pub fn write_frame(&mut self, _pts_ns: i64, data: &[u8]) -> io::Result<()> { if !self.params_written { + // Mark written *before* the write: a partial write that then + // errors must not cause a later re-entry to re-emit the full + // parameter set on top of the bytes the sink already + // received (duplicate/split VPS/SPS/PPS). Callers discard the + // mux on any write error. + self.params_written = true; if let Some(cp) = &self.codec_private { - if let Some(params) = hvcc_to_annex_b(cp) { - self.writer.write_all(¶ms)?; + match hvcc_to_annex_b(cp) { + Some(params) => self.writer.write_all(¶ms)?, + // A non-empty hvcC that yields no NAL is a caller + // contract violation: emitting the stream without + // VPS/SPS/PPS produces undecodable output. Surface it + // rather than dropping the parameter sets silently. + None if !cp.is_empty() => { + return Err(crate::error::Error::HevcParamParse.into()); + } + None => {} } } - self.params_written = true; } let annex_b = length_prefixed_to_annex_b(data); self.writer.write_all(&annex_b) @@ -107,15 +120,22 @@ impl<W: Write> HevcMux<W> { /// /// We don't filter on NAL type — VPS (32), SPS (33), PPS (34), and any /// SEI arrays included in hvcC all get the same Annex B treatment. -fn hvcc_to_annex_b(hvcc: &[u8]) -> Option<Vec<u8>> { +/// +/// This is the single source of truth for hvcC → Annex B across all +/// muxers (HEVC ES, BD-TS, standard MPEG-TS). Do not reimplement it. +pub(crate) fn hvcc_to_annex_b(hvcc: &[u8]) -> Option<Vec<u8>> { if hvcc.len() < 23 { return None; } let num_arrays = hvcc[22] as usize; let mut out = Vec::new(); let mut offset = 23; + // Set when an inner loop exits on truncation so the outer loop stops + // too — otherwise it would re-interpret mid-NAL bytes as the next + // array header and synthesize spurious parameter-set NALs. + let mut truncated = false; for _ in 0..num_arrays { - if offset + 3 > hvcc.len() { + if truncated || offset + 3 > hvcc.len() { break; } offset += 1; // array_completeness + nal_type byte @@ -123,13 +143,20 @@ fn hvcc_to_annex_b(hvcc: &[u8]) -> Option<Vec<u8>> { offset += 2; for _ in 0..num_nalus { if offset + 2 > hvcc.len() { + truncated = true; break; } let nal_len = u16::from_be_bytes([hvcc[offset], hvcc[offset + 1]]) as usize; offset += 2; if offset + nal_len > hvcc.len() { + truncated = true; break; } + // ISO/IEC 14496-15 disallows zero-length NAL entries; emitting + // a bare start code with no RBSP yields an invalid Annex B NAL. + if nal_len == 0 { + continue; + } out.extend_from_slice(&START_CODE); out.extend_from_slice(&hvcc[offset..offset + nal_len]); offset += nal_len; @@ -141,13 +168,42 @@ fn hvcc_to_annex_b(hvcc: &[u8]) -> Option<Vec<u8>> { /// Convert length-prefixed NAL units (`[u32-BE len][NAL]` repeated) to /// Annex B (`00 00 00 01 [NAL]` repeated). /// -/// If the input doesn't parse as length-prefixed (no valid lengths -/// extracted), it's returned unchanged on the assumption that it's -/// already Annex B — some upstream paths (raw HEVC ES from disc) pass -/// Annex B straight through the PES layer. +/// Already-Annex-B input (a buffer beginning with a `00 00 00 01` or +/// `00 00 01` start code) is detected up front and passed through +/// unchanged — some upstream paths (raw HEVC ES from disc) hand Annex B +/// straight through the PES layer, and a genuine start code would +/// otherwise be misread as a u32-BE length prefix. +/// +/// Truncation policy (single source of truth across all muxers): if a +/// length prefix runs past the end of the buffer (e.g. a NAL truncated +/// by a bad disc sector), the truncated trailing NAL is dropped and only +/// the valid Annex-B prefix accumulated so far is emitted. We never emit +/// a half-NAL nor leak raw length-prefixed bytes into the Annex-B stream. pub(crate) fn length_prefixed_to_annex_b(data: &[u8]) -> Vec<u8> { + // Probe for a leading Annex B start code before attempting to parse + // length prefixes: `00 00 00 01` would otherwise parse as length 1. + if starts_with_start_code(data) { + return data.to_vec(); + } let mut out = Vec::with_capacity(data.len() + (data.len() / 32)); + append_length_prefixed_as_annex_b(&mut out, data); + out +} + +/// Append the Annex B form of `data` (length-prefixed NALs) into `out`. +/// +/// Same conversion as [`length_prefixed_to_annex_b`] but writes directly +/// into a caller-owned buffer, avoiding an intermediate allocation on +/// hot paths (e.g. per-frame video muxing). If `data` doesn't parse as +/// length-prefixed (no NALs extracted), it's appended unchanged on the +/// assumption it's already Annex B. +pub(crate) fn append_length_prefixed_as_annex_b(out: &mut Vec<u8>, data: &[u8]) { let mut offset = 0; + // True once we've consumed at least one well-formed length prefix + // (even a zero-length one). Distinguishes "parsed as length-prefixed, + // all NALs empty" (emit nothing) from "not length-prefixed at all" + // (pass through as already-Annex B). + let mut parsed_any = false; while offset + 4 <= data.len() { let len = u32::from_be_bytes([ data[offset], @@ -157,19 +213,39 @@ pub(crate) fn length_prefixed_to_annex_b(data: &[u8]) -> Vec<u8> { ]) as usize; offset += 4; if offset + len > data.len() { - // Mid-NAL truncation — fall through to the pass-through path - // rather than emitting a half-NAL. - return data.to_vec(); + // Mid-NAL truncation (e.g. a NAL cut by a bad disc sector) — + // drop the truncated trailing NAL and emit only the valid + // Annex-B prefix accumulated so far. We never emit a half-NAL + // nor leak raw length-prefixed bytes into the Annex-B stream. + break; + } + parsed_any = true; + if len == 0 { + // A zero-length prefix (e.g. pad bytes read off a damaged + // sector) would otherwise emit a bare start code with no + // RBSP — an invalid empty Annex B NAL. Skip it, mirroring + // the `nal_len == 0` guard in `hvcc_to_annex_b` (ISO/IEC + // 14496-15). + continue; } out.extend_from_slice(&START_CODE); out.extend_from_slice(&data[offset..offset + len]); offset += len; } - if out.is_empty() && !data.is_empty() { - // No length prefixes found — input is likely already Annex B. - return data.to_vec(); + if !parsed_any && !data.is_empty() { + // No length prefixes parsed at all and no leading start code: + // pass the bytes through rather than discard them (recover-100% + // goal — a decoder can attempt its own resync; dropping them + // guarantees loss). This is distinct from "parsed as length- + // prefixed but every NAL was zero-length", which emits nothing. + out.extend_from_slice(data); } - out +} + +/// Whether `data` begins with a 4-byte (`00 00 00 01`) or 3-byte +/// (`00 00 01`) Annex B start code. +fn starts_with_start_code(data: &[u8]) -> bool { + data.starts_with(&START_CODE) || data.starts_with(&[0x00, 0x00, 0x01]) } #[cfg(test)] @@ -203,18 +279,103 @@ mod tests { } #[test] - fn mid_nal_truncation_returns_original() { - // `[u32-BE 100][only 3 bytes]` — length prefix claims 100 bytes - // but the input only has 3 after the prefix. We treat that as - // malformed and pass the original buffer through so receivers - // can attempt their own recovery. + fn mid_nal_truncation_drops_trailing_nal_keeps_prefix() { + // First NAL is valid (2-byte payload), second has a length prefix + // claiming 100 bytes with only 3 present. Policy: emit the valid + // first NAL as Annex B, drop the truncated trailing NAL — never + // leak raw length-prefixed bytes into the Annex B stream. let mut raw = Vec::new(); + raw.extend_from_slice(&2u32.to_be_bytes()); + raw.extend_from_slice(&[0x11, 0x22]); raw.extend_from_slice(&100u32.to_be_bytes()); raw.extend_from_slice(&[0xAA, 0xBB, 0xCC]); let got = length_prefixed_to_annex_b(&raw); + let want = [0x00, 0x00, 0x00, 0x01, 0x11, 0x22]; + assert_eq!(&got[..], &want[..]); + } + + #[test] + fn leading_annex_b_start_code_passes_through() { + // Genuine Annex B beginning with 00 00 00 01 must NOT be reframed: + // the start code would otherwise parse as a u32-BE length of 1. + let raw = [ + 0x00, 0x00, 0x00, 0x01, 0x26, 0x01, 0xDE, 0xAD, // NAL 1 + 0x00, 0x00, 0x00, 0x01, 0x02, 0x01, 0xBE, 0xEF, // NAL 2 + ]; + let got = length_prefixed_to_annex_b(&raw); + assert_eq!( + &got[..], + &raw[..], + "Annex B input must pass through verbatim" + ); + } + + #[test] + fn leading_three_byte_start_code_passes_through() { + let raw = [0x00, 0x00, 0x01, 0x26, 0x01, 0xDE, 0xAD]; + let got = length_prefixed_to_annex_b(&raw); assert_eq!(&got[..], &raw[..]); } + #[test] + fn hvcc_skips_zero_length_nal_entries() { + // hvcC with one array containing a zero-length NAL followed by a + // valid one: the zero-length entry must be skipped, not emitted as + // a bare start code. + let mut hvcc = vec![0u8; 22]; + hvcc.push(1); // numArrays + hvcc.push(33); // SPS + hvcc.extend_from_slice(&2u16.to_be_bytes()); // numNalus = 2 + hvcc.extend_from_slice(&0u16.to_be_bytes()); // NAL 0: length 0 + hvcc.extend_from_slice(&3u16.to_be_bytes()); // NAL 1: length 3 + hvcc.extend_from_slice(&[0x42, 0x01, 0x01]); + let annex_b = hvcc_to_annex_b(&hvcc).expect("one valid NAL"); + let want = [0x00, 0x00, 0x00, 0x01, 0x42, 0x01, 0x01]; + assert_eq!(&annex_b[..], &want[..]); + } + + #[test] + fn write_frame_errors_on_unparseable_non_empty_hvcc() { + // A non-empty hvcC that yields no NAL must surface an error + // instead of silently producing a parameter-set-less stream. + let mut sink: Vec<u8> = Vec::new(); + let mut mux = HevcMux::new(&mut sink); + mux.set_codec_private(vec![0xDE, 0xAD]); // too short to be valid hvcC + let err = mux.write_frame(0, &[]).unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + } + + #[test] + fn zero_length_nal_is_skipped_not_bare_start_code() { + // A zero-length prefix between two real NALs must be skipped, not + // turned into a bare `00 00 00 01` with no RBSP. + let mut buf = Vec::new(); + buf.extend_from_slice(&3u32.to_be_bytes()); + buf.extend_from_slice(&[0xAA, 0xBB, 0xCC]); + buf.extend_from_slice(&0u32.to_be_bytes()); // zero-length NAL + buf.extend_from_slice(&2u32.to_be_bytes()); + buf.extend_from_slice(&[0xDD, 0xEE]); + + let got = length_prefixed_to_annex_b(&buf); + let want = [ + 0x00, 0x00, 0x00, 0x01, 0xAA, 0xBB, 0xCC, // first NAL + 0x00, 0x00, 0x00, 0x01, 0xDD, 0xEE, // second NAL (zero-length skipped) + ]; + assert_eq!(&got[..], &want[..]); + } + + #[test] + fn all_zero_length_nals_emit_nothing() { + // A buffer of only zero-length prefixes parses as length-prefixed + // but yields no NALs — output must be empty, not a pass-through of + // the raw zero bytes. + let mut buf = Vec::new(); + buf.extend_from_slice(&0u32.to_be_bytes()); + buf.extend_from_slice(&0u32.to_be_bytes()); + let got = length_prefixed_to_annex_b(&buf); + assert!(got.is_empty(), "expected empty output, got {got:?}"); + } + #[test] fn hvcc_extracts_vps_sps_pps() { // Build a minimal-but-valid hvcC: 22-byte header, then 3 arrays diff --git a/src/mux/m2ts.rs b/src/mux/m2ts.rs index 5ef206e..a8e11d1 100644 --- a/src/mux/m2ts.rs +++ b/src/mux/m2ts.rs @@ -22,11 +22,13 @@ impl M2tsStream { /// Create for writing PES frames → BD-TS output. /// Writes FMKV metadata header, then muxes PES frames into BD transport stream. pub fn create(mut writer: impl Write + Send + 'static, title: &DiscTitle) -> io::Result<Self> { - // Write FMKV metadata header - if !title.streams.is_empty() { - let m = meta::M2tsMeta::from_title(title); - meta::write_header(&mut writer, &m)?; - } + // Write FMKV metadata header unconditionally. An empty streams + // array is valid JSON and round-trips fine; skipping the header + // for a zero-stream title would make the output indistinguishable + // from a non-FMKV file on read-back (read_header returns + // Ok(None) → PMT fallback) even though M2tsStream produced it. + let m = meta::M2tsMeta::from_title(title); + meta::write_header(&mut writer, &m)?; let pids: Vec<u16> = title .streams .iter() @@ -39,8 +41,14 @@ impl M2tsStream { let boxed: Box<dyn Write + Send> = Box::new(writer); let mut muxer = super::tsmux::TsMuxer::new(boxed, &pids); for (i, cp) in title.codec_privates.iter().enumerate() { + // codec_privates is parallel to streams/pids; ignore any + // trailing entries that exceed the track count rather than + // surfacing a track-range error for a benign metadata overrun. + if i >= pids.len() { + break; + } if let Some(data) = cp { - muxer.set_codec_private(i, data.clone()); + muxer.set_codec_private(i, data.clone())?; } } Ok(Self { @@ -178,8 +186,11 @@ mod tests { let ts_bytes = &buf[header_end..]; // Find first PUSI packet on VIDEO_PID; verify RAI in AF flags. + // chunks_exact drops any partial trailing chunk — only whole + // 192-byte BD-TS packets are valid, and it avoids OOB indexing on a + // short final chunk. let pkt = ts_bytes - .chunks(192) + .chunks_exact(192) .find(|p| { let h = &p[4..]; let pid = (((h[1] & 0x1F) as u16) << 8) | h[2] as u16; diff --git a/src/mux/m2ts_mux/mod.rs b/src/mux/m2ts_mux/mod.rs index 4db2157..7d2a8ca 100644 --- a/src/mux/m2ts_mux/mod.rs +++ b/src/mux/m2ts_mux/mod.rs @@ -76,9 +76,14 @@ const PCR_INTERVAL_PACKETS: u64 = 40; /// the picture it timestamps. 200 ms in 90 kHz ticks. const PCR_LEAD_90KHZ: u64 = 90_000 / 5; -/// Stream-type codes from ISO/IEC 13818-1 Table 2-29 + later amendments. +/// HEVC stream-type code, ISO/IEC 13818-1 Table 2-34 (2015 amendment). const STREAM_TYPE_HEVC: u8 = 0x24; +/// AC-3 / E-AC-3. Not an ISO assignment — sits in the user-private +/// 0x80-0xFF range and is the Blu-ray Disc Association / ATSC A/52 +/// convention. const STREAM_TYPE_AC3: u8 = 0x81; +/// Dolby TrueHD. Also a private/BD-conventional value in the +/// user-private 0x80-0xFF range, not an ISO assignment. const STREAM_TYPE_TRUEHD: u8 = 0x83; /// Audio codec hint for [`M2tsMux::new`] / [`M2tsMux::set_audio`]. The @@ -126,6 +131,10 @@ pub struct M2tsMux<W: Write> { packets_written: u64, /// Video packets written since last PCR, used to gate PCR cadence. video_packets_since_pcr: u64, + /// Set once the first video TS packet has been emitted. Forces a PCR + /// onto the very first video PES so a receiver tuning at stream start + /// has a clock reference (the PMT advertises the video PID as PCR_PID). + first_video_written: bool, } impl<W: Write> M2tsMux<W> { @@ -145,6 +154,7 @@ impl<W: Write> M2tsMux<W> { cc_pmt: 0, packets_written: 0, video_packets_since_pcr: 0, + first_video_written: false, } } @@ -166,7 +176,7 @@ impl<W: Write> M2tsMux<W> { /// PES (and gates codec_private NAL prepending — those only attach to /// the first keyframe). pub fn write_video(&mut self, pts_ns: i64, keyframe: bool, data: &[u8]) -> io::Result<()> { - let pts_90k = self.base_relative_pts(pts_ns); + let pts_90k = self.base_relative_pts(pts_ns, /* may_seed_base */ true); // PCR comes "before" the PTS it timestamps; clamp at 0 for the // first frame so we don't underflow. let pcr = pts_90k.saturating_sub(PCR_LEAD_90KHZ); @@ -177,16 +187,15 @@ impl<W: Write> M2tsMux<W> { let mut es = Vec::with_capacity(data.len() + 64); if keyframe && !self.params_written { if let Some(cp) = &self.video_codec_private { - let payload = hvcc_payload(cp); - if !payload.is_empty() { - let params = super::hevc::length_prefixed_to_annex_b(&payload); + if let Some(params) = super::hevc::hvcc_to_annex_b(cp) { es.extend_from_slice(¶ms); } } self.params_written = true; } - let annex_b = super::hevc::length_prefixed_to_annex_b(data); - es.extend_from_slice(&annex_b); + // Append the Annex-B form directly into the pre-sized `es` + // buffer rather than materializing an intermediate Vec. + super::hevc::append_length_prefixed_as_annex_b(&mut es, data); let pes = build_video_pes(pts_90k, &es); self.write_pes(PID_VIDEO, &pes, Some(pcr), keyframe) @@ -200,7 +209,7 @@ impl<W: Write> M2tsMux<W> { if self.audio.is_none() { return Ok(()); } - let pts_90k = self.base_relative_pts(pts_ns); + let pts_90k = self.base_relative_pts(pts_ns, /* may_seed_base */ false); let pes = build_audio_pes(pts_90k, data); self.write_pes(PID_AUDIO, &pes, None, false) } @@ -212,15 +221,24 @@ impl<W: Write> M2tsMux<W> { } /// Convert input PTS (nanoseconds) to 90 kHz ticks rebased on the - /// first frame's PTS. Saturating at 0 keeps the math friendly when - /// frames arrive slightly out of decode order. - fn base_relative_pts(&mut self, pts_ns: i64) -> u64 { + /// stream's PTS origin. The origin is seeded ONLY by the first video + /// frame (`may_seed_base == true`); audio frames never seed it. This + /// keeps the audio/video offset intact: a leading audio frame can't + /// pull the base up and collapse the first/lowest-PTS video frame to 0. + /// Frames earlier than the base saturate to 0. + fn base_relative_pts(&mut self, pts_ns: i64, may_seed_base: bool) -> u64 { let raw_90k = if pts_ns > 0 { - (pts_ns as u64) * 9 / 100_000 + // Widen to u128 so adversarial timestamps can't overflow the + // intermediate multiply (pts_ns * 9 exceeds u64 above + // ~2.05e18 ns), then clamp to the 33-bit PTS range. + (((pts_ns as u128) * 9 / 100_000) as u64) & 0x1_FFFF_FFFF } else { 0 }; - let base = *self.base_pts_90k.get_or_insert(raw_90k); + if may_seed_base { + self.base_pts_90k.get_or_insert(raw_90k); + } + let base = self.base_pts_90k.unwrap_or(raw_90k); raw_90k.saturating_sub(base) } @@ -255,10 +273,13 @@ impl<W: Write> M2tsMux<W> { while offset < pes.len() { self.maybe_emit_psi()?; + // Force a PCR on the FIRST video PES (PAT+PMT precede it, so + // `packets_written` is never 0 here) so a receiver tuning at + // stream start has the clock reference the PMT promises. let attach_pcr = first && (pid == PID_VIDEO) && (pcr.is_some()) - && (self.packets_written == 0 + && (!self.first_video_written || self.video_packets_since_pcr >= PCR_INTERVAL_PACKETS); // RAI rides only the FIRST packet of a keyframe video PES. @@ -282,9 +303,16 @@ impl<W: Write> M2tsMux<W> { // When AF body is empty we can still skip the AF entirely // and get the full 184 B; only invoke the AF when we'd // otherwise need stuffing. + // + // Per ISO/IEC 13818-1 Table 2-6, when an adaptation field is + // present its first body byte is the mandatory 8-bit flags + // byte. A stuffing-only field still needs that flags byte + // (all flags 0) — omitting it would make a strict decoder + // read the first 0xFF stuffing byte as flags (PCR_flag=1, + // …) and parse a phantom PCR out of the stuffing/payload. let (af_present, payload_len, stuffing): (bool, usize, usize) = if !af_body.is_empty() { - // AF is mandatory (PCR). 1 byte length + body + stuffing - // + payload = 184. + // AF is mandatory (PCR, RAI, …). 1 byte length + body + + // stuffing + payload = 184. let max_payload = 184 - 1 - af_body.len(); let p = remaining.min(max_payload); let s = max_payload - p; @@ -293,11 +321,13 @@ impl<W: Write> M2tsMux<W> { // Full payload packet — no AF at all. (false, 184, 0) } else { - // Last (small) packet — stuff via empty AF. - // 1 byte length + 0 body + stuffing + payload = 184. - let max_payload = 183; + // Last (small) packet — stuff via an AF whose body is a + // single zero-flags byte. 1 byte length + 1 flags byte + + // stuffing + payload = 184, so payload caps at 182. + let max_payload = 182; let p = remaining.min(max_payload); let s = max_payload - p; + af_body.push(0x00); // zero-flags byte (true, p, s) }; @@ -305,14 +335,14 @@ impl<W: Write> M2tsMux<W> { let mut packet = Packet::new(); packet.set_header(pid, first, true, af_present, cc); if af_present { - packet.append_adaptation(&af_body, stuffing); + packet.append_adaptation(&af_body, stuffing)?; } - packet.append_payload(&pes[offset..offset + payload_len]); - debug_assert_eq!(packet.len(), 188, "packet not 188 bytes"); + packet.append_payload(&pes[offset..offset + payload_len])?; self.out.write_packet(&packet)?; self.packets_written += 1; if pid == PID_VIDEO { + self.first_video_written = true; if attach_pcr { self.video_packets_since_pcr = 0; } else { @@ -351,7 +381,7 @@ impl<W: Write> M2tsMux<W> { let cc = self.advance_cc(PID_PAT); let mut packet = Packet::new(); packet.set_header(PID_PAT, true, true, false, cc); - packet.append_payload(&payload); + packet.append_payload(&payload)?; packet.pad_to_188(); self.out.write_packet(&packet)?; self.packets_written += 1; @@ -363,7 +393,7 @@ impl<W: Write> M2tsMux<W> { let cc = self.advance_cc(PID_PMT); let mut packet = Packet::new(); packet.set_header(PID_PMT, true, true, false, cc); - packet.append_payload(&payload); + packet.append_payload(&payload)?; packet.pad_to_188(); self.out.write_packet(&packet)?; self.packets_written += 1; @@ -371,42 +401,6 @@ impl<W: Write> M2tsMux<W> { } } -/// Extract the raw hvcC bytes for handoff to `length_prefixed_to_annex_b`. -/// hvcC layout: 22-byte fixed header, then `numOfArrays` arrays of -/// `(nalType, numNalus, [nalLength:u16, NAL bytes]…)`. We convert this -/// directly to a length-prefixed byte stream (NAL length is u16 in -/// hvcC; widen to u32 for the standard length-prefixed encoding). -fn hvcc_payload(hvcc: &[u8]) -> Vec<u8> { - if hvcc.len() < 23 { - return Vec::new(); - } - let num_arrays = hvcc[22] as usize; - let mut out = Vec::new(); - let mut offset = 23; - for _ in 0..num_arrays { - if offset + 3 > hvcc.len() { - break; - } - offset += 1; - let num_nalus = u16::from_be_bytes([hvcc[offset], hvcc[offset + 1]]) as usize; - offset += 2; - for _ in 0..num_nalus { - if offset + 2 > hvcc.len() { - break; - } - let nal_len = u16::from_be_bytes([hvcc[offset], hvcc[offset + 1]]) as usize; - offset += 2; - if offset + nal_len > hvcc.len() { - break; - } - out.extend_from_slice(&(nal_len as u32).to_be_bytes()); - out.extend_from_slice(&hvcc[offset..offset + nal_len]); - offset += nal_len; - } - } - out -} - /// Build a PES packet for a video access unit. fn build_video_pes(pts_90k: u64, es: &[u8]) -> Vec<u8> { build_pes_packet(0xE0, pts_90k, es, /* length_in_header */ false) @@ -414,9 +408,11 @@ fn build_video_pes(pts_90k: u64, es: &[u8]) -> Vec<u8> { /// Build a PES packet for an audio access unit. fn build_audio_pes(pts_90k: u64, es: &[u8]) -> Vec<u8> { - // Audio PES: length is fillable when it fits in u16. We always - // write the length so receivers don't have to scan for the next - // start code. + // Audio PES: a bounded length is written whenever the PES fits in a + // u16 (the common case), so receivers don't have to scan for the next + // start code. For an access unit larger than ~64 KiB (rare — e.g. a + // large TrueHD frame) the length field falls back to the unbounded + // (0x0000) form, which most demuxers tolerate for private_stream_1. build_pes_packet(0xBD, pts_90k, es, /* length_in_header */ true) } @@ -631,6 +627,30 @@ mod tests { assert!(pids.iter().any(|p| *p == PID_VIDEO)); } + #[test] + fn first_video_pes_carries_pcr() { + // A receiver tuning at stream start needs the clock reference the + // PMT promises (video PID = PCR_PID). The very first video PES must + // therefore carry a PCR even though PAT+PMT precede it. + let mut sink: Vec<u8> = Vec::new(); + let mut mux = M2tsMux::new(&mut sink); + let mut frame = Vec::new(); + frame.extend_from_slice(&4u32.to_be_bytes()); + frame.extend_from_slice(&[0x40, 0x01, 0x0C, 0x01]); + mux.write_video(0, true, &frame).unwrap(); + mux.finish().unwrap(); + drop(mux); + + // First PUSI video packet must carry an AF with the PCR flag (0x10). + let pkt = sink + .chunks(188) + .find(|p| u16::from_be_bytes([p[1] & 0x1F, p[2]]) == PID_VIDEO && (p[1] & 0x40) != 0) + .expect("video PUSI packet exists"); + let af = af_body(pkt).expect("first video PES must carry an adaptation field"); + assert!(!af.is_empty(), "AF flags byte present"); + assert_eq!(af[0] & 0x10, 0x10, "PCR flag set on first video PES"); + } + #[test] fn audio_track_appears_in_pmt_and_stream() { let mut sink: Vec<u8> = Vec::new(); @@ -714,6 +734,69 @@ mod tests { Some(packet[5..5 + af_len].to_vec()) } + #[test] + fn stuffing_only_tail_packet_is_spec_valid() { + // A short final PES packet must stuff via an adaptation field + // whose first body byte is the mandatory zero-flags byte (per + // ISO/IEC 13818-1 Table 2-6), never a bare 0xFF stuffing byte + // that a decoder would misread as PCR/OPCR/etc. flags. + // + // Use an AUDIO PES (no PCR, no RAI on its tail) so the only AF on + // the last packet is the stuffing-only field under test. The PES + // is sized so its final TS packet is short (< 184 payload bytes). + let mut sink: Vec<u8> = Vec::new(); + let mut mux = M2tsMux::new(&mut sink); + mux.set_audio(AudioCodec::Ac3); + // Drive a keyframe first so the stream is well-formed, then the + // audio frame whose tail is short. + let mut vframe = Vec::new(); + vframe.extend_from_slice(&4u32.to_be_bytes()); + vframe.extend_from_slice(&[0x40, 0x01, 0x0C, 0x01]); + mux.write_video(0, true, &vframe).unwrap(); + // 200-byte audio payload → PES > 184 → spills into a short tail. + let audio: Vec<u8> = (0..200u32).map(|i| (i & 0xFF) as u8).collect(); + mux.write_audio(20_000_000, &audio).unwrap(); + mux.finish().unwrap(); + drop(mux); + + assert_ts_well_formed(&sink); + + // Find every audio packet that carries an adaptation field; the + // short tail packet is one of them. Each such AF must have a + // length >= 1 and a zero-flags body byte (not 0xFF). + let mut saw_stuffing_af = false; + for pkt in sink.chunks(188) { + let pid = u16::from_be_bytes([pkt[1] & 0x1F, pkt[2]]); + if pid != PID_AUDIO { + continue; + } + let afc = (pkt[3] >> 4) & 0x03; + if afc & 0b10 == 0 { + continue; // no AF on this packet + } + let af_len = pkt[4] as usize; + assert!( + af_len >= 1, + "stuffing AF must include the mandatory flags byte" + ); + // First AF body byte is the flags byte — must be zero, never + // a 0xFF stuffing byte masquerading as flags. + assert_eq!( + pkt[5], 0x00, + "stuffing-only AF flags byte must be 0x00, not 0x{:02X}", + pkt[5] + ); + // adaptation_field_length + payload must fill exactly 184. + // (4 header + 1 AF-length + af_len + payload = 188.) + assert!(af_len <= 183, "AF length overflows the 184-byte body"); + saw_stuffing_af = true; + } + assert!( + saw_stuffing_af, + "expected at least one audio packet with a stuffing AF" + ); + } + #[test] fn rai_set_on_keyframe_pes_packet() { let mut sink: Vec<u8> = Vec::new(); @@ -792,11 +875,12 @@ mod tests { #[test] fn keyframe_video_with_pcr_combines_flags() { - // PCR attaches only when video_packets_since_pcr >= - // PCR_INTERVAL_PACKETS (40). The very first video packet emits a - // PAT+PMT first, so packets_written != 0 and attach_pcr is false on - // frame 0. We push: keyframe (no PCR) → many non-key (drives the - // PCR counter past the interval) → second keyframe (PCR + RAI). + // The first video PES carries a PCR (and RAI) and resets the PCR + // counter. After that, PCR re-attaches only when + // video_packets_since_pcr >= PCR_INTERVAL_PACKETS (40). We push: + // keyframe (PCR+RAI, counter reset) → many non-key (drives the + // counter past the interval) → second keyframe whose PUSI combines + // RAI (keyframe) and PCR (counter exceeded). let mut sink: Vec<u8> = Vec::new(); let mut mux = M2tsMux::new(&mut sink); let mut small = Vec::new(); diff --git a/src/mux/m2ts_mux/packet.rs b/src/mux/m2ts_mux/packet.rs index 81ca3ba..014fb3e 100644 --- a/src/mux/m2ts_mux/packet.rs +++ b/src/mux/m2ts_mux/packet.rs @@ -4,26 +4,50 @@ //! byte layout so the parent module can compose PSI / PCR / PES bytes //! without each caller re-implementing the 188-byte boundary math. +use crate::error::Error; use std::io::{self, Write}; const TS_PACKET_SIZE: usize = 188; +/// Header is 4 bytes, leaving 184 bytes for the adaptation field area +/// plus payload. With a 1-byte `adaptation_field_length` prefix the +/// field body + stuffing can be at most 183 bytes. +const MAX_AF_LEN: usize = TS_PACKET_SIZE - 4 - 1; const SYNC_BYTE: u8 = 0x47; const STUFF_BYTE: u8 = 0xFF; -/// One TS packet under construction. Always emits 188 bytes when -/// [`pad_to_188`](Self::pad_to_188) is called; if it's not called the -/// caller is responsible for filling the packet exactly. +/// One TS packet under construction. Backed by a fixed 188-byte array +/// with a write cursor — no per-packet heap allocation. Always emits 188 +/// bytes when [`pad_to_188`](Self::pad_to_188) is called; if it's not +/// called the caller is responsible for filling the packet exactly. pub(super) struct Packet { - buf: Vec<u8>, + buf: [u8; TS_PACKET_SIZE], + len: usize, } impl Packet { pub(super) fn new() -> Self { Self { - buf: Vec::with_capacity(TS_PACKET_SIZE), + buf: [0u8; TS_PACKET_SIZE], + len: 0, } } + /// Push a byte, saturating at the packet boundary. The boundary is + /// never reached by the sole caller (mod.rs sizes every field to sum + /// to 188); the bound prevents a future caller from corrupting memory. + fn push(&mut self, b: u8) { + if self.len < TS_PACKET_SIZE { + self.buf[self.len] = b; + self.len += 1; + } + } + + fn extend(&mut self, bytes: &[u8]) { + let n = bytes.len().min(TS_PACKET_SIZE - self.len); + self.buf[self.len..self.len + n].copy_from_slice(&bytes[..n]); + self.len += n; + } + /// Write the 4-byte TS packet header. /// /// * `pid` — 13-bit PID @@ -39,12 +63,12 @@ impl Packet { has_adaptation: bool, cc: u8, ) { - self.buf.clear(); - self.buf.push(SYNC_BYTE); + self.len = 0; + self.push(SYNC_BYTE); let pus_bit = if payload_unit_start { 0x40 } else { 0 }; // transport_error_indicator(1)=0 | payload_unit_start(1) | transport_priority(1)=0 | PID(5 high) - self.buf.push(pus_bit | ((pid >> 8) as u8 & 0x1F)); - self.buf.push(pid as u8); + self.push(pus_bit | ((pid >> 8) as u8 & 0x1F)); + self.push(pid as u8); // transport_scrambling_control(2)=0 | adaptation_field_control(2) | continuity_counter(4) let afc = match (has_adaptation, has_payload) { (false, false) => 0b00, // reserved — should not happen @@ -52,7 +76,7 @@ impl Packet { (true, false) => 0b10, // adaptation only (true, true) => 0b11, // both }; - self.buf.push((afc << 4) | (cc & 0x0F)); + self.push((afc << 4) | (cc & 0x0F)); } /// Append the adaptation field after the header. @@ -62,20 +86,36 @@ impl Packet { /// append after the body. The first byte of the field /// (`adaptation_field_length`) is computed here from /// `body.len() + stuffing`. - pub(super) fn append_adaptation(&mut self, body: &[u8], stuffing: usize) { + /// + /// Returns [`Error::M2tsPacketMalformed`] if the computed + /// `adaptation_field_length` would exceed `MAX_AF_LEN` — the length + /// byte and the bytes actually written must always agree, so an + /// over-long field is rejected rather than written with a clamped + /// (and therefore lying) length byte. + pub(super) fn append_adaptation(&mut self, body: &[u8], stuffing: usize) -> io::Result<()> { let af_len = body.len() + stuffing; - debug_assert!(af_len <= 183, "adaptation field overflow"); - self.buf.push(af_len as u8); - self.buf.extend_from_slice(body); - for _ in 0..stuffing { - self.buf.push(STUFF_BYTE); + if af_len > MAX_AF_LEN { + return Err(Error::M2tsPacketMalformed.into()); } + self.push(af_len as u8); + self.extend(body); + for _ in 0..stuffing { + self.push(STUFF_BYTE); + } + Ok(()) } /// Append payload bytes. - pub(super) fn append_payload(&mut self, payload: &[u8]) { - self.buf.extend_from_slice(payload); - debug_assert!(self.buf.len() <= TS_PACKET_SIZE, "packet overflow"); + /// + /// Returns [`Error::M2tsPacketMalformed`] if doing so would push the + /// packet past 188 bytes — overflow is a muxer invariant break, not + /// something to silently emit. + pub(super) fn append_payload(&mut self, payload: &[u8]) -> io::Result<()> { + if self.len + payload.len() > TS_PACKET_SIZE { + return Err(Error::M2tsPacketMalformed.into()); + } + self.extend(payload); + Ok(()) } /// Pad the packet to exactly 188 bytes with `0xFF` bytes — used by @@ -83,21 +123,24 @@ impl Packet { /// For PSI packets only — payload-carrying packets reserve room for /// stuffing via `append_adaptation`. pub(super) fn pad_to_188(&mut self) { - while self.buf.len() < TS_PACKET_SIZE { - self.buf.push(STUFF_BYTE); + while self.len < TS_PACKET_SIZE { + self.push(STUFF_BYTE); } } pub(super) fn bytes(&self) -> &[u8] { - &self.buf + &self.buf[..self.len] } pub(super) fn len(&self) -> usize { - self.buf.len() + self.len } } -/// Buffered writer for assembled TS packets. Owns the underlying sink. +/// Writer for assembled TS packets. Owns the underlying sink and writes +/// each 188-byte packet straight through — it adds no buffering of its +/// own, so callers that need buffering should wrap the sink in a +/// `BufWriter`. pub(super) struct PacketWriter<W: Write> { inner: W, } @@ -109,7 +152,12 @@ impl<W: Write> PacketWriter<W> { pub(super) fn write_packet(&mut self, packet: &Packet) -> io::Result<()> { let bytes = packet.bytes(); - debug_assert_eq!(bytes.len(), TS_PACKET_SIZE); + // Hard check, not a debug_assert: a non-188-byte packet would + // corrupt the transport stream, so refuse to write it in any + // build rather than emitting a short/long packet silently. + if bytes.len() != TS_PACKET_SIZE { + return Err(Error::M2tsPacketMalformed.into()); + } self.inner.write_all(bytes) } @@ -126,7 +174,7 @@ mod tests { fn pad_fills_to_188() { let mut p = Packet::new(); p.set_header(0x100, true, true, false, 0); - p.append_payload(&[1, 2, 3]); + p.append_payload(&[1, 2, 3]).unwrap(); p.pad_to_188(); assert_eq!(p.bytes().len(), 188); assert_eq!(p.bytes()[0], SYNC_BYTE); @@ -134,6 +182,36 @@ mod tests { assert_eq!(p.bytes()[7], STUFF_BYTE); } + #[test] + fn append_adaptation_rejects_overflow() { + let mut p = Packet::new(); + p.set_header(0x100, true, true, true, 0); + // body(1) + stuffing(MAX_AF_LEN) = MAX_AF_LEN + 1 > MAX_AF_LEN. + let err = p.append_adaptation(&[0x00], MAX_AF_LEN).unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + } + + #[test] + fn append_payload_rejects_overflow() { + let mut p = Packet::new(); + p.set_header(0x100, true, true, false, 0); + // 4-byte header + 185 payload = 189 > 188. + let err = p.append_payload(&[0u8; 185]).unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + } + + #[test] + fn write_packet_rejects_short_packet() { + let mut p = Packet::new(); + p.set_header(0x100, true, true, false, 0); + p.append_payload(&[1, 2, 3]).unwrap(); // only 7 bytes, not padded + let mut sink: Vec<u8> = Vec::new(); + let mut w = PacketWriter::new(&mut sink); + let err = w.write_packet(&p).unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(sink.is_empty(), "short packet must not be written"); + } + #[test] fn header_pid_round_trips() { let mut p = Packet::new(); diff --git a/src/mux/meta.rs b/src/mux/meta.rs index 8c66e2e..664b6be 100644 --- a/src/mux/meta.rs +++ b/src/mux/meta.rs @@ -21,9 +21,17 @@ fn color_space_from_hdr(hdr: HdrFormat) -> ColorSpace { } } -/// Magic bytes: "FMKV" + version 1 + 2 reserved bytes. +/// Magic bytes: "FMKV" + 1 reserved byte + version (=1) + 2 reserved bytes. const MAGIC: [u8; 8] = [b'F', b'M', b'K', b'V', 0x00, 0x01, 0x00, 0x00]; +/// Highest header format version this build understands. A header tagged with +/// a newer version is rejected so older readers cleanly refuse incompatible +/// formats instead of silently mis-parsing them as v1. +const SUPPORTED_VERSION: u8 = 1; + +/// Index of the version byte within [`MAGIC`]. +const VERSION_BYTE: usize = 5; + /// BD-TS packet size (header must be padded to this boundary). const PACKET_SIZE: usize = 192; @@ -83,6 +91,11 @@ pub enum MetaStream { label: String, #[serde(default)] secondary: bool, + /// Base64-encoded codec initialization data. Absent for codecs that + /// carry none. Without this, a remux driven from an FMKV header would + /// emit audio tracks missing their init data versus a direct rip. + #[serde(default, skip_serializing_if = "Option::is_none")] + codec_private: Option<String>, }, #[serde(rename = "subtitle")] Subtitle { @@ -92,6 +105,9 @@ pub enum MetaStream { language: String, #[serde(default)] forced: bool, + /// Base64-encoded codec initialization data (e.g. VobSub idx palette). + #[serde(default, skip_serializing_if = "Option::is_none")] + codec_private: Option<String>, }, } @@ -99,6 +115,16 @@ impl M2tsMeta { /// Build metadata from a DiscTitle. Codec privates come from title.codec_privates. pub fn from_title(title: &DiscTitle) -> Self { use base64::Engine; + // Per-stream codec init data, base64-encoded. Preserved for ALL stream + // kinds (video/audio/subtitle) so an FMKV-header-driven remux matches a + // direct disc rip — previously only video round-tripped. + let codec_private_b64 = |i: usize| -> Option<String> { + title + .codec_privates + .get(i) + .and_then(|cp| cp.as_ref()) + .map(|cp| base64::engine::general_purpose::STANDARD.encode(cp)) + }; let streams = title .streams .iter() @@ -113,11 +139,7 @@ impl M2tsMeta { color_space: v.color_space.id().into(), label: v.label.clone(), secondary: v.secondary, - codec_private: title - .codec_privates - .get(i) - .and_then(|cp| cp.as_ref()) - .map(|cp| base64::engine::general_purpose::STANDARD.encode(cp)), + codec_private: codec_private_b64(i), }, Stream::Audio(a) => MetaStream::Audio { pid: a.pid, @@ -127,12 +149,14 @@ impl M2tsMeta { sample_rate: a.sample_rate.to_string(), label: a.label.clone(), secondary: a.secondary, + codec_private: codec_private_b64(i), }, Stream::Subtitle(s) => MetaStream::Subtitle { pid: s.pid, codec: s.codec.id().into(), language: s.language.clone(), forced: s.forced, + codec_private: codec_private_b64(i), }, }) .collect(); @@ -197,6 +221,7 @@ impl M2tsMeta { sample_rate, label, secondary, + codec_private: _, } => Stream::Audio(AudioStream { pid: *pid, codec: codec.parse().unwrap_or(crate::disc::Codec::Unknown(0)), @@ -216,13 +241,14 @@ impl M2tsMeta { codec, language, forced, + codec_private, } => Stream::Subtitle(SubtitleStream { pid: *pid, codec: codec.parse().unwrap_or(crate::disc::Codec::Unknown(0)), language: language.clone(), forced: *forced, qualifier: crate::disc::LabelQualifier::None, - codec_data: None, + codec_data: decode_codec_private(codec_private), }), }) .collect(); @@ -243,32 +269,44 @@ impl M2tsMeta { /// Extract codec_private data per stream (from FMKV header). /// Returns a Vec matching stream order — None for streams without codec_private. + /// Covers all three stream kinds so audio/subtitle init data round-trips, + /// not just video. pub fn codec_privates(&self) -> Vec<Option<Vec<u8>>> { self.streams .iter() .map(|s| { - if let MetaStream::Video { - codec_private: Some(b64), - .. - } = s - { - { - use base64::Engine; - base64::engine::general_purpose::STANDARD.decode(b64).ok() - } - } else { - None - } + let b64 = match s { + MetaStream::Video { codec_private, .. } + | MetaStream::Audio { codec_private, .. } + | MetaStream::Subtitle { codec_private, .. } => codec_private, + }; + decode_codec_private(b64) }) .collect() } } +/// Decode an optional base64 codec_private string into raw bytes. Invalid +/// base64 decodes to `None` (treated as absent) rather than erroring — a +/// corrupt init blob shouldn't fail the whole metadata parse. +fn decode_codec_private(b64: &Option<String>) -> Option<Vec<u8>> { + use base64::Engine; + b64.as_ref() + .and_then(|s| base64::engine::general_purpose::STANDARD.decode(s).ok()) +} + /// Write the metadata header to a writer. Padded to 192-byte boundary. pub fn write_header(w: &mut impl Write, meta: &M2tsMeta) -> io::Result<()> { - let json = serde_json::to_vec(meta).map_err(io::Error::other)?; + // Serializing our own struct effectively cannot fail, but map the + // error to a numeric crate variant rather than embedding serde's + // English string into an io::Error (no-English rule). + let json = serde_json::to_vec(meta).map_err(|_| crate::error::Error::NoMetadata)?; - let json_len = json.len() as u32; + // Guard the length field against truncation: the read side rejects + // anything over MAX_JSON_SIZE, and `as u32` would silently wrap a + // >=4 GiB JSON into a wrong, smaller length. Near-impossible for + // real stream metadata, but a v1.0 primitive shouldn't truncate. + let json_len = u32::try_from(json.len()).map_err(|_| crate::error::Error::NoMetadata)?; let raw_len = 8 + 4 + json.len(); // magic + len + json let padded_len = raw_len.div_ceil(PACKET_SIZE) * PACKET_SIZE; let padding = padded_len - raw_len; @@ -277,7 +315,9 @@ pub fn write_header(w: &mut impl Write, meta: &M2tsMeta) -> io::Result<()> { w.write_all(&json_len.to_be_bytes())?; w.write_all(&json)?; if padding > 0 { - w.write_all(&vec![0u8; padding])?; + // Padding is at most PACKET_SIZE-1 bytes — stack buffer, no heap alloc. + let pad = [0u8; PACKET_SIZE]; + w.write_all(&pad[..padding])?; } Ok(()) } @@ -288,14 +328,37 @@ pub fn write_header(w: &mut impl Write, meta: &M2tsMeta) -> io::Result<()> { pub fn read_header(r: &mut impl Read) -> io::Result<Option<M2tsMeta>> { const MAX_JSON_SIZE: usize = 10 * 1024 * 1024; // 10 MB - let mut magic = [0u8; 8]; - if r.read_exact(&mut magic).is_err() { - return Ok(None); + // Read the first byte alone so a zero-byte stream (a legitimate + // headerless file) stays Ok(None), while a stream that begins with some + // magic bytes then truncates mid-magic surfaces as an error rather than + // being masked as "no header". + let mut first = [0u8; 1]; + if let Err(e) = r.read_exact(&mut first) { + // A clean EOF (no header at all) means "no FMKV header" — the caller + // falls back to a PMT scan. Any OTHER I/O failure (broken pipe, + // permission denied, mid-read disc error) is a real error and must + // propagate, not masquerade as a headerless stream. + if e.kind() == io::ErrorKind::UnexpectedEof { + return Ok(None); + } + return Err(e); } + if first[0] != MAGIC[0] { + return Ok(None); // not an FMKV stream + } + let mut rest = [0u8; 7]; + r.read_exact(&mut rest)?; // started with 'F' but truncated → error + let magic = [ + first[0], rest[0], rest[1], rest[2], rest[3], rest[4], rest[5], rest[6], + ]; if magic[..4] != MAGIC[..4] { return Ok(None); } + if magic[VERSION_BYTE] > SUPPORTED_VERSION { + // Newer, incompatible format — refuse rather than mis-parse as v1. + return Err(crate::error::Error::NoMetadata.into()); + } let mut len_buf = [0u8; 4]; r.read_exact(&mut len_buf)?; @@ -307,16 +370,17 @@ pub fn read_header(r: &mut impl Read) -> io::Result<Option<M2tsMeta>> { let mut json_buf = vec![0u8; json_len]; r.read_exact(&mut json_buf)?; - let meta: M2tsMeta = serde_json::from_slice(&json_buf) - .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + let meta: M2tsMeta = + serde_json::from_slice(&json_buf).map_err(|_| crate::error::Error::NoMetadata)?; - // Skip padding to next 192-byte boundary + // Skip padding to next 192-byte boundary (at most PACKET_SIZE-1 bytes → + // a stack buffer, no heap allocation). let raw_len = 8 + 4 + json_len; let padded_len = raw_len.div_ceil(PACKET_SIZE) * PACKET_SIZE; let padding = padded_len - raw_len; if padding > 0 { - let mut skip = vec![0u8; padding]; - r.read_exact(&mut skip)?; + let mut skip = [0u8; PACKET_SIZE]; + r.read_exact(&mut skip[..padding])?; } Ok(Some(meta)) @@ -406,6 +470,133 @@ mod tests { } } + #[test] + fn read_header_empty_is_none_not_error() { + // No bytes at all → clean EOF on the magic read → Ok(None), the + // "no FMKV header, fall back" signal. + let empty: &[u8] = &[]; + let mut cursor = io::Cursor::new(empty); + let got = read_header(&mut cursor).expect("clean EOF must be Ok(None)"); + assert!(got.is_none()); + } + + #[test] + fn read_header_propagates_non_eof_error() { + // A reader that fails with a non-EOF error must surface that + // error, not be swallowed as Ok(None). + struct BrokenReader; + impl Read for BrokenReader { + fn read(&mut self, _: &mut [u8]) -> io::Result<usize> { + Err(io::Error::from(io::ErrorKind::BrokenPipe)) + } + } + let mut r = BrokenReader; + let err = read_header(&mut r).expect_err("broken pipe must propagate"); + assert_eq!(err.kind(), io::ErrorKind::BrokenPipe); + } + + #[test] + fn write_header_then_read_header_round_trips() { + let title = video_title(HdrFormat::Hdr10, ColorSpace::Bt2020); + let meta = M2tsMeta::from_title(&title); + let mut buf = Vec::new(); + write_header(&mut buf, &meta).expect("write"); + let mut cursor = io::Cursor::new(&buf); + let back = read_header(&mut cursor) + .expect("read") + .expect("header present"); + assert_eq!(back.streams.len(), 1); + // Header is padded to a 192-byte boundary; the cursor must land + // exactly there so the following BD-TS data stays aligned. + assert_eq!(cursor.position() as usize % PACKET_SIZE, 0); + } + + #[test] + fn audio_and_subtitle_codec_private_round_trip() { + use crate::disc::{AudioChannels, AudioStream, LabelPurpose, SampleRate, SubtitleStream}; + let mut t = DiscTitle::empty(); + t.streams.push(Stream::Audio(AudioStream { + pid: 0x1100, + codec: Codec::Dts, + channels: AudioChannels::Surround51, + language: "eng".into(), + sample_rate: SampleRate::S48, + secondary: false, + purpose: LabelPurpose::Normal, + label: String::new(), + })); + t.streams.push(Stream::Subtitle(SubtitleStream { + pid: 0x1200, + codec: Codec::DvdSub, + language: "eng".into(), + forced: false, + qualifier: crate::disc::LabelQualifier::None, + codec_data: None, + })); + // codec_privates: index 0 = audio init data, index 1 = subtitle init data. + t.codec_privates = vec![Some(vec![0xAA, 0xBB, 0xCC]), Some(vec![0x01, 0x02])]; + + let meta = M2tsMeta::from_title(&t); + // Must serialize for both audio and subtitle (not just video). + let cps = meta.codec_privates(); + assert_eq!(cps[0].as_deref(), Some(&[0xAA, 0xBB, 0xCC][..])); + assert_eq!(cps[1].as_deref(), Some(&[0x01, 0x02][..])); + + // And to_title restores the subtitle codec_data from the header. + let back = meta.to_title(); + match &back.streams[1] { + Stream::Subtitle(s) => { + assert_eq!(s.codec_data.as_deref(), Some(&[0x01, 0x02][..])) + } + _ => panic!("expected subtitle stream"), + } + // The round-tripped title also carries all codec_privates. + assert_eq!( + back.codec_privates[0].as_deref(), + Some(&[0xAA, 0xBB, 0xCC][..]) + ); + assert_eq!(back.codec_privates[1].as_deref(), Some(&[0x01, 0x02][..])); + } + + #[test] + fn newer_version_header_rejected() { + // A header tagged with a version above SUPPORTED_VERSION must be + // refused, not silently parsed as v1. + let meta = M2tsMeta::from_title(&video_title(HdrFormat::Sdr, ColorSpace::Bt709)); + let mut buf = Vec::new(); + write_header(&mut buf, &meta).unwrap(); + buf[VERSION_BYTE] = SUPPORTED_VERSION + 1; // bump version byte + let mut cur = io::Cursor::new(buf); + let err = read_header(&mut cur).unwrap_err(); + // NoMetadata (E9008) maps to InvalidInput. + assert_eq!(err.kind(), io::ErrorKind::InvalidInput); + } + + #[test] + fn empty_stream_is_clean_none_but_partial_magic_errors() { + // Zero bytes → no header (Ok(None)). + let mut empty = io::Cursor::new(Vec::<u8>::new()); + assert!(read_header(&mut empty).unwrap().is_none()); + + // Begins with 'F' (MAGIC[0]) then truncates → error, not None. + let mut partial = io::Cursor::new(vec![b'F', b'M', b'K']); + assert!(read_header(&mut partial).is_err()); + + // Does not begin with the FMKV magic at all → Ok(None) (headerless). + let mut other = io::Cursor::new(vec![0x47u8; 16]); + assert!(read_header(&mut other).unwrap().is_none()); + } + + #[test] + fn header_round_trips_through_write_read() { + let meta = M2tsMeta::from_title(&video_title(HdrFormat::Hdr10, ColorSpace::Bt2020)); + let mut buf = Vec::new(); + write_header(&mut buf, &meta).unwrap(); + let mut cur = io::Cursor::new(buf); + let back = read_header(&mut cur).unwrap().expect("header present"); + assert_eq!(back.streams.len(), 1); + } + #[test] fn legacy_sdr_without_color_space_derives_bt709() { let json = r#"{ diff --git a/src/mux/mkv.rs b/src/mux/mkv.rs index 6ede83d..d084c3d 100644 --- a/src/mux/mkv.rs +++ b/src/mux/mkv.rs @@ -55,6 +55,11 @@ pub fn dolby_vision_config(profile: u8, level: u8, bl_compat_id: u8) -> Vec<u8> } impl MkvTrack { + /// Build a video track from a [`VideoStream`]. Language defaults to `"und"`; + /// colour metadata is derived from the stream's colour space and HDR format + /// (PQ for HDR10/HDR10+/DV, HLG for HLG). When `hdr == DolbyVision` a dvcC + /// BlockAdditionMapping is attached automatically so players recognise the + /// Dolby Vision layer. pub fn video(v: &VideoStream) -> Self { let codec_id = match v.codec { Codec::H264 => "V_MPEG4/ISO/AVC", @@ -111,6 +116,9 @@ impl MkvTrack { } } + /// Build an audio track from an [`AudioStream`]. The codec ID follows the + /// Matroska registry; every DTS family member (core, DTS-HD HR, DTS-HD MA) + /// maps to the single registered `A_DTS` ID (see the note below). pub fn audio(a: &AudioStream) -> Self { // The Matroska codec-ID registry defines `A_DTS` for the entire // DTS family — the spec text for `A_DTS` explicitly states it @@ -160,6 +168,10 @@ impl MkvTrack { } } + /// Build a subtitle track from a [`SubtitleStream`]. PGS maps to + /// `S_HDMV/PGS` and DVD VobSub to `S_VOBSUB`; the stream's `codec_data` + /// (the VobSub `.idx` palette header for DVD) becomes the track's + /// CodecPrivate. The forced-display flag is propagated from the stream. pub fn subtitle(s: &SubtitleStream) -> Self { let codec_id = match s.codec { Codec::DvdSub => "S_VOBSUB", @@ -219,6 +231,12 @@ pub struct MkvMuxer<W: Write + Seek> { last_pts_ms: std::collections::HashMap<usize, i64>, cues: Vec<CuePoint>, frame_count: u64, + /// Frames handed to `write_frame` that were dropped because no cluster was + /// open yet (a cluster only opens on a track-0 video keyframe). If this is + /// non-zero at `finish()` and not a single frame was ever written, the + /// caller produced an empty MKV — surfaced as an error rather than a + /// silently empty file. See `write_frame` for the track-0 invariant. + dropped_pre_cluster: u64, seek_fixups: Vec<SeekPositionFixup>, info_offset: u64, tracks_offset: u64, @@ -229,10 +247,15 @@ pub struct MkvMuxer<W: Write + Seek> { const CLUSTER_DURATION_MS: i64 = 5000; /// Maximum block-relative timestamp expressible in the signed 16-bit -/// SimpleBlock/Block field (`i16::MAX` ms). A frame further than this from -/// the open cluster's timestamp forces a new cluster (see `write_frame`) so -/// the `as i16` cast can never wrap. +/// SimpleBlock/Block field (`i16::MAX` ms). A frame whose offset from the open +/// cluster's timestamp falls outside `i16::MIN..=i16::MAX` ms forces a new +/// cluster (see `write_frame`) so the `as i16` cast can never wrap — in EITHER +/// direction. PES timestamps come from untrusted disc/file bytes and can +/// back-jump on discontinuities, so the lower bound matters as much as the +/// upper one. const MAX_BLOCK_REL_MS: i64 = i16::MAX as i64; +/// Minimum block-relative timestamp expressible in the signed 16-bit field. +const MIN_BLOCK_REL_MS: i64 = i16::MIN as i64; /// Force a per-track block timestamp to be strictly later than the previous one /// written for that track. `prev` is the last timestamp for the track (`None` @@ -247,6 +270,28 @@ fn monotonic_ts(prev: Option<i64>, pts_ms: i64) -> i64 { } } +/// Encode a Matroska track number as an EBML VINT into a stack buffer, +/// returning the buffer and the used length. Track numbers are small (1-based, +/// a handful of tracks), so 1 byte covers `< 0x80` and 2 bytes covers the rest; +/// no heap allocation, called once per block on the mux hot path. +/// +/// The 2-byte form holds 14 payload bits (max 0x3FFF). The `debug_assert` +/// guards the 0x4000 bound: at or above it, `(track_num >> 8)` is >= 0x40 and +/// OR-ing the 0x40 length marker would clobber it, corrupting the track +/// number. Not reachable today (track numbers are `i+1` over a few streams), +/// so this documents the bound rather than handling 3-byte VINTs. +fn track_vint(track_num: usize) -> ([u8; 2], usize) { + if track_num < 0x80 { + ([(track_num as u8) | 0x80, 0], 1) + } else { + debug_assert!( + track_num < 0x4000, + "track number {track_num} exceeds the 14-bit 2-byte EBML VINT range" + ); + ([0x40 | ((track_num >> 8) as u8), track_num as u8], 2) + } +} + impl<W: Write + Seek> MkvMuxer<W> { /// Create a new MKV muxer: writes EBML header, Segment start, Info, Tracks, Chapters. pub fn new( @@ -446,6 +491,7 @@ impl<W: Write + Seek> MkvMuxer<W> { last_pts_ms: std::collections::HashMap::new(), cues: Vec::new(), frame_count: 0, + dropped_pre_cluster: 0, seek_fixups, info_offset, tracks_offset, @@ -469,8 +515,40 @@ impl<W: Write + Seek> MkvMuxer<W> { duration_ns: Option<u64>, ) -> io::Result<()> { let raw_ms = pts_ns / 1_000_000; - let base = *self.base_pts_ms.get_or_insert(raw_ms); - let pts_ms = raw_ms - base; + + // Cluster boundaries normally coincide with a video keyframe so every + // Cues entry resolves to a seekable IDR at the cluster start. + let is_video_key = keyframe && track_idx == 0; + + // Derive the timestamp base from the first *kept* keyframe (the frame + // that opens the first cluster), NOT the first frame merely seen. The + // first frame seen can have a higher display PTS than the subsequent + // I-frame (B-frame reordering / a PTS discontinuity), which would make + // later cluster/cue timestamps negative and wrap to ~u64::MAX on the + // `as u64` cast in `start_cluster`/`finish`. Anchoring on the first kept + // keyframe guarantees the open cluster's timestamp is 0 and all later + // relative offsets are computed from a frame we actually wrote. + let base = match self.base_pts_ms { + Some(b) => b, + None => { + if !is_video_key { + // No cluster can open yet (clusters start on a track-0 + // keyframe). Drop this frame as before, but count it so an + // all-dropped run surfaces as an error at finish(). + self.dropped_pre_cluster += 1; + return Ok(()); + } + self.base_pts_ms = Some(raw_ms); + raw_ms + } + }; + // Floor at 0: base is the first kept keyframe, so any frame with an + // earlier PTS (audio/subtitle arriving with a pre-keyframe timestamp, or + // a back-jump on a stream discontinuity) would compute negative here, + // which would wrap to ~u64::MAX on the `as u64` cluster/cue write and + // could overflow the i16 block-relative cast. Frames before the first + // kept keyframe are clamped to t=0 rather than corrupting the timeline. + let pts_ms = (raw_ms - base).max(0); // Enforce strictly-monotonic per-track block timestamps. Some audio PES // PTS truncate to the same millisecond as the previous frame (or, rarely, @@ -479,14 +557,17 @@ impl<W: Write + Seek> MkvMuxer<W> { // and A/V sync is unaffected at millisecond granularity. let pts_ms = monotonic_ts(self.last_pts_ms.get(&track_idx).copied(), pts_ms); - // Cluster boundaries normally coincide with a video keyframe so every - // Cues entry resolves to a seekable IDR at the cluster start. - let is_video_key = keyframe && track_idx == 0; let needs_new_cluster = !self.cluster_open || (is_video_key && (pts_ms - self.cluster_ts_ms) >= CLUSTER_DURATION_MS); if needs_new_cluster { if !is_video_key { + // A cluster is open but this non-keyframe wants a fresh one only + // because !cluster_open is false here — so this branch is the + // "no cluster open and not a keyframe" case. Drop and count. + if !self.cluster_open { + self.dropped_pre_cluster += 1; + } return Ok(()); } self.start_cluster(pts_ms)?; @@ -495,17 +576,25 @@ impl<W: Write + Seek> MkvMuxer<W> { track: track_idx + 1, cluster_pos: self.cluster_pos - self.segment_start, }); - } else if (pts_ms - self.cluster_ts_ms) > MAX_BLOCK_REL_MS { - // The block-relative timestamp is a signed 16-bit value, so a - // frame more than i16::MAX ms (~32.767 s) past the current - // cluster's timestamp would silently wrap on the `as i16` cast, - // corrupting A/V sync. The keyframe-driven boundary above only - // fires on a video keyframe — a long audio-only stretch, or a - // very long GOP with no intervening keyframe, can drift past the - // i16 range. Force a fresh cluster here even without a keyframe - // to keep the cast in range. This cluster is not keyframe-aligned - // so it gets no Cues entry (Cues stay IDR-only for seekability). - self.start_cluster(pts_ms)?; + } else { + let rel = pts_ms - self.cluster_ts_ms; + if !(MIN_BLOCK_REL_MS..=MAX_BLOCK_REL_MS).contains(&rel) { + // The block-relative timestamp is a signed 16-bit value, so a + // frame whose offset from the current cluster's timestamp falls + // outside i16::MIN..=i16::MAX ms (~±32.767 s) would silently wrap + // on the `as i16` cast, corrupting A/V sync. The keyframe-driven + // boundary above only fires on a video keyframe — a long + // audio-only stretch, a very long GOP with no intervening + // keyframe (positive direction), or an audio/subtitle PES whose + // PTS back-jumps below the open cluster (negative direction, e.g. + // a stream discontinuity) can drift past the i16 range. Force a + // fresh cluster here even without a keyframe to keep the cast in + // range. pts_ms is already floored at 0 above, so the new + // cluster timestamp never wraps on the `as u64` write in + // start_cluster. This cluster is not keyframe-aligned so it gets + // no Cues entry (Cues stay IDR-only for seekability). + self.start_cluster(pts_ms)?; + } } // Committed to writing this frame — record its (monotonic) timestamp so @@ -528,7 +617,24 @@ impl<W: Write + Seek> MkvMuxer<W> { } /// Finish the MKV file: write Cues element. + /// + /// # Track-0 invariant + /// + /// A cluster only opens on a track-0 video keyframe, so the caller must + /// supply track 0 as the video track and deliver a keyframe on it before + /// (or alongside) other-track data. If no track-0 keyframe ever arrives, + /// every `write_frame` is silently dropped; rather than emit a structurally + /// valid but empty MKV (zero clusters, zero frames), `finish` returns + /// `Error::MkvInvalid` when frames were submitted but none were written. pub fn finish(mut self) -> io::Result<()> { + // A title that produced no frames (e.g. fully unreadable, or every + // frame dropped before the first track-0 keyframe opened a cluster) + // would otherwise yield a structurally-empty MKV with no clusters or + // cues. Surface that as an error rather than writing valid-but-empty + // output. + if self.frame_count == 0 { + return Err(crate::error::Error::MkvInvalid.into()); + } // Close final cluster self.end_cluster()?; @@ -558,7 +664,9 @@ impl<W: Write + Seek> MkvMuxer<W> { let offset = match fixup.target_id { ebml::INFO => self.info_offset, ebml::TRACKS => self.tracks_offset, - ebml::CHAPTERS => self.chapters_offset.unwrap_or(0), + ebml::CHAPTERS => self + .chapters_offset + .expect("CHAPTERS seek fixup present => chapters_offset is Some"), ebml::CUES => cues_offset, _ => 0, }; @@ -601,19 +709,15 @@ impl<W: Write + Seek> MkvMuxer<W> { data: &[u8], ) -> io::Result<()> { // SimpleBlock: [track_number VINT] [relative_ts i16] [flags u8] [data] - // Track number as EBML VINT - let track_vint = if track_num < 0x80 { - vec![(track_num as u8) | 0x80] - } else { - vec![0x40 | ((track_num >> 8) as u8), track_num as u8] - }; + let (tv, tv_len) = track_vint(track_num); + let track_vint = &tv[..tv_len]; let flags: u8 = if keyframe { 0x80 } else { 0x00 }; let block_size = track_vint.len() + 2 + 1 + data.len(); // vint + ts(2) + flags(1) + data ebml::write_id(&mut self.writer, ebml::SIMPLE_BLOCK)?; ebml::write_size(&mut self.writer, block_size as u64)?; - self.writer.write_all(&track_vint)?; + self.writer.write_all(track_vint)?; self.writer.write_all(&relative_ts.to_be_bytes())?; self.writer.write_all(&[flags])?; self.writer.write_all(data)?; @@ -629,18 +733,21 @@ impl<W: Write + Seek> MkvMuxer<W> { data: &[u8], duration_ms: u64, ) -> io::Result<()> { - let track_vint = if track_num < 0x80 { - vec![(track_num as u8) | 0x80] - } else { - vec![0x40 | ((track_num >> 8) as u8), track_num as u8] - }; - let flags: u8 = if keyframe { 0x80 } else { 0x00 }; + let (tv, tv_len) = track_vint(track_num); + let track_vint = &tv[..tv_len]; + // The 0x80 Keyframe flag is defined only for SimpleBlock; inside a + // Block within a BlockGroup that high bit is reserved and MUST be 0 + // (keyframe-ness is signalled by the absence of a ReferenceBlock + // child). `keyframe` is intentionally unused here — every Block this + // path emits is intra (PGS subtitle frames carrying a duration). + let _ = keyframe; + let flags: u8 = 0x00; let block_size = track_vint.len() + 2 + 1 + data.len(); let bg_pos = ebml::start_master(&mut self.writer, ebml::BLOCK_GROUP)?; ebml::write_id(&mut self.writer, ebml::BLOCK)?; ebml::write_size(&mut self.writer, block_size as u64)?; - self.writer.write_all(&track_vint)?; + self.writer.write_all(track_vint)?; self.writer.write_all(&relative_ts.to_be_bytes())?; self.writer.write_all(&[flags])?; self.writer.write_all(data)?; @@ -820,29 +927,10 @@ mod tests { #[test] fn mkv_finish_writes_cues_element() { - // Use a Vec wrapped in Cursor, then check after finish + // finish() consumes self and flushes the writer, so use the + // module-level SharedWriter to inspect the buffer afterwards. use std::sync::{Arc, Mutex}; - // We'll write to a Cursor, but finish() consumes self. - // The trick: Cursor<Vec<u8>> - we can get data back via into_inner chain. - // But MkvMuxer::finish consumes self and flushes writer. - // We need a way to inspect the output. Let's use a wrapper. - - struct SharedWriter(Arc<Mutex<Cursor<Vec<u8>>>>); - impl Write for SharedWriter { - fn write(&mut self, buf: &[u8]) -> io::Result<usize> { - self.0.lock().unwrap().write(buf) - } - fn flush(&mut self) -> io::Result<()> { - self.0.lock().unwrap().flush() - } - } - impl Seek for SharedWriter { - fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> { - self.0.lock().unwrap().seek(pos) - } - } - let shared = Arc::new(Mutex::new(Cursor::new(Vec::new()))); let writer = SharedWriter(shared.clone()); let tracks = [make_video_track()]; @@ -1597,4 +1685,114 @@ mod tests { } assert_eq!(sb_count, 1, "expected exactly one SimpleBlock in output"); } + + #[test] + fn no_track0_keyframe_yields_error_not_empty_file() { + // If track 0 never delivers a keyframe, every frame is dropped. finish() + // must surface this rather than emitting a structurally valid empty MKV. + let tracks = [make_video_track(), make_audio_track()]; + let shared = Arc::new(Mutex::new(Cursor::new(Vec::new()))); + let writer = SharedWriter(shared.clone()); + let mut muxer = MkvMuxer::new(writer, &tracks, None, 0.0, &[]).unwrap(); + // Audio frames (track 1) and non-keyframe video — no track-0 keyframe. + muxer.write_frame(1, 0, true, &[0xAA; 8], None).unwrap(); + muxer + .write_frame(0, 10_000_000, false, &[0xBB; 8], None) + .unwrap(); + muxer + .write_frame(1, 20_000_000, true, &[0xCC; 8], None) + .unwrap(); + let err = muxer.finish().unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + } + + #[test] + fn finish_with_no_frames_errors() { + // A muxer that received no frames at all must surface MkvInvalid on + // finish() rather than writing a structurally-empty MKV. + let buf = Cursor::new(Vec::new()); + let tracks = [make_video_track()]; + let muxer = MkvMuxer::new(buf, &tracks, None, 60.0, &[]).unwrap(); + let err = muxer.finish().unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + } + + #[test] + fn negative_relative_audio_forces_new_cluster_no_i16_wrap() { + // An audio frame whose PTS back-jumps far below the open cluster (a + // discontinuity) must force a fresh cluster rather than wrap the i16 + // block-relative cast. Build: keyframe at t=0 opening a cluster, a video + // keyframe far later (so cluster ts is large), then an audio frame whose + // PTS lands before that cluster's start by more than i16::MIN ms. + let tracks = [make_video_track(), make_audio_track()]; + // base = 0 (first kept keyframe). Cluster opens at 0; a later keyframe at + // 40s opens a second cluster at ts=40000. Then audio at t=0 → relative + // 0-40000 = -40000 ms, below i16::MIN (-32768) → must open a new cluster. + let frames = vec![ + (0usize, 0i64, true, vec![0x01; 16]), + (0usize, 40_000_000_000i64, true, vec![0x02; 16]), // 40s + (1usize, 0i64, true, vec![0x03; 16]), // back-jumped audio + ]; + let (data, frame_count) = mux_to_bytes(&tracks, &[], &frames); + assert_eq!(frame_count, 3); + let clusters = find_clusters(&data); + // Three clusters: t=0 (video kf), t=40000 (video kf), t=0 (forced for the + // back-jumped audio, no Cues entry). + assert!( + clusters.len() >= 3, + "back-jumped audio must force a fresh cluster, got {} clusters", + clusters.len() + ); + // Every SimpleBlock's relative timestamp must round-trip through i16 + // without the block landing outside the cluster (verified implicitly by + // the muxer never panicking on the `as i16` cast; here we assert the + // forced cluster's timestamp is non-negative so the `as u64` write is + // also safe). + for (_, _, ts) in &clusters { + assert!(*ts <= i64::MAX as u64, "cluster ts must not have wrapped"); + } + } + + #[test] + fn negative_pts_audio_after_keyframe_does_not_wrap() { + // Stream order: video keyframe at 5s (anchors base=5000ms, opens cluster + // at ts 0), then an audio frame with raw PTS 4s — earlier than base. + // raw_ms - base = -1000ms (negative). It must be floored to 0 rather + // than wrapping the `as u64` cluster/cue write or overflowing the i16 + // relative cast. + let tracks = [make_video_track(), make_audio_track()]; + let frames_in_order = [ + (0usize, 5_000_000_000i64, true, vec![0xBB; 16]), // video kf at 5s + (1usize, 4_000_000_000i64, true, vec![0xAA; 8]), // audio at 4s (< base) + ]; + // Do NOT sort — preserve the out-of-order arrival. + let shared = Arc::new(Mutex::new(Cursor::new(Vec::new()))); + let writer = SharedWriter(shared.clone()); + let mut muxer = MkvMuxer::new(writer, &tracks, None, 0.0, &[]).unwrap(); + for (t, pts, kf, data) in &frames_in_order { + muxer.write_frame(*t, *pts, *kf, data, None).unwrap(); + } + muxer.finish().unwrap(); + let data = shared.lock().unwrap().clone().into_inner(); + let clusters = find_clusters(&data); + assert!(!clusters.is_empty()); + for (_, _, ts) in &clusters { + // A wrapped negative would be a huge near-u64::MAX value. + assert!(*ts < 1_000_000_000, "cluster timestamp wrapped: {}", ts); + } + } + + #[test] + fn track_vint_encodes_one_and_two_byte_forms() { + // 1-byte form for track numbers < 0x80, high bit set. + let (b, n) = track_vint(1); + assert_eq!(&b[..n], &[0x81]); + let (b, n) = track_vint(0x7F); + assert_eq!(&b[..n], &[0xFF]); + // 2-byte form at/above 0x80, 0x40 length marker in the top byte. + let (b, n) = track_vint(0x80); + assert_eq!(&b[..n], &[0x40, 0x80]); + let (b, n) = track_vint(0x3FFF); + assert_eq!(&b[..n], &[0x7F, 0xFF]); + } } diff --git a/src/mux/mkvstream.rs b/src/mux/mkvstream.rs index 98c2337..eb28510 100644 --- a/src/mux/mkvstream.rs +++ b/src/mux/mkvstream.rs @@ -6,7 +6,9 @@ use super::mkv::{MkvMuxer, MkvTrack}; use super::{WriteSeek, ebml}; -type MkvHeaderResult = io::Result<(crate::disc::DiscTitle, Vec<(u16, Vec<u8>)>)>; +/// (title, codec_privates, ts_scale_ns) — `ts_scale_ns` is the +/// TimestampScale in nanoseconds per tick, threaded into the frame read path. +type MkvHeaderResult = io::Result<(crate::disc::DiscTitle, Vec<(u16, Vec<u8>)>, i64)>; /// Skip `n` bytes on a forward-only reader (no Seek required). fn skip_bytes(r: &mut impl Read, n: u64) -> io::Result<()> { @@ -62,7 +64,14 @@ use std::io::{self, Read}; struct ReadState { reader: Box<dyn Read + Send>, - cluster_ts_ms: i64, + /// Current cluster timestamp in TimestampScale *ticks* (not ms). Combined + /// with each block's relative tick offset and scaled to nanoseconds via + /// `ts_scale_ns`. + cluster_ts_ticks: i64, + /// TimestampScale in nanoseconds per tick (Matroska INFO/TimestampScale, + /// default 1_000_000 = 1 ms). Foreign MKVs may use a different scale; the + /// frame PTS must honour it, not assume milliseconds. + ts_scale_ns: i64, /// Codec private data per track (track_number, hvcC/avcC bytes). codec_privates: Vec<(u16, Vec<u8>)>, } @@ -123,12 +132,13 @@ impl MkvStream { /// Open an MKV file for reading → PES frames. pub fn open(mut reader: impl Read + Send + 'static) -> io::Result<Self> { - let (disc_title, codec_privates) = parse_mkv_header(&mut reader)?; + let (disc_title, codec_privates, ts_scale_ns) = parse_mkv_header(&mut reader)?; Ok(Self { disc_title, mode: Mode::Read(ReadState { reader: Box::new(reader), - cluster_ts_ms: 0, + cluster_ts_ticks: 0, + ts_scale_ns, codec_privates, }), }) @@ -137,6 +147,7 @@ impl MkvStream { impl crate::pes::Stream for MkvStream { fn read(&mut self) -> io::Result<Option<crate::pes::PesFrame>> { + let streams_len = self.disc_title.streams.len(); let rs = match self.mode { Mode::Read(ref mut rs) => rs, Mode::Write { .. } => return Err(crate::error::Error::StreamWriteOnly.into()), @@ -145,47 +156,95 @@ impl crate::pes::Stream for MkvStream { loop { let (id, size, _) = match ebml::read_element_header(&mut rs.reader) { Ok(h) => h, - Err(_) => return Ok(None), + // Only a genuine premature/clean EOF ends the stream. Any other + // error (disc read failure, corrupt sector, network drop) must + // propagate, or a mid-mux I/O failure would silently truncate + // the output with no error signal. + Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None), + Err(e) => return Err(e), }; match id { ebml::CLUSTER => continue, ebml::CLUSTER_TIMESTAMP => { - rs.cluster_ts_ms = read_uint_bounded(&mut rs.reader, size)? as i64; + let raw = read_uint_bounded(&mut rs.reader, size)?; + // The cluster timestamp is an untrusted u64; a value above + // i64::MAX would cast to a large negative i64 and poison + // every block PTS in the cluster. Reject it, mirroring the + // EBML-size guard in parse_mkv_header. + if raw > i64::MAX as u64 { + return Err(crate::error::Error::MkvInvalid.into()); + } + rs.cluster_ts_ticks = raw as i64; continue; } ebml::SIMPLE_BLOCK => { let block = ebml::read_binary_val(&mut rs.reader, checked_size(size, MAX_BLOCK_SIZE)?)?; - if block.len() < 4 { - continue; + if let Some(frame) = parse_block( + &block, + rs.cluster_ts_ticks, + rs.ts_scale_ns, + streams_len, + None, + ) { + return Ok(Some(frame)); } - - let (track, vl) = block_vint(&block); - if vl + 3 > block.len() { - continue; + continue; + } + ebml::BLOCK_GROUP => { + // MkvMuxer emits a BlockGroup (BLOCK + BLOCK_DURATION) for + // every frame carrying a duration — i.e. all AC3 audio and + // PGS subtitle frames. Descend into the group, read the + // inner BLOCK (0xA1) and BLOCK_DURATION (0x9B), and yield a + // frame so a round-trip through this muxer does not silently + // drop those tracks. A non-u64::MAX size bounds the children. + if size == u64::MAX { + return Err(crate::error::Error::MkvInvalid.into()); } - - let rel_ts = i16::from_be_bytes([block[vl], block[vl + 1]]); - let keyframe = block[vl + 2] & 0x80 != 0; - let data = block[vl + 3..].to_vec(); - let pts_ms = rs.cluster_ts_ms + rel_ts as i64; - let track_idx = (track as usize).saturating_sub(1); // MKV tracks are 1-based - - // Skip blocks for non-existent tracks - if track_idx >= self.disc_title.streams.len() { - continue; + let mut remaining = size; + let mut block: Option<Vec<u8>> = None; + let mut duration_ms: Option<u64> = None; + while remaining > 0 { + let (cid, cs, hlen) = ebml::read_element_header(&mut rs.reader)?; + if cs == u64::MAX { + return Err(crate::error::Error::MkvInvalid.into()); + } + remaining = remaining.saturating_sub(hlen as u64 + cs); + match cid { + ebml::BLOCK => { + block = Some(ebml::read_binary_val( + &mut rs.reader, + checked_size(cs, MAX_BLOCK_SIZE)?, + )?); + } + ebml::BLOCK_DURATION => { + duration_ms = Some(read_uint_bounded(&mut rs.reader, cs)?); + } + _ => skip_bytes(&mut rs.reader, cs)?, + } } - - return Ok(Some(crate::pes::PesFrame { - track: track_idx, - pts: pts_ms * 1_000_000, // ms → ns - keyframe, - data, - duration_ns: None, - })); + if let Some(block) = block { + let dur_ns = duration_ms.map(|ms| ms.saturating_mul(1_000_000)); + if let Some(frame) = parse_block( + &block, + rs.cluster_ts_ticks, + rs.ts_scale_ns, + streams_len, + dur_ns, + ) { + return Ok(Some(frame)); + } + } + continue; } _ => { + // An unknown-size element here would drain the whole stream + // (take(u64::MAX)) and silently drop all later frames; + // reject it like the rest of the parser. + if size == u64::MAX { + return Err(crate::error::Error::MkvInvalid.into()); + } skip_bytes(&mut rs.reader, size)?; continue; } @@ -272,14 +331,27 @@ fn parse_mkv_header(r: &mut impl Read) -> MkvHeaderResult { } let (id, size, _) = match ebml::read_element_header(r) { Ok(h) => h, - Err(_) => break, + Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break, + Err(e) => return Err(e), }; match id { ebml::INFO => { + // An unknown-size (u64::MAX) parent would drain children until + // an EOF read error instead of a clean MkvInvalid; reject it for + // parity with the segment loop guard below. + if size == u64::MAX { + return Err(crate::error::Error::MkvInvalid.into()); + } let mut remaining = size; while remaining > 0 { let (cid, cs, hlen) = ebml::read_element_header(r)?; + // An inner child declaring EBML unknown size (cs == u64::MAX) + // would overflow `hlen + cs` (debug panic) and is meaningless + // for a sized parent — reject it. + if cs == u64::MAX { + return Err(crate::error::Error::MkvInvalid.into()); + } remaining = remaining.saturating_sub(hlen as u64 + cs); match cid { ebml::TIMESTAMP_SCALE => ts_scale = read_uint_bounded(r, cs)?, @@ -293,9 +365,15 @@ fn parse_mkv_header(r: &mut impl Read) -> MkvHeaderResult { got_info = true; } ebml::TRACKS => { + if size == u64::MAX { + return Err(crate::error::Error::MkvInvalid.into()); + } let mut remaining = size; while remaining > 0 { let (cid, cs, hlen) = ebml::read_element_header(r)?; + if cs == u64::MAX { + return Err(crate::error::Error::MkvInvalid.into()); + } remaining = remaining.saturating_sub(hlen as u64 + cs); if cid == ebml::TRACK_ENTRY { let (stream, tnum, cp) = parse_track(r, cs)?; @@ -325,7 +403,38 @@ fn parse_mkv_header(r: &mut impl Read) -> MkvHeaderResult { streams, ..DiscTitle::empty() }; - Ok((disc_title, codec_privates)) + // Clamp the (untrusted) scale to a positive i64 for the tick→ns multiply on + // the read path; default to 1 ms if absent or absurd. + let ts_scale_ns = if ts_scale == 0 || ts_scale > i64::MAX as u64 { + 1_000_000 + } else { + ts_scale as i64 + }; + Ok((disc_title, codec_privates, ts_scale_ns)) +} + +/// Largest valid 13-bit MPEG-TS PID. +const MAX_TS_PID: u32 = 0x1FFF; + +/// Map an MKV track number to a synthetic BD-TS PID, rejecting any value that +/// would overflow the 13-bit PID space. Track 1 is the video PID (0x1011); +/// every other track maps to `0x1100 + (tnum - 2)`. Computed in `u32` so the +/// addition can never wrap, unlike the prior `u16` arithmetic. +fn ts_pid_for_track(tnum: u16) -> io::Result<u16> { + // MKV track numbers are 1-based; 0 is invalid (and would underflow the + // `tnum - 2` below). + if tnum == 0 { + return Err(crate::error::Error::MkvInvalid.into()); + } + let pid: u32 = if tnum == 1 { + 0x1011 + } else { + 0x1100u32 + (tnum as u32 - 2) + }; + if pid > MAX_TS_PID { + return Err(crate::error::Error::MkvInvalid.into()); + } + Ok(pid as u16) } /// Returns (stream, track_number, codec_private_bytes) @@ -341,9 +450,21 @@ fn parse_track( let mut remaining = size; while remaining > 0 { let (cid, cs, hlen) = ebml::read_element_header(r)?; + if cs == u64::MAX { + return Err(crate::error::Error::MkvInvalid.into()); + } remaining = remaining.saturating_sub(hlen as u64 + cs); match cid { - ebml::TRACK_NUMBER => tnum = read_uint_bounded(r, cs)? as u16, + ebml::TRACK_NUMBER => { + // Reject a TRACK_NUMBER above u16::MAX rather than truncating + // with `as u16` (which would alias 65536→0, 65537→1, … onto + // existing small track numbers and corrupt PID/codec lookup). + let n = read_uint_bounded(r, cs)?; + if n > u16::MAX as u64 { + return Err(crate::error::Error::MkvInvalid.into()); + } + tnum = n as u16; + } ebml::TRACK_TYPE => ttype = read_uint_bounded(r, cs)?, ebml::CODEC_ID => codec_id = read_string_bounded(r, cs)?, ebml::CODEC_PRIVATE => { @@ -359,6 +480,9 @@ fn parse_track( let mut vrem = cs; while vrem > 0 { let (vid, vs, vhlen) = ebml::read_element_header(r)?; + if vs == u64::MAX { + return Err(crate::error::Error::MkvInvalid.into()); + } vrem = vrem.saturating_sub(vhlen as u64 + vs); if vid == ebml::PIXEL_HEIGHT { ph = read_uint_bounded(r, vs)? as u32; @@ -371,6 +495,9 @@ fn parse_track( let mut arem = cs; while arem > 0 { let (aid, as_, ahlen) = ebml::read_element_header(r)?; + if as_ == u64::MAX { + return Err(crate::error::Error::MkvInvalid.into()); + } arem = arem.saturating_sub(ahlen as u64 + as_); match aid { ebml::SAMPLING_FREQUENCY => sr = ebml::read_float_val(r, as_ as usize)?, @@ -409,12 +536,11 @@ fn parse_track( SampleRate::S48 }; - // Map MKV track numbers to BD-TS PIDs - let ts_pid = if tnum == 1 { - 0x1011 - } else { - 0x1100 + (tnum - 2) - }; + // Map MKV track numbers to BD-TS PIDs. A 13-bit TS PID tops out at + // 0x1FFF; compute in u32 so the `0x1100 + (tnum - 2)` arithmetic can't + // wrap u16 for large track numbers, and reject anything that would land + // outside the valid PID space. + let ts_pid = ts_pid_for_track(tnum)?; let stream = match ttype { 1 => { @@ -453,6 +579,58 @@ fn parse_track( Ok((stream, tnum, codec_priv)) } +/// Parse a (Simple)Block payload into a PesFrame, or `None` if it should be +/// skipped (too short, track 0, or a track index out of range). +/// +/// `cluster_ts_ticks` is the open cluster's timestamp in TimestampScale ticks +/// and `ts_scale_ns` is that scale (ns per tick); the block PTS is computed as +/// `(cluster_ts_ticks + rel_ts) * ts_scale_ns` so foreign MKVs whose scale +/// isn't 1 ms are honoured (freemkv's own output uses 1_000_000 and round-trips +/// unchanged). `streams_len` bounds the resolved track index; `duration_ns` is +/// propagated for BlockGroup blocks (None for SimpleBlock). +fn parse_block( + block: &[u8], + cluster_ts_ticks: i64, + ts_scale_ns: i64, + streams_len: usize, + duration_ns: Option<u64>, +) -> Option<crate::pes::PesFrame> { + if block.len() < 4 { + return None; + } + let (track, vl) = block_vint(block); + if vl + 3 > block.len() { + return None; + } + // Track 0 is invalid (MKV track numbers are 1-based). block_vint also + // returns 0 for an unsupported 5+ byte VINT, so a corrupt/zero-track block + // must be skipped rather than attributed to the first stream. + if track == 0 { + return None; + } + + let rel_ts = i16::from_be_bytes([block[vl], block[vl + 1]]); + let keyframe = block[vl + 2] & 0x80 != 0; + let data = block[vl + 3..].to_vec(); + let pts_ticks = cluster_ts_ticks + rel_ts as i64; + let track_idx = (track as usize) - 1; // track >= 1 checked above + + // Skip blocks for non-existent tracks. + if track_idx >= streams_len { + return None; + } + + Some(crate::pes::PesFrame { + track: track_idx, + // saturating_mul: a hostile CLUSTER_TIMESTAMP could push pts_ticks near + // i64::MAX, where ticks→ns would overflow and panic in debug builds. + pts: pts_ticks.saturating_mul(ts_scale_ns), + keyframe, + data, + duration_ns, + }) +} + fn block_vint(d: &[u8]) -> (u64, usize) { if d.is_empty() { return (0, 0); @@ -489,12 +667,38 @@ mod tests { // `From<Error> for io::Error` encodes the numeric code into the // Display string as "E{code}: ...". Check the prefix. + /// Extract the error from a `MkvStream::open` result without requiring + /// `MkvStream: Debug` (which `unwrap_err` would). + fn open_err(r: io::Result<MkvStream>) -> io::Error { + match r { + Ok(_) => panic!("expected MkvStream::open to fail"), + Err(e) => e, + } + } + fn is_mkv_invalid(e: &io::Error) -> bool { e.kind() == io::ErrorKind::InvalidData && e.to_string() .starts_with(&format!("E{}", crate::error::E_MKV_INVALID)) } + #[test] + fn ts_pid_for_track_maps_and_rejects_overflow() { + // Track 1 → video PID; track 2 → first audio PID base. + assert_eq!(ts_pid_for_track(1).unwrap(), 0x1011); + assert_eq!(ts_pid_for_track(2).unwrap(), 0x1100); + assert_eq!(ts_pid_for_track(3).unwrap(), 0x1101); + // Highest track that still lands inside the 13-bit PID space. + // 0x1100 + (tnum-2) <= 0x1FFF ⇒ tnum <= 0xF01. + assert_eq!(ts_pid_for_track(0xF01).unwrap(), 0x1FFF); + // One past the edge must be rejected, not wrap u16. + assert!(is_mkv_invalid(&ts_pid_for_track(0xF02).unwrap_err())); + // Former overflow case (debug panic / release garbage PID) is rejected. + assert!(is_mkv_invalid(&ts_pid_for_track(u16::MAX).unwrap_err())); + // Track 0 is invalid (1-based) and would underflow tnum-2. + assert!(is_mkv_invalid(&ts_pid_for_track(0).unwrap_err())); + } + #[test] fn checked_size_rejects_over_cap() { // Within cap → Ok with usize value. @@ -614,4 +818,274 @@ mod tests { assert!(frame.keyframe); assert_eq!(frame.data, vec![0xAA, 0xBB, 0xCC, 0xDD]); } + #[test] + fn truncated_simple_block_body_errors_not_panics() { + // A SIMPLE_BLOCK that declares a 64-byte payload but supplies none. + // read_exact_bounded must surface a clean typed MkvInvalid error + // (a truncated declared element is malformed input), never panic, + // and never allocate the full declared size up front. + let mut cluster = Vec::new(); + ebml::write_id(&mut cluster, ebml::CLUSTER).unwrap(); + ebml::write_unknown_size(&mut cluster).unwrap(); + ebml::write_id(&mut cluster, ebml::SIMPLE_BLOCK).unwrap(); + ebml::write_size(&mut cluster, 64).unwrap(); + // No body bytes follow → short read. + let bytes = minimal_mkv_with_cluster(&cluster); + + let mut stream = MkvStream::open(Cursor::new(bytes)).unwrap(); + let e = stream.read().unwrap_err(); + assert!(is_mkv_invalid(&e)); + } + + /// Build a minimal MKV header + Segment + Info, then a Tracks element with a + /// single TRACK_ENTRY of the given track number/type, then the cluster bytes. + fn mkv_with_track_and_cluster(tnum: u64, ttype: u64, cluster_body: &[u8]) -> Vec<u8> { + let mut out = Vec::new(); + ebml::write_id(&mut out, ebml::EBML).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + ebml::write_id(&mut out, ebml::SEGMENT).unwrap(); + ebml::write_unknown_size(&mut out).unwrap(); + ebml::write_id(&mut out, ebml::INFO).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + + let mut entry = Vec::new(); + ebml::write_uint(&mut entry, ebml::TRACK_NUMBER, tnum).unwrap(); + ebml::write_uint(&mut entry, ebml::TRACK_TYPE, ttype).unwrap(); + let mut track_entry = Vec::new(); + ebml::write_id(&mut track_entry, ebml::TRACK_ENTRY).unwrap(); + ebml::write_size(&mut track_entry, entry.len() as u64).unwrap(); + track_entry.extend_from_slice(&entry); + ebml::write_id(&mut out, ebml::TRACKS).unwrap(); + ebml::write_size(&mut out, track_entry.len() as u64).unwrap(); + out.extend_from_slice(&track_entry); + + out.extend_from_slice(cluster_body); + out + } + + #[test] + fn oversized_codec_private_is_rejected() { + // A TRACK_ENTRY whose CODEC_PRIVATE declares a payload above + // MAX_CODEC_PRIVATE must be rejected (MkvInvalid) before any + // multi-MB allocation, while parsing the header. + let mut entry = Vec::new(); + ebml::write_uint(&mut entry, ebml::TRACK_NUMBER, 1).unwrap(); + ebml::write_uint(&mut entry, ebml::TRACK_TYPE, 1).unwrap(); + // CODEC_PRIVATE header claiming a huge size (no body needed — the + // size check fires first). + ebml::write_id(&mut entry, ebml::CODEC_PRIVATE).unwrap(); + ebml::write_size(&mut entry, MAX_CODEC_PRIVATE + 1).unwrap(); + let mut track_entry = Vec::new(); + ebml::write_id(&mut track_entry, ebml::TRACK_ENTRY).unwrap(); + ebml::write_size(&mut track_entry, entry.len() as u64).unwrap(); + track_entry.extend_from_slice(&entry); + + let mut out = Vec::new(); + ebml::write_id(&mut out, ebml::EBML).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + ebml::write_id(&mut out, ebml::SEGMENT).unwrap(); + ebml::write_unknown_size(&mut out).unwrap(); + ebml::write_id(&mut out, ebml::INFO).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + ebml::write_id(&mut out, ebml::TRACKS).unwrap(); + ebml::write_size(&mut out, track_entry.len() as u64).unwrap(); + out.extend_from_slice(&track_entry); + + let e = match MkvStream::open(Cursor::new(out)) { + Ok(_) => panic!("expected MkvInvalid, got Ok"), + Err(e) => e, + }; + assert!(is_mkv_invalid(&e)); + } + + #[test] + fn block_group_frame_round_trips_with_duration() { + // MkvMuxer emits AC3/PGS frames as a BlockGroup (BLOCK + BLOCK_DURATION). + // The reader must descend into the group and yield the frame (with its + // duration) rather than skipping it — otherwise every AC3/PGS frame this + // muxer writes is lost on read-back. + let block = [0x82u8, 0x00, 0x05, 0x00, 0x11, 0x22, 0x33]; // track 2, rel 5, not-kf, 3 data + let mut bg_body = Vec::new(); + ebml::write_id(&mut bg_body, ebml::BLOCK).unwrap(); + ebml::write_size(&mut bg_body, block.len() as u64).unwrap(); + bg_body.extend_from_slice(&block); + ebml::write_uint(&mut bg_body, ebml::BLOCK_DURATION, 40).unwrap(); // 40 ms + + let mut cluster = Vec::new(); + ebml::write_id(&mut cluster, ebml::CLUSTER).unwrap(); + ebml::write_unknown_size(&mut cluster).unwrap(); + // CLUSTER_TIMESTAMP = 100 ms so pts = (100 + 5) ms. + ebml::write_uint(&mut cluster, ebml::CLUSTER_TIMESTAMP, 100).unwrap(); + ebml::write_id(&mut cluster, ebml::BLOCK_GROUP).unwrap(); + ebml::write_size(&mut cluster, bg_body.len() as u64).unwrap(); + cluster.extend_from_slice(&bg_body); + + // Track 2 (audio) so track_idx 1 needs two streams; give two TRACK_ENTRYs. + // Reuse the helper for track 1, then a manual second entry would be + // simpler — instead build directly with two entries. + let mut out = Vec::new(); + ebml::write_id(&mut out, ebml::EBML).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + ebml::write_id(&mut out, ebml::SEGMENT).unwrap(); + ebml::write_unknown_size(&mut out).unwrap(); + ebml::write_id(&mut out, ebml::INFO).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + let mut tracks = Vec::new(); + for (n, t) in [(1u64, 1u64), (2u64, 2u64)] { + let mut entry = Vec::new(); + ebml::write_uint(&mut entry, ebml::TRACK_NUMBER, n).unwrap(); + ebml::write_uint(&mut entry, ebml::TRACK_TYPE, t).unwrap(); + ebml::write_id(&mut tracks, ebml::TRACK_ENTRY).unwrap(); + ebml::write_size(&mut tracks, entry.len() as u64).unwrap(); + tracks.extend_from_slice(&entry); + } + ebml::write_id(&mut out, ebml::TRACKS).unwrap(); + ebml::write_size(&mut out, tracks.len() as u64).unwrap(); + out.extend_from_slice(&tracks); + out.extend_from_slice(&cluster); + + let mut stream = MkvStream::open(Cursor::new(out)).unwrap(); + let frame = stream + .read() + .unwrap() + .expect("BlockGroup frame must be read"); + assert_eq!(frame.track, 1, "track 2 → index 1"); + assert!(!frame.keyframe); + assert_eq!(frame.data, vec![0x11, 0x22, 0x33]); + assert_eq!(frame.pts, 105 * 1_000_000, "pts = (cluster 100 + rel 5) ms"); + assert_eq!(frame.duration_ns, Some(40 * 1_000_000)); + } + + #[test] + fn track_number_zero_is_rejected() { + // A TRACK_ENTRY with TRACK_NUMBER 0 must be rejected (the ts_pid + // computation would underflow `tnum - 2`). + let bytes = mkv_with_track_and_cluster(0, 1, &[]); + let e = open_err(MkvStream::open(Cursor::new(bytes))); + assert!(is_mkv_invalid(&e)); + } + + #[test] + fn track_number_above_u16_is_rejected() { + // 65536 would truncate to 0 via `as u16` and then underflow. + let bytes = mkv_with_track_and_cluster(65536, 1, &[]); + let e = open_err(MkvStream::open(Cursor::new(bytes))); + assert!(is_mkv_invalid(&e)); + } + + #[test] + fn unknown_size_inner_child_in_tracks_is_rejected() { + // A TRACK_ENTRY child declaring EBML unknown size (cs == u64::MAX) must + // be rejected, not used in `hlen + cs` (which would overflow → debug + // panic). Hand-build a TRACK_ENTRY whose first child carries the + // unknown-size marker. + let mut entry = Vec::new(); + ebml::write_id(&mut entry, ebml::TRACK_NUMBER).unwrap(); + ebml::write_unknown_size(&mut entry).unwrap(); // child size = unknown + + let mut tracks = Vec::new(); + ebml::write_id(&mut tracks, ebml::TRACK_ENTRY).unwrap(); + ebml::write_size(&mut tracks, entry.len() as u64).unwrap(); + tracks.extend_from_slice(&entry); + + let mut out = Vec::new(); + ebml::write_id(&mut out, ebml::EBML).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + ebml::write_id(&mut out, ebml::SEGMENT).unwrap(); + ebml::write_unknown_size(&mut out).unwrap(); + ebml::write_id(&mut out, ebml::INFO).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + ebml::write_id(&mut out, ebml::TRACKS).unwrap(); + ebml::write_size(&mut out, tracks.len() as u64).unwrap(); + out.extend_from_slice(&tracks); + + let e = open_err(MkvStream::open(Cursor::new(out))); + assert!(is_mkv_invalid(&e)); + } + + #[test] + fn oversized_title_string_is_rejected() { + // INFO/TITLE declaring a string above MAX_STRING_LEN must be + // rejected during header parse, not allocated. + let mut info = Vec::new(); + ebml::write_id(&mut info, ebml::TITLE).unwrap(); + ebml::write_size(&mut info, MAX_STRING_LEN + 1).unwrap(); + + let mut out = Vec::new(); + ebml::write_id(&mut out, ebml::EBML).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + ebml::write_id(&mut out, ebml::SEGMENT).unwrap(); + ebml::write_unknown_size(&mut out).unwrap(); + ebml::write_id(&mut out, ebml::INFO).unwrap(); + ebml::write_size(&mut out, info.len() as u64).unwrap(); + out.extend_from_slice(&info); + + let e = match MkvStream::open(Cursor::new(out)) { + Ok(_) => panic!("expected MkvInvalid, got Ok"), + Err(e) => e, + }; + assert!(is_mkv_invalid(&e)); + } + + #[test] + fn read_uint_val_len_nine_errors_not_panics() { + // Direct helper test: an EBML uint cannot exceed 8 bytes. len=9 + // would index past the fixed 8-byte stack buffer and panic on + // untrusted input; it must return MkvInvalid instead. + let mut data = Cursor::new(vec![0u8; 16]); + let e = ebml::read_uint_val(&mut data, 9).unwrap_err(); + assert!(is_mkv_invalid(&e)); + } + + #[test] + fn read_float_val_bad_width_errors() { + // EBML floats are exactly 0, 4, or 8 bytes. Any other width is + // malformed and must error rather than over- or under-read. + let mut data = Cursor::new(vec![0u8; 16]); + let e = ebml::read_float_val(&mut data, 5).unwrap_err(); + assert!(is_mkv_invalid(&e)); + // 0/4/8 remain valid widths. + let mut z = Cursor::new(vec![0u8; 16]); + assert_eq!(ebml::read_float_val(&mut z, 0).unwrap(), 0.0); + let mut f4 = Cursor::new(vec![0u8; 16]); + assert!(ebml::read_float_val(&mut f4, 4).is_ok()); + let mut f8 = Cursor::new(vec![0u8; 16]); + assert!(ebml::read_float_val(&mut f8, 8).is_ok()); + } + + #[test] + fn non_utf8_string_element_is_rejected() { + // A string element with invalid UTF-8 bytes must surface a numeric + // MkvInvalid error, not an io::Error wrapping the FromUtf8Error + // English message (library no-English rule). + let mut data = Cursor::new(vec![0xFF, 0xFE, 0xFD, 0xFC]); + let e = ebml::read_string_val(&mut data, 4).unwrap_err(); + assert!(is_mkv_invalid(&e)); + } + + #[test] + fn simple_block_track_zero_is_skipped() { + // A SimpleBlock with track vint 0 must be skipped, not attributed to + // track 0. Build one track, then a cluster whose only block is track 0 + // followed by a valid track-1 block; read() must return the track-1 one. + let mut cluster = Vec::new(); + ebml::write_id(&mut cluster, ebml::CLUSTER).unwrap(); + ebml::write_unknown_size(&mut cluster).unwrap(); + // track vint 0 is not directly encodable (0x80 is track 0 → block_vint + // returns (0,1)); use 0x80 as the track byte. + let bad = [0x80u8, 0x00, 0x00, 0x80, 0xEE]; + ebml::write_id(&mut cluster, ebml::SIMPLE_BLOCK).unwrap(); + ebml::write_size(&mut cluster, bad.len() as u64).unwrap(); + cluster.extend_from_slice(&bad); + let good = [0x81u8, 0x00, 0x00, 0x80, 0xAB, 0xCD]; + ebml::write_id(&mut cluster, ebml::SIMPLE_BLOCK).unwrap(); + ebml::write_size(&mut cluster, good.len() as u64).unwrap(); + cluster.extend_from_slice(&good); + + let bytes = mkv_with_track_and_cluster(1, 1, &cluster); + let mut stream = MkvStream::open(Cursor::new(bytes)).unwrap(); + let frame = stream.read().unwrap().expect("track-1 frame expected"); + assert_eq!(frame.track, 0); + assert_eq!(frame.data, vec![0xAB, 0xCD]); + } } diff --git a/src/mux/mod.rs b/src/mux/mod.rs index 77a35a8..f3a5aa5 100644 --- a/src/mux/mod.rs +++ b/src/mux/mod.rs @@ -1,7 +1,15 @@ //! Stream-based I/O pipeline. //! -//! All formats are PES streams. Read from a format → PES frames. -//! Write PES frames → a format. +//! Two muxer families live here: +//! +//! 1. **Bidirectional PES streams** (`disc`, `mkv`, `m2ts`, `network`, +//! `stdio`, `null`) implement the [`crate::pes::Stream`] interface: +//! read a format → PES frames, or write PES frames → a format. +//! 2. **Write-only sequential-sink muxers** (`fmp4`, `hevc`, +//! `m2ts_mux`) consume PES frames and write a container to a +//! `SequentialSink`; they do not implement the read loop below. +//! +//! The bidirectional family is driven like this: //! //! ```text //! let mut input = input("iso://Disc.iso", &opts)?; @@ -16,12 +24,26 @@ //! For disc→ISO (raw sector copy), use `Disc::copy()` instead. // Public modules — types here are intentionally part of the consumable API. -pub mod codec; -pub mod demux_thread; pub mod disc; pub mod pipelined_stream; pub mod resolve; +// Internal-only modules. Every reference is via `crate::mux::…` / +// `super::…` from inside the crate; nothing in the downstream crates or +// integration tests imports them and lib.rs re-exports nothing from +// them, so they are not part of the stable public API. +// +// `#[allow(dead_code)]`: narrowing these from `pub` to `pub(crate)` +// surfaces a handful of helpers/accessors that were only ever reachable +// as (unused) public API — e.g. the MPEG-2 resolution/frame-rate +// accessors and an alternate `DemuxThread` spawn path. They are kept as +// part of the parser/demux surface and covered by unit tests; allow the +// dead-code lint rather than delete still-relevant scaffolding. +#[allow(dead_code)] +pub(crate) mod codec; +#[allow(dead_code)] +pub(crate) mod demux_thread; + // Internal modules — implementation details. Their *types* are re-exported // where appropriate (`MkvStream`, `M2tsStream`, etc. surface from `lib.rs`), // but the module paths themselves are not part of the API. Pre-0.13 these @@ -35,16 +57,30 @@ pub(crate) mod m2ts; /// Exposed for integration tests that exercise the wire format directly. pub mod meta; -// ── Phase 3 sequential muxers ────────────────────────────────────────────── +// ── Sequential-sink muxers ────────────────────────────────────────────────── // -// New container muxers that consume PES frames and write to a -// `SequentialSink`. They are NOT refactors of the existing `MkvStream` / -// `M2tsStream` (which round-trip via the legacy `Stream` trait + the -// BD-TS framing); they're sequential-only and target the Phase 2 sink -// split end-to-end. -pub mod fmp4; -pub mod hevc; -pub mod m2ts_mux; +// Container muxers that consume PES frames and write to a +// `SequentialSink`. They are NOT the bidirectional `MkvStream` / +// `M2tsStream` (which round-trip via the `Stream` trait + BD-TS +// framing); these are write-only and sequential. +// +// `pub(crate)`: these have no external callers and are not re-exported +// from lib.rs. `fmp4` is an explicit STUB (`Fmp4Mux::write_video` +// accumulates and discards) — shipping it as `pub` would lock a +// half-built type into the v1.0 stability contract via the +// `libfreemkv::mux::fmp4::Fmp4Mux` path. `m2ts_mux` is the plain +// MPEG-TS sequential muxer and `hevc` is its Annex-B helper; both are +// staged scaffolding for the sink split and are not yet wired into a +// live pipeline (the production paths use `tsmux` / `mkv`). +// `#[allow(dead_code)]`: retained intentionally until the sink split +// lands; they are exercised by their own unit tests. If any becomes a +// public muxer, re-export its concrete type from lib.rs instead. +#[allow(dead_code)] +pub(crate) mod fmp4; +#[allow(dead_code)] +pub(crate) mod hevc; +#[allow(dead_code)] +pub(crate) mod m2ts_mux; pub(crate) mod mkv; pub(crate) mod mkvstream; pub(crate) mod network; diff --git a/src/mux/network.rs b/src/mux/network.rs index 5be4bfe..e32f2cb 100644 --- a/src/mux/network.rs +++ b/src/mux/network.rs @@ -35,6 +35,11 @@ impl NetworkStream { /// Sends FMKV metadata header on first write. pub fn connect(addr: &str) -> io::Result<Self> { let stream = TcpStream::connect(addr)?; + // The sender is the latency-sensitive side; set nodelay here too + // (the listen side already does) so the final sub-MSS flush after + // finish() isn't held by Nagle. The 256 KB BufWriter coalesces + // bulk writes, so this only affects the tail. + stream.set_nodelay(true)?; Ok(Self { disc_title: DiscTitle::empty(), mode: Mode::Write { @@ -44,7 +49,13 @@ impl NetworkStream { }) } - /// Set stream metadata (for write side). Returns self for chaining. + /// Set stream metadata (write side only). Returns self for chaining. + /// + /// Only meaningful on a [`connect`](Self::connect)-constructed + /// (write) stream — the title is sent in the FMKV header on first + /// write. On a [`listen`](Self::listen)-constructed (read) stream + /// the stored title is immediately overwritten by the header read in + /// `listen()`, so calling `meta()` there is a silent no-op. pub fn meta(mut self, dt: &DiscTitle) -> Self { self.disc_title = dt.clone(); self @@ -52,8 +63,19 @@ impl NetworkStream { /// Listen for an incoming connection and read from it. /// Extracts FMKV metadata header from the sender. + /// + /// Accepts exactly one connection; the listening socket is dropped after + /// `accept`, so the bound port is freed and any subsequent connection + /// attempt to the same address is refused. pub fn listen(addr: &str) -> io::Result<Self> { - let listener = TcpListener::bind(addr)?; + Self::accept_from(TcpListener::bind(addr)?) + } + + /// Accept one connection from an already-bound listener and read from it. + /// Lets a caller bind first (learning the actual port for an ephemeral + /// `:0` bind) and hand the listener in, closing the bind/drop/re-bind race + /// that `listen(addr)` would otherwise have. + pub fn accept_from(listener: TcpListener) -> io::Result<Self> { let (stream, _peer) = listener.accept()?; stream.set_nodelay(true)?; let mut reader = BufReader::with_capacity(NET_BUF_SIZE, stream); @@ -70,6 +92,23 @@ impl NetworkStream { } } +/// Write the FMKV metadata header exactly once, before any frames. Always +/// writes (even when the title has no streams) so the receiver's +/// `read_header()` always finds the magic and never falls into the +/// NoMetadata path on a zero-frame stream. +fn ensure_header_written( + writer: &mut BufWriter<TcpStream>, + header_written: &mut bool, + disc_title: &DiscTitle, +) -> io::Result<()> { + if !*header_written { + let m = meta::M2tsMeta::from_title(disc_title); + meta::write_header(writer, &m)?; + *header_written = true; + } + Ok(()) +} + impl crate::pes::Stream for NetworkStream { fn read(&mut self) -> io::Result<Option<crate::pes::PesFrame>> { match &mut self.mode { @@ -82,22 +121,24 @@ impl crate::pes::Stream for NetworkStream { Mode::Write { writer, header_written, - .. } => { - if !*header_written { - if !self.disc_title.streams.is_empty() { - let m = meta::M2tsMeta::from_title(&self.disc_title); - meta::write_header(&mut *writer, &m)?; - } - *header_written = true; - } + ensure_header_written(writer, header_written, &self.disc_title)?; frame.serialize(writer) } _ => Err(crate::error::Error::StreamReadOnly.into()), } } fn finish(&mut self) -> io::Result<()> { - if let Mode::Write { writer, .. } = &mut self.mode { + if let Mode::Write { + writer, + header_written, + } = &mut self.mode + { + // Always emit the FMKV header before shutdown, even for a + // zero-frame stream (e.g. a title that produced no PES frames). + // Without it the receiver's read_header() sees a clean EOF and + // rejects the stream with NoMetadata. + ensure_header_written(writer, header_written, &self.disc_title)?; writer.flush()?; writer.get_ref().shutdown(std::net::Shutdown::Write)?; } @@ -156,19 +197,21 @@ mod tests { } #[test] - #[ignore] // Requires TCP; may be flaky in CI environments fn network_pes_roundtrip() { use crate::pes; + use std::sync::mpsc; + // The listener thread owns the bound socket and reports its actual + // local address back over a channel before accept(). The main thread + // connects only after receiving the address — no bind/drop/re-bind + // window, no sleep-as-synchronisation. let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let port = listener.local_addr().unwrap().port(); - drop(listener); - - let addr = format!("127.0.0.1:{}", port); - let addr_clone = addr.clone(); + let addr = listener.local_addr().unwrap(); + let (addr_tx, addr_rx) = mpsc::channel(); let handle = std::thread::spawn(move || { - let mut ns = NetworkStream::listen(&addr_clone).unwrap(); + addr_tx.send(addr).unwrap(); + let mut ns = NetworkStream::accept_from(listener).unwrap(); let info = pes::Stream::info(&ns).clone(); let mut frames = Vec::new(); while let Ok(Some(f)) = pes::Stream::read(&mut ns) { @@ -177,10 +220,9 @@ mod tests { (info, frames) }); - std::thread::sleep(std::time::Duration::from_millis(50)); - + let addr = addr_rx.recv().unwrap(); let dt = sample_title(); - let mut writer = NetworkStream::connect(&addr).unwrap().meta(&dt); + let mut writer = NetworkStream::connect(&addr.to_string()).unwrap().meta(&dt); let frame = pes::PesFrame { track: 0, pts: 90000, @@ -199,6 +241,38 @@ mod tests { assert_eq!(frames[0].pts, 90000); } + #[test] + fn network_zero_frame_finish_still_sends_header() { + use crate::pes; + use std::sync::mpsc; + + // A title that produces no PES frames must still send the FMKV header + // on finish(), so the receiver gets the metadata instead of rejecting + // the stream with NoMetadata on a clean EOF. + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let (addr_tx, addr_rx) = mpsc::channel(); + + let handle = std::thread::spawn(move || { + addr_tx.send(addr).unwrap(); + // listen()'s read_header must succeed (header present), not error. + let ns = NetworkStream::accept_from(listener).unwrap(); + pes::Stream::info(&ns).playlist.clone() + }); + + let addr = addr_rx.recv().unwrap(); + let dt = sample_title(); + let mut writer = NetworkStream::connect(&addr.to_string()).unwrap().meta(&dt); + // No write() at all — straight to finish(). + pes::Stream::finish(&mut writer).unwrap(); + + let playlist = handle.join().unwrap(); + assert_eq!( + playlist, "NetworkTest", + "zero-frame finish() must still deliver the metadata header" + ); + } + #[test] fn network_empty_addr_errors() { let result = NetworkStream::connect(""); diff --git a/src/mux/null.rs b/src/mux/null.rs index 340487e..e271e34 100644 --- a/src/mux/null.rs +++ b/src/mux/null.rs @@ -18,7 +18,10 @@ impl NullStream { impl crate::pes::Stream for NullStream { fn read(&mut self) -> io::Result<Option<crate::pes::PesFrame>> { - Ok(None) + // Write-only sink: per the Stream trait contract, read() on a + // write-opened stream returns StreamWriteOnly. Returning Ok(None) + // would be misread as a legitimate empty stream. + Err(crate::error::Error::StreamWriteOnly.into()) } fn write(&mut self, _: &crate::pes::PesFrame) -> io::Result<()> { Ok(()) @@ -53,4 +56,14 @@ mod tests { let _ = sink.info(); sink.finish().unwrap(); } + + /// read() on the write-only NullStream must return StreamWriteOnly, + /// not Ok(None) (which a caller would misread as an empty stream). + #[test] + fn read_returns_write_only_error() { + let title = DiscTitle::empty(); + let mut sink = NullStream::new(&title); + let err = Stream::read(&mut sink).expect_err("read on a sink must error"); + assert_eq!(err.kind(), io::ErrorKind::Unsupported); + } } diff --git a/src/mux/pipelined_stream.rs b/src/mux/pipelined_stream.rs index c7a95c1..ee77f5b 100644 --- a/src/mux/pipelined_stream.rs +++ b/src/mux/pipelined_stream.rs @@ -19,11 +19,14 @@ //! recycled buffer pools — no allocations or memcpys in the steady- //! state hot loop. //! -//! This is the *only* read-side `Stream` impl in tree. Both ISO file -//! mux ([`crate::mux::resolve`]) and BD-TS file mux ([`crate::mux::M2tsStream`]) -//! return a `PipelinedPesStream`; the differences are in how the -//! producer thread (A) is configured — sector-aligned reads with -//! AACS decrypt for ISO, raw byte reads for M2TS. +//! This is the *only* read-side `Stream` impl in tree. Both the ISO +//! file mux and the BD-TS (`m2ts://`) file mux input paths are built by +//! [`crate::mux::resolve`] (`build_iso_pipeline` / the m2ts pipeline +//! builder) and hand back a `PipelinedPesStream`; the differences are +//! in how the producer thread (A) is configured — sector-aligned reads +//! with AACS decrypt for ISO, raw byte reads for M2TS. +//! ([`crate::mux::M2tsStream`] itself is a write-only sink and does not +//! construct this type.) use super::codec::CodecParser; use super::demux_thread::{DemuxBatch, DemuxThread}; @@ -48,6 +51,11 @@ pub struct PipelinedPesStream { pending_frames: std::collections::VecDeque<PesFrame>, eof: bool, + /// Cached `FREEMKV_SKIP_PARSE` profiling flag. Read once in `new()` + /// — the env var cannot change for the life of the stream, and + /// `std::env::var_os` takes a process-wide lock, so the per-batch / + /// per-poll reads it replaces were needless hot-path overhead. + skip_parse: bool, } impl PipelinedPesStream { @@ -55,7 +63,12 @@ impl PipelinedPesStream { /// `DemuxThread` (which in turn owns the producer); we take the /// receiver end + the join handle bundle so cleanup is bounded /// on drop. - pub fn new( + /// + /// `pub(crate)`: the signature takes the internal `DemuxThread` / + /// `DemuxBatch` / `CodecParser` types, so external callers reach this + /// stream via [`super::resolve::input`] / `build_iso_pipeline` + /// instead. + pub(crate) fn new( demux_thread: DemuxThread, demux_rx: Receiver<DemuxBatch>, title: DiscTitle, @@ -70,6 +83,7 @@ impl PipelinedPesStream { demux_thread, pending_frames: std::collections::VecDeque::new(), eof: false, + skip_parse: std::env::var_os("FREEMKV_SKIP_PARSE").is_some(), } } @@ -88,12 +102,19 @@ impl PipelinedPesStream { Ok(true) } Ok(DemuxBatch::Err(e)) => Err(e), - Err(_) => Ok(false), + // Explicit clean-completion sentinel from the demux worker. + Ok(DemuxBatch::Eof) => Ok(false), + // The channel disconnected WITHOUT the worker first sending + // an `Eof` (or `Err`) sentinel — the worker panicked or was + // dropped mid-stream. Surface this as an error so a parser / + // demux panic is never reported to the caller as a clean + // end-of-stream (which would silently truncate output). + Err(_) => Err(crate::error::Error::DemuxThreadPanicked.into()), } } fn consume_ts(&mut self, packets: Vec<PesPacket>) { - let skip_parse = std::env::var_os("FREEMKV_SKIP_PARSE").is_some(); + let skip_parse = self.skip_parse; for pes in packets { if let Some((_, track)) = self .pid_to_track @@ -218,7 +239,7 @@ impl Stream for PipelinedPesStream { // codec_private before the consumer can write the container // header. FREEMKV_SKIP_PARSE forces ready (no parser ever // populates codec_private in that mode). - if std::env::var_os("FREEMKV_SKIP_PARSE").is_some() { + if self.skip_parse { return true; } for (idx, s) in self.title.streams.iter().enumerate() { diff --git a/src/mux/ps.rs b/src/mux/ps.rs index ae916b0..619b592 100644 --- a/src/mux/ps.rs +++ b/src/mux/ps.rs @@ -11,6 +11,8 @@ //! - 0xC0-0xDF: MPEG audio //! - 0xBD: private stream 1 (AC3, DTS, LPCM, subtitles via sub-stream ID) +use super::codec::startcode::find_start_code; + /// Pack header start code suffix. const PACK_HEADER_ID: u8 = 0xBA; @@ -23,6 +25,15 @@ const PROGRAM_END_ID: u8 = 0xB9; /// Private stream 1 (AC3, DTS, LPCM, subtitles). const PRIVATE_STREAM_1: u8 = 0xBD; +/// Hard cap on the demuxer's reassembly buffer. A length-0 (unbounded) video +/// PES is delimited by the next PS-layer boundary; if a corrupt stream declares +/// an unbounded PES and never follows it with a boundary, `feed()` would +/// otherwise accumulate the entire input. Past this cap we force the in-progress +/// unbounded PES to flush at the buffer end so untrusted input cannot drive +/// unbounded allocation. A real DVD pack/PES is at most a few KB; this leaves +/// generous slack while still bounding worst-case memory. +const MAX_PS_BUFFER: usize = 4 * 1024 * 1024; + /// A demuxed PES packet from the Program Stream. #[derive(Debug, Clone)] pub struct PsPacket { @@ -39,16 +50,48 @@ pub struct PsPacket { pub data: Vec<u8>, } +/// Canonical DVD video PID. DVD-Video carries a single MPEG-2 video +/// elementary stream; both the scanner and the muxer use this PID. +pub const DVD_VIDEO_PID: u16 = 0xE0; + +/// Canonical PID for a `private_stream_1` audio stream identified by its +/// on-wire sub-stream id. Returns `None` for sub-ids outside the AC-3 / +/// DTS / LPCM audio ranges. +/// +/// The PID is `0xBD00 | sub_stream_id`, which is unique per sub-stream id +/// (AC-3 / DTS `0x80..=0x8F`, LPCM `0xA0..=0xA7`). Unlike the old +/// per-codec relative arithmetic, distinct sub-ids therefore always yield +/// distinct PIDs — so a mixed-codec title (e.g. AC-3 + DTS, whose sub-ids +/// are 0x80 and 0x88) can never collide on one PID. This is the single +/// source of truth shared with `Disc::scan_dvd_titles` +/// (`src/disc/dvd.rs`), which sets each `AudioStream.pid` from the same +/// function so demuxer output routes through the title's `pid_to_track`. +pub fn dvd_audio_pid(sub_stream_id: u8) -> Option<u16> { + match sub_stream_id { + 0x80..=0x8F | 0xA0..=0xA7 => Some(0xBD00 | sub_stream_id as u16), + _ => None, + } +} + +/// Canonical PID for a VobSub subtitle stream identified by its on-wire +/// sub-stream id (`0x20..=0x3F`). The PID is the sub-id itself (identity), +/// which never overlaps the `0xBD..` audio PID space. +pub fn dvd_subtitle_pid(sub_stream_id: u8) -> Option<u16> { + match sub_stream_id { + 0x20..=0x3F => Some(sub_stream_id as u16), + _ => None, + } +} + impl PsPacket { /// Map this packet to the canonical DVD PID assigned by /// `Disc::scan_dvd_titles` (`src/disc/dvd.rs`), so demux output can /// be looked up in the title's `pid_to_track` map. /// - /// The PID space mirrors `dvd.rs` exactly: - /// - video stream id `0xE0..=0xEF` → `0xE0` - /// - private-stream-1 audio sub-id `0x80..=0x87` (AC-3), - /// `0x88..=0x8F` (DTS), `0xA0..=0xA7` (LPCM) → `0xBD00 + index` - /// - private-stream-1 subtitle sub-id `0x20..=0x3F` → `0x20 + index` + /// Routes by the REAL on-wire `(stream_id, sub_stream_id)` via the + /// shared [`dvd_audio_pid`] / [`dvd_subtitle_pid`] tables the scanner + /// also uses — never per-codec relative arithmetic, which collided on + /// mixed-codec audio (AC-3 0x80 and DTS 0x88 both mapping to 0xBD00). /// /// Returns `None` for stream/sub-stream combinations the DVD title /// scanner does not assign a PID to (e.g. MPEG audio 0xC0-0xDF, @@ -57,16 +100,11 @@ impl PsPacket { /// mis-routing the packet. pub fn dvd_pid(&self) -> Option<u16> { match self.stream_id { - 0xE0..=0xEF => Some(0xE0), - 0xBD => match self.sub_stream_id? { - // AC-3 / DTS / LPCM audio → 0xBD00 + audio index. - s @ 0x80..=0x87 => Some(0xBD00 + (s - 0x80) as u16), - s @ 0x88..=0x8F => Some(0xBD00 + (s - 0x88) as u16), - s @ 0xA0..=0xA7 => Some(0xBD00 + (s - 0xA0) as u16), - // VobSub subtitle sub-id 0x20+j → PID 0x20+j (identity). - s @ 0x20..=0x3F => Some(s as u16), - _ => None, - }, + 0xE0..=0xEF => Some(DVD_VIDEO_PID), + 0xBD => { + let sub = self.sub_stream_id?; + dvd_audio_pid(sub).or_else(|| dvd_subtitle_pid(sub)) + } _ => None, } } @@ -97,20 +135,26 @@ impl PsDemuxer { /// Feed raw MPEG-2 PS bytes, returning any completely parsed PES packets. pub fn feed(&mut self, data: &[u8]) -> Vec<PsPacket> { self.buffer.extend_from_slice(data); - self.extract_packets() + self.extract_packets(false) } /// Flush remaining buffered data, returning any final PES packets. pub fn flush(&mut self) -> Vec<PsPacket> { - // Try to extract whatever remains. If the buffer contains an incomplete - // PES packet we cannot parse, it will be discarded. - let packets = self.extract_packets(); + // At EOF, an unbounded (length 0) PES with no trailing start code is + // a complete-but-unterminated final packet — emit it rather than + // dropping the tail of the last frame. Genuinely incomplete packets + // (a length-bounded PES short of its declared size) are still + // discarded. + let packets = self.extract_packets(true); self.buffer.clear(); packets } - /// Scan the buffer for complete start-code-delimited units and parse them. - fn extract_packets(&mut self) -> Vec<PsPacket> { + /// Scan the buffer for complete start-code-delimited units and parse + /// them. When `flushing` is true, a trailing unbounded PES that has no + /// following start code is emitted using the rest of the buffer as its + /// payload (EOF terminates it). + fn extract_packets(&mut self, flushing: bool) -> Vec<PsPacket> { let mut packets = Vec::with_capacity(4); let mut pos = 0; @@ -132,7 +176,10 @@ impl PsDemuxer { if sc + 14 > self.buffer.len() { break; // wait for more data } - // MPEG-2 packs have bit pattern 01 in bits 7-6 of byte 4. + // DVD-Video is always MPEG-2 PS, so every 0xBA is treated + // as a 14-byte MPEG-2 pack: the low 3 bits of byte 13 are + // pack_stuffing_length. (An MPEG-1 pack would be 12 bytes + // with no stuffing field, but DVD never emits one.) let stuffing = (self.buffer[sc + 13] & 0x07) as usize; let pack_len = 14 + stuffing; if sc + pack_len > self.buffer.len() { @@ -162,13 +209,30 @@ impl PsDemuxer { ((self.buffer[sc + 4] as usize) << 8) | self.buffer[sc + 5] as usize; // Total bytes = 6 (start code + stream_id + length) + pes_packet_len. - // A length of 0 means unbounded (video streams); in that case we need - // to find the next start code to delimit the packet. + // A length of 0 means unbounded (video streams); in that + // case the packet runs to the next PS-LAYER boundary (pack / + // system header / program end / next PES), NOT the next raw + // start code — the video ES payload is itself full of + // 00 00 01 xx codes that would otherwise cut the PES short. let end = if pes_packet_len == 0 { - // Find the next start code after this one. - match find_start_code(&self.buffer, sc + 4) { - Some(next_sc) => next_sc, - None => break, // wait for more data + match find_ps_boundary(&self.buffer, sc + 4) { + Some(next) => next, + // At EOF the rest of the buffer is this PES's + // payload — emit it. + None if flushing => self.buffer.len(), + None => { + // No boundary buffered yet. Normally wait for + // more data, but a corrupt stream could declare + // an unbounded PES followed by endless non- + // boundary bytes — bounding the buffer here + // stops untrusted input forcing unbounded + // allocation. Past the cap, flush what we have. + if self.buffer.len() - sc > MAX_PS_BUFFER { + self.buffer.len() + } else { + break; // wait for more data + } + } } } else { let e = sc + 6 + pes_packet_len; @@ -198,6 +262,36 @@ impl PsDemuxer { } } +/// Find the next PS-layer unit boundary at or after `from`: a start code whose +/// ID byte is a pack (0xBA), system header (0xBB), program-end (0xB9), or a +/// payload-carrying PES stream ID (0xBD..=0xEF). +/// +/// A length-0 (unbounded) video PES must be delimited by the next PS-layer unit +/// — NOT by the next raw `00 00 01`. The MPEG-2 video elementary stream inside +/// the PES is itself full of `00 00 01 xx` start codes (picture 0x00, slices +/// 0x01..=0xAF, GOP 0xB8, sequence 0xB3); a plain start-code scan would cut the +/// PES inside its own payload and re-scan the discarded video bytes as bogus PS +/// units. Restricting the search to PS-layer IDs (>= 0xB9, excluding the video +/// ES codes below it) frames the unbounded PES at the right boundary. +fn find_ps_boundary(data: &[u8], from: usize) -> Option<usize> { + let mut pos = from; + while let Some(sc) = find_start_code(data, pos) { + if sc + 3 >= data.len() { + return None; + } + let id = data[sc + 3]; + if id == PACK_HEADER_ID + || id == SYSTEM_HEADER_ID + || id == PROGRAM_END_ID + || is_pes_stream_id(id) + { + return Some(sc); + } + pos = sc + 4; + } + None +} + /// Check whether a start code byte is a valid PES stream ID that carries payload. fn is_pes_stream_id(id: u8) -> bool { // Video: 0xE0-0xEF, MPEG audio: 0xC0-0xDF, private stream 1: 0xBD, @@ -251,10 +345,15 @@ fn parse_pes_packet(data: &[u8]) -> Option<PsPacket> { let mut pts = None; let mut dts = None; - if pts_dts_flags >= 2 && data.len() >= 14 { + // The PTS (5 bytes at data[9..14]) and DTS (5 bytes at data[14..19]) + // live INSIDE the PES header, so gate on header_data_len covering them + // (>=5 for PTS, >=10 for PTS+DTS), not merely on total length. A + // non-conformant packet that sets the flags but declares a too-short + // header would otherwise read payload bytes as a bogus timestamp. + if pts_dts_flags >= 2 && header_data_len >= 5 && data.len() >= 14 { pts = Some(parse_pts(&data[9..14])); } - if pts_dts_flags == 3 && data.len() >= 19 { + if pts_dts_flags == 3 && header_data_len >= 10 && data.len() >= 19 { dts = Some(parse_pts(&data[14..19])); } @@ -286,13 +385,13 @@ fn parse_pes_packet(data: &[u8]) -> Option<PsPacket> { /// Parse a 5-byte PTS/DTS timestamp field (33 bits at 90kHz). /// -/// Layout: +/// Layout (ISO/IEC 13818-1 Table 2-17): /// ```text -/// byte0: [marker_4bits][bit32][marker_1] -/// byte1: [bits 31..24] -/// byte2: [bits 23..15][marker_1] -/// byte3: [bits 14..7] -/// byte4: [bits 6..0][marker_1] +/// byte0: [prefix:4][pts 32..30:3][marker:1] +/// byte1: [pts 29..22:8] +/// byte2: [pts 21..15:7][marker:1] +/// byte3: [pts 14..7:8] +/// byte4: [pts 6..0:7][marker:1] /// ``` fn parse_pts(buf: &[u8]) -> u64 { debug_assert!(buf.len() >= 5); @@ -305,14 +404,6 @@ fn parse_pts(buf: &[u8]) -> u64 { ((b0 >> 1) & 0x07) << 30 | b1 << 22 | (b2 >> 1) << 15 | b3 << 7 | b4 >> 1 } -/// Find the position of the next start code (00 00 01) at or after `from`. -fn find_start_code(data: &[u8], from: usize) -> Option<usize> { - if data.len() < from + 3 { - return None; - } - (from..data.len() - 2).find(|&i| data[i] == 0x00 && data[i + 1] == 0x00 && data[i + 2] == 0x01) -} - #[cfg(test)] mod tests { use super::*; @@ -537,6 +628,25 @@ mod tests { assert_eq!(p2[0].data, vec![0xAA, 0xBB, 0xCC]); } + #[test] + fn flush_emits_trailing_unbounded_video_pes() { + let mut demuxer = PsDemuxer::new(); + // Unbounded (length 0) video PES with no trailing start code — the + // common EOF case. feed() must not emit it (awaiting a delimiter), + // but flush() must emit the tail rather than discarding it. + let data = vec![ + 0x00, 0x00, 0x01, 0xE0, 0x00, 0x00, // video, length 0 (unbounded) + 0x80, 0x00, 0x00, // no PTS, header_data_len = 0 + 0xAA, 0xBB, 0xCC, 0xDD, + ]; + let fed = demuxer.feed(&data); + assert!(fed.is_empty(), "unbounded PES not emitted until delimited"); + let flushed = demuxer.flush(); + assert_eq!(flushed.len(), 1, "flush emits the trailing PES"); + assert_eq!(flushed[0].stream_id, 0xE0); + assert_eq!(flushed[0].data, vec![0xAA, 0xBB, 0xCC, 0xDD]); + } + // --- Multiple PES packets --- #[test] @@ -564,6 +674,74 @@ mod tests { assert_eq!(packets[1].stream_id, 0xC0); } + // --- unbounded (length-0) video PES framing --- + + #[test] + fn unbounded_video_pes_not_cut_by_embedded_start_codes() { + // A length-0 video PES whose ES payload contains embedded MPEG start + // codes (picture 0x00, slice 0x01, GOP 0xB8, sequence 0xB3) must be + // delimited by the NEXT PS-layer boundary (here a program-end 0xB9), + // not by the first embedded 00 00 01 inside the payload. + let mut demuxer = PsDemuxer::new(); + + let mut data = vec![ + 0x00, 0x00, 0x01, 0xE0, // video stream + 0x00, 0x00, // length = 0 (unbounded) + 0x80, 0x00, 0x00, // flags: no PTS, header_data_len = 0 + ]; + // ES payload with embedded MPEG-2 start codes. + let payload = [ + 0x00, 0x00, 0x01, 0xB3, // sequence header + 0x11, 0x22, 0x00, 0x00, 0x01, 0x00, // picture start code + 0x33, 0x44, 0x00, 0x00, 0x01, 0x01, // slice + 0x55, 0x66, + ]; + data.extend_from_slice(&payload); + // PS-layer boundary that closes the unbounded PES. + data.extend_from_slice(&[0x00, 0x00, 0x01, 0xB9]); + + let packets = demuxer.feed(&data); + assert_eq!(packets.len(), 1, "one PES, not several payload fragments"); + assert_eq!(packets[0].stream_id, 0xE0); + // The whole ES payload survives — none of it discarded as bogus units. + assert_eq!(packets[0].data, payload.to_vec()); + } + + #[test] + fn unbounded_video_pes_waits_for_boundary() { + // Without a following PS-layer boundary the unbounded PES is held + // (waiting for more data), not emitted truncated. + let mut demuxer = PsDemuxer::new(); + let mut data = vec![0x00, 0x00, 0x01, 0xE0, 0x00, 0x00, 0x80, 0x00, 0x00]; + data.extend_from_slice(&[0x00, 0x00, 0x01, 0x00, 0xAA, 0xBB]); // picture SC, no PS boundary + let packets = demuxer.feed(&data); + assert!(packets.is_empty(), "no PS boundary yet → hold the PES"); + } + + #[test] + fn unbounded_video_pes_buffer_is_bounded() { + // A corrupt stream declaring an unbounded PES followed by endless + // non-boundary bytes must not grow the buffer without limit. + let mut demuxer = PsDemuxer::new(); + let header = vec![0x00, 0x00, 0x01, 0xE0, 0x00, 0x00, 0x80, 0x00, 0x00]; + let packets = demuxer.feed(&header); + assert!(packets.is_empty()); + // Feed >MAX_PS_BUFFER of bytes containing no PS-layer boundary. + let chunk = vec![0x55u8; 1024 * 1024]; + let mut emitted = 0; + for _ in 0..(MAX_PS_BUFFER / chunk.len() + 4) { + emitted += demuxer.feed(&chunk).len(); + } + assert!( + demuxer.buffer.len() <= MAX_PS_BUFFER + chunk.len(), + "buffer grew to {} (cap {})", + demuxer.buffer.len(), + MAX_PS_BUFFER + ); + // The force-flush emits the over-long PES rather than accumulating it. + assert!(emitted >= 1, "over-cap unbounded PES is force-flushed"); + } + // --- PTS parsing edge cases --- #[test] @@ -597,14 +775,13 @@ mod tests { #[test] fn dvd_pid_matches_scanner_assignment() { // Video → 0xE0 (matches dvd.rs VideoStream pid). - assert_eq!(mk(0xE0, None).dvd_pid(), Some(0xE0)); - // AC-3 audio stream 0/1 → 0xBD00 / 0xBD01 (matches 0xBD00 + i). - assert_eq!(mk(0xBD, Some(0x80)).dvd_pid(), Some(0xBD00)); - assert_eq!(mk(0xBD, Some(0x81)).dvd_pid(), Some(0xBD01)); - // DTS / LPCM audio indices. - assert_eq!(mk(0xBD, Some(0x88)).dvd_pid(), Some(0xBD00)); - assert_eq!(mk(0xBD, Some(0xA0)).dvd_pid(), Some(0xBD00)); - // VobSub subtitle 0x20/0x21 → 0x20 / 0x21 (matches 0x20 + j). + assert_eq!(mk(0xE0, None).dvd_pid(), Some(DVD_VIDEO_PID)); + // PID = 0xBD00 | sub_stream_id — unique per sub-id, no collision. + assert_eq!(mk(0xBD, Some(0x80)).dvd_pid(), Some(0xBD80)); // AC-3 #0 + assert_eq!(mk(0xBD, Some(0x81)).dvd_pid(), Some(0xBD81)); // AC-3 #1 + assert_eq!(mk(0xBD, Some(0x88)).dvd_pid(), Some(0xBD88)); // DTS #0 + assert_eq!(mk(0xBD, Some(0xA0)).dvd_pid(), Some(0xBDA0)); // LPCM #0 + // VobSub subtitle 0x20/0x21 → 0x20 / 0x21 (identity). assert_eq!(mk(0xBD, Some(0x20)).dvd_pid(), Some(0x20)); assert_eq!(mk(0xBD, Some(0x21)).dvd_pid(), Some(0x21)); // Unmappable: MPEG audio, private stream 2, bogus sub-id. @@ -613,23 +790,61 @@ mod tests { assert_eq!(mk(0xBD, Some(0x10)).dvd_pid(), None); } + #[test] + fn mixed_codec_audio_does_not_collide() { + // The core regression: a title mixing AC-3 (0x80), DTS (0x88) and + // LPCM (0xA0) audio. The old per-codec relative arithmetic mapped + // all three to 0xBD00. They must now get distinct PIDs that match + // what dvd.rs assigns from the same dvd_audio_pid() table. + let ac3 = mk(0xBD, Some(0x80)).dvd_pid().unwrap(); + let dts = mk(0xBD, Some(0x88)).dvd_pid().unwrap(); + let lpcm = mk(0xBD, Some(0xA0)).dvd_pid().unwrap(); + assert_ne!(ac3, dts, "AC-3 and DTS must not collide"); + assert_ne!(ac3, lpcm, "AC-3 and LPCM must not collide"); + assert_ne!(dts, lpcm, "DTS and LPCM must not collide"); + + // Scanner side uses the same table; build a pid_to_track for a + // mixed-codec title [video, AC-3, DTS, LPCM, sub] and route every + // PS packet to its own distinct track. + let pid_to_track: Vec<(u16, usize)> = vec![ + (DVD_VIDEO_PID, 0), + (dvd_audio_pid(0x80).unwrap(), 1), + (dvd_audio_pid(0x88).unwrap(), 2), + (dvd_audio_pid(0xA0).unwrap(), 3), + (dvd_subtitle_pid(0x20).unwrap(), 4), + ]; + let route = |p: PsPacket| -> Option<usize> { + let pid = p.dvd_pid()?; + pid_to_track + .iter() + .find(|(x, _)| *x == pid) + .map(|(_, t)| *t) + }; + assert_eq!(route(mk(0xE0, None)), Some(0)); + assert_eq!(route(mk(0xBD, Some(0x80))), Some(1)); // AC-3 → its own track + assert_eq!(route(mk(0xBD, Some(0x88))), Some(2)); // DTS → its own track + assert_eq!(route(mk(0xBD, Some(0xA0))), Some(3)); // LPCM → its own track + assert_eq!(route(mk(0xBD, Some(0x20))), Some(4)); // sub → its own track + } + #[test] fn subtitle_does_not_collide_with_audio_track() { - // Regression for the (sub_id & 0x1F)+1 bug: subtitle sub-id 0x20 - // used to alias audio track 1. With the real PID it routes to its - // own subtitle PID (0x20), distinct from audio (0xBD00+). - let audio0 = mk(0xBD, Some(0x80)).dvd_pid().unwrap(); // 0xBD00 + // Subtitle sub-id 0x20 routes to its own subtitle PID (0x20), + // distinct from any audio PID (0xBD80+). + let audio0 = mk(0xBD, Some(0x80)).dvd_pid().unwrap(); // 0xBD80 let sub0 = mk(0xBD, Some(0x20)).dvd_pid().unwrap(); // 0x20 assert_ne!( audio0, sub0, "subtitle sub-id 0x20 must NOT map to the audio PID" ); - // Mirror dvd.rs PID assignment for a title with [video, audio0, - // audio1, sub0, sub1] and confirm each PS packet lands on its - // own track via pid_to_track. - let pid_to_track: Vec<(u16, usize)> = - vec![(0xE0, 0), (0xBD00, 1), (0xBD01, 2), (0x20, 3), (0x21, 4)]; + let pid_to_track: Vec<(u16, usize)> = vec![ + (DVD_VIDEO_PID, 0), + (dvd_audio_pid(0x80).unwrap(), 1), + (dvd_audio_pid(0x81).unwrap(), 2), + (dvd_subtitle_pid(0x20).unwrap(), 3), + (dvd_subtitle_pid(0x21).unwrap(), 4), + ]; let route = |p: PsPacket| -> Option<usize> { let pid = p.dvd_pid()?; pid_to_track diff --git a/src/mux/resolve.rs b/src/mux/resolve.rs index 1979916..9d1a805 100644 --- a/src/mux/resolve.rs +++ b/src/mux/resolve.rs @@ -14,6 +14,11 @@ //! //! Bare paths without a scheme are rejected. //! For disc→ISO (raw sector copy), use `Disc::copy()` instead. +//! +//! Note: `disc://` cannot be opened through [`input`]; it returns +//! [`crate::error::Error::DiscUrlNotDirect`]. Live-disc input must go +//! through `Drive::open()` + `Disc::scan()` + `DiscStream::new()`, not +//! the URL resolver. use super::network::NetworkStream; use super::null::NullStream; @@ -29,6 +34,7 @@ use std::path::{Path, PathBuf}; const IO_BUF_SIZE: usize = 4 * 1024 * 1024; /// Parsed stream URL. +#[derive(Debug, Clone)] pub enum StreamUrl { /// Optical disc drive. Device path is optional (auto-detect if None). Disc { device: Option<PathBuf> }, @@ -109,11 +115,18 @@ pub fn parse_url(url: &str) -> StreamUrl { addr: rest.to_string(), }; } - if url == "null://" || url.starts_with("null://") { - return StreamUrl::Null; + if let Some(rest) = url.strip_prefix("null://") { + // null:// / stdio:// are scheme-only; a trailing path is + // malformed and must fall through to Unknown rather than be + // silently discarded. + if rest.is_empty() { + return StreamUrl::Null; + } } - if url == "stdio://" || url.starts_with("stdio://") { - return StreamUrl::Stdio; + if let Some(rest) = url.strip_prefix("stdio://") { + if rest.is_empty() { + return StreamUrl::Stdio; + } } if let Some(rest) = url.strip_prefix("iso://") { return StreamUrl::Iso { @@ -150,6 +163,16 @@ fn validate_network_addr(addr: &str) -> io::Result<()> { } .into()); } + // A bare IPv6 literal ("::1", "2001:db8::1") contains ':' yet has no port, + // so the simple `contains(':')` check would wrongly pass it and TcpListener + // would later return an untyped io::Error. Treat anything that parses as a + // bare IpAddr (v4 or v6) as port-less. + if addr.parse::<std::net::IpAddr>().is_ok() { + return Err(crate::error::Error::StreamUrlMissingPort { + addr: addr.to_string(), + } + .into()); + } if !addr.contains(':') { return Err(crate::error::Error::StreamUrlMissingPort { addr: addr.to_string(), @@ -160,13 +183,15 @@ fn validate_network_addr(addr: &str) -> io::Result<()> { } /// Options for opening an input stream. -#[derive(Default)] +#[derive(Debug, Clone, Default)] pub struct InputOptions { /// Caller-resolved per-CPS-unit AACS keys to apply to the scanned disc /// (`(cps_unit, 16-byte key)`). Empty for an unencrypted disc or when the /// caller has no key. The library does no lookup — a key source resolves /// these and the caller passes them here. pub unit_keys: Vec<(u32, [u8; 16])>, + /// 0-based title index to open; `None` selects title 0. An + /// out-of-range index yields [`crate::error::Error::DiscTitleRange`]. pub title_index: Option<usize>, /// Skip decryption — return raw encrypted bytes. pub raw: bool, @@ -253,19 +278,36 @@ pub fn input(url: &str, opts: &InputOptions) -> io::Result<Box<dyn crate::pes::S } // Correct TrueHD channel counts (MPLS understates 7.1/Atmos as 5.1) // by probing the first DECRYPTED access units of the chosen title. - // A fresh reader avoids disturbing the mux reader below. + // A fresh reader avoids disturbing the mux reader below. Skipped in + // --raw mode: the probe would re-open + decrypt for nothing (on an + // AACS disc with no key the correction is a no-op on ciphertext, and + // raw output isn't decoded anyway). let keys = disc.decrypt_keys(); - if let Ok(probe) = crate::io::file_sector_source::FileSectorSource::open(path) { - let mut dec = crate::sector::DecryptingSectorSource::new(probe, keys.clone()); - crate::disc::correct_truehd_channels(&mut dec, &mut disc.titles[idx]); + if !opts.raw { + match crate::io::file_sector_source::FileSectorSource::open(path) { + Ok(probe) => { + let mut dec = + crate::sector::DecryptingSectorSource::new(probe, keys.clone()); + crate::disc::correct_truehd_channels(&mut dec, &mut disc.titles[idx]); + } + Err(e) => { + // Non-fatal: a failed re-open just leaves MPLS 7.1/Atmos + // channel counts uncorrected (understated as 5.1). Log so + // the uncorrected path is diagnosable rather than silent. + tracing::debug!( + target: "mux", + "TrueHD channel-correction probe re-open failed: {e}" + ); + } + } } let title = disc.titles[idx].clone(); let format = disc.content_format; - // ISO file: 16 MiB batch — sequential read from fast - // storage, no bad sectors. Measured optimum on the rip1 - // testbed; bumping to 32 MiB regressed (more cache - // pressure, longer per-batch latency starves the consumer - // between iterations). Physical drives keep smaller + // ISO file: 8192-sector batch (16 MiB at 2048 B/sector) — + // sequential read from fast storage, no bad sectors. Measured + // optimum on the rip1 testbed; bumping to 16384 sectors (32 MiB) + // regressed (more cache pressure, longer per-batch latency starves + // the consumer between iterations). Physical drives keep smaller // batches for adaptive error handling. const ISO_MUX_BATCH_SECTORS: u16 = 8192; @@ -286,7 +328,7 @@ pub fn input(url: &str, opts: &InputOptions) -> io::Result<Box<dyn crate::pes::S format, None, None, - ); + )?; Ok(Box::new(stream)) } StreamUrl::M2ts { ref path } => { @@ -403,6 +445,21 @@ fn build_demux_state(title: &DiscTitle, format: ContentFormat) -> DemuxState { /// Assemble the ISO mux pipeline (read+decrypt → demux → parse) for /// a `FileSectorSource`-backed reader. Returns the resulting /// `PipelinedPesStream`. +/// +/// # Parameters +/// - `reader`: the sector source to read from (typically a +/// `FileSectorSource` over the ISO image). +/// - `title`: the selected title; its `extents` drive the read range and its +/// `streams` build the demux/parse tables. +/// - `keys`: decryption keys applied per sector batch. Pass +/// [`crate::decrypt::DecryptKeys::None`] for raw / unencrypted reads (the +/// decrypt decorator then becomes a pass-through). +/// - `batch_sectors`: read batch size in logical (2048-byte) sectors — a +/// throughput/latency tuning knob, not a correctness parameter. +/// - `format`: container format (`BdTs` → TS demuxer, `MpegPs` → PS demuxer). +/// - `halt`: cooperative cancel token (not a timeout); when cancelled the +/// pipeline stops at the next boundary. `None` disables cancellation. +/// - `event_fn`: optional progress/event callback invoked by the prefetcher. pub fn build_iso_pipeline<S: SectorSource + Send + 'static>( reader: S, title: DiscTitle, @@ -411,7 +468,7 @@ pub fn build_iso_pipeline<S: SectorSource + Send + 'static>( format: ContentFormat, halt: Option<crate::halt::Halt>, event_fn: Option<crate::sector::prefetched::EventFn>, -) -> PipelinedPesStream { +) -> io::Result<PipelinedPesStream> { let extents = title.extents.clone(); let decrypting = crate::sector::DecryptingSectorSource::new(Box::new(reader) as Box<dyn SectorSource>, keys); @@ -421,13 +478,21 @@ pub fn build_iso_pipeline<S: SectorSource + Send + 'static>( batch_sectors, halt.clone(), event_fn, - ); + ) + .map_err(|e| -> io::Error { e.into() })?; let (rx, recycle_tx, shell) = prefetched.into_channels(); let (parsers, pid_to_track, ts, ps) = build_demux_state(&title, format); let (demux_thread, demux_rx) = - super::demux_thread::DemuxThread::spawn_zero_copy(rx, recycle_tx, shell, halt, ts, ps); - PipelinedPesStream::new(demux_thread, demux_rx, title, parsers, pid_to_track) + super::demux_thread::DemuxThread::spawn_zero_copy(rx, recycle_tx, shell, halt, ts, ps) + .map_err(|e| -> io::Error { e.into() })?; + Ok(PipelinedPesStream::new( + demux_thread, + demux_rx, + title, + parsers, + pid_to_track, + )) } /// Assemble the M2TS file mux pipeline (read → demux → parse) for a @@ -455,19 +520,32 @@ fn build_m2ts_pipeline<R: std::io::Read + Send + 'static>( }; head.truncate(head_len); - // Try FMKV metadata header first; fall back to PMT scan. + // Try FMKV metadata header first; fall back to PMT scan. Only a + // genuine absence of the FMKV magic (`Ok(None)`) falls through to + // the PMT path — a corrupt/truncated FMKV header (`Err`) propagates + // instead of being misreported as a PMT-derived title or NoStreams. let mut cursor = io::Cursor::new(&head); - let (title, head_consumed) = if let Ok(Some(m)) = meta::read_header(&mut cursor) { - (m.to_title(), cursor.position() as usize) - } else { - let streams = super::ts::scan_streams(&head) - .ok_or_else(|| -> io::Error { crate::error::Error::NoStreams.into() })?; - let t = DiscTitle { - duration_secs: 0.0, - streams, - ..DiscTitle::empty() - }; - (t, 0) + let (title, head_consumed) = match meta::read_header(&mut cursor)? { + Some(m) => { + let t = m.to_title(); + // Guard the FMKV branch the same way the ISO and PMT paths + // do: a header carrying zero streams yields an empty title + // that would mux nothing — surface NoStreams instead. + if t.streams.is_empty() { + return Err(crate::error::Error::NoStreams.into()); + } + (t, cursor.position() as usize) + } + None => { + let streams = super::ts::scan_streams(&head) + .ok_or_else(|| -> io::Error { crate::error::Error::NoStreams.into() })?; + let t = DiscTitle { + duration_secs: 0.0, + streams, + ..DiscTitle::empty() + }; + (t, 0) + } }; // Chain: any un-consumed head bytes + the remainder of the @@ -479,12 +557,13 @@ fn build_m2ts_pipeline<R: std::io::Read + Send + 'static>( chained, crate::io::byte_prefetcher::DEFAULT_CHUNK_BYTES, None, - ); + )?; let (rx, recycle_tx, shell) = prefetcher.into_channels(); let (parsers, pid_to_track, ts, ps) = build_demux_state(&title, ContentFormat::BdTs); let (demux_thread, demux_rx) = - super::demux_thread::DemuxThread::spawn_zero_copy(rx, recycle_tx, shell, None, ts, ps); + super::demux_thread::DemuxThread::spawn_zero_copy(rx, recycle_tx, shell, None, ts, ps) + .map_err(|e| -> io::Error { e.into() })?; Ok(PipelinedPesStream::new( demux_thread, demux_rx, @@ -497,8 +576,21 @@ fn build_m2ts_pipeline<R: std::io::Read + Send + 'static>( #[cfg(test)] mod tests { use super::aacs_key_missing; + use super::validate_network_addr; use crate::decrypt::DecryptKeys; + #[test] + fn validate_network_addr_rejects_portless() { + // Empty, bare IPv4, and bare IPv6 (which contains ':') must all fail. + assert!(validate_network_addr("").is_err()); + assert!(validate_network_addr("127.0.0.1").is_err()); + assert!(validate_network_addr("::1").is_err()); + assert!(validate_network_addr("2001:db8::1").is_err()); + // host:port and ip:port forms pass. + assert!(validate_network_addr("127.0.0.1:9000").is_ok()); + assert!(validate_network_addr("host:9000").is_ok()); + } + fn aacs_keys() -> DecryptKeys { DecryptKeys::Aacs { unit_keys: vec![(1, [0x11u8; 16])], diff --git a/src/mux/stdio.rs b/src/mux/stdio.rs index 116e275..88de5b6 100644 --- a/src/mux/stdio.rs +++ b/src/mux/stdio.rs @@ -15,7 +15,12 @@ pub struct StdioStream { writer: Option<io::BufWriter<io::Stdout>>, header_written: bool, header_read: bool, - stored_codec_privates: Vec<Option<Vec<u8>>>, + /// True once an FMKV header was actually parsed on the read side + /// (set only inside the `Some(meta)` arm). Distinct from + /// `header_read`, which is true after the first read attempt even + /// when no header was present — `headers_ready()` must gate on the + /// metadata actually being available, not merely on having looked. + meta_parsed: bool, } impl StdioStream { @@ -27,7 +32,7 @@ impl StdioStream { writer: None, header_written: false, header_read: false, - stored_codec_privates: Vec::new(), + meta_parsed: false, } } @@ -39,10 +44,25 @@ impl StdioStream { writer: Some(io::BufWriter::new(io::stdout())), header_written: false, header_read: false, - stored_codec_privates: Vec::new(), + meta_parsed: false, } } + /// Write the FMKV metadata header to stdout exactly once, before any + /// frames. Always writes (even when the title has no streams) so a + /// zero-frame output stream still emits the magic + metadata header, + /// keeping the wire protocol symmetric with the read side's read_header(). + fn ensure_header_written(&mut self) -> io::Result<()> { + if let Some(w) = &mut self.writer { + if !self.header_written { + let m = meta::M2tsMeta::from_title(&self.disc_title); + meta::write_header(w, &m)?; + self.header_written = true; + } + } + Ok(()) + } + /// Read the FMKV metadata header from stdin on first read. fn ensure_header_read(&mut self) -> io::Result<()> { if self.header_read { @@ -50,10 +70,16 @@ impl StdioStream { } self.header_read = true; if let Some(ref mut r) = self.reader { - if let Ok(Some(m)) = meta::read_header(r) { - let title = m.to_title(); - self.stored_codec_privates = title.codec_privates.clone(); - self.disc_title = title; + // Propagate real header errors. read_header consumes bytes + // from the unbuffered stdin BEFORE it can fail (oversized + // length, bad JSON, partial read), so swallowing the Err + // would leave the stream misaligned and PesFrame::deserialize + // would then read garbage. `?` surfaces the true error; + // Ok(None) (genuine magic mismatch / clean EOF) stays a + // non-error and leaves the empty default title in place. + if let Some(m) = meta::read_header(r)? { + self.disc_title = m.to_title(); + self.meta_parsed = true; } } Ok(()) @@ -69,21 +95,20 @@ impl crate::pes::Stream for StdioStream { } } fn write(&mut self, frame: &crate::pes::PesFrame) -> io::Result<()> { + if self.writer.is_none() { + return Err(crate::error::Error::StreamReadOnly.into()); + } + self.ensure_header_written()?; match &mut self.writer { - Some(w) => { - if !self.header_written { - if !self.disc_title.streams.is_empty() { - let m = meta::M2tsMeta::from_title(&self.disc_title); - meta::write_header(w, &m)?; - } - self.header_written = true; - } - frame.serialize(w) - } + Some(w) => frame.serialize(w), None => Err(crate::error::Error::StreamReadOnly.into()), } } fn finish(&mut self) -> io::Result<()> { + // Emit the header even when write() was never called, so a zero-frame + // title still produces the FMKV magic + metadata header on stdout + // (symmetric with the read side's read_header()). + self.ensure_header_written()?; if let Some(w) = &mut self.writer { w.flush()?; } @@ -94,13 +119,23 @@ impl crate::pes::Stream for StdioStream { } fn codec_private(&self, track: usize) -> Option<Vec<u8>> { - self.stored_codec_privates + // Single source of truth: the title's own codec_privates. (The + // previous `stored_codec_privates` field was a redundant clone + // of exactly this, populated from the same header.) + self.disc_title + .codec_privates .get(track) .and_then(|c| c.clone()) } fn headers_ready(&self) -> bool { - // After first read(), header is parsed and codec_privates populated - self.header_read || self.writer.is_some() + // Write side: caller supplied the title up front, so headers are + // always ready. Read side: ready only once an FMKV header was + // actually parsed — gating on `header_read` alone would claim + // readiness for a headerless stream whose codec_private() is None + // for every track, starving the downstream MKV writer of init + // data. A genuinely headerless stream never flips ready (the + // caller must then fall back to its own codec detection). + self.writer.is_some() || self.meta_parsed } } diff --git a/src/mux/ts.rs b/src/mux/ts.rs index 31e26e9..42e6c7f 100644 --- a/src/mux/ts.rs +++ b/src/mux/ts.rs @@ -128,6 +128,14 @@ impl TsDemuxer { /// limits. Empty `pids` yields max_pid 0; the floor still produces a /// valid (wholly-unused) table. pub fn new(pids: &[u16]) -> Self { + // The PID→assembler index is stored as i16 (-1 = untracked), so a + // 32768th+ tracked PID would truncate to a negative value and be + // silently treated as untracked. Callers pass a handful of PIDs + // (BD-TS has at most ~8192), so this is a programmer-error guard. + debug_assert!( + pids.len() <= i16::MAX as usize, + "TsDemuxer: too many PIDs for an i16 index table" + ); let max_pid = pids.iter().copied().max().unwrap_or(0) as usize; let table_size = (max_pid + 1).max(8192); let mut pid_index = vec![-1i16; table_size]; @@ -222,6 +230,13 @@ impl TsDemuxer { if idx < 0 { return; } + // adaptation_field_control == 0b00 is reserved (ISO 13818-1) and + // carries no payload; discard so a corrupt/desynced packet can't + // inject its 184 bytes into the PES assembler. + if adaptation == 0x00 { + return; + } + let asm = &mut self.assemblers[idx as usize]; let payload_start = if adaptation == 0x03 || adaptation == 0x02 { @@ -309,23 +324,24 @@ fn parse_pes_header(data: &[u8]) -> (Option<i64>, Option<i64>, usize) { let stream_id = data[3]; - // Some stream IDs don't have the standard PES header extension - // (program_stream_map, padding, private_stream_2, ECM, EMM, etc.) + // Some stream IDs don't carry the standard PES header extension + // (ISO 13818-1 Table 2-22: program_stream_map, padding, private_stream_2, + // ECM, EMM, DSMCC_stream 0xF2, H.222.1 type E 0xF8, program_stream_directory). if stream_id == 0xBC || stream_id == 0xBE || stream_id == 0xBF || stream_id == 0xF0 || stream_id == 0xF1 + || stream_id == 0xF2 + || stream_id == 0xF8 || stream_id == 0xFF { return (None, None, 6); } - // Standard PES header: [6] = flags1, [7] = flags2, [8] = header_data_length - if data.len() < 9 { - return (None, None, 6); - } - + // Standard PES header: [6] = flags1, [7] = flags2, [8] = header_data_length. + // The `data.len() < 9` precondition was already checked at the top of + // this function and nothing shrinks `data` since, so no re-check here. let pts_dts_flags = (data[7] >> 6) & 0x03; let header_data_len = data[8] as usize; // Full, uncapped header length. PTS/DTS (if present) live in the @@ -352,8 +368,9 @@ fn parse_timestamp(data: &[u8]) -> Option<i64> { if data.len() < 5 { return None; } - // Validate marker bits: byte 2 bit 0 and byte 4 bit 0 must be 1 - if (data[2] & 0x01) == 0 || (data[4] & 0x01) == 0 { + // Validate marker bits: per MPEG-2 Systems (Table 2-17) bit 0 of + // bytes 0, 2 and 4 of the 5-byte PTS/DTS field must all be 1. + if (data[0] & 0x01) == 0 || (data[2] & 0x01) == 0 || (data[4] & 0x01) == 0 { return None; } let b0 = data[0] as i64; @@ -369,167 +386,252 @@ fn parse_timestamp(data: &[u8]) -> Option<i64> { // Stream scanning (PAT/PMT → stream list) // ============================================================ +/// Whether `offset` is a credible BD-TS packet boundary in the PSI scanner. +/// +/// Requires the sync byte at `data[offset + 4]`, and — to avoid latching onto +/// a stray 0x47 inside a TP_extra_header or payload during a desync — also +/// requires the next 192-spaced position to carry a sync byte when one exists +/// in the buffer. A lone trailing packet (no follower in range) is accepted on +/// its single sync byte. +fn is_resync_point(data: &[u8], offset: usize) -> bool { + if data.get(offset + 4) != Some(&SYNC_BYTE) { + return false; + } + match data.get(offset + BD_TS_PACKET_SIZE + 4) { + Some(&b) => b == SYNC_BYTE, + None => true, // last packet in the buffer — no follower to corroborate + } +} + +/// Compute the byte offset of the PSI payload (the pointer_field) for a BD-TS +/// packet starting at `pkt` (the 4-byte TP_extra_header + 188-byte TS packet). +/// +/// Accounts for the adaptation_field_control (bits 5:4 of the 4th TS header +/// byte). Returns `None` when the packet carries no payload (AFC 0b10 = AF +/// only, or the reserved 0b00) or when the adaptation field length runs past +/// the packet. `pkt` must be at least [`BD_TS_PACKET_SIZE`] bytes. +fn psi_payload_base(pkt: &[u8]) -> Option<usize> { + // TS header is pkt[4..]; byte pkt[7] holds AFC in bits 5:4. + let afc = (pkt[7] >> 4) & 0x03; + match afc { + 0x01 => Some(8), // payload only: 4 (TP_extra) + 4 (TS header) + 0x03 => { + // Adaptation field present + payload. AF length byte is pkt[8]; + // payload starts after it. + let af_len = pkt[8] as usize; + let base = 9 + af_len; // 4 + 4 + 1(length byte) + af_len + if base < BD_TS_PACKET_SIZE { + Some(base) + } else { + None // AF overruns the packet + } + } + // 0x02 = AF only (no payload), 0x00 = reserved. + _ => None, + } +} + +/// Reassemble a single PSI section (PAT / PMT) for `target_pid` with +/// the expected `table_id`, respecting TS-packet boundaries. +/// +/// The section pointed at by `pointer_field` in the PUSI packet may be +/// longer than the 184-byte TS payload (PSI sections can reach 1021 +/// bytes; a PMT with many ES entries spans 2+ packets). Reading a flat +/// slice of the input would walk straight through the next packet's +/// TP_extra_header + TS header as if it were table content, yielding a +/// wrong PID / garbage stream_type. This walks the PUSI packet, applies +/// `pointer_field` bounded to within that packet's payload, then appends +/// the payload of each subsequent continuation packet (same PID, no +/// PUSI) until `3 + section_length` bytes have been collected. +/// +/// The PUSI packet's payload base is computed with [`psi_payload_base`] +/// so a PSI section carried behind an adaptation field is located +/// correctly rather than assuming the payload starts at `offset + 8`. +/// +/// Returns the section bytes (starting at the table_id) or `None` if no +/// matching section is found. +fn collect_psi_section(data: &[u8], target_pid: u16, table_id: u8) -> Option<Vec<u8>> { + let mut offset = 0; + while offset + BD_TS_PACKET_SIZE <= data.len() { + if !is_resync_point(data, offset) { + offset += 1; + continue; + } + let pid = (((data[offset + 5] & 0x1F) as u16) << 8) | data[offset + 6] as u16; + let pusi = data[offset + 5] & 0x40 != 0; + + if pid == target_pid && pusi { + // Locate the payload (pointer_field) accounting for any + // adaptation field. A packet with no payload (AF only) or an + // AF that overruns the packet is skipped. + let Some(payload_off) = psi_payload_base(&data[offset..offset + BD_TS_PACKET_SIZE]) + else { + offset += BD_TS_PACKET_SIZE; + continue; + }; + let payload = &data[offset + payload_off..offset + BD_TS_PACKET_SIZE]; + // pointer_field is the FIRST payload byte; the section starts + // pointer_field bytes after it. Bound the start to within + // THIS packet's payload — a pointer that runs into the next + // packet is malformed. + let pointer = payload[0] as usize; + let sec_start = 1 + pointer; + if sec_start + 3 > payload.len() || payload[sec_start] != table_id { + offset += BD_TS_PACKET_SIZE; + continue; + } + let section_len = + (((payload[sec_start + 1] & 0x0F) as usize) << 8) | payload[sec_start + 2] as usize; + let total = 3 + section_len; // table_id + 2 length bytes + body + let mut section = Vec::with_capacity(total); + section.extend_from_slice(&payload[sec_start..]); + if section.len() >= total { + section.truncate(total); + return Some(section); + } + // Need continuation packets: same PID, no PUSI. + let mut scan = offset + BD_TS_PACKET_SIZE; + while scan + BD_TS_PACKET_SIZE <= data.len() && section.len() < total { + if data[scan + 4] != SYNC_BYTE { + scan += 1; + continue; + } + let cpid = (((data[scan + 5] & 0x1F) as u16) << 8) | data[scan + 6] as u16; + let cpusi = data[scan + 5] & 0x40 != 0; + if cpid == target_pid && !cpusi { + // Continuation packets may also carry an adaptation + // field; compute their payload base the same way. + if let Some(cbase) = psi_payload_base(&data[scan..scan + BD_TS_PACKET_SIZE]) { + section.extend_from_slice(&data[scan + cbase..scan + BD_TS_PACKET_SIZE]); + } + } + scan += BD_TS_PACKET_SIZE; + } + if section.len() >= total { + section.truncate(total); + return Some(section); + } + // Incomplete section (truncated input) — stop looking. + return None; + } + offset += BD_TS_PACKET_SIZE; + } + None +} + /// Scan BD-TS data for streams by parsing PAT and PMT tables. /// Returns None if no valid program is found. pub fn scan_streams(data: &[u8]) -> Option<Vec<crate::disc::Stream>> { use crate::disc::*; - // Pass 1: find PMT PID from PAT + // Pass 1: find PMT PID from PAT (table_id 0x00 on PID 0). + let pat = collect_psi_section(data, 0, 0x00)?; + let pat_section_len = (((pat[1] & 0x0F) as usize) << 8) | pat[2] as usize; + if pat_section_len < 4 { + return None; + } let mut pat_pmt_pid: Option<u16> = None; - let mut offset = 0; - while offset + BD_TS_PACKET_SIZE <= data.len() { - if data[offset + 4] != SYNC_BYTE { - offset += 1; - continue; - } - let pid = (((data[offset + 5] & 0x1F) as u16) << 8) | data[offset + 6] as u16; - let pusi = data[offset + 5] & 0x40 != 0; - - if pid == 0 && pusi { - let payload_start = offset + 4 + 4; - if payload_start + 12 < data.len() { - let pointer = data[payload_start] as usize; - let pat_start = payload_start + 1 + pointer; - if pat_start + 12 < data.len() && data[pat_start] == 0x00 { - let section_len = (((data[pat_start + 1] & 0x0F) as usize) << 8) - | data[pat_start + 2] as usize; - let entries_start = pat_start + 8; - if section_len < 4 { - offset += BD_TS_PACKET_SIZE; - continue; - } - let entries_end = pat_start + 3 + section_len - 4; - let mut e = entries_start; - while e + 4 <= data.len() && e < entries_end { - let prog_num = ((data[e] as u16) << 8) | data[e + 1] as u16; - let p = (((data[e + 2] & 0x1F) as u16) << 8) | data[e + 3] as u16; - if prog_num != 0 { - pat_pmt_pid = Some(p); - break; - } - e += 4; - } - } + { + let entries_start = 8; + // section_length counts bytes after the length field, incl. the + // 4-byte CRC; the program loop stops before the CRC. + let entries_end = (3 + pat_section_len - 4).min(pat.len()); + let mut e = entries_start; + while e + 4 <= entries_end { + let prog_num = ((pat[e] as u16) << 8) | pat[e + 1] as u16; + let p = (((pat[e + 2] & 0x1F) as u16) << 8) | pat[e + 3] as u16; + if prog_num != 0 { + pat_pmt_pid = Some(p); + break; } + e += 4; } - offset += BD_TS_PACKET_SIZE; } let pmt_pid = pat_pmt_pid?; - // Pass 2: parse PMT for stream entries + // Pass 2: parse PMT for stream entries (table_id 0x02 on pmt_pid). let mut streams = Vec::new(); - offset = 0; - while offset + BD_TS_PACKET_SIZE <= data.len() { - if data[offset + 4] != SYNC_BYTE { - offset += 1; - continue; + let pmt = collect_psi_section(data, pmt_pid, 0x02)?; + if pmt.len() >= 12 { + let section_len = (((pmt[1] & 0x0F) as usize) << 8) | pmt[2] as usize; + // section_length counts the bytes after this field, including the + // trailing 4-byte CRC; `< 4` would underflow `end` below. + if section_len < 4 { + return None; } - let pid = (((data[offset + 5] & 0x1F) as u16) << 8) | data[offset + 6] as u16; - let pusi = data[offset + 5] & 0x40 != 0; + let prog_info_len = (((pmt[10] & 0x0F) as usize) << 8) | pmt[11] as usize; + let mut pos = 12 + prog_info_len; + // Clamp the section end to the reassembled bytes; a malformed + // section_len or prog_info_len must never drive reads past `pmt`. + let end = (3 + section_len - 4).min(pmt.len()); - if pid == pmt_pid && pusi { - let payload_start = offset + 4 + 4; - if payload_start + 1 >= data.len() { - offset += BD_TS_PACKET_SIZE; - continue; - } - let pointer = data[payload_start] as usize; - let pmt_start = payload_start + 1 + pointer; - if pmt_start + 12 >= data.len() { - offset += BD_TS_PACKET_SIZE; - continue; - } - if data[pmt_start] != 0x02 { - offset += BD_TS_PACKET_SIZE; - continue; - } + while pos + 5 <= end { + let stream_type = pmt[pos]; + let es_pid = (((pmt[pos + 1] & 0x1F) as u16) << 8) | pmt[pos + 2] as u16; + let es_info_len = (((pmt[pos + 3] & 0x0F) as usize) << 8) | pmt[pos + 4] as usize; - let section_len = - (((data[pmt_start + 1] & 0x0F) as usize) << 8) | data[pmt_start + 2] as usize; - // section_length counts the bytes after this field, including the - // trailing 4-byte CRC; `< 4` would underflow `end` below. Guard it - // exactly like the PAT parser above. - if section_len < 4 { - offset += BD_TS_PACKET_SIZE; - continue; - } - let prog_info_len = - (((data[pmt_start + 10] & 0x0F) as usize) << 8) | data[pmt_start + 11] as usize; - let mut pos = pmt_start + 12 + prog_info_len; - // Clamp the section end to the buffer; a malformed section_len or - // prog_info_len must never drive reads past `data`. - let end = (pmt_start + 3 + section_len - 4).min(data.len()); - - while pos + 5 <= data.len() && pos < end { - let stream_type = data[pos]; - let es_pid = (((data[pos + 1] & 0x1F) as u16) << 8) | data[pos + 2] as u16; - let es_info_len = (((data[pos + 3] & 0x0F) as usize) << 8) | data[pos + 4] as usize; - - // Single source of truth for stream_type → Codec: reuse - // `Codec::from_coding_type` (the same table the BD STN / - // disc scanner uses) so the two mappings can never drift. - // We only retain the category (video/audio/subtitle) and - // per-kind default attribute logic here. - let codec = Codec::from_coding_type(stream_type); - let stream = match codec.kind() { - CodecKind::Video => { - // Default resolution by codec generation (HEVC → - // UHD, MPEG-2 → 1080i, else 1080p); refined later - // from the actual elementary stream. - let resolution = match codec { - Codec::Hevc => Resolution::R2160p, - Codec::Mpeg2 => Resolution::R1080i, - _ => Resolution::R1080p, - }; - Some(Stream::Video(VideoStream { - pid: es_pid, - codec, - resolution, - frame_rate: FrameRate::Unknown, - hdr: HdrFormat::Sdr, - color_space: ColorSpace::Bt709, - secondary: false, - label: String::new(), - })) - } - CodecKind::Audio => Some(Stream::Audio(AudioStream { + // Single source of truth for stream_type → Codec: reuse + // `Codec::from_coding_type` (the same table the BD STN / + // disc scanner uses) so the two mappings can never drift. + // We only retain the category (video/audio/subtitle) and + // per-kind default attribute logic here. + let codec = Codec::from_coding_type(stream_type); + let stream = match codec.kind() { + CodecKind::Video => { + // Default resolution by codec generation (HEVC → + // UHD, MPEG-2 → 1080i, else 1080p); refined later + // from the actual elementary stream. + let resolution = match codec { + Codec::Hevc => Resolution::R2160p, + Codec::Mpeg2 => Resolution::R1080i, + _ => Resolution::R1080p, + }; + Some(Stream::Video(VideoStream { pid: es_pid, codec, - channels: AudioChannels::Surround51, - language: "und".into(), - sample_rate: SampleRate::S48, + resolution, + frame_rate: FrameRate::Unknown, + hdr: HdrFormat::Sdr, + color_space: ColorSpace::Bt709, secondary: false, - purpose: crate::disc::LabelPurpose::Normal, label: String::new(), - })), - CodecKind::Subtitle => Some(Stream::Subtitle(SubtitleStream { - pid: es_pid, - codec, - language: "und".into(), - forced: false, - qualifier: crate::disc::LabelQualifier::None, - codec_data: None, - })), - CodecKind::Unknown => { - tracing::warn!( - target: "mux", - "dropping PMT stream entry with unknown stream_type {:#04x} (PID {:#06x})", - stream_type, - es_pid, - ); - None - } - }; - - if let Some(s) = stream { - streams.push(s); + })) } - pos += 5 + es_info_len; + CodecKind::Audio => Some(Stream::Audio(AudioStream { + pid: es_pid, + codec, + channels: AudioChannels::Surround51, + language: "und".into(), + sample_rate: SampleRate::S48, + secondary: false, + purpose: crate::disc::LabelPurpose::Normal, + label: String::new(), + })), + CodecKind::Subtitle => Some(Stream::Subtitle(SubtitleStream { + pid: es_pid, + codec, + language: "und".into(), + forced: false, + qualifier: crate::disc::LabelQualifier::None, + codec_data: None, + })), + CodecKind::Unknown => { + tracing::warn!( + target: "mux", + "dropping PMT stream entry with unknown stream_type {:#04x} (PID {:#06x})", + stream_type, + es_pid, + ); + None + } + }; + + if let Some(s) = stream { + streams.push(s); } - break; + pos += 5 + es_info_len; } - offset += BD_TS_PACKET_SIZE; } if streams.is_empty() { @@ -646,6 +748,139 @@ mod tests { bdts_packet(body, pmt_pid, true) } + /// Build a 192-byte BD-TS data packet on `pid` carrying `payload` + /// (payload-only adaptation, truncated/padded to fit one packet). + fn data_packet(pid: u16, pusi: bool, payload: &[u8]) -> Vec<u8> { + let mut pkt = vec![0u8; BD_TS_PACKET_SIZE]; + pkt[4] = SYNC_BYTE; + pkt[5] = ((pid >> 8) as u8) & 0x1F; + if pusi { + pkt[5] |= 0x40; + } + pkt[6] = (pid & 0xFF) as u8; + pkt[7] = 0x10; // payload only, no adaptation field + let room = TS_PACKET_SIZE - 4; // 184 ES bytes after the 4-byte TS header + let n = payload.len().min(room); + pkt[8..8 + n].copy_from_slice(&payload[..n]); + pkt + } + + /// Like `pmt_packet` but with a 2-byte adaptation field (AFC=0b11) of + /// stuffing before the payload, to exercise the adaptation-field-aware + /// payload base computation in scan_streams. + fn pmt_packet_with_af(pmt_pid: u16, entries: &[(u8, u16)]) -> Vec<u8> { + let af_len: u8 = 2; // 1 flags byte + 1 stuffing byte + let mut pkt = vec![0u8; BD_TS_PACKET_SIZE]; + pkt[4] = SYNC_BYTE; + pkt[5] = (((pmt_pid >> 8) as u8) & 0x1F) | 0x40; // PUSI set + pkt[6] = (pmt_pid & 0xFF) as u8; + pkt[7] = 0x30; // AFC = 0b11 (adaptation + payload) + pkt[8] = af_len; // adaptation_field_length + pkt[9] = 0x00; // AF flags + pkt[10] = 0xFF; // stuffing + // Payload (PSI) begins at 4 + 4 + 1 + af_len = 11. + let payload_off = 4 + 4 + 1 + af_len as usize; + let mut body = vec![0xFFu8; BD_TS_PACKET_SIZE - payload_off]; + body[0] = 0x00; // pointer_field + let s = 1; + body[s] = 0x02; // table_id = PMT + let entries_len = entries.len() * 5; + let section_length = 9 + entries_len + 4; + body[s + 1] = 0xB0 | (((section_length >> 8) as u8) & 0x0F); + body[s + 2] = (section_length & 0xFF) as u8; + body[s + 3] = 0x00; + body[s + 4] = 0x01; + body[s + 5] = 0xC1; + body[s + 6] = 0x00; + body[s + 7] = 0x00; + body[s + 8] = 0xE0; + body[s + 9] = 0x00; + body[s + 10] = 0xF0; + body[s + 11] = 0x00; + let mut p = s + 12; + for &(stype, es_pid) in entries { + body[p] = stype; + body[p + 1] = 0xE0 | (((es_pid >> 8) as u8) & 0x1F); + body[p + 2] = (es_pid & 0xFF) as u8; + body[p + 3] = 0xF0; + body[p + 4] = 0x00; + p += 5; + } + pkt[payload_off..].copy_from_slice(&body); + pkt + } + + #[test] + fn short_pes_payload_injects_no_header_bytes() { + // A PUSI packet whose payload is NOT a valid PES start + // (no 00 00 01 start code / too short) must contribute ZERO bytes to + // the assembled elementary stream — otherwise a stray 00 00 01 in the + // garbage masquerades as an Annex-B NAL / PES start code in the codec + // parser. Only the following well-formed continuation bytes survive. + let pid = 0x1011; + let mut demux = TsDemuxer::new(&[pid]); + + // Garbage PUSI payload with NO valid PES start code (no leading + // 00 00 01). It must parse as malformed → header_len 0 → nothing + // pushed. The bytes include a 00 00 01 03 sequence mid-payload that, + // if leaked, would masquerade as an Annex-B NAL / PES start code. + let mut garbage = vec![0xAAu8; 32]; + garbage[8] = 0x00; + garbage[9] = 0x00; + garbage[10] = 0x01; + garbage[11] = 0x03; + let mut stream = demux.feed(&data_packet(pid, true, &garbage)); + assert!( + stream.is_empty(), + "garbage PUSI packet must not complete a PES on its own" + ); + + // Continuation packet (no PUSI) carrying real ES bytes. + let es = [0xDEu8, 0xAD, 0xBE, 0xEF]; + stream.extend(demux.feed(&data_packet(pid, false, &es))); + stream.extend(demux.flush()); + + assert_eq!(stream.len(), 1, "one PES assembled from the continuation"); + let pes = &stream[0]; + // The continuation ES bytes survive… + assert!( + pes.data.windows(es.len()).any(|w| w == es), + "continuation ES bytes present, got {:02X?}", + pes.data + ); + // …but none of the garbage PUSI payload leaked in. In particular the + // 0xAA filler and the embedded 00 00 01 sequence must be absent — the + // malformed PES header contributed ZERO bytes to the elementary stream. + assert!( + !pes.data.iter().any(|&b| b == 0xAA), + "garbage PES-header bytes must not appear in the elementary stream" + ); + assert!( + !pes.data.windows(3).any(|w| w == [0x00, 0x00, 0x01]), + "no injected start code leaked from the malformed PES header" + ); + } + + #[test] + fn scan_streams_handles_adaptation_field_in_pmt() { + use crate::disc::{Codec, Stream}; + let pmt_pid = 0x0100; + let mut data = pat_packet(pmt_pid); + // PMT carried in a packet with an adaptation field — payload base must + // account for af_len, not assume offset+8. + data.extend(pmt_packet_with_af(pmt_pid, &[(0x1B, 0x1011)])); + // Follower sync byte so is_resync_point corroborates the PMT packet. + data.extend(pat_packet(pmt_pid)); + + let streams = scan_streams(&data).expect("PMT with AF should parse"); + assert!( + streams + .iter() + .any(|s| matches!(s, Stream::Video(v) if v.codec == Codec::H264)), + "H.264 video must be found past the adaptation field" + ); + } + #[test] fn scan_streams_maps_lpcm_via_from_coding_type() { use crate::disc::{Codec, Stream}; @@ -673,4 +908,86 @@ mod tests { "H.264 video present" ); } + + /// Build a PMT whose reassembled section spans MORE than one 184-byte + /// TS payload, returned as two BD-TS packets: a PUSI packet carrying + /// the section head and a continuation (no-PUSI) packet carrying the + /// tail. The reassembler must stitch them back together; a flat-slice + /// parser would read the continuation packet's TS header as table + /// content and mis-type or drop the trailing entries. + fn pmt_two_packets(pmt_pid: u16, entries: &[(u8, u16)]) -> Vec<u8> { + // Assemble the raw PSI section (table_id + length + body + CRC). + let entries_len = entries.len() * 5; + let section_length = 9 + entries_len + 4; // fixed PMT fields + entries + CRC + let mut section = Vec::new(); + section.push(0x02); // table_id + section.push(0xB0 | (((section_length >> 8) as u8) & 0x0F)); + section.push((section_length & 0xFF) as u8); + section.extend_from_slice(&[0x00, 0x01]); // program_number + section.push(0xC1); // version/current_next + section.push(0x00); // section_number + section.push(0x00); // last_section_number + section.extend_from_slice(&[0xE0, 0x00]); // PCR PID + section.extend_from_slice(&[0xF0, 0x00]); // program_info_length = 0 + for &(stype, es_pid) in entries { + section.push(stype); + section.push(0xE0 | (((es_pid >> 8) as u8) & 0x1F)); + section.push((es_pid & 0xFF) as u8); + section.extend_from_slice(&[0xF0, 0x00]); // ES_info_length = 0 + } + section.extend_from_slice(&[0xFF, 0xFF, 0xFF, 0xFF]); // CRC (unchecked) + + // First packet payload: pointer_field(0) + as much section as fits. + let first_cap = 184 - 1; // minus pointer_field + let head_len = first_cap.min(section.len()); + let mut p0 = [0xFFu8; 184]; + p0[0] = 0x00; // pointer_field + p0[1..1 + head_len].copy_from_slice(§ion[..head_len]); + let pkt0 = bdts_packet(p0, pmt_pid, true); + + // Continuation packet (no PUSI) carries the rest. + let mut p1 = [0xFFu8; 184]; + let tail = §ion[head_len..]; + assert!(!tail.is_empty(), "test must actually span two packets"); + p1[..tail.len()].copy_from_slice(tail); + let pkt1 = bdts_packet(p1, pmt_pid, false); + + let mut out = pkt0; + out.extend(pkt1); + out + } + + #[test] + fn scan_streams_reassembles_pmt_across_packets() { + use crate::disc::{Codec, Stream}; + let pmt_pid = 0x0100; + // Enough entries that the section exceeds one 183-byte payload: + // 12 fixed + 4*N*... at 5 bytes/entry; 40 entries = 200 bytes of + // entries alone, forcing a continuation packet. + let mut entries: Vec<(u8, u16)> = Vec::new(); + entries.push((0x1B, 0x1011)); // H.264 video + for i in 0..40u16 { + entries.push((0x80, 0x1100 + i)); // LPCM audio tracks + } + let mut data = pat_packet(pmt_pid); + data.extend(pmt_two_packets(pmt_pid, &entries)); + + let streams = scan_streams(&data).expect("multi-packet PMT should parse"); + // All entries must survive reassembly (video + 40 audio). + assert_eq!(streams.len(), entries.len(), "every PMT entry reassembled"); + assert!( + streams + .iter() + .any(|s| matches!(s, Stream::Video(v) if v.codec == Codec::H264)), + "video survives the split" + ); + // The LAST audio entry lives in the continuation packet — proves + // the tail was stitched in, not read from a TS header. + assert!( + streams.iter().any( + |s| matches!(s, Stream::Audio(a) if a.pid == 0x1100 + 39 && a.codec == Codec::Lpcm) + ), + "trailing audio entry from the continuation packet survives" + ); + } } diff --git a/src/mux/tsmux.rs b/src/mux/tsmux.rs index cbdb801..83e7e26 100644 --- a/src/mux/tsmux.rs +++ b/src/mux/tsmux.rs @@ -4,17 +4,47 @@ //! packets. Each frame is wrapped in a PES header, split into TS packets, //! and prepended with the 4-byte TP_extra_header. +use super::hevc::{hvcc_to_annex_b, length_prefixed_to_annex_b}; use std::io::{self, Write}; const SYNC_BYTE: u8 = 0x47; const TS_PAYLOAD: usize = 184; +/// PID range treated as video (HEVC, triggers Annex-B conversion + RAI +/// on keyframes). Both `write_frame` and `build_pes_header` consult this +/// so a PID's stream_id and its NAL handling can never disagree. +const VIDEO_PID_RANGE: std::ops::RangeInclusive<u16> = 0x1011..=0x101F; + +/// Largest PES payload that fits a bounded `PES_packet_length` (u16) on a +/// `0xBD` (private_stream_1) stream after the 8 PES-header bytes. Frames +/// larger than this are split into multiple PES so the length field stays +/// spec-conformant (the unbounded `0` length is only legal for video). +const MAX_BD_PES_PAYLOAD: usize = u16::MAX as usize - 8; + +fn is_video_pid(pid: u16) -> bool { + VIDEO_PID_RANGE.contains(&pid) +} + +/// BD-TS muxer: PES frames in, 192-byte BD-TS packets out. +/// +/// Constructed over an output writer and a slice of per-track PIDs. The +/// `track` index passed to [`TsMuxer::write_frame`] and +/// [`TsMuxer::set_codec_private`] is the position in that PID slice; all +/// per-track state vectors are sized to `pids.len()`. PIDs in +/// `0x1011..=0x101F` are treated as video (length-prefixed NALUs in, +/// Annex B out, with parameter-set prepend and RAI on keyframes); every +/// other PID is carried as `private_stream_1` (`0xBD`) audio/subtitle. +/// All tracks share one PTS origin seeded from the first video frame, so +/// audio/video PTS offsets are preserved. pub struct TsMuxer<W: Write> { writer: W, pids: Vec<u16>, continuity: Vec<u8>, // per-PID continuity counter (0-15) codec_privates: Vec<Option<Vec<u8>>>, // per-track codec_private (for video parameter sets) params_written: Vec<bool>, // per-track: have we written parameter sets? + /// Global PTS origin (nanoseconds), seeded by the FIRST video frame so + /// the audio/video offset is preserved. Frames that arrive before it + /// is set saturate to 0. base_pts_ns: Option<i64>, } @@ -33,15 +63,29 @@ impl<W: Write> TsMuxer<W> { /// Set codec_private data for a track. Used to prepend VPS/SPS/PPS /// as Annex B NALs before the first keyframe in the transport stream. - pub fn set_codec_private(&mut self, track: usize, data: Vec<u8>) { - if track < self.codec_privates.len() { - self.codec_privates[track] = Some(data); + /// + /// `track` is the index into the PID slice passed to [`TsMuxer::new`]. + /// Returns [`Error::MuxTrackRange`](crate::error::Error::MuxTrackRange) + /// for an out-of-range index. + pub fn set_codec_private(&mut self, track: usize, data: Vec<u8>) -> io::Result<()> { + if track >= self.codec_privates.len() { + return Err(crate::error::Error::MuxTrackRange { + track, + tracks: self.codec_privates.len(), + } + .into()); } + self.codec_privates[track] = Some(data); + Ok(()) } /// Write a PES frame as BD-TS packets. /// Video frame data is expected as length-prefixed NALUs (MKV/PES format) /// and is converted to Annex B for transport stream. + /// + /// `track` is the index into the PID slice passed to [`TsMuxer::new`]. + /// Returns [`Error::MuxTrackRange`](crate::error::Error::MuxTrackRange) + /// for an out-of-range index. pub fn write_frame( &mut self, track: usize, @@ -50,10 +94,14 @@ impl<W: Write> TsMuxer<W> { data: &[u8], ) -> io::Result<()> { if track >= self.pids.len() { - return Ok(()); // unknown track, skip + return Err(crate::error::Error::MuxTrackRange { + track, + tracks: self.pids.len(), + } + .into()); } let pid = self.pids[track]; - let is_video = (0x1011..=0x101F).contains(&pid); + let is_video = is_video_pid(pid); // Drop non-key video before any keyframe — decoder has no IDR or // parameter sets to anchor on. @@ -61,12 +109,26 @@ impl<W: Write> TsMuxer<W> { return Ok(()); } - let base = *self.base_pts_ns.get_or_insert(pts_ns); - let pts_ns = pts_ns - base; + // Seed the global PTS origin from the FIRST video frame only, so the + // audio/video offset is preserved. A leading audio frame must not + // pull the base up and collapse the first video IDR to t=0. + if is_video { + self.base_pts_ns.get_or_insert(pts_ns); + } + let base = self.base_pts_ns.unwrap_or(pts_ns); + let pts_ns = pts_ns.saturating_sub(base); // For video: convert length-prefixed NALUs to Annex B (start codes). // Prepend codec_private parameter sets on the FIRST keyframe only. - let es_data = if is_video && !data.is_empty() { + // + // Arm `params_written` on the first video keyframe regardless of + // whether it carries data: an empty-data keyframe still anchors + // the stream, and leaving the flag unset would make every later + // non-key frame fail the drop guard above and silently vanish. + // For non-video the ES bytes pass through unchanged, so borrow + // `data` directly rather than copying it; only video needs an + // owned Annex-B conversion buffer. + let es_data: std::borrow::Cow<'_, [u8]> = if is_video { let mut annex_b = Vec::new(); if keyframe && !self.params_written[track] { if let Some(ref cp) = self.codec_privates[track] { @@ -77,25 +139,58 @@ impl<W: Write> TsMuxer<W> { self.params_written[track] = true; } annex_b.extend_from_slice(&length_prefixed_to_annex_b(data)); - annex_b + std::borrow::Cow::Owned(annex_b) } else { - data.to_vec() + std::borrow::Cow::Borrowed(data) }; - // Build PES packet: header + data let pts_90k = if pts_ns >= 0 { (pts_ns as u64).saturating_mul(9) / 100_000 } else { 0 }; - let pes_header = build_pes_header(pid, pts_90k, es_data.len()); - let pes_packet = [&pes_header[..], &es_data[..]].concat(); - // Split into TS packets + // Video PES may be unbounded (length 0); a 0xBD private_stream_1 + // PES must carry a bounded length, so split oversized audio/sub + // access units into multiple PES packets. Each emitted PES carries + // the same PTS and starts on its own PUSI packet (only the keyframe + // RAI rides the first packet of the first PES). + if is_video || es_data.len() <= MAX_BD_PES_PAYLOAD { + self.write_pes_chain(track, pid, pts_90k, is_video, keyframe, &es_data)?; + } else { + let mut first_pes = true; + for chunk in es_data.chunks(MAX_BD_PES_PAYLOAD) { + self.write_pes_chain(track, pid, pts_90k, is_video, keyframe && first_pes, chunk)?; + first_pes = false; + } + } + Ok(()) + } + + /// Wrap `es_data` in a PES header and split it into 192-byte BD-TS + /// packets. `keyframe` drives the RAI bit on the first packet (video + /// only). The PES header and ES bytes are sliced in place — no second + /// full-frame copy. + fn write_pes_chain( + &mut self, + track: usize, + pid: u16, + pts_90k: u64, + is_video: bool, + keyframe: bool, + es_data: &[u8], + ) -> io::Result<()> { + let pes_header = build_pes_header(pid, pts_90k, es_data.len()); + + // Logical PES packet = header bytes followed by es_data. It is + // indexed (and written) in place, without materializing the + // concatenation, to avoid a second full-frame copy on the hot path. + let pes_len = pes_header.len() + es_data.len(); + let mut offset = 0; let mut first = true; - while offset < pes_packet.len() { - let remaining = pes_packet.len() - offset; + while offset < pes_len { + let remaining = pes_len - offset; // Invariant: TP_extra(4) + TS_header(4) + AF(af_bytes) + payload(payload_len) = 192, // i.e. af_bytes + payload_len = TS_PAYLOAD (184). @@ -165,8 +260,20 @@ impl<W: Write> TsMuxer<W> { } } - self.writer - .write_all(&pes_packet[offset..offset + payload_len])?; + // Write the payload span [offset, offset+payload_len), which may + // straddle the header/es_data boundary — emit each side in one + // write_all rather than copying the whole frame again. + let end = offset + payload_len; + let hdr_len = pes_header.len(); + if offset < hdr_len { + let hdr_end = end.min(hdr_len); + self.writer.write_all(&pes_header[offset..hdr_end])?; + } + if end > hdr_len { + let es_start = offset.max(hdr_len) - hdr_len; + let es_end = end - hdr_len; + self.writer.write_all(&es_data[es_start..es_end])?; + } offset += payload_len; first = false; @@ -175,6 +282,8 @@ impl<W: Write> TsMuxer<W> { Ok(()) } + /// Flush the underlying writer. BD-TS needs no stream trailer, so this + /// only drains buffering; the muxer remains usable afterwards. pub fn finish(&mut self) -> io::Result<()> { self.writer.flush() } @@ -183,7 +292,7 @@ impl<W: Write> TsMuxer<W> { /// Build a PES packet header for a BD stream. fn build_pes_header(pid: u16, pts_90k: u64, data_len: usize) -> Vec<u8> { // Determine stream_id from PID range - let stream_id: u8 = if (0x1011..=0x101F).contains(&pid) { + let stream_id: u8 = if is_video_pid(pid) { 0xE0 // video } else { 0xBD // audio, PGS subtitle, or default (private stream 1) @@ -198,7 +307,10 @@ fn build_pes_header(pid: u16, pts_90k: u64, data_len: usize) -> Vec<u8> { header.push(0x01); header.push(stream_id); - // PES packet length (0 = unbounded for video or if too large for u16) + // PES packet length. The unbounded form (0) is only spec-legal for + // video; `write_frame` splits oversized 0xBD access units so a private + // stream always fits a bounded u16 length here. The `> 65535` arm + // remains a defensive fallback for video only. if stream_id == 0xE0 || pes_data_len > 65535 { header.push(0x00); header.push(0x00); @@ -226,72 +338,6 @@ fn build_pes_header(pid: u16, pts_90k: u64, data_len: usize) -> Vec<u8> { header } -/// Extract NAL arrays from HEVCDecoderConfigurationRecord and convert to Annex B. -/// Returns VPS + SPS + PPS as Annex B NAL units (00 00 00 01 + NAL). -fn hvcc_to_annex_b(hvcc: &[u8]) -> Option<Vec<u8>> { - // HEVCDecoderConfigurationRecord: 22 bytes header, then NAL arrays - if hvcc.len() < 23 { - return None; - } - let num_arrays = hvcc[22] as usize; - let mut out = Vec::new(); - let mut offset = 23; - - for _ in 0..num_arrays { - if offset + 3 > hvcc.len() { - break; - } - // array: 1 byte (completeness + NAL type), 2 bytes (numNalus) - let _nal_type = hvcc[offset] & 0x3F; - let num_nalus = u16::from_be_bytes([hvcc[offset + 1], hvcc[offset + 2]]) as usize; - offset += 3; - - for _ in 0..num_nalus { - if offset + 2 > hvcc.len() { - break; - } - let nal_len = u16::from_be_bytes([hvcc[offset], hvcc[offset + 1]]) as usize; - offset += 2; - if offset + nal_len > hvcc.len() { - break; - } - out.extend_from_slice(&[0x00, 0x00, 0x00, 0x01]); - out.extend_from_slice(&hvcc[offset..offset + nal_len]); - offset += nal_len; - } - } - - if out.is_empty() { None } else { Some(out) } -} - -/// Convert length-prefixed NALUs (4-byte BE length + NAL) to Annex B -/// (00 00 00 01 + NAL). Used for video elementary streams in TS. -fn length_prefixed_to_annex_b(data: &[u8]) -> Vec<u8> { - let mut out = Vec::with_capacity(data.len()); - let mut offset = 0; - while offset + 4 <= data.len() { - let len = u32::from_be_bytes([ - data[offset], - data[offset + 1], - data[offset + 2], - data[offset + 3], - ]) as usize; - offset += 4; - if offset + len > data.len() { - break; - } - out.extend_from_slice(&[0x00, 0x00, 0x00, 0x01]); - out.extend_from_slice(&data[offset..offset + len]); - offset += len; - } - // If data doesn't look like length-prefixed NALs (no valid parse), - // return original data unchanged — it may already be Annex B. - if out.is_empty() && !data.is_empty() { - return data.to_vec(); - } - out -} - #[cfg(test)] mod tests { use super::*; @@ -444,7 +490,7 @@ mod tests { let mut sink: Vec<u8> = Vec::new(); { let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]); - mux.set_codec_private(0, hvcc); + mux.set_codec_private(0, hvcc).unwrap(); // Non-IDR before any IDR: should be dropped. let p = fake_hevc_nal(1, 50); mux.write_frame(0, 0, false, &p).unwrap(); @@ -477,6 +523,38 @@ mod tests { ); } + #[test] + fn empty_data_keyframe_arms_params_so_later_frames_survive() { + // An empty-data keyframe must still arm params_written; otherwise + // every subsequent non-key frame would be dropped by the + // pre-keyframe guard and the track would emit no real frames. + let mut sink: Vec<u8> = Vec::new(); + { + let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]); + // Keyframe with empty payload (e.g. a frame whose NALs were + // all stripped upstream) — anchors the stream. + mux.write_frame(0, 0, true, &[]).unwrap(); + // Now a real non-key frame; it must NOT be dropped. + let p = fake_hevc_nal(1, 80); + mux.write_frame(0, 41_000_000, false, &p).unwrap(); + mux.finish().unwrap(); + } + let packets = parse_bd_ts(&sink); + // The non-key frame's NAL body byte (0x02 = (1<<1)) must appear in + // a video payload — proof it wasn't dropped. + let video_bytes: Vec<u8> = packets + .iter() + .filter(|p| p.pid == VIDEO_PID) + .flat_map(|p| p.payload.clone()) + .collect(); + assert!( + video_bytes + .windows(4) + .any(|w| w == [0x00, 0x00, 0x00, 0x01]), + "later non-key frame must survive after an empty-data keyframe" + ); + } + #[test] fn non_key_before_first_keyframe_dropped() { let mut sink: Vec<u8> = Vec::new(); @@ -493,4 +571,96 @@ mod tests { "non-key before first keyframe must be dropped" ); } + + const AUDIO_PID: u16 = 0x1100; + + /// Decode the 33-bit PTS from the first PUSI packet on `pid`. Assumes + /// the PES header carries PTS (flags 0x80 at PES byte 7). + fn first_pts_90k(packets: &[TsPacket], pid: u16) -> u64 { + let pkt = packets + .iter() + .find(|p| p.pid == pid && p.pusi) + .expect("PUSI packet present"); + // PES payload starts the packet payload: 00 00 01 stream_id len len + // flags1 flags2 hdr_len then 5 PTS bytes. + let p = &pkt.payload; + let pts = &p[9..14]; + ((((pts[0] >> 1) & 0x07) as u64) << 30) + | ((pts[1] as u64) << 22) + | (((pts[2] >> 1) as u64) << 15) + | ((pts[3] as u64) << 7) + | ((pts[4] >> 1) as u64) + } + + #[test] + fn av_offset_preserved_with_audio_before_first_video() { + // Audio at t=0 arrives BEFORE the first video keyframe at t=1s. + // The global base must be seeded from the VIDEO frame so the + // audio/video PTS offset is preserved (audio earlier ⇒ saturates to + // 0, video lands at +1s = 90000 ticks), not both collapsed to 0. + let mut sink: Vec<u8> = Vec::new(); + { + let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID, AUDIO_PID]); + // Audio frame first, at PTS 0. + mux.write_frame(1, 0, false, &[0x0B, 0x77, 0x00, 0x00]) + .unwrap(); + // Video keyframe at PTS 1s — seeds the base. + let idr = fake_hevc_nal(19, 100); + mux.write_frame(0, 1_000_000_000, true, &idr).unwrap(); + mux.finish().unwrap(); + } + let packets = parse_bd_ts(&sink); + let video_pts = first_pts_90k(&packets, VIDEO_PID); + let audio_pts = first_pts_90k(&packets, AUDIO_PID); + // Video keyframe is the base ⇒ its relative PTS is 0. + assert_eq!(video_pts, 0, "video keyframe seeds the base at t=0"); + // Audio arrived 1s earlier ⇒ saturates to 0, NOT lifted past video. + assert_eq!(audio_pts, 0, "earlier audio saturates to 0"); + assert!( + audio_pts <= video_pts, + "audio must not be pulled ahead of the video base" + ); + } + + #[test] + fn out_of_range_track_errors() { + let mut sink: Vec<u8> = Vec::new(); + let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]); + let err = mux.write_frame(5, 0, true, &[0xAA]).unwrap_err(); + assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput); + let err2 = mux.set_codec_private(5, vec![0u8; 4]).unwrap_err(); + assert_eq!(err2.kind(), std::io::ErrorKind::InvalidInput); + } + + #[test] + fn oversized_bd_audio_pes_is_split_and_bounded() { + // A private_stream_1 (0xBD) audio frame larger than the bounded PES + // limit must be split into multiple PES, each with a non-zero + // PES_packet_length (never the unbounded 0 form, which is illegal + // for 0xBD). + let mut sink: Vec<u8> = Vec::new(); + let big: Vec<u8> = (0..(MAX_BD_PES_PAYLOAD + 5000)) + .map(|i| (i & 0xFF) as u8) + .collect(); + { + let mut mux = TsMuxer::new(&mut sink, &[AUDIO_PID]); + mux.write_frame(0, 0, false, &big).unwrap(); + mux.finish().unwrap(); + } + let packets = parse_bd_ts(&sink); + let pusi: Vec<&TsPacket> = packets + .iter() + .filter(|p| p.pid == AUDIO_PID && p.pusi) + .collect(); + assert!( + pusi.len() >= 2, + "oversized audio must span ≥2 PES, got {}", + pusi.len() + ); + for p in pusi { + // PES length field at payload bytes [4..6] must be non-zero. + let len = u16::from_be_bytes([p.payload[4], p.payload[5]]); + assert_ne!(len, 0, "0xBD PES must carry a bounded length"); + } + } } diff --git a/src/pes.rs b/src/pes.rs index 60a260b..939e68f 100644 --- a/src/pes.rs +++ b/src/pes.rs @@ -6,6 +6,12 @@ //! disc.read() → PES frame (sectors → decrypt → demux internally) //! mkv.write(frame) → MKV file (mux internally) +/// Maximum frame payload size, shared by `serialize` and `deserialize` +/// so the wire format round-trips: any frame that serializes can be read +/// back. A frame larger than this is rejected on write rather than written +/// and then hard-erroring mid-stream on read. +const MAX_FRAME_SIZE: usize = 256 * 1024 * 1024; // 256 MiB + /// One frame of elementary stream data. #[derive(Debug, Clone)] pub struct PesFrame { @@ -27,9 +33,11 @@ impl PesFrame { /// Serialize to bytes: track(1) | pts(8) | keyframe(1) | len(4) | data pub fn serialize(&self, w: &mut dyn std::io::Write) -> std::io::Result<()> { if self.track > 255 { - return Err(crate::error::Error::PesInvalidMagic.into()); + return Err(crate::error::Error::PesTrackTooLarge { track: self.track }.into()); } - if self.data.len() > u32::MAX as usize { + // Enforce the same ceiling the reader uses, so a frame that writes + // can always be read back (round-trippable wire format). + if self.data.len() > MAX_FRAME_SIZE { return Err(crate::error::Error::PesFrameTooLarge { size: self.data.len(), } @@ -42,16 +50,35 @@ impl PesFrame { w.write_all(&self.data) } - /// Deserialize from bytes. Returns None at EOF. + /// Deserialize from bytes. Returns None at a clean end of stream. + /// + /// A clean EOF is exactly zero bytes available before the next frame. + /// A partial header (1-13 bytes, e.g. a crash or short write) is a real + /// error (`UnexpectedEof`), not silently treated as EOF — otherwise + /// truncated `.pes` data would be accepted as a graceful end. pub fn deserialize(r: &mut dyn std::io::Read) -> std::io::Result<Option<Self>> { - const MAX_FRAME_SIZE: usize = 256 * 1024 * 1024; // 256 MB - - let mut header = [0u8; 14]; // 1 + 8 + 1 + 4 - match r.read_exact(&mut header) { + // Probe one byte first to distinguish clean EOF from a truncated + // header. + let mut first = [0u8; 1]; + match r.read(&mut first) { + Ok(0) => return Ok(None), // clean EOF, no frame started Ok(_) => {} - Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None), + Err(e) if e.kind() == std::io::ErrorKind::Interrupted => { + // Retry-once on EINTR before committing to the header read. + match r.read(&mut first) { + Ok(0) => return Ok(None), + Ok(_) => {} + Err(e) => return Err(e), + } + } Err(e) => return Err(e), } + + let mut header = [0u8; 14]; // 1 + 8 + 1 + 4 + header[0] = first[0]; + // The remaining 13 header bytes must be present; a short read here is + // a truncated frame, propagated as UnexpectedEof. + r.read_exact(&mut header[1..])?; let track = header[0] as usize; let pts = i64::from_le_bytes([ header[1], header[2], header[3], header[4], header[5], header[6], header[7], header[8], @@ -73,7 +100,10 @@ impl PesFrame { } /// Create from a codec::Frame with a track index. - pub fn from_codec_frame(track: usize, frame: crate::mux::codec::Frame) -> Self { + /// + /// `pub(crate)`: takes the internal `mux::codec::Frame` type, so it + /// can't be part of the public API surface. + pub(crate) fn from_codec_frame(track: usize, frame: crate::mux::codec::Frame) -> Self { Self { track, pts: frame.pts_ns, @@ -167,8 +197,11 @@ impl Stream for CountingStream { } fn write(&mut self, frame: &PesFrame) -> std::io::Result<()> { + // Only count bytes that actually made it to the inner sink, so a + // failed write doesn't permanently inflate bytes_written(). + self.inner.write(frame)?; self.written += frame.data.len() as u64; - self.inner.write(frame) + Ok(()) } fn finish(&mut self) -> std::io::Result<()> { @@ -186,6 +219,10 @@ impl Stream for CountingStream { fn headers_ready(&self) -> bool { self.inner.headers_ready() } + + fn errors(&self) -> u64 { + self.inner.errors() + } } #[cfg(test)] @@ -280,4 +317,87 @@ mod tests { let _ = s.info(); s.finish().unwrap(); } + + #[test] + fn frame_roundtrips_through_bytes() { + let frame = make_frame(3, 123_456); + let mut buf = Vec::new(); + frame.serialize(&mut buf).expect("serialize"); + let mut cursor = std::io::Cursor::new(buf); + let got = PesFrame::deserialize(&mut cursor) + .expect("deserialize") + .expect("frame present"); + assert_eq!(got.track, frame.track); + assert_eq!(got.pts, frame.pts); + assert_eq!(got.keyframe, frame.keyframe); + assert_eq!(got.data, frame.data); + // Next read is a clean EOF. + assert!(PesFrame::deserialize(&mut cursor).unwrap().is_none()); + } + + #[test] + fn empty_input_is_clean_eof() { + let mut cursor = std::io::Cursor::new(Vec::new()); + assert!(PesFrame::deserialize(&mut cursor).unwrap().is_none()); + } + + #[test] + fn truncated_header_is_error_not_eof() { + // A partial 14-byte header (here 5 bytes) must surface as an error, + // not be swallowed as a graceful end of stream. + let mut cursor = std::io::Cursor::new(vec![1u8, 2, 3, 4, 5]); + let err = PesFrame::deserialize(&mut cursor).expect_err("partial header must error"); + assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof); + } + + #[test] + fn oversize_track_rejected_on_serialize() { + let frame = make_frame(256, 0); + let mut buf = Vec::new(); + let err = frame + .serialize(&mut buf) + .expect_err("track > 255 must fail"); + let code = format!("E{}", crate::error::E_PES_TRACK_TOO_LARGE); + assert!(err.to_string().contains(&code), "got: {err}"); + } + + /// Output stream whose `write` always fails — for CountingStream tests. + struct FailingWriteStream { + title: DiscTitle, + } + + impl Stream for FailingWriteStream { + fn read(&mut self) -> std::io::Result<Option<PesFrame>> { + Ok(None) + } + fn write(&mut self, _frame: &PesFrame) -> std::io::Result<()> { + Err(std::io::Error::from(std::io::ErrorKind::BrokenPipe)) + } + fn finish(&mut self) -> std::io::Result<()> { + Ok(()) + } + fn info(&self) -> &DiscTitle { + &self.title + } + } + + #[test] + fn counting_stream_does_not_count_failed_writes() { + let mut cs = CountingStream::new(Box::new(FailingWriteStream { + title: DiscTitle::empty(), + })); + let frame = make_frame(0, 0); + assert!(cs.write(&frame).is_err()); + // Failed write must not inflate the byte count. + assert_eq!(cs.bytes_written(), 0); + } + + #[test] + fn counting_stream_counts_successful_writes() { + let frame = make_frame(0, 0); + let payload = frame.data.len() as u64; + let mut cs = CountingStream::new(Box::new(MockStream::new(Vec::new()))); + cs.write(&frame).unwrap(); + assert_eq!(cs.bytes_written(), payload); + } } diff --git a/src/platform/fs_type/mod.rs b/src/platform/fs_type/mod.rs index cc6c4ea..ebabdb5 100644 --- a/src/platform/fs_type/mod.rs +++ b/src/platform/fs_type/mod.rs @@ -76,8 +76,17 @@ pub fn detect_fd(fd: std::os::unix::io::RawFd) -> FsType { detect_fd_impl(fd) } +/// Non-Linux stub for [`detect_fd`]. +/// +/// Always returns [`FsType::Unknown`]: only Linux keys its writeback +/// policy off this classification, so other platforms have nothing to +/// detect. The `fd` parameter is a bare `i32` rather than +/// `std::os::unix::io::RawFd` because this arm also compiles on Windows, +/// which has no `RawFd` — the universal integer keeps one signature across +/// all non-Linux targets. Unused on these targets (no caller invokes it), +/// hence `allow(dead_code)`. #[cfg(not(target_os = "linux"))] -#[allow(dead_code)] // API parity with the linux impl; callers cfg-gate. +#[allow(dead_code)] pub fn detect_fd(_fd: i32) -> FsType { FsType::Unknown } diff --git a/src/platform/fs_type/windows.rs b/src/platform/fs_type/windows.rs index b510d17..001e0e4 100644 --- a/src/platform/fs_type/windows.rs +++ b/src/platform/fs_type/windows.rs @@ -2,9 +2,10 @@ //! //! Heuristic-only: any UNC path (`\\server\share\...`) is treated as a //! network mount and bucketed into `Nfs`. Strictly, SMB is not NFS, but -//! the buffering-policy outcome for our purposes is the same — there is -//! no platform `WritebackFile` machinery on Windows yet, so the worst -//! case of a false positive is using `LocalFileSink` regardless. A +//! the buffering-policy outcome is the same here — there is no platform +//! `WritebackFile` machinery on Windows yet, so both `Nfs` and `Local` +//! select `LocalFileSink`. A misclassification in either direction is +//! therefore harmless on Windows today: the sink choice does not change. A //! proper `GetVolumeInformation` query is a Phase 4 concern. use std::path::Path; diff --git a/src/platform/mt1959/mod.rs b/src/platform/mt1959/mod.rs index bce01c5..00d124e 100644 --- a/src/platform/mt1959/mod.rs +++ b/src/platform/mt1959/mod.rs @@ -18,6 +18,8 @@ const BUFFER_ID_B: u8 = 0x77; // ── SCSI opcodes ────────────────────────────────────────────────────── const SCSI_READ_BUFFER: u8 = 0x3C; const SCSI_READ_CAPACITY: u8 = 0x25; +/// Shared by both firmware-upload variants (see `variant_a` / `variant_b`). +pub(super) const SCSI_WRITE_BUFFER: u8 = 0x3B; // ── Sub-commands (shared A/B) ───────────────────────────────────────── const SUB_CMD_UNLOCK: u8 = 0x00; @@ -111,6 +113,14 @@ impl Mt1959 { buf: &mut [u8], expected: usize, ) -> Result<usize> { + // The READ_BUFFER CDB transfer-length is a single byte; an + // `expected` above 255 cannot be expressed and would silently + // truncate. All in-crate callers pass small fixed sizes (4); guard + // the invariant rather than emit a malformed CDB. + debug_assert!( + expected <= u8::MAX as usize, + "read_buffer_probe expected exceeds 1-byte CDB length field" + ); let cdb = self.read_buffer_sub(sub_cmd, address, expected as u8); let result = scsi.execute(&cdb, DataDirection::FromDevice, buf, 5_000)?; if result.bytes_transferred != expected { @@ -146,16 +156,23 @@ impl Mt1959 { 0x00, ]; let mut response = vec![0u8; UNLOCK_RESPONSE_SIZE as usize]; - scsi.execute(&cdb, DataDirection::FromDevice, &mut response, 30_000)?; + let result = scsi.execute(&cdb, DataDirection::FromDevice, &mut response, 30_000)?; - if response.len() >= 4 && response[0..4] != self.profile.signature { + // `response` is a fixed 64-byte buffer, so `response.len()` is + // always >= every offset below — the meaningful bound is how many + // bytes the drive actually delivered. Validate against + // `bytes_transferred` so a short/partial transfer (stale trailing + // zeros) can't be read as if the drive sent real marker bytes. + let n = result.bytes_transferred.min(response.len()); + + if n >= 4 && response[0..4] != self.profile.signature { return Err(Error::SignatureMismatch { expected: self.profile.signature, got: response[0..4].try_into().unwrap_or([0; 4]), }); } - if response.len() >= FIRMWARE_ACTIVE_OFFSET + 4 + if n >= FIRMWARE_ACTIVE_OFFSET + 4 && response[FIRMWARE_ACTIVE_OFFSET..FIRMWARE_ACTIVE_OFFSET + 4] != FIRMWARE_ACTIVE_SIG { return Err(Error::UnlockFailed); @@ -170,7 +187,7 @@ impl Mt1959 { // the rest of the response. Requiring both before we tell the // upper layer "OEM path is live" keeps any partial / corrupted // response from steering us off the cert-auth fallback. - self.unlocked = response.len() >= FIRMWARE_MODE_OFFSET + 4 + self.unlocked = n >= FIRMWARE_MODE_OFFSET + 4 && response[FIRMWARE_ACTIVE_OFFSET..FIRMWARE_ACTIVE_OFFSET + 4] == FIRMWARE_ACTIVE_SIG && response[FIRMWARE_MODE_OFFSET..FIRMWARE_MODE_OFFSET + 4] == FIRMWARE_MODE_SIG; @@ -272,7 +289,11 @@ impl Mt1959 { .execute(&cap_cdb, DataDirection::FromDevice, &mut cap_buf, 5_000) .is_ok() { - u32::from_be_bytes([cap_buf[0], cap_buf[1], cap_buf[2], cap_buf[3]]) + 1 + // last_lba + 1 = sector count. A 0xFFFFFFFF last-LBA is the + // READ CAPACITY(10) "capacity exceeds 32 bits" sentinel; saturate + // rather than wrap to 0 (which would misclassify a huge disc as + // BD). A saturated count stays above the UHD threshold -> UHD. + u32::from_be_bytes([cap_buf[0], cap_buf[1], cap_buf[2], cap_buf[3]]).saturating_add(1) } else { 0 }; diff --git a/src/platform/mt1959/variant_a.rs b/src/platform/mt1959/variant_a.rs index ccfa864..3fdebf5 100644 --- a/src/platform/mt1959/variant_a.rs +++ b/src/platform/mt1959/variant_a.rs @@ -6,17 +6,28 @@ use super::Mt1959; use crate::error::Result; use crate::scsi::{DataDirection, ScsiTransport}; -const SCSI_WRITE_BUFFER: u8 = 0x3B; +use super::SCSI_WRITE_BUFFER; + const VERIFY_BUFFER_ID: u8 = 0x45; +/// WRITE_BUFFER carries a 24-bit transfer length, so a firmware blob +/// larger than this cannot be uploaded in one command. +const WRITE_BUFFER_MAX_LEN: usize = 0x00FF_FFFF; + pub(super) fn load_firmware(mt: &mut Mt1959, scsi: &mut dyn ScsiTransport) -> Result<()> { let firmware = &mt.profile.firmware; if firmware.is_empty() { return Err(crate::error::Error::UnlockFailed); } - // Upload firmware via WRITE_BUFFER + // Upload firmware via WRITE_BUFFER. The CDB's length is a 24-bit field; + // if the blob exceeds that, the encoded length would silently disagree + // with the bytes actually sent (`data`). Reject rather than upload a + // length-mismatched command. let len = firmware.len(); + if len > WRITE_BUFFER_MAX_LEN { + return Err(crate::error::Error::UnlockFailed); + } let cdb = [ SCSI_WRITE_BUFFER, 0x06, @@ -53,8 +64,11 @@ pub(super) fn load_firmware(mt: &mut Mt1959, scsi: &mut dyn ScsiTransport) -> Re 5_000, ); - // Double unlock after firmware upload - mt.do_unlock(scsi)?; + // Double unlock after firmware upload. The first establishes the + // unlock and is fatal on failure; the second is a confirmation pass and + // is best-effort (matching variant B), so a benign hiccup on the + // redundant call doesn't fail an already-successful unlock. mt.do_unlock(scsi)?; + let _ = mt.do_unlock(scsi); Ok(()) } diff --git a/src/platform/mt1959/variant_b.rs b/src/platform/mt1959/variant_b.rs index 2637f09..a28f8b3 100644 --- a/src/platform/mt1959/variant_b.rs +++ b/src/platform/mt1959/variant_b.rs @@ -2,13 +2,11 @@ //! //! MODE SELECT (0x55) → read metadata → WRITE_BUFFER → vendor verify (0xF1) → unlock × 5+1 -use super::Mt1959; +use super::{Mt1959, SCSI_READ_BUFFER, SCSI_WRITE_BUFFER}; use crate::error::Result; use crate::scsi::{DataDirection, ScsiTransport}; const SCSI_MODE_SELECT: u8 = 0x55; -const SCSI_WRITE_BUFFER: u8 = 0x3B; -const SCSI_READ_BUFFER: u8 = 0x3C; const FIRMWARE_MAX_SIZE: usize = 0x9C0; const FIRMWARE_EXTRA: [u8; 16] = [0; 16]; const VENDOR_VERIFY: [u8; 10] = [0xF1, 0x01, 0x02, 0x00, 0x0D, 0x30, 0x01, 0xF3, 0xAD, 0x23]; @@ -19,7 +17,13 @@ pub(super) fn load_firmware(mt: &mut Mt1959, scsi: &mut dyn ScsiTransport) -> Re return Err(crate::error::Error::UnlockFailed); } - // Step 1: Upload firmware via MODE SELECT + // Step 1: Upload firmware via MODE SELECT. Variant-B firmware blobs are + // exactly FIRMWARE_MAX_SIZE; a larger blob means a corrupt/wrong profile, + // and silently truncating it would upload a partial image that can't + // unlock. Reject it explicitly instead. + if firmware.len() > FIRMWARE_MAX_SIZE { + return Err(crate::error::Error::UnlockFailed); + } let write_len = FIRMWARE_MAX_SIZE.min(firmware.len()); let mode_select_cdb = [ SCSI_MODE_SELECT, @@ -77,7 +81,11 @@ pub(super) fn load_firmware(mt: &mut Mt1959, scsi: &mut dyn ScsiTransport) -> Re let mut dummy = [0u8; 0]; let _ = scsi.execute(&VENDOR_VERIFY, DataDirection::None, &mut dummy, 5_000); - // Step 5: Unlock retries (up to 5, then final attempt) + // Step 5: Unlock retries (up to 5, then a final fatal attempt). On a + // successful unlock we issue one confirmation pass; its result is + // intentionally best-effort — the first call already established the + // unlock state, so a hiccup on the redundant confirmation must not fail + // an otherwise-good unlock. for _attempt in 0..5 { if mt.do_unlock(scsi).is_ok() { let _ = mt.do_unlock(scsi); diff --git a/src/profile.rs b/src/profile.rs index aa48ef3..938847f 100644 --- a/src/profile.rs +++ b/src/profile.rs @@ -31,8 +31,14 @@ pub struct Identity { #[derive(Debug, Clone, Deserialize)] pub struct DriveProfile { pub identity: Identity, + /// Expected first 4 bytes of the drive's unlock response — the + /// per-drive signature the platform checks before trusting the + /// extended-access surface. JSON-encoded as 8 lowercase hex chars. #[serde(default, deserialize_with = "deserialize_hex4")] pub signature: [u8; 4], + /// Runtime firmware image uploaded during unlock (variant A/B + /// firmware-load step). JSON-encoded as standard base64; empty when + /// the profile carries no firmware blob. #[serde(default, deserialize_with = "deserialize_base64")] pub firmware: Vec<u8>, @@ -82,31 +88,56 @@ pub enum Platform { } impl Platform { + /// Stable, language-neutral platform identifier. The two MT1959 variants + /// share the chipset but differ in their firmware-upload / unlock + /// sequence, so they get distinct suffixes — callers (and logs) that key + /// off `name()` must be able to tell A from B. pub fn name(&self) -> &'static str { match self { - Platform::Mt1959A => "MediaTek MT1959", - Platform::Mt1959B => "MediaTek MT1959", + Platform::Mt1959A => "MediaTek MT1959-A", + Platform::Mt1959B => "MediaTek MT1959-B", Platform::Renesas => "Renesas", } } } -/// Result of profile lookup. +/// Result of a profile lookup: the matched profile plus the platform +/// (chipset + variant) of the section it was found in. The platform +/// determines which unlock/firmware sequence the driver runs. pub struct ProfileMatch { + /// The matched profile, cloned out of the profiles file. pub profile: DriveProfile, + /// Which platform section the profile came from. pub platform: Platform, } // ── Parsing ──────────────────────────────────────────────────────────── +/// Decode an even-length ASCII hex string into bytes. +/// +/// Operates on raw bytes rather than `&str` char-boundary slices: a +/// non-ASCII input (e.g. a hand-edited profile with a multi-byte char) +/// could otherwise have `&s[i..i+2]` land inside a UTF-8 char boundary and +/// panic. Hex is ASCII, so any non-ASCII or non-hex byte simply fails to +/// decode. The error is a stable, language-neutral token (`"hex"`), not a +/// translatable English message. +fn decode_hex(s: &str) -> std::result::Result<Vec<u8>, &'static str> { + let bytes = s.as_bytes(); + if bytes.len() % 2 != 0 { + return Err("hex"); + } + let mut out = Vec::with_capacity(bytes.len() / 2); + for pair in bytes.chunks_exact(2) { + let hi = (pair[0] as char).to_digit(16).ok_or("hex")?; + let lo = (pair[1] as char).to_digit(16).ok_or("hex")?; + out.push((hi * 16 + lo) as u8); + } + Ok(out) +} + fn parse_hex4(s: &str) -> Result<[u8; 4]> { - if s.len() != 8 { - return Err(Error::ProfileParse); - } - let mut out = [0u8; 4]; - for i in 0..4 { - out[i] = u8::from_str_radix(&s[i * 2..i * 2 + 2], 16).map_err(|_| Error::ProfileParse)?; - } + let bytes = decode_hex(s).map_err(|_| Error::ProfileParse)?; + let out: [u8; 4] = bytes.try_into().map_err(|_| Error::ProfileParse)?; Ok(out) } @@ -141,15 +172,7 @@ where // An empty string / null / missing field decodes as `None`. fn parse_hex_bytes(s: &str) -> std::result::Result<Vec<u8>, &'static str> { - if s.len() % 2 != 0 { - return Err("odd hex length"); - } - let mut out = Vec::with_capacity(s.len() / 2); - for i in (0..s.len()).step_by(2) { - let byte = u8::from_str_radix(&s[i..i + 2], 16).map_err(|_| "invalid hex digit")?; - out.push(byte); - } - Ok(out) + decode_hex(s) } fn deserialize_opt_hex_bytes_10<'de, D>( @@ -164,11 +187,9 @@ where return Ok(None); } let bytes = parse_hex_bytes(&s).map_err(serde::de::Error::custom)?; - if bytes.len() != 10 { - return Err(serde::de::Error::custom("expected 10 bytes")); - } - let mut out = [0u8; 10]; - out.copy_from_slice(&bytes); + let out: [u8; 10] = bytes + .try_into() + .map_err(|_| serde::de::Error::custom("len"))?; Ok(Some(out)) } @@ -184,11 +205,9 @@ where return Ok(None); } let bytes = parse_hex_bytes(&s).map_err(serde::de::Error::custom)?; - if bytes.len() != 12 { - return Err(serde::de::Error::custom("expected 12 bytes")); - } - let mut out = [0u8; 12]; - out.copy_from_slice(&bytes); + let out: [u8; 12] = bytes + .try_into() + .map_err(|_| serde::de::Error::custom("len"))?; Ok(Some(out)) } @@ -211,15 +230,51 @@ where const BUNDLED_PROFILES: &str = include_str!("../profiles.json"); +/// Parse the bundled profiles fresh into an owned [`ProfilesFile`]. +/// +/// Re-parses the embedded JSON (~800 KB) on every call; prefer +/// [`bundled`] for the hot path, which parses once and caches. This +/// owned form is kept for callers that need a mutable / independent copy. pub fn load_bundled() -> Result<ProfilesFile> { load_from_str(BUNDLED_PROFILES) } +/// Borrow the process-wide cached bundled profiles, parsing once on first +/// use. Avoids re-parsing the ~800 KB JSON on every `Drive::open()`. +/// +/// Returns `None` if the embedded JSON fails to parse (a build-time bug — +/// the bundled blob is fixed at compile time, so the first successful call +/// guarantees all later calls succeed too). +pub fn bundled() -> Option<&'static ProfilesFile> { + use std::sync::OnceLock; + static CACHE: OnceLock<Option<ProfilesFile>> = OnceLock::new(); + CACHE + .get_or_init(|| load_from_str(BUNDLED_PROFILES).ok()) + .as_ref() +} + +/// Find a profile for a drive against the cached bundled profiles. +/// +/// Convenience wrapper over [`bundled`] + [`find_by_drive_id`] that skips +/// the per-call re-parse. Returns `None` if no profile matches (or, in the +/// build-bug case, if the bundled JSON failed to parse). +pub fn find_bundled(drive_id: &crate::identity::DriveId) -> Option<ProfileMatch> { + find_by_drive_id(bundled()?, drive_id) +} + fn load_from_str(data: &str) -> Result<ProfilesFile> { serde_json::from_str(data).map_err(|_| Error::ProfileParse) } /// Find a profile matching a drive's INQUIRY fields. +/// +/// Two-pass per platform section (MT1959-A, then MT1959-B, then Renesas): +/// first an exact match including `firmware_date`, then — if none — a +/// looser match on vendor / revision / vendor-specific only. The exact +/// pass wins so a drive with a known firmware date binds to its precise +/// profile; the looser pass lets a drive whose firmware date we don't have +/// on file still match a same-model profile. All comparisons are +/// whitespace-trimmed. Returns the first section that yields a match. pub fn find_by_drive_id( profiles: &ProfilesFile, drive_id: &crate::identity::DriveId, @@ -290,4 +345,42 @@ mod tests { let id = make_drive_id("FAKE-VND", "9.99", "XX12345", ""); assert!(find_by_drive_id(&profiles, &id).is_none()); } + + #[test] + fn decode_hex_rejects_non_ascii_without_panic() { + // A multi-byte char of even byte-length must not slice inside a + // char boundary; it must decode-fail gracefully. + assert!(decode_hex("中中").is_err()); // 6 bytes, none hex + assert!(parse_hex4("中中").is_err()); // 6 bytes != 8 anyway + // An 8-byte non-ASCII string (two 4-byte chars) hits the exact-len + // path of parse_hex4; must still error, not panic. + assert!(parse_hex4("𝕏𝕏").is_err()); + } + + #[test] + fn decode_hex_roundtrips_valid_hex() { + assert_eq!(decode_hex("00ff10").unwrap(), vec![0x00, 0xff, 0x10]); + assert_eq!(parse_hex4("deadbeef").unwrap(), [0xde, 0xad, 0xbe, 0xef]); + assert!(decode_hex("abc").is_err()); // odd length + assert!(decode_hex("zz").is_err()); // non-hex + } + + #[test] + fn bundled_is_cached_and_matches_fresh_parse() { + let cached = bundled().expect("bundled profiles parse"); + let fresh = load_bundled().unwrap(); + // Same data either way (compare section sizes — ProfilesFile isn't Eq). + assert_eq!(cached.mt1959_a.len(), fresh.mt1959_a.len()); + // Cached accessor returns a stable address across calls. + let a = bundled().unwrap() as *const ProfilesFile; + let b = bundled().unwrap() as *const ProfilesFile; + assert_eq!(a, b); + } + + #[test] + fn find_bundled_matches_known_drive() { + let id = make_drive_id("HL-DT-ST", "1.03", "NM00000", "211810241934"); + let m = find_bundled(&id).unwrap(); + assert_eq!(m.platform, Platform::Mt1959A); + } } diff --git a/src/progress.rs b/src/progress.rs index 586cd9c..665ebb5 100644 --- a/src/progress.rs +++ b/src/progress.rs @@ -1,10 +1,10 @@ //! Pipeline-progress reporting for the rip pipeline. //! -//! v0.13.16 architecture rule: ONE progress signal type. Every long-running -//! pipeline operation (`Disc::copy`, `Disc::patch`, `verify_title`) emits the same -//! `PassProgress` shape via the `Progress` trait. Consumers (autorip) compute -//! a single `PipelineStats` derived view and never reach into per-pass -//! internals. +//! Architecture rule: ONE progress signal type. Every long-running +//! pipeline operation (`Disc::copy`, `Disc::patch`, `verify_title`) emits the +//! same [`PassProgress`] shape via the [`Progress`] trait. Consumers (autorip) +//! compute their own single derived view from these fields and never reach +//! into per-pass internals. //! //! Why this matters: pre-0.13.16 the API leaked `pos`, `bytes_good`, //! `work_done`, `bytes_pending`, `Finished/NonTrimmed` mapfile semantics — @@ -26,7 +26,10 @@ pub enum PassKind { /// `Disc::patch` final pass at 1 sector per block. Scrape { reverse: bool }, /// Demux ISO → output (MKV / M2TS / network). Single phase that runs - /// after all rip passes complete. + /// after all rip passes complete. The library's mux pipeline does not + /// currently emit `PassProgress` itself, so this variant exists for + /// consumers (e.g. autorip) that label their own mux phase with the + /// same `PassKind` vocabulary. Mux, /// Sector verification — reads every sector and classifies health. Verify, @@ -70,21 +73,24 @@ impl PassProgress { /// Percentage of work completed for this pass (0..=100). /// /// Returns `100.0` if `work_total` is zero to avoid division by zero. + /// Clamped to `0..=100` so a transient `work_done > work_total` + /// (e.g. a count that briefly overshoots) never reports above 100%. pub fn work_pct(&self) -> f64 { if self.work_total == 0 { return 100.0; } - self.work_done as f64 / self.work_total as f64 * 100.0 + (self.work_done as f64 / self.work_total as f64 * 100.0).clamp(0.0, 100.0) } /// Percentage of the disc that is confirmed clean (0..=100). /// - /// Computed from `bytes_good_total / bytes_total_disc`. + /// Computed from `bytes_good_total / bytes_total_disc`, clamped to + /// `0..=100`. pub fn good_pct(&self) -> f64 { if self.bytes_total_disc == 0 { return 100.0; } - self.bytes_good_total as f64 / self.bytes_total_disc as f64 * 100.0 + (self.bytes_good_total as f64 / self.bytes_total_disc as f64 * 100.0).clamp(0.0, 100.0) } /// Percentage of the disc that is unreadable (0..=100). @@ -92,7 +98,8 @@ impl PassProgress { if self.bytes_total_disc == 0 { return 0.0; } - self.bytes_unreadable_total as f64 / self.bytes_total_disc as f64 * 100.0 + (self.bytes_unreadable_total as f64 / self.bytes_total_disc as f64 * 100.0) + .clamp(0.0, 100.0) } /// Percentage of the disc that is still pending (not yet attempted or needs retry). @@ -100,7 +107,7 @@ impl PassProgress { if self.bytes_total_disc == 0 { return 0.0; } - self.bytes_pending_total as f64 / self.bytes_total_disc as f64 * 100.0 + (self.bytes_pending_total as f64 / self.bytes_total_disc as f64 * 100.0).clamp(0.0, 100.0) } } diff --git a/src/scsi/linux.rs b/src/scsi/linux.rs index 772f80e..c81e9a2 100644 --- a/src/scsi/linux.rs +++ b/src/scsi/linux.rs @@ -10,9 +10,7 @@ //! //! This matches what every reference project does: MakeMKV (8 s sync //! ioctl), sg_dd (60 s sync ioctl), the kernel default for SCSI block -//! devices (30 s `/sys/.../timeout`). See -//! the SCSI architecture audit (2026-04-26) for the full primary-source -//! references. +//! devices (30 s `/sys/.../timeout`). //! //! Pre-0.13.20 we ran an async `write() + poll(1.5s) + close-on-timeout + //! bg reopen` pattern. That abandoned slow-but-alive commands faster than @@ -90,23 +88,10 @@ impl SgIoTransport { }) } - /// Clean up kernel SG_IO state and unlock the tray. NOT a hardware - /// reset — purely software cleanup before this process opens the - /// device for real work. - /// - /// When a previous process is killed (SIGKILL) mid-SG_IO, the kernel - /// may hold queued commands against the dead fd, and `Drop` never - /// ran so the tray may still be locked via PREVENT MEDIUM REMOVAL. - /// This routine handles both: open + close flushes the kernel SG - /// queue (sg_release cancels commands tied to the fd), the 2 s sleep - /// gives the kernel time to finish that cleanup, then a fresh fd - /// sends ALLOW MEDIUM REMOVAL to clear any stale tray lock. - /// - /// We do NOT verify the drive with TUR or escalate to SG_SCSI_RESET / - /// STOP+START UNIT. Both escalations were tried in 0.13.0–0.13.5 - /// against the LG BU40N (Initio USB-SATA bridge); both failed to - /// recover wedged drives and made the wedge worse — see - /// the BU40N wedge recovery postmortem (2026-04-25). + /// Map the current `errno` (from a failed `libc::open`) to a typed + /// [`Error`]: `EACCES`/`EPERM` → [`Error::DevicePermission`], anything + /// else → [`Error::DeviceNotFound`]. The device path is carried in the + /// error; no English commentary (the app layer localizes). fn open_error<T>(device: &Path) -> Result<T> { let err = std::io::Error::last_os_error(); Err(if err.kind() == std::io::ErrorKind::PermissionDenied { @@ -186,6 +171,17 @@ impl Drop for SgIoTransport { let _ = Self::raw_command(self.fd, &[0x1E, 0, 0, 0, 0, 0], 3_000); unsafe { libc::close(self.fd) }; } + // A failed execute() spawns a detached thread that opens a fresh + // fd into fd_recovery; that slot is normally drained at the top of + // the next execute(). If the transport is dropped before another + // execute() runs (the common abort-on-wedge path), the recovered + // fd would otherwise leak. Claim and close it here. + let recovered = self + .fd_recovery + .swap(-1, std::sync::atomic::Ordering::Acquire); + if recovered >= 0 { + unsafe { libc::close(recovered) }; + } } } @@ -224,6 +220,17 @@ impl ScsiTransport for SgIoTransport { data: &mut [u8], timeout_ms: u32, ) -> Result<ScsiResult> { + // Guard the entry point: `ScsiTransport` is a pub trait, so an + // external caller could pass an empty CDB. Indexing cdb[0] below + // (and in the error paths) would panic. In-crate callers always + // pass non-empty literal CDBs. + if cdb.is_empty() { + return Err(Error::ScsiError { + opcode: 0, + status: super::SCSI_STATUS_TRANSPORT_FAILURE, + sense: None, + }); + } let exec_t0 = std::time::Instant::now(); let opcode = cdb[0]; tracing::trace!( @@ -342,14 +349,36 @@ impl ScsiTransport for SgIoTransport { }); std::thread::spawn(move || { - let c_path = std::ffi::CString::new(path.as_os_str().as_encoded_bytes()).unwrap(); + // Don't unwrap: a device path with an interior NUL would + // panic this detached thread (silently swallowed). Bail + // and leave fd_recovery untouched instead. + let c_path = match std::ffi::CString::new(path.as_os_str().as_encoded_bytes()) { + Ok(c) => c, + Err(_) => return, + }; let new_fd = unsafe { libc::open( c_path.as_ptr() as *const libc::c_char, libc::O_RDWR | libc::O_NONBLOCK | libc::O_CLOEXEC, ) }; - recovery.store(new_fd, std::sync::atomic::Ordering::Release); + if new_fd < 0 { + return; + } + // Publish only into an empty (-1) slot. If two recovery + // threads race, the loser closes its own fd rather than + // overwriting (and leaking) the winner's. + if recovery + .compare_exchange( + -1, + new_fd, + std::sync::atomic::Ordering::Release, + std::sync::atomic::Ordering::Relaxed, + ) + .is_err() + { + unsafe { libc::close(new_fd) }; + } }); return Err(Error::ScsiError { @@ -383,7 +412,12 @@ impl ScsiTransport for SgIoTransport { }); } - let bytes_transferred = (data.len() as i32).saturating_sub(hdr.resid).max(0) as usize; + // Compute in usize so transfers in the 2–4 GiB range (permitted by + // the `> u32::MAX` guard above) don't wrap through an i32 cast and + // report a large successful read as ~0 bytes. A negative resid is + // clamped to 0 before subtracting. + let resid = hdr.resid.max(0) as usize; + let bytes_transferred = data.len().saturating_sub(resid); tracing::trace!( target: "freemkv::scsi", phase = "ok", @@ -504,10 +538,15 @@ fn enumerate_sg_names() -> Vec<String> { continue; } let type_path = format!("/sys/class/scsi_generic/{name}/device/type"); + // By design: only type-5 (optical) sg nodes are collected. + // A non-optical `type` value, or an unreadable `type` file + // (race against device teardown, restricted sysfs in a minimal + // container), is silently skipped — neither is a fatal + // enumeration error, the node simply is not an optical target. match std::fs::read_to_string(&type_path) { Ok(s) if s.trim() == SCSI_TYPE_OPTICAL => names.push(name), - Ok(_) => {} // not optical - Err(_) => {} + Ok(_) => {} // not optical + Err(_) => {} // type file unreadable } } } else { @@ -553,12 +592,14 @@ pub(super) fn drive_has_disc(path: &Path) -> Result<bool> { hdr.flags = SG_FLAG_Q_AT_HEAD; let ret = unsafe { libc::ioctl(fd, SG_IO as _, &mut hdr as *mut sg_io_hdr) }; + // Capture the ioctl errno BEFORE close(): POSIX permits close() to + // set errno (e.g. EIO on a flaky USB path), which would otherwise + // clobber the ioctl failure reason reported below. + let ioctl_err = std::io::Error::last_os_error(); unsafe { libc::close(fd) }; if ret < 0 { - return Err(Error::IoError { - source: std::io::Error::last_os_error(), - }); + return Err(Error::IoError { source: ioctl_err }); } let driver_status_real = hdr.driver_status & !super::DRIVER_SENSE; diff --git a/src/scsi/macos.rs b/src/scsi/macos.rs index e98c36e..61aa1c5 100644 --- a/src/scsi/macos.rs +++ b/src/scsi/macos.rs @@ -17,9 +17,21 @@ use super::{DataDirection, ScsiResult, ScsiTransport}; use crate::error::{Error, Result}; use std::path::Path; +use std::sync::atomic::{AtomicBool, Ordering}; const K_SENSE_DATA_SIZE: usize = 32; +/// Max CDB length the SCSI commands this library issues ever use; also +/// the clamp Linux applies. Used to bound the `cdb_len` passed to the +/// shim so a pathological >255-byte slice can't wrap a `u8`. +const K_MAX_CDB_SIZE: usize = 16; + +/// The C shim uses a single global IOKit handle (`g_handle`), so only one +/// [`MacScsiTransport`] may exist at a time — a second `open()` would +/// share that handle and the first `drop()` would tear it down out from +/// under the other. This flag enforces single-instance ownership. +static OPEN: AtomicBool = AtomicBool::new(false); + #[repr(C)] #[derive(Copy, Clone)] struct ShimDriveInfo { @@ -66,13 +78,38 @@ impl MacScsiTransport { dev_str }; + // Enforce single-instance: the shim's global handle can't back two + // live transports safely. Bail rather than corrupt shared state. + if OPEN.swap(true, Ordering::Acquire) { + return Err(Error::DeviceLocked { + path: bsd_name.to_string(), + kr: 0, + }); + } + let mut bsd_c = bsd_name.as_bytes().to_vec(); bsd_c.push(0); let rc = unsafe { shim_open_exclusive(bsd_c.as_ptr()) }; if rc != 0 { - return Err(Error::DeviceNotFound { - path: bsd_name.to_string(), + // Release the single-instance lock taken by the OPEN.swap above; + // a failed open must not leave it held or every later open wedges. + OPEN.store(false, Ordering::Release); + let path = bsd_name.to_string(); + // The shim returns distinct negative sentinels per failure + // stage; map them to the typed variants that already exist + // rather than collapsing every failure to DeviceNotFound. + // These sentinels are not IOReturn codes, so kr is left 0. + return Err(match rc { + // -2/-3/-4: IOCreatePlugInInterfaceForService / + // QueryInterface MMCDeviceInterface / + // GetSCSITaskDeviceInterface failed. + -4..=-2 => Error::IoKitPluginFailed { path, kr: 0 }, + // -5: ObtainExclusiveAccess failed (held by another + // process). + -5 => Error::DeviceLocked { path, kr: 0 }, + // -1 and anything else: device not present. + _ => Error::DeviceNotFound { path }, }); } @@ -85,6 +122,7 @@ impl MacScsiTransport { impl Drop for MacScsiTransport { fn drop(&mut self) { unsafe { shim_close() }; + OPEN.store(false, Ordering::Release); } } @@ -94,8 +132,25 @@ impl ScsiTransport for MacScsiTransport { cdb: &[u8], direction: DataDirection, data: &mut [u8], + // NOTE: timeout_ms is currently ignored on macOS. The C shim + // (`macos_shim.c`) hardcodes `SetTimeoutDuration(task, 30000)`, so + // every command uses a fixed 30 s budget regardless of the + // caller's READ_TIMEOUT_MS / READ_RECOVERY_TIMEOUT_MS / TUR value. + // Plumbing it through the shim signature is tracked separately; + // macOS is dev/test-only per the project rules. _timeout_ms: u32, ) -> Result<ScsiResult> { + // Match the Linux guard: a >=4 GiB buffer would wrap when cast to + // u32 for the shim, producing a short transfer reported as success + // with the wrong byte count. + if data.len() > u32::MAX as usize { + return Err(Error::ScsiError { + opcode: cdb.first().copied().unwrap_or(0), + status: super::SCSI_STATUS_TRANSPORT_FAILURE, + sense: None, + }); + } + let data_in = match direction { DataDirection::FromDevice => 1, DataDirection::ToDevice => 0, @@ -106,10 +161,11 @@ impl ScsiTransport for MacScsiTransport { let mut task_status: u8 = 0xFF; let mut transfer_count: u64 = 0; + let cdb_len = cdb.len().min(K_MAX_CDB_SIZE) as u8; let kr = unsafe { shim_execute( cdb.as_ptr(), - cdb.len() as u8, + cdb_len, data.as_mut_ptr(), data.len() as u32, data_in, @@ -122,7 +178,7 @@ impl ScsiTransport for MacScsiTransport { if kr != 0 { return Err(Error::ScsiError { - opcode: cdb[0], + opcode: cdb.first().copied().unwrap_or(0), status: super::SCSI_STATUS_TRANSPORT_FAILURE, sense: None, }); @@ -131,7 +187,7 @@ impl ScsiTransport for MacScsiTransport { if task_status != 0 { let parsed = super::parse_sense(&sense, K_SENSE_DATA_SIZE as u8); return Err(Error::ScsiError { - opcode: cdb[0], + opcode: cdb.first().copied().unwrap_or(0), status: task_status, sense: Some(parsed), }); @@ -139,7 +195,11 @@ impl ScsiTransport for MacScsiTransport { Ok(ScsiResult { status: 0, - bytes_transferred: transfer_count as usize, + // Clamp to the buffer length, matching the Linux transport's + // structural bound (data.len().saturating_sub(resid)). A lying + // drive/shim can't then produce a bytes_transferred that + // exceeds the buffer a future caller might slice with. + bytes_transferred: (transfer_count as usize).min(data.len()), sense, }) } diff --git a/src/scsi/mod.rs b/src/scsi/mod.rs index 769fafd..341d670 100644 --- a/src/scsi/mod.rs +++ b/src/scsi/mod.rs @@ -56,8 +56,7 @@ pub(crate) const TUR_TIMEOUT_MS: u32 = 5_000; /// /// 10 s catches every legitimate slow read with comfortable margin and /// short-circuits truly bad sectors at ~10 s rather than letting the -/// kernel mid-layer escalate for 30 s+. See the SCSI architecture audit -/// (2026-04-26) for primary-source references. +/// kernel mid-layer escalate for 30 s+. /// /// Pre-0.13.21 this was 1.5 s, which forced the kernel mid-layer to /// time out *normal* reads (cold-start often takes ~1.5 s) and run its @@ -245,9 +244,9 @@ impl ScsiSense { /// `sb_len_wr` is the number of bytes the transport actually wrote into /// `sense`. When the buffer is too short for the relevant fields we /// return [`ScsiSense::NONE`] for the missing pieces rather than reading -/// uninitialised memory. The minimum useful sense reply per SPC-4 is 8 -/// bytes (descriptor) or 14 bytes (fixed, to reach ASC/ASCQ at offsets -/// 12/13). +/// uninitialised memory. The minimum useful sense reply is 4 bytes +/// (descriptor, to reach ASCQ at offset 3) or 14 bytes (fixed, to reach +/// ASC/ASCQ at offsets 12/13). /// /// Pure function — same parse on every platform backend (Linux SG_IO, /// macOS IOKit, Windows SPTI) so a regression here would silently @@ -261,7 +260,9 @@ pub(crate) fn parse_sense(sense: &[u8], sb_len_wr: u8) -> ScsiSense { let descriptor = response_code == 0x72 || response_code == 0x73; if descriptor { // Descriptor format: key/asc/ascq are at fixed offsets 1/2/3. - let asc = if n >= 3 { sense[2] } else { 0 }; + // n >= 3 is guaranteed by the early return above, so byte 2 is + // always in bounds; only ascq (byte 3) needs a length check. + let asc = sense[2]; let ascq = if n >= 4 { sense[3] } else { 0 }; ScsiSense { sense_key: sense[1] & 0x0F, @@ -449,13 +450,15 @@ pub fn list_drives() -> Vec<DriveInfo> { /// other ready/not-ready response → `Ok(true)` or interpreted ready /// state. Suitable for poll-loop tick (~50 ms / drive on a healthy bus). /// -/// **Internal wedge recovery.** When the kernel's response indicates a -/// wedged target — the `0xff` status pattern that means "no answer from -/// the device" — this function transparently escalates: SCSI bus reset -/// → if still wedged → USB device reset (`USBDEVFS_RESET` on Linux) → -/// retry TUR. Callers never see wedge errors and never need to know -/// about the escalation; if even the recovery path can't get a response, -/// `Err(DeviceResetFailed)` surfaces. **No SCSI primitive is exposed to +/// **No internal recovery.** A single TUR is issued; nothing else. When +/// the transport reports a wedged target (the `0xff` "no answer from the +/// device" pattern synthesised by the backend from a non-zero +/// `host_status` / `driver_status`), that failure surfaces directly to +/// the caller as `Err(Error::ScsiError)` with +/// `status == SCSI_STATUS_TRANSPORT_FAILURE (0xFF)` and `sense: None`. No +/// SCSI bus reset, no USB device reset, no retry is attempted in-library +/// (the USB-reset escalation was removed in 0.13.4 after it was shown to +/// deepen rather than clear the wedge). **No SCSI primitive is exposed to /// outside crates** — autorip / freemkv CLI / bdemu use this single /// function for the entire "is there a disc?" decision. pub fn drive_has_disc(path: &Path) -> Result<bool> { @@ -561,8 +564,11 @@ pub fn build_set_cd_speed(read_speed: u16) -> [u8; 12] { ] } -/// Build a READ(10) CDB with the raw read flag. -pub fn build_read10_raw(lba: u32, count: u16) -> [u8; 10] { +/// Build a READ(10) CDB with Force Unit Access (FUA) set — byte 1 bit 3 +/// (0x08). FUA bypasses the drive cache and reads directly from the +/// medium. (Note: this is *not* a "raw" read; raw optical reads require +/// READ CD, opcode 0xBE.) +pub fn build_read10_fua(lba: u32, count: u16) -> [u8; 10] { [ SCSI_READ_10, 0x08, @@ -579,7 +585,7 @@ pub fn build_read10_raw(lba: u32, count: u16) -> [u8; 10] { #[cfg(test)] mod parse_sense_tests { - //! Unit tests for `parse_sense_key`. Covers both SPC-4 sense data + //! Unit tests for [`parse_sense`]. Covers both SPC-4 sense data //! formats (descriptor / fixed) and the short-buffer fallback. The //! same helper runs on every platform backend so a regression here //! would silently miscategorize SCSI errors on Linux, macOS, and diff --git a/src/scsi/windows.rs b/src/scsi/windows.rs index a7be236..4cab95f 100644 --- a/src/scsi/windows.rs +++ b/src/scsi/windows.rs @@ -87,9 +87,9 @@ pub struct SptiTransport { handle: isize, } -// SptiTransport's only field is the isize HANDLE — Send is auto-derived -// and intentional. Sync is NOT: handle mutation in execute() requires -// &mut, enforced by the trait object dispatch. +// SptiTransport's only field is an isize HANDLE, so the compiler +// auto-derives BOTH Send and Sync. Exclusive use of the raw handle is +// enforced by `&mut self` on `execute()`, not by any absence of Sync. /// Normalize a device path to Windows \\.\X: format. /// @@ -307,7 +307,7 @@ impl ScsiTransport for SptiTransport { // is at best redundant and at worst deepens the wedge. Caller // surfaces the failure to UX. return Err(Error::ScsiError { - opcode: cdb[0], + opcode: cdb.first().copied().unwrap_or(0), status: super::SCSI_STATUS_TRANSPORT_FAILURE, sense: None, }); @@ -324,7 +324,7 @@ impl ScsiTransport for SptiTransport { // `ScsiSense::is_medium_error()` etc. let parsed = super::parse_sense(&sptwb.sense, K_SENSE_SIZE as u8); return Err(Error::ScsiError { - opcode: cdb[0], + opcode: cdb.first().copied().unwrap_or(0), status: sptwb.spt.ScsiStatus, sense: Some(parsed), }); diff --git a/src/sector/decrypting.rs b/src/sector/decrypting.rs index 1d086dd..0ca5518 100644 --- a/src/sector/decrypting.rs +++ b/src/sector/decrypting.rs @@ -1,12 +1,12 @@ //! `DecryptingSectorSource` — wrap any [`SectorSource`] to apply //! AACS / CSS in-place decryption on every read. //! -//! This is the 0.18 single-source-of-truth for decrypt-on-read. The -//! actual cipher code lives in [`crate::aacs`] and [`crate::css`]; -//! we just call the existing [`crate::decrypt::decrypt_sectors`] -//! helper that already drives both of them. In follow-up commits -//! `sweep_pipeline` and `DiscStream` migrate onto this decorator -//! and delete their duplicate decrypt call sites. +//! This is the single source of truth for decrypt-on-read: every +//! decrypt-on-read caller (e.g. `DiscStream`) wraps its source in this +//! decorator. The actual cipher code lives in [`crate::aacs`] and +//! [`crate::css`]; we just call the existing +//! [`crate::decrypt::decrypt_sectors`] helper that drives both of them +//! in-place after each read (a no-op for [`DecryptKeys::None`]). //! //! Composition: `Drive` → `DecryptingSectorSource` → caller sees //! plaintext. For `DecryptKeys::None` discs the decorator is a @@ -94,10 +94,8 @@ impl<S: SectorSource> SectorSource for DecryptingSectorSource<S> { recovery: bool, ) -> Result<usize> { let n = self.inner.read_sectors(lba, count, buf, recovery)?; - // Reuse the existing crate-wide decrypt entry point — same - // path the 0.17 sweep_pipeline and DiscStream call, so we - // inherit their AACS / CSS / None semantics verbatim. The - // helper is a no-op for DecryptKeys::None. + // Apply the crate-wide AACS/CSS/None decrypt entry point in-place + // over the bytes just read. No-op for DecryptKeys::None. decrypt_sectors(&mut buf[..n], &self.keys, self.unit_key_idx)?; Ok(n) } diff --git a/src/sector/file.rs b/src/sector/file.rs index 485c4f5..d738808 100644 --- a/src/sector/file.rs +++ b/src/sector/file.rs @@ -57,11 +57,17 @@ impl FileSectorSink { impl SectorSink for FileSectorSink { fn write_sectors(&mut self, lba: u32, buf: &[u8]) -> Result<()> { - debug_assert!( - buf.len() % 2048 == 0, - "FileSectorSink::write_sectors: buf len {} not a multiple of 2048", - buf.len() - ); + // SectorSink's contract requires a 2048-multiple buffer. Enforce + // it in all build modes (a `debug_assert!` is a no-op in release): + // a misaligned buffer would `write_all` partial bytes at + // lba*2048 and silently corrupt the ISO. Current in-tree callers + // always pass aligned buffers; this guards the public trait + // contract against any (including future external) caller. + if buf.len() % 2048 != 0 { + return Err(Error::IoError { + source: std::io::Error::from(std::io::ErrorKind::InvalidInput), + }); + } let offset = lba as u64 * 2048; self.inner .seek(SeekFrom::Start(offset)) diff --git a/src/sector/mod.rs b/src/sector/mod.rs index 67df098..dac27bd 100644 --- a/src/sector/mod.rs +++ b/src/sector/mod.rs @@ -8,7 +8,7 @@ //! - [`SectorSource`] is implemented by `Drive` (hardware) and //! [`FileSectorSource`] (file-backed). //! - [`SectorSink`] is implemented by [`FileSectorSink`] -//! (ISO-backed) and sweep/patch consumer adapters. +//! (ISO-backed). //! - [`DecryptingSectorSource`] is a decorator that wraps any //! `SectorSource` and applies AACS / CSS in-place decrypt to //! yield plaintext sectors. @@ -37,6 +37,14 @@ pub trait SectorSource: Send { /// the flag. /// /// Returns the number of bytes written into `buf` on success. + /// + /// # Panics + /// + /// Implementations may panic if `buf.len() < count * 2048`. This + /// is a caller contract enforced via `debug_assert!` in the + /// primary impl ([`FileSectorSource`]); in release builds an + /// undersized buffer panics on the slice. Callers must size `buf` + /// to at least `count * 2048` bytes. fn read_sectors( &mut self, lba: u32, @@ -102,7 +110,8 @@ impl SectorSource for &mut (dyn SectorSource + '_) { pub trait SectorSink: Send { /// Write the sectors in `buf` starting at `lba`. `buf.len()` /// must be a multiple of 2048; the implementation seeks to - /// `lba * 2048` before writing. + /// `lba as u64 * 2048` before writing (the `u64` cast is required — + /// a bare `u32` `lba * 2048` wraps past ~4 GB on UHD-scale images). fn write_sectors(&mut self, lba: u32, buf: &[u8]) -> Result<()>; /// Flush, fsync, and close. Consumes the sink. Always called diff --git a/src/sector/prefetched.rs b/src/sector/prefetched.rs index ea0c1d0..6120e31 100644 --- a/src/sector/prefetched.rs +++ b/src/sector/prefetched.rs @@ -18,7 +18,8 @@ //! //! ## Lifecycle //! -//! The producer thread is spawned by [`PrefetchedSectorSource::new`]. +//! The producer thread is spawned by [`PrefetchedSectorSource::new`] +//! (which returns `Err` if the OS refuses the thread spawn). //! It walks the supplied extent list in order, reads the configured //! batch size at each LBA, and sends the resulting plaintext buffer //! into a [`crossbeam_channel::bounded`] channel of small depth (so @@ -78,10 +79,11 @@ pub struct PrefetchedSectorSource { recycle_tx: Sender<Vec<u8>>, /// Joined on drop so producer cleanup runs deterministically. producer: Option<JoinHandle<()>>, - /// Cumulative bytes drained by `read_sectors` calls. Exposed via - /// [`capacity_sectors`] indirectly: the consumer-side state needs - /// this to advance its position bookkeeping in lockstep with what - /// the producer fed. + /// Total sector count across all extents, computed once at + /// construction (the sum of each extent's `sector_count`) and + /// returned by [`capacity_sectors`]. Never updated by reads. + /// + /// [`capacity_sectors`]: SectorSource::capacity_sectors total_sectors: u32, } @@ -91,12 +93,24 @@ impl PrefetchedSectorSource { /// [`DecryptingSectorSource`](crate::sector::DecryptingSectorSource)) /// — every byte the producer emits is what the consumer's demux /// will feed to its codec parsers. + /// + /// ## Unit-alignment precondition + /// + /// Each extent's `sector_count` should be a multiple of + /// [`SECTOR_ALIGNMENT`] (3 sectors / one 6144-byte AACS aligned + /// unit). Blu-ray m2ts extents satisfy this by spec. If an extent + /// has a trailing 1-2 sectors that cannot fill a complete unit, + /// the producer surfaces [`Error::ExtentNotUnitAligned`] through + /// the channel rather than handing the decrypt step a sub-unit + /// chunk it would silently leave encrypted. + /// + /// [`Error::ExtentNotUnitAligned`]: crate::error::Error::ExtentNotUnitAligned pub fn new<S>( reader: S, extents: Vec<crate::disc::Extent>, batch_sectors: u16, halt: Option<Halt>, - ) -> Self + ) -> Result<Self> where S: SectorSource + Send + 'static, { @@ -113,11 +127,30 @@ impl PrefetchedSectorSource { batch_sectors: u16, halt: Option<Halt>, event_fn: Option<EventFn>, - ) -> Self + ) -> Result<Self> where S: SectorSource + Send + 'static, { - let total_sectors: u32 = extents.iter().map(|e| e.sector_count).sum(); + // A zero batch would make the producer loop forever emitting + // empty batches (sectors = remaining.min(0) = 0, offset never + // advances). All production callers pass a nonzero constant; a + // public-API caller passing 0 is a programming error, so reject + // it rather than spin a thread that never makes progress. + if batch_sectors == 0 { + return Err(crate::error::Error::IoError { + source: std::io::Error::from(std::io::ErrorKind::InvalidInput), + }); + } + // Accumulate in u64 then clamp: extents can derive from + // untrusted nav/MPLS/UDF data, so a naive u32 `sum()` could + // panic in debug / wrap in release on a hostile total. The + // clamp only affects the advisory `capacity_sectors` figure; + // the producer walks each extent independently below. + let total_sectors: u32 = extents + .iter() + .map(|e| e.sector_count as u64) + .sum::<u64>() + .min(u32::MAX as u64) as u32; let bytes_total_extents: u64 = extents.iter().map(|e| e.sector_count as u64 * 2048).sum(); let (tx, rx) = bounded::<Batch>(PREFETCH_CHANNEL_DEPTH); let (recycle_tx, recycle_rx) = bounded::<Vec<u8>>(PREFETCH_CHANNEL_DEPTH + 1); @@ -148,9 +181,39 @@ impl PrefetchedSectorSource { offset = 0; continue; } + // The AACS aligned unit is SECTOR_ALIGNMENT (3) + // sectors / 6144 bytes; the decrypt step only + // processes full units and silently leaves a + // shorter trailing chunk encrypted. So a batch must + // be a whole number of units — except for the final + // batch of an extent whose `sector_count` is itself + // unit-aligned (then the remaining tail is exactly + // 0 mod 3 and forms full units on its own). + // + // If the trailing sectors of this extent cannot fill + // a complete unit (`remaining < SECTOR_ALIGNMENT` + // with nothing more to read, or a 1-2 sector + // leftover after the last full unit), there is no + // way to hand the decrypt step an aligned chunk — + // surface a typed error instead of emitting + // still-encrypted bytes. + if remaining % SECTOR_ALIGNMENT as u32 != 0 + && remaining < SECTOR_ALIGNMENT as u32 + { + let _ = tx.send(Err(crate::error::Error::ExtentNotUnitAligned.into())); + return; + } let mut sectors = remaining.min(batch_sectors as u32) as u16; + // Trim to a whole number of units. Once trimmed to 0 + // here it means `remaining >= SECTOR_ALIGNMENT` but + // the *batch window* landed on a sub-unit boundary — + // never the trailing-tail case, which the guard + // above already rejected. Clamp to one unit so we + // always make forward progress. if sectors >= SECTOR_ALIGNMENT { sectors -= sectors % SECTOR_ALIGNMENT; + } else { + sectors = SECTOR_ALIGNMENT; } let bytes = sectors as usize * 2048; let mut buf = match recycle_rx.recv() { @@ -159,17 +222,37 @@ impl PrefetchedSectorSource { }; if bytes <= buf.capacity() { // Re-expose `bytes` without zero-filling pages that - // `read_sectors` is about to overwrite. The capacity - // guard makes the `set_len` provably sound even if a - // recycled buffer ever comes back smaller than the + // `read_sectors` is about to overwrite. The enclosing + // capacity guard makes the `set_len` provably sound even + // if a recycled buffer ever comes back smaller than the // `vec![0u8; batch_bytes]` it was born with. + debug_assert!(bytes <= buf.capacity(), "set_len exceeds capacity"); unsafe { buf.set_len(bytes) }; } else { buf.resize(bytes, 0); } - let lba = extent.start_lba + offset; + // `start_lba + offset` derives from untrusted extent + // data — saturate rather than wrap/panic on a + // hostile start_lba near u32::MAX. + let lba = extent.start_lba.saturating_add(offset); match reader.read_sectors(lba, sectors, &mut buf[..bytes], false) { Ok(n) => { + // A short read must not silently desync the + // stream: advance the extent cursor by the + // sectors actually read, not the requested + // count, and reject a byte count that isn't a + // whole number of sectors (it would split a + // sector and leave the decrypt step a partial + // unit). The sole production inner source + // (FileSectorSource) read_exact's the full + // request, so this is belt-and-braces against + // a future short-reading source. + if n % 2048 != 0 { + let _ = + tx.send(Err(crate::error::Error::ExtentNotUnitAligned.into())); + return; + } + let sectors_read = (n / 2048) as u32; buf.truncate(n); bytes_read_total = bytes_read_total.saturating_add(n as u64); if let Some(ref f) = event_fn { @@ -183,7 +266,13 @@ impl PrefetchedSectorSource { if tx.send(Ok(buf)).is_err() { return; // consumer dropped } - offset += sectors as u32; + // A genuine zero-byte read with no error would + // otherwise spin this loop forever; treat it + // as end-of-source. + if sectors_read == 0 { + return; + } + offset = offset.saturating_add(sectors_read); } Err(e) => { let _ = tx.send(Err(e.into())); @@ -193,14 +282,14 @@ impl PrefetchedSectorSource { } // Drop tx implicitly — consumer sees RecvError → EOF. }) - .expect("freemkv-prefetch producer spawn failed"); + .map_err(|e| crate::error::Error::IoError { source: e })?; - Self { + Ok(Self { rx, recycle_tx, producer: Some(producer), total_sectors, - } + }) } /// Peel off the receivers for zero-copy pipeline mode. The @@ -215,14 +304,30 @@ impl PrefetchedSectorSource { /// queries; its `SectorSource` impl becomes invalid after this /// call (data has been moved out). pub fn into_channels(self) -> (Receiver<Batch>, Sender<Vec<u8>>, PrefetchShell) { + // MOVE the three fields out cleanly — never clone. Each of + // `rx` and `recycle_tx` ends up with exactly ONE live copy: + // the one in the returned tuple. The pre-1.0.0 implementation + // cloned both and then `mem::forget`-ed `self`, leaking the + // originals so an extra live receiver + sender survived + // forever. That defeated the channel-disconnection shutdown: + // when the demux consumer exited early (halt, or a `tx.send` + // error in `demux_thread`), the producer's `recycle_rx.recv()` + // and `tx.send()` never saw all-peers-dropped, so the producer + // never returned and `PrefetchShell::drop`'s `join()` hung. + // + // `ManuallyDrop` + `ptr::read` reads each field out by value + // and suppresses `self`'s own `Drop` (which would otherwise + // double-`join`), leaving NO extra live endpoint behind. This + // is the panic-free equivalent of the `Option::take` approach. let total = self.total_sectors; - // Drop the SectorSource side; transfer the producer join - // handle to a shell that just waits on Drop. - let mut me = self; - let producer = me.producer.take(); - let rx = me.rx.clone(); - let recycle = me.recycle_tx.clone(); - std::mem::forget(me); + let me = std::mem::ManuallyDrop::new(self); + // SAFETY: `me` is `ManuallyDrop`, so none of these fields will + // be dropped by `me`. Each `ptr::read` performs exactly one + // bitwise move out; every field is read exactly once and never + // touched again, so there are no double-frees and no aliasing. + let producer = unsafe { std::ptr::read(&me.producer) }; + let rx = unsafe { std::ptr::read(&me.rx) }; + let recycle = unsafe { std::ptr::read(&me.recycle_tx) }; (rx, recycle, PrefetchShell { producer, total }) } } @@ -274,8 +379,29 @@ impl SectorSource for PrefetchedSectorSource { // requested count. match self.rx.recv() { Ok(Ok(filled)) => { - let n = filled.len().min(buf.len()); + // Precondition: the caller's buffer must be large + // enough to hold the producer's batch. If it is not we + // would silently drop `filled[buf.len()..]`, desyncing + // the stream. The production zero-copy path never uses + // this method (it consumes the channel directly via + // `into_channels`), so a too-small buffer here is a + // caller bug — surface it instead of corrupting data. + if filled.len() > buf.len() { + return Err(crate::error::Error::IoError { + source: std::io::Error::from(std::io::ErrorKind::InvalidInput), + }); + } + let n = filled.len(); buf[..n].copy_from_slice(&filled[..n]); + // Return the buffer to the recycle pool so the producer + // can re-fill it. Without this the pool (seeded with + // PREFETCH_CHANNEL_DEPTH+1 buffers) drains after that + // many reads and the producer blocks forever on + // `recycle_rx.recv()` while the consumer blocks on the + // next `rx.recv()` — a permanent deadlock. The + // `into_channels` zero-copy path recycles explicitly; + // this direct-read path must do the same. + let _ = self.recycle_tx.send(filled); Ok(n) } Ok(Err(e)) => Err(crate::error::Error::IoError { source: e }), @@ -284,3 +410,315 @@ impl SectorSource for PrefetchedSectorSource { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::disc::Extent; + use crate::error::Result; + use std::sync::mpsc; + use std::time::Duration; + + /// Endless zero-yielding source: every read succeeds, so the + /// producer keeps trying to push batches forward until the forward + /// channel disconnects. Exactly the shape that wedged the pre-1.0.0 + /// `clone + mem::forget` `into_channels`. + struct EndlessZeroSource; + impl SectorSource for EndlessZeroSource { + fn read_sectors( + &mut self, + _lba: u32, + count: u16, + buf: &mut [u8], + _recovery: bool, + ) -> Result<usize> { + let bytes = count as usize * 2048; + buf[..bytes].fill(0); + Ok(bytes) + } + } + + /// Synthetic source that fills `buf` with a per-sector byte + /// pattern derived from the LBA, always satisfying the full + /// request (mirrors `FileSectorSource`'s read_exact contract). + struct PatternSource { + capacity: u32, + } + + impl SectorSource for PatternSource { + fn capacity_sectors(&self) -> u32 { + self.capacity + } + + fn read_sectors( + &mut self, + lba: u32, + count: u16, + buf: &mut [u8], + _recovery: bool, + ) -> Result<usize> { + let bytes = count as usize * 2048; + for s in 0..count as usize { + let base = s * 2048; + let tag = (lba.wrapping_add(s as u32) & 0xff) as u8; + for b in &mut buf[base..base + 2048] { + *b = tag; + } + } + Ok(bytes) + } + } + + /// Source that returns a short read (fewer sectors than + /// requested) on its very first call, then full reads. Used to + /// prove the producer advances by sectors actually read. + struct ShortFirstSource { + capacity: u32, + first: bool, + } + + impl SectorSource for ShortFirstSource { + fn capacity_sectors(&self) -> u32 { + self.capacity + } + + fn read_sectors( + &mut self, + lba: u32, + count: u16, + buf: &mut [u8], + _recovery: bool, + ) -> Result<usize> { + let give = if self.first { + self.first = false; + // Short read: hand back one aligned unit (3 sectors) + // regardless of the larger request. + SECTOR_ALIGNMENT.min(count) + } else { + count + }; + let bytes = give as usize * 2048; + for s in 0..give as usize { + let base = s * 2048; + let tag = (lba.wrapping_add(s as u32) & 0xff) as u8; + for b in &mut buf[base..base + 2048] { + *b = tag; + } + } + Ok(bytes) + } + } + + fn big_extent() -> Vec<Extent> { + // One huge extent so the producer never reaches EOF on its own; + // the only way it can exit is by observing channel disconnection. + vec![Extent { + start_lba: 0, + sector_count: u32::MAX, + }] + } + + /// Run `f` on a helper thread and fail if it does not finish within + /// `secs`. Used to turn a join-deadlock into a test failure instead + /// of a hung CI run. + fn within<F: FnOnce() + Send + 'static>(secs: u64, f: F) { + let (done_tx, done_rx) = bounded::<()>(1); + std::thread::spawn(move || { + f(); + let _ = done_tx.send(()); + }); + assert!( + done_rx + .recv_timeout(std::time::Duration::from_secs(secs)) + .is_ok(), + "operation did not complete within {secs}s (deadlock)" + ); + } + + /// Run `f` on a worker thread; fail (rather than hang) if it does + /// not finish within `timeout`. Guards the deadlock regression so + /// a reintroduced bug fails the suite instead of wedging it. + fn with_watchdog<F>(timeout: Duration, f: F) + where + F: FnOnce() + Send + 'static, + { + let (done_tx, done_rx) = mpsc::channel::<()>(); + let h = std::thread::spawn(move || { + f(); + let _ = done_tx.send(()); + }); + match done_rx.recv_timeout(timeout) { + Ok(()) => { + let _ = h.join(); + } + Err(_) => panic!("watchdog timeout — likely deadlock/hang in prefetch read path"), + } + } + + /// The CRITICAL regression: after `into_channels`, dropping the + /// returned forward receiver + recycle sender must let the producer + /// observe disconnection and exit, so dropping the `PrefetchShell` + /// (which joins the producer) returns promptly. With the old + /// clone+forget the leaked endpoints kept the producer blocked and + /// this join hung forever. + #[test] + fn into_channels_drop_releases_producer() { + within(10, || { + let src = PrefetchedSectorSource::new(EndlessZeroSource, big_extent(), 3, None) + .expect("spawn"); + let (rx, recycle_tx, shell) = src.into_channels(); + // Consumer goes away early (halt / abort analogue): drop both + // channel endpoints without draining to EOF. + drop(rx); + drop(recycle_tx); + // Joining the producer must not hang. + drop(shell); + }); + } + + /// Same property via the halt path: cancel the token, then the + /// producer must exit and the shell join must complete. + /// + /// The producer parks in a BLOCKING `tx.send` on the forward channel + /// and only checks `halt` at the loop top, so `halt.cancel()` cannot + /// interrupt a send that is already blocked on a full channel. To + /// keep the test deterministic under load, we drain the forward + /// receiver on a background thread: every send then makes progress, + /// the producer reaches the loop top, observes the cancelled halt, + /// and exits — so the shell join completes promptly regardless of + /// scheduling. (Channel-disconnection shutdown is covered separately + /// by `into_channels_drop_releases_producer`.) + #[test] + fn halt_releases_producer() { + within(10, || { + let halt = Halt::new(); + let src = + PrefetchedSectorSource::new(EndlessZeroSource, big_extent(), 3, Some(halt.clone())) + .expect("spawn"); + let (rx, recycle_tx, shell) = src.into_channels(); + // Drain the forward channel so the producer's sends always + // make progress and it can reach the halt check at the loop + // top, recycling buffers so it never blocks on the pool. + let drainer = std::thread::spawn(move || { + while let Ok(item) = rx.recv() { + if let Ok(buf) = item { + let _ = recycle_tx.send(buf); + } + } + }); + halt.cancel(); + drop(shell); + let _ = drainer.join(); + }); + } + + /// `batch_sectors == 0` is rejected rather than spawning a thread + /// that spins forever emitting empty batches. + #[test] + fn zero_batch_rejected() { + let err = PrefetchedSectorSource::new(EndlessZeroSource, big_extent(), 0, None); + assert!(err.is_err(), "zero batch_sectors must be rejected"); + } + + /// >3 sequential direct `read_sectors` calls must succeed. The + /// recycle pool seeds PREFETCH_CHANNEL_DEPTH+1 (3) buffers; before + /// the fix the direct path dropped each drained buffer, so the 4th + /// call deadlocked. Watchdog-guarded. + #[test] + fn direct_reads_past_pool_depth_do_not_deadlock() { + with_watchdog(Duration::from_secs(10), || { + // 24 sectors = 8 aligned units; batch of 3 sectors gives 8 + // sequential batches, well past the 3-buffer pool depth. + let extents = vec![Extent { + start_lba: 0, + sector_count: 24, + }]; + let src = PatternSource { capacity: 24 }; + let mut pf = PrefetchedSectorSource::new(src, extents, 3, None).expect("spawn"); + + let mut buf = vec![0u8; 3 * 2048]; + let mut total = 0usize; + for _ in 0..16 { + let n = pf.read_sectors(0, 3, &mut buf, false).unwrap(); + if n == 0 { + break; // EOF + } + total += n; + } + assert_eq!(total, 24 * 2048, "all 24 sectors should be drained"); + }); + } + + /// An extent whose sector_count is not a multiple of 3 must not + /// emit a still-encrypted sub-unit tail. The producer delivers the + /// readable full units, then surfaces a typed error on the tail + /// instead of a short batch. Watchdog-guarded so a regression that + /// hangs (rather than errors) still fails. + #[test] + fn non_multiple_of_three_extent_errors_on_tail() { + with_watchdog(Duration::from_secs(10), || { + // 8 sectors = 2 full units (6 sectors) + 2 leftover. + let extents = vec![Extent { + start_lba: 100, + sector_count: 8, + }]; + let src = PatternSource { capacity: 200 }; + let mut pf = PrefetchedSectorSource::new(src, extents, 3, None).expect("spawn"); + + let mut buf = vec![0u8; 3 * 2048]; + // First two reads: the 6 unit-aligned sectors come through + // as full 3-sector (6144-byte) batches. + let n0 = pf.read_sectors(0, 3, &mut buf, false).unwrap(); + assert_eq!(n0, 3 * 2048); + let n1 = pf.read_sectors(0, 3, &mut buf, false).unwrap(); + assert_eq!(n1, 3 * 2048); + // Third read hits the 2-sector tail: it must be an error, + // never a 4096-byte (sub-unit) batch that decrypt would + // leave encrypted. + let err = pf.read_sectors(0, 3, &mut buf, false); + assert!( + err.is_err(), + "non-unit-aligned tail must error, got Ok({:?})", + err + ); + }); + } + + /// A short read (inner source returns fewer sectors than + /// requested) must advance the extent cursor by the sectors + /// actually read, not the requested count — otherwise the bytes + /// between the short read and the request size are silently + /// skipped. We verify every sector of the extent is delivered. + #[test] + fn short_read_does_not_desync_stream() { + with_watchdog(Duration::from_secs(10), || { + // 9 sectors = 3 full units. batch of 9 means the first + // request is for 9 sectors; ShortFirstSource hands back + // only 3, so the producer must re-request the remaining 6. + let extents = vec![Extent { + start_lba: 0, + sector_count: 9, + }]; + let src = ShortFirstSource { + capacity: 9, + first: true, + }; + let mut pf = PrefetchedSectorSource::new(src, extents, 9, None).expect("spawn"); + + let mut buf = vec![0u8; 9 * 2048]; + let mut total = 0usize; + for _ in 0..16 { + let n = pf.read_sectors(0, 9, &mut buf, false).unwrap(); + if n == 0 { + break; + } + total += n; + } + assert_eq!( + total, + 9 * 2048, + "short read must not drop sectors; all 9 must be delivered" + ); + }); + } +} diff --git a/src/speed.rs b/src/speed.rs index 6196be8..3489c71 100644 --- a/src/speed.rs +++ b/src/speed.rs @@ -1,7 +1,13 @@ //! Drive speed constants. /// Common optical drive speeds with KB/s values for SET_CD_SPEED. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +/// +/// Ordering is by [`to_kbps`](Self::to_kbps) throughput, not declaration +/// order — `PartialOrd`/`Ord` are implemented manually so e.g. +/// `DVD1x < BD1x` (1385 < 4500 KB/s) holds. A naive derive would have +/// ordered by variant position, making the slow DVD speeds sort above the +/// fast BD speeds. `Max` (0xFFFF) sorts highest, as intended. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum DriveSpeed { BD1x, BD2x, @@ -19,6 +25,8 @@ pub enum DriveSpeed { } impl DriveSpeed { + /// Throughput in KB/s for the SET_CD_SPEED CDB. `Max` maps to the + /// 0xFFFF sentinel that tells the drive to use its maximum speed. pub fn to_kbps(self) -> u16 { match self { DriveSpeed::BD1x => 4_500, @@ -38,8 +46,46 @@ impl DriveSpeed { } } +impl PartialOrd for DriveSpeed { + fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> { + Some(self.cmp(other)) + } +} + +impl Ord for DriveSpeed { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.to_kbps().cmp(&other.to_kbps()) + } +} + impl std::fmt::Display for DriveSpeed { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{:?} ({} KB/s)", self, self.to_kbps()) + // `Max` is the "let the drive pick its maximum" sentinel; printing + // its 0xFFFF KB/s value would read as a real (absurd) throughput. + match self { + DriveSpeed::Max => write!(f, "Max"), + _ => write!(f, "{:?} ({} KB/s)", self, self.to_kbps()), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ordering_is_by_throughput_not_declaration() { + assert!(DriveSpeed::DVD1x < DriveSpeed::BD1x); + assert!(DriveSpeed::DVD16x < DriveSpeed::BD8x); + assert!(DriveSpeed::BD12x < DriveSpeed::Max); + let mut v = [DriveSpeed::Max, DriveSpeed::DVD1x, DriveSpeed::BD4x]; + v.sort(); + assert_eq!(v, [DriveSpeed::DVD1x, DriveSpeed::BD4x, DriveSpeed::Max]); + } + + #[test] + fn max_display_omits_sentinel_value() { + assert_eq!(DriveSpeed::Max.to_string(), "Max"); + assert!(DriveSpeed::BD1x.to_string().contains("4500 KB/s")); } } diff --git a/src/udf.rs b/src/udf.rs index 61364eb..ca9e4b1 100644 --- a/src/udf.rs +++ b/src/udf.rs @@ -21,6 +21,19 @@ use crate::error::{Error, Result}; use crate::sector::SectorSource; +/// Upper bound on a single metadata file read (`read_file`). BD-ROM +/// metadata files (.mpls/.clpi/.inf/.bdmv) are a few KiB to tens of MiB; +/// 64 MiB is a generous ceiling. Caps the allocation so a crafted ICB +/// info_length / extent length cannot force a huge zeroed reservation +/// before any data is read. +const MAX_FILE_BYTES: u64 = 64 * 1024 * 1024; + +/// Upper bound on a single directory's on-disc data. Real BD-ROM +/// directories are a few KiB; 1 MiB is well above any legitimate value. +/// Caps the allocation so a corrupt 30-bit directory ICB allocation +/// length cannot force a ~1 GiB zeroed allocation per recursion level. +const MAX_DIR_BYTES: u32 = 1024 * 1024; + /// A UDF filesystem parsed from disc. #[derive(Debug)] pub struct UdfFs { @@ -80,9 +93,7 @@ impl UdfFs { Some(current) } - /// Read a file by path, returning its raw bytes. - /// Reads sector by sector from disc — no buffering. - /// Get the absolute starting LBA of a file on disc. + /// Get the absolute starting LBA of a file's first data extent on disc. /// Used by the rip pipeline to locate m2ts content sectors. pub fn file_start_lba(&self, reader: &mut dyn SectorSource, path: &str) -> Result<u32> { let parts: Vec<&str> = path.trim_matches('/').split('/').collect(); @@ -112,9 +123,17 @@ impl UdfFs { path: path.to_string(), })?; let (data_lba, _) = self.read_icb_extent(reader, entry.meta_lba)?; - Ok(self.partition_start + data_lba) + self.partition_start + .checked_add(data_lba) + .ok_or(Error::DiscRead { + sector: self.partition_start as u64, + status: None, + sense: None, + }) } + /// Read a file by path, returning its raw bytes. + /// Reads all data extents sector by sector from disc — no buffering. pub fn read_file(&self, reader: &mut dyn SectorSource, path: &str) -> Result<Vec<u8>> { let parts: Vec<&str> = path.trim_matches('/').split('/').collect(); let mut current = &self.root; @@ -147,22 +166,76 @@ impl UdfFs { path: path.to_string(), })?; - // Read the file's ICB to get its data extent - let (data_lba, data_len) = self.read_icb_extent(reader, entry.meta_lba)?; + // Read ALL the file's data extents. Multi-extent files (fragmented or + // split across dual layers) would otherwise be silently truncated to + // the first extent, since the buffer is sized to entry.size and + // truncate() can't grow it. + let extents = self.read_icb_extents(reader, entry.meta_lba)?; - // Read file data sector by sector - // File DATA is in the physical partition (partition_start + lba), - // NOT the metadata partition. ICBs are in metadata, data is in physical. - let sector_count = (data_len as u64).div_ceil(2048) as u32; - let mut data = vec![0u8; (sector_count as usize) * 2048]; - let abs_start = self.partition_start + data_lba; - - for i in 0..sector_count { - let offset = (i as usize) * 2048; - read_sector(reader, abs_start + i, &mut data[offset..offset + 2048])?; + // Reject an oversized declared total before allocating: entry.size is + // a raw u64 off the ICB, so a crafted file could otherwise force a + // multi-hundred-MiB / GiB allocation across its extents. + if entry.size > MAX_FILE_BYTES { + return Err(Error::DiscRead { + sector: self.partition_start as u64, + status: None, + sense: None, + }); } - data.truncate(entry.size as usize); + // Read ALL the file's data extents. File DATA is in the physical + // partition (partition_start + lba), NOT the metadata partition: ICBs + // are in metadata, data is in physical. + let mut data = Vec::with_capacity(entry.size as usize); + let mut sector = [0u8; 2048]; + for (data_lba, data_len) in extents { + // Cumulative guard: entry.size and each per-extent data_len are + // capped individually above, but a crafted ICB can chain many + // small extents (read_icb_extents follows type-3 chains up to + // MAX_AD_BLOCKS) whose running total grows `data` into GiB. Reject + // once the accumulated bytes would exceed MAX_FILE_BYTES. + if data.len() as u64 + data_len as u64 > MAX_FILE_BYTES { + return Err(Error::DiscRead { + sector: self.partition_start as u64, + status: None, + sense: None, + }); + } + // data_len is the disc-controlled 30-bit extent length; reject an + // oversized extent before reading so a crafted ICB can't grow the + // buffer past MAX_FILE_BYTES. + if data_len as u64 > MAX_FILE_BYTES { + return Err(Error::DiscRead { + sector: self.partition_start as u64, + status: None, + sense: None, + }); + } + let abs_start = self + .partition_start + .checked_add(data_lba) + .ok_or(Error::DiscRead { + sector: self.partition_start as u64, + status: None, + sense: None, + })?; + let sector_count = (data_len as u64).div_ceil(2048) as u32; + for i in 0..sector_count { + let abs = abs_start.checked_add(i).ok_or(Error::DiscRead { + sector: abs_start as u64, + status: None, + sense: None, + })?; + read_sector(reader, abs, &mut sector)?; + data.extend_from_slice(§or); + } + } + + // Trim to the real file size; if extents under-covered the file (e.g. + // sparse), leave what we have rather than over-reporting. + if data.len() > entry.size as usize { + data.truncate(entry.size as usize); + } Ok(data) } @@ -170,17 +243,19 @@ impl UdfFs { /// /// Returns a list of (start_lba, sector_count) ranges covering: /// - UDF structure (AVDP, VDS, metadata partition, directories) - /// - BDMV/PLAYLIST/*.mpls, CLIPINF/*.clpi, JAR/*, META/*, *.bdmv - /// - AACS/* (Content*.cer, Unit_Key_RO.inf, CPSUnit*.cci) + /// - every non-STREAM file the tree walk reaches that is <= 50 MB /// - /// Skips: STREAM/ (video), BACKUP/, DUPLICATE/, - /// MKB_RO.inf, ContentHash*, ContentRevocation* + /// Skip policy (actual): directories named `STREAM` (case-insensitive) + /// are not descended, and individual files larger than 50 MB are + /// omitted. Nothing else is filtered by name — `BACKUP`/`DUPLICATE` + /// are traversed, and `MKB_RO.inf` is excluded only because it exceeds + /// the 50 MB cap. pub fn metadata_sector_ranges(&self, reader: &mut dyn SectorSource) -> Result<Vec<(u32, u32)>> { let mut ranges = Vec::new(); // UDF structure: sector 0 through end of metadata partition // Covers AVDP, VDS, partition descriptor, metadata ICB, FSD, all directories - let meta_end = self.metadata_start + self.metadata_sectors; + let meta_end = self.metadata_start.saturating_add(self.metadata_sectors); ranges.push((0, meta_end)); // Walk tree, collect ranges for each metadata file @@ -198,7 +273,7 @@ impl UdfFs { let mut ranges = Vec::new(); // UDF structure sectors - let meta_end = self.metadata_start + self.metadata_sectors; + let meta_end = self.metadata_start.saturating_add(self.metadata_sectors); ranges.push((0, meta_end)); // Walk entire tree including STREAM directories @@ -221,12 +296,15 @@ impl UdfFs { self.collect_all_file_ranges(reader, child, ranges)?; } else { // Include the ICB sector - ranges.push((self.meta_to_abs(child.meta_lba), 1)); + ranges.push((self.meta_to_abs(child.meta_lba)?, 1)); // Include ALL file data extents (large m2ts files have many) if let Ok(extents) = self.read_icb_extents(reader, child.meta_lba) { for (data_lba, data_len) in extents { - let abs_start = self.partition_start + data_lba; + let abs_start = match self.partition_start.checked_add(data_lba) { + Some(v) => v, + None => continue, + }; let sector_count = (data_len as u64).div_ceil(2048) as u32; ranges.push((abs_start, sector_count)); } @@ -251,17 +329,25 @@ impl UdfFs { self.collect_file_ranges(reader, child, ranges)?; } else { // Include the ICB sector itself (in metadata partition) - ranges.push((self.meta_to_abs(child.meta_lba), 1)); + ranges.push((self.meta_to_abs(child.meta_lba)?, 1)); // Include file data — skip only truly huge files (MKB_RO.inf = 134MB) if child.size > 50_000_000 { continue; } - if let Ok((data_lba, data_len)) = self.read_icb_extent(reader, child.meta_lba) { - let abs_start = self.partition_start + data_lba; - let sector_count = (data_len as u64).div_ceil(2048) as u32; - ranges.push((abs_start, sector_count)); + // Push every extent: a fragmented AACS cert / MPLS / CLPI can + // span multiple extents, and key readers downstream need all + // of them (mirror collect_all_file_ranges). + if let Ok(extents) = self.read_icb_extents(reader, child.meta_lba) { + for (data_lba, data_len) in extents { + let abs_start = match self.partition_start.checked_add(data_lba) { + Some(v) => v, + None => continue, + }; + let sector_count = (data_len as u64).div_ceil(2048) as u32; + ranges.push((abs_start, sector_count)); + } } } } @@ -269,8 +355,16 @@ impl UdfFs { } /// Convert a metadata-partition-relative LBA to an absolute sector number. - fn meta_to_abs(&self, meta_lba: u32) -> u32 { - self.metadata_start + meta_lba + /// `meta_lba` is disc-controlled, so the sum is checked to avoid a + /// wrap-to-wrong-sector on a crafted ICB. + fn meta_to_abs(&self, meta_lba: u32) -> Result<u32> { + self.metadata_start + .checked_add(meta_lba) + .ok_or(Error::DiscRead { + sector: self.metadata_start as u64, + status: None, + sense: None, + }) } /// Read an Extended File Entry (tag 266) or File Entry (tag 261) @@ -279,7 +373,10 @@ impl UdfFs { fn read_icb_extent(&self, reader: &mut dyn SectorSource, meta_lba: u32) -> Result<(u32, u32)> { let extents = self.read_icb_extents(reader, meta_lba)?; extents.first().copied().ok_or(Error::DiscRead { - sector: 0, + // Diagnostic sector only; meta_to_abs can overflow on a crafted + // meta_lba, in which case 0 is a harmless placeholder for the + // error-context field. + sector: self.meta_to_abs(meta_lba).unwrap_or(0) as u64, status: None, sense: None, }) @@ -287,14 +384,17 @@ impl UdfFs { /// Read ALL allocation extents for a file from its ICB. /// Returns Vec of (partition_relative_lba, byte_length) pairs. - /// Handles files with many extents (e.g. 88 GB m2ts files have ~90 extents). + /// Handles files with many extents (e.g. 88 GB m2ts files have ~90 extents) + /// including files whose allocation descriptors span multiple blocks via + /// continuation (extent_type 3) descriptors. fn read_icb_extents( &self, reader: &mut dyn SectorSource, meta_lba: u32, ) -> Result<Vec<(u32, u32)>> { + let icb_abs = self.meta_to_abs(meta_lba)?; let mut icb = [0u8; 2048]; - read_sector(reader, self.meta_to_abs(meta_lba), &mut icb)?; + read_sector(reader, icb_abs, &mut icb)?; let tag = u16::from_le_bytes([icb[0], icb[1]]); @@ -307,7 +407,7 @@ impl UdfFs { let ad_offset = 216 + l_ea; if ad_offset + l_ad > icb.len() { return Err(Error::DiscRead { - sector: self.meta_to_abs(meta_lba) as u64, + sector: icb_abs as u64, status: None, sense: None, }); @@ -321,7 +421,7 @@ impl UdfFs { let ad_offset = 176 + l_ea; if ad_offset + l_ad > icb.len() { return Err(Error::DiscRead { - sector: self.meta_to_abs(meta_lba) as u64, + sector: icb_abs as u64, status: None, sense: None, }); @@ -330,7 +430,7 @@ impl UdfFs { } _ => { return Err(Error::DiscRead { - sector: 0, + sector: icb_abs as u64, status: None, sense: None, }); @@ -338,25 +438,72 @@ impl UdfFs { }; let mut extents = Vec::new(); - let num_descriptors = l_ad / 8; // Short Allocation Descriptor = 8 bytes - for i in 0..num_descriptors { - let off = ad_offset + i * 8; - if off + 8 > icb.len() { - break; + // Parse the first allocation-descriptor list from the ICB. A type-3 + // descriptor ("next extent of allocation descriptors") points at a + // continuation block in the metadata partition holding more ADs; we + // follow the chain. The hop count is bounded to avoid looping on a + // crafted/corrupt disc. + let mut block = icb; + let mut ad_start = ad_offset; + let mut ad_bytes = l_ad; + const MAX_AD_BLOCKS: usize = 256; + + for _ in 0..MAX_AD_BLOCKS { + let num_descriptors = ad_bytes / 8; // Short Allocation Descriptor = 8 bytes + let mut next_block: Option<u32> = None; + + for i in 0..num_descriptors { + let off = ad_start + i * 8; + if off + 8 > block.len() { + break; + } + + let raw_len = u32::from_le_bytes([ + block[off], + block[off + 1], + block[off + 2], + block[off + 3], + ]); + let extent_type = raw_len >> 30; + let data_len = raw_len & 0x3FFF_FFFF; + let data_lba = u32::from_le_bytes([ + block[off + 4], + block[off + 5], + block[off + 6], + block[off + 7], + ]); + + match extent_type { + // Recorded and allocated. A zero-length type-0 + // descriptor is the AD-list terminator (continuation + // blocks are scanned to the end of the sector, so the + // trailing zero padding must not be read as extents). + 0 if data_len == 0 => break, + 0 => extents.push((data_lba, data_len)), + 1 => {} // allocated but not recorded (sparse) + 3 => { + // Continuation: the rest of the ADs live in the block + // at data_lba (metadata-partition-relative). Stop + // scanning this block and follow the pointer. + if data_len > 0 { + next_block = Some(data_lba); + } + break; + } + _ => break, + } } - let raw_len = u32::from_le_bytes([icb[off], icb[off + 1], icb[off + 2], icb[off + 3]]); - let extent_type = raw_len >> 30; - let data_len = raw_len & 0x3FFF_FFFF; - let data_lba = - u32::from_le_bytes([icb[off + 4], icb[off + 5], icb[off + 6], icb[off + 7]]); - - match extent_type { - 0 => extents.push((data_lba, data_len)), // recorded and allocated - 1 => {} // allocated but not recorded (sparse) — skip - 3 => break, // next extent of allocation descriptors — TODO - _ => break, + match next_block { + Some(cont_lba) => { + read_sector(reader, self.meta_to_abs(cont_lba)?, &mut block)?; + // A continuation block is a list of Short ADs from byte 0, + // spanning the whole sector. + ad_start = 0; + ad_bytes = block.len(); + } + None => break, } } @@ -398,9 +545,16 @@ impl UdfFs { })?; let alloc_extents = self.read_icb_extents(reader, entry.meta_lba)?; - let mut disc_extents = Vec::new(); + let mut disc_extents = Vec::with_capacity(alloc_extents.len()); for (lba, byte_len) in alloc_extents { - let abs_lba = self.partition_start + lba; + let abs_lba = self + .partition_start + .checked_add(lba) + .ok_or(Error::DiscRead { + sector: self.partition_start as u64, + status: None, + sense: None, + })?; let sectors = (byte_len as u64).div_ceil(2048) as u32; disc_extents.push((abs_lba, sectors)); } @@ -426,7 +580,7 @@ pub fn read_filesystem(reader: &mut dyn SectorSource) -> Result<UdfFs> { let tag_id = u16::from_le_bytes([avdp[0], avdp[1]]); if tag_id != 2 { return Err(Error::DiscRead { - sector: 0, + sector: 256, status: None, sense: None, }); @@ -538,7 +692,11 @@ pub fn read_filesystem(reader: &mut dyn SectorSource) -> Result<UdfFs> { meta_icb[ad_off + 7], ]); // Metadata content starts at partition_start + ad_pos - partition_start + ad_pos + partition_start.checked_add(ad_pos).ok_or(Error::DiscRead { + sector: partition_start as u64, + status: None, + sense: None, + })? } else { // Fallback: no metadata partition, use physical partition directly partition_start @@ -562,7 +720,7 @@ pub fn read_filesystem(reader: &mut dyn SectorSource) -> Result<UdfFs> { let fsd_tag = u16::from_le_bytes([fsd[0], fsd[1]]); if fsd_tag != 256 { return Err(Error::DiscRead { - sector: 0, + sector: metadata_start as u64, status: None, sense: None, }); @@ -586,11 +744,17 @@ pub fn read_filesystem(reader: &mut dyn SectorSource) -> Result<UdfFs> { }) } -/// Read a UDF directory and its children (up to max_depth levels). +/// Maximum directory nesting depth followed when building the tree. +/// Bounds recursion on a corrupt/looping disc; real BD-ROM and DVD trees +/// are far shallower (BDMV/BACKUP/BDJO is the deepest standard path at 3). +const MAX_DIR_DEPTH: u32 = 8; + +/// Read a UDF directory and its children (up to [`MAX_DIR_DEPTH`] levels). /// /// Each directory is an ICB (Extended File Entry) pointing to directory data /// containing File Identifier Descriptors (FIDs). Each FID names a file/subdir -/// and points to its ICB. +/// and points to its ICB. Directories deeper than [`MAX_DIR_DEPTH`] are +/// recorded as entries but not descended into. #[allow(clippy::only_used_in_recursion)] fn read_directory( reader: &mut dyn SectorSource, @@ -601,8 +765,13 @@ fn read_directory( depth: u32, ) -> Result<DirEntry> { // Read ICB for this directory + let icb_abs = meta_start.checked_add(meta_lba).ok_or(Error::DiscRead { + sector: meta_start as u64, + status: None, + sense: None, + })?; let mut icb = [0u8; 2048]; - read_sector(reader, meta_start + meta_lba, &mut icb)?; + read_sector(reader, icb_abs, &mut icb)?; let tag = u16::from_le_bytes([icb[0], icb[1]]); @@ -613,7 +782,7 @@ fn read_directory( let ad_off = 216 + l_ea; if ad_off + 8 > icb.len() { return Err(Error::DiscRead { - sector: (meta_start + meta_lba) as u64, + sector: icb_abs as u64, status: None, sense: None, }); @@ -637,7 +806,7 @@ fn read_directory( let ad_off = 176 + l_ea; if ad_off + 8 > icb.len() { return Err(Error::DiscRead { - sector: (meta_start + meta_lba) as u64, + sector: icb_abs as u64, status: None, sense: None, }); @@ -667,14 +836,36 @@ fn read_directory( } }; + // Reject an oversized directory before allocating: ad_len is the + // disc-controlled 30-bit ICB allocation length, so a corrupt value + // could otherwise force a ~1 GiB zeroed allocation (amplified by + // recursion). Real directories are a few KiB; the 1 MiB cap still + // covers a large STREAM/ dir with thousands of .m2ts FIDs. + if ad_len > MAX_DIR_BYTES { + return Err(Error::DiscRead { + sector: meta_start as u64, + status: None, + sense: None, + }); + } + // Read directory data - let dir_abs = meta_start + ad_pos; - let sector_count = ad_len.div_ceil(2048).min(64); + let dir_abs = meta_start.checked_add(ad_pos).ok_or(Error::DiscRead { + sector: meta_start as u64, + status: None, + sense: None, + })?; + let sector_count = ad_len.div_ceil(2048); let mut dir_data = vec![0u8; sector_count as usize * 2048]; for i in 0..sector_count { + let abs = dir_abs.checked_add(i).ok_or(Error::DiscRead { + sector: dir_abs as u64, + status: None, + sense: None, + })?; read_sector( reader, - dir_abs + i, + abs, &mut dir_data[(i as usize) * 2048..(i as usize + 1) * 2048], )?; } @@ -720,8 +911,11 @@ fn read_directory( // Read the ICB to get file size let file_size = read_file_size(reader, meta_start, icb_lba).unwrap_or(0); - if is_dir && depth < 3 { - // Recurse into subdirectory (max 3 levels: BDMV/PLAYLIST/*.mpls) + if is_dir && depth < MAX_DIR_DEPTH { + // Recurse into subdirectory. The cap guards against + // pathological/looping directory trees on a corrupt disc + // while comfortably covering real BD-ROM nesting + // (e.g. BDMV/BACKUP/BDJO/*.bdjo is 3 levels deep). let subdir = read_directory( reader, part_start, @@ -759,8 +953,13 @@ fn read_directory( /// Read file size (info_length) from an Extended File Entry ICB. fn read_file_size(reader: &mut dyn SectorSource, meta_start: u32, meta_lba: u32) -> Result<u64> { + let abs = meta_start.checked_add(meta_lba).ok_or(Error::DiscRead { + sector: meta_start as u64, + status: None, + sense: None, + })?; let mut icb = [0u8; 2048]; - read_sector(reader, meta_start + meta_lba, &mut icb)?; + read_sector(reader, abs, &mut icb)?; let tag = u16::from_le_bytes([icb[0], icb[1]]); match tag { @@ -814,10 +1013,13 @@ fn merge_ranges(ranges: &[(u32, u32)]) -> Vec<(u32, u32)> { let mut result = vec![ranges[0]]; for &(start, count) in &ranges[1..] { let last = result.last_mut().unwrap(); - let last_end = last.0 + last.1; - if start <= last_end + 1 { + // Saturating arithmetic: ranges derive from disc-controlled ICB + // LBAs/lengths, so a corrupt disc could otherwise overflow u32 + // (panic in debug, wrap in release). + let last_end = last.0.saturating_add(last.1); + if start <= last_end.saturating_add(1) { // Overlapping or adjacent — extend - let new_end = (start + count).max(last_end); + let new_end = start.saturating_add(count).max(last_end); last.1 = new_end - last.0; } else { result.push((start, count)); @@ -868,11 +1070,11 @@ fn parse_dstring(data: &[u8]) -> String { } } -/// Read a single 2048-byte sector from the drive. -/// Uses standard READ(10) — no unlock required. -/// Buffered sector reader — reduces SCSI round-trips by pre-fetching blocks. -/// Each SCSI command has ~500ms overhead on USB drives, so reading 32 sectors -/// at once (one command) is 32x faster than 32 individual reads. +/// Buffered sector reader — reduces SCSI round-trips by coalescing +/// single-sector reads into `batch`-sized SCSI commands. Per-command +/// latency dominates on USB drives, so serving many adjacent single-sector +/// reads from one bulk read is substantially faster than issuing each +/// individually. `batch` is a runtime field, not a fixed count. pub(crate) struct BufferedSectorReader<'a> { inner: &'a mut dyn SectorSource, cache_start: u32, @@ -961,6 +1163,11 @@ impl SectorSource for BufferedSectorReader<'_> { _recovery: bool, ) -> std::result::Result<usize, crate::error::Error> { if count == 1 { + // Contract: a single-sector read needs at least one sector of + // destination. Return an error rather than panicking on the slice. + if buf.len() < 2048 { + return Err(crate::error::Error::UdfBufferTooSmall); + } // Check permanent prefetch cache first (HashMap) if let Some(data) = self.prefetched.get(&lba) { buf[..2048].copy_from_slice(data); @@ -980,7 +1187,12 @@ impl SectorSource for BufferedSectorReader<'_> { self.cache_sectors = block as u32; } Err(_) => { - // Near end of disc or error — single sector fallback + // By design: a `block`-sector batch read that starts + // valid but runs past the last recorded sector fails as + // a unit. Retry the one sector actually requested so a + // batch overrunning the disc tail still serves the live + // LBA instead of erroring; a genuinely bad single sector + // then propagates via `?`. self.cache.resize(2048, 0); self.inner.read_sectors(lba, 1, &mut self.cache, true)?; self.cache_start = lba; @@ -996,7 +1208,367 @@ impl SectorSource for BufferedSectorReader<'_> { } } +/// Read a single 2048-byte sector from the drive. +/// Uses standard READ(10) — no unlock required. fn read_sector(reader: &mut dyn SectorSource, lba: u32, buf: &mut [u8]) -> Result<()> { reader.read_sectors(lba, 1, buf, true)?; Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + /// In-memory SectorSource backed by an explicit absolute-LBA → sector map. + /// Unmapped sectors read as zeroes. + struct MapReader { + sectors: HashMap<u32, [u8; 2048]>, + } + + impl MapReader { + fn new() -> Self { + Self { + sectors: HashMap::new(), + } + } + fn put(&mut self, lba: u32, data: [u8; 2048]) { + self.sectors.insert(lba, data); + } + } + + impl SectorSource for MapReader { + fn read_sectors( + &mut self, + lba: u32, + count: u16, + buf: &mut [u8], + _recovery: bool, + ) -> Result<usize> { + let need = count as usize * 2048; + if buf.len() < need { + return Err(Error::UdfBufferTooSmall); + } + for i in 0..count as u32 { + let off = i as usize * 2048; + let s = self.sectors.get(&(lba + i)).copied().unwrap_or([0u8; 2048]); + buf[off..off + 2048].copy_from_slice(&s); + } + Ok(need) + } + } + + /// Build an Extended File Entry (tag 266) ICB sector with the given + /// info_length and a list of (extent_type, data_len, data_lba) short ADs. + fn build_efe(info_length: u64, ads: &[(u32, u32, u32)]) -> [u8; 2048] { + let mut s = [0u8; 2048]; + s[0..2].copy_from_slice(&266u16.to_le_bytes()); // tag + s[56..64].copy_from_slice(&info_length.to_le_bytes()); // info_length + let l_ea: u32 = 0; + let l_ad: u32 = (ads.len() * 8) as u32; + s[208..212].copy_from_slice(&l_ea.to_le_bytes()); + s[212..216].copy_from_slice(&l_ad.to_le_bytes()); + let mut off = 216 + l_ea as usize; + for &(etype, dlen, dlba) in ads { + let raw_len = (etype << 30) | (dlen & 0x3FFF_FFFF); + s[off..off + 4].copy_from_slice(&raw_len.to_le_bytes()); + s[off + 4..off + 8].copy_from_slice(&dlba.to_le_bytes()); + off += 8; + } + s + } + + /// A continuation block: a bare list of short ADs from byte 0. + fn build_cont_block(ads: &[(u32, u32, u32)]) -> [u8; 2048] { + let mut s = [0u8; 2048]; + let mut off = 0usize; + for &(etype, dlen, dlba) in ads { + let raw_len = (etype << 30) | (dlen & 0x3FFF_FFFF); + s[off..off + 4].copy_from_slice(&raw_len.to_le_bytes()); + s[off + 4..off + 8].copy_from_slice(&dlba.to_le_bytes()); + off += 8; + } + s + } + + fn fs_with(part_start: u32, meta_start: u32, root: DirEntry) -> UdfFs { + UdfFs { + root, + volume_id: String::new(), + partition_start: part_start, + metadata_start: meta_start, + metadata_sectors: 0, + } + } + + fn file_entry(name: &str, meta_lba: u32, size: u64) -> DirEntry { + DirEntry { + name: name.to_string(), + is_dir: false, + meta_lba, + size, + entries: Vec::new(), + } + } + + #[test] + fn icb_extents_follow_type3_continuation() { + let part_start = 1000; + let meta_start = 100; + // ICB at meta_lba 5: one real extent + a type-3 continuation pointer. + let icb = build_efe( + 6144, + &[ + (0, 4096, 10), // recorded extent at part-rel lba 10 + (3, 2048, 50), // continuation block at meta-rel lba 50 + ], + ); + // Continuation block holds the tail extent. + let cont = build_cont_block(&[(0, 2048, 20)]); + + let mut reader = MapReader::new(); + reader.put(meta_start + 5, icb); + reader.put(meta_start + 50, cont); + + let fs = fs_with(part_start, meta_start, file_entry("X", 5, 6144)); + let extents = fs.read_icb_extents(&mut reader, 5).expect("extents"); + assert_eq!(extents, vec![(10, 4096), (20, 2048)]); + } + + #[test] + fn read_file_spans_multiple_extents() { + let part_start = 0; + let meta_start = 0; + // Two extents of one sector each; distinct fill bytes per data sector. + let icb = build_efe(4096, &[(0, 2048, 10), (0, 2048, 30)]); + let mut reader = MapReader::new(); + reader.put(5, icb); + reader.put(10, [0xAA; 2048]); + reader.put(30, [0xBB; 2048]); + + let root = DirEntry { + name: String::new(), + is_dir: true, + meta_lba: 0, + size: 0, + entries: vec![file_entry("F", 5, 4096)], + }; + let fs = fs_with(part_start, meta_start, root); + let data = fs.read_file(&mut reader, "/F").expect("read"); + assert_eq!(data.len(), 4096); + assert!(data[..2048].iter().all(|&b| b == 0xAA)); + assert!(data[2048..].iter().all(|&b| b == 0xBB)); + } + + #[test] + fn merge_ranges_saturates_near_u32_max() { + // Adjacent ranges near u32::MAX must not panic (debug) or wrap. + let ranges = [(u32::MAX - 1, 2), (u32::MAX, 5)]; + let merged = merge_ranges(&ranges); + // No panic; result is a single merged range starting at the first. + assert_eq!(merged.len(), 1); + assert_eq!(merged[0].0, u32::MAX - 1); + } + + #[test] + fn buffered_reader_short_buf_errors_not_panics() { + let mut inner = MapReader::new(); + inner.put(0, [0u8; 2048]); + let mut br = BufferedSectorReader::new(&mut inner, 8); + let mut tiny = [0u8; 100]; + let err = br.read_sectors(0, 1, &mut tiny, true); + assert!(matches!(err, Err(Error::UdfBufferTooSmall))); + } + + /// Minimal in-memory SectorSource that serves pre-loaded 2048-byte + /// sectors by LBA. Unmapped LBAs read as zeros. + struct MemReader { + sectors: HashMap<u32, [u8; 2048]>, + } + + impl MemReader { + fn new() -> Self { + Self { + sectors: HashMap::new(), + } + } + fn put(&mut self, lba: u32, sector: [u8; 2048]) { + self.sectors.insert(lba, sector); + } + } + + impl SectorSource for MemReader { + fn read_sectors( + &mut self, + lba: u32, + count: u16, + buf: &mut [u8], + _recovery: bool, + ) -> Result<usize> { + for i in 0..count as u32 { + let off = i as usize * 2048; + let dst = &mut buf[off..off + 2048]; + match self.sectors.get(&(lba + i)) { + Some(s) => dst.copy_from_slice(s), + None => dst.fill(0), + } + } + Ok(count as usize * 2048) + } + } + + /// Build an Extended File Entry (tag 266) ICB sector with a single + /// short allocation descriptor declaring `data_len` bytes at `data_lba`. + /// `info_length` (offset 56) is set to `info_len`. + fn build_efe_icb(info_len: u64, data_len: u32, data_lba: u32) -> [u8; 2048] { + let mut icb = [0u8; 2048]; + // tag identifier 266 (Extended File Entry) + icb[0..2].copy_from_slice(&266u16.to_le_bytes()); + // info_length at offset 56 + icb[56..64].copy_from_slice(&info_len.to_le_bytes()); + // l_ea = 0 at offset 208, l_ad = 8 (one short AD) at offset 212 + icb[208..212].copy_from_slice(&0u32.to_le_bytes()); + icb[212..216].copy_from_slice(&8u32.to_le_bytes()); + // ad_offset = 216 + l_ea = 216. Short AD: len(4) | lba(4). + // extent_type 0 (recorded) is the top 2 bits = 0, so raw == len. + icb[216..220].copy_from_slice(&(data_len & 0x3FFF_FFFF).to_le_bytes()); + icb[220..224].copy_from_slice(&data_lba.to_le_bytes()); + icb + } + + /// Build a UdfFs with a single file entry under root, for read_file tests. + fn fs_with_file(meta_lba: u32, size: u64) -> UdfFs { + UdfFs { + root: DirEntry { + name: String::new(), + is_dir: true, + meta_lba: 0, + size: 0, + entries: vec![DirEntry { + name: "F".to_string(), + is_dir: false, + meta_lba, + size, + entries: Vec::new(), + }], + }, + volume_id: String::new(), + partition_start: 0, + metadata_start: 0, + metadata_sectors: 0, + } + } + + #[test] + fn read_file_rejects_oversized_extent_before_allocating() { + // data_len just over the 64 MiB cap must error, not allocate. + let oversized = MAX_FILE_BYTES as u32 + 2048; + let icb = build_efe_icb(oversized as u64, oversized, 100); + let mut reader = MemReader::new(); + reader.put(10, icb); // ICB at meta_lba 10 (metadata_start 0) + + let fs = fs_with_file(10, oversized as u64); + let err = fs.read_file(&mut reader, "/F").unwrap_err(); + assert!(matches!(err, Error::DiscRead { .. })); + } + + /// Build an Extended File Entry ICB with multiple inline short ADs, each + /// `(data_len, data_lba)`. Lets a test chain extents whose individual + /// lengths are all under the per-extent cap but whose running total + /// exceeds MAX_FILE_BYTES. + fn build_efe_icb_multi(info_len: u64, ads: &[(u32, u32)]) -> [u8; 2048] { + let mut icb = [0u8; 2048]; + icb[0..2].copy_from_slice(&266u16.to_le_bytes()); + icb[56..64].copy_from_slice(&info_len.to_le_bytes()); + let l_ad = (ads.len() * 8) as u32; + icb[208..212].copy_from_slice(&0u32.to_le_bytes()); + icb[212..216].copy_from_slice(&l_ad.to_le_bytes()); + for (i, (data_len, data_lba)) in ads.iter().enumerate() { + let off = 216 + i * 8; + icb[off..off + 4].copy_from_slice(&(data_len & 0x3FFF_FFFF).to_le_bytes()); + icb[off + 4..off + 8].copy_from_slice(&data_lba.to_le_bytes()); + } + icb + } + + #[test] + fn read_file_rejects_cumulative_extents_over_cap() { + // Two extents, each individually within MAX_FILE_BYTES, that together + // exceed it. The cumulative guard must fire on the second extent + // (before reading it) rather than growing `data` past the cap. + // First extent: a single sector (read, data.len() = 2048). Second + // extent: exactly MAX_FILE_BYTES (passes the per-extent cap) — the + // 2048 already buffered pushes the running total over the cap. + let big = MAX_FILE_BYTES as u32; + let icb = build_efe_icb_multi(MAX_FILE_BYTES * 2, &[(2048, 100), (big, 200_000)]); + let mut reader = MemReader::new(); + reader.put(10, icb); + let mut data_sector = [0u8; 2048]; + data_sector[0] = 0xCD; + reader.put(100, data_sector); + + // entry.size declared small so the entry.size cap passes; the + // cumulative extent total is what must trip the guard. + let fs = fs_with_file(10, 2048); + let err = fs.read_file(&mut reader, "/F").unwrap_err(); + assert!(matches!(err, Error::DiscRead { .. })); + } + + #[test] + fn read_file_rejects_oversized_info_length() { + // Small extent but a crafted huge info_length (entry.size) must also + // be rejected before truncate could be reached. + let icb = build_efe_icb(0, 2048, 100); + let mut reader = MemReader::new(); + reader.put(10, icb); + + let fs = fs_with_file(10, MAX_FILE_BYTES + 1); + let err = fs.read_file(&mut reader, "/F").unwrap_err(); + assert!(matches!(err, Error::DiscRead { .. })); + } + + #[test] + fn read_file_accepts_small_file() { + // A 1-sector file within the cap reads back its declared size. + let icb = build_efe_icb(2048, 2048, 100); + let mut reader = MemReader::new(); + reader.put(10, icb); + // file data sector at partition_start + data_lba = 0 + 100 + let mut data_sector = [0u8; 2048]; + data_sector[0] = 0xAB; + reader.put(100, data_sector); + + let fs = fs_with_file(10, 2048); + let data = fs + .read_file(&mut reader, "/F") + .expect("small file should read"); + assert_eq!(data.len(), 2048); + assert_eq!(data[0], 0xAB); + } + + #[test] + fn read_directory_rejects_oversized_dir_before_allocating() { + // A directory ICB declaring an allocation length above the 1 MiB + // ceiling must error rather than allocate a huge buffer. + let oversized = MAX_DIR_BYTES + 2048; + let icb = build_efe_icb(oversized as u64, oversized, 50); + let mut reader = MemReader::new(); + reader.put(5, icb); // directory ICB at meta_start(0) + meta_lba(5) + + let err = read_directory(&mut reader, 0, 0, 5, "DIR", 0).unwrap_err(); + assert!(matches!(err, Error::DiscRead { .. })); + } + + #[test] + fn read_directory_accepts_small_empty_dir() { + // ad_len within the cap, pointing at zeroed directory data → an empty + // (no valid FID) directory parses without error. + let icb = build_efe_icb(2048, 2048, 50); + let mut reader = MemReader::new(); + reader.put(5, icb); + // directory data at meta_start(0) + ad_pos(50) = 50 reads as zeros. + let dir = read_directory(&mut reader, 0, 0, 5, "DIR", 0).expect("small dir parses"); + assert!(dir.entries.is_empty()); + assert!(dir.is_dir); + } +} diff --git a/src/verify.rs b/src/verify.rs index 607ba4a..57ec52a 100644 --- a/src/verify.rs +++ b/src/verify.rs @@ -1,9 +1,36 @@ //! Disc sector verification — read every sector and classify health. use crate::disc::{Chapter, DiscTitle}; -use crate::progress::Progress; +use crate::progress::{PassProgress, Progress}; use crate::sector::SectorSource; -use std::time::Instant; +use std::time::{Duration, Instant}; + +/// Pause before the single retry attempt on a failed sector, letting a +/// drive that briefly went NOT READY spin back up. Slept in +/// [`RETRY_POLL_INTERVAL`] increments so a cancel request (the progress +/// callback returning `false`) is observed within that interval rather +/// than after the full pause. +const RETRY_PAUSE: Duration = Duration::from_secs(2); + +/// Polling cadence while sleeping out [`RETRY_PAUSE`]. +const RETRY_POLL_INTERVAL: Duration = Duration::from_millis(250); + +/// Sleep up to `total`, returning early with `false` if `cancelled` +/// observes a stop request mid-pause. Returns `true` if the full pause +/// elapsed without cancellation. +fn cancellable_pause(total: Duration, cancelled: &mut dyn FnMut() -> bool) -> bool { + let deadline = Instant::now() + total; + loop { + if cancelled() { + return false; + } + let now = Instant::now(); + if now >= deadline { + return true; + } + std::thread::sleep(RETRY_POLL_INTERVAL.min(deadline - now)); + } +} /// Health status of a single sector read. #[derive(Debug, Clone, Copy, PartialEq)] @@ -43,14 +70,23 @@ pub struct VerifyResult { impl VerifyResult { /// Percentage of sectors that are fully readable (good + slow + recovered). + /// + /// Returns `100.0` when there are no sectors. The subtraction is + /// saturating so a caller-constructed `VerifyResult` with `bad` greater + /// than `total_sectors` cannot wrap in release builds (it pins at 0%). pub fn readable_pct(&self) -> f64 { if self.total_sectors == 0 { return 100.0; } - (self.total_sectors - self.bad) as f64 / self.total_sectors as f64 * 100.0 + self.total_sectors.saturating_sub(self.bad) as f64 / self.total_sectors as f64 * 100.0 } - /// True if every sector read successfully. + /// True if every sector read first-attempt clean — no bad, no + /// recovered-on-retry, and no slow sectors. A `Recovered` or `Slow` + /// sector means data was eventually returned but the surface is + /// degraded, so `is_perfect()` is intentionally stricter than + /// "fully readable" (which [`readable_pct`](Self::readable_pct) + /// reports by counting recovered + slow as readable). pub fn is_perfect(&self) -> bool { self.bad == 0 && self.recovered == 0 && self.slow == 0 } @@ -79,8 +115,20 @@ impl VerifyResult { } /// Verify all sectors in a title's extents. -/// Reads in batches for speed, falls back to single-sector on failure. -/// The progress callback returns false to request early stop. +/// +/// Reads in batches for speed, falling back to single-sector reads on a +/// batch failure. Each sector that fails its single read gets one retry +/// after a [`RETRY_PAUSE`] cool-down (cancellable — see below). +/// +/// `batch_sectors` is the read granularity in 2048-byte sectors; a value +/// of `0` is treated as `1` (a `0` batch would make no forward progress). +/// +/// `on_progress` is both the progress sink and the cancellation channel: +/// returning `false` from [`Progress::report`] requests an early stop. The +/// stop is honoured between sectors *and* mid-retry-pause, so a verify of a +/// disc with many bad sectors stays responsive instead of blocking the full +/// pause per sector. Timing classification uses a 500 ms-per-sector +/// threshold to mark a successful-but-slow read as [`SectorStatus::Slow`]. pub fn verify_title( reader: &mut dyn SectorSource, title: &DiscTitle, @@ -96,6 +144,10 @@ pub fn verify_title( let mut sectors_done: u64 = 0; let mut byte_offset: u64 = 0; + // A zero batch size would never advance `offset` -> infinite loop. + // Clamp to at least one sector per read. + let batch_sectors = batch_sectors.max(1); + let total_sectors: u64 = title.extents.iter().map(|e| e.sector_count as u64).sum(); let mut buf = vec![0u8; batch_sectors as usize * 2048]; @@ -104,7 +156,10 @@ pub fn verify_title( while offset < ext.sector_count { let remaining = ext.sector_count - offset; let count = remaining.min(batch_sectors as u32) as u16; - let lba = ext.start_lba + offset; + // `start_lba + offset` stays within the extent's own LBA span by + // construction, but a crafted/corrupt extent can push the sum past + // u32::MAX; saturate rather than wrap (release) or panic (debug). + let lba = ext.start_lba.saturating_add(offset); let bytes = count as usize * 2048; let batch_start = Instant::now(); @@ -125,12 +180,28 @@ pub fn verify_title( SectorStatus::Good => good += count as u64, SectorStatus::Slow => { slow += count as u64; - ranges.push(SectorRange { - start_lba: lba, - count: count as u32, - status: SectorStatus::Slow, - byte_offset, + // Merge with the previous range when contiguous and the + // same status, mirroring the per-sector path so a run of + // slow batches coalesces into one range instead of one + // entry per batch. + let merged = ranges.last_mut().is_some_and(|last| { + if last.status == SectorStatus::Slow + && last.start_lba.saturating_add(last.count) == lba + { + last.count = last.count.saturating_add(count as u32); + true + } else { + false + } }); + if !merged { + ranges.push(SectorRange { + start_lba: lba, + count: count as u32, + status: SectorStatus::Slow, + byte_offset, + }); + } } _ => {} } @@ -157,7 +228,7 @@ pub fn verify_title( } else { // Batch failed — test each sector individually for i in 0..count { - let sector_lba = lba + i as u32; + let sector_lba = lba.saturating_add(i as u32); let sector_offset = i as usize * 2048; let sector_byte_offset = byte_offset + i as u64 * 2048; @@ -179,8 +250,32 @@ pub fn verify_title( slow += 1; SectorStatus::Slow } else { - // Retry once more after brief pause - std::thread::sleep(std::time::Duration::from_secs(2)); + // Retry once more after a brief cool-down. The pause is + // cancellable via the progress callback so a long run of + // bad sectors doesn't pin the thread for 2s each with no + // way to stop. If cancelled mid-pause, skip the retry and + // count the sector as bad, then bail out of the scan. + let cancelled_during_pause = !cancellable_pause(RETRY_PAUSE, &mut || { + on_progress.is_some_and(|cb| { + !cb.report(&PassProgress { + kind: crate::progress::PassKind::Verify, + work_done: sectors_done, + work_total: total_sectors, + bytes_good_total: (good + slow + recovered) * 2048, + bytes_unreadable_total: bad * 2048, + bytes_pending_total: 0, + bytes_total_disc: total_sectors * 2048, + disc_duration_secs: Some(title.duration_secs), + bytes_bad_in_main_title: 0, + main_title_duration_secs: Some(title.duration_secs), + main_title_size_bytes: Some(total_sectors * 2048), + }) + }) + }); + if cancelled_during_pause { + bad += 1; + break 'outer; + } if reader .read_sectors( sector_lba, @@ -201,7 +296,9 @@ pub fn verify_title( if status != SectorStatus::Good { // Merge with previous range if contiguous and same status if let Some(last) = ranges.last_mut() { - if last.status == status && last.start_lba + last.count == sector_lba { + if last.status == status + && last.start_lba.saturating_add(last.count) == sector_lba + { last.count += 1; } else { ranges.push(SectorRange { @@ -258,3 +355,143 @@ pub fn verify_title( elapsed_secs: start.elapsed().as_secs_f64(), } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::disc::{ContentFormat, DiscTitle, Extent}; + use crate::error::Error; + + /// Synthetic source: every read succeeds, filling the buffer with zeros. + struct AlwaysGood; + impl SectorSource for AlwaysGood { + fn read_sectors( + &mut self, + _lba: u32, + count: u16, + buf: &mut [u8], + _recovery: bool, + ) -> Result<usize, Error> { + let n = count as usize * 2048; + buf[..n].fill(0); + Ok(n) + } + } + + /// Synthetic source: every read fails. + struct AlwaysBad; + impl SectorSource for AlwaysBad { + fn read_sectors( + &mut self, + _lba: u32, + _count: u16, + _buf: &mut [u8], + _recovery: bool, + ) -> Result<usize, Error> { + Err(Error::ScsiError { + opcode: 0x28, + status: crate::scsi::SCSI_STATUS_CHECK_CONDITION, + sense: None, + }) + } + } + + fn title_with_extent(start_lba: u32, sector_count: u32) -> DiscTitle { + DiscTitle { + playlist: String::new(), + playlist_id: 0, + duration_secs: 0.0, + size_bytes: 0, + clips: Vec::new(), + streams: Vec::new(), + chapters: Vec::new(), + extents: vec![Extent { + start_lba, + sector_count, + }], + content_format: ContentFormat::BdTs, + codec_privates: Vec::new(), + } + } + + #[test] + fn batch_zero_does_not_hang() { + // A zero batch size must be clamped to 1 internally; otherwise the + // inner `while offset < sector_count` loop never advances. The + // assertion only matters because the test returns at all. + let title = title_with_extent(0, 4); + let mut src = AlwaysGood; + let r = verify_title(&mut src, &title, 0, None); + assert_eq!(r.total_sectors, 4); + assert_eq!(r.good, 4); + assert!(r.is_perfect()); + } + + #[test] + fn all_good_is_perfect() { + let title = title_with_extent(100, 10); + let mut src = AlwaysGood; + let r = verify_title(&mut src, &title, 4, None); + assert_eq!(r.good, 10); + assert_eq!(r.bad, 0); + assert!(r.ranges.is_empty()); + assert_eq!(r.readable_pct(), 100.0); + } + + #[test] + fn cancel_during_retry_pause_stops_promptly() { + // AlwaysBad forces the single-sector retry path, which pauses before + // its retry. The progress callback requests stop on its first call, + // so the pause must return early and the scan must bail out rather + // than sleeping RETRY_PAUSE * sector_count. A generous wall-clock + // bound (well under a single full pause) proves cancellation worked. + let title = title_with_extent(0, 8); + let mut src = AlwaysBad; + let cb = |_p: &PassProgress| false; // always request stop + let started = Instant::now(); + let r = verify_title(&mut src, &title, 4, Some(&cb)); + assert!( + started.elapsed() < RETRY_PAUSE, + "verify did not honour cancel during retry pause" + ); + // Cancelled mid-first-bad-sector: exactly one sector counted bad. + assert_eq!(r.bad, 1); + } + + #[test] + fn readable_pct_saturates_on_inconsistent_counts() { + let r = VerifyResult { + total_sectors: 10, + good: 0, + slow: 0, + recovered: 0, + bad: 9999, // impossible in practice; must not wrap + ranges: Vec::new(), + elapsed_secs: 0.0, + }; + assert_eq!(r.readable_pct(), 0.0); + } + + #[test] + fn high_lba_extent_does_not_overflow() { + // An extent whose start_lba + offset would exceed u32::MAX must + // saturate, not panic in debug builds. + let title = title_with_extent(u32::MAX - 1, 4); + let mut src = AlwaysGood; + let r = verify_title(&mut src, &title, 2, None); + assert_eq!(r.total_sectors, 4); + } + + #[test] + fn no_progress_callback_still_completes_on_bad_sectors() { + // Without a callback there is no cancel signal, so each bad sector + // takes the full RETRY_PAUSE. Keep the extent tiny (1 sector) so the + // test stays fast while still exercising the bad+retry path. + let title = title_with_extent(0, 1); + let mut src = AlwaysBad; + let r = verify_title(&mut src, &title, 1, None); + assert_eq!(r.bad, 1); + assert_eq!(r.ranges.len(), 1); + assert_eq!(r.ranges[0].status, SectorStatus::Bad); + } +} diff --git a/tests/pass_n_size_aware_skip.rs b/tests/pass_n_size_aware_skip.rs index de19de6..de83bae 100644 --- a/tests/pass_n_size_aware_skip.rs +++ b/tests/pass_n_size_aware_skip.rs @@ -13,6 +13,7 @@ use libfreemkv::disc::CopyOptions; use libfreemkv::disc::DiscRegion; +use libfreemkv::disc::PatchOptions; use libfreemkv::disc::mapfile::{Mapfile, SectorStatus}; use libfreemkv::error::Result; use libfreemkv::{ContentFormat, Disc, DiscFormat, SectorSource}; @@ -223,6 +224,84 @@ fn patch_recovers_good_middle_of_a_bad_range() { ); } +/// Regression: `PatchOptions::block_sectors == Some(0)` must not +/// busy-spin. `block_sectors` is a public `Option<u16>` field; a zero +/// value would compute a zero-length read every iteration, never +/// advance `block_end`, and burn a CPU core until the per-range +/// watchdog fired (up to 30 min on a large range). The entry-point +/// `.max(1)` clamp turns Some(0) into a single-sector batch so the +/// range recovers and the call returns promptly. +#[test] +fn patch_block_sectors_zero_does_not_busy_spin() { + let capacity_sectors: u32 = 256; + let total_bytes: u64 = capacity_sectors as u64 * SECTOR_SIZE as u64; + + // Small NonTrimmed range that is entirely readable (no bad LBAs), so + // single-sector patch reads recover it immediately. Without the + // clamp the loop would never progress regardless of readability. + let (mut reader, _trace) = PatternedSectorReader::new(capacity_sectors, HashSet::new()); + let disc = synthetic_disc(capacity_sectors); + + let tmp = tempfile::NamedTempFile::new().unwrap(); + let iso_path = tmp.path().to_path_buf(); + drop(tmp); + + let finished = [ + (0, 100 * 2048), + (110 * 2048, (capacity_sectors as u64 - 110) * 2048), + ]; + let nontrimmed = [(100 * 2048, 10 * 2048)]; + prep_iso_and_mapfile(&iso_path, total_bytes, &finished, &nontrimmed); + + // A halt watchdog bounds the run: the inner loop polls `halt` every + // iteration, so even a busy-spin regression breaks out within the + // window instead of hanging the test binary. With the clamp the run + // finishes long before the watchdog fires; without it the watchdog + // trips and the bytes_good assertion below fails loudly. + let halt = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let halt_for_watchdog = halt.clone(); + let watchdog = std::thread::spawn(move || { + std::thread::sleep(std::time::Duration::from_secs(20)); + halt_for_watchdog.store(true, std::sync::atomic::Ordering::Relaxed); + }); + + let opts = PatchOptions { + decrypt: false, + block_sectors: Some(0), + full_recovery: false, + reverse: false, + wedged_threshold: 0, + progress: None, + halt: Some(halt.clone()), + }; + + let outcome = disc.patch(&mut reader, &iso_path, &opts); + // Stop the watchdog regardless of outcome. + halt.store(true, std::sync::atomic::Ordering::Relaxed); + let _ = watchdog.join(); + + let map_path = libfreemkv::disc::mapfile_path_for(&iso_path); + let _ = std::fs::remove_file(&iso_path); + let _ = std::fs::remove_file(&map_path); + + let outcome = outcome.expect("patch returns Ok"); + assert!( + !outcome.halted, + "patch with block_sectors=Some(0) must complete on its own \ + (clamped to a 1-sector batch), not be cut off by the watchdog" + ); + let bytes_good = outcome.bytes_good; + // The 10-sector NonTrimmed range was fully readable; clamped to a + // 1-sector batch it must recover. Initial good = 100 + (256-110) = + // 246 sectors; after patch the 10-sector range is also Finished. + let initial_good_sectors: u64 = 100 + (capacity_sectors as u64 - 110); + assert!( + bytes_good >= (initial_good_sectors + 10) * 2048, + "block_sectors=Some(0) clamped to 1 should recover the readable range; \ + bytes_good={bytes_good}" + ); +} + /// A second test: a bad range that's actually 4 small bad sub-zones /// separated by good sectors. Demonstrates the bisection behaviour /// converges when zones are non-uniform.