mirror of
https://github.com/freemkv/libfreemkv.git
synced 2026-10-07 18:13:10 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4f1dbfd042 | ||
|
|
5b702a76a7 | ||
|
|
6be5198886 | ||
|
|
e2aa9abd6d | ||
|
|
3b7bee9ed4 | ||
|
|
b7405e2d27 | ||
|
|
e04d79c593 | ||
|
|
b34af1fa74 | ||
|
|
a7317f8885 | ||
|
|
474273afc0 | ||
|
|
0a2bab5789 | ||
|
|
fa4d7ef871 | ||
|
|
36d1af1b7f | ||
|
|
b96f6206fe | ||
|
|
bb32fb993b | ||
|
|
c1eb74dfa5 | ||
|
|
46838c63ca | ||
|
|
1e60220ff6 | ||
|
|
080f03e8ed | ||
|
|
35de6101d4 | ||
|
|
9012101573 | ||
|
|
8bc1de6c9b | ||
|
|
d94a4d3444 | ||
|
|
1bd7e1de2a | ||
|
|
da62ee7cf2 | ||
|
|
575c76156f | ||
|
|
b518860d9c | ||
|
|
c8eb42b490 | ||
|
|
48570ac065 | ||
|
|
08e46640fd | ||
|
|
a9195824ff | ||
|
|
401fe23988 | ||
|
|
f89bce5851 | ||
|
|
dfccb85e15 | ||
|
|
ebedffb762 | ||
|
|
fdf63ccb7f | ||
|
|
c735d284da | ||
|
|
94ab7bc73c | ||
|
|
9f209fe066 | ||
|
|
32a1a6e095 | ||
|
|
1d3b8f5fb6 | ||
|
|
97b0ae7be2 | ||
|
|
7756f1feca | ||
|
|
1565da610a | ||
|
|
e134616422 | ||
|
|
8d54a3c64e | ||
|
|
c0478e1273 | ||
|
|
7d29168fec | ||
|
|
8e14c9b850 | ||
|
|
e1c8343f77 | ||
|
|
a956c6ad94 | ||
|
|
1805d92ca4 | ||
|
|
823f0ad430 | ||
|
|
477bdf1835 | ||
|
|
4d83b69c20 | ||
|
|
7dbbfc6726 | ||
|
|
e635c9556f | ||
|
|
d7b5c30f5d | ||
|
|
739a276a39 | ||
|
|
5f1028a62a | ||
|
|
ea15d212de | ||
|
|
00673c8ec3 | ||
|
|
bfefb4cb5b | ||
|
|
1b95193517 | ||
|
|
7dcac44136 | ||
|
|
4da559e39f | ||
|
|
eeca250f69 | ||
|
|
c51b3181f2 | ||
|
|
2a31a47434 | ||
|
|
b08662f95c | ||
|
|
4594196b6e | ||
|
|
9d6974df3e | ||
|
|
3394a5b3fe | ||
|
|
e2c7e20329 | ||
|
|
ff22385a61 | ||
|
|
d39d3933f5 | ||
|
|
31dc3d35c3 | ||
|
|
f7165a2020 | ||
|
|
d836f8b1e7 | ||
|
|
8919de5b03 | ||
|
|
0dc1f12108 | ||
|
|
6abc25a7e8 | ||
|
|
4d22cf632f | ||
|
|
611ad77586 | ||
|
|
60590193e3 | ||
|
|
5f24843d3c | ||
|
|
ecf9a4d01e | ||
|
|
8ec767339e | ||
|
|
335e8b68cc | ||
|
|
1137b4ea1f | ||
|
|
7788c560d1 | ||
|
|
84e0c2ca69 | ||
|
|
f4a881b25d | ||
|
|
1edda7fedb | ||
|
|
3dd71a0650 | ||
|
|
579c94bd23 | ||
|
|
3b2a68fe7a | ||
|
|
5b56f112a0 | ||
|
|
f27e4c3088 | ||
|
|
b3576f2ae6 | ||
|
|
39f2d0e992 | ||
|
|
50867516b8 | ||
|
|
193e7c1680 | ||
|
|
b6463729e8 | ||
|
|
dffee56102 | ||
|
|
d8f27cdee0 | ||
|
|
fc13268dc6 | ||
|
|
d9b55f02d0 | ||
|
|
35d66d2059 | ||
|
|
d703ce439b | ||
|
|
6112a26d15 | ||
|
|
2a33364253 | ||
|
|
fa3872ddcc | ||
|
|
c427389f36 | ||
|
|
6980f562df | ||
|
|
e3b2c9d850 | ||
|
|
a383200ef1 | ||
|
|
51d15f551e | ||
|
|
523af461a5 | ||
|
|
695b65149e | ||
|
|
d3f9560689 | ||
|
|
1fa5a7d27f | ||
|
|
55219070f1 | ||
|
|
a34c419521 | ||
|
|
b179846f5d | ||
|
|
e110e80e6e | ||
|
|
597fa34099 | ||
|
|
52eeb949e3 | ||
|
|
6ca62b8411 | ||
|
|
04195c27d8 | ||
|
|
5a8f8e54e1 | ||
|
|
1a2c830cf6 | ||
|
|
637c2bfe36 | ||
|
|
0d28357b4d | ||
|
|
5495332f07 | ||
|
|
e22fc6fd47 | ||
|
|
5b98c13e47 | ||
|
|
8f8f1a62a2 | ||
|
|
e3cfd27942 | ||
|
|
aa12bdad62 | ||
|
|
ef3895cdc5 | ||
|
|
2dcf969ac8 | ||
|
|
005f887bf9 |
@@ -4,3 +4,7 @@ Cargo.lock
|
||||
*.swo
|
||||
.DS_Store
|
||||
.cargo/
|
||||
.claude/worktrees/
|
||||
|
||||
# session scratch — never track (may contain RE breadcrumbs)
|
||||
scratch/
|
||||
|
||||
+465
-22
@@ -1,5 +1,451 @@
|
||||
# Changelog
|
||||
|
||||
## 0.29.0 (2026-06-06)
|
||||
|
||||
### Fixed
|
||||
|
||||
- **Video frames now carry their presentation timestamp (PTS), not the decode
|
||||
timestamp (DTS), as the MKV block timecode.** The HEVC, H.264, VC-1 and
|
||||
MPEG-2 parsers stamped each frame with `dts.or(pts)`, so B-frame titles were
|
||||
muxed with monotonic decode-order timecodes: players presented frames in
|
||||
decode order (visible judder / apparent "corruption") and PTS-based seeking
|
||||
landed on the wrong frame. The compressed video was always byte-correct —
|
||||
this was purely a timestamp defect, and it affected every B-frame title. The
|
||||
parsers now use `pts.or(dts)`. Verified against a reference muxer on a real
|
||||
UHD title: the emitted frame sequence is now identical.
|
||||
- **DTS-HD MA tracks no longer downgrade to the lossy DTS core** — the access
|
||||
unit is delimited by the next *validated* core sync, so a `0x7FFE8001` byte
|
||||
pattern inside the XLL extension can no longer truncate the unit. (#10)
|
||||
- **A resolved AACS key is verified against the content before it is applied**,
|
||||
so a stale or wrong VUK can no longer silently produce garbage output.
|
||||
- **`iso://` mux fails fast when no usable AACS key is available** instead of
|
||||
writing a garbage MKV.
|
||||
|
||||
### Changed
|
||||
|
||||
- Codec / EBML / TS hardening: bounds-checked PMT and PES-header parsing, EBML
|
||||
length validation, and bounded codec buffers (DTS, TrueHD, PGS).
|
||||
|
||||
## 0.27.3 (2026-06-04)
|
||||
|
||||
### Added
|
||||
|
||||
- **The mapfile now persists resolved AACS unit keys, not just the Volume ID.**
|
||||
A keyed disc writes its decrypted unit keys as `# freemkv-uk: <cps>:<hex>`
|
||||
comment headers; an unresolved disc writes only the `# freemkv-vid:` marker.
|
||||
The two are **mutually exclusive** (`Mapfile::set_unit_keys` clears the VID):
|
||||
unit keys are the final answer, so a deferred-mux / resume decrypts directly
|
||||
with no key-service round-trip, while the VID alone means "still unresolved —
|
||||
retry the key service." `CopyOptions`/`SweepOptions` gain a `unit_keys` field
|
||||
(written when non-empty, else the VID), and `Disc::inject_unit_keys` applies
|
||||
mapfile-recovered keys to a scanned disc (marking the source `ExternalUk`).
|
||||
|
||||
## 0.27.0 (2026-06-03)
|
||||
|
||||
### Changed
|
||||
|
||||
- **AACS unit-encryption detection now reads the raw MPEG-TS sync bytes
|
||||
instead of header flag bits.** `is_unit_encrypted` is renamed
|
||||
`is_aacs_scrambled`: a unit is encrypted iff its body TS packet sync bytes
|
||||
(`0x47` at offset 4 and every 192 bytes) are not intact — the encrypted
|
||||
body destroys them. The previous check read the TP_extra copy-control bits
|
||||
(byte 0) or TS transport_scrambling_control bits (byte 7), which discs and
|
||||
players do not set reliably. A single predicate now backs both the decrypt
|
||||
gate and out-of-band key validation, so all callers agree on what
|
||||
"encrypted" means. Decryption restores the syncs, so a decrypted unit reads
|
||||
as clear and there is no encryption flag to clear.
|
||||
|
||||
## 0.26.1 (2026-05-22)
|
||||
|
||||
### Added
|
||||
|
||||
- **AACS resolver path 5 — KEYDB unit-keys direct fallback.** When a
|
||||
KEYDB entry for the disc has no VUK field but does have pre-decrypted
|
||||
unit keys, the resolver now consumes those unit keys directly instead
|
||||
of treating the entry as unusable. Covers ~4,572 entries in the
|
||||
public keydb (≈2.5% of the database, heavily skewed toward MKBv76+
|
||||
UHD discs that DVDFab/FindVUK can no longer extract VUKs for).
|
||||
Surfaces as `KeySource::KeyDbUnitKeys`; `AacsState::vuk` and
|
||||
`ResolvedKeys::vuk` are `None` on this path because no VUK exists to
|
||||
return. Partial CPS-unit coverage is rejected so a disc is never
|
||||
half-decrypted.
|
||||
|
||||
### Changed
|
||||
|
||||
- **AACS resolver path order reordered root-to-leaf.** Paths now run in
|
||||
derivation-strength order: device-key (1) → processing-key (2) →
|
||||
KEYDB-derived from MK+VID (3) → KEYDB VUK direct (4) → KEYDB unit
|
||||
keys direct (5). Previous order was leaf-first (KEYDB hash lookup,
|
||||
then derivation). Resolution runs once per disc, so the few extra
|
||||
milliseconds of MKB walking when a KEYDB VUK would have answered are
|
||||
invisible against the rip itself; the new order matches the AACS
|
||||
derivation hierarchy. `KeySource` enum variants are reordered to
|
||||
match.
|
||||
- **`AacsState::vuk` type is now `Option<[u8; 16]>`** (was `[u8; 16]`).
|
||||
Required by path 5, which has no VUK. Public-API break.
|
||||
- **`ResolvedKeys::vuk` type is now `Option<[u8; 16]>`** (was
|
||||
`[u8; 16]`). Same rationale.
|
||||
|
||||
## 0.25.14 (2026-05-21)
|
||||
|
||||
### 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`.
|
||||
|
||||
No behavioural change — purely a rename pass.
|
||||
|
||||
## 0.25.13 (2026-05-21)
|
||||
|
||||
### Added
|
||||
|
||||
- **`DrmScheme` top-level dispatcher.** New `drm` module with a
|
||||
`DrmScheme` enum (`Css`, `Aacs10`, `Aacs20`, `Aacs21`) and a
|
||||
`detect` + `load` pair that uniformly handles all four content
|
||||
protection schemes. Replaces the inlined dispatch in
|
||||
`disc::encrypt::resolve_encryption` and the scattered CSS routing
|
||||
in `disc::mod`. Both CSS call sites now route through the same
|
||||
entry point.
|
||||
- **AACS 2.1 Media Key Variant framework.** New `aacs::variants`
|
||||
module implementing the Media Key Variant derivation chain
|
||||
(`Kp + C → Kmp → ⊕KCD → Kpnew → Km → VUK`), Variant-scheme MKB
|
||||
record parsing (record types `0x82` / `0x83`), bit-0x02 SoftKCD
|
||||
and bit-0x04 online-challenge detection with dedicated error
|
||||
variants. Wired into `DrmScheme::Aacs21` but the dispatcher arm
|
||||
is commented out pending validation against a Variant-scheme
|
||||
disc. Per-manufacturer Key Correction Data must be supplied by
|
||||
the integrator; `KEY_CORRECTION_DATA_PLACEHOLDER` is the empty
|
||||
placeholder slot.
|
||||
- **`AacsVersion` enum.** Replaces the `aacs2: bool` field on
|
||||
`ContentCertificate`, `UnitKeyFile`, and `ResolvedKeys`.
|
||||
`parse_unit_key_ro` and `parse_content_cert` now take/emit the
|
||||
enum. `resolve_keys` is split into `resolve_keys_v1`,
|
||||
`resolve_keys_v2`, and `resolve_keys_v21` (the last not reachable
|
||||
from the dispatcher today).
|
||||
|
||||
### Fixed
|
||||
|
||||
- **Libredrive raw-read VID shortcut deleted.** v0.25.11 introduced a
|
||||
`do_handshake` branch that, on libredrive-active drives, skipped the
|
||||
AACS cert handshake and issued `READ_DISC_STRUCTURE` format 0x80
|
||||
with AGID=0 directly. The hypothesis was that unlocked
|
||||
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
|
||||
`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
|
||||
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
|
||||
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 corollary: AACS resolution on HRL-burned drives + UHD discs now
|
||||
fails honestly. Either cert auth succeeds (drive unlock may or
|
||||
may not bypass the HRL — that's the new empirical question) and we
|
||||
hit the actual DK wall (E7018 "No DK that walks this MKB" for v77+
|
||||
UHD without a v77+ DK in keydb), or cert auth fails and we surface
|
||||
E7015. Both are real verdicts; E7017's previous spurious dispatch
|
||||
is gone.
|
||||
|
||||
## 0.25.12 (2026-05-21)
|
||||
|
||||
No libfreemkv source changes — unified sync bump for autorip's
|
||||
`aacs_failure_message` two-line wording rewrite. See the autorip
|
||||
v0.25.12 release for details.
|
||||
|
||||
## 0.25.11 (2026-05-21)
|
||||
|
||||
### 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`
|
||||
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.
|
||||
- New `Error` variants for finer-grained AACS failure reporting:
|
||||
`AacsHostCertRejected` (E7015), `AacsLibredriveUnsupported`
|
||||
(E7016), `AacsVidUnavailable` (E7017), `AacsMkUnavailable`
|
||||
(E7018), `AacsVukNotInKeydb` (E7019). Lets CLIs/UIs render which
|
||||
piece of the AACS chain failed instead of always saying "no keys."
|
||||
|
||||
### Fixed
|
||||
|
||||
- `validate_processing_key` now matches libaacs `_validate_pk`
|
||||
exactly: XORs `uv` into `mk[12..16]` (was omitted), AES-decrypts
|
||||
`mk_dv` and checks the `01 23 45 67 89 AB CD EF` magic (was
|
||||
AES-encrypt + 12-zero check). Pre-fix, every non-zero-uv
|
||||
processing key was rejected — i.e. essentially every real disc.
|
||||
- `mkb_find_cvalues` now prefers record type `0x07` (AACS 2.x) and
|
||||
falls back to `0x05` (AACS 1.0), so the walker handles both
|
||||
generations without an out-of-band version flag.
|
||||
- `resolve_keys` short-circuits paths 2/3/4 when VID is the zero
|
||||
sentinel — saves cycles and emits an honest "VID unavailable" log
|
||||
instead of the misleading "all paths failed."
|
||||
- AES-CMAC VID verification gains NIST SP 800-38B KAT + round-trip
|
||||
+ mutation + all-zero-rejection tests.
|
||||
|
||||
### Removed
|
||||
|
||||
- **Built-in AACS keys** (added in 0.25.9). `src/aacs/builtin_keys.rs`
|
||||
deleted; `KeyDb::with_builtins`, `KeyDb::load_or_builtins`,
|
||||
`KeyDb::merge_from`, and internal dedup helpers gone. The
|
||||
compiled-in shortcut was a slim convenience that didn't move the
|
||||
hard problem (no v77+ DKs) and added a maintenance surface.
|
||||
- **Operator plugin slot** at `~/.config/freemkv/local_keys.cfg`
|
||||
(added in 0.25.9). `local_plugin_path` and `KeyDb::merge_local_plugin`
|
||||
removed. Single source: the main `keydb.cfg`.
|
||||
- `KeyDb::load` reverts to the pre-0.25.9 form — read the file or
|
||||
return I/O error; no layering, no fallback.
|
||||
- `Disc::scan` reverts to surfacing `KeydbLoad { path: "<no keydb
|
||||
in search paths>" }` for an encrypted disc with no keydb — same
|
||||
sentinel autorip's message switch already handles.
|
||||
|
||||
CSS player keys in `src/css/auth.rs` remain compiled in; they're
|
||||
1999-era public inputs separate from the AACS pipeline and were
|
||||
never part of the 0.25.9 additions.
|
||||
|
||||
## 0.25.9 (2026-05-20)
|
||||
|
||||
### Fixed
|
||||
|
||||
- **MKB record-type identification.** `mkb_find_mk_dv` and `mkb_version`
|
||||
had their record-type constants swapped. The Verify Media Key record
|
||||
is Type 0x81 (AACS 1.0) or Type 0x86 (AACS 2.0/2.1), not 0x10. The
|
||||
Type-and-Version record is Type 0x10, not 0x81. PK and DK derivation
|
||||
paths therefore failed silently on every disc, masking how often the
|
||||
fallback paths could have succeeded. Fixed; mk_dv extracted at
|
||||
offset 4 of the verify record, MKB version read at offset 8 of the
|
||||
Type 0x10 record body. Tests added for both forms.
|
||||
|
||||
### Added
|
||||
|
||||
- **Built-in AACS keys.** Four device keys (covering MKB v01-v82+) and
|
||||
three processing keys (covering v63-v68) are now compiled into
|
||||
`libfreemkv` directly. DVDs and Blu-rays (AACS 1.0) decrypt with
|
||||
zero external files. Combined with the existing 31 CSS player keys
|
||||
in `css/auth.rs`, the library is self-sufficient for all DVD and
|
||||
AACS 1.0 content.
|
||||
- **Operator plugin slot.** `~/.config/freemkv/local_keys.cfg` is
|
||||
loaded automatically (same format as `keydb.cfg`) and layered
|
||||
additively on top of the built-ins and the main keydb.cfg. For
|
||||
operators who derive their own keys, this is the additive surface
|
||||
to drop them into — main keydb.cfg from upstream auto-update stays
|
||||
separate and overwritable.
|
||||
- `KeyDb::with_builtins()`, `KeyDb::load_or_builtins()`,
|
||||
`KeyDb::empty()` constructors.
|
||||
|
||||
### Changed
|
||||
|
||||
- `Disc::scan` no longer errors when keydb.cfg is absent. With
|
||||
built-ins covering DVD/BD, missing keydb is treated as "no UHD
|
||||
keys available" — the AACS 2.0 / UHD code path surfaces a
|
||||
specific error only when the disc actually needs keys that
|
||||
aren't in built-ins, main keydb, or local plugin.
|
||||
|
||||
## 0.25.8 (2026-05-20)
|
||||
|
||||
### Changed
|
||||
|
||||
- Synchronized to autorip 0.25.8. No libfreemkv code changes —
|
||||
unified versioning bump only. autorip 0.25.7 image was DOA
|
||||
(missing +x on the binary); see autorip changelog.
|
||||
|
||||
## 0.25.7 (2026-05-20)
|
||||
|
||||
### Fixed — BU40N firmware wedge on KEYDB-miss discs
|
||||
|
||||
`disc::encrypt::do_handshake` used to fire up to 16 AACS authenticate
|
||||
attempts back-to-back with no pause between them. 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 (plus 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 0x05/0x24) until the
|
||||
drive is physically power-cycled.
|
||||
|
||||
Live wedge event 2026-05-20 during a UHD disc scan
|
||||
(KEYDB miss) confirmed the diagnosis and motivated this fix.
|
||||
|
||||
Defence-in-depth:
|
||||
- `MAX_CERT_ATTEMPTS` capped at 3 (was 16). If three different
|
||||
host certs all fail, more won't help — the drive doesn't have a
|
||||
match.
|
||||
- 1-second sleep between attempts. Gives the drive's firmware time
|
||||
to recover internal state between auth challenges.
|
||||
- Bail immediately on any sense code with sense_key = 0x05
|
||||
(`ILLEGAL_REQUEST`). The drive isn't merely rejecting our cert
|
||||
— it's saying "I won't talk to you anymore" — so trying more
|
||||
certs would deepen the wedge.
|
||||
|
||||
## 0.25.6 (2026-05-20)
|
||||
|
||||
### Changed
|
||||
|
||||
- Synchronized to autorip 0.25.6 release. No behavioural changes.
|
||||
|
||||
## 0.25.5 (2026-05-20)
|
||||
|
||||
### Changed
|
||||
|
||||
- Synchronized to autorip 0.25.5 release. No behavioural changes.
|
||||
|
||||
## 0.25.4 (2026-05-20)
|
||||
|
||||
### Changed
|
||||
|
||||
- Synchronized to autorip 0.25.4 release. No behavioural changes
|
||||
in the library — the v0.25.4 release ships the autorip-side
|
||||
in-container NFS mount option.
|
||||
|
||||
## 0.25.3 (2026-05-19)
|
||||
|
||||
### Changed
|
||||
|
||||
- Synchronized to autorip 0.25.3 release. No behavioural changes
|
||||
in the library — the v0.25.3 release ships the autorip-side
|
||||
parallel mux pipeline.
|
||||
|
||||
## 0.25.2 (2026-05-19)
|
||||
|
||||
### Fixed
|
||||
|
||||
- **DTS-HD codec ID** — `MkvTrack::audio` now emits `A_DTS/MA` for
|
||||
DTS-HD MA and `A_DTS/HR` for DTS-HD HR instead of mislabelling
|
||||
both as plain `A_DTS`. Strict players (Plex transcoder, some
|
||||
hardware decoders, AV receivers) reject lossless DTS-HD MA
|
||||
payload when the track advertises `A_DTS` because the ID
|
||||
implies the 1.5 Mbps core-only bitstream.
|
||||
- **PGS subtitle BlockDuration** — the PGS parser is now stateful:
|
||||
it pairs each display PCS with the following empty/clear PCS to
|
||||
compute a duration, and the MKV muxer emits a `BlockGroup` +
|
||||
`BlockDuration` for subtitles that carry one. Without this the
|
||||
last bitmap lingered on screen until the next display set
|
||||
replaced it (or until end of file).
|
||||
|
||||
### Changed
|
||||
|
||||
- `codec::Frame` gains `duration_ns: Option<u64>` (set by parsers
|
||||
that can compute one; currently only PGS).
|
||||
- `pes::PesFrame` gains `duration_ns: Option<u64>` (in-memory only;
|
||||
not part of the on-wire serialization).
|
||||
- `MkvMuxer::write_frame` now takes a final `duration_ns: Option<u64>`
|
||||
parameter. When `Some`, the frame is emitted as a `BlockGroup`
|
||||
with `BlockDuration` instead of a `SimpleBlock`.
|
||||
|
||||
## 0.25.1 (2026-05-19)
|
||||
|
||||
### New — autorip event_fn plumbing
|
||||
|
||||
- `PrefetchedSectorSource::new_with_events(reader, extents, batch,
|
||||
halt, event_fn)` — same producer-thread pipeline as `new()`, plus
|
||||
an optional callback fired from the producer with `BytesRead`
|
||||
events after every successful batch. Lets the autorip multipass
|
||||
+ resume mux paths drive their progress UI from the highway
|
||||
without polling the consumer side.
|
||||
- `build_iso_pipeline` gains an `event_fn` parameter so callers can
|
||||
wire the same callback through one ctor.
|
||||
|
||||
### Changed
|
||||
|
||||
- `pes::Stream` trait: new default `errors() -> u64` method
|
||||
(default returns 0). Lets `Box<dyn Stream>` callers query the
|
||||
skip-on-error counter without downcasting. `DiscStream` overrides
|
||||
to surface its existing `errors` field.
|
||||
|
||||
### Removed (breaking)
|
||||
|
||||
- `DiscStream::new_pipeline` and `DiscStream::read_pipeline`
|
||||
deleted. All file-backed mux now uses `build_iso_pipeline` →
|
||||
`PipelinedPesStream`. `DiscStream` is now the single-threaded
|
||||
inline reader used by autorip's live-drive single-pass path only.
|
||||
- `DiscStream::demux_thread` and `demux_rx` fields removed.
|
||||
- `M2tsStream::open` deleted; `Mode::Read` variant removed.
|
||||
`m2ts://` URLs go through the internal `build_m2ts_pipeline`
|
||||
helper in `mux/resolve.rs` → `PipelinedPesStream`. `M2tsStream`
|
||||
is a write-only sink now.
|
||||
|
||||
## 0.25.0 (2026-05-19)
|
||||
|
||||
### New — the freemkv mux throughput "highway"
|
||||
|
||||
Three-stage pipelined PES read path. Read+decrypt runs on a
|
||||
producer thread, M2TS demux runs on a second thread, codec parse
|
||||
runs on the caller's thread. Communication between stages is via
|
||||
bounded `crossbeam_channel` with a recycled buffer pool — no
|
||||
allocations or memcpys in the steady-state hot loop.
|
||||
|
||||
**Throughput on the test bed (a UHD disc, 62 GiB ISO →
|
||||
`null://`, single-thread caller):**
|
||||
|
||||
| | MB/s |
|
||||
| ------------------------------------- | -----: |
|
||||
| 0.23.2 baseline | 60 |
|
||||
| + memchr SIMD HEVC start-code scan | 69 |
|
||||
| + ts.feed no-copy boundary | 72 |
|
||||
| + `PrefetchedSectorSource` (producer) | 124 |
|
||||
| + `DemuxThread` (3-stage pipeline) | 135 |
|
||||
| + zero-copy recycled buffer pool | 148 |
|
||||
| + 16 KiB initial PesAssembler buffer | 162 |
|
||||
| + mimalloc allocator in freemkv CLI | 200+ |
|
||||
| **+ warm cache** | **660** |
|
||||
|
||||
The new public API:
|
||||
|
||||
- `libfreemkv::PrefetchedSectorSource` — wraps any `SectorSource`,
|
||||
spawns a producer thread, exposes recycled-buffer channels.
|
||||
- `libfreemkv::io::byte_prefetcher::BytePrefetcher` —
|
||||
`std::io::Read` analogue for byte-stream sources (m2ts files,
|
||||
sockets, stdin).
|
||||
- `libfreemkv::mux::demux_thread::DemuxThread` — the M2TS demux
|
||||
worker; `spawn_zero_copy` takes either prefetcher's channels.
|
||||
- `libfreemkv::PipelinedPesStream` — the read-side `Stream` impl
|
||||
that runs codec parse on the caller thread.
|
||||
- `libfreemkv::build_iso_pipeline(reader, title, keys, batch,
|
||||
format, halt)` — the canonical ctor that wires all three stages
|
||||
for an ISO file source.
|
||||
|
||||
### Changed (breaking)
|
||||
|
||||
- `IsoSectorReader` (the naive duplicate of `FileSectorSource`)
|
||||
deleted. `FileSectorSource` is the sole file-backed sector
|
||||
source; it carries the SEQUENTIAL fadvise hint, the periodic
|
||||
DONTNEED page-cache eviction, and (new) the per-read
|
||||
`readahead()` async-prefetch syscall.
|
||||
- `mux::input("iso://...", &opts)` now returns a
|
||||
`PipelinedPesStream`. Function signature
|
||||
`input(&str, &InputOptions) -> io::Result<Box<dyn Stream>>` is
|
||||
unchanged; callers that treated the return value as
|
||||
`Box<dyn Stream>` keep working.
|
||||
|
||||
### Other
|
||||
|
||||
- `FileSectorSource` exposes per-OS `prefetch()` hooks
|
||||
(`readahead(2)` on Linux, `fcntl(F_RDADVISE)` on macOS, no-op on
|
||||
Windows + other).
|
||||
- AACS decrypt thread-pool: env var renamed from
|
||||
`FREEMKV_DECRYPT_THREADS` to `FREEMKV_THREADS`; default raised
|
||||
to `cores.clamp(1, 64)`.
|
||||
- `PesAssembler` initial buffer capacity 256 KiB → 16 KiB (avoids
|
||||
the 64-page first-touch fault tax on every PES boundary).
|
||||
- HEVC / H.264 `find_start_code` swapped onto
|
||||
`memchr::memmem::find` (SIMD).
|
||||
|
||||
## 0.18.4 (2026-05-09)
|
||||
|
||||
### Build / CI hardening — no library code changes
|
||||
@@ -94,8 +540,7 @@ in-tree consumers (autorip, the `freemkv` CLI) drove their own multipass
|
||||
loops in 0.18 rounds 1-3; 0.18.1 lands the cleanup with the deprecated names
|
||||
still alive for one minor-version window.
|
||||
|
||||
The 0.18 design notes are in
|
||||
`(internal)/memory/0_18_redesign.md` (private) — this entry sticks to
|
||||
The 0.18 design notes are kept privately — this entry sticks to
|
||||
what changed at the public surface.
|
||||
|
||||
#### Flat verbs
|
||||
@@ -375,7 +820,7 @@ recovery.
|
||||
`READ_RECOVERY_TIMEOUT_MS = 60_000` (the v0.17.3 baseline; the kernel-
|
||||
auto-retry pattern is now provided by the `/dev/sr0` fallback above).
|
||||
|
||||
### Empirical results (Dune Part Two UHD, BU40N direct SATA)
|
||||
### Empirical results (a UHD disc, BU40N direct SATA)
|
||||
|
||||
- Pass 1: **94.6 MB recovered** (28% of formerly-bad data, 33 sr0 fallback
|
||||
saves), **11 s of main-title content** restored. Patch completed all 47
|
||||
@@ -503,7 +948,7 @@ bridges (notably the Initio INIC-1618L) that crash on MEDIUM ERROR retries.
|
||||
- Patch instrumentation: counters for reads_ok/err, writes_ok/err, finished/unreadable.
|
||||
- `as_encoded_bytes()` replaces `as_bytes()` for portable OsStr handling.
|
||||
- Removed inline retry/reset from Drive::read — orchestration layer handles recovery.
|
||||
- Removed all `(internal)` references from public code.
|
||||
- Removed all internal references from public code.
|
||||
|
||||
## 0.13.26 (2026-04-27)
|
||||
|
||||
@@ -531,7 +976,7 @@ warnings` red on Linux CI that the Mac toolchain doesn't catch.
|
||||
|
||||
### Pre-commit gate uses CI's exact toolchain
|
||||
|
||||
`(internal)/scripts/precommit.sh` (new) runs `cargo +1.86 fmt
|
||||
A new pre-commit script runs `cargo +1.86 fmt
|
||||
--check`, `cargo +1.86 clippy -- -D warnings`, and `cargo +1.86 test
|
||||
--tests` across all 5 freemkv crates. Mirrors each repo's
|
||||
`.github/workflows/ci.yml` step-for-step. Use it as a pre-commit hook
|
||||
@@ -573,7 +1018,7 @@ Through the entire 0.13.x line, every CHECK CONDITION reply from the
|
||||
drive (the standard way SCSI reports a sector failure) was being
|
||||
collapsed into a synthetic `status=0xFF, sense_key=0` "transport
|
||||
wedge" sentinel and the real sense data was thrown away. Live tracing
|
||||
on the BU40N reading Dune 2 on 2026-04-27 confirmed it: the drive was
|
||||
on the BU40N reading a UHD disc on 2026-04-27 confirmed it: the drive was
|
||||
returning `host_status=0, driver_status=8, status=2, exec_elapsed_ms=1416`
|
||||
on every bad sector — a clean CHECK CONDITION carrying full sense
|
||||
data — and the library was misclassifying it as a wedge and bailing.
|
||||
@@ -645,7 +1090,7 @@ recovered or marked Unreadable. Calibration data
|
||||
|
||||
### Replace bisect-on-fail with hysteresis state machine (Block ↔ Single)
|
||||
|
||||
Live test on Dune 2 v0.13.21 showed bisect-on-fail recovered every
|
||||
Live test on a UHD disc v0.13.21 showed bisect-on-fail recovered every
|
||||
recoverable sector, but spent ~30 sec per damaged 60-block (paying a
|
||||
~5 sec kernel timeout at every bisection level). Each level descended
|
||||
log₂(60) ≈ 6 times on the failing branch.
|
||||
@@ -674,7 +1119,7 @@ Per-block cost on a 60-sector damaged block with 1 bad sector:
|
||||
|
||||
Plus inside a damaged cluster spanning many 60-blocks, hysteresis
|
||||
pays the bpt=batch fail cost ONCE on entry; bisection paid it every
|
||||
60 sectors. For Dune 2's ~1248-sector boundary cluster that's ~21
|
||||
60 sectors. For the ~1248-sector boundary cluster that's ~21
|
||||
fewer 5-sec waits = ~100 sec saved.
|
||||
|
||||
Telemetry: new `phase=mode_change` trace event with `from`, `to`,
|
||||
@@ -692,8 +1137,7 @@ synthetic BU40N-pattern reader; same 100% recovery expectation.
|
||||
### Fix: Disc::copy bisect-on-fail (replaces skip-forward)
|
||||
|
||||
Empirical live-hardware testing on the LG BU40N (see
|
||||
`(internal)/docs/audits/2026-04-26-test-plan-audit.md` and the
|
||||
TEST_PLAN.md run log) revealed that the drive often **fails
|
||||
the internal test-plan audit and run log) revealed that the drive often **fails
|
||||
multi-sector READ commands** in damaged regions but **succeeds when
|
||||
asked one sector at a time**. The old skip-forward strategy responded
|
||||
to multi-sector failures by jumping up to 1 % of the disc forward,
|
||||
@@ -705,7 +1149,7 @@ half, retry each half, recurse down to single-sector reads. Sectors
|
||||
the drive can read individually are recovered in Pass 1; only sectors
|
||||
that fail at bpt=1 are marked NonTrimmed for the patch passes.
|
||||
|
||||
Empirical results on Dune 2 UHD on the BU40N:
|
||||
Empirical results on a UHD disc on the BU40N:
|
||||
- Old algorithm: 25 GB read in Pass 1, then ~6 GB skip-forwarded;
|
||||
retry passes failed to recover most of the skipped zone.
|
||||
- New algorithm: ~99 % of disc recovered in Pass 1; only the truly
|
||||
@@ -747,7 +1191,7 @@ LUN/BUS/HOST reset. `READ_RECOVERY_TIMEOUT_MS` (60 s) unchanged.
|
||||
`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)/docs/audits/2026-04-26-scsi-architecture-research.md`,
|
||||
internal SCSI-architecture research,
|
||||
no reference project (MakeMKV / sg_dd / ddrescue) 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 /
|
||||
@@ -874,7 +1318,7 @@ clean reads (Finished sectors) and freezes during skip-forward bad zones,
|
||||
which made every previous version's UI look hung at the bad-zone boundary.
|
||||
This was the v0.13.9 stall-guard origin bug.
|
||||
|
||||
Live trace from v0.13.14: Pass 1 hit a Dune 2 bad zone at 24 GB and
|
||||
Live trace from v0.13.14: Pass 1 hit a bad zone at 24 GB and
|
||||
appeared "stuck" for 14 minutes per autorip's UI (`bytes_good = 23.97 GB`
|
||||
unchanged). Disc trace events showed `pos` actually advanced from 25.8 GB
|
||||
to 70 GB during that window — Pass 1 was 83 % through the disc, marking
|
||||
@@ -922,7 +1366,7 @@ trace events were silently dropped by the default `libfreemkv=warn` rule.
|
||||
|
||||
### Telemetry: instrument the rip pipeline for in-flight diagnosis
|
||||
|
||||
v0.13.12 shipped Fix 1+2+4 + cross-platform parity but a live test on Dune 2
|
||||
v0.13.12 shipped Fix 1+2+4 + cross-platform parity but a live test
|
||||
showed Pass 1 sat for 14 minutes with `bytes_good=0` while the inner loop
|
||||
appeared to iterate (the throttled `on_progress` log fired every 78s). The
|
||||
async fd_recovery design at §7 said each `execute()` call should bound at
|
||||
@@ -966,7 +1410,7 @@ The v0.13.9 stall guard at `disc/mod.rs` exited Pass 1 early when
|
||||
`bytes_good` was flat for `stall_secs` (default 120s). This violated the
|
||||
ddrescue model: Pass 1 must sweep end-to-end, marking failed reads
|
||||
NonTrimmed for Pass 2 retry. The guard caused Pass 1 to bail at 30% on
|
||||
Dune 2 with 56 GB still NonTried, leaving Pass 2 nothing useful to do.
|
||||
a UHD disc with 56 GB still NonTried, leaving Pass 2 nothing useful to do.
|
||||
|
||||
- Deleted the stall-guard state vars and the `if cur_good != ...
|
||||
break 'outer;` block.
|
||||
@@ -1021,7 +1465,7 @@ same observable recovery contract as Linux:
|
||||
|
||||
`PatchResult` now reports `blocks_attempted`, `blocks_read_ok`,
|
||||
`blocks_read_failed`. Pass 2's "100 minutes recovered 0 bytes" mystery
|
||||
(Dune 2) becomes diagnosable from these counters: distinguish "drive
|
||||
(a UHD disc) becomes diagnosable from these counters: distinguish "drive
|
||||
returned Ok but write/record dropped data" from "every read was Err for
|
||||
the entire range" without instrumenting from outside the lib.
|
||||
|
||||
@@ -1050,7 +1494,7 @@ against the spawned close()). The intent was to escape the 60-s
|
||||
blocking reopen.
|
||||
|
||||
The cost was too high: a single transient poll timeout permanently
|
||||
killed the transport. Live test on Dune 2 (post-replug):
|
||||
killed the transport. Live test on a UHD disc (post-replug):
|
||||
- Pass 1 ran for **45 ms** then returned with 0 GB good and 80 GB
|
||||
pending.
|
||||
- The first SCSI READ timed out, fd went to -1, every subsequent
|
||||
@@ -1083,7 +1527,7 @@ Sync bump for the autorip-side fix in 0.13.10 (Pass 1 batch reporting).
|
||||
### Fix: Disc::copy silent stall + SgIoTransport reopen-after-timeout serialization
|
||||
|
||||
Two correlated fixes for a hang observed live on the LG BU40N during a
|
||||
v0.13.8 rip of Dune: Part Two. At ~30 % progress through Pass 1
|
||||
v0.13.8 rip of a UHD disc. At ~30 % progress through Pass 1
|
||||
(disc → ISO), `bytes_good` froze for 10+ minutes with `errs=0`,
|
||||
no error surfaced, drive not wedged.
|
||||
|
||||
@@ -1970,10 +2414,9 @@ Rip recovery rewritten. The old binary-search-per-bad-sector model paid the full
|
||||
- **Kernel transfer limit detection**: auto-detect `max_hw_sectors_kb` via sysfs, resolve sg→block device. Previously hardcoded to 510 sectors (1MB) which exceeded the 120KB kernel limit, causing all reads to error and fall back to 6KB reads at 4.8 MB/s. Now auto-tunes to 48 sectors (96KB) or whatever the device supports.
|
||||
- **Result: 12.5 MB/s sustained, 23 MB/s peak** (was 4.8 MB/s)
|
||||
|
||||
### LibreDrive — full init pipeline
|
||||
### Drive init pipeline
|
||||
|
||||
- **All 10 ARM handlers translated**: unlock, firmware upload (A: WRITE_BUFFER, B: MODE SELECT), calibrate (256 zones), register reads, status, probe, set_read_speed, keepalive, timing
|
||||
- **Cold boot firmware upload**: WRITE_BUFFER 1888B (A variant) or MODE SELECT 2496B (B variant) proven on hardware
|
||||
- **Full drive init pipeline**: unlock, calibrate (256 zones), register reads, status, probe, set_read_speed, keepalive, timing
|
||||
- **Speed calibration**: 256+ disc surface probes, 64-entry speed table, triple SET_CD_SPEED
|
||||
- **Platform trait locked down**: `pub(crate)`, 3 methods only (init, set_read_speed, is_ready)
|
||||
- **Init guard**: prevents double-init, signature mismatch aborts early
|
||||
@@ -1988,7 +2431,7 @@ Rip recovery rewritten. The old binary-search-per-bad-sector model paid the full
|
||||
### Profiles
|
||||
|
||||
- **206 profiles with full per-drive data**: ld_microcode (base64), all CDBs, speed tables, signatures
|
||||
- **Automated pipeline**: `sdf_unpack --profiles` → profiles.json (no manual merging)
|
||||
- **Automated pipeline**: profile import → profiles.json (no manual merging)
|
||||
|
||||
## 0.4.0 (2026-04-07)
|
||||
|
||||
|
||||
+19
-1
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "libfreemkv"
|
||||
version = "0.20.2"
|
||||
version = "0.29.0"
|
||||
edition = "2024"
|
||||
rust-version = "1.86"
|
||||
license = "AGPL-3.0-only"
|
||||
@@ -8,6 +8,8 @@ description = "Open source raw disc access library for optical drives"
|
||||
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 = []
|
||||
|
||||
[dependencies]
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
@@ -28,6 +30,22 @@ base64 = "0.22.1"
|
||||
# under project docs ("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`
|
||||
# so the halt-aware send/finish loops can BLOCK on consumer drain instead of
|
||||
# polling — the 50 ms poll cadence of the previous mpsc-based impl capped mux
|
||||
# throughput at ~1 MB/s (0.21.7).
|
||||
crossbeam-channel = "0.5"
|
||||
# Persistent work-stealing thread pool for parallel AACS unit
|
||||
# decryption. Per-call std::thread::scope spawned fresh OS threads
|
||||
# and that overhead dominated for typical batch sizes (60 units).
|
||||
# rayon's global pool initialises once on first use.
|
||||
rayon = "1"
|
||||
# SIMD-accelerated bytestring search. Drives the HEVC/H.264 start-code
|
||||
# scan in `mux::codec::h264::find_start_code` — naive byte-by-byte
|
||||
# walk is ~500 MB/s single-thread on x86_64; memchr's vectorised
|
||||
# `memmem::find` for the 3-byte `00 00 01` needle hits ~5 GB/s on
|
||||
# AVX2-capable hosts.
|
||||
memchr = "2"
|
||||
|
||||
[target.'cfg(target_os = "linux")'.dependencies]
|
||||
libc = "0.2"
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
# libfreemkv — local dev helper.
|
||||
# Mirrors the cross-crate scripts in (internal)/scripts/test-all.sh
|
||||
# but scoped to this single crate.
|
||||
# Mirrors the workspace-wide CI checks but scoped to this single crate.
|
||||
|
||||
.PHONY: test build check ci clean
|
||||
|
||||
|
||||
@@ -6,6 +6,8 @@
|
||||
|
||||
Rust library for 4K UHD / Blu-ray / DVD optical drives. Drive access, disc scanning, stream labels, AACS decryption, CSS decryption, KEYDB updates, and content reading in one crate. Bundled drive profiles — no external files needed.
|
||||
|
||||
Built-in keys cover DVDs and Blu-rays (AACS 1.0). For UHD (AACS 2.0 / 2.1) discs, an optional `keydb.cfg` supplies disc-specific volume unique keys.
|
||||
|
||||
**12+ MB/s** sustained read speeds on BD. Full init: unlock, firmware upload, speed calibration — all from pure Rust.
|
||||
|
||||
Multi-lingual by design — the library outputs structured data and numeric error codes, never English text. Build any UI or localization on top.
|
||||
@@ -18,7 +20,7 @@ Part of the [freemkv](https://github.com/freemkv) project.
|
||||
|
||||
```toml
|
||||
[dependencies]
|
||||
libfreemkv = "0.18"
|
||||
libfreemkv = "0.25"
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
@@ -56,8 +58,7 @@ output.finish()?;
|
||||
For damaged discs the library exposes two flat verbs — `Disc::sweep` for the
|
||||
forward Pass 1 and `Disc::patch` for retrying bad ranges. The library never
|
||||
loops; the multipass policy is the caller's job. See
|
||||
[`docs/rip-recovery.md`](docs/rip-recovery.md) and the design notes in
|
||||
`(internal)/memory/0_18_redesign.md`.
|
||||
[`docs/rip-recovery.md`](docs/rip-recovery.md).
|
||||
|
||||
```rust
|
||||
use libfreemkv::{SweepOptions, PatchOptions};
|
||||
@@ -122,7 +123,11 @@ loop {
|
||||
|
||||
Streams implement `FrameSource` (read) and/or `FrameSink` (write); direction is type-checked. `input()` / `output()` resolve URL strings to PES stream instances. All URLs use the `scheme://path` format — bare paths are rejected.
|
||||
|
||||
AACS decryption requires a KEYDB.cfg file. If available at `~/.config/aacs/KEYDB.cfg` or passed via `ScanOptions`, the library handles everything — handshake, key derivation, and per-sector decryption — without the application needing to know anything about encryption.
|
||||
### Keys
|
||||
|
||||
DVDs (CSS) decrypt out of the box — the 1999-era public player keys are compiled into the library.
|
||||
|
||||
Blu-rays and UHD (AACS) require a `keydb.cfg` at `~/.config/freemkv/keydb.cfg` (or passed via `ScanOptions`). The file holds all DKs, PKs, host certs, and per-disc VUKs. No AACS key material is compiled into the binary.
|
||||
|
||||
## Architecture
|
||||
|
||||
|
||||
+4
-4
@@ -85,10 +85,10 @@ All URLs require a `scheme://path` format. Bare paths are rejected.
|
||||
// PES pipeline (frame-level) — input() returns Box<dyn FrameSource>,
|
||||
// output() returns Box<dyn FrameSink>.
|
||||
let input = libfreemkv::input("disc:///dev/sg4", &opts)?; // DiscStream
|
||||
let input = libfreemkv::input("iso://Dune.iso", &opts)?; // IsoStream
|
||||
let output = libfreemkv::output("mkv://Dune.mkv", &title)?; // MkvOutputStream
|
||||
let output = libfreemkv::output("m2ts://Dune.m2ts", &title)?; // M2tsOutputStream
|
||||
let output = libfreemkv::output("network://10.0.0.1:9000", &title)?; // NetworkOutputStream
|
||||
let input = libfreemkv::input("iso://Movie.iso", &opts)?; // IsoStream
|
||||
let output = libfreemkv::output("mkv://Movie.mkv", &title)?; // MkvOutputStream
|
||||
let output = libfreemkv::output("m2ts://Movie.m2ts", &title)?; // M2tsOutputStream
|
||||
let output = libfreemkv::output("network://192.0.2.10:9000", &title)?; // NetworkOutputStream
|
||||
let output = libfreemkv::output("null://", &title)?; // NullOutputStream
|
||||
```
|
||||
|
||||
|
||||
+2464
-206
File diff suppressed because it is too large
Load Diff
+50
-24
@@ -69,16 +69,29 @@ pub(crate) fn aes_cbc_decrypt(key: &[u8; 16], data: &mut [u8]) {
|
||||
|
||||
// ── Content decryption ──────────────────────────────────────────────────────
|
||||
|
||||
/// Check if a 6144-byte aligned unit is encrypted (copy_permission_indicator bits).
|
||||
pub fn is_unit_encrypted(unit: &[u8]) -> bool {
|
||||
unit.len() >= ALIGNED_UNIT_LEN && (unit[0] & 0xC0) != 0
|
||||
/// True if a 6144-byte aligned unit is AACS-scrambled on disc.
|
||||
///
|
||||
/// AACS encrypts the unit body, which destroys the MPEG-TS sync bytes (`0x47`)
|
||||
/// a clear unit carries at offsets 4, 196, 388, … (one per 192-byte source
|
||||
/// packet). So "scrambled" = "the TS syncs are NOT intact". This is
|
||||
/// flag-independent: it does NOT read the TP_extra copy-control bits (byte 0)
|
||||
/// or the TS scrambling-control bits (byte 7) — AACS sets neither reliably
|
||||
/// across discs/players.
|
||||
///
|
||||
/// This is the single shared definition of "encrypted" for the whole ecosystem
|
||||
/// — libfreemkv's decrypt gate, autorip's sample selection, and the online key
|
||||
/// service's validation gate all call THIS, so they always agree on what is
|
||||
/// encrypted. A correctly-decrypted (or natively-clear) unit reports `false`,
|
||||
/// so the decrypt path never double-decrypts and there is no flag to clear.
|
||||
pub fn is_aacs_scrambled(unit: &[u8]) -> bool {
|
||||
unit.len() >= ALIGNED_UNIT_LEN && !ts_syncs_intact(unit)
|
||||
}
|
||||
|
||||
/// Verify decrypted unit by checking TS sync bytes at expected offsets.
|
||||
fn verify_ts(unit: &[u8]) -> bool {
|
||||
// In a 6144-byte unit, TS packets start at byte 0 with 4-byte TP_extra_header
|
||||
// then 188-byte TS packet, repeating every 192 bytes.
|
||||
// Sync byte 0x47 should appear at offset 4, 196, 388, ...
|
||||
/// 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).
|
||||
fn ts_syncs_intact(unit: &[u8]) -> bool {
|
||||
let mut count = 0;
|
||||
let mut offset = 4;
|
||||
while offset < unit.len() {
|
||||
@@ -87,11 +100,20 @@ fn verify_ts(unit: &[u8]) -> bool {
|
||||
}
|
||||
offset += TS_PACKET_LEN;
|
||||
}
|
||||
// Expect at least most packets to have sync bytes
|
||||
let total = (unit.len() - 4) / TS_PACKET_LEN + 1;
|
||||
// One sync byte is checked per 192-byte BD-TS packet (at offset 4 of
|
||||
// each). `total` is exactly that packet count; the old
|
||||
// `(len - 4) / TS_PACKET_LEN + 1` over-counted by one for lengths of
|
||||
// the form `4 + k·192` (harmless for the always-6144 aligned unit, but
|
||||
// wrong in general and it biased the majority threshold).
|
||||
let total = unit.len() / TS_PACKET_LEN;
|
||||
count > total / 2
|
||||
}
|
||||
|
||||
/// Verify a decrypted unit looks like clear MPEG-TS (sync bytes intact).
|
||||
fn verify_ts(unit: &[u8]) -> bool {
|
||||
ts_syncs_intact(unit)
|
||||
}
|
||||
|
||||
/// Decrypt one AACS aligned unit (6144 bytes) in-place.
|
||||
/// Returns true if decryption succeeded (verified by TS sync bytes).
|
||||
///
|
||||
@@ -99,12 +121,14 @@ fn verify_ts(unit: &[u8]) -> bool {
|
||||
/// 1. AES-128-ECB encrypt first 16 bytes with unit_key → derived
|
||||
/// 2. XOR derived with original 16 bytes → unit_decrypt_key
|
||||
/// 3. AES-128-CBC decrypt bytes 16..6143 with unit_decrypt_key and AACS IV
|
||||
/// 4. Clear encryption flag bits
|
||||
///
|
||||
/// Decryption restores the TS sync bytes, so the unit reads as clear afterward;
|
||||
/// there is no flag to clear.
|
||||
pub fn decrypt_unit(unit: &mut [u8], unit_key: &[u8; 16]) -> bool {
|
||||
if unit.len() < ALIGNED_UNIT_LEN {
|
||||
return false;
|
||||
}
|
||||
if !is_unit_encrypted(unit) {
|
||||
if !is_aacs_scrambled(unit) {
|
||||
return true; // not encrypted
|
||||
}
|
||||
|
||||
@@ -124,16 +148,13 @@ pub fn decrypt_unit(unit: &mut [u8], unit_key: &[u8; 16]) -> bool {
|
||||
// Step 3: Decrypt bytes 16..6143 with AES-CBC
|
||||
aes_cbc_decrypt(&decrypt_key, &mut unit[16..ALIGNED_UNIT_LEN]);
|
||||
|
||||
// Step 4: Clear encryption flag
|
||||
unit[0] &= !0xC0;
|
||||
|
||||
// Verify
|
||||
// Decryption restored the TS syncs; verify the unit now looks like clear TS.
|
||||
verify_ts(unit)
|
||||
}
|
||||
|
||||
/// 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<usize> {
|
||||
if !is_unit_encrypted(unit) {
|
||||
if !is_aacs_scrambled(unit) {
|
||||
return Some(0);
|
||||
}
|
||||
|
||||
@@ -173,7 +194,7 @@ pub fn decrypt_unit_full(
|
||||
unit_key: &[u8; 16],
|
||||
read_data_key: Option<&[u8; 16]>,
|
||||
) -> bool {
|
||||
if !is_unit_encrypted(unit) {
|
||||
if !is_aacs_scrambled(unit) {
|
||||
return true;
|
||||
}
|
||||
if let Some(rdk) = read_data_key {
|
||||
@@ -200,10 +221,15 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_decrypt_unit_unencrypted() {
|
||||
// Unit with 0xC0 bits clear should pass through unchanged
|
||||
// A clear unit (TS syncs intact) is not scrambled → passes through.
|
||||
let mut unit = vec![0u8; ALIGNED_UNIT_LEN];
|
||||
unit[0] = 0x00; // not encrypted
|
||||
let mut off = 4;
|
||||
while off < ALIGNED_UNIT_LEN {
|
||||
unit[off] = TS_SYNC;
|
||||
off += TS_PACKET_LEN;
|
||||
}
|
||||
let key = [0u8; 16];
|
||||
assert!(!is_aacs_scrambled(&unit));
|
||||
assert!(decrypt_unit(&mut unit, &key));
|
||||
}
|
||||
|
||||
@@ -253,8 +279,8 @@ mod tests {
|
||||
plain[offset] = TS_SYNC;
|
||||
offset += TS_PACKET_LEN;
|
||||
}
|
||||
// Set encryption flag
|
||||
plain[0] |= 0xC0;
|
||||
// No flag set: CBC-encrypting the body below scrambles packets 1..31's
|
||||
// TS syncs, which is exactly what `is_aacs_scrambled` (raw-sync) detects.
|
||||
|
||||
// Now encrypt bytes 16..6143 using the AACS algorithm (reverse of decrypt)
|
||||
let header: [u8; 16] = plain[..16].try_into().unwrap();
|
||||
@@ -281,9 +307,9 @@ mod tests {
|
||||
|
||||
// Now plain contains encrypted data. Decrypt it.
|
||||
let mut unit = plain;
|
||||
assert!(is_unit_encrypted(&unit));
|
||||
assert!(is_aacs_scrambled(&unit));
|
||||
assert!(decrypt_unit(&mut unit, &unit_key));
|
||||
assert!(!is_unit_encrypted(&unit)); // flag should be cleared
|
||||
assert!(!is_aacs_scrambled(&unit)); // decrypted: TS syncs restored
|
||||
|
||||
// Verify TS sync bytes
|
||||
let mut count = 0;
|
||||
|
||||
+89
-2
@@ -1235,8 +1235,8 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_aes_cmac() {
|
||||
// Basic CMAC test — at minimum verify it produces consistent output
|
||||
fn test_aes_cmac_deterministic() {
|
||||
// Same (data, key) must always produce the same MAC.
|
||||
let key = [
|
||||
0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
|
||||
0x4f, 0x3c,
|
||||
@@ -1248,6 +1248,93 @@ mod tests {
|
||||
assert_ne!(mac1, [0u8; 16]); // shouldn't be all zeros
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_aes_cmac_nist_kat_full_block() {
|
||||
// NIST SP 800-38B Appendix D.1, Example 2 (Mlen = 128):
|
||||
// K = 2b7e1516 28aed2a6 abf71588 09cf4f3c
|
||||
// M = 6bc1bee2 2e409f96 e93d7e11 7393172a
|
||||
// T = 070a16b4 6b4d4144 f79bdd9d d04a287c
|
||||
let key = [
|
||||
0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
|
||||
0x4f, 0x3c,
|
||||
];
|
||||
let data = [
|
||||
0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96, 0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93,
|
||||
0x17, 0x2a,
|
||||
];
|
||||
let expected = [
|
||||
0x07, 0x0a, 0x16, 0xb4, 0x6b, 0x4d, 0x41, 0x44, 0xf7, 0x9b, 0xdd, 0x9d, 0xd0, 0x4a,
|
||||
0x28, 0x7c,
|
||||
];
|
||||
let mac = aes_cmac_16(&data, &key);
|
||||
assert_eq!(mac, expected, "AES-CMAC-128 must match NIST SP 800-38B KAT");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_vid_mac_verify_roundtrip() {
|
||||
// Simulate the drive-side: pick a (bus_key, vid), compute the MAC, and
|
||||
// verify the host-side check accepts it. Then mutate VID and MAC each
|
||||
// in turn and verify both mutations cause a mismatch (the path that
|
||||
// would yield Error::AacsVidMac in read_volume_id).
|
||||
let bus_key = [
|
||||
0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef, 0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54,
|
||||
0x32, 0x10,
|
||||
];
|
||||
let vid = [
|
||||
0xde, 0xad, 0xbe, 0xef, 0xca, 0xfe, 0xba, 0xbe, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66,
|
||||
0x77, 0x88,
|
||||
];
|
||||
|
||||
// Drive returns vid + mac where mac == AES-CMAC-128(bus_key, vid).
|
||||
let drive_mac = aes_cmac_16(&vid, &bus_key);
|
||||
let calc_mac = aes_cmac_16(&vid, &bus_key);
|
||||
assert_eq!(calc_mac, drive_mac, "honest drive: MACs must match");
|
||||
|
||||
// Mutate the MAC: a malicious drive that swapped VID but returned its
|
||||
// original MAC would produce a mismatch here.
|
||||
let mut bad_mac = drive_mac;
|
||||
bad_mac[0] ^= 0x01;
|
||||
assert_ne!(calc_mac, bad_mac, "mutated MAC must be rejected");
|
||||
|
||||
// Mutate the VID: even one bit of VID drift produces a wildly different
|
||||
// CMAC (this is what catches a substituted VID with a stale MAC).
|
||||
let mut bad_vid = vid;
|
||||
bad_vid[15] ^= 0x01;
|
||||
let calc_for_bad_vid = aes_cmac_16(&bad_vid, &bus_key);
|
||||
assert_ne!(
|
||||
calc_for_bad_vid, drive_mac,
|
||||
"MAC over mutated VID must not match original MAC"
|
||||
);
|
||||
|
||||
// Wrong bus key (e.g. handshake replayed against the wrong session)
|
||||
// also produces a different MAC over the same VID.
|
||||
let mut wrong_key = bus_key;
|
||||
wrong_key[0] ^= 0xff;
|
||||
let calc_with_wrong_key = aes_cmac_16(&vid, &wrong_key);
|
||||
assert_ne!(
|
||||
calc_with_wrong_key, drive_mac,
|
||||
"MAC under wrong bus key must not match"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_vid_mac_all_zero_mac_rejected() {
|
||||
// Defensive: a buggy or hostile drive that returns all-zero MAC must
|
||||
// be rejected (the real MAC over any non-trivial VID is nearly never
|
||||
// 0...0). This guards against a class of "drive returned garbage"
|
||||
// failures masquerading as success.
|
||||
let bus_key = [
|
||||
0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
|
||||
0x4f, 0x3c,
|
||||
];
|
||||
let vid = [
|
||||
0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96, 0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93,
|
||||
0x17, 0x2a,
|
||||
];
|
||||
let calc_mac = aes_cmac_16(&vid, &bus_key);
|
||||
assert_ne!(calc_mac, [0u8; 16], "real CMAC must not be all zeros");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_verify_host_cert_from_keydb() {
|
||||
// Verify the host cert from our KEYDB
|
||||
|
||||
+162
-17
@@ -89,6 +89,17 @@ pub(crate) fn parse_hex20(s: &str) -> Option<[u8; 20]> {
|
||||
}
|
||||
|
||||
impl KeyDb {
|
||||
/// Construct an empty KeyDb. Used by unit tests; production code
|
||||
/// reaches a populated KeyDb via [`KeyDb::load`] or [`KeyDb::parse`].
|
||||
pub fn empty() -> Self {
|
||||
KeyDb {
|
||||
device_keys: Vec::new(),
|
||||
processing_keys: Vec::new(),
|
||||
host_certs: Vec::new(),
|
||||
disc_entries: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a KEYDB.cfg file from a string.
|
||||
pub fn parse(data: &str) -> Self {
|
||||
let mut db = KeyDb {
|
||||
@@ -106,10 +117,19 @@ impl KeyDb {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Device Key
|
||||
// Device Key.
|
||||
// Two shapes are accepted:
|
||||
// 1. Positioned DK: `| DK | DEVICE_KEY 0x... | DEVICE_NODE 0x... | KEY_UV 0x... | KEY_U_MASK_SHIFT 0x...`
|
||||
// → loaded into `device_keys` (deterministic tree walk via `calc_pk_from_dk`).
|
||||
// 2. Orphan DK: `| DK | DEVICE_KEY 0x...` with no position fields.
|
||||
// → loaded into `processing_keys` (brute walker / terminal validation).
|
||||
// Per AACS spec a "PK" IS a DK at terminal position, so both row types
|
||||
// are DKs in the unified model; only the metadata differs.
|
||||
if line.starts_with("| DK") {
|
||||
if let Some(dk) = Self::parse_device_key(line) {
|
||||
db.device_keys.push(dk);
|
||||
} else if let Some(key) = Self::parse_orphan_dk(line) {
|
||||
db.processing_keys.push(key);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
@@ -152,7 +172,7 @@ impl KeyDb {
|
||||
db
|
||||
}
|
||||
|
||||
/// Load KEYDB.cfg from a file path.
|
||||
/// Load a KEYDB.cfg from disk.
|
||||
pub fn load(path: &std::path::Path) -> std::io::Result<Self> {
|
||||
let data = std::fs::read_to_string(path)?;
|
||||
Ok(Self::parse(&data))
|
||||
@@ -184,8 +204,53 @@ impl KeyDb {
|
||||
.or_else(|| self.disc_entries.get(&hash))
|
||||
}
|
||||
|
||||
// ── Parsers ─────────────────────────────────────────────────────────────
|
||||
/// Iterate every disc entry. Used by Path 3 (scan for matching VID).
|
||||
pub fn iter_disc_entries(&self) -> impl Iterator<Item = &DiscEntry> {
|
||||
self.disc_entries.values()
|
||||
}
|
||||
}
|
||||
|
||||
// ── KeyProvider impl ──────────────────────────────────────────────────────────
|
||||
//
|
||||
// Lets `KeyDb` plug into `resolve_keys` via the trait. Cloning happens in the
|
||||
// bulk methods because the trait returns owned `Vec`s (so HTTP-backed providers
|
||||
// don't need to retain state across calls).
|
||||
|
||||
impl super::provider::KeyProvider for KeyDb {
|
||||
fn device_keys(&self) -> Vec<DeviceKey> {
|
||||
self.device_keys.clone()
|
||||
}
|
||||
fn processing_keys(&self) -> Vec<[u8; 16]> {
|
||||
self.processing_keys.clone()
|
||||
}
|
||||
fn media_keys(&self) -> Vec<[u8; 16]> {
|
||||
// Every per-disc Media Key in the db. The resolver dedups; MKs are
|
||||
// MKB-scoped so the same value recurs across a pressing's discs.
|
||||
self.iter_disc_entries()
|
||||
.filter_map(|e| e.media_key)
|
||||
.collect()
|
||||
}
|
||||
fn host_certs(&self) -> Vec<HostCert> {
|
||||
self.host_certs.clone()
|
||||
}
|
||||
fn lookup_disc_by_hash(&self, disc_hash: &[u8; 20]) -> Option<DiscEntry> {
|
||||
let mut hex = String::with_capacity(42);
|
||||
hex.push_str("0x");
|
||||
for b in disc_hash {
|
||||
hex.push_str(&format!("{b:02X}"));
|
||||
}
|
||||
self.find_disc(&hex).cloned()
|
||||
}
|
||||
fn lookup_disc_by_vid(&self, volume_id: &[u8; 16]) -> Option<DiscEntry> {
|
||||
self.iter_disc_entries()
|
||||
.find(|e| matches!(e.disc_id, Some(id) if &id == volume_id))
|
||||
.cloned()
|
||||
}
|
||||
}
|
||||
|
||||
// ── Private parsers (re-open the inherent impl) ─────────────────────────────
|
||||
|
||||
impl KeyDb {
|
||||
fn parse_device_key(line: &str) -> Option<DeviceKey> {
|
||||
// | DK | DEVICE_KEY 0x... | DEVICE_NODE 0x... | KEY_UV 0x... | KEY_U_MASK_SHIFT 0x...
|
||||
let key_str = line.split("DEVICE_KEY").nth(1)?.split('|').next()?.trim();
|
||||
@@ -218,6 +283,30 @@ impl KeyDb {
|
||||
None
|
||||
}
|
||||
|
||||
/// Parse an orphan DK row: a `| DK |` line carrying only the
|
||||
/// `DEVICE_KEY` field (no position metadata). The key is then
|
||||
/// treated like a terminal/unpositioned label by the resolver
|
||||
/// (Path 2's brute walker). Returns `None` if the line carries
|
||||
/// any position field — those are positioned DKs and parsed by
|
||||
/// [`Self::parse_device_key`] instead.
|
||||
fn parse_orphan_dk(line: &str) -> Option<[u8; 16]> {
|
||||
if line.contains("DEVICE_NODE")
|
||||
|| line.contains("KEY_UV")
|
||||
|| line.contains("KEY_U_MASK_SHIFT")
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let key_str = line
|
||||
.split("DEVICE_KEY")
|
||||
.nth(1)?
|
||||
.split('|')
|
||||
.next()?
|
||||
.split(';')
|
||||
.next()?
|
||||
.trim();
|
||||
parse_hex16(key_str)
|
||||
}
|
||||
|
||||
fn parse_host_cert(line: &str) -> Option<HostCert> {
|
||||
// | HC | HOST_PRIV_KEY 0x... | HOST_CERT 0x...
|
||||
let priv_str = line
|
||||
@@ -366,28 +455,84 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_parse_disc_entry() {
|
||||
let line = r#"***REMOVED*** = DUNE_PART_TWO (Dune: Part Two) | D | 2024-04-02 | M | ***REMOVED*** | I | ***REMOVED*** | V | ***REMOVED*** | U | 1-***REMOVED*** ; MKBv77"#;
|
||||
let entry = KeyDb::parse_disc_entry(line).unwrap();
|
||||
assert_eq!(entry.title, "Dune: Part Two");
|
||||
// All-zero placeholders — synthetic; no real key material in code.
|
||||
let z40 = "00".repeat(20);
|
||||
let z32 = "00".repeat(16);
|
||||
let line = format!(
|
||||
"0x{z40} = SAMPLE_FILM (Sample Film) | D | 2024-01-01 | M | 0x{z32} | I | 0x{z32} | V | 0x{z32} | U | 1-0x{z32} ; MKBv77"
|
||||
);
|
||||
let entry = KeyDb::parse_disc_entry(&line).unwrap();
|
||||
assert_eq!(entry.title, "Sample Film");
|
||||
assert!(entry.media_key.is_some());
|
||||
assert!(entry.vuk.is_some());
|
||||
assert_eq!(entry.unit_keys.len(), 1);
|
||||
assert_eq!(entry.unit_keys[0].0, 1);
|
||||
}
|
||||
|
||||
// NOTE: key fields below use obvious repeated-byte / zero placeholders
|
||||
// (0x01.., 0x02.., 0x03.., 0x00..). NEVER put real — or real-looking — host,
|
||||
// device, or processing key material in code; these tests exercise the
|
||||
// parser's field-splitting only, not any genuine key.
|
||||
|
||||
#[test]
|
||||
fn test_parse_device_key() {
|
||||
let line = "| DK | DEVICE_KEY ***REMOVED*** | DEVICE_NODE 0x0800 | KEY_UV 0x00000400 | KEY_U_MASK_SHIFT 0x17 ; MKBv01-MKBv48";
|
||||
let line = "| DK | DEVICE_KEY 0x00000000000000000000000000000000 | DEVICE_NODE 0x0800 | KEY_UV 0x00000400 | KEY_U_MASK_SHIFT 0x17 ; MKBv01-MKBv48";
|
||||
let dk = KeyDb::parse_device_key(line).unwrap();
|
||||
assert_eq!(dk.node, 0x0800);
|
||||
assert_eq!(dk.u_mask_shift, 0x17);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_orphan_dk_row_loads_into_processing_keys() {
|
||||
// `| DK |` row without position fields = an orphan DK. Per the
|
||||
// unified model the resolver treats it like a terminal/PK
|
||||
// candidate: it lands in `processing_keys` and the brute walker
|
||||
// handles it.
|
||||
let cfg = r#"
|
||||
| DK | DEVICE_KEY 0x01010101010101010101010101010101 ; orphan, no position fields
|
||||
| DK | DEVICE_KEY 0x02020202020202020202020202020202 | DEVICE_NODE 0x0800 | KEY_UV 0x00000400 | KEY_U_MASK_SHIFT 0x17 ; positioned MKBv01-MKBv48
|
||||
| PK | 0x03030303030303030303030303030303 ; legacy PK row still works
|
||||
"#;
|
||||
let db = KeyDb::parse(cfg);
|
||||
assert_eq!(
|
||||
db.device_keys.len(),
|
||||
1,
|
||||
"positioned DK row should land in device_keys"
|
||||
);
|
||||
// Orphan DK + legacy PK row both end up in processing_keys.
|
||||
assert_eq!(
|
||||
db.processing_keys.len(),
|
||||
2,
|
||||
"orphan DK row + legacy PK row both belong in processing_keys"
|
||||
);
|
||||
assert_eq!(db.processing_keys[0][..4], [0x01, 0x01, 0x01, 0x01]);
|
||||
assert_eq!(db.processing_keys[1][..4], [0x03, 0x03, 0x03, 0x03]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_orphan_dk_rejects_lines_with_position_fields() {
|
||||
// The parser must NOT pick up a positioned DK row as an orphan
|
||||
// (that would double-count). parse_orphan_dk explicitly checks.
|
||||
let positioned = "| DK | DEVICE_KEY 0x02020202020202020202020202020202 | DEVICE_NODE 0x0800 | KEY_UV 0x00000400 | KEY_U_MASK_SHIFT 0x17";
|
||||
assert!(
|
||||
KeyDb::parse_orphan_dk(positioned).is_none(),
|
||||
"positioned DK must not match orphan parser"
|
||||
);
|
||||
let orphan = "| DK | DEVICE_KEY 0x01010101010101010101010101010101";
|
||||
let key = KeyDb::parse_orphan_dk(orphan).expect("orphan should parse");
|
||||
assert_eq!(key[..4], [0x01, 0x01, 0x01, 0x01]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_host_cert() {
|
||||
let line = "| HC | HOST_PRIV_KEY ***REMOVED*** | HOST_CERT ***REMOVED*** ; Revoked";
|
||||
let hc = KeyDb::parse_host_cert(line).unwrap();
|
||||
assert_eq!(hc.private_key[0], 0x90);
|
||||
// 20-byte priv + 92-byte cert, all zeros — placeholders, not a key.
|
||||
let line = format!(
|
||||
"| HC | HOST_PRIV_KEY 0x{} | HOST_CERT 0x{} ; Revoked",
|
||||
"00".repeat(20),
|
||||
"00".repeat(92)
|
||||
);
|
||||
let hc = KeyDb::parse_host_cert(&line).unwrap();
|
||||
assert_eq!(hc.private_key, [0u8; 20]);
|
||||
assert_eq!(hc.certificate.len(), 92);
|
||||
}
|
||||
|
||||
@@ -405,15 +550,15 @@ mod tests {
|
||||
assert!(!db.host_certs.is_empty());
|
||||
assert!(db.disc_entries.len() > 170000);
|
||||
|
||||
// Look up Dune: Part Two
|
||||
let dune = db
|
||||
// Look up any disc entry carrying a full key set.
|
||||
let entry = db
|
||||
.disc_entries
|
||||
.values()
|
||||
.find(|e| e.title.contains("Dune: Part Two") && e.vuk.is_some())
|
||||
.expect("Dune: Part Two not found");
|
||||
assert!(dune.media_key.is_some());
|
||||
assert!(dune.vuk.is_some());
|
||||
assert!(!dune.unit_keys.is_empty());
|
||||
.find(|e| e.vuk.is_some() && e.media_key.is_some() && !e.unit_keys.is_empty())
|
||||
.expect("no disc entry with a full key set");
|
||||
assert!(entry.media_key.is_some());
|
||||
assert!(entry.vuk.is_some());
|
||||
assert!(!entry.unit_keys.is_empty());
|
||||
|
||||
eprintln!(
|
||||
"Parsed {} disc entries, {} DK, {} PK",
|
||||
|
||||
+1241
-188
File diff suppressed because it is too large
Load Diff
+14
-4
@@ -17,16 +17,26 @@ pub mod decrypt;
|
||||
pub mod handshake;
|
||||
pub mod keydb;
|
||||
pub mod keys;
|
||||
pub mod provider;
|
||||
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_unit_encrypted,
|
||||
is_aacs_scrambled,
|
||||
};
|
||||
pub use keydb::{DeviceKey, DiscEntry, HostCert, KeyDb};
|
||||
pub use keys::probe;
|
||||
pub use keys::{
|
||||
ContentCert, ResolvedKeys, UnitKeyFile, decrypt_unit_key, derive_media_key_from_dk,
|
||||
derive_media_key_from_pk, derive_vuk, disc_hash, disc_hash_hex, mkb_version,
|
||||
parse_content_cert, parse_unit_key_ro, read_mkb_from_drive, resolve_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,
|
||||
};
|
||||
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,
|
||||
};
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
//! Key source abstraction for the AACS resolve chain.
|
||||
//!
|
||||
//! libfreemkv keeps all crypto (AES-G primitives, SD-tree walking,
|
||||
//! validation, MK/VUK/TK derivation) but accepts key material from
|
||||
//! arbitrary backends via [`KeyProvider`].
|
||||
//!
|
||||
//! Methods come in two flavors:
|
||||
//!
|
||||
//! - **Bulk material** ([`device_keys`], [`processing_keys`],
|
||||
//! [`host_certs`]) — the resolver unions 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.
|
||||
//!
|
||||
//! 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.
|
||||
//!
|
||||
//! Calls may block (disk I/O, network round-trips). The resolver
|
||||
//! invokes each method at most a handful of times per scan; for
|
||||
//! per-disc memoization, implementations should cache internally.
|
||||
//!
|
||||
//! [`device_keys`]: KeyProvider::device_keys
|
||||
//! [`processing_keys`]: KeyProvider::processing_keys
|
||||
//! [`host_certs`]: KeyProvider::host_certs
|
||||
//! [`lookup_disc_by_hash`]: KeyProvider::lookup_disc_by_hash
|
||||
//! [`lookup_disc_by_vid`]: KeyProvider::lookup_disc_by_vid
|
||||
|
||||
use super::keydb::{DeviceKey, DiscEntry, HostCert};
|
||||
|
||||
/// Source of AACS key material.
|
||||
///
|
||||
/// Implementors return raw material only — the resolver in
|
||||
/// `aacs::keys` owns all the crypto (DK→PK walking, PK validation,
|
||||
/// MK→VUK→TK derivation). See module docs for method semantics.
|
||||
pub trait KeyProvider: Send + Sync {
|
||||
/// Device keys (top-of-tree, walked by the resolver).
|
||||
fn device_keys(&self) -> Vec<DeviceKey> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
/// Processing keys — terminal PKs or walk-input PKs. The
|
||||
/// resolver tries each as a terminal first (cheap validate).
|
||||
fn processing_keys(&self) -> Vec<[u8; 16]> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
/// Every Media Key this provider holds, regardless of which disc it was
|
||||
/// filed under. An MK is MKB-scoped (shared across a pressing/MKB-family),
|
||||
/// so the resolver can verify each against the disc's MKB (`km_verifies`)
|
||||
/// and resolve a disc whose own hash/VID isn't directly keyed.
|
||||
fn media_keys(&self) -> Vec<[u8; 16]> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
/// AACS host certificates (with their private keys) for drive
|
||||
/// authentication. Multiple in case some are revoked.
|
||||
fn host_certs(&self) -> Vec<HostCert> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
/// Direct per-disc lookup by SHA-1 of `Unit_Key_RO.inf`. Returns
|
||||
/// `Some(entry)` if this provider has pre-computed material for
|
||||
/// the disc (paths 4 and 5). Short-circuits the resolver.
|
||||
fn lookup_disc_by_hash(&self, _disc_hash: &[u8; 20]) -> Option<DiscEntry> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Lookup by Volume ID (path 3 — pre-computed MK + matching
|
||||
/// VID). Short-circuits the resolver on hit.
|
||||
fn lookup_disc_by_vid(&self, _volume_id: &[u8; 16]) -> Option<DiscEntry> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
pub(crate) struct Providers<'a>(pub &'a [&'a dyn KeyProvider]);
|
||||
|
||||
impl Providers<'_> {
|
||||
/// Union — gather DKs from every provider.
|
||||
pub fn device_keys(&self) -> Vec<DeviceKey> {
|
||||
self.0.iter().flat_map(|p| p.device_keys()).collect()
|
||||
}
|
||||
|
||||
/// Union — gather PKs from every provider.
|
||||
pub fn processing_keys(&self) -> Vec<[u8; 16]> {
|
||||
self.0.iter().flat_map(|p| p.processing_keys()).collect()
|
||||
}
|
||||
|
||||
/// Union of distinct Media Keys across every provider, for the MK-pool
|
||||
/// brute (`km_verifies` against the disc's MKB).
|
||||
pub fn media_keys(&self) -> Vec<[u8; 16]> {
|
||||
let mut v: Vec<[u8; 16]> = self.0.iter().flat_map(|p| p.media_keys()).collect();
|
||||
v.sort_unstable();
|
||||
v.dedup();
|
||||
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.
|
||||
#[allow(dead_code)]
|
||||
pub fn host_certs(&self) -> Vec<HostCert> {
|
||||
self.0.iter().flat_map(|p| p.host_certs()).collect()
|
||||
}
|
||||
|
||||
/// Short-circuit — query providers in array order, first hit wins.
|
||||
pub fn lookup_disc_by_hash(&self, disc_hash: &[u8; 20]) -> Option<DiscEntry> {
|
||||
self.0.iter().find_map(|p| p.lookup_disc_by_hash(disc_hash))
|
||||
}
|
||||
|
||||
/// Short-circuit — query providers in array order, first hit wins.
|
||||
pub fn lookup_disc_by_vid(&self, volume_id: &[u8; 16]) -> Option<DiscEntry> {
|
||||
self.0.iter().find_map(|p| p.lookup_disc_by_vid(volume_id))
|
||||
}
|
||||
}
|
||||
|
||||
/// A [`KeyProvider`] backed by a single caller-supplied key's raw material —
|
||||
/// the bridge for [`crate::disc::Disc::decrypt_with`].
|
||||
///
|
||||
/// The application's key source did the lookup and handed in material at one
|
||||
/// level (DK / PK / MK / VUK). This exposes exactly that material to the
|
||||
/// version-dispatched resolver, which owns ALL derivation — so a source never
|
||||
/// derives, and the lib remains the single home for the AACS chain across
|
||||
/// 1.0 / 2.0 / 2.1 / 2.x.
|
||||
///
|
||||
/// Each level fills only its own field; the rest stay empty, so the resolver
|
||||
/// naturally runs the matching path (DK→…, PK→…, MK-pool brute, or a
|
||||
/// disc-keyed VUK hit). `decrypt_with` already knows the disc, so the
|
||||
/// `lookup_disc_by_*` hash/VID arguments are irrelevant — a present
|
||||
/// `disc_entry` is returned for any query.
|
||||
pub(crate) struct SuppliedKey {
|
||||
pub device_keys: Vec<DeviceKey>,
|
||||
pub processing_keys: Vec<[u8; 16]>,
|
||||
pub media_keys: Vec<[u8; 16]>,
|
||||
pub disc_entry: Option<DiscEntry>,
|
||||
}
|
||||
|
||||
impl KeyProvider for SuppliedKey {
|
||||
fn device_keys(&self) -> Vec<DeviceKey> {
|
||||
self.device_keys.clone()
|
||||
}
|
||||
fn processing_keys(&self) -> Vec<[u8; 16]> {
|
||||
self.processing_keys.clone()
|
||||
}
|
||||
fn media_keys(&self) -> Vec<[u8; 16]> {
|
||||
self.media_keys.clone()
|
||||
}
|
||||
fn lookup_disc_by_hash(&self, _disc_hash: &[u8; 20]) -> Option<DiscEntry> {
|
||||
self.disc_entry.clone()
|
||||
}
|
||||
fn lookup_disc_by_vid(&self, _volume_id: &[u8; 16]) -> Option<DiscEntry> {
|
||||
self.disc_entry.clone()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,679 @@
|
||||
//! AACS Media Key Variant chain.
|
||||
//!
|
||||
//! On AACS 2.1 the Media Key derivation gains a second stage on top of
|
||||
//! the classical subset-difference walk. The classical walk yields a
|
||||
//! Media Key Precursor (Kmp) rather than the final Media Key; the
|
||||
//! Precursor combines with disc-supplied Variant Key Data (VKD) and an
|
||||
//! integrator-supplied Key Correction Data (KCD) constant to produce
|
||||
//! the Media Key.
|
||||
//!
|
||||
//! This module is wiring only — `resolve_keys` is not aware of it. The
|
||||
//! entry point is [`derive_media_key_variant`]. The Variant scheme is
|
||||
//! detected via the new MKB record types `0x82` (Encrypted Media Key
|
||||
//! Variant Data + Variant Key Data) and `0x83` (Variant Number). When
|
||||
//! a disc carries neither, callers should fall back to the classical
|
||||
//! single-stage derivation in [`super::keys`].
|
||||
//!
|
||||
//! The chain follows the published spec:
|
||||
//!
|
||||
//! ```text
|
||||
//! Kmp = AES-128D(Kp, C) XOR uv
|
||||
//! Kpnew = Kmp XOR KCD
|
||||
//! Kvn = AES-G(Kp, Nonce) & 0xFFFF (low 16 bits, BE)
|
||||
//! VKD_idx = Kvn XOR VARIANTS[uv]
|
||||
//! VKD = vkd_table[VKD_idx * 16 .. +16]
|
||||
//! Km = AES-128D(Kpnew, VKD) XOR uv
|
||||
//! ```
|
||||
//!
|
||||
//! 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.
|
||||
|
||||
use super::decrypt::aes_ecb_decrypt;
|
||||
use super::keydb::DeviceKey;
|
||||
|
||||
// ── Public constants ──────────────────────────────────────────────────────
|
||||
|
||||
/// Placeholder Key Correction Data. Sixteen zero bytes.
|
||||
///
|
||||
/// Integrators MUST supply a non-placeholder KCD via the `kcd` argument
|
||||
/// to [`derive_media_key_variant`]; the chain refuses to operate when
|
||||
/// the supplied KCD compares equal to this placeholder.
|
||||
pub const KEY_CORRECTION_DATA_PLACEHOLDER: [u8; 16] = [0u8; 16];
|
||||
|
||||
// ── MKB record walking ────────────────────────────────────────────────────
|
||||
|
||||
/// A single MKB record produced by [`walk_mkb`].
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MkbRecord {
|
||||
/// Byte offset of the record within the MKB.
|
||||
pub offset: usize,
|
||||
/// Record type byte.
|
||||
pub rec_type: u8,
|
||||
/// Record length in bytes (includes the 4-byte header).
|
||||
pub rec_len: usize,
|
||||
/// Record body (the bytes after the 4-byte header).
|
||||
pub body: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Walk an MKB into a flat list of records.
|
||||
///
|
||||
/// MKB record framing per AACS: 1 byte type, 3 bytes BE length
|
||||
/// INCLUDING the 4-byte header, followed by payload. The walker stops
|
||||
/// at the first `(type=0, len=0)` end marker or at end of buffer.
|
||||
pub fn walk_mkb(mkb: &[u8]) -> Vec<MkbRecord> {
|
||||
let mut out = Vec::new();
|
||||
let mut pos = 0;
|
||||
while pos + 4 <= mkb.len() {
|
||||
let rec_type = mkb[pos];
|
||||
let rec_len = ((mkb[pos + 1] as usize) << 16)
|
||||
| ((mkb[pos + 2] as usize) << 8)
|
||||
| (mkb[pos + 3] as usize);
|
||||
if rec_type == 0 && rec_len == 0 {
|
||||
break;
|
||||
}
|
||||
if rec_len < 4 || pos + rec_len > mkb.len() {
|
||||
break;
|
||||
}
|
||||
let body = mkb[pos + 4..pos + rec_len].to_vec();
|
||||
out.push(MkbRecord {
|
||||
offset: pos,
|
||||
rec_type,
|
||||
rec_len,
|
||||
body,
|
||||
});
|
||||
pos += rec_len;
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// True iff `records` contains at least one Media Key Variant record
|
||||
/// (type `0x82` or `0x83`).
|
||||
pub fn is_variant_mkb(records: &[MkbRecord]) -> bool {
|
||||
records.iter().any(|r| matches!(r.rec_type, 0x82 | 0x83))
|
||||
}
|
||||
|
||||
/// Body of the Encrypted Media Key Variant Data record (type `0x82`).
|
||||
pub fn variant_data_record(records: &[MkbRecord]) -> Option<&[u8]> {
|
||||
records
|
||||
.iter()
|
||||
.find(|r| r.rec_type == 0x82)
|
||||
.map(|r| r.body.as_slice())
|
||||
}
|
||||
|
||||
/// 16-byte Nonce from the Variant Number record (type `0x83`). Returns
|
||||
/// the first 16 bytes of the body.
|
||||
pub fn variant_nonce(records: &[MkbRecord]) -> Option<[u8; 16]> {
|
||||
let r = records.iter().find(|r| r.rec_type == 0x83)?;
|
||||
if r.body.len() < 16 {
|
||||
return None;
|
||||
}
|
||||
let mut out = [0u8; 16];
|
||||
out.copy_from_slice(&r.body[..16]);
|
||||
Some(out)
|
||||
}
|
||||
|
||||
/// 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]> {
|
||||
records
|
||||
.iter()
|
||||
.find(|r| r.rec_type == 0x82 && !r.body.is_empty() && r.body.len() % 16 == 0)
|
||||
.map(|r| r.body.as_slice())
|
||||
}
|
||||
|
||||
// ── AES-G ────────────────────────────────────────────────────────────────
|
||||
|
||||
/// AES-G(x1, x2) = AES-128D(x1, x2) XOR x2.
|
||||
///
|
||||
/// The Media Key Variant chain uses AES-G to derive both the variant
|
||||
/// number (`Kvn = AES-G(Kp, Nonce)`) and the Volume Unique Key
|
||||
/// (`Kvu = AES-G(Km, VID)`). See [`super::keys::derive_vuk`] for the
|
||||
/// classical VUK form — the math is identical, this exposes it as a
|
||||
/// neutral primitive for the variant chain.
|
||||
fn aes_g(x1: &[u8; 16], x2: &[u8; 16]) -> [u8; 16] {
|
||||
let mut out = aes_ecb_decrypt(x1, x2);
|
||||
for i in 0..16 {
|
||||
out[i] ^= x2[i];
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
// ── Subset-difference walk that exposes (Kp, uv) ──────────────────────────
|
||||
|
||||
/// AES-G3 seed register initial value.
|
||||
const AESG3_SEED: [u8; 16] = [
|
||||
0x7B, 0x10, 0x3C, 0x5D, 0xCB, 0x08, 0xC4, 0xE5, 0x1A, 0x27, 0xB0, 0x17, 0x99, 0x05, 0x3B, 0xD9,
|
||||
];
|
||||
|
||||
/// AES-G3 single step: AES-G against the seed register at offset `inc`.
|
||||
fn aesg3_step(key: &[u8; 16], inc: u8) -> [u8; 16] {
|
||||
let mut seed = AESG3_SEED;
|
||||
seed[15] = seed[15].wrapping_add(inc);
|
||||
aes_g(key, &seed)
|
||||
}
|
||||
|
||||
fn calc_v_mask(uv: u32) -> u32 {
|
||||
let mut v_mask: u32 = 0xFFFF_FFFF;
|
||||
while (uv & !v_mask) == 0 && v_mask != 0 {
|
||||
v_mask <<= 1;
|
||||
}
|
||||
v_mask
|
||||
}
|
||||
|
||||
fn calc_pk_from_dk(dk: &[u8; 16], uv: u32, v_mask: u32, dev_key_v_mask: u32) -> [u8; 16] {
|
||||
let mut left_child = aesg3_step(dk, 0);
|
||||
let mut pk = aesg3_step(dk, 1);
|
||||
let mut right_child = aesg3_step(dk, 2);
|
||||
let mut current_v_mask = dev_key_v_mask;
|
||||
|
||||
while current_v_mask != v_mask {
|
||||
let mut bit_pos: i32 = -1;
|
||||
for i in (0..32).rev() {
|
||||
if (current_v_mask & (1u32 << i)) == 0 {
|
||||
bit_pos = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let curr_key = if bit_pos < 0 || (uv & (1u32 << bit_pos as u32)) == 0 {
|
||||
left_child
|
||||
} else {
|
||||
right_child
|
||||
};
|
||||
|
||||
left_child = aesg3_step(&curr_key, 0);
|
||||
pk = aesg3_step(&curr_key, 1);
|
||||
right_child = aesg3_step(&curr_key, 2);
|
||||
|
||||
current_v_mask = ((current_v_mask as i32) >> 1) as u32;
|
||||
}
|
||||
|
||||
pk
|
||||
}
|
||||
|
||||
/// Outcome of a subset-difference walk against an MKB. Carries the
|
||||
/// processing key and the matching `uv` slot — both needed as inputs
|
||||
/// to the variant chain.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct ProcessingKeyMatch {
|
||||
/// Processing Key.
|
||||
pub kp: [u8; 16],
|
||||
/// Subset-difference node number that matched.
|
||||
pub uv: u32,
|
||||
/// 16-byte cvalue that the matched uv selected.
|
||||
pub cvalue: [u8; 16],
|
||||
/// Index of the matching cvalue within the cvalues record.
|
||||
pub cvalue_index: usize,
|
||||
}
|
||||
|
||||
fn mkb_find_body(records: &[MkbRecord], rec_type: u8) -> Option<&[u8]> {
|
||||
records
|
||||
.iter()
|
||||
.find(|r| r.rec_type == rec_type && !r.body.is_empty())
|
||||
.map(|r| r.body.as_slice())
|
||||
}
|
||||
|
||||
fn mkb_find_mk_dv(records: &[MkbRecord]) -> Option<[u8; 16]> {
|
||||
let r = records
|
||||
.iter()
|
||||
.find(|r| (r.rec_type == 0x81 || r.rec_type == 0x86) && r.body.len() >= 16)?;
|
||||
let mut out = [0u8; 16];
|
||||
out.copy_from_slice(&r.body[..16]);
|
||||
Some(out)
|
||||
}
|
||||
|
||||
/// Walk an MKB and return the first `(Kp, uv, cvalue)` that
|
||||
/// `device_keys` covers. Returns `None` if no DK walks any uv.
|
||||
pub fn walk_processing_key(
|
||||
records: &[MkbRecord],
|
||||
device_keys: &[DeviceKey],
|
||||
) -> Option<ProcessingKeyMatch> {
|
||||
let mk_dv = mkb_find_mk_dv(records)?;
|
||||
let uvs = mkb_find_body(records, 0x04)?;
|
||||
let cvalues = mkb_find_body(records, 0x07).or_else(|| mkb_find_body(records, 0x05))?;
|
||||
|
||||
let num_uvs = uvs
|
||||
.chunks(5)
|
||||
.take_while(|c| c.len() == 5 && (c[0] & 0xC0) == 0)
|
||||
.count();
|
||||
|
||||
for dk in device_keys {
|
||||
let device_number = dk.node as u32;
|
||||
|
||||
for uvs_idx in 0..num_uvs {
|
||||
let p_uv = &uvs[1 + 5 * uvs_idx..];
|
||||
let u_mask_shift = uvs[5 * uvs_idx];
|
||||
|
||||
if u_mask_shift & 0xC0 != 0 {
|
||||
break;
|
||||
}
|
||||
|
||||
let uv = u32::from_be_bytes([p_uv[0], p_uv[1], p_uv[2], p_uv[3]]);
|
||||
if uv == 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let u_mask: u32 = 0xFFFF_FFFFu32.wrapping_shl(u_mask_shift as u32);
|
||||
let v_mask = calc_v_mask(uv);
|
||||
|
||||
if ((device_number & u_mask) == (uv & u_mask))
|
||||
&& ((device_number & v_mask) != (uv & v_mask))
|
||||
{
|
||||
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);
|
||||
|
||||
if u_mask == dev_key_u_mask && (uv & dev_key_v_mask) == (dk.uv & dev_key_v_mask) {
|
||||
let pk = calc_pk_from_dk(&dk.key, uv, v_mask, dev_key_v_mask);
|
||||
|
||||
if uvs_idx >= cvalues.len() / 16 {
|
||||
continue;
|
||||
}
|
||||
let mut cv = [0u8; 16];
|
||||
cv.copy_from_slice(&cvalues[uvs_idx * 16..(uvs_idx + 1) * 16]);
|
||||
|
||||
// Validate: AES-D(Kp, cv), XOR uv into low 4 bytes,
|
||||
// then AES-D(.., mk_dv) must reveal the verify magic.
|
||||
let mut km_candidate = aes_ecb_decrypt(&pk, &cv);
|
||||
let uv_bytes = uv.to_be_bytes();
|
||||
for i in 0..4 {
|
||||
km_candidate[12 + i] ^= uv_bytes[i];
|
||||
}
|
||||
let dec_vd = aes_ecb_decrypt(&km_candidate, &mk_dv);
|
||||
const VERIFY_MAGIC: [u8; 8] = [0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF];
|
||||
// On a classical (non-variant) MKB this magic must
|
||||
// match. On a variant MKB it won't — `km_candidate`
|
||||
// is really Kmp and the magic check is moot. We
|
||||
// still gate the walk on cvalue indexing being
|
||||
// sane; the chain itself enforces the variant
|
||||
// semantics downstream.
|
||||
let classical_ok = dec_vd[..8] == VERIFY_MAGIC;
|
||||
let variant_present = is_variant_mkb(records);
|
||||
if !(classical_ok || variant_present) {
|
||||
continue;
|
||||
}
|
||||
|
||||
return Some(ProcessingKeyMatch {
|
||||
kp: pk,
|
||||
uv,
|
||||
cvalue: cv,
|
||||
cvalue_index: uvs_idx,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
// ── Error reporting ───────────────────────────────────────────────────────
|
||||
|
||||
/// Outcome of [`derive_media_key_variant`] when the chain cannot
|
||||
/// produce a Media Key. Every variant is a classification only — no
|
||||
/// strings, no Display impl beyond the error code.
|
||||
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
|
||||
pub enum MediaKeyVariantError {
|
||||
/// MKB carries no Variant records. Caller should fall back to the
|
||||
/// classical single-stage derivation.
|
||||
NotVariantMkb,
|
||||
/// MKB is missing a required record (mk_dv, subset-difference,
|
||||
/// cvalues, variant data, or variant nonce).
|
||||
MkbIncomplete,
|
||||
/// `device_keys` did not cover any uv slot in this MKB.
|
||||
ProcessingKeyUnavailable,
|
||||
/// `Kmp[15]` carries bit `0x02`: the soft-correction path applies
|
||||
/// for this Precursor. Out of scope for the hardcoded-KCD chain.
|
||||
SoftCorrectionRequired,
|
||||
/// `Kmp[15]` carries bit `0x04`: the online-challenge path applies
|
||||
/// for this Precursor. Out of scope for the hardcoded-KCD chain.
|
||||
OnlineChallengeRequired,
|
||||
/// Supplied KCD equals [`KEY_CORRECTION_DATA_PLACEHOLDER`]. The
|
||||
/// derivation refuses to run with the all-zero placeholder.
|
||||
KcdNotProvided,
|
||||
/// `VARIANTS[uv]` lookup for the matched uv is not implemented.
|
||||
VariantsTableUnavailable,
|
||||
/// VKD index resolved out of the supplied `vkd_table`.
|
||||
VkdIndexOutOfRange,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for MediaKeyVariantError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
let code: u16 = match self {
|
||||
MediaKeyVariantError::NotVariantMkb => 7100,
|
||||
MediaKeyVariantError::MkbIncomplete => 7101,
|
||||
MediaKeyVariantError::ProcessingKeyUnavailable => 7102,
|
||||
MediaKeyVariantError::SoftCorrectionRequired => 7103,
|
||||
MediaKeyVariantError::OnlineChallengeRequired => 7104,
|
||||
MediaKeyVariantError::KcdNotProvided => 7105,
|
||||
MediaKeyVariantError::VariantsTableUnavailable => 7106,
|
||||
MediaKeyVariantError::VkdIndexOutOfRange => 7107,
|
||||
};
|
||||
write!(f, "E{code}")
|
||||
}
|
||||
}
|
||||
|
||||
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> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Run the Media Key Variant chain on an MKB.
|
||||
///
|
||||
/// Inputs:
|
||||
///
|
||||
/// - `mkb_records` : MKB pre-walked via [`walk_mkb`].
|
||||
/// - `device_keys` : pool of device keys; the chain runs against the
|
||||
/// first uv slot any DK covers.
|
||||
/// - `kcd` : integrator-supplied Key Correction Data. Must not
|
||||
/// equal [`KEY_CORRECTION_DATA_PLACEHOLDER`].
|
||||
/// - `vid` : 16-byte Volume ID for the disc. Used to derive
|
||||
/// the final VUK alongside the Media Key.
|
||||
///
|
||||
/// Returns `(Km, Kvu)` on success.
|
||||
pub fn derive_media_key_variant(
|
||||
mkb_records: &[MkbRecord],
|
||||
device_keys: &[DeviceKey],
|
||||
kcd: &[u8; 16],
|
||||
vid: &[u8; 16],
|
||||
) -> Result<([u8; 16], [u8; 16]), MediaKeyVariantError> {
|
||||
if !is_variant_mkb(mkb_records) {
|
||||
return Err(MediaKeyVariantError::NotVariantMkb);
|
||||
}
|
||||
|
||||
let pkm = walk_processing_key(mkb_records, device_keys)
|
||||
.ok_or(MediaKeyVariantError::ProcessingKeyUnavailable)?;
|
||||
|
||||
let nonce = variant_nonce(mkb_records).ok_or(MediaKeyVariantError::MkbIncomplete)?;
|
||||
let vkd_table = variant_key_data(mkb_records).ok_or(MediaKeyVariantError::MkbIncomplete)?;
|
||||
let c_value = variant_data_record(mkb_records).ok_or(MediaKeyVariantError::MkbIncomplete)?;
|
||||
if c_value.len() < 16 {
|
||||
return Err(MediaKeyVariantError::MkbIncomplete);
|
||||
}
|
||||
let mut c_block = [0u8; 16];
|
||||
c_block.copy_from_slice(&c_value[..16]);
|
||||
|
||||
// Step: Kmp = AES-128D(Kp, C) XOR uv (uv into low 4 bytes).
|
||||
let mut kmp = aes_ecb_decrypt(&pkm.kp, &c_block);
|
||||
let uv_bytes = pkm.uv.to_be_bytes();
|
||||
for i in 0..4 {
|
||||
kmp[12 + i] ^= uv_bytes[i];
|
||||
}
|
||||
|
||||
// Condition bits on Kmp[15] route off the hardcoded-KCD path.
|
||||
if kmp[15] & 0b0000_0010 != 0 {
|
||||
return Err(MediaKeyVariantError::SoftCorrectionRequired);
|
||||
}
|
||||
if kmp[15] & 0b0000_0100 != 0 {
|
||||
return Err(MediaKeyVariantError::OnlineChallengeRequired);
|
||||
}
|
||||
if kcd == &KEY_CORRECTION_DATA_PLACEHOLDER {
|
||||
return Err(MediaKeyVariantError::KcdNotProvided);
|
||||
}
|
||||
|
||||
// Step: Kpnew = Kmp XOR KCD.
|
||||
let mut kpnew = [0u8; 16];
|
||||
for i in 0..16 {
|
||||
kpnew[i] = kmp[i] ^ kcd[i];
|
||||
}
|
||||
|
||||
// Step: Kvn = AES-G(Kp, Nonce) & 0xFFFF (low 16 bits, BE).
|
||||
let kvn_block = aes_g(&pkm.kp, &nonce);
|
||||
let kvn = u16::from_be_bytes([kvn_block[14], kvn_block[15]]);
|
||||
|
||||
// Step: VKD_idx = Kvn XOR VARIANTS[uv].
|
||||
let v_for_uv = variants_for_uv(mkb_records, pkm.cvalue_index)
|
||||
.ok_or(MediaKeyVariantError::VariantsTableUnavailable)?;
|
||||
let vkd_idx = kvn ^ v_for_uv;
|
||||
|
||||
// Step: VKD = vkd_table[VKD_idx * 16 .. +16].
|
||||
let off = (vkd_idx as usize) * 16;
|
||||
if off + 16 > vkd_table.len() {
|
||||
return Err(MediaKeyVariantError::VkdIndexOutOfRange);
|
||||
}
|
||||
let mut vkd = [0u8; 16];
|
||||
vkd.copy_from_slice(&vkd_table[off..off + 16]);
|
||||
|
||||
// Step: Km = AES-128D(Kpnew, VKD) XOR uv.
|
||||
let mut km = aes_ecb_decrypt(&kpnew, &vkd);
|
||||
for i in 0..4 {
|
||||
km[12 + i] ^= uv_bytes[i];
|
||||
}
|
||||
|
||||
// Step: Kvu = AES-G(Km, VID).
|
||||
let kvu = aes_g(&km, vid);
|
||||
|
||||
Ok((km, kvu))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// ── Helpers ──
|
||||
|
||||
fn synthetic_mkb_classical() -> Vec<u8> {
|
||||
// Minimal MKB: type/version record + cvalues + mk_dv. No variant
|
||||
// records.
|
||||
let mut mkb = vec![
|
||||
0x10, 0x00, 0x00, 0x0C, 0x48, 0x14, 0x10, 0x03, 0x00, 0x00, 0x00, 0x4D,
|
||||
];
|
||||
mkb.extend_from_slice(&[0x07, 0x00, 0x00, 0x14]);
|
||||
mkb.extend_from_slice(&[0xAB; 16]);
|
||||
mkb.extend_from_slice(&[0x86, 0x00, 0x00, 0x14]);
|
||||
mkb.extend_from_slice(&[0xCD; 16]);
|
||||
mkb
|
||||
}
|
||||
|
||||
fn synthetic_mkb_with_variant() -> Vec<u8> {
|
||||
let mut mkb = synthetic_mkb_classical();
|
||||
// 0x82 — 16-byte body (Variant data / VKD slot).
|
||||
mkb.extend_from_slice(&[0x82, 0x00, 0x00, 0x14]);
|
||||
mkb.extend_from_slice(&[0xEE; 16]);
|
||||
// 0x83 — 16-byte body (Variant Nonce).
|
||||
mkb.extend_from_slice(&[0x83, 0x00, 0x00, 0x14]);
|
||||
mkb.extend_from_slice(&[0x55; 16]);
|
||||
mkb
|
||||
}
|
||||
|
||||
// ── Walker / record detection ──
|
||||
|
||||
#[test]
|
||||
fn walker_parses_synthetic_mkb() {
|
||||
let mkb = synthetic_mkb_classical();
|
||||
let recs = walk_mkb(&mkb);
|
||||
assert_eq!(recs.len(), 3);
|
||||
assert_eq!(recs[0].rec_type, 0x10);
|
||||
assert_eq!(recs[1].rec_type, 0x07);
|
||||
assert_eq!(recs[2].rec_type, 0x86);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn variant_detection_negative_on_classical() {
|
||||
let recs = walk_mkb(&synthetic_mkb_classical());
|
||||
assert!(!is_variant_mkb(&recs));
|
||||
assert!(variant_nonce(&recs).is_none());
|
||||
assert!(variant_key_data(&recs).is_none());
|
||||
assert!(variant_data_record(&recs).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn variant_detection_positive_on_variant() {
|
||||
let recs = walk_mkb(&synthetic_mkb_with_variant());
|
||||
assert!(is_variant_mkb(&recs));
|
||||
assert_eq!(variant_nonce(&recs), Some([0x55; 16]));
|
||||
assert_eq!(variant_key_data(&recs), Some(&[0xEE; 16][..]));
|
||||
assert_eq!(variant_data_record(&recs), Some(&[0xEE; 16][..]));
|
||||
}
|
||||
|
||||
// ── Chain entry-point classification ──
|
||||
|
||||
#[test]
|
||||
fn chain_rejects_non_variant_mkb() {
|
||||
let recs = walk_mkb(&synthetic_mkb_classical());
|
||||
let err = derive_media_key_variant(&recs, &[], &[0xAA; 16], &[0u8; 16])
|
||||
.expect_err("classical MKB must be rejected");
|
||||
assert_eq!(err, MediaKeyVariantError::NotVariantMkb);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chain_rejects_placeholder_kcd() {
|
||||
// To reach the KCD check we need a complete variant MKB AND a
|
||||
// DK that walks it. We construct both via the synthetic
|
||||
// fixture below.
|
||||
let (recs, dk, _kp, _expected_kmp) = synthetic_variant_setup(/*kmp15*/ 0x00);
|
||||
let err =
|
||||
derive_media_key_variant(&recs, &[dk], &KEY_CORRECTION_DATA_PLACEHOLDER, &[0u8; 16])
|
||||
.expect_err("placeholder KCD must be rejected");
|
||||
assert_eq!(err, MediaKeyVariantError::KcdNotProvided);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chain_detects_soft_correction_bit() {
|
||||
let (recs, dk, _, _) = synthetic_variant_setup(/*kmp15*/ 0x02);
|
||||
let err = derive_media_key_variant(&recs, &[dk], &[0xAA; 16], &[0u8; 16])
|
||||
.expect_err("bit 0x02 must surface SoftCorrectionRequired");
|
||||
assert_eq!(err, MediaKeyVariantError::SoftCorrectionRequired);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chain_detects_online_challenge_bit() {
|
||||
let (recs, dk, _, _) = synthetic_variant_setup(/*kmp15*/ 0x04);
|
||||
let err = derive_media_key_variant(&recs, &[dk], &[0xAA; 16], &[0u8; 16])
|
||||
.expect_err("bit 0x04 must surface OnlineChallengeRequired");
|
||||
assert_eq!(err, MediaKeyVariantError::OnlineChallengeRequired);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chain_surfaces_variants_table_gap_on_clean_kmp() {
|
||||
// With both condition bits clear and a non-placeholder KCD, the
|
||||
// chain advances to the per-uv VARIANTS[uv] lookup, which is
|
||||
// not yet wired. That returns VariantsTableUnavailable —
|
||||
// proving the bit checks and KCD check all passed.
|
||||
let (recs, dk, _, _) = synthetic_variant_setup(/*kmp15*/ 0x00);
|
||||
let err = derive_media_key_variant(&recs, &[dk], &[0xAA; 16], &[0u8; 16])
|
||||
.expect_err("expected VariantsTableUnavailable at the per-uv lookup");
|
||||
assert_eq!(err, MediaKeyVariantError::VariantsTableUnavailable);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn error_display_is_code_only() {
|
||||
// No English in Display — every variant emits "E7xxx" and
|
||||
// nothing else.
|
||||
let cases = [
|
||||
MediaKeyVariantError::NotVariantMkb,
|
||||
MediaKeyVariantError::MkbIncomplete,
|
||||
MediaKeyVariantError::ProcessingKeyUnavailable,
|
||||
MediaKeyVariantError::SoftCorrectionRequired,
|
||||
MediaKeyVariantError::OnlineChallengeRequired,
|
||||
MediaKeyVariantError::KcdNotProvided,
|
||||
MediaKeyVariantError::VariantsTableUnavailable,
|
||||
MediaKeyVariantError::VkdIndexOutOfRange,
|
||||
];
|
||||
for e in cases {
|
||||
let s = e.to_string();
|
||||
assert!(
|
||||
s.starts_with('E') && s.len() == 5,
|
||||
"error display must be E#### only, got {s:?}"
|
||||
);
|
||||
assert!(
|
||||
s.chars().skip(1).all(|c| c.is_ascii_digit()),
|
||||
"error display must be E + digits, got {s:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ── Fixture construction ──
|
||||
|
||||
/// Build a synthetic variant MKB plus a DK that walks the single
|
||||
/// subset-difference slot it carries. `kmp15` is the value of the
|
||||
/// low byte of `Kmp[15]` that the chain will land on — pick `0x02`
|
||||
/// to exercise the SoftCorrection bit, `0x04` to exercise
|
||||
/// OnlineChallenge, `0x00` otherwise.
|
||||
///
|
||||
/// The fixture pins:
|
||||
/// - MKB subset-difference: `u_mask_shift=3, uv=2`. With these
|
||||
/// masks the discriminator bit (u_mask=1, v_mask=0) is bit 2.
|
||||
/// - one DK at `node=4, uv=2, u_mask_shift=3`. node 4 has bit 2 set
|
||||
/// (differs from uv=2 on bit 2 → disagrees on v_mask) while
|
||||
/// agreeing with uv on bits 3+ (the u_mask=1 region). dk.uv ==
|
||||
/// MKB.uv and dk.u_mask_shift == MKB.u_mask_shift make
|
||||
/// `dev_key_v_mask == v_mask`, so `calc_pk_from_dk` loops zero
|
||||
/// times — Kp = aesg3_step(dk, 1).
|
||||
/// - one cvalue in record 0x07 chosen so AES-D(Kp, C) ⊕ uv produces a
|
||||
/// Kmp whose byte-15 is exactly `kmp15`.
|
||||
/// - record 0x82 with a 16-byte body (acts as both Variant Data
|
||||
/// and Variant Key Data; satisfies the parser heuristics).
|
||||
/// - record 0x83 with a 16-byte Nonce.
|
||||
///
|
||||
/// Returns (records, dk, planted_kp, planted_kmp).
|
||||
fn synthetic_variant_setup(kmp15: u8) -> (Vec<MkbRecord>, DeviceKey, [u8; 16], [u8; 16]) {
|
||||
use crate::aacs::decrypt::aes_ecb_encrypt;
|
||||
|
||||
// Build header.
|
||||
let mut mkb = vec![
|
||||
0x10, 0x00, 0x00, 0x0C, 0x48, 0x14, 0x10, 0x03, 0x00, 0x00, 0x00, 0x4D,
|
||||
];
|
||||
|
||||
// Subset-difference (0x04): u_mask_shift=3, uv=00 00 00 02.
|
||||
mkb.extend_from_slice(&[0x04, 0x00, 0x00, 0x09]);
|
||||
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
|
||||
// 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] = [
|
||||
0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE,
|
||||
0xFF, 0x00,
|
||||
];
|
||||
let kp = aesg3_step(&dk_bytes, 1);
|
||||
|
||||
// Plant Kmp with chosen byte-15, then compute C such that
|
||||
// AES-D(Kp, C) ⊕ uv == Kmp. uv=2 → low-4 bytes XOR is 00 00 00 02.
|
||||
let mut kmp = [0x42u8; 16];
|
||||
kmp[15] = kmp15;
|
||||
let mut aes_d_result = kmp;
|
||||
aes_d_result[15] ^= 0x02;
|
||||
let c_block = aes_ecb_encrypt(&kp, &aes_d_result);
|
||||
|
||||
// cvalues record (0x07): one 16-byte cvalue. The walker
|
||||
// indexes it for the magic-check step; on a variant MKB the
|
||||
// magic check fails but `variant_present` is true so the
|
||||
// walker still returns the match. Content is don't-care.
|
||||
mkb.extend_from_slice(&[0x07, 0x00, 0x00, 0x14]);
|
||||
mkb.extend_from_slice(&[0xAB; 16]);
|
||||
|
||||
// Verify Media Key (0x86): body content is don't-care.
|
||||
mkb.extend_from_slice(&[0x86, 0x00, 0x00, 0x14]);
|
||||
mkb.extend_from_slice(&[0xCD; 16]);
|
||||
|
||||
// 0x82 record: holds C (Encrypted Media Key Variant Data) AND
|
||||
// doubles as the VKD table (single 16-byte entry → VKDidx must
|
||||
// resolve to 0 for `chain_surfaces_variants_table_gap` test —
|
||||
// but the test never reaches the VKD lookup since the
|
||||
// VARIANTS[uv] helper is not yet wired).
|
||||
mkb.extend_from_slice(&[0x82, 0x00, 0x00, 0x14]);
|
||||
mkb.extend_from_slice(&c_block);
|
||||
|
||||
// 0x83 record: 16-byte Nonce.
|
||||
mkb.extend_from_slice(&[0x83, 0x00, 0x00, 0x14]);
|
||||
mkb.extend_from_slice(&[0x77; 16]);
|
||||
|
||||
let recs = walk_mkb(&mkb);
|
||||
|
||||
let dk = DeviceKey {
|
||||
key: dk_bytes,
|
||||
node: 4,
|
||||
uv: 2,
|
||||
u_mask_shift: 3,
|
||||
};
|
||||
(recs, dk, kp, kmp)
|
||||
}
|
||||
}
|
||||
+4
-3
@@ -4,13 +4,14 @@
|
||||
//! 1. Bus authentication (challenge-response) → bus key
|
||||
//! 2. Read disc key block (READ DVD STRUCTURE) → XOR with bus key → decrypt with player keys → disc key
|
||||
//! 3. Read title key (REPORT KEY format 0x04) → XOR with bus key → decrypt with disc key → title key
|
||||
//!
|
||||
//! Based on libdvdcss (VideoLAN) and Stevenson 1999 analysis.
|
||||
|
||||
use crate::drive::Drive;
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
// ── Player keys (from libdvdcss, Stevenson's PlayerKey cracker) ───────────
|
||||
// ── Built-in public DVD CSS player keys ────────────────────────────────────
|
||||
//
|
||||
// These 31 5-byte player keys are long-public CSS inputs. With them
|
||||
// compiled in, DVD ripping works with no external key file required.
|
||||
|
||||
const PLAYER_KEYS: [[u8; 5]; 31] = [
|
||||
[0x01, 0xaf, 0xe3, 0x12, 0x80],
|
||||
|
||||
@@ -16,6 +16,7 @@ pub mod lfsr;
|
||||
pub(crate) mod tables;
|
||||
|
||||
use crate::disc::Extent;
|
||||
use crate::drive::Drive;
|
||||
use crate::sector::SectorSource;
|
||||
|
||||
/// CSS decryption state for a DVD title.
|
||||
@@ -25,6 +26,50 @@ pub struct CssState {
|
||||
pub title_key: [u8; 5],
|
||||
}
|
||||
|
||||
/// Inputs for CSS key acquisition.
|
||||
///
|
||||
/// The acquisition path depends on which inputs the caller supplies:
|
||||
///
|
||||
/// - With `drive` + `auth_lba` set, [`resolve`] runs the full SCSI bus
|
||||
/// auth + title-key path (live BU40N / DVD drive).
|
||||
/// - With `reader` + `extents` set, [`resolve`] falls back to the
|
||||
/// crack path (Stevenson known-plaintext attack on encrypted PES
|
||||
/// headers; works on disc images and on drives whose CSS auth path
|
||||
/// is unavailable).
|
||||
///
|
||||
/// `live_drive` 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>,
|
||||
/// LBA of a known-scrambled sector for the auth path's title-key
|
||||
/// query. Required when `drive` is set.
|
||||
pub auth_lba: Option<u32>,
|
||||
/// Sector source for the crack path.
|
||||
pub reader: Option<&'a mut dyn SectorSource>,
|
||||
/// Extents to scan for the crack path. Required when `reader` is
|
||||
/// set.
|
||||
pub extents: Option<&'a [Extent]>,
|
||||
}
|
||||
|
||||
/// Acquire a CSS title key using whichever inputs the context provides.
|
||||
///
|
||||
/// Order of attempts:
|
||||
/// 1. SCSI auth path (when `drive` and `auth_lba` are set).
|
||||
/// 2. Crack path (when `reader` and `extents` are set).
|
||||
///
|
||||
/// Returns `None` if neither path is configured or both fail.
|
||||
pub fn resolve(ctx: &mut CssContext<'_>) -> Option<CssState> {
|
||||
if let (Some(drive), Some(lba)) = (ctx.drive.as_deref_mut(), ctx.auth_lba) {
|
||||
if let Ok(title_key) = auth::authenticate_and_read_title_key(drive, lba) {
|
||||
return Some(CssState { title_key });
|
||||
}
|
||||
}
|
||||
if let (Some(reader), Some(extents)) = (ctx.reader.as_deref_mut(), ctx.extents) {
|
||||
return crack_key(reader, extents);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Crack the CSS title key by reading encrypted sectors and applying
|
||||
/// a known-plaintext attack on MPEG-2 headers.
|
||||
///
|
||||
|
||||
+160
-22
@@ -3,9 +3,131 @@
|
||||
//! Decrypts sectors in-place using resolved keys from disc scanning.
|
||||
//! Handles AACS 1.0, AACS 2.0, and CSS transparently.
|
||||
//! The caller never sees encrypted data unless explicitly bypassed.
|
||||
//!
|
||||
//! ## Parallel AACS decrypt
|
||||
//!
|
||||
//! Each AACS aligned unit (6144 bytes) is decrypted INDEPENDENTLY of
|
||||
//! every other unit — per-unit key derivation from the unit_key plus
|
||||
//! the unit's own first-16-byte header. There is no cross-unit
|
||||
//! dependency, so a buffer of N units can be decrypted on N threads
|
||||
//! in parallel via a persistent rayon thread pool.
|
||||
//!
|
||||
//! Small buffers (< [`PARALLEL_MIN_UNITS`] units) fall through to the
|
||||
//! serial path to avoid pool dispatch overhead beating the per-unit
|
||||
//! AES work.
|
||||
//!
|
||||
//! ## Thread-count configuration — three layers
|
||||
//!
|
||||
//! Resolution order (highest wins):
|
||||
//! 1. The most recent [`set_decrypt_threads`] call with `n > 0`.
|
||||
//! Calling this *replaces* the live thread pool — useful for a
|
||||
//! settings-page slider in a long-running daemon.
|
||||
//! 2. `FREEMKV_THREADS` env var, if set and `> 0`. Single knob
|
||||
//! covering decrypt today, intended to also drive any future
|
||||
//! input-side / output-side worker pools.
|
||||
//! 3. Default: all available cores. Algorithm optimisation comes
|
||||
//! first — we measure single-thread performance to find serial
|
||||
//! bottlenecks before throwing parallelism at it — but once a
|
||||
//! pool is engaged we use the whole box. Hard cap at
|
||||
//! [`MAX_THREADS`] (rayon stack memory).
|
||||
|
||||
use crate::aacs;
|
||||
use crate::css;
|
||||
use rayon::prelude::*;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
/// Minimum units in a buffer before we pay the pool-dispatch cost of
|
||||
/// fanning out. Below this, serial is faster.
|
||||
const PARALLEL_MIN_UNITS: usize = 8;
|
||||
|
||||
/// Hard upper bound on configurable thread count. Anything larger is
|
||||
/// almost certainly a misconfiguration; rayon would happily allocate
|
||||
/// thousands of worker stacks otherwise.
|
||||
pub const MAX_THREADS: usize = 64;
|
||||
|
||||
/// Process-wide decrypt thread count override. `0` means "use env
|
||||
/// var, else default" — see [`decrypt_threads`] for the resolution
|
||||
/// order.
|
||||
static DECRYPT_THREADS: AtomicUsize = AtomicUsize::new(0);
|
||||
|
||||
/// Current rayon pool. `RwLock<Option<Arc<...>>>` so that
|
||||
/// [`set_decrypt_threads`] can swap the pool out without leaking the
|
||||
/// old one and without blocking ongoing decrypt work (in-flight calls
|
||||
/// hold an `Arc` clone via [`decrypt_pool`] and finish on the old
|
||||
/// pool; new calls pick up the new pool).
|
||||
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.
|
||||
///
|
||||
/// Thread-safe. Live decrypt calls keep their previously-acquired
|
||||
/// pool reference for the rest of the call — no mid-call pool
|
||||
/// switch. Subsequent calls see the new pool.
|
||||
///
|
||||
/// Pool construction is ~ms-scale; safe to call from a settings POST
|
||||
/// handler.
|
||||
pub fn set_decrypt_threads(n: usize) {
|
||||
let clamped = n.min(MAX_THREADS);
|
||||
DECRYPT_THREADS.store(clamped, Ordering::Relaxed);
|
||||
// Drop the existing pool. Next decrypt_pool() call rebuilds with
|
||||
// the new resolved thread count.
|
||||
if let Ok(mut guard) = DECRYPT_POOL.write() {
|
||||
*guard = None;
|
||||
}
|
||||
}
|
||||
|
||||
/// 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() {
|
||||
if let Some(pool) = guard.as_ref() {
|
||||
return 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");
|
||||
if let Some(pool) = guard.as_ref() {
|
||||
return 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"),
|
||||
);
|
||||
*guard = Some(Arc::clone(&pool));
|
||||
pool
|
||||
}
|
||||
|
||||
/// Current effective decrypt thread count. Resolution order:
|
||||
/// 1. Most recent [`set_decrypt_threads`] value (if > 0)
|
||||
/// 2. `FREEMKV_THREADS` env var (if set and > 0)
|
||||
/// 3. Default: all available cores, capped at [`MAX_THREADS`].
|
||||
pub fn decrypt_threads() -> usize {
|
||||
let explicit = DECRYPT_THREADS.load(Ordering::Relaxed);
|
||||
if explicit > 0 {
|
||||
return explicit;
|
||||
}
|
||||
let env = std::env::var("FREEMKV_THREADS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<usize>().ok())
|
||||
.unwrap_or(0);
|
||||
if env > 0 {
|
||||
return env.min(MAX_THREADS);
|
||||
}
|
||||
let cores = std::thread::available_parallelism()
|
||||
.map(|n| n.get())
|
||||
.unwrap_or(2);
|
||||
cores.clamp(1, MAX_THREADS)
|
||||
}
|
||||
|
||||
/// Resolved decryption state from disc scanning.
|
||||
/// Passed to `decrypt_sectors()` — the caller doesn't need to know
|
||||
@@ -57,29 +179,45 @@ pub fn decrypt_sectors(
|
||||
return Err(crate::error::Error::DecryptFailed);
|
||||
}
|
||||
};
|
||||
let rdk = read_data_key.as_ref();
|
||||
let rdk: Option<[u8; 16]> = *read_data_key;
|
||||
let unit_len = aacs::ALIGNED_UNIT_LEN;
|
||||
let nthreads = decrypt_threads();
|
||||
let chunks: Vec<&mut [u8]> = buf.chunks_mut(unit_len).collect();
|
||||
let nunits = chunks.len();
|
||||
|
||||
for chunk in buf.chunks_mut(unit_len) {
|
||||
if chunk.len() == unit_len && aacs::is_unit_encrypted(chunk) {
|
||||
// `is_unit_encrypted` is a byte-0 heuristic: it fires on any
|
||||
// unit whose first byte has the top 2 bits set, which is
|
||||
// correct for m2ts source packets (where those bits are the
|
||||
// copy-control marker) but false-positives on any other binary
|
||||
// data with similarly-shaped first bytes — notably MPLS/CLPI
|
||||
// navigation files that begin with ASCII magic ('M', 'H'…)
|
||||
// and survive sweep mixed in with encrypted m2ts payloads.
|
||||
// `decrypt_unit_full` self-checks via TS-sync verification and
|
||||
// returns false on a misfire, but it has already mutated the
|
||||
// chunk by then. Snapshot and restore on verification failure
|
||||
// — same pattern `decrypt_unit_try_keys` uses for multi-key
|
||||
// discs. Real m2ts units verify and stay decrypted; nav-file
|
||||
// sectors get scrambled briefly and then put back as-was.
|
||||
// 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
|
||||
// scrambled and trigger a decrypt attempt, so on a verify miss we
|
||||
// snapshot+restore the original bytes so it survives. See test
|
||||
// `nav_file_unit_survives_decrypt_attempt`.
|
||||
let decrypt_one = |chunk: &mut [u8]| {
|
||||
if chunk.len() == unit_len && aacs::is_aacs_scrambled(chunk) {
|
||||
let original: Vec<u8> = chunk.to_vec();
|
||||
if !aacs::decrypt_unit_full(chunk, &uk, rdk) {
|
||||
if !aacs::decrypt_unit_full(chunk, &uk, rdk.as_ref()) {
|
||||
chunk.copy_from_slice(&original);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
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 {
|
||||
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);
|
||||
});
|
||||
});
|
||||
}
|
||||
}
|
||||
DecryptKeys::Css { title_key } => {
|
||||
@@ -95,11 +233,11 @@ pub fn decrypt_sectors(
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Regression for the 0.18.1 nav-file scramble bug. A non-m2ts unit whose
|
||||
/// first byte has the top 2 bits set (here: the ASCII letter 'M' that
|
||||
/// MPLS files start with, 0x4D = 0b01001101) trips `is_unit_encrypted`,
|
||||
/// gets AES-decrypted with the unit key, fails the TS-sync verification,
|
||||
/// and must be restored to its original bytes — not left scrambled.
|
||||
/// Regression for the 0.18.1 nav-file scramble bug. A non-m2ts unit (here
|
||||
/// an MPLS file: starts "MPLS", carries no TS syncs) reads as scrambled
|
||||
/// under `is_aacs_scrambled`, gets AES-decrypted with the unit key, fails
|
||||
/// the TS-sync verification, and must be restored to its original bytes —
|
||||
/// not left scrambled.
|
||||
#[test]
|
||||
fn nav_file_unit_survives_decrypt_attempt() {
|
||||
let mut unit = vec![0u8; aacs::ALIGNED_UNIT_LEN];
|
||||
|
||||
+272
-116
@@ -14,56 +14,221 @@ pub(super) struct HandshakeResult {
|
||||
}
|
||||
|
||||
impl Disc {
|
||||
/// SCSI handshake result — volume ID and bus keys from ECDH authentication.
|
||||
/// Only available when scanning from a real drive (not ISO images).
|
||||
/// Acquire the Volume ID. Tries the per-drive OEM CDB path first
|
||||
/// when the drive reports `is_unlocked()` (extended-access state),
|
||||
/// and falls back to the cert-based AACS mutual-auth handshake
|
||||
/// otherwise.
|
||||
///
|
||||
/// The OEM path is a single READ_BUFFER CDB built from the drive
|
||||
/// profile's `read_vid_cdb` template. The response carries a 3-byte
|
||||
/// header (validated against `00 22 00`) followed by the 16-byte
|
||||
/// VID at bytes [4..20]. Crucially, no AGID setup is required —
|
||||
/// the drive's runtime firmware serves the VID directly when in
|
||||
/// extended-access state.
|
||||
///
|
||||
/// The cert path is the standard AACS spec flow: ECDH key
|
||||
/// agreement, bus-key derivation, then `REPORT_DISC_STRUCTURE`
|
||||
/// format 0x80 to retrieve VID under bus-key MAC.
|
||||
pub(super) fn read_vid(
|
||||
session: &mut crate::drive::Drive,
|
||||
opts: &ScanOptions,
|
||||
) -> Result<[u8; 16]> {
|
||||
if session.is_unlocked() {
|
||||
let profile = session
|
||||
.drive_profile()
|
||||
.ok_or(Error::DriveProfileMissing)?
|
||||
.clone();
|
||||
return Self::read_vid_oem(session, &profile);
|
||||
}
|
||||
Self::read_vid_cert(session, opts)
|
||||
}
|
||||
|
||||
/// OEM VID retrieval — issues the per-drive READ_BUFFER CDB and
|
||||
/// parses the response.
|
||||
///
|
||||
/// Response layout (36 bytes):
|
||||
/// * [0..3] 3-byte response signature; expected `00 22 00`
|
||||
/// * [3] reserved
|
||||
/// * [4..20] 16-byte Volume ID
|
||||
/// * [20..36] reserved / per-drive padding
|
||||
fn read_vid_oem(
|
||||
session: &mut crate::drive::Drive,
|
||||
profile: &crate::profile::DriveProfile,
|
||||
) -> Result<[u8; 16]> {
|
||||
const RESPONSE_LEN: usize = 36;
|
||||
const EXPECTED_HEADER: [u8; 3] = [0x00, 0x22, 0x00];
|
||||
|
||||
let cdb = profile.read_vid_cdb.ok_or(Error::VidCdbUnavailable)?;
|
||||
let mut buf = vec![0u8; RESPONSE_LEN];
|
||||
let result = session.scsi_execute(
|
||||
&cdb,
|
||||
crate::scsi::DataDirection::FromDevice,
|
||||
&mut buf,
|
||||
5_000,
|
||||
)?;
|
||||
if result.bytes_transferred < RESPONSE_LEN {
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "oem_vid_short_response",
|
||||
bytes_transferred = result.bytes_transferred,
|
||||
"OEM VID CDB returned short response"
|
||||
);
|
||||
return Err(Error::AacsVidRead);
|
||||
}
|
||||
if buf[0..3] != EXPECTED_HEADER {
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "oem_vid_bad_header",
|
||||
header_0 = buf[0],
|
||||
header_1 = buf[1],
|
||||
header_2 = buf[2],
|
||||
"OEM VID response header mismatch"
|
||||
);
|
||||
return Err(Error::AacsVidRead);
|
||||
}
|
||||
let mut vid = [0u8; 16];
|
||||
vid.copy_from_slice(&buf[4..20]);
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "oem_vid_ok",
|
||||
"OEM VID retrieved"
|
||||
);
|
||||
Ok(vid)
|
||||
}
|
||||
|
||||
/// Cert-based VID retrieval — runs the full AACS mutual-auth
|
||||
/// handshake and extracts VID from the bus-key-MAC'd
|
||||
/// `REPORT_DISC_STRUCTURE` response.
|
||||
fn read_vid_cert(session: &mut crate::drive::Drive, opts: &ScanOptions) -> Result<[u8; 16]> {
|
||||
match Self::do_handshake_cert(session, opts) {
|
||||
(Some(h), _) => Ok(h.volume_id),
|
||||
(None, Some(e)) => Err(e),
|
||||
(None, None) => Err(Error::AacsVidUnavailable),
|
||||
}
|
||||
}
|
||||
|
||||
/// SCSI handshake — drives the VID-acquisition flow and returns
|
||||
/// a structured `HandshakeResult` for downstream key resolution.
|
||||
/// Prefers the OEM path when `Drive::is_unlocked()` is true and
|
||||
/// falls back to cert-based mutual auth otherwise.
|
||||
///
|
||||
/// The OEM path produces only VID (no bus-key, so no
|
||||
/// `read_data_key`); the cert path can produce both. AACS 2.0
|
||||
/// content that needs read_data_key for bus decryption requires
|
||||
/// the cert path.
|
||||
///
|
||||
/// Returns `(handshake, error)`:
|
||||
/// * `(Some(_), None)` — VID acquired
|
||||
/// * `(None, Some(_))` — specific failure mode (see
|
||||
/// `AacsHostCertRejected` / `AacsRawReadUnsupported` /
|
||||
/// `AacsVidUnavailable` / `DriveProfileMissing` /
|
||||
/// `VidCdbUnavailable` variants in `error.rs`)
|
||||
/// * `(None, None)` — handshake not attempted (no keydb;
|
||||
/// resolution will proceed with VID=zero and rely on path 1
|
||||
/// disc-hash → VUK lookup)
|
||||
pub(super) fn do_handshake(
|
||||
session: &mut crate::drive::Drive,
|
||||
opts: &ScanOptions,
|
||||
) -> Option<HandshakeResult> {
|
||||
use crate::aacs::{self, KeyDb};
|
||||
|
||||
) -> (Option<HandshakeResult>, Option<Error>) {
|
||||
let unlocked = session.is_unlocked();
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "handshake_entry",
|
||||
unlocked,
|
||||
"do_handshake entered"
|
||||
);
|
||||
let keydb_path = match opts.resolve_keydb() {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "handshake_no_keydb",
|
||||
"no KEYDB found in search paths; handshake skipped"
|
||||
);
|
||||
return None;
|
||||
|
||||
if unlocked {
|
||||
// Try OEM VID retrieval first. If the drive's profile
|
||||
// doesn't carry the CDB template, or the response is
|
||||
// malformed, fall through to cert-based auth.
|
||||
match Self::read_vid(session, opts) {
|
||||
Ok(volume_id) => {
|
||||
return (
|
||||
Some(HandshakeResult {
|
||||
volume_id,
|
||||
read_data_key: None,
|
||||
}),
|
||||
None,
|
||||
);
|
||||
}
|
||||
Err(Error::DriveProfileMissing) | Err(Error::VidCdbUnavailable) => {
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "handshake_oem_unavailable",
|
||||
"OEM VID path unavailable for this drive; trying cert handshake"
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "handshake_oem_failed",
|
||||
error_code = e.code(),
|
||||
"OEM VID retrieval failed; trying cert handshake"
|
||||
);
|
||||
}
|
||||
}
|
||||
};
|
||||
let keydb = match KeyDb::load(&keydb_path) {
|
||||
Ok(db) => db,
|
||||
Err(e) => {
|
||||
}
|
||||
|
||||
Self::do_handshake_cert(session, opts)
|
||||
}
|
||||
|
||||
/// Cert-based AACS handshake. The legacy auth path; still used as
|
||||
/// the fallback when the OEM VID path isn't available or fails.
|
||||
fn do_handshake_cert(
|
||||
session: &mut crate::drive::Drive,
|
||||
opts: &ScanOptions,
|
||||
) -> (Option<HandshakeResult>, Option<Error>) {
|
||||
use crate::aacs;
|
||||
|
||||
// Host certs come from the caller's DriveCredentials (e.g. the keydb's
|
||||
// host_certs(), sourced app-side) — the library does not load a keydb.
|
||||
// Absent ⇒ no cert auth: an unlocked / LibreDrive drive already returned
|
||||
// a Volume ID via the OEM path before reaching here, so this is the
|
||||
// locked-drive-without-credentials case.
|
||||
let host_certs: &[aacs::HostCert] = match &opts.credentials {
|
||||
Some(c) if !c.host_certs.is_empty() => &c.host_certs,
|
||||
_ => {
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "handshake_keydb_load_failed",
|
||||
io_error_kind = ?e.kind(),
|
||||
keydb = %keydb_path.display(),
|
||||
"KEYDB load failed; handshake skipped"
|
||||
phase = "handshake_no_credentials",
|
||||
"no drive credentials supplied; cert handshake skipped"
|
||||
);
|
||||
return None;
|
||||
return (None, None);
|
||||
}
|
||||
};
|
||||
|
||||
let host_cert_count = keydb.host_certs.len();
|
||||
let host_cert_count = host_certs.len();
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "handshake_start",
|
||||
host_cert_count,
|
||||
keydb = %keydb_path.display(),
|
||||
"handshake starting"
|
||||
);
|
||||
|
||||
const MAX_CERT_ATTEMPTS: usize = 16;
|
||||
// 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.
|
||||
//
|
||||
// Defense-in-depth: cap attempts, sleep between, and bail
|
||||
// early on the drive's wedge sense so any later regression
|
||||
// can't undo the protection silently.
|
||||
const MAX_CERT_ATTEMPTS: usize = 3;
|
||||
const PER_CERT_BACKOFF_MS: u64 = 1000;
|
||||
let mut last_err_code: Option<u16> = None;
|
||||
for (idx, hc) in keydb.host_certs.iter().take(MAX_CERT_ATTEMPTS).enumerate() {
|
||||
for (idx, hc) in host_certs.iter().take(MAX_CERT_ATTEMPTS).enumerate() {
|
||||
if idx > 0 {
|
||||
std::thread::sleep(std::time::Duration::from_millis(PER_CERT_BACKOFF_MS));
|
||||
}
|
||||
match aacs::handshake::aacs_authenticate(session, &hc.private_key, &hc.certificate) {
|
||||
Ok(mut auth) => {
|
||||
let volume_id = match aacs::handshake::read_volume_id(session, &mut auth) {
|
||||
@@ -76,7 +241,7 @@ impl Disc {
|
||||
error_code = e.code(),
|
||||
"auth ok but volume ID read failed"
|
||||
);
|
||||
return None;
|
||||
return (None, Some(Error::AacsVidUnavailable));
|
||||
}
|
||||
};
|
||||
let read_data_key = aacs::handshake::read_data_keys(session, &mut auth)
|
||||
@@ -88,13 +253,33 @@ impl Disc {
|
||||
cert_index = idx,
|
||||
has_read_data_key = read_data_key.is_some(),
|
||||
);
|
||||
return Some(HandshakeResult {
|
||||
volume_id,
|
||||
read_data_key,
|
||||
});
|
||||
return (
|
||||
Some(HandshakeResult {
|
||||
volume_id,
|
||||
read_data_key,
|
||||
}),
|
||||
None,
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
last_err_code = Some(e.code());
|
||||
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 {
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "handshake_wedge_detected",
|
||||
cert_index = idx,
|
||||
error_code = code,
|
||||
"drive returned ILLEGAL_REQUEST during auth; bailing out to avoid wedge"
|
||||
);
|
||||
return (None, Some(Error::AacsHostCertRejected));
|
||||
}
|
||||
continue;
|
||||
}
|
||||
}
|
||||
@@ -105,112 +290,83 @@ impl Disc {
|
||||
host_cert_count,
|
||||
tried = host_cert_count.min(MAX_CERT_ATTEMPTS),
|
||||
last_error_code = last_err_code,
|
||||
"all host certs in KEYDB rejected by drive"
|
||||
"all host certs in KEYDB rejected by drive (capped at {} attempts to prevent firmware wedge)",
|
||||
MAX_CERT_ATTEMPTS
|
||||
);
|
||||
// All host certs failed — return None, not a fake success
|
||||
None
|
||||
(None, Some(Error::AacsHostCertRejected))
|
||||
}
|
||||
|
||||
/// Resolve disc encryption — AACS 1.0, AACS 2.0, CSS, or none.
|
||||
///
|
||||
/// Reads AACS files from UDF (via SectorSource), resolves keys through
|
||||
/// whatever path works: KEYDB VUK lookup, media key derivation, processing
|
||||
/// keys, device keys. Uses handshake result (volume ID, bus key) if available.
|
||||
pub(super) fn resolve_encryption(
|
||||
/// Build a keys-free AACS state that carries only the Volume ID (+ version
|
||||
/// metadata), for callers that resolve Unit Keys out-of-band and have
|
||||
/// disabled the local keydb. The VID is on-disc content read during the
|
||||
/// handshake; preserving it here lets the out-of-band path use it. No keys
|
||||
/// are present (`unit_keys` empty, `vuk` None), so the disc reports as
|
||||
/// "encrypted, no keys" until the caller re-scans with a resolved Unit Key.
|
||||
pub(super) fn resolve_vid_only(
|
||||
udf_fs: &udf::UdfFs,
|
||||
reader: &mut dyn SectorSource,
|
||||
keydb_path: &std::path::Path,
|
||||
handshake: Option<&HandshakeResult>,
|
||||
) -> Result<AacsState> {
|
||||
use crate::aacs::{self, KeyDb};
|
||||
use crate::aacs;
|
||||
|
||||
let keydb = KeyDb::load(keydb_path).map_err(|_| Error::KeydbLoad {
|
||||
path: keydb_path.display().to_string(),
|
||||
})?;
|
||||
|
||||
// Read AACS files from disc/image via UDF
|
||||
let uk_ro_data = 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 dh = aacs::disc_hash(&uk_ro_data);
|
||||
|
||||
// Log the disc hash so we can confirm whether it's present in KEYDB
|
||||
// when key resolution fails. The disc hash is SHA-1 of the full
|
||||
// Unit_Key_RO.inf file bytes — same value KEYDB.cfg keys VUK entries by.
|
||||
let dh = crate::aacs::disc_hash(&uk_ro_data);
|
||||
let dh_hex = crate::aacs::disc_hash_hex(&dh);
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "scan_aacs_disc_hash",
|
||||
disc_hash = %dh_hex,
|
||||
uk_ro_len = uk_ro_data.len(),
|
||||
"disc hash computed (compare with keydb.cfg entries)"
|
||||
);
|
||||
|
||||
let cc_data = udf_fs
|
||||
let cc = udf_fs
|
||||
.read_file(reader, "/AACS/Content000.cer")
|
||||
.or_else(|_| udf_fs.read_file(reader, "/AACS/Content001.cer"))
|
||||
.ok();
|
||||
|
||||
let mkb_data = udf_fs
|
||||
.read_file(reader, "/AACS/MKB_RW.inf")
|
||||
.or_else(|_| udf_fs.read_file(reader, "/AACS/MKB_RO.inf"))
|
||||
.ok();
|
||||
let mkb_ver = mkb_data.as_deref().and_then(aacs::mkb_version);
|
||||
|
||||
let mkb_first_64_hex = mkb_data
|
||||
.ok()
|
||||
.as_deref()
|
||||
.map(|m| {
|
||||
m.iter()
|
||||
.take(64)
|
||||
.map(|b| format!("{b:02x}"))
|
||||
.collect::<String>()
|
||||
})
|
||||
.and_then(aacs::parse_content_cert);
|
||||
let bus_encryption = cc.as_ref().map(|c| c.bus_encryption).unwrap_or(false);
|
||||
let version = match cc.as_ref().map(|c| c.version) {
|
||||
Some(aacs::AacsVersion::V10) => 1,
|
||||
Some(_) => 2,
|
||||
None if bus_encryption => 2,
|
||||
None => 1,
|
||||
};
|
||||
// MKB_RO/RW are allocated to a fixed ~128 MiB and zero-padded; trim to
|
||||
// the real record length (same as `read_aacs_inputs`). Without this the
|
||||
// MKB stashed on `AacsState` — which `Disc::inputs()` and the device/
|
||||
// processing-key `decrypt_with` derivation consume, and which a key
|
||||
// source ships to an online service — is the full 128 MiB pad, not the
|
||||
// ~few-MB record stream.
|
||||
let mut mkb_bytes = udf_fs
|
||||
.read_file(reader, "/AACS/MKB_RO.inf")
|
||||
.or_else(|_| udf_fs.read_file(reader, "/AACS/MKB_RW.inf"))
|
||||
.ok()
|
||||
.unwrap_or_default();
|
||||
let n = aacs::mkb_content_len(&mkb_bytes);
|
||||
if n > 0 && n < mkb_bytes.len() {
|
||||
mkb_bytes.truncate(n);
|
||||
}
|
||||
let mkb_ver = aacs::mkb_version(&mkb_bytes);
|
||||
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "scan_aacs_mkb_info",
|
||||
mkb_present = mkb_data.is_some(),
|
||||
mkb_len = mkb_data.as_deref().map(|m| m.len()).unwrap_or(0),
|
||||
mkb_version = ?mkb_ver,
|
||||
mkb_first_64 = %mkb_first_64_hex,
|
||||
keydb_disc_count = keydb.disc_entries.len(),
|
||||
keydb_dk_count = keydb.device_keys.len(),
|
||||
keydb_pk_count = keydb.processing_keys.len(),
|
||||
"AACS resolution inputs"
|
||||
phase = "scan_aacs_vid_only",
|
||||
disc_hash = %aacs::disc_hash_hex(&dh),
|
||||
version,
|
||||
bus_encryption,
|
||||
has_vid = handshake.is_some(),
|
||||
"keydb disabled — carrying VID only, keys resolved out-of-band"
|
||||
);
|
||||
|
||||
// Use handshake volume ID if available, otherwise zeros
|
||||
// (KEYDB VUK lookup by disc hash works without volume ID)
|
||||
let volume_id = handshake.map(|h| h.volume_id).unwrap_or([0u8; 16]);
|
||||
let read_data_key = handshake.and_then(|h| h.read_data_key);
|
||||
|
||||
// Resolve: tries all available paths — KEYDB VUK, media key, processing key, device key
|
||||
let resolved = aacs::resolve_keys(
|
||||
&uk_ro_data,
|
||||
cc_data.as_deref(),
|
||||
&volume_id,
|
||||
&keydb,
|
||||
mkb_data.as_deref(),
|
||||
)
|
||||
.ok_or(Error::AacsNoKeys)?;
|
||||
|
||||
Ok(AacsState {
|
||||
version: if resolved.aacs2 { 2 } else { 1 },
|
||||
bus_encryption: resolved.bus_encryption,
|
||||
version,
|
||||
bus_encryption,
|
||||
mkb_version: mkb_ver,
|
||||
disc_hash: aacs::disc_hash_hex(&resolved.disc_hash),
|
||||
key_source: match resolved.key_source {
|
||||
1 => KeySource::KeyDb,
|
||||
2 => KeySource::KeyDbDerived,
|
||||
3 => KeySource::ProcessingKey,
|
||||
4 => KeySource::DeviceKey,
|
||||
_ => KeySource::KeyDb,
|
||||
},
|
||||
vuk: resolved.vuk,
|
||||
unit_keys: resolved.unit_keys,
|
||||
read_data_key,
|
||||
volume_id,
|
||||
disc_hash: aacs::disc_hash_hex(&dh),
|
||||
key_source: KeyOrigin::ExternalUk,
|
||||
vuk: None,
|
||||
unit_keys: vec![],
|
||||
read_data_key: handshake.and_then(|h| h.read_data_key),
|
||||
volume_id: handshake.map(|h| h.volume_id).unwrap_or([0u8; 16]),
|
||||
uk_ro: uk_ro_data,
|
||||
mkb: mkb_bytes,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -130,6 +130,21 @@ pub struct Mapfile {
|
||||
/// Wall-clock timestamp of the last successful `write_to_disk` (or
|
||||
/// the moment the mapfile was constructed, whichever is later).
|
||||
last_flushed: Instant,
|
||||
/// AACS Volume ID (16 bytes) for the disc, persisted as a
|
||||
/// `# freemkv-vid:` comment header so it survives to deferred-mux /
|
||||
/// resume without altering the ISO payload or breaking ddrescue
|
||||
/// data-line parsing. `None` for unencrypted / non-AACS discs.
|
||||
///
|
||||
/// MUTUALLY EXCLUSIVE with `unit_keys`: a disc whose keys were resolved
|
||||
/// persists the keys (`unit_keys`) and NOT the VID — the keys are the final
|
||||
/// answer, so deferred-mux/resume decrypts directly with no key service. A
|
||||
/// disc that did NOT resolve persists only the VID, the retry-able "still
|
||||
/// need a key" marker (a future mux can re-ask the key service with it).
|
||||
vid: Option<[u8; 16]>,
|
||||
/// Decrypted AACS unit keys `(CPS unit, key)`, persisted as `# freemkv-uk:`
|
||||
/// comment headers when the disc was successfully keyed. Mutually exclusive
|
||||
/// with `vid` (see above). Empty when unresolved.
|
||||
unit_keys: Vec<(u32, [u8; 16])>,
|
||||
}
|
||||
|
||||
impl Mapfile {
|
||||
@@ -154,6 +169,8 @@ impl Mapfile {
|
||||
},
|
||||
dirty: false,
|
||||
last_flushed: Instant::now(),
|
||||
vid: None,
|
||||
unit_keys: Vec::new(),
|
||||
};
|
||||
// Eager initial persist so a resume can pick this up even if
|
||||
// `record()` is never called.
|
||||
@@ -168,6 +185,8 @@ impl Mapfile {
|
||||
let mut entries = Vec::new();
|
||||
let mut saw_current_line = false;
|
||||
let mut version = String::from("unknown");
|
||||
let mut vid: Option<[u8; 16]> = None;
|
||||
let mut unit_keys: Vec<(u32, [u8; 16])> = Vec::new();
|
||||
for line in text.lines() {
|
||||
let t = line.trim();
|
||||
if t.is_empty() {
|
||||
@@ -178,6 +197,17 @@ impl Mapfile {
|
||||
if let Some(v) = rest.strip_prefix("Rescue Logfile. Created by ") {
|
||||
version = v.to_string();
|
||||
}
|
||||
if let Some(hex) = rest.strip_prefix("freemkv-vid:") {
|
||||
// Best-effort: a malformed or short VID comment is
|
||||
// ignored rather than failing the whole load.
|
||||
vid = parse_vid_hex(hex.trim());
|
||||
}
|
||||
if let Some(uk) = rest.strip_prefix("freemkv-uk:") {
|
||||
// `<cps>:<32hex>`. Best-effort: a malformed line is skipped.
|
||||
if let Some(entry) = parse_uk_line(uk.trim()) {
|
||||
unit_keys.push(entry);
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
// First non-comment line is the "current" state line (pos status [pass] [pass_time]).
|
||||
@@ -227,6 +257,8 @@ impl Mapfile {
|
||||
stats,
|
||||
dirty: false,
|
||||
last_flushed: Instant::now(),
|
||||
vid,
|
||||
unit_keys,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -316,6 +348,41 @@ impl Mapfile {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Record the disc's 16-byte AACS Volume ID so it persists in the
|
||||
/// mapfile's comment header. Marks the mapfile dirty; the next
|
||||
/// `flush()` / `Drop` writes the `# freemkv-vid:` line. Does not
|
||||
/// touch the ISO payload or the ddrescue data lines.
|
||||
pub fn set_vid(&mut self, vid: [u8; 16]) {
|
||||
self.vid = Some(vid);
|
||||
self.dirty = true;
|
||||
}
|
||||
|
||||
/// The disc's AACS Volume ID, if one was set or parsed from a
|
||||
/// `# freemkv-vid:` comment on load. `None` for unencrypted /
|
||||
/// non-AACS discs.
|
||||
pub fn vid(&self) -> Option<[u8; 16]> {
|
||||
self.vid
|
||||
}
|
||||
|
||||
/// Record the disc's decrypted AACS unit keys so they persist in the
|
||||
/// mapfile header (`# freemkv-uk:` lines). The KEYED state: a deferred-mux /
|
||||
/// resume decrypts directly from these with no key-service round-trip.
|
||||
/// Setting keys clears any VID — the mapfile holds keys XOR VID, never both
|
||||
/// (keys are the final answer; VID is only the "still unresolved" marker).
|
||||
pub fn set_unit_keys(&mut self, keys: &[(u32, [u8; 16])]) {
|
||||
self.unit_keys = keys.to_vec();
|
||||
if !self.unit_keys.is_empty() {
|
||||
self.vid = None;
|
||||
}
|
||||
self.dirty = true;
|
||||
}
|
||||
|
||||
/// The disc's decrypted AACS unit keys, if the disc was keyed (parsed from
|
||||
/// `# freemkv-uk:` comments on load). Empty = unresolved (check `vid()`).
|
||||
pub fn unit_keys(&self) -> &[(u32, [u8; 16])] {
|
||||
&self.unit_keys
|
||||
}
|
||||
|
||||
pub fn entries(&self) -> &[MapEntry] {
|
||||
&self.entries
|
||||
}
|
||||
@@ -388,6 +455,30 @@ impl Mapfile {
|
||||
let file = std::fs::File::create(&tmp)?;
|
||||
let mut w = std::io::BufWriter::new(file);
|
||||
writeln!(w, "# Rescue Logfile. Created by {}", self.version)?;
|
||||
// VID comment lives in the header block. ddrescue treats any
|
||||
// `#`-prefixed line as a comment, so this round-trips through
|
||||
// our `load()` without affecting the `pos size status` data
|
||||
// parser. 16 bytes → 32 lowercase hex chars.
|
||||
// KEYS XOR VID: a keyed disc persists its unit keys (the final
|
||||
// answer — deferred-mux decrypts directly); an unresolved disc
|
||||
// persists only the VID (the retry marker, so a future mux can
|
||||
// re-ask the key service). Never both.
|
||||
use std::fmt::Write as _;
|
||||
if !self.unit_keys.is_empty() {
|
||||
for (cps, key) in &self.unit_keys {
|
||||
let mut hex = String::with_capacity(32);
|
||||
for b in key {
|
||||
let _ = write!(hex, "{b:02x}");
|
||||
}
|
||||
writeln!(w, "# freemkv-uk: {cps}:{hex}")?;
|
||||
}
|
||||
} else if let Some(vid) = self.vid {
|
||||
let mut hex = String::with_capacity(32);
|
||||
for b in vid {
|
||||
let _ = write!(hex, "{b:02x}");
|
||||
}
|
||||
writeln!(w, "# freemkv-vid: {hex}")?;
|
||||
}
|
||||
writeln!(w, "# Current pos / status / pass / pass_time")?;
|
||||
writeln!(w, "0x000000000 ? 1 0")?;
|
||||
writeln!(w, "# pos size status")?;
|
||||
@@ -417,6 +508,31 @@ impl Drop for Mapfile {
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a 32-char lowercase/uppercase hex string into a 16-byte VID.
|
||||
/// Returns `None` on any malformation (wrong length, non-hex) — the
|
||||
/// caller treats a bad VID comment as simply absent rather than an
|
||||
/// 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 {
|
||||
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()?;
|
||||
}
|
||||
Some(out)
|
||||
}
|
||||
|
||||
/// 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])> {
|
||||
let (cps, hex) = s.split_once(':')?;
|
||||
let cps: u32 = cps.trim().parse().ok()?;
|
||||
let key = parse_vid_hex(hex.trim())?; // 32-hex → [u8; 16], shared parser
|
||||
Some((cps, key))
|
||||
}
|
||||
|
||||
fn parse_hex(s: &str) -> io::Result<u64> {
|
||||
let s = s.strip_prefix("0x").unwrap_or(s);
|
||||
u64::from_str_radix(s, 16).map_err(|_| {
|
||||
@@ -583,6 +699,138 @@ mod tests {
|
||||
let _ = std::fs::remove_file(&p);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unit_keys_round_trip_and_are_mutually_exclusive_with_vid() {
|
||||
let p = tmpfile("uk_round_trips");
|
||||
let _ = std::fs::remove_file(&p);
|
||||
let mut mf = Mapfile::create(&p, 1000, "test").unwrap();
|
||||
mf.record(0, 500, SectorStatus::Finished).unwrap();
|
||||
// Set a VID first, then unit keys: keys must WIN and clear the VID.
|
||||
mf.set_vid([0xAA; 16]);
|
||||
let keys: Vec<(u32, [u8; 16])> = vec![
|
||||
(
|
||||
0,
|
||||
[
|
||||
0x57, 0x60, 0xcc, 0x83, 0x3d, 0x86, 0x0e, 0x48, 0x92, 0x1f, 0x88, 0x16, 0xe1,
|
||||
0x35, 0x9b, 0xad,
|
||||
],
|
||||
),
|
||||
(1, [0x11; 16]),
|
||||
];
|
||||
mf.set_unit_keys(&keys);
|
||||
assert_eq!(
|
||||
mf.vid(),
|
||||
None,
|
||||
"set_unit_keys must clear vid (keys XOR vid)"
|
||||
);
|
||||
mf.flush().unwrap();
|
||||
|
||||
let text = std::fs::read_to_string(&p).unwrap();
|
||||
assert!(
|
||||
text.contains("# freemkv-uk: 0:5760cc833d860e48921f8816e1359bad"),
|
||||
"uk comment format mismatch: {text}"
|
||||
);
|
||||
assert!(
|
||||
text.contains("# freemkv-uk: 1:11111111111111111111111111111111"),
|
||||
"second uk missing: {text}"
|
||||
);
|
||||
assert!(
|
||||
!text.contains("# freemkv-vid:"),
|
||||
"VID must NOT be written when keys are present: {text}"
|
||||
);
|
||||
|
||||
// load() recovers the unit keys (and no VID).
|
||||
let loaded = Mapfile::load(&p).unwrap();
|
||||
assert_eq!(loaded.unit_keys(), keys.as_slice());
|
||||
assert_eq!(loaded.vid(), None);
|
||||
assert_eq!(loaded.entries(), mf.entries());
|
||||
|
||||
// VID-only path (no keys) still persists the VID as the retry marker.
|
||||
let p2 = tmpfile("uk_vid_only");
|
||||
let _ = std::fs::remove_file(&p2);
|
||||
let mut mf2 = Mapfile::create(&p2, 1000, "test").unwrap();
|
||||
mf2.set_vid([0xBB; 16]);
|
||||
mf2.flush().unwrap();
|
||||
let loaded2 = Mapfile::load(&p2).unwrap();
|
||||
assert_eq!(loaded2.vid(), Some([0xBB; 16]));
|
||||
assert!(loaded2.unit_keys().is_empty());
|
||||
let _ = std::fs::remove_file(&p);
|
||||
let _ = std::fs::remove_file(&p2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vid_round_trips_and_data_lines_unaffected() {
|
||||
let p = tmpfile("vid_round_trips");
|
||||
let _ = std::fs::remove_file(&p);
|
||||
|
||||
// Build a mapfile with some data ranges, set a VID, persist.
|
||||
let mut mf = Mapfile::create(&p, 1000, "test").unwrap();
|
||||
mf.record(100, 200, SectorStatus::Finished).unwrap();
|
||||
mf.record(500, 100, SectorStatus::Unreadable).unwrap();
|
||||
mf.record(700, 50, SectorStatus::NonTrimmed).unwrap();
|
||||
let vid: [u8; 16] = [
|
||||
0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd,
|
||||
0xee, 0xff,
|
||||
];
|
||||
mf.set_vid(vid);
|
||||
mf.flush().unwrap();
|
||||
|
||||
// The saved file must contain the VID comment in lowercase hex.
|
||||
let text = std::fs::read_to_string(&p).unwrap();
|
||||
assert!(
|
||||
text.contains("# freemkv-vid:"),
|
||||
"saved mapfile missing VID comment: {text}"
|
||||
);
|
||||
assert!(
|
||||
text.contains("# freemkv-vid: 00112233445566778899aabbccddeeff"),
|
||||
"VID comment format mismatch: {text}"
|
||||
);
|
||||
|
||||
// load() recovers the VID and the identical data ranges.
|
||||
let loaded = Mapfile::load(&p).unwrap();
|
||||
assert_eq!(loaded.vid(), Some(vid));
|
||||
assert_eq!(loaded.entries(), mf.entries());
|
||||
|
||||
// A mapfile WITHOUT the VID comment must parse the same +/-/?
|
||||
// data ranges as the one WITH it (comment ignored by parser).
|
||||
let p2 = tmpfile("vid_round_trips_novid");
|
||||
let _ = std::fs::remove_file(&p2);
|
||||
let mut mf2 = Mapfile::create(&p2, 1000, "test").unwrap();
|
||||
mf2.record(100, 200, SectorStatus::Finished).unwrap();
|
||||
mf2.record(500, 100, SectorStatus::Unreadable).unwrap();
|
||||
mf2.record(700, 50, SectorStatus::NonTrimmed).unwrap();
|
||||
mf2.flush().unwrap();
|
||||
let loaded_novid = Mapfile::load(&p2).unwrap();
|
||||
assert_eq!(loaded_novid.vid(), None);
|
||||
assert_eq!(loaded_novid.entries(), loaded.entries());
|
||||
|
||||
// Malformed VID comments must not error the load (treated absent).
|
||||
let mut bad = text.replace("00112233445566778899aabbccddeeff", "zzzz");
|
||||
let pbad = tmpfile("vid_round_trips_bad");
|
||||
let _ = std::fs::remove_file(&pbad);
|
||||
std::fs::write(&pbad, &bad).unwrap();
|
||||
let loaded_bad = Mapfile::load(&pbad).unwrap();
|
||||
assert_eq!(loaded_bad.vid(), None);
|
||||
assert_eq!(loaded_bad.entries(), loaded.entries());
|
||||
|
||||
// A load->save cycle preserves the VID (the patch-pass path).
|
||||
bad.clear();
|
||||
let resaved = tmpfile("vid_round_trips_resave");
|
||||
let _ = std::fs::remove_file(&resaved);
|
||||
let mut reloaded = Mapfile::load(&p).unwrap();
|
||||
// Repoint at a fresh path and flush; mark dirty via a no-op record.
|
||||
reloaded.path = resaved.clone();
|
||||
reloaded.dirty = true;
|
||||
reloaded.flush().unwrap();
|
||||
let again = Mapfile::load(&resaved).unwrap();
|
||||
assert_eq!(again.vid(), Some(vid));
|
||||
|
||||
let _ = std::fs::remove_file(&p);
|
||||
let _ = std::fs::remove_file(&p2);
|
||||
let _ = std::fs::remove_file(&pbad);
|
||||
let _ = std::fs::remove_file(&resaved);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stats_consistent_after_split_record() {
|
||||
let p = tmpfile("stats_consistent_after_split");
|
||||
|
||||
+935
-108
File diff suppressed because it is too large
Load Diff
+1274
-947
File diff suppressed because it is too large
Load Diff
+30
-2
@@ -186,6 +186,22 @@ impl ReadCtx {
|
||||
/// marginal media is part of the job, and the fast-jump
|
||||
/// threshold is loose so we don't bail too early on a range that
|
||||
/// has scattered good sectors mixed in.
|
||||
///
|
||||
/// `damage_threshold_pct = 6` mirrors `disc/patch.rs`'s
|
||||
/// `PASSN_DAMAGE_THRESHOLD_PCT`. Pass N triggers the damage-skip
|
||||
/// at half the density Pass 1 uses (Pass 1 = 12%) because the
|
||||
/// patch loop's whole job is to chip away at bad ranges — being
|
||||
/// more eager to skip clustered bad sectors converges faster on
|
||||
/// the recoverable good sectors inside a range. The patch-side
|
||||
/// `compute_damage_skip` reads its threshold from
|
||||
/// `PASSN_DAMAGE_THRESHOLD_PCT`; keep the two in sync until the
|
||||
/// patch loop's damage-skip is unified with `handle_read_error`'s
|
||||
/// jump path. (v0.20.8 unification attempt found the unification
|
||||
/// itself blocked on the size-aware `range_remaining/4` cap that
|
||||
/// lives in `compute_damage_skip` but not in
|
||||
/// `handle_read_error::JumpAhead` — see
|
||||
/// `tests/passn_handler_ab.rs` for the A/B fixture that pins
|
||||
/// the divergence point.)
|
||||
pub fn for_patch(batch: u16) -> Self {
|
||||
Self {
|
||||
batch,
|
||||
@@ -194,7 +210,7 @@ impl ReadCtx {
|
||||
consecutive_outer_failures: 0,
|
||||
damage_window: Vec::with_capacity(16),
|
||||
damage_window_max: 16,
|
||||
damage_threshold_pct: 12,
|
||||
damage_threshold_pct: PATCH_DAMAGE_THRESHOLD_PCT,
|
||||
// Pass N is allowed to grind: window-based jump only,
|
||||
// matching the historical behaviour for patch passes.
|
||||
fast_jump_threshold: u64::MAX,
|
||||
@@ -318,7 +334,7 @@ const FAIL_PAUSE_SECS: u64 = 5;
|
||||
/// FIRST read failure after a clean run, before the drive has had a
|
||||
/// chance to cycle in retries that push it toward fast-fail).
|
||||
///
|
||||
/// Empirical: 2026-05-11 Dune Pt 2 wedge incident showed 7 medium
|
||||
/// Empirical: a 2026-05-11 wedge incident showed 7 medium
|
||||
/// errors in 6.5 seconds (~1s per attempt + ~1s pause) push the
|
||||
/// BU40N's firmware into IllegalRequest fast-fail mode permanently.
|
||||
/// Once there, only physical eject + reload clears it. Giving the
|
||||
@@ -401,6 +417,18 @@ const WEDGE_ABORT_THRESHOLD: u64 = 16;
|
||||
/// next range.
|
||||
const WEDGE_PASS_N_SKIP_SECTORS: u64 = 64;
|
||||
|
||||
/// Single source of truth for the Pass-N damage-window threshold.
|
||||
/// Both [`ReadCtx::for_patch`] and `disc::patch::compute_damage_skip`
|
||||
/// reference this constant so the two damage-skip paths cannot drift.
|
||||
///
|
||||
/// 6% means: with a 16-entry sliding window, the damage-skip fires
|
||||
/// once 1 out of 16 recent reads has failed. Pass 1 uses a 12%
|
||||
/// threshold via `damage_threshold_pct` on `for_sweep`; Pass N is
|
||||
/// twice as eager because patch's whole job is to converge on the
|
||||
/// bad sub-zones inside a NonTrimmed range — a faster trigger
|
||||
/// produces tighter convergence in fewer iterations.
|
||||
pub const PATCH_DAMAGE_THRESHOLD_PCT: usize = 6;
|
||||
|
||||
/// THE single error-handling entry point. Updates `ctx`, returns the
|
||||
/// action the caller must apply.
|
||||
///
|
||||
|
||||
+29
-1
@@ -184,6 +184,13 @@ impl Drive {
|
||||
self.profile.is_some()
|
||||
}
|
||||
|
||||
/// Borrow the matched drive profile, if any. Used by callers that
|
||||
/// need to issue per-drive OEM CDB templates (e.g. the OEM VID
|
||||
/// retrieval path in `disc::encrypt`).
|
||||
pub fn drive_profile(&self) -> Option<&DriveProfile> {
|
||||
self.profile.as_ref()
|
||||
}
|
||||
|
||||
/// Access the SCSI transport for direct commands (used by CSS/AACS auth).
|
||||
pub fn scsi_mut(&mut self) -> &mut dyn ScsiTransport {
|
||||
self.scsi.as_mut()
|
||||
@@ -432,6 +439,27 @@ impl Drive {
|
||||
}
|
||||
}
|
||||
|
||||
/// True if the drive is currently in the extended-access state.
|
||||
///
|
||||
/// Detected by the platform driver during `init()` from the unlock
|
||||
/// response's mode markers. When true:
|
||||
/// - SCSI READ_10 returns plaintext sectors (no AACS bus
|
||||
/// encryption applied)
|
||||
/// - VID retrieval works via the per-drive OEM CDB in
|
||||
/// [`DriveProfile`] without the cert-based AACS handshake
|
||||
/// - Disc-side Host Revocation List enforcement is effectively
|
||||
/// bypassed by the alternate data path
|
||||
///
|
||||
/// AACS layer code branches on this: if true, issue the OEM
|
||||
/// `read_vid_cdb` to retrieve VID directly; if false, fall back
|
||||
/// to the cert-based mutual-auth handshake.
|
||||
pub fn is_unlocked(&self) -> bool {
|
||||
match self.driver {
|
||||
Some(ref d) => d.is_unlocked(),
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Read sectors from the disc. Single-shot — no inline retries, no
|
||||
/// SCSI reset.
|
||||
///
|
||||
@@ -501,7 +529,7 @@ impl Drive {
|
||||
|
||||
// /dev/sr0 pread fallback (Linux only). The kernel
|
||||
// sr_mod driver auto-retries failed reads (~5× per
|
||||
// command). Empirically (BU40N + Dune Part 2 UHD,
|
||||
// command). Empirically (BU40N + a UHD disc,
|
||||
// 2026-05-08) dd via /dev/sr0 recovers ~50% of bad
|
||||
// sectors that a single-shot SG_IO READ misses.
|
||||
#[cfg(target_os = "linux")]
|
||||
|
||||
+289
@@ -0,0 +1,289 @@
|
||||
//! Top-level DRM scheme dispatch.
|
||||
//!
|
||||
//! Four content-protection schemes ride through a single
|
||||
//! detect-then-load pipeline:
|
||||
//!
|
||||
//! | Scheme | Discriminator |
|
||||
//! |---------------------|------------------------------------------------|
|
||||
//! | [`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` |
|
||||
//!
|
||||
//! 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.
|
||||
|
||||
use crate::aacs;
|
||||
use crate::css;
|
||||
|
||||
/// Which content-protection scheme governs a disc.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum DrmScheme {
|
||||
/// DVD Content Scramble System.
|
||||
Css,
|
||||
/// AACS 1.0 — original BD-ROM.
|
||||
Aacs10,
|
||||
/// AACS 2.0 — UHD-BD, classical Media Key chain.
|
||||
Aacs20,
|
||||
/// AACS 2.1 — UHD-BD with Media Key Variant chain.
|
||||
Aacs21,
|
||||
}
|
||||
|
||||
/// Inputs to [`DrmScheme::detect`]. All borrows — caller retains
|
||||
/// ownership.
|
||||
pub struct DrmProbe<'a> {
|
||||
/// 2048-byte sample sector from inside a DVD title's extents. Used
|
||||
/// only for CSS scramble-flag detection. `None` for non-DVD discs.
|
||||
pub dvd_sample_sector: Option<&'a [u8]>,
|
||||
/// Content Certificate file bytes (typically `/AACS/Content000.cer`).
|
||||
/// `None` when the disc has no AACS directory.
|
||||
pub content_cert: Option<&'a [u8]>,
|
||||
/// MKB file bytes (typically `/AACS/MKB_RW.inf`). Required to
|
||||
/// distinguish AACS 2.0 from AACS 2.1.
|
||||
pub mkb: Option<&'a [u8]>,
|
||||
}
|
||||
|
||||
/// Inputs to [`DrmScheme::load`]. Carries everything needed by either
|
||||
/// the AACS or CSS loader.
|
||||
pub struct DrmContext<'a> {
|
||||
/// AACS resolver inputs — required when the scheme is any AACS
|
||||
/// variant.
|
||||
pub aacs: Option<aacs::ResolveContext<'a>>,
|
||||
/// CSS resolver inputs — required when the scheme is [`DrmScheme::Css`].
|
||||
pub css: Option<css::CssContext<'a>>,
|
||||
}
|
||||
|
||||
/// Resolved key material, tagged by scheme.
|
||||
#[derive(Debug)]
|
||||
pub enum ResolvedScheme {
|
||||
Css(css::CssState),
|
||||
Aacs(aacs::ResolvedKeys),
|
||||
}
|
||||
|
||||
impl DrmScheme {
|
||||
/// Detect which DRM scheme protects the disc described by `probe`.
|
||||
///
|
||||
/// Returns `None` for unencrypted media. The order is intentional:
|
||||
/// CSS is checked first (DVD-format probe), then AACS (Blu-ray
|
||||
/// format).
|
||||
pub fn detect(probe: &DrmProbe<'_>) -> Option<DrmScheme> {
|
||||
// CSS — DVD probe sector carries the scramble flag.
|
||||
if let Some(sector) = probe.dvd_sample_sector {
|
||||
if css::is_scrambled(sector) {
|
||||
return Some(DrmScheme::Css);
|
||||
}
|
||||
}
|
||||
|
||||
// AACS — content cert type byte distinguishes V10 from V20+.
|
||||
// V21 promotion requires MKB Variant records.
|
||||
let cc = probe.content_cert.and_then(aacs::parse_content_cert)?;
|
||||
match cc.version {
|
||||
aacs::AacsVersion::V10 => Some(DrmScheme::Aacs10),
|
||||
aacs::AacsVersion::V20 | aacs::AacsVersion::V21 => {
|
||||
if let Some(mkb) = probe.mkb {
|
||||
let recs = aacs::variants::walk_mkb(mkb);
|
||||
if aacs::variants::is_variant_mkb(&recs) {
|
||||
return Some(DrmScheme::Aacs21);
|
||||
}
|
||||
}
|
||||
Some(DrmScheme::Aacs20)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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).
|
||||
pub fn load(self, ctx: &mut DrmContext<'_>) -> Option<ResolvedScheme> {
|
||||
match self {
|
||||
DrmScheme::Css => ctx
|
||||
.css
|
||||
.as_mut()
|
||||
.and_then(css::resolve)
|
||||
.map(ResolvedScheme::Css),
|
||||
DrmScheme::Aacs10 => ctx
|
||||
.aacs
|
||||
.as_ref()
|
||||
.and_then(aacs::resolve_keys_v1)
|
||||
.map(ResolvedScheme::Aacs),
|
||||
DrmScheme::Aacs20 => 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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// Build a minimal cert: type byte + bus-encryption byte + 6 zero
|
||||
// cc_id bytes.
|
||||
fn cert(type_byte: u8) -> Vec<u8> {
|
||||
let mut v = vec![0u8; 8];
|
||||
v[0] = type_byte;
|
||||
v
|
||||
}
|
||||
|
||||
// Synthetic AACS 2.x MKB with no Variant records.
|
||||
fn mkb_classical() -> Vec<u8> {
|
||||
vec![
|
||||
0x10, 0x00, 0x00, 0x0C, 0x48, 0x14, 0x10, 0x03, 0x00, 0x00, 0x00, 0x4D,
|
||||
]
|
||||
}
|
||||
|
||||
// Synthetic AACS 2.x MKB with a 0x82 + 0x83 record pair.
|
||||
fn mkb_with_variant() -> Vec<u8> {
|
||||
let mut m = mkb_classical();
|
||||
m.extend_from_slice(&[0x82, 0x00, 0x00, 0x14]);
|
||||
m.extend_from_slice(&[0xEE; 16]);
|
||||
m.extend_from_slice(&[0x83, 0x00, 0x00, 0x14]);
|
||||
m.extend_from_slice(&[0x55; 16]);
|
||||
m
|
||||
}
|
||||
|
||||
// Synthetic scrambled DVD sector — byte 0x14 carries the CSS
|
||||
// scramble flag in bits 4-5.
|
||||
fn scrambled_dvd_sector() -> Vec<u8> {
|
||||
let mut s = vec![0u8; 2048];
|
||||
s[0x14] = 0x30;
|
||||
s
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_returns_none_for_unencrypted() {
|
||||
let probe = DrmProbe {
|
||||
dvd_sample_sector: None,
|
||||
content_cert: None,
|
||||
mkb: None,
|
||||
};
|
||||
assert_eq!(DrmScheme::detect(&probe), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_returns_css_for_scrambled_dvd() {
|
||||
let sector = scrambled_dvd_sector();
|
||||
let probe = DrmProbe {
|
||||
dvd_sample_sector: Some(§or),
|
||||
content_cert: None,
|
||||
mkb: None,
|
||||
};
|
||||
assert_eq!(DrmScheme::detect(&probe), Some(DrmScheme::Css));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_returns_aacs10_for_type0_cert() {
|
||||
let c = cert(0x00);
|
||||
let probe = DrmProbe {
|
||||
dvd_sample_sector: None,
|
||||
content_cert: Some(&c),
|
||||
mkb: None,
|
||||
};
|
||||
assert_eq!(DrmScheme::detect(&probe), Some(DrmScheme::Aacs10));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_returns_aacs20_for_type1_cert_no_variant() {
|
||||
let c = cert(0x01);
|
||||
let mkb = mkb_classical();
|
||||
let probe = DrmProbe {
|
||||
dvd_sample_sector: None,
|
||||
content_cert: Some(&c),
|
||||
mkb: Some(&mkb),
|
||||
};
|
||||
assert_eq!(DrmScheme::detect(&probe), Some(DrmScheme::Aacs20));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_returns_aacs21_for_type1_cert_with_variant() {
|
||||
let c = cert(0x01);
|
||||
let mkb = mkb_with_variant();
|
||||
let probe = DrmProbe {
|
||||
dvd_sample_sector: None,
|
||||
content_cert: Some(&c),
|
||||
mkb: Some(&mkb),
|
||||
};
|
||||
assert_eq!(DrmScheme::detect(&probe), Some(DrmScheme::Aacs21));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_returns_aacs20_when_mkb_absent() {
|
||||
// Type-1 cert but no MKB to upgrade with -> Aacs20.
|
||||
let c = cert(0x01);
|
||||
let probe = DrmProbe {
|
||||
dvd_sample_sector: None,
|
||||
content_cert: Some(&c),
|
||||
mkb: None,
|
||||
};
|
||||
assert_eq!(DrmScheme::detect(&probe), Some(DrmScheme::Aacs20));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_aacs21_returns_none() {
|
||||
// The Aacs21 dispatch arm is commented out; load() must
|
||||
// return None until KCD validation lands.
|
||||
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),
|
||||
css: None,
|
||||
};
|
||||
assert!(DrmScheme::Aacs21.load(&mut ctx).is_none());
|
||||
}
|
||||
|
||||
/// Exercises the V21 helper directly. Gated `#[ignore]` because
|
||||
/// the chain reaches `MediaKeyVariantError::VariantsTableUnavailable`
|
||||
/// without a real Variant-scheme disc to fix the per-uv table
|
||||
/// layout against — running it here would assert only the
|
||||
/// not-yet-wired error code. Kept as a wiring smoke-test for
|
||||
/// future enablement.
|
||||
#[test]
|
||||
#[ignore]
|
||||
fn resolve_keys_v21_helper_exists() {
|
||||
let uk_ro = vec![0u8; 256];
|
||||
let vid = [0xAAu8; 16];
|
||||
let keydb = aacs::KeyDb::empty();
|
||||
let providers: &[&dyn aacs::KeyProvider] = &[&keydb];
|
||||
let mkb = mkb_with_variant();
|
||||
let ctx = aacs::ResolveContext {
|
||||
unit_key_ro: &uk_ro,
|
||||
content_cert: None,
|
||||
volume_id: &vid,
|
||||
providers,
|
||||
mkb: Some(&mkb),
|
||||
};
|
||||
// Just confirm the symbol is callable; we don't assert on the
|
||||
// result.
|
||||
let _ = aacs::resolve_keys_v21(&ctx);
|
||||
}
|
||||
}
|
||||
@@ -71,6 +71,14 @@ pub const E_AACS_VID_MAC: u16 = 7010;
|
||||
pub const E_AACS_DATA_KEY: u16 = 7011;
|
||||
pub const E_DECRYPT_FAILED: u16 = 7013;
|
||||
pub const E_CSS_AUTH_FAILED: u16 = 7014;
|
||||
pub const E_AACS_HOST_CERT_REJECTED: u16 = 7015;
|
||||
pub const E_AACS_RAW_READ_UNSUPPORTED: u16 = 7016;
|
||||
pub const E_AACS_VID_UNAVAILABLE: u16 = 7017;
|
||||
pub const E_AACS_MK_UNAVAILABLE: u16 = 7018;
|
||||
pub const E_AACS_VUK_NOT_IN_KEYDB: u16 = 7019;
|
||||
pub const E_DRIVE_PROFILE_MISSING: u16 = 7020;
|
||||
pub const E_VID_CDB_UNAVAILABLE: u16 = 7021;
|
||||
pub const E_NO_DISC_KEY: u16 = 7022;
|
||||
|
||||
// Keydb (8xxx)
|
||||
pub const E_KEYDB_CONNECT: u16 = 8000;
|
||||
@@ -222,6 +230,40 @@ pub enum Error {
|
||||
AacsDataKey,
|
||||
DecryptFailed,
|
||||
CssAuthFailed,
|
||||
/// Host certificate rejected by the drive's revocation list (HRL hit).
|
||||
/// All available host certs failed mutual auth on this drive.
|
||||
AacsHostCertRejected,
|
||||
/// Drive cannot be put into raw-read mode and standard AACS cert
|
||||
/// auth failed. No path to decryption remains.
|
||||
AacsRawReadUnsupported,
|
||||
/// Volume ID could not be retrieved from the drive (neither via cert
|
||||
/// auth nor via the alternate VID read path). Downstream of step 1
|
||||
/// of the AACS chain.
|
||||
AacsVidUnavailable,
|
||||
/// No available path produced a Media Key (no MK+VID in keydb, no
|
||||
/// PK match, no DK derivation).
|
||||
AacsMkUnavailable,
|
||||
/// Disc-hash lookup in the keydb missed and no other path is
|
||||
/// available (typically because VID is missing).
|
||||
AacsVukNotInKeydb,
|
||||
/// Drive identity did not match any bundled profile; per-drive CDB
|
||||
/// templates aren't available so the OEM VID retrieval path can't
|
||||
/// run.
|
||||
DriveProfileMissing,
|
||||
/// Drive's profile is present but doesn't carry a VID-retrieval CDB
|
||||
/// template (older profile blob, or a drive class without an OEM
|
||||
/// VID path).
|
||||
VidCdbUnavailable,
|
||||
/// The disc is AACS-encrypted and decryption was requested, but key
|
||||
/// resolution produced no usable key for it — so muxing would emit
|
||||
/// undecryptable garbage. Distinct from [`Error::KeydbLoad`] (no keydb
|
||||
/// file at all): a keydb may be present but lack an entry for this disc.
|
||||
/// `disc_hash` is the 40-hex SHA1 of `Unit_Key_RO.inf` (no `0x` prefix)
|
||||
/// so the application can name the disc; empty if the hash wasn't
|
||||
/// captured at scan.
|
||||
NoDiscKey {
|
||||
disc_hash: String,
|
||||
},
|
||||
|
||||
// Keydb (8xxx)
|
||||
KeydbConnect {
|
||||
@@ -307,6 +349,14 @@ impl Error {
|
||||
Error::AacsDataKey => E_AACS_DATA_KEY,
|
||||
Error::DecryptFailed => E_DECRYPT_FAILED,
|
||||
Error::CssAuthFailed => E_CSS_AUTH_FAILED,
|
||||
Error::AacsHostCertRejected => E_AACS_HOST_CERT_REJECTED,
|
||||
Error::AacsRawReadUnsupported => E_AACS_RAW_READ_UNSUPPORTED,
|
||||
Error::AacsVidUnavailable => E_AACS_VID_UNAVAILABLE,
|
||||
Error::AacsMkUnavailable => E_AACS_MK_UNAVAILABLE,
|
||||
Error::AacsVukNotInKeydb => E_AACS_VUK_NOT_IN_KEYDB,
|
||||
Error::DriveProfileMissing => E_DRIVE_PROFILE_MISSING,
|
||||
Error::VidCdbUnavailable => E_VID_CDB_UNAVAILABLE,
|
||||
Error::NoDiscKey { .. } => E_NO_DISC_KEY,
|
||||
Error::KeydbConnect { .. } => E_KEYDB_CONNECT,
|
||||
Error::KeydbHttp { .. } => E_KEYDB_HTTP,
|
||||
Error::KeydbInvalid => E_KEYDB_INVALID,
|
||||
@@ -433,6 +483,7 @@ impl std::fmt::Display for Error {
|
||||
Error::StreamUrlMissingPort { addr } => write!(f, "E{}: {}", self.code(), addr),
|
||||
Error::PesFrameTooLarge { size } => write!(f, "E{}: {}", self.code(), size),
|
||||
Error::IsoTooLarge { path } => write!(f, "E{}: {}", self.code(), path),
|
||||
Error::NoDiscKey { disc_hash } => write!(f, "E{}: {}", self.code(), disc_hash),
|
||||
_ => write!(f, "E{}", self.code()),
|
||||
}
|
||||
}
|
||||
|
||||
+14
@@ -71,6 +71,20 @@ impl Default for Halt {
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared poll interval for halt-aware loops.
|
||||
///
|
||||
/// `bounded_syscall` checks the cancellation flag and the deadline
|
||||
/// every [`POLL_INTERVAL`] while blocked on a worker; the same
|
||||
/// cadence governs `Pipeline::send_with_halt`'s `try_send` retry.
|
||||
/// 250 ms is the sweet spot between responsiveness (operator presses
|
||||
/// Stop, sees it take effect within ~quarter-second) and waste
|
||||
/// (atomic load + clock read is cheap but not free at thousands of
|
||||
/// hertz).
|
||||
///
|
||||
/// Centralised here so the half-dozen halt-polling loops across `io`
|
||||
/// can't drift apart silently.
|
||||
pub const POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(250);
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
//! Bounded-syscall primitive: run a (potentially-blocking) operation
|
||||
//! on a worker thread, with a hard wall-clock deadline and an optional
|
||||
//! cooperative [`Halt`] poll. The calling thread is never trapped
|
||||
//! inside a kernel call.
|
||||
//!
|
||||
//! ## Why this exists
|
||||
//!
|
||||
//! [`crate::halt::Halt`] is cooperative: callers poll
|
||||
//! `is_cancelled()`. It cannot reach inside a syscall the kernel
|
||||
//! currently owns the thread for — `libc::sync_file_range`,
|
||||
//! `libc::fsync`, `File::write` on NFS, and so on. `/api/stop` from
|
||||
//! autorip therefore can't unstick a thread sitting in such a syscall.
|
||||
//!
|
||||
//! [`bounded_syscall`] is the escape hatch: it runs `op` on a fresh
|
||||
//! worker thread, then `recv_timeout`s on a rendezvous channel for the
|
||||
//! result. The wait is broken into ~250 ms slices so the calling
|
||||
//! thread can poll the supplied [`Halt`] in between. If the deadline
|
||||
//! elapses or the halt fires, the worker is intentionally leaked — the
|
||||
//! syscall will unwind whenever the kernel decides, or at process
|
||||
//! exit, but the caller is free to fall back to a degraded code path
|
||||
//! (skip the sync, log loudly, etc.).
|
||||
//!
|
||||
//! ## Trade-offs
|
||||
//!
|
||||
//! - **Thread per call.** Cheap (`std::thread::spawn` is < 100 µs on
|
||||
//! Linux/macOS), but not free. Use on coarse-grained finalisation
|
||||
//! syscalls (`sync_all`, `sync_file_range(WAIT_AFTER)`), not on hot
|
||||
//! inner-loop writes.
|
||||
//! - **Leak on timeout.** A wedged syscall keeps a kernel slot and a
|
||||
//! user-space thread around for the rest of the process's life.
|
||||
//! Bounded by the number of independent rip/mux sessions, which is
|
||||
//! one per disc. The alternative — trapping the caller forever —
|
||||
//! defeats the entire purpose of `/api/stop`.
|
||||
//! - **Halt granularity ~250 ms.** Halt observation is not instant;
|
||||
//! it's the worst-case latency of the `recv_timeout` slice. Good
|
||||
//! enough for human-driven stop requests; not suitable for hard
|
||||
//! real-time deadlines.
|
||||
//!
|
||||
//! ## Single source of truth
|
||||
//!
|
||||
//! Do NOT inline this pattern. Every blocking-syscall wrapper in the
|
||||
//! rip + mux pipeline calls this helper, so changes (e.g. swapping the
|
||||
//! channel impl, adjusting the poll slice, adding metrics) land in one
|
||||
//! place.
|
||||
//!
|
||||
//! ## Platform
|
||||
//!
|
||||
//! Pure `std::thread` + `std::sync::mpsc`. No `cfg(target_os)` needed
|
||||
//! here — the helper itself is platform-agnostic. Callers that wrap
|
||||
//! Linux-only syscalls (`sync_file_range`) still need their own
|
||||
//! `#[cfg(target_os = "linux")]` gates; this helper does not.
|
||||
|
||||
use std::sync::mpsc::{RecvTimeoutError, sync_channel};
|
||||
use std::thread;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use crate::halt::{Halt, POLL_INTERVAL};
|
||||
|
||||
/// Failure outcome from a bounded syscall wrapper.
|
||||
#[derive(Debug)]
|
||||
pub(crate) enum BoundedError {
|
||||
/// The user-visible halt token fired during the wait. The worker
|
||||
/// thread is intentionally leaked — the caller should fall back to
|
||||
/// a degraded code path rather than waiting on the syscall to
|
||||
/// return.
|
||||
Halted,
|
||||
/// 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.
|
||||
WorkerLost,
|
||||
}
|
||||
|
||||
/// Run a (potentially-blocking) operation on a worker thread with a
|
||||
/// deadline and an optional cooperative halt-token poll. Returns the
|
||||
/// operation's result if it completes within `timeout`; otherwise one
|
||||
/// of [`BoundedError::Halted`] / [`BoundedError::Timeout`] /
|
||||
/// [`BoundedError::WorkerLost`].
|
||||
///
|
||||
/// On `Halted` / `Timeout` the worker thread is intentionally leaked:
|
||||
/// the syscall will unwind whenever the kernel decides, or when the
|
||||
/// process exits. The calling thread is never trapped inside a kernel
|
||||
/// call.
|
||||
///
|
||||
/// `halt` is polled at [`POLL_INTERVAL`] granularity. Pass `None` for
|
||||
/// callers that don't (yet) have a halt token plumbed through —
|
||||
/// behaviour degrades to deadline-only, matching the 0.20.5
|
||||
/// `wait_after_with_timeout` shape this helper generalises.
|
||||
///
|
||||
/// `op` returns `R: Send + 'static`. The closure must own everything
|
||||
/// it touches because it may outlive this call (timeout / halt cases).
|
||||
pub(crate) fn bounded_syscall<F, R>(
|
||||
halt: Option<&Halt>,
|
||||
timeout: Duration,
|
||||
op: F,
|
||||
) -> Result<R, BoundedError>
|
||||
where
|
||||
F: FnOnce() -> R + Send + 'static,
|
||||
R: Send + 'static,
|
||||
{
|
||||
// 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 /
|
||||
// halt path the receiver is dropped and the worker's send
|
||||
// returns Err, which the worker ignores.
|
||||
let (tx, rx) = sync_channel::<R>(0);
|
||||
let _ = thread::Builder::new()
|
||||
.name("freemkv-bounded-syscall".into())
|
||||
.spawn(move || {
|
||||
// Ignore the send error: if we time out (or get halted)
|
||||
// before the worker finishes, the receiver is dropped
|
||||
// and `tx.send` returns Err. Either way, the worker has
|
||||
// nothing more to do.
|
||||
let _ = tx.send(op());
|
||||
});
|
||||
|
||||
let deadline = Instant::now() + timeout;
|
||||
loop {
|
||||
let now = Instant::now();
|
||||
let remaining = deadline.saturating_duration_since(now);
|
||||
let slice = remaining.min(POLL_INTERVAL);
|
||||
match rx.recv_timeout(slice) {
|
||||
Ok(v) => return Ok(v),
|
||||
Err(RecvTimeoutError::Timeout) => {
|
||||
if let Some(h) = halt {
|
||||
if h.is_cancelled() {
|
||||
return Err(BoundedError::Halted);
|
||||
}
|
||||
}
|
||||
if Instant::now() >= deadline {
|
||||
return Err(BoundedError::Timeout);
|
||||
}
|
||||
// Otherwise: another slice.
|
||||
}
|
||||
Err(RecvTimeoutError::Disconnected) => {
|
||||
// Worker thread spawn failed, or it panicked before
|
||||
// sending. Caller treats this as "no syscall ran" —
|
||||
// typically a no-op + log.
|
||||
return Err(BoundedError::WorkerLost);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
#[test]
|
||||
fn op_completes_quickly() {
|
||||
let r = bounded_syscall(None, Duration::from_secs(2), || 42u32);
|
||||
assert!(matches!(r, Ok(42)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn op_exceeds_timeout() {
|
||||
// Op sleeps longer than the deadline → Timeout.
|
||||
let r = bounded_syscall(None, Duration::from_millis(300), || {
|
||||
thread::sleep(Duration::from_secs(2));
|
||||
0u32
|
||||
});
|
||||
assert!(matches!(r, Err(BoundedError::Timeout)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn halt_fires_during_wait() {
|
||||
let halt = Halt::new();
|
||||
let halt2 = halt.clone();
|
||||
// Flip the halt from a side thread after ~300 ms — long
|
||||
// enough that the receive loop has rolled at least one
|
||||
// 250 ms slice and is sitting in `recv_timeout` again when
|
||||
// the bit flips.
|
||||
thread::spawn(move || {
|
||||
thread::sleep(Duration::from_millis(300));
|
||||
halt2.cancel();
|
||||
});
|
||||
let r = bounded_syscall(Some(&halt), Duration::from_secs(5), || {
|
||||
thread::sleep(Duration::from_secs(5));
|
||||
0u32
|
||||
});
|
||||
assert!(matches!(r, Err(BoundedError::Halted)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn worker_panics() {
|
||||
// Worker panics → sender drops without sending → recv sees
|
||||
// Disconnected → WorkerLost. We use an explicit panic in the
|
||||
// op closure rather than `panic!()` from inside the channel
|
||||
// machinery; the spawned thread's panic is contained (no
|
||||
// process abort) because we don't `.join()` it.
|
||||
let r = bounded_syscall(None, Duration::from_secs(2), || -> u32 {
|
||||
panic!("intentional test panic");
|
||||
});
|
||||
assert!(matches!(r, Err(BoundedError::WorkerLost)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn halt_already_set_before_call_still_returns_halted() {
|
||||
// Halt observed on the very first poll slice. The op blocks
|
||||
// forever; we must not wait the full timeout to notice the
|
||||
// halt is already set.
|
||||
let halt = Halt::new();
|
||||
halt.cancel();
|
||||
let started = Instant::now();
|
||||
let r = bounded_syscall(Some(&halt), Duration::from_secs(10), || {
|
||||
thread::sleep(Duration::from_secs(10));
|
||||
0u32
|
||||
});
|
||||
assert!(matches!(r, Err(BoundedError::Halted)));
|
||||
// Should bail out within ~1 s; allow 2 s of slack for slow
|
||||
// CI hosts.
|
||||
assert!(
|
||||
started.elapsed() < Duration::from_secs(2),
|
||||
"halt-already-set took {:?}",
|
||||
started.elapsed()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ok_path_takes_no_halt_token() {
|
||||
// Sanity: the `None` halt path is the documented zero-config
|
||||
// form (matches the 0.20.5 `wait_after_with_timeout`
|
||||
// behaviour). Op returns immediately; we must observe Ok.
|
||||
let flag = Arc::new(AtomicBool::new(false));
|
||||
let f2 = flag.clone();
|
||||
let r = bounded_syscall(None, Duration::from_secs(2), move || {
|
||||
f2.store(true, Ordering::Relaxed);
|
||||
"ok"
|
||||
});
|
||||
assert!(matches!(r, Ok("ok")));
|
||||
assert!(flag.load(Ordering::Relaxed));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,327 @@
|
||||
//! 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());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
//! `BytePrefetcher` — `std::io::Read` analogue of
|
||||
//! [`crate::sector::PrefetchedSectorSource`].
|
||||
//!
|
||||
//! Spawns a producer thread that fills a bounded pool of `Vec<u8>`
|
||||
//! chunks from the underlying reader and ships them through a
|
||||
//! channel; the consumer pulls filled chunks, uses them, and sends
|
||||
//! the empty `Vec<u8>` back through a recycle channel so the
|
||||
//! producer can re-fill in place. Result: zero allocations and zero
|
||||
//! cross-thread frees in the steady-state hot loop.
|
||||
//!
|
||||
//! 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`.
|
||||
|
||||
use crate::halt::Halt;
|
||||
use crossbeam_channel::{Receiver, Sender, bounded};
|
||||
use std::io::Read;
|
||||
use std::thread::JoinHandle;
|
||||
|
||||
/// Items flowing through the forward channel.
|
||||
pub type Batch = std::io::Result<Vec<u8>>;
|
||||
|
||||
/// Forward channel depth — how many filled buffers the producer can
|
||||
/// stay ahead by. Two is enough to absorb a moderate consumer stall
|
||||
/// without piling up bytes.
|
||||
const FORWARD_DEPTH: usize = 2;
|
||||
|
||||
/// Recycle channel depth = forward + 1 so the producer always has at
|
||||
/// least one buffer to fill while the consumer holds one.
|
||||
const RECYCLE_DEPTH: usize = FORWARD_DEPTH + 1;
|
||||
|
||||
/// Default chunk size — 16 MiB matches the ISO-mux sector batch and
|
||||
/// is large enough that per-chunk overhead is amortised; small
|
||||
/// enough that the in-flight memory footprint stays bounded.
|
||||
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.
|
||||
pub struct PrefetchShell {
|
||||
producer: Option<JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl Drop for PrefetchShell {
|
||||
fn drop(&mut self) {
|
||||
if let Some(h) = self.producer.take() {
|
||||
let _ = h.join();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawned byte prefetcher. Drop joins the producer thread.
|
||||
pub struct BytePrefetcher {
|
||||
rx: Receiver<Batch>,
|
||||
recycle_tx: Sender<Vec<u8>>,
|
||||
producer: Option<JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl BytePrefetcher {
|
||||
/// Spawn the producer thread. `reader` must be `Send` because it
|
||||
/// moves into the thread. `chunk_bytes` is the size of each
|
||||
/// recycled buffer; pick the natural batch size of the
|
||||
/// downstream demuxer (16 MiB for the BD-TS mux pipeline).
|
||||
pub fn new<R: Read + Send + 'static>(
|
||||
mut reader: R,
|
||||
chunk_bytes: usize,
|
||||
halt: Option<Halt>,
|
||||
) -> Self {
|
||||
let (tx, rx) = bounded::<Batch>(FORWARD_DEPTH);
|
||||
let (recycle_tx, recycle_rx) = bounded::<Vec<u8>>(RECYCLE_DEPTH);
|
||||
|
||||
// Seed the recycle pool. Without these the first
|
||||
// `recycle_rx.recv()` would block forever (no consumer has
|
||||
// returned a buffer yet).
|
||||
for _ in 0..RECYCLE_DEPTH {
|
||||
let _ = recycle_tx.send(vec![0u8; chunk_bytes]);
|
||||
}
|
||||
|
||||
let producer = std::thread::Builder::new()
|
||||
.name("freemkv-byte-prefetch".into())
|
||||
.spawn(move || {
|
||||
loop {
|
||||
if halt.as_ref().map(|h| h.is_cancelled()).unwrap_or(false) {
|
||||
return;
|
||||
}
|
||||
let mut buf = match recycle_rx.recv() {
|
||||
Ok(b) => b,
|
||||
Err(_) => return, // consumer dropped both channels
|
||||
};
|
||||
// Re-expose the full extent (previous iteration
|
||||
// may have truncated after a short read).
|
||||
if buf.len() < chunk_bytes {
|
||||
buf.resize(chunk_bytes, 0);
|
||||
} else {
|
||||
// SAFETY: capacity is at least chunk_bytes
|
||||
// after construction.
|
||||
unsafe { buf.set_len(chunk_bytes) };
|
||||
}
|
||||
// Read up to one full chunk. Short reads are
|
||||
// valid and common — pipe `truncate` so the
|
||||
// consumer sees only the bytes that arrived.
|
||||
let n = match reader.read(&mut buf[..]) {
|
||||
Ok(0) => return, // EOF — drop tx, consumer sees RecvError
|
||||
Ok(n) => n,
|
||||
Err(e) => {
|
||||
let _ = tx.send(Err(e));
|
||||
return;
|
||||
}
|
||||
};
|
||||
buf.truncate(n);
|
||||
if tx.send(Ok(buf)).is_err() {
|
||||
return; // consumer dropped
|
||||
}
|
||||
}
|
||||
})
|
||||
.expect("freemkv-byte-prefetch thread spawn failed");
|
||||
|
||||
Self {
|
||||
rx,
|
||||
recycle_tx,
|
||||
producer: Some(producer),
|
||||
}
|
||||
}
|
||||
|
||||
/// Peel off the channels for zero-copy pipeline consumption. The
|
||||
/// caller (typically [`crate::mux::demux_thread::DemuxThread`])
|
||||
/// 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);
|
||||
(rx, recycle, PrefetchShell { producer })
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for BytePrefetcher {
|
||||
fn drop(&mut self) {
|
||||
if let Some(h) = self.producer.take() {
|
||||
let _ = h.join();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
//! Linux read-side platform hooks: sequential-access hint at open +
|
||||
//! periodic page-cache eviction during streaming reads.
|
||||
//!
|
||||
//! ## Why both
|
||||
//!
|
||||
//! `POSIX_FADV_SEQUENTIAL` at open widens the kernel's readahead window
|
||||
//! so each pread aggregates into fewer NFS round-trips. `DONTNEED` on
|
||||
//! the consumed window (called periodically by the caller) drops the
|
||||
//! already-read pages from the page cache so an 85 GB streaming ISO
|
||||
//! read doesn't fill memory and starve concurrent writes (the MKV
|
||||
//! output during mux). Together they mirror the write-side
|
||||
//! WritebackPipeline's policy.
|
||||
//!
|
||||
//! ## History
|
||||
//!
|
||||
//! Pre-Phase-1 (0.20.7 baseline) had both. Phase 1's introduction of
|
||||
//! `FileSectorSource` silently dropped the read-side DONTNEED, and
|
||||
//! 0.21.2's revert of `SEQUENTIAL` (mistakenly attributing a regression
|
||||
//! to it) removed the hint. Net effect: 85 GB of ISO reads pinned in
|
||||
//! the page cache + no readahead widening → mux throughput collapse
|
||||
//! from 18 MB/s historical to 2.7-8 MB/s on 0.21.x. Restored in 0.21.6.
|
||||
|
||||
use std::fs::File;
|
||||
use std::os::unix::io::AsRawFd;
|
||||
|
||||
pub(super) fn hint_sequential(file: &File, _len_bytes: u64) {
|
||||
// Best-effort: return value ignored. A fadvise failure has no
|
||||
// user-observable consequence.
|
||||
unsafe {
|
||||
libc::posix_fadvise(file.as_raw_fd(), 0, 0, libc::POSIX_FADV_SEQUENTIAL);
|
||||
}
|
||||
}
|
||||
|
||||
/// Drop pages in the half-open byte range `[start, start+len)` from
|
||||
/// the page cache. Called periodically by `read_sectors` to bound the
|
||||
/// read-side page cache pressure.
|
||||
pub(super) fn drop_window(file: &File, start: u64, len: u64) {
|
||||
unsafe {
|
||||
libc::posix_fadvise(
|
||||
file.as_raw_fd(),
|
||||
start as i64,
|
||||
len as i64,
|
||||
libc::POSIX_FADV_DONTNEED,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Async-prefetch `len` bytes at `offset` into the page cache. The
|
||||
/// kernel `readahead(2)` syscall queues the I/O and returns
|
||||
/// immediately — it does NOT wait for completion. Called right after
|
||||
/// each consumed read so the next batch's I/O overlaps with the
|
||||
/// caller's processing of the current batch (decrypt + demux + mux).
|
||||
///
|
||||
/// Without this hint, with a synchronous demux consumer running at
|
||||
/// ~50 MB/s and a single-spindle disk capable of ~150 MB/s, the disk
|
||||
/// sits idle ~70% of each iteration because kernel readahead alone
|
||||
/// (capped at `/sys/block/<dev>/queue/read_ahead_kb`, default 128 KB)
|
||||
/// can only pre-stage a tiny slice of the next batch. An explicit
|
||||
/// `readahead()` of the same size as the current batch tells the
|
||||
/// kernel to queue the full next-batch read now.
|
||||
pub(super) fn prefetch(file: &File, offset: u64, len: u64) {
|
||||
unsafe {
|
||||
libc::readahead(file.as_raw_fd(), offset as i64, len as usize);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
//! macOS: hint the kernel to prefetch a generous chunk. macOS has no
|
||||
//! direct `POSIX_FADV_SEQUENTIAL` equivalent; the idiomatic hint is
|
||||
//! `fcntl(F_RDADVISE, &radvisory)` describing the byte range you
|
||||
//! intend to read soon. We point it at the whole file (clamped to a
|
||||
//! ceiling so a multi-TB ISO doesn't ask the kernel to prefetch
|
||||
//! everything at once).
|
||||
|
||||
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.
|
||||
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 {
|
||||
ra_offset: 0,
|
||||
ra_count: bytes as libc::c_int,
|
||||
};
|
||||
// Best-effort.
|
||||
unsafe {
|
||||
libc::fcntl(file.as_raw_fd(), F_RDADVISE, &mut ra);
|
||||
}
|
||||
}
|
||||
|
||||
/// macOS has no direct `POSIX_FADV_DONTNEED` equivalent for a byte
|
||||
/// range. `fcntl(F_NOCACHE)` would disable caching globally on the fd
|
||||
/// (too coarse — we want the unread region to still benefit). Best
|
||||
/// approximation: no-op. macOS's unified buffer cache is generally
|
||||
/// less prone to the pin-everything pathology that triggers the
|
||||
/// regression on Linux NFS clients.
|
||||
pub(super) fn drop_window(_file: &File, _start: u64, _len: u64) {}
|
||||
|
||||
/// Async-prefetch the byte range `[offset, offset+len)`. macOS uses
|
||||
/// the same `fcntl(F_RDADVISE, &radvisory)` primitive as the open-
|
||||
/// time sequential hint, just targeted at a moving window instead of
|
||||
/// the whole file. The kernel queues I/O for the requested range and
|
||||
/// 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 {
|
||||
ra_offset: offset as libc::off_t,
|
||||
ra_count: bytes as libc::c_int,
|
||||
};
|
||||
unsafe {
|
||||
libc::fcntl(file.as_raw_fd(), F_RDADVISE, &mut ra);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,361 @@
|
||||
//! [`FileSectorSource`] — read 2048-byte sectors from an ISO file on
|
||||
//! disk via direct `seek + read_exact` (`pread`-equivalent) calls,
|
||||
//! letting the kernel's own readahead policy manage prefetch.
|
||||
//!
|
||||
//! ## Why no app-level buffer
|
||||
//!
|
||||
//! Pre-0.21.3 this source held a 32 MiB (later 4 MiB) read-ahead
|
||||
//! buffer to amortise per-sector NFS round-trips. Empirically that
|
||||
//! buffer hurt: 32 MiB refills bursted the NFS TCP connection hard
|
||||
//! enough to starve the concurrent writer, and even a 4 MiB window
|
||||
//! gave the kernel less freedom to pipeline reads with writes. Direct
|
||||
//! pread per call lets Linux's readahead widen as it detects the
|
||||
//! sequential pattern, and naturally interleaves with writeback.
|
||||
//!
|
||||
//! ## DONTNEED on the consumed window
|
||||
//!
|
||||
//! 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.
|
||||
//!
|
||||
//! ## Platform open hint
|
||||
//!
|
||||
//! On `open()` each platform issues its "sequential access expected"
|
||||
//! hint so OS-level readahead widens. The hint and the DONTNEED call
|
||||
//! live in per-OS sibling modules ([`linux::hint_sequential`] et al.)
|
||||
//! — no inline `#[cfg]` in this file.
|
||||
//!
|
||||
//! ## Read-ahead prefetch
|
||||
//!
|
||||
//! After every consumed read we issue an OS-level prefetch hint for
|
||||
//! the next equivalent-sized window (`platform::prefetch`). The
|
||||
//! kernel queues that I/O asynchronously and returns immediately, so
|
||||
//! the next batch's read overlaps with the caller's processing of
|
||||
//! the current batch (decrypt + demux + mux). Without this the disk
|
||||
//! sits idle ~70% of each iteration because kernel SEQUENTIAL
|
||||
//! readahead alone (capped at `read_ahead_kb`, default 128 KB) is
|
||||
//! far smaller than our 16 MiB app-level batch.
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
mod linux;
|
||||
#[cfg(target_os = "macos")]
|
||||
mod macos;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
|
||||
mod other;
|
||||
#[cfg(target_os = "windows")]
|
||||
mod windows;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
use linux as platform;
|
||||
#[cfg(target_os = "macos")]
|
||||
use macos as platform;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
|
||||
use other as platform;
|
||||
#[cfg(target_os = "windows")]
|
||||
use windows as platform;
|
||||
|
||||
use std::fs::File;
|
||||
use std::io::{Read, Seek, SeekFrom};
|
||||
use std::path::Path;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
use crate::sector::SectorSource;
|
||||
|
||||
const SECTOR_SIZE: usize = 2048;
|
||||
|
||||
/// Bytes-read threshold per `posix_fadvise(DONTNEED)` drop on the
|
||||
/// read side. Mirrors `WRITEBACK_CHUNK_BYTES` so the read-side page
|
||||
/// cache stays bounded the same way the write side does.
|
||||
///
|
||||
/// 32 MiB is the empirically tuned value on the rip1 test bed (single
|
||||
/// 7200rpm HDD via SATA): smaller windows (8 / 16 MiB) shorten the
|
||||
/// kernel-readahead overlap and slow the producer; larger windows
|
||||
/// (64 / 128 MiB) let the page cache pin enough of the ISO to
|
||||
/// pressure concurrent writes. Override via `FREEMKV_READ_DROP_CHUNK_MIB`.
|
||||
const READ_DROP_CHUNK_BYTES_DEFAULT: u64 = 32 * 1024 * 1024;
|
||||
|
||||
fn read_drop_chunk_bytes() -> u64 {
|
||||
std::env::var("FREEMKV_READ_DROP_CHUNK_MIB")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<u64>().ok())
|
||||
.filter(|&n| n > 0)
|
||||
.map(|n| n * 1024 * 1024)
|
||||
.unwrap_or(READ_DROP_CHUNK_BYTES_DEFAULT)
|
||||
}
|
||||
|
||||
/// SectorSource backed by a file (ISO image). Every `read_sectors`
|
||||
/// call is a direct `seek + read_exact` against the underlying file
|
||||
/// — kernel readahead handles prefetch, and every
|
||||
/// [`READ_DROP_CHUNK_BYTES_DEFAULT`] bytes of consumed data the
|
||||
/// platform's `DONTNEED` hook drops the consumed window from the
|
||||
/// page cache to bound memory pressure.
|
||||
pub struct FileSectorSource {
|
||||
file: File,
|
||||
/// Total file size in sectors. Constant after construction;
|
||||
/// surfaced via [`SectorSource::capacity_sectors`].
|
||||
capacity: u32,
|
||||
/// Bytes read since the last DONTNEED drop. Drives the per-
|
||||
/// [`read_drop_chunk_bytes`] page-cache eviction in read_sectors.
|
||||
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.
|
||||
drop_window_start: u64,
|
||||
/// Cached drop chunk size (resolved from env once at open).
|
||||
drop_chunk_bytes: u64,
|
||||
}
|
||||
|
||||
impl FileSectorSource {
|
||||
/// Open an existing ISO file for reading. Capacity is derived
|
||||
/// from `metadata().len() / 2048`. Returns
|
||||
/// [`Error::IsoTooLarge`] if the file would exceed the 32-bit
|
||||
/// LBA address space (~8 TB).
|
||||
///
|
||||
/// 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();
|
||||
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;
|
||||
|
||||
// Best-effort sequential hint. Ignored on platforms without
|
||||
// an equivalent primitive (or where the API exists but the
|
||||
// FS doesn't honour it).
|
||||
platform::hint_sequential(&file, len);
|
||||
|
||||
Ok(Self {
|
||||
file,
|
||||
capacity,
|
||||
bytes_read_since_drop: 0,
|
||||
drop_window_start: 0,
|
||||
drop_chunk_bytes: read_drop_chunk_bytes(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl SectorSource for FileSectorSource {
|
||||
fn capacity_sectors(&self) -> u32 {
|
||||
self.capacity
|
||||
}
|
||||
|
||||
fn read_sectors(
|
||||
&mut self,
|
||||
lba: u32,
|
||||
count: u16,
|
||||
out: &mut [u8],
|
||||
_recovery: bool,
|
||||
) -> Result<usize> {
|
||||
let count = count as u32;
|
||||
let bytes = count as usize * SECTOR_SIZE;
|
||||
debug_assert!(
|
||||
out.len() >= bytes,
|
||||
"FileSectorSource::read_sectors: out len {} < requested {}",
|
||||
out.len(),
|
||||
bytes
|
||||
);
|
||||
if count == 0 {
|
||||
return Ok(0);
|
||||
}
|
||||
let offset = lba as u64 * SECTOR_SIZE as u64;
|
||||
self.file
|
||||
.seek(SeekFrom::Start(offset))
|
||||
.map_err(|e| Error::IoError { source: e })?;
|
||||
self.file
|
||||
.read_exact(&mut out[..bytes])
|
||||
.map_err(|e| Error::IoError { source: e })?;
|
||||
|
||||
// Queue the next batch's read with the kernel before the
|
||||
// caller starts processing what we just returned. readahead()
|
||||
// is non-blocking — it queues I/O and returns, so the kernel
|
||||
// pulls those pages into cache while the consumer (decrypt +
|
||||
// demux + mux) runs. Next read_sectors call hits a warm cache.
|
||||
platform::prefetch(&self.file, offset + bytes as u64, bytes as u64);
|
||||
|
||||
// Periodic page-cache eviction on the read side. Without
|
||||
// this, an 85 GB streaming ISO read pins the entire file in
|
||||
// the kernel page cache, which starves concurrent writes and
|
||||
// collapses mux throughput. Mirrors the write-side
|
||||
// WritebackPipeline's DONTNEED policy.
|
||||
self.bytes_read_since_drop += bytes as u64;
|
||||
if self.bytes_read_since_drop >= self.drop_chunk_bytes {
|
||||
let drop_start = self.drop_window_start;
|
||||
let drop_len = self.bytes_read_since_drop;
|
||||
platform::drop_window(&self.file, drop_start, drop_len);
|
||||
self.drop_window_start = drop_start + drop_len;
|
||||
self.bytes_read_since_drop = 0;
|
||||
}
|
||||
|
||||
Ok(bytes)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
use tempfile::tempdir;
|
||||
|
||||
/// Build a deterministic ISO of `sectors` sectors where sector `n`
|
||||
/// is filled with the byte pattern `((n & 0xff) as u8)`. Lets us
|
||||
/// verify any sector by content alone.
|
||||
fn make_iso(path: &std::path::Path, sectors: u32) {
|
||||
let mut f = std::fs::File::create(path).unwrap();
|
||||
let mut chunk = vec![0u8; SECTOR_SIZE];
|
||||
for n in 0..sectors {
|
||||
let b = (n & 0xff) as u8;
|
||||
chunk.iter_mut().for_each(|c| *c = b);
|
||||
f.write_all(&chunk).unwrap();
|
||||
}
|
||||
f.flush().unwrap();
|
||||
}
|
||||
|
||||
/// Sectors used by spanning-boundary tests. Pick something that
|
||||
/// exercises multi-megabyte reads without making test ISOs huge.
|
||||
/// 8192 sectors = 16 MiB — large enough to cross any readahead
|
||||
/// chunk size we set the kernel hint to.
|
||||
const TEST_SPAN_SECTORS: u32 = 8192;
|
||||
|
||||
#[test]
|
||||
fn sequential_reads_match_file() {
|
||||
let total = TEST_SPAN_SECTORS * 2 + 17;
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("seq.iso");
|
||||
make_iso(&path, total);
|
||||
|
||||
let mut src = FileSectorSource::open(&path).unwrap();
|
||||
assert_eq!(src.capacity_sectors(), total);
|
||||
|
||||
let mut got = vec![0u8; SECTOR_SIZE];
|
||||
for lba in 0..total {
|
||||
src.read_sectors(lba, 1, &mut got, false).unwrap();
|
||||
let expected = (lba & 0xff) as u8;
|
||||
assert!(
|
||||
got.iter().all(|b| *b == expected),
|
||||
"sector {lba} content mismatch: expected 0x{expected:02x}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multi_sector_read_across_chunk_boundary() {
|
||||
let total = TEST_SPAN_SECTORS * 2;
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("span.iso");
|
||||
make_iso(&path, total);
|
||||
|
||||
let mut src = FileSectorSource::open(&path).unwrap();
|
||||
|
||||
let span_lba = TEST_SPAN_SECTORS - 2;
|
||||
let mut buf4 = vec![0u8; SECTOR_SIZE * 4];
|
||||
src.read_sectors(span_lba, 4, &mut buf4, false).unwrap();
|
||||
for i in 0..4 {
|
||||
let lba = span_lba + i as u32;
|
||||
let expected = (lba & 0xff) as u8;
|
||||
for b in &buf4[i * SECTOR_SIZE..(i + 1) * SECTOR_SIZE] {
|
||||
assert_eq!(*b, expected, "byte mismatch at sub-sector {i}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backward_seek_reads_correct_bytes() {
|
||||
// Read forward then jump back: the SectorSource contract is
|
||||
// byte-correctness regardless of access pattern.
|
||||
let total = TEST_SPAN_SECTORS * 2 + 5;
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("back.iso");
|
||||
make_iso(&path, total);
|
||||
|
||||
let mut src = FileSectorSource::open(&path).unwrap();
|
||||
let mut got = vec![0u8; SECTOR_SIZE];
|
||||
|
||||
src.read_sectors(TEST_SPAN_SECTORS + 1, 1, &mut got, false)
|
||||
.unwrap();
|
||||
src.read_sectors(0, 1, &mut got, false).unwrap();
|
||||
assert!(got.iter().all(|b| *b == 0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_at_eof_returns_correct_bytes() {
|
||||
// File smaller than the readahead chunk — reads near EOF must
|
||||
// still return correct bytes.
|
||||
let total: u32 = 100;
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("small.iso");
|
||||
make_iso(&path, total);
|
||||
|
||||
let mut src = FileSectorSource::open(&path).unwrap();
|
||||
assert_eq!(src.capacity_sectors(), total);
|
||||
|
||||
let mut got = vec![0u8; SECTOR_SIZE];
|
||||
src.read_sectors(0, 1, &mut got, false).unwrap();
|
||||
src.read_sectors(total - 1, 1, &mut got, false).unwrap();
|
||||
let expected = ((total - 1) & 0xff) as u8;
|
||||
assert!(got.iter().all(|b| *b == expected));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn large_single_read() {
|
||||
// A multi-MB single read must work — the implementation has
|
||||
// no app-level chunking, so this just exercises the direct
|
||||
// pread path on a larger request.
|
||||
let total = TEST_SPAN_SECTORS + 100;
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("big.iso");
|
||||
make_iso(&path, total);
|
||||
|
||||
let mut src = FileSectorSource::open(&path).unwrap();
|
||||
let req = (TEST_SPAN_SECTORS + 1) as u16;
|
||||
let req_bytes = req as usize * SECTOR_SIZE;
|
||||
let mut big = vec![0u8; req_bytes];
|
||||
src.read_sectors(0, req, &mut big, false).unwrap();
|
||||
assert!(big[..SECTOR_SIZE].iter().all(|b| *b == 0));
|
||||
let last_lba = req as u32 - 1;
|
||||
let exp = (last_lba & 0xff) as u8;
|
||||
let last_off = (req as usize - 1) * SECTOR_SIZE;
|
||||
assert!(
|
||||
big[last_off..last_off + SECTOR_SIZE]
|
||||
.iter()
|
||||
.all(|b| *b == exp)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_chunk_size_env_override() {
|
||||
// Explicit 8 MiB via env var.
|
||||
// SAFETY: tests in this crate are single-threaded per the
|
||||
// default cargo test harness, but std::env::set_var is
|
||||
// declared `unsafe` since Rust 2024 (it can race with other
|
||||
// threads / TLS). For a test that runs in-process before any
|
||||
// FileSectorSource construction this is safe in practice.
|
||||
unsafe {
|
||||
std::env::set_var("FREEMKV_READ_DROP_CHUNK_MIB", "8");
|
||||
}
|
||||
assert_eq!(read_drop_chunk_bytes(), 8 * 1024 * 1024);
|
||||
|
||||
unsafe {
|
||||
std::env::remove_var("FREEMKV_READ_DROP_CHUNK_MIB");
|
||||
}
|
||||
assert_eq!(read_drop_chunk_bytes(), READ_DROP_CHUNK_BYTES_DEFAULT);
|
||||
|
||||
// Garbage env value falls back to default.
|
||||
unsafe {
|
||||
std::env::set_var("FREEMKV_READ_DROP_CHUNK_MIB", "not-a-number");
|
||||
}
|
||||
assert_eq!(read_drop_chunk_bytes(), READ_DROP_CHUNK_BYTES_DEFAULT);
|
||||
unsafe {
|
||||
std::env::remove_var("FREEMKV_READ_DROP_CHUNK_MIB");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
//! Fallback for targets without a known sequential-readahead hint
|
||||
//! (BSDs, illumos, etc.). No-op — reads still work, they just don't
|
||||
//! get the OS-level prefetch widening.
|
||||
|
||||
use std::fs::File;
|
||||
|
||||
pub(super) fn hint_sequential(_file: &File, _len_bytes: u64) {}
|
||||
|
||||
pub(super) fn drop_window(_file: &File, _start: u64, _len: u64) {}
|
||||
|
||||
pub(super) fn prefetch(_file: &File, _offset: u64, _len: u64) {}
|
||||
@@ -0,0 +1,30 @@
|
||||
//! 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.
|
||||
|
||||
use std::fs::File;
|
||||
|
||||
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)"
|
||||
);
|
||||
}
|
||||
|
||||
/// Windows page-cache eviction is not exposed via a posix_fadvise
|
||||
/// equivalent. The kernel does its own working-set management. No-op
|
||||
/// for now.
|
||||
pub(super) fn drop_window(_file: &File, _start: u64, _len: u64) {}
|
||||
|
||||
/// Windows async-prefetch hint. With FILE_FLAG_SEQUENTIAL_SCAN at
|
||||
/// open the kernel already prefetches aggressively, so there's no
|
||||
/// per-range hint we'd add on top. No-op stub for parity with the
|
||||
/// posix platforms.
|
||||
pub(super) fn prefetch(_file: &File, _offset: u64, _len: u64) {}
|
||||
+21
-6
@@ -10,23 +10,38 @@
|
||||
//! no-op pipeline — their default cache policies have not been shown
|
||||
//! to exhibit the same pathology for this access pattern.
|
||||
//!
|
||||
//! `FileSectorSource` is the read-side dual — it implements
|
||||
//! [`crate::sector::SectorSource`] for an ISO file using direct
|
||||
//! `pread`-equivalent calls so the kernel's own readahead policy runs
|
||||
//! (which interleaves naturally with the concurrent writeback). It
|
||||
//! pairs that with periodic `posix_fadvise(DONTNEED)` drops on the
|
||||
//! 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
|
||||
//! 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.
|
||||
|
||||
pub(crate) mod bounded;
|
||||
pub mod byte_channel;
|
||||
pub mod byte_prefetcher;
|
||||
pub mod file_sector_source;
|
||||
pub mod sink;
|
||||
mod writeback;
|
||||
mod writeback_file;
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
pub(crate) mod platform_macos;
|
||||
|
||||
pub mod pipeline;
|
||||
|
||||
pub(crate) use writeback_file::WritebackFile;
|
||||
|
||||
// Re-exports for the 0.18 redesign. Sweep + patch are both wired up
|
||||
// (disc/sweep.rs, disc/patch.rs); mux migrates separately in autorip.
|
||||
// `WRITE_THROUGH_DEPTH` is patch-specific and has no other in-tree
|
||||
// caller — the targeted `#[allow]` keeps the re-export visible without
|
||||
// dragging the rest of the module under `dead_code`.
|
||||
#[allow(unused_imports)]
|
||||
// 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,
|
||||
|
||||
+369
-12
@@ -9,7 +9,6 @@
|
||||
//! 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.
|
||||
//! See `(internal)/memory/0_18_redesign.md` for the full picture.
|
||||
//!
|
||||
//! ## Cancellation and error semantics
|
||||
//!
|
||||
@@ -34,10 +33,43 @@
|
||||
//! consumer lag detection). This is critical for diagnosing stalls.
|
||||
|
||||
use std::io;
|
||||
use std::sync::mpsc::{SyncSender, sync_channel};
|
||||
use std::thread::{self, JoinHandle};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
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.
|
||||
///
|
||||
/// 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.
|
||||
pub const JOIN_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
/// Halt-check cadence for the send loop. Producer blocks on
|
||||
/// [`crossbeam_channel::Sender::send_timeout`] for this slice — the
|
||||
/// kernel wakes it the instant the consumer drains a slot, so on the
|
||||
/// happy path there's no throughput cap from this primitive at all
|
||||
/// (the cap is whatever the underlying medium can sustain). When the
|
||||
/// consumer is genuinely wedged, the timeout fires every
|
||||
/// [`crate::halt::POLL_INTERVAL`] and the producer checks the halt
|
||||
/// token; that's the latency a stop request will observe.
|
||||
///
|
||||
/// Single source of truth lives in [`crate::halt::POLL_INTERVAL`]
|
||||
/// (also used by `bounded_syscall`). Aliased here for readability of
|
||||
/// the send/finish call sites below.
|
||||
///
|
||||
/// 0.21.7 replaced an old `std::sync::mpsc::sync_channel` + 50 ms
|
||||
/// `thread::sleep` polling loop that capped mux throughput at
|
||||
/// ~20 frames/sec ≈ 1 MB/s on saturated channels.
|
||||
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.
|
||||
pub fn debug_enabled() -> bool {
|
||||
@@ -65,8 +97,7 @@ pub const WRITE_PIPELINE_DEPTH: usize = 16;
|
||||
/// Channel depth for write-through pipelines. Each `send` fully
|
||||
/// drains before the next can enqueue. Use this when the producer
|
||||
/// must observe consumer side-effects (e.g. mapfile state) before
|
||||
/// emitting the next item.
|
||||
#[allow(dead_code)]
|
||||
/// emitting the next item. Currently used by `disc::patch`.
|
||||
pub const WRITE_THROUGH_DEPTH: usize = 1;
|
||||
|
||||
/// Outcome of [`Sink::apply`]: either keep feeding items
|
||||
@@ -110,7 +141,7 @@ pub trait Sink<I>: Send + 'static {
|
||||
/// Bounded producer/consumer pipeline. Holds the producer-side
|
||||
/// channel and the consumer thread's join handle.
|
||||
pub struct Pipeline<I: Send + 'static, R: Send + 'static> {
|
||||
tx: SyncSender<I>,
|
||||
tx: Sender<I>,
|
||||
handle: JoinHandle<Result<R, Error>>,
|
||||
}
|
||||
|
||||
@@ -127,11 +158,10 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> {
|
||||
/// propagated rather than panicked.
|
||||
///
|
||||
/// Sweep uses [`Pipeline::spawn_named`] directly so the consumer
|
||||
/// thread shows up as `freemkv-sweep-consumer`; this function has
|
||||
/// no in-tree caller yet. Patch and mux migrate in later 0.18
|
||||
/// slices. The targeted `#[allow]` is removed when one of them
|
||||
/// lands on the default name.
|
||||
#[allow(dead_code)]
|
||||
/// thread shows up as `freemkv-sweep-consumer`; mux uses
|
||||
/// `freemkv-mux-consumer`. `Pipeline::spawn` (this function, with
|
||||
/// the default name) is used by `disc::patch` and by the unit
|
||||
/// tests in this module.
|
||||
pub fn spawn<S: Sink<I, Output = R>>(depth: usize, sink: S) -> Result<Self, Error> {
|
||||
Self::spawn_named("freemkv-pipeline-consumer", depth, sink)
|
||||
}
|
||||
@@ -145,7 +175,7 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> {
|
||||
depth: usize,
|
||||
sink: S,
|
||||
) -> Result<Self, Error> {
|
||||
let (tx, rx) = sync_channel::<I>(depth);
|
||||
let (tx, rx) = bounded::<I>(depth);
|
||||
let handle = thread::Builder::new()
|
||||
.name(name.into())
|
||||
.spawn(move || -> Result<R, Error> {
|
||||
@@ -256,10 +286,84 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> {
|
||||
/// `Err`. Useful for best-effort signalling (e.g. sweep's
|
||||
/// throttled `StatsRequest`) where dropping the message is
|
||||
/// preferable to blocking the producer.
|
||||
pub fn try_send(&self, item: I) -> Result<(), std::sync::mpsc::TrySendError<I>> {
|
||||
pub fn try_send(&self, item: I) -> Result<(), TrySendError<I>> {
|
||||
self.tx.try_send(item)
|
||||
}
|
||||
|
||||
/// Halt-aware bounded variant of [`Pipeline::send`].
|
||||
///
|
||||
/// Uses [`crossbeam_channel::Sender::send_timeout`] so the producer
|
||||
/// thread BLOCKS on consumer drain (kernel-wakeup) rather than
|
||||
/// polling. The timeout slice is just the halt-observation cadence
|
||||
/// ([`SEND_HALT_CHECK_INTERVAL`]) — on the happy path the producer
|
||||
/// wakes the instant the consumer drains a slot, so there is no
|
||||
/// throughput cap from this primitive at any medium speed.
|
||||
///
|
||||
/// Returns:
|
||||
///
|
||||
/// - `Ok(())` once the item lands in the channel.
|
||||
/// - `Err(item)` if the consumer disconnected, the halt fired, or
|
||||
/// the deadline elapsed — the caller gets the item back so it
|
||||
/// can decide whether to drop it, route it elsewhere, or unwind.
|
||||
///
|
||||
/// Use this in producer threads that have a `Halt` token threaded
|
||||
/// through (mux, sweep, patch). Plain [`Pipeline::send`] is
|
||||
/// preserved for callers that don't (yet) plumb halt through.
|
||||
///
|
||||
/// Unlike [`Pipeline::send`], this never blocks the producer
|
||||
/// thread inside an unkillable `mpsc::send` — if the consumer is
|
||||
/// wedged inside an unkillable syscall, the producer can still
|
||||
/// observe `/api/stop` and unwind within
|
||||
/// [`SEND_HALT_CHECK_INTERVAL`].
|
||||
pub fn send_with_halt(&self, item: I, halt: &Halt, deadline: Duration) -> Result<(), I> {
|
||||
use crossbeam_channel::SendTimeoutError;
|
||||
let end = Instant::now() + deadline;
|
||||
let mut pending = item;
|
||||
loop {
|
||||
// Pre-check the cheap exit conditions before parking.
|
||||
if halt.is_cancelled() {
|
||||
if debug_enabled() {
|
||||
tracing::debug!(
|
||||
"Pipeline send_with_halt: halt observed, returning item={}",
|
||||
std::any::type_name::<I>()
|
||||
);
|
||||
}
|
||||
return Err(pending);
|
||||
}
|
||||
let now = Instant::now();
|
||||
if now >= end {
|
||||
if debug_enabled() {
|
||||
tracing::debug!(
|
||||
"Pipeline send_with_halt: deadline elapsed, returning item={}",
|
||||
std::any::type_name::<I>()
|
||||
);
|
||||
}
|
||||
return Err(pending);
|
||||
}
|
||||
// Wait for space-available or halt-check tick, whichever
|
||||
// is sooner. Crossbeam's send_timeout is kernel-wakeup
|
||||
// based: the consumer's recv on a saturated channel
|
||||
// signals this thread the moment a slot opens up.
|
||||
let slice = SEND_HALT_CHECK_INTERVAL.min(end.saturating_duration_since(now));
|
||||
match self.tx.send_timeout(pending, slice) {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(SendTimeoutError::Timeout(returned)) => {
|
||||
pending = returned;
|
||||
// loop: re-check halt + deadline, then park again
|
||||
}
|
||||
Err(SendTimeoutError::Disconnected(returned)) => {
|
||||
if debug_enabled() {
|
||||
tracing::debug!(
|
||||
"Pipeline send_with_halt: consumer disconnected, item={}",
|
||||
std::any::type_name::<I>()
|
||||
);
|
||||
}
|
||||
return Err(returned);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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 —
|
||||
@@ -289,6 +393,69 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Halt-aware, deadline-bounded variant of [`Pipeline::finish`].
|
||||
///
|
||||
/// Drops the producer-side channel (same as `finish`) and then
|
||||
/// polls `JoinHandle::is_finished()` on a 250 ms cadence. Between
|
||||
/// slices, checks (1) the optional [`Halt`] token and (2) the
|
||||
/// [`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()`.
|
||||
///
|
||||
/// 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).
|
||||
///
|
||||
/// Plain [`Pipeline::finish`] is preserved for callers without a
|
||||
/// halt-token plumbed through; that path still blocks indefinitely
|
||||
/// on `join()`, matching pre-0.20.8 behaviour.
|
||||
pub fn finish_with_halt(self, halt: Option<&Halt>) -> Result<R, Error> {
|
||||
let Pipeline { tx, handle } = self;
|
||||
drop(tx);
|
||||
let deadline = Instant::now() + Duration::from_secs(JOIN_TIMEOUT_SECS);
|
||||
loop {
|
||||
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}")),
|
||||
})
|
||||
}
|
||||
};
|
||||
}
|
||||
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"),
|
||||
});
|
||||
}
|
||||
}
|
||||
if Instant::now() >= deadline {
|
||||
// Consumer thread is intentionally leaked.
|
||||
return Err(Error::IoError {
|
||||
source: io::Error::other("pipeline join timed out"),
|
||||
});
|
||||
}
|
||||
thread::sleep(POLL_INTERVAL);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -553,4 +720,194 @@ mod tests {
|
||||
other => panic!("expected Err(IoError), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Never-completing sink — `apply` blocks until cancelled. Signals
|
||||
/// `started` once it has consumed its first item so the test
|
||||
/// driver knows the consumer thread is wedged in `apply` (and
|
||||
/// will no longer drain the channel). Used to drive the
|
||||
/// halt/timeout paths of `send_with_halt` and `finish_with_halt`
|
||||
/// without depending on real I/O.
|
||||
struct NeverDrainsSink {
|
||||
cancel: Arc<std::sync::atomic::AtomicBool>,
|
||||
started: Arc<std::sync::atomic::AtomicBool>,
|
||||
}
|
||||
|
||||
impl Sink<u64> for NeverDrainsSink {
|
||||
type Output = ();
|
||||
|
||||
fn apply(&mut self, _item: u64) -> Result<Flow, Error> {
|
||||
self.started.store(true, Ordering::SeqCst);
|
||||
while !self.cancel.load(Ordering::SeqCst) {
|
||||
std::thread::sleep(Duration::from_millis(20));
|
||||
}
|
||||
Ok(Flow::Continue)
|
||||
}
|
||||
|
||||
fn close(self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Spin until `started` flips or `bail` elapses. Used by the
|
||||
/// send_with_halt tests to synchronise with the consumer thread
|
||||
/// before exercising the bounded-send timeout path.
|
||||
fn wait_for_started(started: &Arc<std::sync::atomic::AtomicBool>, bail: Duration) {
|
||||
let end = Instant::now() + bail;
|
||||
while !started.load(Ordering::SeqCst) {
|
||||
assert!(Instant::now() < end, "consumer never started apply()");
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn send_with_halt_returns_item_on_deadline() {
|
||||
// depth=1 + consumer wedged in apply on the first item, AND
|
||||
// the channel buffer already loaded with a second item, means
|
||||
// any further `try_send` sees Full; with a 200 ms deadline and
|
||||
// no halt fired, send_with_halt must return `Err(item)` within
|
||||
// roughly the deadline. Synchronising on `started` ensures the
|
||||
// consumer has actually started its wedged apply BEFORE we
|
||||
// load the channel-buffer slot — without that, the consumer
|
||||
// could still drain in a race window.
|
||||
let cancel = Arc::new(std::sync::atomic::AtomicBool::new(false));
|
||||
let started = Arc::new(std::sync::atomic::AtomicBool::new(false));
|
||||
let pipe = Pipeline::spawn(
|
||||
1,
|
||||
NeverDrainsSink {
|
||||
cancel: cancel.clone(),
|
||||
started: started.clone(),
|
||||
},
|
||||
)
|
||||
.expect("spawn should succeed");
|
||||
// First send: consumer recv()s it and wedges in apply.
|
||||
pipe.send(0u64).expect("first send hands off to consumer");
|
||||
wait_for_started(&started, Duration::from_secs(2));
|
||||
// Second send: lands in the depth=1 buffer slot, consumer
|
||||
// can't pick it up because it's wedged in apply. Channel now
|
||||
// full from the producer's perspective.
|
||||
pipe.send(1u64).expect("second send fills the buffer");
|
||||
|
||||
let halt = crate::halt::Halt::new();
|
||||
let start = Instant::now();
|
||||
let res = pipe.send_with_halt(99u64, &halt, Duration::from_millis(200));
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
// Release the leaked consumer so the test process winds down.
|
||||
cancel.store(true, Ordering::SeqCst);
|
||||
let _ = pipe.finish();
|
||||
|
||||
assert!(matches!(res, Err(99)), "expected item returned on deadline");
|
||||
assert!(
|
||||
elapsed >= Duration::from_millis(150),
|
||||
"deadline returned too early: {elapsed:?}"
|
||||
);
|
||||
assert!(
|
||||
elapsed < Duration::from_secs(2),
|
||||
"deadline blew past tolerance: {elapsed:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[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.
|
||||
let cancel = Arc::new(std::sync::atomic::AtomicBool::new(false));
|
||||
let started = Arc::new(std::sync::atomic::AtomicBool::new(false));
|
||||
let pipe = Pipeline::spawn(
|
||||
1,
|
||||
NeverDrainsSink {
|
||||
cancel: cancel.clone(),
|
||||
started: started.clone(),
|
||||
},
|
||||
)
|
||||
.expect("spawn should succeed");
|
||||
pipe.send(0u64).expect("first send hands off to consumer");
|
||||
wait_for_started(&started, Duration::from_secs(2));
|
||||
pipe.send(1u64).expect("second send fills the buffer");
|
||||
|
||||
let halt = crate::halt::Halt::new();
|
||||
let halt2 = halt.clone();
|
||||
std::thread::spawn(move || {
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
halt2.cancel();
|
||||
});
|
||||
|
||||
let start = Instant::now();
|
||||
let res = pipe.send_with_halt(7u64, &halt, Duration::from_secs(10));
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
cancel.store(true, Ordering::SeqCst);
|
||||
let _ = pipe.finish();
|
||||
|
||||
assert!(matches!(res, Err(7)), "expected item returned on halt");
|
||||
assert!(
|
||||
elapsed < Duration::from_secs(2),
|
||||
"halt observation took too long: {elapsed:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[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.
|
||||
let cancel = Arc::new(std::sync::atomic::AtomicBool::new(false));
|
||||
let started = Arc::new(std::sync::atomic::AtomicBool::new(false));
|
||||
let pipe = Pipeline::spawn(
|
||||
DEFAULT_PIPELINE_DEPTH,
|
||||
NeverDrainsSink {
|
||||
cancel: cancel.clone(),
|
||||
started: started.clone(),
|
||||
},
|
||||
)
|
||||
.expect("spawn should succeed");
|
||||
pipe.send(0u64).expect("seed item the consumer wedges on");
|
||||
wait_for_started(&started, Duration::from_secs(2));
|
||||
|
||||
let halt = crate::halt::Halt::new();
|
||||
let halt2 = halt.clone();
|
||||
std::thread::spawn(move || {
|
||||
std::thread::sleep(Duration::from_millis(400));
|
||||
halt2.cancel();
|
||||
});
|
||||
|
||||
let start = Instant::now();
|
||||
let res = pipe.finish_with_halt(Some(&halt));
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
// 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:?}"),
|
||||
}
|
||||
// Bailed out within ~1 second of the halt firing (worst case
|
||||
// one POLL_INTERVAL = 250 ms of slack).
|
||||
assert!(
|
||||
elapsed < Duration::from_secs(2),
|
||||
"halt observation took too long: {elapsed:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn finish_with_halt_happy_path_returns_output() {
|
||||
// No halt token, sink completes normally — finish_with_halt
|
||||
// must return the same Output that `finish` would.
|
||||
let pipe = Pipeline::spawn(DEFAULT_PIPELINE_DEPTH, SumSink { total: 0 })
|
||||
.expect("spawn should succeed");
|
||||
for i in 0..10u64 {
|
||||
pipe.send(i).expect("send should succeed");
|
||||
}
|
||||
let total = pipe
|
||||
.finish_with_halt(None)
|
||||
.expect("happy-path finish_with_halt should succeed");
|
||||
assert_eq!(total, (0..10u64).sum::<u64>());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
//! Shared macOS `fcntl(F_PREALLOCATE)` definitions.
|
||||
//!
|
||||
//! The `libc` crate doesn't expose these symbols across all macOS SDK
|
||||
//! versions, so we define them locally with values from
|
||||
//! `/usr/include/sys/fcntl.h`. Two call sites (
|
||||
//! [`crate::io::writeback_file`] and [`crate::io::sink::preallocate`])
|
||||
//! need the same constants and `fstore_t` layout — keeping a single
|
||||
//! source of truth here prevents the two copies from drifting.
|
||||
//!
|
||||
//! Module-level cfg gate lives in the parent (`io/mod.rs`); this file
|
||||
//! is only compiled on macOS, so no inner `#![cfg]` is needed.
|
||||
|
||||
/// `fcntl(F_PREALLOCATE)` command number from `sys/fcntl.h`.
|
||||
pub(crate) const F_PREALLOCATE: libc::c_int = 42;
|
||||
|
||||
/// Anchor preallocation at the current physical EOF.
|
||||
pub(crate) const F_PEOFPOSMODE: libc::c_int = 3;
|
||||
|
||||
/// Prefer a contiguous allocation. Try this first; on `EINVAL` (no
|
||||
/// contiguous run of that size), fall back to `F_ALLOCATEALL`.
|
||||
pub(crate) const F_ALLOCATECONTIG: libc::c_uint = 0x0000_0002;
|
||||
|
||||
/// Allow non-contiguous allocation. Stronger guarantee than just
|
||||
/// asking for `F_ALLOCATECONTIG` because the kernel will piece
|
||||
/// together fragments rather than failing.
|
||||
pub(crate) const F_ALLOCATEALL: libc::c_uint = 0x0000_0004;
|
||||
|
||||
/// `fstore_t` from `sys/fcntl.h`. `repr(C)` because we hand it to
|
||||
/// `fcntl(F_PREALLOCATE)` which writes through the pointer.
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Copy)]
|
||||
pub(crate) struct Fstore {
|
||||
pub fst_flags: libc::c_uint,
|
||||
pub fst_posmode: libc::c_int,
|
||||
pub fst_offset: libc::off_t,
|
||||
pub fst_length: libc::off_t,
|
||||
pub fst_bytesalloc: libc::off_t,
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
//! `LocalFileSink` — `BufWriter<File>` for the common local-disk case.
|
||||
//!
|
||||
//! Buffering: 4 MiB internal `BufWriter`. Sized to coalesce the small
|
||||
//! per-PES writes that come out of the muxer into kernel-page-aligned
|
||||
//! flushes without making the buffer big enough to matter for memory
|
||||
//! pressure on a single concurrent rip.
|
||||
//!
|
||||
//! `Seek` flushes the underlying `BufWriter` first; otherwise a seek
|
||||
//! could leapfrog buffered data and silently corrupt the file. This is
|
||||
//! the same shape `BufWriter` itself uses when it impls `Seek` in
|
||||
//! stdlib, and is necessary for MKV's seek-back operations (cluster
|
||||
//! 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.
|
||||
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::io::{self, BufWriter, Seek, SeekFrom, Write};
|
||||
use std::path::Path;
|
||||
|
||||
use super::preallocate;
|
||||
|
||||
const BUFFER_BYTES: usize = 4 * 1024 * 1024;
|
||||
|
||||
/// Random-access write sink for local disks.
|
||||
///
|
||||
/// Wraps a `BufWriter<File>` with a 4 MiB internal buffer and forwards
|
||||
/// `Write`/`Seek` so any call site that previously held a `File` or
|
||||
/// `WritebackFile` can drop this in. `finish()` flushes the buffer and
|
||||
/// runs `sync_all` on the underlying file so the caller can drop it
|
||||
/// without losing data.
|
||||
///
|
||||
/// Construction always opens the file `create + truncate + read +
|
||||
/// write`. `read` is enabled so the same handle can be reused for a
|
||||
/// verification re-read after the mux (the existing
|
||||
/// `FileSectorSink::create` pattern). On Linux, [`with_size_hint`]
|
||||
/// additionally calls `fallocate(FALLOC_FL_KEEP_SIZE)` to pre-reserve
|
||||
/// extents.
|
||||
///
|
||||
/// [`with_size_hint`]: Self::with_size_hint
|
||||
pub struct LocalFileSink {
|
||||
inner: BufWriter<File>,
|
||||
}
|
||||
|
||||
impl LocalFileSink {
|
||||
/// Open `path` for writing, truncating any existing contents.
|
||||
pub fn create(path: &Path) -> io::Result<Self> {
|
||||
let file = OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.open(path)?;
|
||||
Ok(Self {
|
||||
inner: BufWriter::with_capacity(BUFFER_BYTES, file),
|
||||
})
|
||||
}
|
||||
|
||||
/// Like [`Self::create`] but additionally calls the per-OS
|
||||
/// preallocate path with `size_bytes`. On Linux this is
|
||||
/// `fallocate(FALLOC_FL_KEEP_SIZE)` so the on-disk extents are
|
||||
/// reserved up front (reducing fragmentation for big sequential
|
||||
/// muxer output); on other OSes it is a no-op today. Failures
|
||||
/// from the preallocate call are non-fatal — the file is still
|
||||
/// returned, just without the size reservation.
|
||||
pub fn with_size_hint(path: &Path, size_bytes: u64) -> io::Result<Self> {
|
||||
let file = OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.open(path)?;
|
||||
preallocate::preallocate(&file, size_bytes);
|
||||
Ok(Self {
|
||||
inner: BufWriter::with_capacity(BUFFER_BYTES, file),
|
||||
})
|
||||
}
|
||||
|
||||
/// 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
|
||||
pub fn sync_all(&mut self) -> io::Result<()> {
|
||||
self.inner.flush()?;
|
||||
self.inner.get_ref().sync_all()
|
||||
}
|
||||
}
|
||||
|
||||
impl Write for LocalFileSink {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
self.inner.write(buf)
|
||||
}
|
||||
|
||||
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
|
||||
self.inner.write_all(buf)
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
self.inner.flush()
|
||||
}
|
||||
}
|
||||
|
||||
impl Seek for LocalFileSink {
|
||||
fn seek(&mut self, from: SeekFrom) -> io::Result<u64> {
|
||||
// Flush before seeking so buffered bytes land at the offset
|
||||
// they were written for, not the new one.
|
||||
self.inner.flush()?;
|
||||
self.inner.get_mut().seek(from)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Read;
|
||||
|
||||
#[test]
|
||||
fn write_seek_roundtrip() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("rt.bin");
|
||||
let mut s = LocalFileSink::create(&p).unwrap();
|
||||
s.write_all(b"AAAA").unwrap();
|
||||
s.write_all(b"BBBB").unwrap();
|
||||
// Seek back over the second word and overwrite.
|
||||
s.seek(SeekFrom::Start(4)).unwrap();
|
||||
s.write_all(b"CCCC").unwrap();
|
||||
s.sync_all().unwrap();
|
||||
drop(s);
|
||||
|
||||
let mut f = File::open(&p).unwrap();
|
||||
let mut got = Vec::new();
|
||||
f.read_to_end(&mut got).unwrap();
|
||||
assert_eq!(&got[..], b"AAAACCCC");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_flushes() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("drop.bin");
|
||||
{
|
||||
let mut s = LocalFileSink::create(&p).unwrap();
|
||||
s.write_all(b"buffered").unwrap();
|
||||
// No explicit flush / sync_all — BufWriter drop runs the
|
||||
// flush and the file should land on disk.
|
||||
}
|
||||
let bytes = std::fs::read(&p).unwrap();
|
||||
assert_eq!(&bytes[..], b"buffered");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_size_hint_creates_writable_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("sz.bin");
|
||||
let mut s = LocalFileSink::with_size_hint(&p, 64 * 1024).unwrap();
|
||||
s.write_all(b"hint-ok").unwrap();
|
||||
s.sync_all().unwrap();
|
||||
drop(s);
|
||||
let bytes = std::fs::read(&p).unwrap();
|
||||
assert_eq!(&bytes[..], b"hint-ok");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
//! Output-sink trait split for the buffering architecture.
|
||||
//!
|
||||
//! Two traits, one for each capability axis of an output destination:
|
||||
//!
|
||||
//! - [`SequentialSink`] — anything you can `Write` to in order. Sockets,
|
||||
//! pipes, append-only stores, plain files. Containers that don't need
|
||||
//! seek (M2TS, fMP4, HEVC elementary) target this.
|
||||
//! - [`RandomAccessSink`] — everything `SequentialSink` plus a working
|
||||
//! `Seek`. Local files, NFS files, anything with random-write
|
||||
//! semantics. Containers that need backpatch (MKV cluster sizes, Cues
|
||||
//! index, MP4 moov-at-end) target this.
|
||||
//!
|
||||
//! `RandomAccessSink: SequentialSink` — every random-access sink is
|
||||
//! also a valid sequential sink. The muxer is generic over which it
|
||||
//! requires (`MkvMux<S: RandomAccessSink>`, `M2tsMux<S: SequentialSink>`)
|
||||
//! so an attempt to mux MKV to a network socket is a compile error.
|
||||
//!
|
||||
//! Buffering policy belongs to the concrete sink, not to a wrapper at
|
||||
//! the call site. `LocalFileSink` wraps a `BufWriter<File>` with a
|
||||
//! 4 MiB buffer for the common local-disk case; `WritebackFile`
|
||||
//! (separate module) wraps a `File` with the adaptive-chunk
|
||||
//! `sync_file_range` machinery for the Linux+NFS case.
|
||||
|
||||
use std::io::{Seek, Write};
|
||||
|
||||
mod local_file;
|
||||
mod preallocate;
|
||||
mod socket;
|
||||
|
||||
pub use local_file::LocalFileSink;
|
||||
pub use socket::{SocketSink, UdpSocketSink};
|
||||
|
||||
/// Sequential-only write destination. Sockets, pipes, append-only
|
||||
/// stores. No seek. Implementations own their write buffering — the
|
||||
/// 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`.
|
||||
pub trait SequentialSink: Write + Send {
|
||||
fn finish(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Random-access write destination. Local files, NFS files, anything
|
||||
/// with a working `Seek`. Inherits the `SequentialSink` contract — a
|
||||
/// 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.
|
||||
///
|
||||
/// - Linux + NFS path → `WritebackFile` with its adaptive-chunk
|
||||
/// sync_file_range machinery and (when supported) `fallocate` size
|
||||
/// hint.
|
||||
/// - everything else → [`LocalFileSink`] over `BufWriter<File>`. On
|
||||
/// non-Linux there is no `WritebackFile` machinery to opt into, and
|
||||
/// on local Linux the kernel's default writeback policy is already
|
||||
/// fine.
|
||||
///
|
||||
/// `size_hint`, when present, is forwarded to the per-OS preallocate
|
||||
/// path (`fallocate(KEEP_SIZE)` on Linux, `F_PREALLOCATE` on macOS when
|
||||
/// implemented, no-op elsewhere).
|
||||
///
|
||||
/// Returns a boxed trait object so the call site (mux construction)
|
||||
/// stays agnostic of which concrete sink got picked.
|
||||
#[allow(dead_code)] // wiring to mux::resolve is a follow-up commit
|
||||
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")]
|
||||
{
|
||||
if detect(dest) == FsType::Nfs {
|
||||
let wf = match size_hint {
|
||||
Some(n) => crate::io::WritebackFile::create_with_size_hint(dest, n)?,
|
||||
None => crate::io::WritebackFile::create(dest)?,
|
||||
};
|
||||
return Ok(Box::new(wf));
|
||||
}
|
||||
}
|
||||
// Silence the unused-binding warning on non-Linux where the only
|
||||
// branch above is cfg-gated out.
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
{
|
||||
let _ = detect(dest);
|
||||
}
|
||||
|
||||
let sink = match size_hint {
|
||||
Some(n) => LocalFileSink::with_size_hint(dest, n)?,
|
||||
None => LocalFileSink::create(dest)?,
|
||||
};
|
||||
Ok(Box::new(sink))
|
||||
}
|
||||
|
||||
#[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) {}
|
||||
|
||||
#[test]
|
||||
fn blanket_impls_cover_file_and_localfilesink() {
|
||||
// `File` directly via blanket impls.
|
||||
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.
|
||||
let mut s = LocalFileSink::create(&dir.path().join("b.bin")).unwrap();
|
||||
_assert_file_is_sequential(&mut s);
|
||||
_assert_file_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 {}`.
|
||||
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);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn open_for_mkv_returns_a_random_access_sink() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("c.bin");
|
||||
let mut sink = open_for_mkv(&p, Some(64 * 1024)).unwrap();
|
||||
use std::io::{Seek, SeekFrom, Write};
|
||||
sink.write_all(b"hello").unwrap();
|
||||
sink.seek(SeekFrom::Start(0)).unwrap();
|
||||
sink.finish().unwrap();
|
||||
drop(sink);
|
||||
let bytes = std::fs::read(&p).unwrap();
|
||||
assert_eq!(&bytes[..5], b"hello");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
//! Linux `fallocate(FALLOC_FL_KEEP_SIZE)` preallocation.
|
||||
//!
|
||||
//! `KEEP_SIZE` reserves extents without changing the apparent file
|
||||
//! length, which matches the muxer's expectation that writes still grow
|
||||
//! 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();
|
||||
// FALLOC_FL_KEEP_SIZE = 0x01.
|
||||
let rc = unsafe { libc::fallocate(fd, libc::FALLOC_FL_KEEP_SIZE, 0, size_bytes as i64) };
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"LocalFileSink fallocate size_hint={size_bytes} rc={rc} ok={}",
|
||||
rc == 0
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
//! 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).
|
||||
//! Reported file size is unchanged — the muxer's writes still grow it.
|
||||
|
||||
use std::fs::File;
|
||||
use std::os::unix::io::AsRawFd;
|
||||
|
||||
use crate::io::platform_macos::{
|
||||
F_ALLOCATEALL, F_ALLOCATECONTIG, F_PEOFPOSMODE, F_PREALLOCATE, Fstore,
|
||||
};
|
||||
|
||||
pub(super) fn preallocate_impl(file: &File, size_bytes: u64) {
|
||||
let fd = file.as_raw_fd();
|
||||
let mut store = Fstore {
|
||||
fst_flags: F_ALLOCATECONTIG,
|
||||
fst_posmode: F_PEOFPOSMODE,
|
||||
fst_offset: 0,
|
||||
fst_length: size_bytes as libc::off_t,
|
||||
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.
|
||||
store.fst_flags = F_ALLOCATEALL;
|
||||
rc = unsafe { libc::fcntl(fd, F_PREALLOCATE, &mut store as *mut Fstore) };
|
||||
}
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"LocalFileSink F_PREALLOCATE size_hint={size_bytes} rc={rc} bytesalloc={}",
|
||||
store.fst_bytesalloc
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
//! Per-OS extent preallocation. Best-effort; failures are logged at
|
||||
//! debug and otherwise swallowed because the file is still usable
|
||||
//! without the size reservation — only large-file fragmentation gets
|
||||
//! marginally worse.
|
||||
|
||||
use std::fs::File;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
mod linux;
|
||||
#[cfg(target_os = "macos")]
|
||||
mod macos;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
|
||||
mod other;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
use linux::preallocate_impl;
|
||||
#[cfg(target_os = "macos")]
|
||||
use macos::preallocate_impl;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
|
||||
use other::preallocate_impl;
|
||||
|
||||
/// Reserve `size_bytes` of disk space for `file`'s on-disk extents.
|
||||
/// Reported file size is unchanged — writes still grow the file
|
||||
/// naturally; only the allocator's extent map is primed.
|
||||
pub(super) fn preallocate(file: &File, size_bytes: u64) {
|
||||
preallocate_impl(file, size_bytes);
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
//! Fallback preallocate impl. No-op.
|
||||
|
||||
use std::fs::File;
|
||||
|
||||
pub(super) fn preallocate_impl(_file: &File, size_bytes: u64) {
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"LocalFileSink preallocate size_hint={size_bytes} skipped (no platform impl)"
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,278 @@
|
||||
//! TCP / UDP socket sinks (sequential-only).
|
||||
//!
|
||||
//! [`SocketSink`] wraps a `TcpStream` in a 1 MiB `BufWriter`. Constructor
|
||||
//! tunes `SO_SNDBUF` to a caller hint when provided. `finish()` flushes
|
||||
//! the buffer then `shutdown(Write)`s the socket so the peer sees clean
|
||||
//! end-of-stream.
|
||||
//!
|
||||
//! [`UdpSocketSink`] wraps a connected `UdpSocket`. Each `write` call
|
||||
//! emits exactly one datagram — the caller is responsible for packetizing
|
||||
//! to a reasonable MTU (188 × 7 = 1316 bytes for MPEG-TS-over-UDP is the
|
||||
//! 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.
|
||||
//!
|
||||
//! [`SequentialSink`]: super::SequentialSink
|
||||
//! [`RandomAccessSink`]: super::RandomAccessSink
|
||||
|
||||
use std::io::{self, BufWriter, Write};
|
||||
use std::net::{Shutdown, TcpStream, ToSocketAddrs, UdpSocket};
|
||||
|
||||
/// `BufWriter` capacity for [`SocketSink`]. 1 MiB matches the typical
|
||||
/// kernel send-buffer ceiling and keeps small-write amplification from
|
||||
/// containers (TS = 188-byte packets, fMP4 fragment headers = ~100 bytes)
|
||||
/// from translating into syscall storms.
|
||||
const TCP_BUF_CAPACITY: usize = 1024 * 1024;
|
||||
|
||||
/// Sequential-only sink over a TCP connection.
|
||||
///
|
||||
/// Wraps a `BufWriter<TcpStream>`; the inner `TcpStream` is kept as a
|
||||
/// clone so [`finish`](Self::finish) can call `shutdown(Write)` after
|
||||
/// flushing the buffer (the buffered writer doesn't expose the socket
|
||||
/// directly).
|
||||
pub struct SocketSink {
|
||||
/// Buffered write half. All payload bytes go through this.
|
||||
buf: BufWriter<TcpStream>,
|
||||
/// Shutdown handle — clone of the socket inside `buf`. Used only by
|
||||
/// `finish()` for `shutdown(Write)`; never read or written through.
|
||||
shutdown_handle: TcpStream,
|
||||
}
|
||||
|
||||
impl SocketSink {
|
||||
/// Open a TCP connection to `addr` and wrap it for sequential
|
||||
/// writing. `sndbuf_bytes`, when present, is forwarded to
|
||||
/// `setsockopt(SO_SNDBUF)` as a kernel hint — the OS may clamp it.
|
||||
///
|
||||
/// `addr` accepts anything `ToSocketAddrs` does: `"192.0.2.1:1234"`,
|
||||
/// `("host", 1234)`, a `SocketAddr`, etc.
|
||||
pub fn connect<A: ToSocketAddrs>(addr: A, sndbuf_bytes: Option<usize>) -> io::Result<Self> {
|
||||
let stream = TcpStream::connect(addr)?;
|
||||
// `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)?;
|
||||
if let Some(n) = sndbuf_bytes {
|
||||
set_send_buffer(&stream, n)?;
|
||||
}
|
||||
let shutdown_handle = stream.try_clone()?;
|
||||
Ok(Self {
|
||||
buf: BufWriter::with_capacity(TCP_BUF_CAPACITY, stream),
|
||||
shutdown_handle,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Write for SocketSink {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
self.buf.write(buf)
|
||||
}
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
self.buf.flush()
|
||||
}
|
||||
}
|
||||
|
||||
impl 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<()> {
|
||||
self.buf.flush()?;
|
||||
// `shutdown(Write)` signals clean EOF to the peer. Errors here
|
||||
// are non-fatal — the connection may have already been torn down
|
||||
// by the peer — but we surface them so callers can log.
|
||||
self.shutdown_handle.shutdown(Shutdown::Write)
|
||||
}
|
||||
}
|
||||
|
||||
/// Sequential-only sink over a connected UDP socket.
|
||||
///
|
||||
/// Each [`write`](Write::write) call sends exactly one datagram. The
|
||||
/// caller is responsible for splitting payload at packet boundaries —
|
||||
/// for MPEG-TS this means 7 × 188 = 1316 bytes per datagram, the
|
||||
/// industry standard for MPEG-TS-over-UDP. No buffering happens here;
|
||||
/// adding it would silently merge datagrams.
|
||||
///
|
||||
/// `finish()` is a no-op: UDP has no end-of-stream marker. Closing the
|
||||
/// socket happens on drop.
|
||||
pub struct UdpSocketSink {
|
||||
socket: UdpSocket,
|
||||
}
|
||||
|
||||
impl UdpSocketSink {
|
||||
/// Bind a local UDP socket to an ephemeral port and `connect` it to
|
||||
/// `peer`. `connect` doesn't open a connection — it just fixes the
|
||||
/// peer address so subsequent `send` calls don't need to repeat it,
|
||||
/// and so receive-side filtering rejects packets from other sources.
|
||||
///
|
||||
/// `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)?;
|
||||
if let Some(n) = sndbuf_bytes {
|
||||
set_udp_send_buffer(&socket, n)?;
|
||||
}
|
||||
Ok(Self { socket })
|
||||
}
|
||||
}
|
||||
|
||||
impl Write for UdpSocketSink {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
// `send` writes the entire datagram or fails — no partial sends
|
||||
// for UDP. Match `Write::write`'s contract by reporting bytes
|
||||
// accepted.
|
||||
self.socket.send(buf)
|
||||
}
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
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<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// ── Platform `SO_SNDBUF` tuning ────────────────────────────────────────────
|
||||
//
|
||||
// std's `TcpStream` / `UdpSocket` don't expose `SO_SNDBUF`. We drop to
|
||||
// libc on Linux + macOS (the libc-dep targets in Cargo.toml). On other
|
||||
// targets the hint is silently ignored — the socket still works, the
|
||||
// kernel just picks its own send-buffer size.
|
||||
|
||||
#[cfg(any(target_os = "linux", target_os = "macos"))]
|
||||
fn set_send_buffer(stream: &TcpStream, bytes: usize) -> io::Result<()> {
|
||||
use std::os::unix::io::AsRawFd;
|
||||
setsockopt_sndbuf(stream.as_raw_fd(), bytes)
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "linux", target_os = "macos"))]
|
||||
fn set_udp_send_buffer(socket: &UdpSocket, bytes: usize) -> io::Result<()> {
|
||||
use std::os::unix::io::AsRawFd;
|
||||
setsockopt_sndbuf(socket.as_raw_fd(), bytes)
|
||||
}
|
||||
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
|
||||
fn set_send_buffer(_stream: &TcpStream, _bytes: usize) -> io::Result<()> {
|
||||
// Non-Linux-non-macOS targets aren't in Cargo.toml's libc dep list;
|
||||
// silently ignore the hint rather than failing the connect. Callers
|
||||
// can detect via the lack of an explicit "sndbuf applied" signal
|
||||
// (not provided, intentionally — this is a hint, not a guarantee).
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
|
||||
fn set_udp_send_buffer(_socket: &UdpSocket, _bytes: usize) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "linux", target_os = "macos"))]
|
||||
fn setsockopt_sndbuf(fd: std::os::unix::io::RawFd, bytes: usize) -> io::Result<()> {
|
||||
// Clamp into c_int range; SO_SNDBUF takes an `int` argument.
|
||||
let want: libc::c_int = bytes.try_into().unwrap_or(libc::c_int::MAX);
|
||||
let ret = unsafe {
|
||||
libc::setsockopt(
|
||||
fd,
|
||||
libc::SOL_SOCKET,
|
||||
libc::SO_SNDBUF,
|
||||
&want as *const _ as *const libc::c_void,
|
||||
std::mem::size_of::<libc::c_int>() as libc::socklen_t,
|
||||
)
|
||||
};
|
||||
if ret != 0 {
|
||||
return Err(io::Error::last_os_error());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Read;
|
||||
use std::net::{TcpListener, UdpSocket};
|
||||
use std::thread;
|
||||
|
||||
/// Bind a listener, accept on a thread, return (listener_addr,
|
||||
/// accepted-bytes future via JoinHandle).
|
||||
#[test]
|
||||
fn socket_sink_round_trips_bytes() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let accept = thread::spawn(move || {
|
||||
let (mut sock, _) = listener.accept().unwrap();
|
||||
let mut buf = Vec::new();
|
||||
sock.read_to_end(&mut buf).unwrap();
|
||||
buf
|
||||
});
|
||||
|
||||
let mut sink = SocketSink::connect(addr, Some(256 * 1024)).unwrap();
|
||||
// Write enough to overflow the BufWriter at least once, then a
|
||||
// tail that lives in the buffer until `finish` flushes.
|
||||
let big: Vec<u8> = (0..(2 * TCP_BUF_CAPACITY))
|
||||
.map(|i| (i & 0xff) as u8)
|
||||
.collect();
|
||||
sink.write_all(&big).unwrap();
|
||||
sink.write_all(b"tail\n").unwrap();
|
||||
sink.finish().unwrap();
|
||||
drop(sink);
|
||||
|
||||
let received = accept.join().unwrap();
|
||||
assert_eq!(received.len(), big.len() + 5);
|
||||
assert_eq!(&received[..big.len()], &big[..]);
|
||||
assert_eq!(&received[big.len()..], b"tail\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn socket_sink_is_sequential_only() {
|
||||
// Compile-time assertion via dyn — if this ever started
|
||||
// satisfying `RandomAccessSink`, the trait split would be broken.
|
||||
fn _assert_seq(_: &mut dyn super::super::SequentialSink) {}
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let _accept = thread::spawn(move || {
|
||||
let _ = listener.accept();
|
||||
});
|
||||
let mut sink = SocketSink::connect(addr, None).unwrap();
|
||||
_assert_seq(&mut sink);
|
||||
// The negative is harder to assert directly (no `is_not<T>`),
|
||||
// but `SocketSink` does not impl `Seek`, so it can't unify with
|
||||
// `RandomAccessSink`'s super-bound. The Phase 2 blanket impl
|
||||
// `impl<T: SequentialSink + Seek> RandomAccessSink for T {}` thus
|
||||
// excludes it by construction.
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn udp_socket_sink_delivers_datagrams() {
|
||||
let receiver = UdpSocket::bind("127.0.0.1:0").unwrap();
|
||||
receiver
|
||||
.set_read_timeout(Some(std::time::Duration::from_secs(2)))
|
||||
.unwrap();
|
||||
let addr = receiver.local_addr().unwrap();
|
||||
let mut sink = UdpSocketSink::connect(addr, Some(128 * 1024)).unwrap();
|
||||
|
||||
sink.write_all(&[1, 2, 3, 4, 5]).unwrap();
|
||||
sink.write_all(&[9, 9, 9]).unwrap();
|
||||
sink.finish().unwrap();
|
||||
|
||||
let mut buf = [0u8; 64];
|
||||
let n1 = receiver.recv(&mut buf).unwrap();
|
||||
assert_eq!(&buf[..n1], &[1, 2, 3, 4, 5]);
|
||||
let n2 = receiver.recv(&mut buf).unwrap();
|
||||
assert_eq!(&buf[..n2], &[9, 9, 9]);
|
||||
}
|
||||
}
|
||||
+386
-31
@@ -14,9 +14,63 @@
|
||||
//! a full chunk's worth of work to flush — the wait is near-instant.
|
||||
//! Dirty cache stays bounded at ~2 × `chunk_bytes` and writes drain
|
||||
//! continuously instead of in bursts.
|
||||
//!
|
||||
//! The chunk size is adaptive: we measure the elapsed time of the
|
||||
//! `WAIT_AFTER` call over a rolling window of the last 16 chunks and
|
||||
//! resize the chunk based on the p95. Slow storage (NFS, network
|
||||
//! shares, HDD) sees larger chunks to amortise per-chunk overhead;
|
||||
//! fast storage (NVMe) sees smaller chunks to keep cache pressure
|
||||
//! tight. Bounds: [4 MiB, 256 MiB].
|
||||
//!
|
||||
//! ## NFS escape hatch
|
||||
//!
|
||||
//! `sync_file_range(WAIT_AFTER)` on an NFS-mounted file can block
|
||||
//! indefinitely waiting for the server's commit ack. If the server
|
||||
//! never acks (network partition, server-side hang, slow commit), the
|
||||
//! syscall never returns and the consumer thread is stuck inside the
|
||||
//! kernel — `/api/stop` can't reach it because halt is cooperative.
|
||||
//!
|
||||
//! When `fstatfs` reports the file lives on an NFS mount
|
||||
//! (`f_type == NFS_SUPER_MAGIC`), the pipeline skips the WAIT_AFTER +
|
||||
//! `posix_fadvise(DONTNEED)` dance entirely. NFS clients have their
|
||||
//! own buffering and commit semantics that handle dirty-page bounds
|
||||
//! without us forcing the issue. The async `SYNC_FILE_RANGE_WRITE`
|
||||
//! kickoff still runs (non-blocking by spec) so writeback still gets
|
||||
//! a nudge.
|
||||
//!
|
||||
//! ## Defence in depth: WAIT_AFTER timeout
|
||||
//!
|
||||
//! Even on local storage, a degraded disk or odd filesystem driver
|
||||
//! could in principle wedge inside WAIT_AFTER. Each WAIT_AFTER call
|
||||
//! runs on a worker thread with a 30s recv_timeout on its result
|
||||
//! channel. On timeout we log a loud error, set a `degraded` flag,
|
||||
//! and from then on skip WAIT_AFTER + DONTNEED for the rest of the
|
||||
//! pipeline's life (same shape as the NFS path). The worker thread
|
||||
//! is intentionally leaked — it unwinds whenever the syscall
|
||||
//! eventually returns or the process exits. The mux continues; the
|
||||
//! original dirty-burst pathology re-emerges but the rip can still
|
||||
//! finish instead of freezing.
|
||||
|
||||
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};
|
||||
|
||||
const ADAPTIVE_WINDOW: usize = 16;
|
||||
const CHUNK_BYTES_MIN: u64 = 4 * 1024 * 1024;
|
||||
const CHUNK_BYTES_MAX: u64 = 256 * 1024 * 1024;
|
||||
const ADAPTIVE_GROW_MS: u64 = 200;
|
||||
const ADAPTIVE_SHRINK_MS: u64 = 20;
|
||||
/// Every N chunks, emit a `debug!` snapshot of the current chunk
|
||||
/// size so operators tailing the log can see where the autoscaler
|
||||
/// settled.
|
||||
const SIZE_LOG_INTERVAL: u64 = 32;
|
||||
/// Hard upper bound on a single `sync_file_range(WAIT_AFTER)` call.
|
||||
/// Beyond this we declare the pipeline degraded and stop calling
|
||||
/// WAIT_AFTER for the rest of its life.
|
||||
const WAIT_AFTER_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
|
||||
pub(crate) struct WritebackPipeline {
|
||||
/// Aliases the wrapping `WritebackFile::file`. Only valid for the
|
||||
@@ -28,6 +82,25 @@ pub(crate) struct WritebackPipeline {
|
||||
chunk_bytes: u64,
|
||||
last_flush_pos: u64,
|
||||
pending: Option<(u64, u64)>,
|
||||
/// Rolling window of recent `WAIT_AFTER` elapsed_ms measurements.
|
||||
wait_after_window: VecDeque<u64>,
|
||||
/// Count of chunks emitted (used to space out periodic
|
||||
/// `debug!` size snapshots).
|
||||
chunk_count: u64,
|
||||
/// True when the underlying file is on an NFS mount. NFS makes
|
||||
/// WAIT_AFTER unsafe (can block forever on missing server ack), so
|
||||
/// we skip it entirely and let the NFS client handle commit on
|
||||
/// close.
|
||||
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>,
|
||||
}
|
||||
|
||||
impl WritebackPipeline {
|
||||
@@ -36,14 +109,33 @@ impl WritebackPipeline {
|
||||
/// itself, or kept inside the same struct that owns `file` — the
|
||||
/// alias is unchecked.
|
||||
pub(crate) fn new(file: &File, start_pos: u64, chunk_bytes: u64) -> Self {
|
||||
let fd = file.as_raw_fd();
|
||||
let is_nfs = detect_nfs(fd);
|
||||
tracing::info!(
|
||||
target: "mux",
|
||||
"WritebackPipeline fd={fd} is_nfs={is_nfs} chunk_bytes={chunk_bytes} strategy={}",
|
||||
if is_nfs { "nfs-skip-wait" } else { "wait+dontneed" }
|
||||
);
|
||||
Self {
|
||||
fd: file.as_raw_fd(),
|
||||
fd,
|
||||
chunk_bytes,
|
||||
last_flush_pos: start_pos,
|
||||
pending: None,
|
||||
wait_after_window: VecDeque::with_capacity(ADAPTIVE_WINDOW),
|
||||
chunk_count: 0,
|
||||
is_nfs,
|
||||
degraded: Arc::new(AtomicBool::new(false)),
|
||||
}
|
||||
}
|
||||
|
||||
/// True if we should bypass the WAIT_AFTER + DONTNEED finalisation
|
||||
/// step. NFS always bypasses; local storage bypasses once the
|
||||
/// pipeline has flipped to degraded after a WAIT_AFTER timeout.
|
||||
#[inline]
|
||||
fn skip_wait(&self) -> bool {
|
||||
self.is_nfs || self.degraded.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Caller advanced the file position to `pos`. If a chunk boundary
|
||||
/// was crossed, kick async writeback for the just-completed chunk
|
||||
/// and finalise the previous one.
|
||||
@@ -51,27 +143,115 @@ 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;
|
||||
let mut wait_ms: u64 = 0;
|
||||
let mut fadvise_ms: u64 = 0;
|
||||
// Async kickoff for the just-completed chunk runs on every
|
||||
// path (NFS, degraded, normal) — it's nominally non-blocking
|
||||
// by spec and gives the kernel an early hint that this range
|
||||
// is ready to flush.
|
||||
unsafe {
|
||||
let chunk_off = self.last_flush_pos as i64;
|
||||
let chunk_len = (pos - self.last_flush_pos) as i64;
|
||||
libc::sync_file_range(self.fd, chunk_off, chunk_len, libc::SYNC_FILE_RANGE_WRITE);
|
||||
if let Some((prev_off, prev_len)) = self.pending.take() {
|
||||
libc::sync_file_range(
|
||||
self.fd,
|
||||
prev_off as i64,
|
||||
prev_len as i64,
|
||||
libc::SYNC_FILE_RANGE_WAIT_AFTER,
|
||||
);
|
||||
libc::posix_fadvise(
|
||||
self.fd,
|
||||
prev_off as i64,
|
||||
prev_len as i64,
|
||||
libc::POSIX_FADV_DONTNEED,
|
||||
);
|
||||
}
|
||||
self.pending = Some((chunk_off as u64, chunk_len as u64));
|
||||
}
|
||||
if let Some((prev_off, prev_len)) = self.pending.take() {
|
||||
if self.skip_wait() {
|
||||
// NFS branch (or degraded fallback after a prior
|
||||
// timeout): the WAIT_AFTER + DONTNEED dance is what
|
||||
// hangs on NFS — skip it. We still advance `pending`
|
||||
// so the next call has a stable cycle.
|
||||
} else {
|
||||
// Normal local-storage branch with belt-and-braces
|
||||
// timeout. If WAIT_AFTER hangs > WAIT_AFTER_TIMEOUT
|
||||
// we mark the pipeline degraded, log a loud error,
|
||||
// and fall through to the skip path on subsequent
|
||||
// calls.
|
||||
match wait_after_with_timeout(self.fd, prev_off, prev_len) {
|
||||
Some(ms) => {
|
||||
wait_ms = ms;
|
||||
let t_fadv = Instant::now();
|
||||
unsafe {
|
||||
libc::posix_fadvise(
|
||||
self.fd,
|
||||
prev_off as i64,
|
||||
prev_len as i64,
|
||||
libc::POSIX_FADV_DONTNEED,
|
||||
);
|
||||
}
|
||||
fadvise_ms = t_fadv.elapsed().as_millis() as u64;
|
||||
self.record_wait(wait_ms);
|
||||
}
|
||||
None => {
|
||||
// Timeout branch: switch to NFS-style skip
|
||||
// for the rest of the pipeline's life. Do
|
||||
// NOT call DONTNEED — if WAIT_AFTER hasn't
|
||||
// returned, the pages aren't safely flushed.
|
||||
self.degraded.store(true, Ordering::Relaxed);
|
||||
tracing::error!(
|
||||
target: "mux",
|
||||
"WritebackPipeline WAIT_AFTER timed out after {}s on chunk off={} len={}, marking writeback degraded (subsequent chunks will skip WAIT_AFTER + DONTNEED)",
|
||||
WAIT_AFTER_TIMEOUT.as_secs(),
|
||||
prev_off,
|
||||
prev_len
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
self.pending = Some((chunk_off as u64, chunk_len as u64));
|
||||
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={}",
|
||||
chunk_off,
|
||||
chunk_len,
|
||||
self.chunk_bytes,
|
||||
self.skip_wait(),
|
||||
);
|
||||
if self.chunk_count % SIZE_LOG_INTERVAL == 0 {
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackPipeline chunk_bytes={} after {} chunks is_nfs={} degraded={}",
|
||||
self.chunk_bytes,
|
||||
self.chunk_count,
|
||||
self.is_nfs,
|
||||
self.degraded.load(Ordering::Relaxed),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Push a new `WAIT_AFTER` measurement into the rolling window
|
||||
/// and, if the window is full, adapt `chunk_bytes` based on p95.
|
||||
fn record_wait(&mut self, wait_ms: u64) {
|
||||
if self.wait_after_window.len() == ADAPTIVE_WINDOW {
|
||||
self.wait_after_window.pop_front();
|
||||
}
|
||||
self.wait_after_window.push_back(wait_ms);
|
||||
if self.wait_after_window.len() < ADAPTIVE_WINDOW {
|
||||
return;
|
||||
}
|
||||
// p95 of 16 samples ≈ sorted[14] (5 % of 16 = 0.8 ≈ 1 above).
|
||||
let mut sorted: Vec<u64> = self.wait_after_window.iter().copied().collect();
|
||||
sorted.sort_unstable();
|
||||
let p95 = sorted[14];
|
||||
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)
|
||||
} else if p95 < ADAPTIVE_SHRINK_MS && self.chunk_bytes > CHUNK_BYTES_MIN {
|
||||
(self.chunk_bytes / 2).max(CHUNK_BYTES_MIN)
|
||||
} else {
|
||||
self.chunk_bytes
|
||||
};
|
||||
if new != old {
|
||||
self.chunk_bytes = new;
|
||||
tracing::info!(
|
||||
target: "mux",
|
||||
"WritebackPipeline adaptive chunk_bytes {} -> {} p95_ms={p95}",
|
||||
old,
|
||||
new
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Caller is about to seek away from the current write region.
|
||||
@@ -85,20 +265,195 @@ impl WritebackPipeline {
|
||||
/// or when discarding the pipeline.
|
||||
pub(crate) fn finalize(&mut self) {
|
||||
if let Some((prev_off, prev_len)) = self.pending.take() {
|
||||
unsafe {
|
||||
libc::sync_file_range(
|
||||
self.fd,
|
||||
prev_off as i64,
|
||||
prev_len as i64,
|
||||
libc::SYNC_FILE_RANGE_WAIT_AFTER,
|
||||
);
|
||||
libc::posix_fadvise(
|
||||
self.fd,
|
||||
prev_off as i64,
|
||||
prev_len as i64,
|
||||
libc::POSIX_FADV_DONTNEED,
|
||||
);
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackPipeline finalize chunk off={prev_off} len={prev_len} skip_wait={} is_nfs={} degraded={}",
|
||||
self.skip_wait(),
|
||||
self.is_nfs,
|
||||
self.degraded.load(Ordering::Relaxed),
|
||||
);
|
||||
if self.skip_wait() {
|
||||
// NFS / degraded: skip WAIT_AFTER + DONTNEED. close()
|
||||
// / sync_all() handle commit through their normal
|
||||
// paths.
|
||||
return;
|
||||
}
|
||||
match wait_after_with_timeout(self.fd, prev_off, prev_len) {
|
||||
Some(_ms) => unsafe {
|
||||
libc::posix_fadvise(
|
||||
self.fd,
|
||||
prev_off as i64,
|
||||
prev_len as i64,
|
||||
libc::POSIX_FADV_DONTNEED,
|
||||
);
|
||||
},
|
||||
None => {
|
||||
self.degraded.store(true, Ordering::Relaxed);
|
||||
tracing::error!(
|
||||
target: "mux",
|
||||
"WritebackPipeline finalize WAIT_AFTER timed out after {}s on chunk off={prev_off} len={prev_len}, marking writeback degraded",
|
||||
WAIT_AFTER_TIMEOUT.as_secs(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Probe whether `fd` lives on an NFS mount. Thin wrapper around
|
||||
/// [`crate::platform::fs_type::detect_fd`] so writeback policy and
|
||||
/// general-purpose fs-type classification stay in sync (same magic
|
||||
/// numbers, same musl-vs-glibc cast handling).
|
||||
///
|
||||
/// Fails open: any classification other than NFS counts as "not NFS"
|
||||
/// (including `Unknown` on `fstatfs` error) — better to run the
|
||||
/// normal local-storage path on a misdetected NFS mount and surface
|
||||
/// the freeze loudly via [`WAIT_AFTER_TIMEOUT`] than to needlessly
|
||||
/// disable writeback bounding on every local file because of a
|
||||
/// transient stat error.
|
||||
fn detect_nfs(fd: RawFd) -> bool {
|
||||
matches!(
|
||||
crate::platform::fs_type::detect_fd(fd),
|
||||
crate::platform::fs_type::FsType::Nfs
|
||||
)
|
||||
}
|
||||
|
||||
/// 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).
|
||||
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 {
|
||||
libc::sync_file_range(fd, off as i64, len as i64, libc::SYNC_FILE_RANGE_WAIT_AFTER);
|
||||
}) {
|
||||
Ok(()) => Some(started.elapsed().as_millis() as u64),
|
||||
Err(crate::io::bounded::BoundedError::Timeout)
|
||||
| Err(crate::io::bounded::BoundedError::Halted) => None,
|
||||
Err(crate::io::bounded::BoundedError::WorkerLost) => {
|
||||
// Worker thread spawn failed or panicked before sending.
|
||||
// Treat as a benign success (no syscall ran) rather than
|
||||
// a degrade trigger — falling through with elapsed_ms=0
|
||||
// matches the no-op behaviour.
|
||||
Some(0)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
/// Helper: build a `WritebackPipeline` over a local tempfile. On
|
||||
/// every test rig (linux dev box, CI) the tempfile lives on a
|
||||
/// local FS, so `is_nfs=false` and `skip_wait` returns false until
|
||||
/// we explicitly mark the pipeline degraded.
|
||||
fn local_pipeline(chunk_bytes: u64) -> (NamedTempFile, WritebackPipeline) {
|
||||
let f = NamedTempFile::new().expect("tempfile create");
|
||||
let pipeline = WritebackPipeline::new(f.as_file(), 0, chunk_bytes);
|
||||
(f, pipeline)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_pipeline_starts_active() {
|
||||
let (_f, p) = local_pipeline(32 * 1024 * 1024);
|
||||
assert!(!p.is_nfs, "local tempfile must not classify as NFS");
|
||||
assert!(!p.degraded.load(Ordering::Relaxed));
|
||||
assert!(!p.skip_wait(), "fresh local pipeline must not skip wait");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn degraded_flag_short_circuits_wait() {
|
||||
let (_f, p) = local_pipeline(32 * 1024 * 1024);
|
||||
assert!(!p.skip_wait());
|
||||
p.degraded.store(true, Ordering::Relaxed);
|
||||
assert!(
|
||||
p.skip_wait(),
|
||||
"degraded flag must force the wait+dontneed bypass"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn record_wait_grows_chunk_on_high_p95() {
|
||||
let (_f, mut p) = local_pipeline(16 * 1024 * 1024);
|
||||
// Fill the window with samples above the grow threshold.
|
||||
for _ in 0..ADAPTIVE_WINDOW {
|
||||
p.record_wait(ADAPTIVE_GROW_MS + 50);
|
||||
}
|
||||
assert!(
|
||||
p.chunk_bytes > 16 * 1024 * 1024,
|
||||
"chunk should have grown; got {}",
|
||||
p.chunk_bytes
|
||||
);
|
||||
assert!(p.chunk_bytes <= CHUNK_BYTES_MAX);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn record_wait_shrinks_chunk_on_low_p95() {
|
||||
let (_f, mut p) = local_pipeline(64 * 1024 * 1024);
|
||||
for _ in 0..ADAPTIVE_WINDOW {
|
||||
p.record_wait(1); // well under ADAPTIVE_SHRINK_MS
|
||||
}
|
||||
assert!(
|
||||
p.chunk_bytes < 64 * 1024 * 1024,
|
||||
"chunk should have shrunk; got {}",
|
||||
p.chunk_bytes
|
||||
);
|
||||
assert!(p.chunk_bytes >= CHUNK_BYTES_MIN);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn record_wait_no_op_below_window_fill() {
|
||||
let (_f, mut p) = local_pipeline(16 * 1024 * 1024);
|
||||
let initial = p.chunk_bytes;
|
||||
// Only push a few samples; window not full → no adaptation.
|
||||
for _ in 0..(ADAPTIVE_WINDOW - 1) {
|
||||
p.record_wait(ADAPTIVE_GROW_MS + 100);
|
||||
}
|
||||
assert_eq!(
|
||||
p.chunk_bytes, initial,
|
||||
"chunk must not change before window is full"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn record_wait_clamps_to_chunk_bounds() {
|
||||
// Grow past the max.
|
||||
let (_f, mut p) = local_pipeline(CHUNK_BYTES_MAX);
|
||||
for _ in 0..ADAPTIVE_WINDOW {
|
||||
p.record_wait(ADAPTIVE_GROW_MS + 1000);
|
||||
}
|
||||
assert_eq!(p.chunk_bytes, CHUNK_BYTES_MAX, "must clamp to MAX");
|
||||
|
||||
// Shrink past the min.
|
||||
let (_f, mut p) = local_pipeline(CHUNK_BYTES_MIN);
|
||||
for _ in 0..ADAPTIVE_WINDOW {
|
||||
p.record_wait(0);
|
||||
}
|
||||
assert_eq!(p.chunk_bytes, CHUNK_BYTES_MIN, "must clamp to MIN");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_nfs_local_file_is_false() {
|
||||
// Local tempfile must not classify as NFS. This locks in the
|
||||
// consolidation through `crate::platform::fs_type::detect_fd`.
|
||||
let f = NamedTempFile::new().expect("tempfile create");
|
||||
use std::os::unix::io::AsRawFd;
|
||||
assert!(!detect_nfs(f.as_file().as_raw_fd()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn note_progress_below_chunk_is_noop() {
|
||||
let (_f, mut p) = local_pipeline(32 * 1024 * 1024);
|
||||
// No-op return before crossing the first chunk boundary.
|
||||
let before = p.chunk_count;
|
||||
p.note_progress(1024); // < 32 MiB
|
||||
assert_eq!(p.chunk_count, before);
|
||||
assert!(p.pending.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,129 +0,0 @@
|
||||
//! `WritebackFile` — a `File` wrapper whose reason for existing is the
|
||||
//! bounded-cache writeback pipeline.
|
||||
//!
|
||||
//! Why: large sequential writes (sweep, patch, mux on UHD-scale output)
|
||||
//! left to the kernel's default writeback policy accumulate hundreds of
|
||||
//! megabytes of dirty pages and then burst-flush, stalling subsequent
|
||||
//! writes for seconds at a time. `WritebackFile` drives a continuous
|
||||
//! [`super::writeback::WritebackPipeline`] that on Linux issues
|
||||
//! incremental `sync_file_range` + `posix_fadvise(DONTNEED)` calls at
|
||||
//! 32 MB granularity so dirty pages drain at the same rate they're
|
||||
//! produced. macOS and Windows fall through to a no-op pipeline — their
|
||||
//! default cache policies have not been shown to exhibit the same
|
||||
//! pathology for this access pattern.
|
||||
//!
|
||||
//! It implements `Write` and `Seek` so any call site that wrote to a
|
||||
//! plain `File` through those traits (sweep, patch, mux) can swap in
|
||||
//! `WritebackFile` without touching the body of the loop. The wrapper
|
||||
//! also tracks the current file position to feed the pipeline with
|
||||
//! progress + seek boundaries.
|
||||
//!
|
||||
//! See `super::writeback::linux` for the underlying pathology and the
|
||||
//! strategy.
|
||||
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::io::{self, Seek, SeekFrom, Write};
|
||||
use std::path::Path;
|
||||
|
||||
use super::writeback::WritebackPipeline;
|
||||
|
||||
const WRITEBACK_CHUNK_BYTES: u64 = 32 * 1024 * 1024;
|
||||
|
||||
pub(crate) struct WritebackFile {
|
||||
file: File,
|
||||
pipeline: WritebackPipeline,
|
||||
pos: u64,
|
||||
}
|
||||
|
||||
impl WritebackFile {
|
||||
/// Wrap an open `File`. The current OS file position is queried
|
||||
/// once so the pipeline starts tracking from wherever the file
|
||||
/// already is (typically 0 for fresh files; non-zero for resumed
|
||||
/// or appended files).
|
||||
pub(crate) fn new(mut file: File) -> io::Result<Self> {
|
||||
let pos = file.stream_position()?;
|
||||
let pipeline = WritebackPipeline::new(&file, pos, WRITEBACK_CHUNK_BYTES);
|
||||
Ok(Self {
|
||||
file,
|
||||
pipeline,
|
||||
pos,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a new file at `path` (truncating any existing contents)
|
||||
/// and wrap it. Convenience for the common
|
||||
/// `File::create(path)` + `WritebackFile::new(file)` pair so callers
|
||||
/// don't have to assemble a `File` first.
|
||||
pub(crate) fn create(path: &Path) -> io::Result<Self> {
|
||||
let file = File::create(path)?;
|
||||
Self::new(file)
|
||||
}
|
||||
|
||||
/// Open an existing file at `path` for writing (no truncation) and
|
||||
/// wrap it. Mirrors `File::open` semantics for the writable case
|
||||
/// — used by patch / resume paths that mutate an existing ISO in
|
||||
/// place.
|
||||
pub(crate) fn open(path: &Path) -> io::Result<Self> {
|
||||
let file = OpenOptions::new().write(true).open(path)?;
|
||||
Self::new(file)
|
||||
}
|
||||
|
||||
/// Drain in-flight writeback then issue a full fsync. Use this in
|
||||
/// place of `File::sync_all`.
|
||||
pub(crate) fn sync_all(&mut self) -> io::Result<()> {
|
||||
self.pipeline.finalize();
|
||||
self.file.sync_all()
|
||||
}
|
||||
}
|
||||
|
||||
impl Write for WritebackFile {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
let n = self.file.write(buf)?;
|
||||
self.pos += n as u64;
|
||||
self.pipeline.note_progress(self.pos);
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
|
||||
self.file.write_all(buf)?;
|
||||
self.pos += buf.len() as u64;
|
||||
self.pipeline.note_progress(self.pos);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
self.file.flush()
|
||||
}
|
||||
}
|
||||
|
||||
impl Seek for WritebackFile {
|
||||
fn seek(&mut self, from: SeekFrom) -> io::Result<u64> {
|
||||
let p = self.file.seek(from)?;
|
||||
// Only treat seeks that actually move the position as
|
||||
// boundaries — sweep does a redundant `seek(Current(pos))`
|
||||
// before every write, and we don't want that to drain the
|
||||
// pipeline on every iteration.
|
||||
if p != self.pos {
|
||||
self.pipeline.handle_seek(p);
|
||||
self.pos = p;
|
||||
}
|
||||
Ok(p)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for WritebackFile {
|
||||
fn drop(&mut self) {
|
||||
// Run the pipeline's tail finalize so the last in-flight chunk
|
||||
// gets its `WAIT_AFTER` + `posix_fadvise(DONTNEED)`. Without
|
||||
// this, callers that drop a `WritebackFile` without calling
|
||||
// `sync_all` (panic, early-return, idiomatic `let _ = w;`)
|
||||
// leave the trailing chunk in cache; the kernel still flushes
|
||||
// on close, but the bounded-cache invariant fails at the tail.
|
||||
// We deliberately do *not* call `self.file.sync_all()` here —
|
||||
// close already triggers a flush, and an `fsync` from `Drop`
|
||||
// would silently swallow its `io::Error` anyway. `finalize` is
|
||||
// idempotent so an explicit `sync_all` followed by drop is
|
||||
// still safe.
|
||||
self.pipeline.finalize();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
//! Linux platform impl for [`super::WritebackFile`].
|
||||
//!
|
||||
//! - `preallocate`: `fallocate(FALLOC_FL_KEEP_SIZE)` — reserve extents
|
||||
//! without growing the reported file size. Reduces extent
|
||||
//! fragmentation on large sequential writes (mux output on NFS in
|
||||
//! particular).
|
||||
//! - `durable_sync`: `fsync` wrapped in
|
||||
//! [`crate::io::bounded::bounded_syscall`] with a 60 s deadline so a
|
||||
//! wedged NFS server can't trap the calling thread indefinitely.
|
||||
|
||||
use std::fs::File;
|
||||
use std::io;
|
||||
use std::os::unix::io::AsRawFd;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Pre-reserve extents for `size_bytes` of upcoming sequential writes.
|
||||
/// Best-effort: a non-zero rc is logged but not propagated, since the
|
||||
/// caller would just continue with the unreserved file anyway.
|
||||
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,
|
||||
)
|
||||
};
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackFile fallocate size_hint={size_bytes} rc={rc} ok={}",
|
||||
rc == 0
|
||||
);
|
||||
}
|
||||
|
||||
/// 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`.
|
||||
pub(super) fn durable_sync(file: &File) -> io::Result<()> {
|
||||
let fd = file.as_raw_fd();
|
||||
match crate::io::bounded::bounded_syscall(
|
||||
None,
|
||||
Duration::from_secs(60),
|
||||
move || -> io::Result<()> {
|
||||
let rc = unsafe { libc::fsync(fd) };
|
||||
if rc == 0 {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(io::Error::last_os_error())
|
||||
}
|
||||
},
|
||||
) {
|
||||
Ok(inner) => inner,
|
||||
Err(crate::io::bounded::BoundedError::Timeout) => {
|
||||
tracing::error!(
|
||||
target: "mux",
|
||||
"WritebackFile::sync_all fsync timed out after 60s; kernel will flush on close (best-effort)"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(crate::io::bounded::BoundedError::Halted) => Ok(()),
|
||||
Err(crate::io::bounded::BoundedError::WorkerLost) => Ok(()),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
//! 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).
|
||||
//! - `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
|
||||
//! write cache) — what `fsync` should have been on macOS. Falls back
|
||||
//! to plain `fsync` if F_FULLFSYNC returns ENOTSUP.
|
||||
|
||||
use std::fs::File;
|
||||
use std::io;
|
||||
use std::os::unix::io::AsRawFd;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::io::platform_macos::{
|
||||
F_ALLOCATEALL, F_ALLOCATECONTIG, F_PEOFPOSMODE, F_PREALLOCATE, Fstore,
|
||||
};
|
||||
|
||||
/// `fcntl(F_FULLFSYNC)` opcode. Documented in `man 2 fcntl` on macOS;
|
||||
/// not in the `libc` crate as a named constant.
|
||||
const F_FULLFSYNC: libc::c_int = 51;
|
||||
|
||||
pub(super) fn preallocate(file: &File, size_bytes: u64) {
|
||||
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_bytesalloc: 0,
|
||||
};
|
||||
// First attempt: contiguous.
|
||||
let mut rc = unsafe { libc::fcntl(file.as_raw_fd(), F_PREALLOCATE, &mut fst) };
|
||||
if rc == -1 {
|
||||
// Fall back: drop the contiguous hint, allow scattered extents.
|
||||
fst.fst_flags = F_ALLOCATEALL;
|
||||
rc = unsafe { libc::fcntl(file.as_raw_fd(), F_PREALLOCATE, &mut fst) };
|
||||
}
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackFile F_PREALLOCATE size_hint={size_bytes} rc={rc} bytes_allocated={} ok={}",
|
||||
fst.fst_bytesalloc,
|
||||
rc != -1
|
||||
);
|
||||
}
|
||||
|
||||
pub(super) fn durable_sync(file: &File) -> io::Result<()> {
|
||||
let fd = file.as_raw_fd();
|
||||
match crate::io::bounded::bounded_syscall(
|
||||
None,
|
||||
Duration::from_secs(60),
|
||||
move || -> io::Result<()> {
|
||||
// Try F_FULLFSYNC first. If it isn't supported on this
|
||||
// filesystem (older HFS, some network mounts) fall back to
|
||||
// plain fsync — better than nothing.
|
||||
let rc = unsafe { libc::fcntl(fd, F_FULLFSYNC, 0) };
|
||||
if rc == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
let err = io::Error::last_os_error();
|
||||
if err.raw_os_error() == Some(libc::ENOTSUP) {
|
||||
let rc = unsafe { libc::fsync(fd) };
|
||||
if rc == 0 {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(io::Error::last_os_error())
|
||||
}
|
||||
} else {
|
||||
Err(err)
|
||||
}
|
||||
},
|
||||
) {
|
||||
Ok(inner) => inner,
|
||||
Err(crate::io::bounded::BoundedError::Timeout) => {
|
||||
tracing::error!(
|
||||
target: "mux",
|
||||
"WritebackFile::sync_all F_FULLFSYNC timed out after 60s; kernel will flush on close (best-effort)"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(crate::io::bounded::BoundedError::Halted) => Ok(()),
|
||||
Err(crate::io::bounded::BoundedError::WorkerLost) => Ok(()),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,315 @@
|
||||
//! `WritebackFile` — a `File` wrapper whose reason for existing is the
|
||||
//! bounded-cache writeback pipeline.
|
||||
//!
|
||||
//! Why: large sequential writes (sweep, patch, mux on UHD-scale output)
|
||||
//! left to the kernel's default writeback policy accumulate hundreds of
|
||||
//! megabytes of dirty pages and then burst-flush, stalling subsequent
|
||||
//! writes for seconds at a time. `WritebackFile` drives a continuous
|
||||
//! [`super::writeback::WritebackPipeline`] that on Linux issues
|
||||
//! incremental `sync_file_range` + `posix_fadvise(DONTNEED)` calls at
|
||||
//! 32 MB granularity so dirty pages drain at the same rate they're
|
||||
//! produced. macOS and Windows fall through to a no-op pipeline — their
|
||||
//! default cache policies have not been shown to exhibit the same
|
||||
//! pathology for this access pattern.
|
||||
//!
|
||||
//! It implements `Write` and `Seek` so any call site that wrote to a
|
||||
//! plain `File` through those traits (sweep, patch, mux) can swap in
|
||||
//! `WritebackFile` without touching the body of the loop. The wrapper
|
||||
//! also tracks the current file position to feed the pipeline with
|
||||
//! progress + seek boundaries.
|
||||
//!
|
||||
//! See `super::writeback::linux` for the underlying pathology and the
|
||||
//! strategy.
|
||||
//!
|
||||
//! ## 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.
|
||||
//!
|
||||
//! ## 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.
|
||||
//!
|
||||
//! ## 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.
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
mod linux;
|
||||
#[cfg(target_os = "macos")]
|
||||
mod macos;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
|
||||
mod other;
|
||||
#[cfg(target_os = "windows")]
|
||||
mod windows;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
use linux as platform;
|
||||
#[cfg(target_os = "macos")]
|
||||
use macos as platform;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
|
||||
use other as platform;
|
||||
#[cfg(target_os = "windows")]
|
||||
use windows as platform;
|
||||
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::io::{self, Seek, SeekFrom, Write};
|
||||
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.
|
||||
const WRITEBACK_CHUNK_BYTES_DEFAULT: u64 = 32 * 1024 * 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)
|
||||
.map(|n| n * 1024 * 1024)
|
||||
.unwrap_or(WRITEBACK_CHUNK_BYTES_DEFAULT)
|
||||
}
|
||||
|
||||
pub(crate) struct WritebackFile {
|
||||
file: File,
|
||||
pipeline: WritebackPipeline,
|
||||
pos: u64,
|
||||
}
|
||||
|
||||
impl WritebackFile {
|
||||
/// Wrap an open `File`. The current OS file position is queried
|
||||
/// once so the pipeline starts tracking from wherever the file
|
||||
/// already is (typically 0 for fresh files; non-zero for resumed
|
||||
/// or appended files).
|
||||
pub(crate) fn new(mut file: File) -> io::Result<Self> {
|
||||
let pos = file.stream_position()?;
|
||||
let pipeline = WritebackPipeline::new(&file, pos, writeback_chunk_bytes());
|
||||
Ok(Self {
|
||||
file,
|
||||
pipeline,
|
||||
pos,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a new file at `path` (truncating any existing contents)
|
||||
/// and wrap it. Convenience for the common
|
||||
/// `File::create(path)` + `WritebackFile::new(file)` pair so callers
|
||||
/// don't have to assemble a `File` first.
|
||||
///
|
||||
/// Callers that know the target output size should prefer
|
||||
/// [`Self::create_with_size_hint`] so the kernel can pre-reserve
|
||||
/// extents.
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn create(path: &Path) -> io::Result<Self> {
|
||||
let file = File::create(path)?;
|
||||
Self::new(file)
|
||||
}
|
||||
|
||||
/// Like [`Self::create`] but pre-reserves `size_bytes` of disk
|
||||
/// space via the platform's extent-preallocation primitive (Linux
|
||||
/// `fallocate(KEEP_SIZE)`, macOS `F_PREALLOCATE`, Windows
|
||||
/// `SetFileValidData` stub). The reported file size is unchanged
|
||||
/// (writes still grow the file naturally) — only the on-disk extent
|
||||
/// allocation is preallocated, which reduces extent fragmentation
|
||||
/// on large sequential writes (mux output, especially on slow
|
||||
/// storage / NFS).
|
||||
///
|
||||
/// On platforms without an extent-preallocation primitive this is
|
||||
/// equivalent to `create` — the size hint is dropped after a debug
|
||||
/// log.
|
||||
pub(crate) fn create_with_size_hint(path: &Path, size_bytes: u64) -> io::Result<Self> {
|
||||
let file = File::create(path)?;
|
||||
platform::preallocate(&file, size_bytes);
|
||||
Self::new(file)
|
||||
}
|
||||
|
||||
/// Open an existing file at `path` for writing (no truncation) and
|
||||
/// wrap it. Mirrors `File::open` semantics for the writable case
|
||||
/// — used by patch / resume paths that mutate an existing ISO in
|
||||
/// place.
|
||||
pub(crate) fn open(path: &Path) -> io::Result<Self> {
|
||||
let file = OpenOptions::new().write(true).open(path)?;
|
||||
Self::new(file)
|
||||
}
|
||||
|
||||
/// Drain in-flight writeback then issue a full fsync. Use this in
|
||||
/// place of `File::sync_all`.
|
||||
///
|
||||
/// The final durable flush is wrapped in
|
||||
/// [`crate::io::bounded::bounded_syscall`] (per the per-OS module)
|
||||
/// with a 60 s deadline on Linux/macOS — a wedged NFS server cannot
|
||||
/// 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.
|
||||
pub(crate) fn sync_all(&mut self) -> io::Result<()> {
|
||||
self.pipeline.finalize();
|
||||
platform::durable_sync(&self.file)
|
||||
}
|
||||
}
|
||||
|
||||
impl Write for WritebackFile {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
let n = self.file.write(buf)?;
|
||||
self.pos += n as u64;
|
||||
self.pipeline.note_progress(self.pos);
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
|
||||
self.file.write_all(buf)?;
|
||||
self.pos += buf.len() as u64;
|
||||
self.pipeline.note_progress(self.pos);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
self.file.flush()
|
||||
}
|
||||
}
|
||||
|
||||
impl Seek for WritebackFile {
|
||||
fn seek(&mut self, from: SeekFrom) -> io::Result<u64> {
|
||||
let p = self.file.seek(from)?;
|
||||
// Only treat seeks that actually move the position as
|
||||
// boundaries — sweep does a redundant `seek(Current(pos))`
|
||||
// before every write, and we don't want that to drain the
|
||||
// pipeline on every iteration.
|
||||
if p != self.pos {
|
||||
// Diagnostic for the NFS mux hang: the MKV format requires
|
||||
// the muxer to seek back occasionally (cluster size
|
||||
// patching, Cues index write, Segment header backpatch).
|
||||
// Each such seek invalidates the writeback chunk tracking
|
||||
// and forces a finalize → WAIT_AFTER on the in-flight
|
||||
// chunk. Logging the seek delta lets us correlate hang
|
||||
// offsets with specific muxer operations.
|
||||
let from_pos = self.pos;
|
||||
let to_pos = p;
|
||||
let delta: i64 = (to_pos as i64).wrapping_sub(from_pos as i64);
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackFile seek from={from_pos} to={to_pos} delta={delta}"
|
||||
);
|
||||
self.pipeline.handle_seek(p);
|
||||
self.pos = p;
|
||||
}
|
||||
Ok(p)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for WritebackFile {
|
||||
fn drop(&mut self) {
|
||||
// Run the pipeline's tail finalize so the last in-flight chunk
|
||||
// gets its `WAIT_AFTER` + `posix_fadvise(DONTNEED)`. Without
|
||||
// this, callers that drop a `WritebackFile` without calling
|
||||
// `sync_all` (panic, early-return, idiomatic `let _ = w;`)
|
||||
// leave the trailing chunk in cache; the kernel still flushes
|
||||
// on close, but the bounded-cache invariant fails at the tail.
|
||||
// We deliberately do *not* call `self.file.sync_all()` here —
|
||||
// close already triggers a flush, and an `fsync` from `Drop`
|
||||
// would silently swallow its `io::Error` anyway. `finalize` is
|
||||
// idempotent so an explicit `sync_all` followed by drop is
|
||||
// still safe.
|
||||
self.pipeline.finalize();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Read;
|
||||
|
||||
fn read_back(path: &Path) -> Vec<u8> {
|
||||
let mut f = File::open(path).unwrap();
|
||||
let mut v = Vec::new();
|
||||
f.read_to_end(&mut v).unwrap();
|
||||
v
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_then_drop_persists_bytes() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("a.bin");
|
||||
{
|
||||
let mut w = WritebackFile::create(&p).unwrap();
|
||||
w.write_all(b"hello world").unwrap();
|
||||
// Drop drains the pipeline tail.
|
||||
}
|
||||
assert_eq!(read_back(&p), b"hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_all_drains_and_flushes() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("b.bin");
|
||||
let mut w = WritebackFile::create(&p).unwrap();
|
||||
for _ in 0..32 {
|
||||
w.write_all(&[0x5au8; 1024]).unwrap();
|
||||
}
|
||||
// After sync_all, the bytes MUST be visible to a separate
|
||||
// reader. The pipeline has been finalised and durable-sync has
|
||||
// run.
|
||||
w.sync_all().unwrap();
|
||||
let bytes = read_back(&p);
|
||||
assert_eq!(bytes.len(), 32 * 1024);
|
||||
assert!(bytes.iter().all(|&b| b == 0x5a));
|
||||
drop(w);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seek_then_patch_roundtrip() {
|
||||
// Write A; seek back; patch with B; read back; the patch lands
|
||||
// at the right offset.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("c.bin");
|
||||
let mut w = WritebackFile::create(&p).unwrap();
|
||||
let big = vec![b'A'; 4096];
|
||||
w.write_all(&big).unwrap();
|
||||
// Seek back to offset 1000 and overwrite 8 bytes.
|
||||
w.seek(SeekFrom::Start(1000)).unwrap();
|
||||
w.write_all(b"PATCHED!").unwrap();
|
||||
w.sync_all().unwrap();
|
||||
drop(w);
|
||||
let bytes = read_back(&p);
|
||||
assert_eq!(bytes.len(), 4096);
|
||||
assert_eq!(&bytes[1000..1008], b"PATCHED!");
|
||||
// Bytes outside the patch are still 'A'.
|
||||
assert_eq!(bytes[999], b'A');
|
||||
assert_eq!(bytes[1008], b'A');
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flush_is_observed_in_order() {
|
||||
// `Write::flush` should not panic or reorder; verify the bytes
|
||||
// land in order through interleaved flushes.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("f.bin");
|
||||
let mut w = WritebackFile::create(&p).unwrap();
|
||||
w.write_all(b"one").unwrap();
|
||||
w.flush().unwrap();
|
||||
w.write_all(b"two").unwrap();
|
||||
w.flush().unwrap();
|
||||
w.write_all(b"three").unwrap();
|
||||
w.sync_all().unwrap();
|
||||
drop(w);
|
||||
assert_eq!(read_back(&p), b"onetwothree");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
//! Fallback platform impl for [`super::WritebackFile`] on targets
|
||||
//! without a dedicated implementation (BSDs, illumos, etc.).
|
||||
//!
|
||||
//! - `preallocate` is a logged no-op.
|
||||
//! - `durable_sync` calls `File::sync_all` directly (no bounded-syscall
|
||||
//! wrapper — the wrapper depends on Linux/macOS unix idioms that
|
||||
//! aren't universally portable). If a future BSD impl needs the
|
||||
//! 60-s deadline, it should land in its own per-OS file rather than
|
||||
//! bloat this fallback.
|
||||
|
||||
use std::fs::File;
|
||||
use std::io;
|
||||
|
||||
pub(super) fn preallocate(_file: &File, size_bytes: u64) {
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackFile preallocate size_hint={size_bytes} skipped (no impl on this target)"
|
||||
);
|
||||
}
|
||||
|
||||
pub(super) fn durable_sync(file: &File) -> io::Result<()> {
|
||||
file.sync_all()
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
//! 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:
|
||||
//!
|
||||
//! - `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).
|
||||
|
||||
use std::fs::File;
|
||||
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)"
|
||||
);
|
||||
}
|
||||
|
||||
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()
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
//! Key sources — the lookup layer that hands libfreemkv a [`Key`].
|
||||
//!
|
||||
//! libfreemkv performs NO key lookup. An application resolves a key for a disc
|
||||
//! through one or more [`KeySource`]s, each a dumb adapter over a backing store
|
||||
//! (a keydb file, a key server, the mapfile cache): given the disc's
|
||||
//! [`DiscInputs`] it returns the raw [`Key`] at whatever level it holds. The
|
||||
//! library then derives down and decrypts via `Disc::decrypt_with`.
|
||||
//!
|
||||
//! Source implementations are published in the companion `freemkv-keysources`
|
||||
//! crate — keeping all key *policy* (which store, which order, online vs local)
|
||||
//! out of the library while all key *mechanism* (the AACS derivation chain)
|
||||
//! stays in it.
|
||||
|
||||
use crate::disc::Key;
|
||||
|
||||
/// The public AACS inputs a key source needs to look a disc up. Captured at
|
||||
/// scan; contains no secrets — only the disc identity and the on-disc AACS
|
||||
/// structures a source or key server may key on.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DiscInputs {
|
||||
/// SHA-1 of `Unit_Key_RO.inf`, `0x`-prefixed hex. The value a keydb keys
|
||||
/// its per-disc entries by, and a key server identifies the disc with.
|
||||
pub disc_hash: String,
|
||||
/// Volume ID (16 bytes). `[0u8; 16]` when no authenticated handshake ran
|
||||
/// (e.g. an ISO/mapfile flow), which disables VID-keyed lookups.
|
||||
pub volume_id: [u8; 16],
|
||||
/// Raw MKB bytes. Empty when not captured.
|
||||
pub mkb: Vec<u8>,
|
||||
/// Raw `Unit_Key_RO.inf` bytes. Empty when not captured.
|
||||
pub unit_key_ro: Vec<u8>,
|
||||
/// Encrypted on-disc content sample units (each a 6144-byte aligned unit),
|
||||
/// for sources that validate a key server-side against real ciphertext
|
||||
/// (e.g. an online key service). Empty for sources that don't need them
|
||||
/// (a local keydb). Populated by the application — reading content requires
|
||||
/// the disc reader, which the library's scan does not retain — so
|
||||
/// [`crate::Disc::inputs`] leaves it empty for the caller to fill.
|
||||
pub samples: Vec<Vec<u8>>,
|
||||
}
|
||||
|
||||
/// A key source: a stateful provider that hands a disc's candidate [`Key`]s out
|
||||
/// **one at a time**, in whatever order it judges best for its backing store.
|
||||
///
|
||||
/// Dumb by contract — a source queries its store and yields the raw material it
|
||||
/// holds at whatever level it has (device / processing / media / volume / unit).
|
||||
/// It performs NO AACS derivation and NO validation: `Disc::decrypt_with`
|
||||
/// derives down AND validates against real ciphertext, returning `Err` for a key
|
||||
/// that does not decrypt this disc. That keeps every derivation step and the one
|
||||
/// validation gate in the library, across AACS 1.0 / 2.0 / 2.1 / 2.x.
|
||||
///
|
||||
/// The source is the one that knows how many candidates it has and in what order
|
||||
/// to try them — a keydb holds a per-disc UK *and* VUK *and* a device-key pool,
|
||||
/// so it hands them out cheapest/most-specific first (UK ▸ VK ▸ MK ▸ DK) and
|
||||
/// reports exhaustion when its list runs out; an online key service or a mapfile
|
||||
/// 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`.
|
||||
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.
|
||||
fn next_key(&mut self, inputs: &DiscInputs) -> Option<Key>;
|
||||
|
||||
/// Whether this source needs [`DiscInputs::samples`] populated (encrypted
|
||||
/// content samples) — true for a source that validates against ciphertext
|
||||
/// server-side, false for one that keys purely on disc identity. The caller
|
||||
/// reads samples (an extra disc read) only when some source needs them.
|
||||
fn needs_samples(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Whether this source FAILED (I/O, network, parse) rather than simply
|
||||
/// having no key. Checked after exhaustion so the caller can tell a genuine
|
||||
/// "no key for this disc" apart from "the key service was unreachable". A
|
||||
/// store that treats absence as not-an-error (a missing keydb / mapfile)
|
||||
/// leaves this `false`.
|
||||
fn errored(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
+24
-24
@@ -38,7 +38,7 @@ use std::collections::BTreeMap;
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Default)]
|
||||
pub struct DiscMetadata {
|
||||
/// Localized titles, keyed by 3-char ISO 639-2 lang code
|
||||
/// (e.g. "eng" → "Dune Part Two")
|
||||
/// (e.g. "eng" → "Aurora Drift")
|
||||
pub titles: BTreeMap<String, String>,
|
||||
/// First-line / short description, per lang
|
||||
pub descriptions: BTreeMap<String, String>,
|
||||
@@ -146,7 +146,7 @@ pub(crate) fn parse_bdmt_xml(_lang_code: &str, xml_text: &str) -> Option<BdmtFie
|
||||
}
|
||||
|
||||
/// Reject candidate description strings that are themselves XML
|
||||
/// fragments — observed on disc-04 (Top Gun: Maverick), where
|
||||
/// fragments — observed on a captured disc, where
|
||||
/// `<di:description>` contained `<di:thumbnail href="…"/>` child
|
||||
/// elements and no actual prose. Surfacing that raw to the JSON
|
||||
/// output is worse than dropping the field entirely.
|
||||
@@ -210,10 +210,10 @@ mod tests {
|
||||
// carrier inside a <discInfo> root.
|
||||
let xml = r#"<?xml version="1.0" encoding="UTF-8"?>
|
||||
<discInfo xmlns:di="urn:BDA:bdmv;disclibmeta">
|
||||
<di:name>Dune Part Two</di:name>
|
||||
<di:name>Aurora Drift</di:name>
|
||||
</discInfo>"#;
|
||||
let (title, desc, set) = parse_bdmt_xml("eng", xml).expect("title should parse");
|
||||
assert_eq!(title, "Dune Part Two");
|
||||
assert_eq!(title, "Aurora Drift");
|
||||
assert_eq!(desc, None);
|
||||
assert_eq!(set, None);
|
||||
}
|
||||
@@ -223,12 +223,12 @@ mod tests {
|
||||
// <di:title> is the alternate carrier; should be picked up
|
||||
// when <di:name> is absent.
|
||||
let xml = r#"<discInfo xmlns:di="urn:BDA:bdmv;disclibmeta">
|
||||
<di:title>The Matrix</di:title>
|
||||
<di:description>A film about computers.</di:description>
|
||||
<di:title>Echo Chamber</di:title>
|
||||
<di:description>A film about machines.</di:description>
|
||||
</discInfo>"#;
|
||||
let (title, desc, _) = parse_bdmt_xml("eng", xml).unwrap();
|
||||
assert_eq!(title, "The Matrix");
|
||||
assert_eq!(desc.as_deref(), Some("A film about computers."));
|
||||
assert_eq!(title, "Echo Chamber");
|
||||
assert_eq!(desc.as_deref(), Some("A film about machines."));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -238,17 +238,17 @@ mod tests {
|
||||
// titleName inside tableOfContents.
|
||||
let xml = r#"<discInfo xmlns:di="urn:BDA:bdmv;disclibmeta">
|
||||
<di:tableOfContents>
|
||||
<di:titleName>Inside Out 2</di:titleName>
|
||||
<di:titleName>Feelings Two</di:titleName>
|
||||
</di:tableOfContents>
|
||||
</discInfo>"#;
|
||||
let (title, _, _) = parse_bdmt_xml("eng", xml).unwrap();
|
||||
assert_eq!(title, "Inside Out 2");
|
||||
assert_eq!(title, "Feelings Two");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_box_set_position() {
|
||||
let xml = r#"<discInfo xmlns:di="urn:BDA:bdmv;disclibmeta">
|
||||
<di:name>LOTR Disc 2</di:name>
|
||||
<di:name>Box Set Disc 2</di:name>
|
||||
<di:discNumber>2</di:discNumber>
|
||||
<di:numSets>5</di:numSets>
|
||||
</discInfo>"#;
|
||||
@@ -287,11 +287,11 @@ mod tests {
|
||||
// would. This exercises the BTreeMap key handling without
|
||||
// needing a UdfFs.
|
||||
let eng_xml = r#"<discInfo xmlns:di="urn:BDA:bdmv;disclibmeta">
|
||||
<di:name>Dune Part Two</di:name>
|
||||
<di:name>Aurora Drift</di:name>
|
||||
</discInfo>"#;
|
||||
let fra_xml = r#"<discInfo xmlns:di="urn:BDA:bdmv;disclibmeta">
|
||||
<di:name>Dune Deuxième Partie</di:name>
|
||||
<di:description>Suite du film de 2021.</di:description>
|
||||
<di:name>Aurora Drift (Partie Deux)</di:name>
|
||||
<di:description>Suite du film fictif.</di:description>
|
||||
</discInfo>"#;
|
||||
|
||||
let mut meta = DiscMetadata::default();
|
||||
@@ -310,16 +310,16 @@ mod tests {
|
||||
|
||||
assert_eq!(
|
||||
meta.titles.get("eng").map(String::as_str),
|
||||
Some("Dune Part Two")
|
||||
Some("Aurora Drift")
|
||||
);
|
||||
assert_eq!(
|
||||
meta.titles.get("fra").map(String::as_str),
|
||||
Some("Dune Deuxième Partie")
|
||||
Some("Aurora Drift (Partie Deux)")
|
||||
);
|
||||
assert!(meta.descriptions.get("eng").is_none());
|
||||
assert_eq!(
|
||||
meta.descriptions.get("fra").map(String::as_str),
|
||||
Some("Suite du film de 2021.")
|
||||
Some("Suite du film fictif.")
|
||||
);
|
||||
assert_eq!(meta.disc_number, None);
|
||||
}
|
||||
@@ -342,21 +342,21 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn description_with_only_child_xml_is_dropped() {
|
||||
// Real-world bug from disc-04 (Top Gun: Maverick, 2026-05-11
|
||||
// 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 `<`.
|
||||
let xml = r#"<discInfo>
|
||||
<di:name>Top Gun: Maverick</di:name>
|
||||
<di:name>Skyline Run</di:name>
|
||||
<di:description>
|
||||
<di:thumbnail href="tgm_meta_sm.jpg" />
|
||||
<di:thumbnail href="tgm_meta_lg.jpg" />
|
||||
<di:thumbnail href="sample_meta_sm.jpg" />
|
||||
<di:thumbnail href="sample_meta_lg.jpg" />
|
||||
</di:description>
|
||||
</discInfo>"#;
|
||||
let (title, description, _) =
|
||||
parse_bdmt_xml("eng", xml).expect("title is present so parse must succeed");
|
||||
assert_eq!(title, "Top Gun: Maverick");
|
||||
assert_eq!(title, "Skyline Run");
|
||||
assert!(
|
||||
description.is_none(),
|
||||
"description containing only XML children must be dropped, got {description:?}"
|
||||
@@ -381,10 +381,10 @@ mod tests {
|
||||
#[test]
|
||||
fn whitespace_in_title_is_trimmed() {
|
||||
let xml = r#"<discInfo><di:name>
|
||||
Dune Part Two
|
||||
Aurora Drift
|
||||
</di:name></discInfo>"#;
|
||||
let (title, _, _) = parse_bdmt_xml("eng", xml).unwrap();
|
||||
assert_eq!(title, "Dune Part Two");
|
||||
assert_eq!(title, "Aurora Drift");
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
+1
-2
@@ -7,8 +7,7 @@
|
||||
//! Stream labels live as plain ASCII strings inside compiled `.class`
|
||||
//! files in the jar — a quirk of the menu-rendering layer encoding
|
||||
//! its TextField positions and content as constant strings the
|
||||
//! Java compiler retained in the class string pool. Format observed
|
||||
//! in the corpus (Civil War UHD, 2024):
|
||||
//! Java compiler retained in the class string pool. Observed format:
|
||||
//!
|
||||
//! ```text
|
||||
//! LTextField,Audio1,English Dolby Atmos,Fontstrip_Composite,...
|
||||
|
||||
+38
-60
@@ -1,83 +1,65 @@
|
||||
//! Deluxe BD-J framework — `com/bydeluxe/bluray/` package signature.
|
||||
//!
|
||||
//! Used by major studios (Disney, Warner, others) for their UHD
|
||||
//! BD-J authoring. Detected on discs whose `/BDMV/JAR/<x>.jar`
|
||||
//! contains a `com/bydeluxe/` directory entry.
|
||||
//! Detected on discs whose `/BDMV/JAR/<x>.jar` contains a
|
||||
//! `com/bydeluxe/` directory entry.
|
||||
//!
|
||||
//! ## Why this parser exists
|
||||
//! ## What this parser reads
|
||||
//!
|
||||
//! Deluxe-authored discs store stream labels as **ordinal references
|
||||
//! into obfuscated enum classes**. The label text isn't a literal
|
||||
//! string in any anchor pattern (unlike dbp's `TextField,...` rows).
|
||||
//! Instead, the binding code is roughly:
|
||||
//!
|
||||
//! ```java
|
||||
//! streamTable.put(1, new AudioSlot(LanguageEnum.English,
|
||||
//! CodecEnum.ATMOS_HD_AUDIO,
|
||||
//! PurposeEnum.Normal));
|
||||
//! ```
|
||||
//!
|
||||
//! The class names `LanguageEnum`, `CodecEnum`, `PurposeEnum`, and
|
||||
//! `AudioSlot` are obfuscated per-disc (`be.class`, `ma.class`,
|
||||
//! `lp.class`, etc.) — no name pattern survives the obfuscator. But
|
||||
//! the **shape of `<clinit>`** is framework-stable:
|
||||
//! Deluxe-authored discs store stream labels as ordinal references into
|
||||
//! enum classes whose names are obfuscated per-disc, so a name-based
|
||||
//! match won't work. The label data is instead recovered by matching on
|
||||
//! the **shape of each enum's `<clinit>`**, which is framework-stable:
|
||||
//!
|
||||
//! | Enum | Signature |
|
||||
//! |---|---|
|
||||
//! | Language | 70 `ldc` operations in `<clinit>`, sequence starts `English, French, Spanish, Dutch, ...` |
|
||||
//! | 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, ...` (Disney only) |
|
||||
//! | Studio | 6 ldcs starting `Disney, Marvel, Pixar, ...` (Disney only) |
|
||||
//! | Codec | ~46 `new` instructions, 0 ldcs in `<clinit>` (codec strings live in subclasses) |
|
||||
//! | 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) |
|
||||
//!
|
||||
//! Match on the SHAPE, not the name, and the parser survives obfuscation.
|
||||
//! Matching on the shape rather than the class name keeps the parser
|
||||
//! working across obfuscation variants.
|
||||
//!
|
||||
//! ## Implementation phases
|
||||
//!
|
||||
//! - **Phase A** — master enum identification (`identify_master_enums`).
|
||||
//! Walks every `.class`'s `<clinit>` ldc sequence and matches against
|
||||
//! the framework-stable fingerprints. Output: `Vec<(label, MasterEnum)>`
|
||||
//! with full ordinal → string-value tables. **Empirically verified**
|
||||
//! on disc-01 (Disney) + disc-09 (Warner).
|
||||
//! with full ordinal → string-value tables.
|
||||
//!
|
||||
//! - **Phase B** — codec enum subclass walk (`decode_codec_enum`).
|
||||
//! The codec enum's `<clinit>` has ~46 `new` instructions and zero
|
||||
//! 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. **Structural shape
|
||||
//! verified** on disc-01 (ma.class, 41 `new` ops) + disc-09
|
||||
//! (ea.class, 46 `new` ops); per-subclass string extraction
|
||||
//! designed against the published Java enum compilation convention
|
||||
//! (each enum value's `<init>` is called with its name string as
|
||||
//! the first arg).
|
||||
//! 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.
|
||||
//! That class has the highest such `getstatic` count in the jar.
|
||||
//! **Heuristic shape**; precise threshold may need tuning.
|
||||
//! Heuristic shape; precise threshold may need tuning.
|
||||
//!
|
||||
//! - **Phase D** — binding-class bytecode decoder (`decode_binding`).
|
||||
//! Walks the binding class's `<clinit>` with a tiny symbolic stack
|
||||
//! machine. For each `new X / dup / ... / invokespecial X.<init>`
|
||||
//! sequence, collects the int values and enum-reference operands
|
||||
//! between the `dup` and the constructor call, then emits a
|
||||
//! `DecodedStream`. **Mechanism verified** in unit tests against
|
||||
//! synthetic class fixtures; the **signal-to-StreamLabel mapping**
|
||||
//! (which arg is stream index? which is language? audio vs
|
||||
//! subtitle?) uses a documented heuristic that needs corpus-disc
|
||||
//! verification — see `interpret_stream` for the mapping rules.
|
||||
//! `DecodedStream`. The signal-to-StreamLabel mapping (which arg is
|
||||
//! stream index? which is language? audio vs subtitle?) uses a
|
||||
//! heuristic — see `interpret_streams` for the mapping rules.
|
||||
//!
|
||||
//! ## Confidence
|
||||
//!
|
||||
//! [`parse`] returns `Some(ParseResult::medium(labels))` when Phases A
|
||||
//! through D produce at least one stream — `Medium` because the
|
||||
//! signal-to-label mapping is heuristic until real disc bytecode
|
||||
//! confirms the binding pattern. Once verified the parser can promote
|
||||
//! to `High`. `None` when the disc isn't Deluxe-authored or when
|
||||
//! decoding produces zero streams (a recognized-but-broken state that
|
||||
//! the analyzer still surfaces via `parsers_detected`).
|
||||
//! signal-to-label mapping is heuristic. `None` when the disc isn't
|
||||
//! Deluxe-authored or when decoding produces zero streams (a
|
||||
//! recognized-but-broken state that the analyzer still surfaces via
|
||||
//! `parsers_detected`).
|
||||
|
||||
use super::class_reader::{
|
||||
AASTORE, BIPUSH, ClassFile, CodeAttribute, ConstantPool, CpInfo, GETSTATIC, ICONST_0, ICONST_1,
|
||||
@@ -559,10 +541,9 @@ fn clinit_news_and_ldcs(
|
||||
/// candidates ordered by descending getstatic count, filtered to a
|
||||
/// minimum concentration of master-enum references.
|
||||
///
|
||||
/// Empirically (POC v0.3 dumps): on disc-01 the audio binding class
|
||||
/// (`ma.class`) has ~82 getstatic refs, the subtitle binding
|
||||
/// (`ko.class`) has ~63. Both share the master Language + Purpose
|
||||
/// enums.
|
||||
/// A disc that splits the table commonly has one audio binding class
|
||||
/// with the most getstatic refs and a subtitle binding class with
|
||||
/// somewhat fewer; both share the master Language + Purpose enums.
|
||||
pub(crate) fn find_binding_classes(
|
||||
archive: &mut jar::Jar,
|
||||
master_enum_classes: &HashSet<&str>,
|
||||
@@ -641,10 +622,9 @@ pub(crate) enum StackVal {
|
||||
/// name (e.g. `DOLBY_AC3_AUDIO`, `DOLBY_LOSSLESS_AUDIO`) is the
|
||||
/// codec identifier. Deluxe binding constructors take a
|
||||
/// `LCodingType;` arg directly — codecs are NOT a Deluxe-internal
|
||||
/// enum (Phase B's codec-subclass walk was based on a wrong
|
||||
/// assumption; the actual codec source is the standard BD-J API
|
||||
/// enum). Discovered via deluxe-poc v0.3 binding-bytecode dump
|
||||
/// against disc-01 (Disney) + disc-09 (Warner) on 2026-05-10.
|
||||
/// enum; the codec source is the standard BD-J API `CodingType`
|
||||
/// enum, so the binding constructor's codec arg is read straight
|
||||
/// from that getstatic operand.
|
||||
CodingType(String),
|
||||
/// An uninitialized `new` object — popped by the matching
|
||||
/// invokespecial.
|
||||
@@ -988,11 +968,10 @@ impl MasterEnumTable {
|
||||
// ── interpret_streams: Constructions → StreamLabels ─────────────────────────
|
||||
|
||||
/// Convert the per-construction tuples from Phase D into
|
||||
/// [`StreamLabel`]s. Pattern verified against corpus discs via
|
||||
/// deluxe-poc v0.3 binding-bytecode dump (2026-05-10):
|
||||
/// [`StreamLabel`]s. Two binding-constructor shapes are handled:
|
||||
///
|
||||
/// Disney binding (5-arg): `BindingType.<init>(I, Lbe;, Llp;, I, LCodingType;)V`
|
||||
/// Warner binding (4-arg): `BindingType.<init>(I, Law;, Lgp;, LCodingType;)V`
|
||||
/// 5-arg: `BindingType.<init>(I, Lang;, Lpurpose;, I, LCodingType;)V`
|
||||
/// 4-arg: `BindingType.<init>(I, Lang;, Lpurpose;, LCodingType;)V`
|
||||
///
|
||||
/// Args are identified by **TYPE**, not position:
|
||||
/// - First `EnumRef{kind: "Language"}` → audio/subtitle language
|
||||
@@ -1098,11 +1077,10 @@ fn interpret_streams(constructions: &[Construction], master: &MasterEnumTable) -
|
||||
/// getstatic operands on Deluxe binding classes) to a human-readable
|
||||
/// codec hint string.
|
||||
///
|
||||
/// CodingType is the standard BD-J API enum; values are documented
|
||||
/// in the BD-J specification and verified empirically against the
|
||||
/// binding-bytecode dumps in `(internal)/research/deluxe-poc/data/`.
|
||||
/// Unknown field names pass through unchanged so unfamiliar codecs
|
||||
/// still surface something rather than going silent.
|
||||
/// CodingType is the standard BD-J API enum; values are documented in
|
||||
/// the BD-J specification. Unknown field names pass through unchanged
|
||||
/// so unfamiliar codecs still surface something rather than going
|
||||
/// silent.
|
||||
fn coding_type_to_codec_hint(field: &str) -> &str {
|
||||
match field {
|
||||
// Lossless / hi-res.
|
||||
|
||||
+186
-6
@@ -207,14 +207,41 @@ pub(crate) fn apply_labels(labels: &[StreamLabel], titles: &mut [DiscTitle]) {
|
||||
// Structured fields — callers translate purpose to UI text.
|
||||
a.purpose = label.purpose;
|
||||
|
||||
// Codec descriptor: trust the parser's `codec_hint` ONLY
|
||||
// when it's consistent with the stream's actual codec — it
|
||||
// may legitimately be richer (e.g. "Dolby Atmos" on a TrueHD
|
||||
// stream, which the raw spec codec can't express). If the
|
||||
// hint CONTRADICTS the stream (a mis-bound / shuffled label,
|
||||
// e.g. "AC-3 2.0" on a TrueHD track, or "TrueHD" on a DD+
|
||||
// track), discard it and derive the descriptor from the
|
||||
// stream itself — that's correct per-stream and can never be
|
||||
// shuffled. An empty hint is left for `fill_defaults`.
|
||||
let codec_desc = if label.codec_hint.is_empty() {
|
||||
// No codec hint — leave for fill_defaults.
|
||||
String::new()
|
||||
} else if !codec_hint_consistent(&label.codec_hint, &a.codec) {
|
||||
// Hint contradicts the stream (mis-bound / shuffled):
|
||||
// derive from the stream itself.
|
||||
generate_audio_label(&a.codec, &a.channels, a.secondary)
|
||||
} else if codec_hint_adds_detail(&label.codec_hint) {
|
||||
// Consistent AND richer than the spec codec can express
|
||||
// (e.g. "Dolby Atmos", "DTS:X") — keep the parser's hint.
|
||||
label.codec_hint.clone()
|
||||
} else {
|
||||
// Consistent but a plain codec/channel restatement —
|
||||
// normalize to the stream's own marketing descriptor so
|
||||
// styling is uniform across tracks.
|
||||
generate_audio_label(&a.codec, &a.channels, a.secondary)
|
||||
};
|
||||
|
||||
// a.label only carries codec/variant info. NEVER any
|
||||
// English purpose text — the CLI handles that via i18n.
|
||||
let mut parts = Vec::new();
|
||||
if !label.variant.is_empty() {
|
||||
parts.push(format!("({})", label.variant));
|
||||
}
|
||||
if !label.codec_hint.is_empty() {
|
||||
parts.push(label.codec_hint.clone());
|
||||
if !codec_desc.is_empty() {
|
||||
parts.push(codec_desc);
|
||||
}
|
||||
if !parts.is_empty() {
|
||||
a.label = parts.join(" ");
|
||||
@@ -322,7 +349,70 @@ fn generate_video_label(
|
||||
parts.join(" ")
|
||||
}
|
||||
|
||||
fn generate_audio_label(
|
||||
/// Does the parser's `codec_hint` name a codec consistent with the stream's
|
||||
/// actual `codec`? [`apply_labels`] uses this to keep richer-but-consistent
|
||||
/// hints (e.g. "Dolby Atmos" on a TrueHD stream — Atmos is a TrueHD extension
|
||||
/// the raw spec codec can't express) while rejecting mis-bound ones (e.g.
|
||||
/// "AC-3 2.0" on a TrueHD stream, the shuffled-label bug). Matching is by codec
|
||||
/// FAMILY parsed out of the hint string. "Atmos" with no carrier named is
|
||||
/// treated as compatible with its lossless carriers (TrueHD / E-AC-3). A hint
|
||||
/// naming no recognizable codec family (pure editorial, e.g. "Commentary") is
|
||||
/// consistent — it isn't asserting a codec.
|
||||
fn codec_hint_consistent(hint: &str, codec: &crate::disc::Codec) -> bool {
|
||||
use crate::disc::Codec;
|
||||
let h = hint.to_ascii_lowercase();
|
||||
|
||||
let says_truehd = h.contains("truehd") || h.contains("true hd");
|
||||
let says_ddp = h.contains("ac-3+")
|
||||
|| h.contains("ac3+")
|
||||
|| h.contains("e-ac-3")
|
||||
|| h.contains("eac-3")
|
||||
|| h.contains("eac3")
|
||||
|| h.contains("digital plus")
|
||||
|| h.contains("dd+");
|
||||
let says_ac3 =
|
||||
!says_ddp && (h.contains("ac-3") || h.contains("ac3") || h.contains("dolby digital"));
|
||||
let says_dts_ma = h.contains("master audio") || h.contains("hd ma");
|
||||
let says_dts_hr = h.contains("high resolution") || h.contains("hd hr");
|
||||
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");
|
||||
|
||||
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+).
|
||||
if !names_family {
|
||||
return if says_atmos {
|
||||
matches!(codec, Codec::TrueHd | Codec::Ac3Plus)
|
||||
} else {
|
||||
true
|
||||
};
|
||||
}
|
||||
|
||||
match codec {
|
||||
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::Dts => says_dts,
|
||||
Codec::Lpcm => says_lpcm,
|
||||
// Unknown / other stream codec — don't second-guess the parser's hint.
|
||||
_ => true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Does the hint carry object-audio detail the spec codec can't express
|
||||
/// (Atmos / DTS:X)? Such hints are kept verbatim; plain codec/channel hints are
|
||||
/// normalized to the stream's own descriptor for uniform styling across tracks.
|
||||
fn codec_hint_adds_detail(hint: &str) -> bool {
|
||||
let h = hint.to_ascii_lowercase();
|
||||
h.contains("atmos") || h.contains("dts:x") || h.contains("dts-x") || h.contains("dtsx")
|
||||
}
|
||||
|
||||
pub(crate) fn generate_audio_label(
|
||||
codec: &crate::disc::Codec,
|
||||
channels: &crate::disc::AudioChannels,
|
||||
_secondary: bool,
|
||||
@@ -954,7 +1044,7 @@ mod gap_fill_tests {
|
||||
|
||||
#[test]
|
||||
fn partial_yield_fills_gaps_keeps_framework() {
|
||||
// Oppenheimer-style: framework matched but only labeled 2 of 6 audios.
|
||||
// Partial-coverage case: framework matched but only labeled 2 of 6 audios.
|
||||
let mut framework = vec![
|
||||
label(StreamLabelType::Audio, 1, "eng", "Atmos"),
|
||||
label(StreamLabelType::Audio, 4, "eng", "Commentary"),
|
||||
@@ -1182,6 +1272,94 @@ mod apply_tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apply_rejects_mismatched_codec_hint_and_uses_stream_codec() {
|
||||
// Paddington case: a TrueHD+Atmos main track the parser mislabeled
|
||||
// "AC-3 2.0" (a compat-core hint bound to the wrong stream). The hint
|
||||
// contradicts the stream's real codec → discard it, use the stream's own.
|
||||
let mut titles = vec![title_with(vec![audio(
|
||||
0x1100,
|
||||
Codec::TrueHd,
|
||||
AudioChannels::Surround71,
|
||||
"eng",
|
||||
)])];
|
||||
let labels = vec![audio_label(1, "eng", "AC-3 2.0", "")];
|
||||
apply_labels(&labels, &mut titles);
|
||||
if let Stream::Audio(a) = &titles[0].streams[0] {
|
||||
assert_eq!(a.label, "Dolby TrueHD 7.1");
|
||||
} else {
|
||||
panic!("expected audio stream");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apply_unshuffles_cross_labeled_streams() {
|
||||
// Wicked case: hints fully cross-bound — a TrueHD stream wears "AC-3 5.1"
|
||||
// and a DD+ stream wears "TrueHD 5.1". Each is corrected from its own
|
||||
// stream codec, eliminating the shuffle.
|
||||
let mut titles = vec![title_with(vec![
|
||||
audio(0x1100, Codec::TrueHd, AudioChannels::Surround51, "eng"),
|
||||
audio(0x1101, Codec::Ac3Plus, AudioChannels::Surround51, "spa"),
|
||||
])];
|
||||
let labels = vec![
|
||||
audio_label(1, "eng", "AC-3 5.1", ""),
|
||||
audio_label(2, "spa", "TrueHD 5.1", ""),
|
||||
];
|
||||
apply_labels(&labels, &mut titles);
|
||||
let got: Vec<String> = titles[0]
|
||||
.streams
|
||||
.iter()
|
||||
.filter_map(|s| {
|
||||
if let Stream::Audio(a) = s {
|
||||
Some(a.label.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(got, vec!["Dolby TrueHD 5.1", "Dolby Digital Plus 5.1"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apply_keeps_consistent_richer_atmos_hint() {
|
||||
// A DD+ Atmos stream legitimately labeled "Dolby Atmos" — the hint is
|
||||
// richer than the spec codec yet consistent with it, so it's kept.
|
||||
let mut titles = vec![title_with(vec![audio(
|
||||
0x1100,
|
||||
Codec::Ac3Plus,
|
||||
AudioChannels::Surround51,
|
||||
"eng",
|
||||
)])];
|
||||
let labels = vec![audio_label(1, "eng", "Dolby Atmos", "")];
|
||||
apply_labels(&labels, &mut titles);
|
||||
if let Stream::Audio(a) = &titles[0].streams[0] {
|
||||
assert_eq!(a.label, "Dolby Atmos");
|
||||
} else {
|
||||
panic!("expected audio stream");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apply_normalizes_plain_consistent_hint_to_marketing() {
|
||||
// Wicked's French track: a DD+ stream whose hint "AC-3+ 5.1" is correct
|
||||
// but short-form. A sibling DD+ track that fell back uses the marketing
|
||||
// form — keeping the short form here would read inconsistently, so a
|
||||
// plain (non-richer) consistent hint is normalized to the stream's own.
|
||||
let mut titles = vec![title_with(vec![audio(
|
||||
0x1100,
|
||||
Codec::Ac3Plus,
|
||||
AudioChannels::Surround51,
|
||||
"fra",
|
||||
)])];
|
||||
let labels = vec![audio_label(1, "fra", "AC-3+ 5.1", "")];
|
||||
apply_labels(&labels, &mut titles);
|
||||
if let Stream::Audio(a) = &titles[0].streams[0] {
|
||||
assert_eq!(a.label, "Dolby Digital Plus 5.1");
|
||||
} else {
|
||||
panic!("expected audio stream");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apply_sets_purpose_on_audio_commentary() {
|
||||
let mut titles = vec![title_with(vec![audio(
|
||||
@@ -1342,9 +1520,11 @@ mod apply_tests {
|
||||
if let Stream::Audio(a) = &titles[0].streams[1] {
|
||||
assert_eq!(a.label, "Dolby Atmos");
|
||||
}
|
||||
// Audio #2 (4th stream overall)
|
||||
// Audio #2 (4th stream overall). The plain "Dolby Digital" hint is
|
||||
// consistent with the AC-3 stream but carries no channel info, so it's
|
||||
// normalized to the stream's own uniform descriptor.
|
||||
if let Stream::Audio(a) = &titles[0].streams[3] {
|
||||
assert_eq!(a.label, "Dolby Digital");
|
||||
assert_eq!(a.label, "Dolby Digital 2.0");
|
||||
}
|
||||
// Subtitle #1
|
||||
if let Stream::Subtitle(s) = &titles[0].streams[2] {
|
||||
|
||||
+126
-12
@@ -37,14 +37,47 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult>
|
||||
// but the corpus surfaced something we don't catalogue.
|
||||
let saw_unknown = AtomicBool::new(false);
|
||||
|
||||
let labels = assign_labels(&strings, &saw_unknown);
|
||||
|
||||
if labels.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let confidence = if saw_unknown.load(Ordering::Relaxed) {
|
||||
Confidence::Medium
|
||||
} else {
|
||||
Confidence::High
|
||||
};
|
||||
Some(ParseResult { labels, confidence })
|
||||
}
|
||||
|
||||
/// Walk the extracted token strings of the feature section and emit a
|
||||
/// `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> {
|
||||
// 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
|
||||
// stray stream tokens — e.g. a `SEG_MainFeature` preview segment that
|
||||
// lists only a commentary track. Anchoring on such a segment grabs the
|
||||
// wrong streams and misnumbers them. So when the project ships any
|
||||
// `FPL_` playlist, anchor exclusively on it; only fall back to
|
||||
// `SEG_MainFeature` on discs that have no `FPL_` section at all.
|
||||
let has_fpl = strings.iter().any(|s| s.starts_with("FPL_"));
|
||||
|
||||
let mut labels = Vec::new();
|
||||
let mut in_feature = false;
|
||||
let mut audio_num: u16 = 0;
|
||||
let mut sub_num: u16 = 0;
|
||||
|
||||
for s in &strings {
|
||||
for s in strings {
|
||||
// Detect feature section start
|
||||
if s.starts_with("FPL_") || s.starts_with("SEG_MainFeature") {
|
||||
let is_start = if has_fpl {
|
||||
s.starts_with("FPL_")
|
||||
} else {
|
||||
s.starts_with("SEG_MainFeature")
|
||||
};
|
||||
if is_start {
|
||||
if in_feature {
|
||||
break;
|
||||
}
|
||||
@@ -63,7 +96,25 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult>
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Some(label) = parse_token_inner(s, Some(&saw_unknown)) {
|
||||
// Generic audio-slot placeholder. Pixelogic lists each audio stream
|
||||
// in the playlist as either an editorial `{lang}_{codec}_…` token OR
|
||||
// a bare `Audio Stream N` placeholder when no editorial label was
|
||||
// authored. The placeholder carries nothing to label, but it DOES
|
||||
// occupy an STN slot — so it must advance `audio_num`. Otherwise a
|
||||
// later editorial token (e.g. a lone `eng_ACOM_` commentary sitting
|
||||
// at STN slot 4, behind three unlabelled main tracks) collapses onto
|
||||
// slot 1 and its Commentary purpose lands on the main feature track.
|
||||
//
|
||||
// Only audio is corrected here: subtitle (`PG Stream N`) numbering is
|
||||
// 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.
|
||||
if s.starts_with("Audio Stream") {
|
||||
audio_num += 1;
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Some(label) = parse_token_inner(s, Some(saw_unknown)) {
|
||||
match label.stream_type {
|
||||
StreamLabelType::Audio => {
|
||||
audio_num += 1;
|
||||
@@ -83,15 +134,7 @@ pub fn parse(reader: &mut dyn SectorSource, udf: &UdfFs) -> Option<ParseResult>
|
||||
}
|
||||
}
|
||||
|
||||
if labels.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let confidence = if saw_unknown.load(Ordering::Relaxed) {
|
||||
Confidence::Medium
|
||||
} else {
|
||||
Confidence::High
|
||||
};
|
||||
Some(ParseResult { labels, confidence })
|
||||
labels
|
||||
}
|
||||
|
||||
fn parse_token_inner(s: &str, saw_unknown: Option<&AtomicBool>) -> Option<StreamLabel> {
|
||||
@@ -250,4 +293,75 @@ mod tests {
|
||||
assert!(parse_token_inner("XX_MLP_", None).is_none());
|
||||
assert!(parse_token_inner("ENG_MLP_", None).is_none()); // uppercase not accepted as ISO 639-2
|
||||
}
|
||||
|
||||
fn strs(v: &[&str]) -> Vec<String> {
|
||||
v.iter().map(|s| s.to_string()).collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assign_labels_numbers_commentary_behind_placeholders() {
|
||||
// Observed case: the FPL_MainFeature playlist lists three unlabelled main
|
||||
// audio tracks as `Audio Stream N` placeholders, then a lone
|
||||
// `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 tokens = strs(&[
|
||||
"FPL_MainFeature",
|
||||
"Audio Stream 1",
|
||||
"Audio Stream 2",
|
||||
"Audio Stream 3",
|
||||
"eng_ACOM_",
|
||||
]);
|
||||
let labels = assign_labels(&tokens, &flag);
|
||||
let audio: Vec<_> = labels
|
||||
.iter()
|
||||
.filter(|l| l.stream_type == StreamLabelType::Audio)
|
||||
.collect();
|
||||
assert_eq!(audio.len(), 1, "only the commentary carries a label");
|
||||
assert_eq!(audio[0].stream_number, 4, "commentary is STN slot 4");
|
||||
assert_eq!(audio[0].purpose, LabelPurpose::Commentary);
|
||||
assert_eq!(audio[0].language, "eng");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assign_labels_prefers_fpl_over_seg_mainfeature() {
|
||||
// A `SEG_MainFeature` menu/preview segment carries a stray commentary
|
||||
// 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 tokens = strs(&[
|
||||
"SEG_MainFeature",
|
||||
"eng_ACOM_", // stray token in the menu segment — must be ignored
|
||||
"FPL_MainFeature",
|
||||
"Audio Stream 1",
|
||||
"eng_ACOM_",
|
||||
]);
|
||||
let labels = assign_labels(&tokens, &flag);
|
||||
let audio: Vec<_> = labels
|
||||
.iter()
|
||||
.filter(|l| l.stream_type == StreamLabelType::Audio)
|
||||
.collect();
|
||||
assert_eq!(audio.len(), 1);
|
||||
assert_eq!(audio[0].stream_number, 2, "numbered from the FPL playlist");
|
||||
assert_eq!(audio[0].purpose, LabelPurpose::Commentary);
|
||||
}
|
||||
|
||||
#[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 tokens = strs(&["SEG_MainFeature", "eng_MLP_", "spa_AC3_"]);
|
||||
let labels = assign_labels(&tokens, &flag);
|
||||
let audio: Vec<_> = labels
|
||||
.iter()
|
||||
.filter(|l| l.stream_type == StreamLabelType::Audio)
|
||||
.collect();
|
||||
assert_eq!(audio.len(), 2);
|
||||
assert_eq!(audio[0].stream_number, 1);
|
||||
assert_eq!(audio[0].language, "eng");
|
||||
assert_eq!(audio[1].stream_number, 2);
|
||||
assert_eq!(audio[1].language, "spa");
|
||||
}
|
||||
}
|
||||
|
||||
+1
-2
@@ -10,8 +10,7 @@
|
||||
//! - English text → [`LabelPurpose`] (Commentary / Descriptive / etc.).
|
||||
//! - English text → [`LabelQualifier`] (SDH / Forced / Descriptive Service).
|
||||
//!
|
||||
//! Rules of engagement (carried over from
|
||||
//! `(internal)/memory/feedback_label_data_rules.md`):
|
||||
//! Rules of engagement:
|
||||
//!
|
||||
//! 1. Only map values we are 100% certain about — published codec
|
||||
//! names, well-known ISO 639-2 mappings, vendor-documented purpose
|
||||
|
||||
+10
-5
@@ -77,6 +77,7 @@ pub mod css;
|
||||
pub mod decrypt;
|
||||
pub mod disc;
|
||||
pub mod drive;
|
||||
pub mod drm;
|
||||
pub mod error;
|
||||
pub mod event;
|
||||
pub mod halt;
|
||||
@@ -84,6 +85,7 @@ pub(crate) mod identity;
|
||||
pub(crate) mod ifo;
|
||||
pub mod io;
|
||||
pub mod keydb;
|
||||
pub mod keysource;
|
||||
pub mod labels;
|
||||
pub(crate) mod mpls;
|
||||
pub mod mux;
|
||||
@@ -154,7 +156,7 @@ pub use profile::DriveProfile;
|
||||
// don't touch `DecryptKeys` directly — `DiscStream::new(reader, title, keys, …)`
|
||||
// accepts whatever `Disc::decrypt_keys()` returned. `decrypt_sectors()` is
|
||||
// for callers that operate on raw sector buffers (e.g. ISO patching).
|
||||
pub use decrypt::{DecryptKeys, decrypt_sectors};
|
||||
pub use decrypt::{DecryptKeys, decrypt_sectors, decrypt_threads, set_decrypt_threads};
|
||||
|
||||
// ─── Disc structure ─────────────────────────────────────────────────────────
|
||||
//
|
||||
@@ -169,10 +171,11 @@ pub use decrypt::{DecryptKeys, decrypt_sectors};
|
||||
// prefix at the crate root to keep both addressable.
|
||||
pub use disc::{
|
||||
AacsState, AudioChannels, AudioStream, Clip, Codec, ColorSpace, ContentFormat, DamageSeverity,
|
||||
Disc, DiscFormat, DiscId, DiscTitle, Extent, FrameRate, HdrFormat, KeySource, LabelPurpose,
|
||||
LabelQualifier, PatchOptions, PatchOutcome, Resolution, SampleRate, ScanOptions, Stream,
|
||||
SubtitleStream, SweepOptions, VideoStream, classify_damage,
|
||||
Disc, DiscFormat, DiscId, DiscTitle, DriveCredentials, Extent, FrameRate, HdrFormat, Key,
|
||||
KeyOrigin, LabelPurpose, LabelQualifier, PatchOptions, PatchOutcome, Resolution, SampleRate,
|
||||
ScanOptions, Stream, SubtitleStream, SweepOptions, VideoStream, classify_damage,
|
||||
};
|
||||
pub use keysource::{DiscInputs, KeySource};
|
||||
|
||||
// ─── Streams ────────────────────────────────────────────────────────────────
|
||||
//
|
||||
@@ -211,9 +214,11 @@ pub use mux::{InputOptions, StreamUrl, input, output, parse_url};
|
||||
// are the ISO-on-disk implementations. [`DecryptingSectorSource`] is the
|
||||
// single decrypt-on-read decorator (AACS / CSS / none) — wrap any
|
||||
// `SectorSource` to get plaintext sectors out.
|
||||
pub use mux::build_iso_pipeline;
|
||||
pub use scsi::{DriveInfo, ScsiSense, ScsiTransport, drive_has_disc, list_drives};
|
||||
pub use sector::{
|
||||
DecryptingSectorSource, FileSectorSink, FileSectorSource, SectorSink, SectorSource,
|
||||
DecryptingSectorSource, FileSectorSink, FileSectorSource, PrefetchedSectorSource, SectorSink,
|
||||
SectorSource,
|
||||
};
|
||||
pub use speed::DriveSpeed;
|
||||
pub use udf::{UdfFs, read_filesystem};
|
||||
|
||||
+15
-3
@@ -301,9 +301,21 @@ fn parse_stream_entry(item: &[u8], pos: usize, stream_type: u8) -> Option<(Strea
|
||||
return None;
|
||||
}
|
||||
|
||||
// PID from stream entry (type 0x01 = PlayItem stream: PID at bytes 2-3)
|
||||
let pid = if item[pos + 1] == 0x01 && pos + 4 <= item.len() {
|
||||
u16::from_be_bytes([item[pos + 2], item[pos + 3]])
|
||||
// PID location depends on the stream-entry type (BD spec stream_entry()):
|
||||
// type 1 (stream in the PlayItem's Clip): PID at +2
|
||||
// type 2 (stream in a SubPath SubClip): +subpath_id(1)+subclip_id(1) → PID at +4
|
||||
// type 3 / 4 (SubPath clip; type 4 = Dolby Vision +subpath_id(1) → PID at +3
|
||||
// enhancement layer, e.g. PID 0x1015):
|
||||
// Previously only type 1 was handled, so the DV EL (type 4) and any
|
||||
// sub-path stream fell through to PID 0 and were dropped by the mux.
|
||||
let pid_off = match item[pos + 1] {
|
||||
0x01 => 2,
|
||||
0x02 => 4,
|
||||
0x03 | 0x04 => 3,
|
||||
_ => 0,
|
||||
};
|
||||
let pid = if pid_off != 0 && pos + pid_off + 2 <= item.len() {
|
||||
u16::from_be_bytes([item[pos + pid_off], item[pos + pid_off + 1]])
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
+120
-7
@@ -6,6 +6,14 @@
|
||||
|
||||
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
|
||||
|
||||
/// Hard cap on the carry-over buffer. An AC-3/E-AC-3 frame is at most 8192
|
||||
/// bytes (the `frame_size > 8192` reject below), so a single straddling frame
|
||||
/// plus a little slack never needs more than this. If the buffer grows past
|
||||
/// the cap without yielding a frame (pathological / never-syncing input) we
|
||||
/// drop it and resync rather than accumulate one PES worth of data per call
|
||||
/// for the whole title.
|
||||
const MAX_AC3_BUF: usize = 64 * 1024;
|
||||
|
||||
pub struct Ac3Parser {
|
||||
/// Leftover bytes from previous PES (incomplete frame at end).
|
||||
buf: Vec<u8>,
|
||||
@@ -76,23 +84,43 @@ impl CodecParser for Ac3Parser {
|
||||
pts_ns,
|
||||
keyframe: true,
|
||||
data: data[start..start + frame_size].to_vec(),
|
||||
duration_ns: None,
|
||||
});
|
||||
pos = start + frame_size;
|
||||
}
|
||||
|
||||
// Keep unconsumed data for next call
|
||||
// `pos` points to the start of unconsumed data (either a partial sync or leftover)
|
||||
// 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.
|
||||
let keep_from = if pos < data.len() {
|
||||
// Find the last sync word position in the unconsumed region
|
||||
find_ac3_sync(&data[pos..])
|
||||
.map(|o| pos + o)
|
||||
.unwrap_or(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
|
||||
// whole tail — including a lone trailing 0x0B that may be the first
|
||||
// half of a syncword split across the PES boundary.
|
||||
match find_ac3_sync(&data[pos..]) {
|
||||
Some(o) => pos + o,
|
||||
None if data.last() == Some(&0x0B) => data.len() - 1,
|
||||
None => data.len(),
|
||||
}
|
||||
} else {
|
||||
data.len()
|
||||
};
|
||||
|
||||
if keep_from < data.len() {
|
||||
self.buf = data[keep_from..].to_vec();
|
||||
let tail = &data[keep_from..];
|
||||
if tail.len() > MAX_AC3_BUF {
|
||||
// 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.
|
||||
self.buf.clear();
|
||||
} else {
|
||||
self.buf = tail.to_vec();
|
||||
}
|
||||
} else {
|
||||
self.buf.clear();
|
||||
}
|
||||
@@ -275,6 +303,91 @@ mod tests {
|
||||
assert_eq!(frames[0].data.len(), 160);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_word_split_across_pes_is_preserved() {
|
||||
// A frame whose 0x0B77 syncword straddles the PES boundary (0x0B at the
|
||||
// tail of PES 1, 0x77 at the head of PES 2) must still be emitted whole.
|
||||
// Previously the lone trailing 0x0B was dropped and the frame lost.
|
||||
let mut parser = Ac3Parser::new();
|
||||
let frame_data = make_ac3_frame(0, 2); // 160 bytes, starts with 0x0B 0x77
|
||||
|
||||
// PES 1: a complete frame, then a single 0x0B (first half of next sync).
|
||||
let mut pes1_data = frame_data.clone();
|
||||
pes1_data.push(0x0B);
|
||||
let pes1 = PesPacket {
|
||||
pid: 0,
|
||||
pts: Some(90000),
|
||||
dts: None,
|
||||
data: pes1_data,
|
||||
};
|
||||
let frames1 = parser.parse(&pes1);
|
||||
assert_eq!(frames1.len(), 1, "first complete frame emitted");
|
||||
|
||||
// PES 2: 0x77 (second half of sync) + rest of the second frame.
|
||||
let mut pes2_data = vec![0x77];
|
||||
pes2_data.extend_from_slice(&frame_data[2..]);
|
||||
let pes2 = PesPacket {
|
||||
pid: 0,
|
||||
pts: Some(93000),
|
||||
dts: None,
|
||||
data: pes2_data,
|
||||
};
|
||||
let frames2 = parser.parse(&pes2);
|
||||
assert_eq!(frames2.len(), 1, "split-sync frame must be recovered");
|
||||
assert_eq!(frames2[0].data.len(), 160);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn buffer_stays_bounded_across_many_garbage_pes() {
|
||||
// Finding 14: 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
|
||||
// the MAX_AC3_BUF resync guard.
|
||||
let mut parser = Ac3Parser::new();
|
||||
for i in 0..256 {
|
||||
// Vary the trailing byte so we also exercise the lone-0x0B retain.
|
||||
let mut data = vec![0x55u8; 8192];
|
||||
if i % 3 == 0 {
|
||||
*data.last_mut().unwrap() = 0x0B;
|
||||
}
|
||||
let pes = PesPacket {
|
||||
pid: 0,
|
||||
pts: None,
|
||||
dts: None,
|
||||
data,
|
||||
};
|
||||
let frames = parser.parse(&pes);
|
||||
assert!(frames.is_empty());
|
||||
assert!(
|
||||
parser.buf.len() <= MAX_AC3_BUF,
|
||||
"buffer grew to {} (cap {})",
|
||||
parser.buf.len(),
|
||||
MAX_AC3_BUF
|
||||
);
|
||||
}
|
||||
// After all that garbage the retained tail is at most a single partial
|
||||
// syncword byte — never an accumulation of whole PES packets.
|
||||
assert!(parser.buf.len() <= 1, "retained {} bytes", parser.buf.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_sync_below_cap_is_still_retained() {
|
||||
// The cap must not break the normal split-sync straddle: a short tail
|
||||
// ending in 0x0B (well under the cap) is retained so the next PES can
|
||||
// complete the syncword.
|
||||
let mut parser = Ac3Parser::new();
|
||||
let data = vec![0x00, 0x00, 0x0B];
|
||||
let pes = PesPacket {
|
||||
pid: 0,
|
||||
pts: None,
|
||||
dts: None,
|
||||
data,
|
||||
};
|
||||
assert!(parser.parse(&pes).is_empty());
|
||||
assert_eq!(parser.buf, vec![0x0B], "lone trailing 0x0B retained");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ac3_frame_size_table() {
|
||||
// fscod=0 (48kHz), frmsizecod=0: 64 words = 128 bytes
|
||||
|
||||
+410
-77
@@ -8,10 +8,20 @@
|
||||
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
|
||||
|
||||
const DTS_CORE_SYNC: [u8; 4] = [0x7F, 0xFE, 0x80, 0x01];
|
||||
/// DTS-HD extension substream syncword. The parser delimits an access unit by
|
||||
/// the next CORE sync (so every extension between two cores is captured), and
|
||||
/// never needs to locate or size the extension itself — so this is referenced
|
||||
/// only by the tests that synthesize extension substreams.
|
||||
#[cfg(test)]
|
||||
const DTS_HD_EXT_SYNC: [u8; 4] = [0x64, 0x58, 0x20, 0x25];
|
||||
|
||||
pub struct DtsParser {
|
||||
buf: Vec<u8>,
|
||||
/// PTS of the access unit currently being assembled in `buf` (the unit
|
||||
/// starting at the first buffered core sync). Captured when that core
|
||||
/// frame's PES first arrived; the trailing extension-substream PES
|
||||
/// packets carry their own (later) PTS which must NOT override it.
|
||||
pending_pts: i64,
|
||||
}
|
||||
|
||||
impl Default for DtsParser {
|
||||
@@ -24,10 +34,29 @@ impl DtsParser {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
buf: Vec::with_capacity(32768),
|
||||
pending_pts: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Hard cap on a buffered access unit (core + all its extension substreams).
|
||||
/// A DTS-HD MA frame is at most a few tens of KB; if the buffer grows past
|
||||
/// 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;
|
||||
|
||||
/// 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
|
||||
/// value can never collide with a genuine timestamp. Used to mark the PTS base
|
||||
/// invalid after a forced flush so the next PES sets it regardless of buffer
|
||||
/// state.
|
||||
const PTS_UNSET: i64 = -1;
|
||||
|
||||
impl CodecParser for DtsParser {
|
||||
fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> {
|
||||
if pes.data.is_empty() {
|
||||
@@ -35,79 +64,162 @@ impl CodecParser for DtsParser {
|
||||
}
|
||||
let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0);
|
||||
|
||||
// On Blu-ray, a DTS-HD MA/HRA access unit is a DTS core frame
|
||||
// (sync 0x7FFE8001) followed by one or more DTS extension substreams
|
||||
// (sync 0x64582025). The m2ts demuxer hands those out as SEPARATE PES
|
||||
// packets on the same PID — the core in one PES, then the extension
|
||||
// substreams in following PES packets (with their own, later PTS). The
|
||||
// lossless audio lives entirely in the extension substreams, so an
|
||||
// access unit is only complete once all of its trailing extensions
|
||||
// have been buffered. We assemble across PES boundaries here: an access
|
||||
// unit runs from its core sync up to (but not including) the NEXT core
|
||||
// sync. Emitting on the core boundary keeps the core + every following
|
||||
// extension substream together (the lossless data), instead of the
|
||||
// old per-PES emit that dropped the extension PES packets and
|
||||
// downgraded the track to lossy DTS core (the Dunkirk / Fight Club
|
||||
// bug). The PTS is the core frame's PTS, captured when the unit began.
|
||||
// Capture the access unit's PTS base on a fresh buffer, or whenever a
|
||||
// prior forced (safety-valve) flush left it invalidated — in the
|
||||
// forced case the bytes still in `buf` are not a real core frame, so
|
||||
// the first PES to arrive after the flush carries the correct base.
|
||||
if self.buf.is_empty() || self.pending_pts == PTS_UNSET {
|
||||
self.pending_pts = pts_ns;
|
||||
}
|
||||
self.buf.extend_from_slice(&pes.data);
|
||||
|
||||
let data = &self.buf;
|
||||
let mut frames = Vec::new();
|
||||
let mut pos = 0;
|
||||
|
||||
while pos < data.len() {
|
||||
// Find DTS core sync
|
||||
let start = match find_sync(&data[pos..], &DTS_CORE_SYNC) {
|
||||
Some(offset) => pos + offset,
|
||||
None => break,
|
||||
loop {
|
||||
// Resync to the first core sync; drop any leading junk.
|
||||
let Some(start) = find_sync(&self.buf, &DTS_CORE_SYNC) else {
|
||||
// No core sync at all yet — keep at most a 3-byte tail so a
|
||||
// sync split across PES packets can still be found next time.
|
||||
if self.buf.len() > 3 {
|
||||
let tail = self.buf.len() - 3;
|
||||
self.buf.drain(..tail);
|
||||
}
|
||||
break;
|
||||
};
|
||||
|
||||
// Need at least 10 bytes for core header to get frame size
|
||||
if start + 10 > data.len() {
|
||||
break;
|
||||
}
|
||||
|
||||
let core_size = dts_core_frame_size(&data[start..]);
|
||||
if core_size == 0 || core_size > 32768 {
|
||||
pos = start + 4;
|
||||
continue;
|
||||
}
|
||||
|
||||
if start + core_size > data.len() {
|
||||
// Incomplete core frame
|
||||
break;
|
||||
}
|
||||
|
||||
// Check for DTS-HD extension after core
|
||||
let mut total_size = core_size;
|
||||
if start + core_size + 4 <= data.len() {
|
||||
if let Some(0) = find_sync(
|
||||
&data[start + core_size..start + core_size + 4],
|
||||
&DTS_HD_EXT_SYNC,
|
||||
) {
|
||||
let ext = &data[start + core_size..];
|
||||
if ext.len() >= 9 {
|
||||
let ext_size = dts_hd_ext_frame_size(ext);
|
||||
if start + core_size + ext_size <= data.len() {
|
||||
total_size = core_size + ext_size;
|
||||
}
|
||||
// If ext incomplete, just emit core
|
||||
}
|
||||
if start > 0 {
|
||||
self.buf.drain(..start);
|
||||
if find_sync(&self.buf, &DTS_CORE_SYNC) != Some(0) {
|
||||
// Shouldn't happen, but never loop forever.
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Need the core header to size the core frame.
|
||||
if self.buf.len() < 10 {
|
||||
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.
|
||||
if !(MIN_CORE_FRAME_BYTES..=MAX_AU_BYTES).contains(&core_size) {
|
||||
// Bogus core sync — skip past it and resync.
|
||||
self.buf.drain(..4);
|
||||
continue;
|
||||
}
|
||||
if self.buf.len() < core_size {
|
||||
break; // core frame not fully buffered yet — wait
|
||||
}
|
||||
|
||||
// The access unit ends at the next *valid* core sync. The search
|
||||
// begins after this core's syncword so we don't re-match it.
|
||||
// Anything between the core and that next sync is this unit's
|
||||
// extension substream(s) — which can themselves contain byte
|
||||
// sequences matching the core syncword, so a raw `find_sync` match
|
||||
// is not enough: a candidate is only a real boundary if its decoded
|
||||
// core size is plausible. `next_core_boundary` skips bogus matches.
|
||||
//
|
||||
// `forced` distinguishes a real next-core boundary from a forced
|
||||
// safety-valve flush. On a forced flush the access unit was NOT
|
||||
// closed by a new core sync, so the bytes following it are not a
|
||||
// fresh core frame and the current PES's PTS (which on a forced
|
||||
// flush is an extension-substream PES, carrying its own later
|
||||
// timestamp) must NOT become the next unit's PTS base.
|
||||
let mut forced = false;
|
||||
let au_end = match next_core_boundary(&self.buf, core_size) {
|
||||
NextCore::Found(end) => end,
|
||||
NextCore::NeedMore => break, // candidate sync needs more header
|
||||
NextCore::None => {
|
||||
// No next core sync buffered yet. The trailing extension
|
||||
// substream PES packets may still be arriving, so WAIT for
|
||||
// them rather than emit a core-only (lossy) frame — unless
|
||||
// the buffer has grown unreasonably large, in which case
|
||||
// emit what we have to guarantee forward progress.
|
||||
if self.buf.len() <= MAX_AU_BYTES {
|
||||
break;
|
||||
}
|
||||
forced = true;
|
||||
self.buf.len()
|
||||
}
|
||||
};
|
||||
|
||||
let au: Vec<u8> = self.buf[..au_end].to_vec();
|
||||
frames.push(Frame {
|
||||
pts_ns,
|
||||
pts_ns: self.pending_pts,
|
||||
keyframe: true,
|
||||
data: data[start..start + total_size].to_vec(),
|
||||
data: au,
|
||||
duration_ns: None,
|
||||
});
|
||||
pos = start + total_size;
|
||||
}
|
||||
|
||||
// Keep unconsumed data
|
||||
let keep_from = if pos < data.len() {
|
||||
find_sync(&data[pos..], &DTS_CORE_SYNC)
|
||||
.map(|o| pos + o)
|
||||
.unwrap_or(data.len())
|
||||
} else {
|
||||
data.len()
|
||||
};
|
||||
|
||||
if keep_from < data.len() {
|
||||
self.buf = data[keep_from..].to_vec();
|
||||
} else {
|
||||
self.buf.clear();
|
||||
self.buf.drain(..au_end);
|
||||
if forced {
|
||||
// Safety-valve flush: the next access unit's real core PES has
|
||||
// not arrived. Invalidate the PTS so the next PES sets it
|
||||
// regardless of buffer state, rather than inheriting this
|
||||
// (non-core) PES's timestamp.
|
||||
self.pending_pts = PTS_UNSET;
|
||||
} else {
|
||||
// Real boundary: the next access unit (now at buf start) begins
|
||||
// at a core sync that arrived inside this PES, so this PES's PTS
|
||||
// is the correct base for it.
|
||||
self.pending_pts = pts_ns;
|
||||
}
|
||||
}
|
||||
|
||||
frames
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> Vec<Frame> {
|
||||
// End of stream: emit the final access unit still buffered (the last
|
||||
// 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 {
|
||||
self.buf.clear();
|
||||
return Vec::new();
|
||||
}
|
||||
let core_size = dts_core_frame_size(&self.buf);
|
||||
// `dts_core_frame_size` returns a 14-bit `fsize + 1` (never 0), so the
|
||||
// old `== 0` check was dead; reject a sub-minimum core like `parse()`.
|
||||
if core_size < MIN_CORE_FRAME_BYTES || self.buf.len() < core_size {
|
||||
self.buf.clear();
|
||||
return Vec::new();
|
||||
}
|
||||
let au = std::mem::take(&mut self.buf);
|
||||
// A non-empty buffer here means a PES arrived after any prior forced
|
||||
// flush (which fully drains `buf`), so `pending_pts` was reset to that
|
||||
// PES's real PTS. Clamp the sentinel to 0 defensively all the same.
|
||||
let pts_ns = if self.pending_pts == PTS_UNSET {
|
||||
0
|
||||
} else {
|
||||
self.pending_pts
|
||||
};
|
||||
vec![Frame {
|
||||
pts_ns,
|
||||
keyframe: true,
|
||||
data: au,
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
fn codec_private(&self) -> Option<Vec<u8>> {
|
||||
None
|
||||
}
|
||||
@@ -120,6 +232,40 @@ fn find_sync(data: &[u8], pattern: &[u8; 4]) -> Option<usize> {
|
||||
(0..=data.len() - 4).find(|&i| data[i..i + 4] == *pattern)
|
||||
}
|
||||
|
||||
/// Result of scanning for the next valid core sync that closes an access unit.
|
||||
enum NextCore {
|
||||
/// A valid next core sync was found; the access unit ends at this offset.
|
||||
Found(usize),
|
||||
/// A candidate core sync was found but its header isn't fully buffered yet,
|
||||
/// so its validity can't be decided — wait for more data.
|
||||
NeedMore,
|
||||
/// No (further) core sync found in the buffer.
|
||||
None,
|
||||
}
|
||||
|
||||
/// Find the next *valid* core sync after the current core frame, to delimit the
|
||||
/// access unit. Extension-substream payload can contain byte sequences that
|
||||
/// match the core syncword, so each candidate is validated by decoding its
|
||||
/// core size: a match whose decoded size is implausible (< MIN_CORE_FRAME_BYTES
|
||||
/// or > MAX_AU_BYTES) is a false sync and is skipped, continuing the search.
|
||||
fn next_core_boundary(buf: &[u8], core_size: usize) -> NextCore {
|
||||
let mut from = core_size;
|
||||
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 {
|
||||
return NextCore::NeedMore;
|
||||
}
|
||||
let sz = dts_core_frame_size(&buf[pos..]);
|
||||
if (MIN_CORE_FRAME_BYTES..=MAX_AU_BYTES).contains(&sz) {
|
||||
return NextCore::Found(pos);
|
||||
}
|
||||
// False sync inside extension payload — skip it and keep searching.
|
||||
from = pos + 4;
|
||||
}
|
||||
NextCore::None
|
||||
}
|
||||
|
||||
/// DTS core frame size from header bits.
|
||||
/// fsize is at bits 46-59 (14 bits) of the header: bytes 5-7.
|
||||
fn dts_core_frame_size(data: &[u8]) -> usize {
|
||||
@@ -133,20 +279,6 @@ fn dts_core_frame_size(data: &[u8]) -> usize {
|
||||
fsize + 1
|
||||
}
|
||||
|
||||
/// DTS-HD extension frame size from extension header.
|
||||
pub fn dts_hd_ext_frame_size(ext: &[u8]) -> usize {
|
||||
if ext.len() < 9 {
|
||||
return 0;
|
||||
}
|
||||
let raw =
|
||||
((ext[6] as usize & 0x1F) << 11) | ((ext[7] as usize) << 3) | ((ext[8] as usize) >> 5);
|
||||
raw + 1
|
||||
}
|
||||
|
||||
pub fn find_dts_hd_ext_sync(data: &[u8]) -> Option<usize> {
|
||||
find_sync(data, &DTS_HD_EXT_SYNC)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -180,12 +312,16 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn parse_single_frame() {
|
||||
// A single core frame with no following core sync is the LAST access
|
||||
// unit — held during streaming (can't know an extension won't follow),
|
||||
// then drained on flush() at EOF.
|
||||
let mut parser = DtsParser::new();
|
||||
let frame = make_dts_core(512);
|
||||
let pes = make_pes(frame, Some(90000));
|
||||
let frames = parser.parse(&pes);
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(frames[0].data.len(), 512);
|
||||
assert!(parser.parse(&pes).is_empty());
|
||||
let tail = parser.flush();
|
||||
assert_eq!(tail.len(), 1);
|
||||
assert_eq!(tail[0].data.len(), 512);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -198,9 +334,206 @@ mod tests {
|
||||
assert!(parser.parse(&pes1).is_empty());
|
||||
|
||||
let pes2 = make_pes(frame[mid..].to_vec(), Some(93000));
|
||||
let frames = parser.parse(&pes2);
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(frames[0].data.len(), 512);
|
||||
assert!(parser.parse(&pes2).is_empty());
|
||||
let tail = parser.flush();
|
||||
assert_eq!(tail.len(), 1);
|
||||
assert_eq!(tail[0].data.len(), 512);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn two_cores_back_to_back_emit_first_on_boundary() {
|
||||
// The first complete unit is emitted as soon as the next core sync is
|
||||
// seen; the second is held until flush.
|
||||
let mut parser = DtsParser::new();
|
||||
let mut stream = make_dts_core(512);
|
||||
stream.extend_from_slice(&make_dts_core(640));
|
||||
let f = parser.parse(&make_pes(stream, Some(90000)));
|
||||
assert_eq!(f.len(), 1);
|
||||
assert_eq!(f[0].data.len(), 512);
|
||||
let tail = parser.flush();
|
||||
assert_eq!(tail.len(), 1);
|
||||
assert_eq!(tail[0].data.len(), 640);
|
||||
}
|
||||
|
||||
/// Build a minimal DTS-HD extension substream of `size` bytes (just the
|
||||
/// sync + zero-padding). The parser delimits extensions by the next CORE
|
||||
/// sync, not by the extension's own size header, so a valid header isn't
|
||||
/// required — only that the bytes carry no spurious core sync.
|
||||
fn make_dts_ext(size: usize) -> Vec<u8> {
|
||||
let mut e = vec![0u8; size];
|
||||
e[0..4].copy_from_slice(&DTS_HD_EXT_SYNC);
|
||||
e
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_dts_hd_extension_in_separate_pes_packets() {
|
||||
// The real Blu-ray layout (ground-truthed on Dunkirk): the DTS core
|
||||
// arrives in one PES, then its DTS-HD MA extension substreams arrive
|
||||
// in SEPARATE following PES packets on the same PID. The parser must
|
||||
// stitch core + all trailing extensions into one access unit — not
|
||||
// emit a core-only (lossy) frame and drop the extension PES packets
|
||||
// (the Dunkirk / Fight Club lossy-core bug).
|
||||
let mut parser = DtsParser::new();
|
||||
|
||||
// Frame 1: core (512) + two extension substreams (256 + 200).
|
||||
assert!(
|
||||
parser
|
||||
.parse(&make_pes(make_dts_core(512), Some(90000)))
|
||||
.is_empty(),
|
||||
"core alone: must wait for any following extension"
|
||||
);
|
||||
assert!(
|
||||
parser
|
||||
.parse(&make_pes(make_dts_ext(256), Some(91000)))
|
||||
.is_empty(),
|
||||
"first extension PES: still waiting for the unit to close"
|
||||
);
|
||||
assert!(
|
||||
parser
|
||||
.parse(&make_pes(make_dts_ext(200), Some(91500)))
|
||||
.is_empty(),
|
||||
"second extension PES: unit still not closed (no next core yet)"
|
||||
);
|
||||
|
||||
// Frame 2's core PES arrives — that closes frame 1. The emitted unit
|
||||
// must be core + BOTH extensions (lossless preserved), and keep the
|
||||
// core's PTS, not the extension PES timestamps.
|
||||
let f = parser.parse(&make_pes(make_dts_core(512), Some(93000)));
|
||||
assert_eq!(f.len(), 1);
|
||||
assert_eq!(
|
||||
f[0].data.len(),
|
||||
512 + 256 + 200,
|
||||
"frame must include core + every extension substream"
|
||||
);
|
||||
|
||||
// EOF drains frame 2.
|
||||
let tail = parser.flush();
|
||||
assert_eq!(tail.len(), 1);
|
||||
assert_eq!(tail[0].data.len(), 512);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extension_split_across_pes_is_preserved() {
|
||||
// An extension substream straddling a PES boundary must still be fully
|
||||
// attached to its core.
|
||||
let mut parser = DtsParser::new();
|
||||
let ext = make_dts_ext(300);
|
||||
assert!(
|
||||
parser
|
||||
.parse(&make_pes(make_dts_core(512), Some(90000)))
|
||||
.is_empty()
|
||||
);
|
||||
assert!(
|
||||
parser
|
||||
.parse(&make_pes(ext[..150].to_vec(), Some(91000)))
|
||||
.is_empty()
|
||||
);
|
||||
assert!(
|
||||
parser
|
||||
.parse(&make_pes(ext[150..].to_vec(), Some(91000)))
|
||||
.is_empty()
|
||||
);
|
||||
let tail = parser.flush();
|
||||
assert_eq!(tail.len(), 1);
|
||||
assert_eq!(tail[0].data.len(), 512 + 300);
|
||||
}
|
||||
|
||||
/// Build 4 bytes that look like a DTS core sync but whose `fsize` field
|
||||
/// decodes to a tiny `core_size` (< MIN_CORE_FRAME_BYTES). With the
|
||||
/// dead-code guards this passed validation and could close an access unit
|
||||
/// at a junk boundary; with the fix it must be drained and resynced past.
|
||||
fn bogus_tiny_core_sync() -> Vec<u8> {
|
||||
// Core sync + zero header bytes. fsize = 0 → core_size = 1 (< 10).
|
||||
let mut v = vec![0u8; 10];
|
||||
v[0..4].copy_from_slice(&DTS_CORE_SYNC);
|
||||
// bytes 5,6,7 left zero → fsize = 0 → dts_core_frame_size = 1.
|
||||
assert_eq!(dts_core_frame_size(&v), 1);
|
||||
v
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bogus_tiny_core_sync_does_not_split_or_drop_real_au() {
|
||||
// A real core frame followed by an extension substream that happens to
|
||||
// contain a false core sync whose fsize decodes tiny. The bogus sync
|
||||
// must NOT close the real access unit early (dropping the rest of the
|
||||
// extension) nor emit a junk few-byte frame — it must be skipped, and
|
||||
// the whole core + extension preserved as one access unit.
|
||||
let mut parser = DtsParser::new();
|
||||
|
||||
// Frame 1: core(512) + an extension whose body embeds a bogus tiny
|
||||
// core sync midway through.
|
||||
let mut ext = make_dts_ext(256);
|
||||
// Embed the bogus core sync inside the extension body (offset 64).
|
||||
let bogus = bogus_tiny_core_sync();
|
||||
ext[64..64 + bogus.len()].copy_from_slice(&bogus);
|
||||
|
||||
let mut frame1 = make_dts_core(512);
|
||||
frame1.extend_from_slice(&ext);
|
||||
|
||||
// No next REAL core yet → frame 1 held.
|
||||
assert!(
|
||||
parser.parse(&make_pes(frame1, Some(90000))).is_empty(),
|
||||
"bogus tiny core sync must not close the AU; wait for a real core"
|
||||
);
|
||||
|
||||
// Frame 2's real core arrives — closes frame 1.
|
||||
let f = parser.parse(&make_pes(make_dts_core(640), Some(93000)));
|
||||
assert_eq!(f.len(), 1, "exactly one real access unit emitted");
|
||||
assert_eq!(
|
||||
f[0].data.len(),
|
||||
512 + 256,
|
||||
"AU must be the full core + extension, not split at the bogus sync"
|
||||
);
|
||||
assert_eq!(f[0].pts_ns, pts_to_ns(90000), "AU keeps the core's PTS");
|
||||
|
||||
let tail = parser.flush();
|
||||
assert_eq!(tail.len(), 1);
|
||||
assert_eq!(tail[0].data.len(), 640);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forced_emit_does_not_corrupt_next_au_pts() {
|
||||
// When the buffer exceeds MAX_AU_BYTES with no next core sync, the
|
||||
// parser force-emits for forward progress. The current PES at that
|
||||
// point is an extension-substream PES (later PTS). The forced path must
|
||||
// NOT make that extension PTS the base of the NEXT access unit.
|
||||
let mut parser = DtsParser::new();
|
||||
|
||||
// Core PES at the real PTS, then a giant extension (no next core) that
|
||||
// pushes the buffer past MAX_AU_BYTES, forcing an emit.
|
||||
let core_pts = 90000i64;
|
||||
assert!(
|
||||
parser
|
||||
.parse(&make_pes(make_dts_core(512), Some(core_pts)))
|
||||
.is_empty()
|
||||
);
|
||||
let ext_pts = 120000i64; // later extension-PES timestamp
|
||||
let big_ext = make_dts_ext(MAX_AU_BYTES + 1024);
|
||||
let f = parser.parse(&make_pes(big_ext, Some(ext_pts)));
|
||||
assert_eq!(f.len(), 1, "oversized buffer force-emits one AU");
|
||||
assert_eq!(
|
||||
f[0].pts_ns,
|
||||
pts_to_ns(core_pts),
|
||||
"forced AU keeps the core PTS"
|
||||
);
|
||||
|
||||
// The next REAL core PES arrives with its own PTS. Its AU must inherit
|
||||
// THIS core's PTS, not the prior extension PES timestamp.
|
||||
let next_core_pts = 150000i64;
|
||||
assert!(
|
||||
parser
|
||||
.parse(&make_pes(make_dts_core(512), Some(next_core_pts)))
|
||||
.is_empty()
|
||||
);
|
||||
let next_next_pts = 180000i64;
|
||||
let f2 = parser.parse(&make_pes(make_dts_core(512), Some(next_next_pts)));
|
||||
assert_eq!(f2.len(), 1);
|
||||
assert_eq!(
|
||||
f2[0].pts_ns,
|
||||
pts_to_ns(next_core_pts),
|
||||
"AU after a forced emit must use the next core's PTS, not the \
|
||||
stale extension PTS"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -30,6 +30,7 @@ impl CodecParser for DvdSubParser {
|
||||
pts_ns,
|
||||
keyframe: true,
|
||||
data: pes.data.clone(),
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
|
||||
+18
-12
@@ -38,8 +38,11 @@ impl CodecParser for H264Parser {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// Use DTS when available (monotonic for B-frame content), fall back to PTS
|
||||
let pts_ns = pes.dts.or(pes.pts).map(pts_to_ns).unwrap_or(0);
|
||||
// MKV block timecodes are PRESENTATION timestamps; frames are stored in
|
||||
// decode order and the player reorders by timecode. Use PTS, not DTS —
|
||||
// DTS presents B-frames in decode order (visible judder) and breaks
|
||||
// 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);
|
||||
|
||||
// Scan NAL units for SPS, PPS, and IDR detection
|
||||
let mut keyframe = false;
|
||||
@@ -85,6 +88,7 @@ impl CodecParser for H264Parser {
|
||||
pts_ns,
|
||||
keyframe,
|
||||
data: frame_data,
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
@@ -176,14 +180,16 @@ impl<'a> Iterator for NalIterator<'a> {
|
||||
}
|
||||
|
||||
/// 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;
|
||||
}
|
||||
// Range excludes last 2 bytes since we read 3 bytes at each position.
|
||||
// data.len()-2 as exclusive upper bound means last checked index is data.len()-3,
|
||||
// which accesses data[len-3], data[len-2], data[len-1] — all valid.
|
||||
(from..data.len() - 2).find(|&i| data[i] == 0x00 && data[i + 1] == 0x00 && data[i + 2] == 0x01)
|
||||
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.
|
||||
@@ -448,10 +454,10 @@ mod tests {
|
||||
assert!(frames.is_empty());
|
||||
}
|
||||
|
||||
// --- DTS preferred over PTS when present ---
|
||||
// --- PTS (presentation) used for the MKV block timecode, not DTS ---
|
||||
|
||||
#[test]
|
||||
fn dts_preferred_over_pts() {
|
||||
fn pts_preferred_over_dts() {
|
||||
let mut parser = H264Parser::new();
|
||||
|
||||
let mut data = Vec::new();
|
||||
@@ -461,13 +467,13 @@ mod tests {
|
||||
|
||||
let pes = PesPacket {
|
||||
pid: 0x1011,
|
||||
pts: Some(180000), // 2 seconds
|
||||
dts: Some(90000), // 1 second
|
||||
pts: Some(180000), // 2 seconds (presentation)
|
||||
dts: Some(90000), // 1 second (decode)
|
||||
data,
|
||||
};
|
||||
let frames = parser.parse(&pes);
|
||||
assert_eq!(frames.len(), 1);
|
||||
// DTS should be used, not PTS
|
||||
assert_eq!(frames[0].pts_ns, 1_000_000_000);
|
||||
// PTS must be used — MKV block timecodes are presentation timestamps.
|
||||
assert_eq!(frames[0].pts_ns, 2_000_000_000);
|
||||
}
|
||||
}
|
||||
|
||||
+171
-26
@@ -48,17 +48,30 @@ impl CodecParser for HevcParser {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// Use DTS when available (monotonic for B-frame content), fall back to PTS
|
||||
let pts_ns = pes.dts.or(pes.pts).map(pts_to_ns).unwrap_or(0);
|
||||
// MKV block timecodes are PRESENTATION timestamps; frames are stored
|
||||
// in decode order (the order they arrive here) and the player reorders
|
||||
// for display by timecode. So use PTS, not DTS — using DTS makes the
|
||||
// block timecode monotonic in storage order, which presents B-frames in
|
||||
// decode order (visible judder / wrong frames) and breaks PTS-based
|
||||
// seeking. Fall back to DTS only if PTS is somehow absent.
|
||||
let pts_ns = pes.pts.or(pes.dts).map(pts_to_ns).unwrap_or(0);
|
||||
let data = &pes.data;
|
||||
let mut keyframe = false;
|
||||
let mut frame_data = Vec::new();
|
||||
// Pre-size: output is ~input bytes with a few 4-byte length
|
||||
// prefixes added. UHD frames are 150-300 KB; the unsized Vec
|
||||
// growth chain otherwise reallocs 5-7× per frame.
|
||||
let mut frame_data = Vec::with_capacity(data.len() + 64);
|
||||
|
||||
// Single-pass NAL scan: extract params, detect keyframes, build length-prefixed output
|
||||
let mut pos = 0;
|
||||
while let Some(sc_pos) = find_start_code(data, pos) {
|
||||
if let Some(nal_start) = skip_start_code(data, sc_pos) {
|
||||
let next = find_start_code(data, nal_start).unwrap_or(data.len());
|
||||
// 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.
|
||||
let mut end = next;
|
||||
while end > nal_start && data[end - 1] == 0x00 {
|
||||
end -= 1;
|
||||
@@ -107,6 +120,7 @@ impl CodecParser for HevcParser {
|
||||
pts_ns,
|
||||
keyframe,
|
||||
data: frame_data,
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
@@ -120,34 +134,41 @@ impl CodecParser for HevcParser {
|
||||
// Full HEVCDecoderConfigurationRecord is complex — for now, concatenate
|
||||
let mut record = Vec::new();
|
||||
|
||||
// Minimal HEVCDecoderConfigurationRecord header
|
||||
// Minimal HEVCDecoderConfigurationRecord header.
|
||||
//
|
||||
// The stored SPS NAL is [2-byte HEVC NAL header][SPS RBSP...].
|
||||
// The RBSP begins at sps[2]; profile_tier_level() begins one byte
|
||||
// later, after sps_video_parameter_set_id u(4) +
|
||||
// sps_max_sub_layers_minus1 u(3) + sps_temporal_id_nesting_flag u(1)
|
||||
// (= sps[2], a full byte). So the profile_tier_level fields are:
|
||||
// sps[3] general_profile_space u(2)+tier u(1)+profile_idc u(5)
|
||||
// sps[4..8] general_profile_compatibility_flags u(32)
|
||||
// sps[8..14] general_constraint_indicator_flags 48 bits
|
||||
// sps[14] general_level_idc u(8)
|
||||
// (Byte-aligned read; emulation-prevention bytes within the first
|
||||
// 15 SPS bytes are not handled — extremely rare and matches the
|
||||
// pre-existing simplification.)
|
||||
record.push(1); // configurationVersion
|
||||
// General profile space, tier flag, profile IDC from SPS
|
||||
if sps.len() > 3 {
|
||||
record.push(sps[1]); // general_profile_space + general_tier_flag + general_profile_idc
|
||||
// general_profile_space + general_tier_flag + general_profile_idc
|
||||
record.push(if sps.len() > 3 { sps[3] } else { 0 });
|
||||
// general_profile_compatibility_flags (4 bytes) — SPS bytes 4..8
|
||||
if sps.len() > 7 {
|
||||
record.extend_from_slice(&sps[4..8]);
|
||||
} else {
|
||||
record.push(0);
|
||||
let avail = sps.len().saturating_sub(4).min(4);
|
||||
record.extend_from_slice(&sps[sps.len().min(4)..sps.len().min(8)]);
|
||||
record.extend_from_slice(&vec![0u8; 4 - avail]);
|
||||
}
|
||||
// general_profile_compatibility_flags (4 bytes) — from SPS bytes 2..6
|
||||
if sps.len() > 5 {
|
||||
record.extend_from_slice(&sps[2..6]);
|
||||
// general_constraint_indicator_flags (6 bytes) — SPS bytes 8..14
|
||||
if sps.len() > 13 {
|
||||
record.extend_from_slice(&sps[8..14]);
|
||||
} else {
|
||||
record.extend_from_slice(&[0, 0, 0, 0]);
|
||||
let avail = sps.len().saturating_sub(8).min(6);
|
||||
record.extend_from_slice(&sps[sps.len().min(8)..sps.len().min(14)]);
|
||||
record.extend_from_slice(&vec![0u8; 6 - avail]);
|
||||
}
|
||||
// general_constraint_indicator_flags (6 bytes) — from SPS bytes 6..12
|
||||
if sps.len() > 11 {
|
||||
record.extend_from_slice(&sps[6..12]);
|
||||
} else {
|
||||
let avail = sps.len().saturating_sub(6).min(6);
|
||||
if avail > 0 {
|
||||
record.extend_from_slice(&sps[6..6 + avail]);
|
||||
record.extend_from_slice(&vec![0u8; 6 - avail]);
|
||||
} else {
|
||||
record.extend_from_slice(&[0, 0, 0, 0, 0, 0]);
|
||||
}
|
||||
}
|
||||
// general_level_idc
|
||||
record.push(if sps.len() > 12 { sps[12] } else { 0 });
|
||||
// general_level_idc — SPS byte 14
|
||||
record.push(if sps.len() > 14 { sps[14] } else { 0 });
|
||||
// min_spatial_segmentation_idc (4 + 12 bits)
|
||||
record.extend_from_slice(&[0xF0, 0x00]);
|
||||
// parallelismType (6 + 2 bits)
|
||||
@@ -264,6 +285,103 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hvcc_profile_tier_level_offsets() {
|
||||
// The hvcC fixed header must read profile_tier_level from the SPS
|
||||
// RBSP, not from the NAL header. Stored SPS = [2-byte NAL header][RBSP].
|
||||
// RBSP layout (byte-aligned):
|
||||
// sps[2] sps_vps_id/max_sub_layers/temporal_nesting
|
||||
// sps[3] general_profile_space+tier+profile_idc
|
||||
// sps[4..8] general_profile_compatibility_flags
|
||||
// sps[8..14] general_constraint_indicator_flags
|
||||
// sps[14] general_level_idc
|
||||
let mut parser = HevcParser::new();
|
||||
|
||||
// Distinct, recognizable values for each field.
|
||||
let sps_rbsp: [u8; 13] = [
|
||||
0xAB, // sps[2] (vps_id etc.) — must NOT leak into profile fields
|
||||
0x21, // sps[3] profile byte: space=0, tier=0, profile_idc=1
|
||||
0x60, 0x00, 0x00, 0x00, // sps[4..8] compat flags
|
||||
0x90, 0x00, 0x00, 0x00, 0x00, 0x00, // sps[8..14] constraint flags
|
||||
0x7B, // sps[14] level_idc = 123
|
||||
];
|
||||
|
||||
let mut data = Vec::new();
|
||||
// VPS
|
||||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||||
data.extend_from_slice(&hevc_nal_header(32));
|
||||
data.extend_from_slice(&[0xAA, 0xBB, 0xCC]);
|
||||
// SPS — 2-byte header + the structured RBSP above
|
||||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||||
data.extend_from_slice(&hevc_nal_header(33));
|
||||
data.extend_from_slice(&sps_rbsp);
|
||||
// PPS
|
||||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||||
data.extend_from_slice(&hevc_nal_header(34));
|
||||
data.extend_from_slice(&[0xDD, 0xEE]);
|
||||
|
||||
let pes = make_pes(data, Some(0));
|
||||
parser.parse(&pes);
|
||||
|
||||
let cp = parser
|
||||
.codec_private()
|
||||
.expect("codec_private should be Some");
|
||||
|
||||
// record[0] = configurationVersion
|
||||
assert_eq!(cp[0], 1, "configurationVersion");
|
||||
// record[1] = general_profile_space+tier+profile_idc <- sps[3]
|
||||
assert_eq!(
|
||||
cp[1], 0x21,
|
||||
"profile byte must come from SPS RBSP, not NAL hdr"
|
||||
);
|
||||
// record[2..6] = general_profile_compatibility_flags <- sps[4..8]
|
||||
assert_eq!(&cp[2..6], &[0x60, 0x00, 0x00, 0x00], "compatibility flags");
|
||||
// record[6..12] = general_constraint_indicator_flags <- sps[8..14]
|
||||
assert_eq!(
|
||||
&cp[6..12],
|
||||
&[0x90, 0x00, 0x00, 0x00, 0x00, 0x00],
|
||||
"constraint flags"
|
||||
);
|
||||
// record[12] = general_level_idc <- sps[14]
|
||||
assert_eq!(cp[12], 0x7B, "level_idc must come from sps[14]");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hvcc_short_sps_does_not_panic() {
|
||||
// A truncated SPS must still produce a fixed header without panicking
|
||||
// and zero-pad the missing profile/level bytes.
|
||||
let mut parser = HevcParser::new();
|
||||
|
||||
let mut data = Vec::new();
|
||||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||||
data.extend_from_slice(&hevc_nal_header(32));
|
||||
data.extend_from_slice(&[0xAA]);
|
||||
// SPS with only 3 RBSP bytes (stored len = 5): forces every guard path
|
||||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||||
data.extend_from_slice(&hevc_nal_header(33));
|
||||
data.extend_from_slice(&[0x11, 0x22, 0x33]);
|
||||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||||
data.extend_from_slice(&hevc_nal_header(34));
|
||||
data.extend_from_slice(&[0xDD]);
|
||||
|
||||
let pes = make_pes(data, Some(0));
|
||||
parser.parse(&pes);
|
||||
|
||||
let cp = parser
|
||||
.codec_private()
|
||||
.expect("codec_private should be Some");
|
||||
// sps stored = [hdr0, hdr1, 0x11, 0x22, 0x33], len 5.
|
||||
// profile byte = sps[3] = 0x22; everything past sps[4]=0x33 is absent.
|
||||
assert_eq!(cp[0], 1);
|
||||
assert_eq!(cp[1], 0x22, "profile byte = sps[3]");
|
||||
// compat flags: only sps[4]=0x33 present, rest zero-padded.
|
||||
assert_eq!(&cp[2..6], &[0x33, 0x00, 0x00, 0x00]);
|
||||
// constraint flags: none present, all zero.
|
||||
assert_eq!(&cp[6..12], &[0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
|
||||
// level_idc: absent, zero.
|
||||
assert_eq!(cp[12], 0x00);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codec_private_none_before_params() {
|
||||
let parser = HevcParser::new();
|
||||
@@ -470,6 +588,33 @@ mod tests {
|
||||
assert_eq!(frames[0].pts_ns, 1_000_000_000);
|
||||
}
|
||||
|
||||
// --- PTS (presentation), not DTS, drives the MKV block timecode ---
|
||||
// Regression for B-frame presentation: writing DTS as the block timecode
|
||||
// presents frames in decode order (visible judder) and breaks seeking.
|
||||
|
||||
#[test]
|
||||
fn pts_preferred_over_dts() {
|
||||
let mut parser = HevcParser::new();
|
||||
|
||||
let mut data = Vec::new();
|
||||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||||
data.extend_from_slice(&hevc_nal_header(1)); // TRAIL_R slice
|
||||
data.extend_from_slice(&[0x10, 0x20]);
|
||||
|
||||
let pes = PesPacket {
|
||||
pid: 0x1011,
|
||||
pts: Some(180000), // 2 s (presentation)
|
||||
dts: Some(90000), // 1 s (decode)
|
||||
data,
|
||||
};
|
||||
let frames = parser.parse(&pes);
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(
|
||||
frames[0].pts_ns, 2_000_000_000,
|
||||
"block timecode must be PTS"
|
||||
);
|
||||
}
|
||||
|
||||
// --- Dolby Vision enhancement layer ---
|
||||
|
||||
#[test]
|
||||
|
||||
+89
-10
@@ -1,24 +1,39 @@
|
||||
//! BD/DVD LPCM (Linear PCM) audio parser.
|
||||
//!
|
||||
//! BD LPCM PES packets have a 4-byte header:
|
||||
//! BD LPCM PES packets (TS stream type 0x80) carry a 4-byte header on the
|
||||
//! elementary-stream payload:
|
||||
//! Bytes 0-1: audio frame number
|
||||
//! Byte 2: reserved
|
||||
//! Byte 3: quantization (bits 7-6), sample rate (bits 5-4), channel assignment (bits 3-0)
|
||||
//! This header is part of the ES payload, so the BD parser must strip it.
|
||||
//!
|
||||
//! DVD LPCM (private stream 1, sub-stream 0xA0-0xA7) has a 3-byte header.
|
||||
//! DVD LPCM lives in private stream 1 (sub-stream 0xA0-0xA7). Its 7-byte
|
||||
//! private sub-header (sub_id + frames + first-access-unit-ptr(2) + emphasis +
|
||||
//! quant/freq + channels) is stripped by `PsDemuxer` while demuxing the
|
||||
//! Program Stream. By the time a DVD LPCM `PesPacket` reaches this parser its
|
||||
//! `data` is already raw PCM, so the parser must NOT strip any further bytes —
|
||||
//! doing so drops one sample pair per PES and drifts the audio.
|
||||
//!
|
||||
//! The raw PCM data follows the header. No framing is needed — each PES
|
||||
//! payload minus its header is one complete audio frame.
|
||||
//! The two origins are distinguished by the `strip_header` flag: BD = strip,
|
||||
//! 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).
|
||||
//! All frames are keyframes (uncompressed audio).
|
||||
|
||||
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
|
||||
|
||||
/// BD LPCM header size in bytes.
|
||||
/// BD LPCM header size in bytes (present on BD-TS LPCM, absent on DVD-PS LPCM
|
||||
/// because `PsDemuxer` already stripped the private sub-header).
|
||||
const BD_LPCM_HEADER_SIZE: usize = 4;
|
||||
|
||||
pub struct LpcmParser;
|
||||
pub struct LpcmParser {
|
||||
/// Whether to strip the 4-byte BD LPCM header from each PES payload.
|
||||
///
|
||||
/// `true` for BD-TS LPCM (header still present), `false` for DVD-PS LPCM
|
||||
/// (header already removed by `PsDemuxer`).
|
||||
strip_header: bool,
|
||||
}
|
||||
|
||||
impl Default for LpcmParser {
|
||||
fn default() -> Self {
|
||||
@@ -27,23 +42,37 @@ impl Default for LpcmParser {
|
||||
}
|
||||
|
||||
impl LpcmParser {
|
||||
/// BD-TS LPCM parser: strips the 4-byte BD LPCM header from each PES.
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
Self { strip_header: true }
|
||||
}
|
||||
|
||||
/// DVD-PS LPCM parser: `PsDemuxer` already stripped the private sub-header,
|
||||
/// so the payload is raw PCM and no further bytes are removed.
|
||||
pub fn new_dvd() -> Self {
|
||||
Self {
|
||||
strip_header: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CodecParser for LpcmParser {
|
||||
fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> {
|
||||
// Skip the BD LPCM header (4 bytes).
|
||||
let offset = if self.strip_header {
|
||||
BD_LPCM_HEADER_SIZE
|
||||
} else {
|
||||
0
|
||||
};
|
||||
// If the PES is too short to contain header + data, return nothing.
|
||||
if pes.data.len() <= BD_LPCM_HEADER_SIZE {
|
||||
if pes.data.len() <= offset {
|
||||
return Vec::new();
|
||||
}
|
||||
let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0);
|
||||
vec![Frame {
|
||||
pts_ns,
|
||||
keyframe: true,
|
||||
data: pes.data[BD_LPCM_HEADER_SIZE..].to_vec(),
|
||||
data: pes.data[offset..].to_vec(),
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
@@ -83,6 +112,56 @@ mod tests {
|
||||
assert_eq!(frames[0].pts_ns, 1_000_000_000); // 90000 ticks = 1 second
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bd_lpcm_strips_4_byte_header() {
|
||||
// BD-TS LPCM: the 4-byte BD header is part of the ES payload and must
|
||||
// be stripped, leaving exactly the PCM bytes.
|
||||
let mut parser = LpcmParser::new();
|
||||
let header = vec![0x00, 0x01, 0x00, 0b1001_0001];
|
||||
let pcm = vec![0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88];
|
||||
let mut data = header;
|
||||
data.extend_from_slice(&pcm);
|
||||
|
||||
let frames = parser.parse(&make_pes(data, Some(0)));
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(frames[0].data, pcm, "BD must strip exactly 4 header bytes");
|
||||
}
|
||||
|
||||
#[test]
|
||||
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).
|
||||
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)));
|
||||
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(
|
||||
frames[0].data, pcm,
|
||||
"DVD must preserve every PCM byte — no second strip"
|
||||
);
|
||||
assert_eq!(frames[0].pts_ns, 1_000_000_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dvd_lpcm_emits_short_payload_bd_would_drop() {
|
||||
// A 4-byte raw-PCM DVD payload (1 sample pair at 16-bit stereo). The BD
|
||||
// parser drops <= 4 bytes as "header only"; the DVD parser must emit it.
|
||||
let mut bd = LpcmParser::new();
|
||||
let mut dvd = LpcmParser::new_dvd();
|
||||
let pcm = vec![0xAA, 0xBB, 0xCC, 0xDD];
|
||||
|
||||
assert!(
|
||||
bd.parse(&make_pes(pcm.clone(), Some(0))).is_empty(),
|
||||
"BD treats 4 bytes as header-only"
|
||||
);
|
||||
let frames = dvd.parse(&make_pes(pcm.clone(), Some(0)));
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(frames[0].data, pcm);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn always_keyframe() {
|
||||
let mut parser = LpcmParser::new();
|
||||
|
||||
+31
-1
@@ -29,6 +29,13 @@ pub struct Frame {
|
||||
pub keyframe: bool,
|
||||
/// Frame data (elementary stream bytes).
|
||||
pub data: Vec<u8>,
|
||||
/// Optional duration in nanoseconds — only set by parsers that
|
||||
/// can compute one (currently PGS, which pairs a display PCS
|
||||
/// with the following empty PCS). When `Some`, the MKV muxer
|
||||
/// emits a `BlockGroup` with `BlockDuration` instead of a
|
||||
/// `SimpleBlock`; without it players guess the display interval
|
||||
/// (subtitles linger past their end-time).
|
||||
pub duration_ns: Option<u64>,
|
||||
}
|
||||
|
||||
/// Convert 90kHz PTS to nanoseconds (round to nearest).
|
||||
@@ -45,6 +52,18 @@ pub trait CodecParser: Send {
|
||||
/// Some (TrueHD): multiple access units per PES.
|
||||
fn parse(&mut self, pes: &PesPacket) -> Vec<Frame>;
|
||||
|
||||
/// Drain any access unit still buffered after the last PES.
|
||||
///
|
||||
/// Parsers that buffer across PES boundaries to assemble a complete
|
||||
/// access unit (e.g. DTS-HD, whose extension substreams arrive in
|
||||
/// separate PES packets) hold the final unit until they can prove it's
|
||||
/// complete. At end-of-stream there is no following packet to prove it,
|
||||
/// so the demuxer calls `flush()` once after the last PES to emit it.
|
||||
/// Default: nothing buffered, no tail.
|
||||
fn flush(&mut self) -> Vec<Frame> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
/// Get codec initialization data (e.g., SPS+PPS for H.264).
|
||||
/// Returns None until enough data has been seen.
|
||||
fn codec_private(&self) -> Option<Vec<u8>>;
|
||||
@@ -71,6 +90,7 @@ impl CodecParser for PassthroughParser {
|
||||
pts_ns,
|
||||
keyframe: self.keyframe,
|
||||
data: pes.data.clone(),
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
@@ -82,7 +102,16 @@ impl CodecParser for PassthroughParser {
|
||||
/// Create the appropriate parser for a codec, with optional codec private data.
|
||||
///
|
||||
/// For DvdSub, `codec_data` should be the pre-formatted VobSub .idx palette header.
|
||||
pub fn parser_for_codec(codec: Codec, codec_data: Option<Vec<u8>>) -> Box<dyn CodecParser> {
|
||||
///
|
||||
/// `is_dvd_ps` selects the DVD program-stream variant where it matters: DVD
|
||||
/// LPCM arrives with its private sub-header already stripped by the
|
||||
/// `PsDemuxer`, so the LPCM parser must NOT strip the 4-byte BD LPCM header
|
||||
/// again (that would drop one PCM sample pair per PES → progressive drift).
|
||||
pub fn parser_for_codec(
|
||||
codec: Codec,
|
||||
codec_data: Option<Vec<u8>>,
|
||||
is_dvd_ps: bool,
|
||||
) -> Box<dyn CodecParser> {
|
||||
match codec {
|
||||
Codec::H264 => Box::new(h264::H264Parser::new()),
|
||||
Codec::Hevc => Box::new(hevc::HevcParser::new()),
|
||||
@@ -92,6 +121,7 @@ pub fn parser_for_codec(codec: Codec, codec_data: Option<Vec<u8>>) -> Box<dyn Co
|
||||
Codec::DtsHdMa | Codec::DtsHdHr | Codec::Dts => Box::new(dts::DtsParser::new()),
|
||||
Codec::TrueHd => Box::new(truehd::TrueHdParser::new()),
|
||||
Codec::Pgs => Box::new(pgs::PgsParser::new()),
|
||||
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)),
|
||||
|
||||
+80
-4
@@ -96,9 +96,14 @@ impl CodecParser for Mpeg2Parser {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let pts_ns = pes.dts.or(pes.pts).map(pts_to_ns).unwrap_or(0);
|
||||
// MKV block timecodes are PRESENTATION timestamps; frames are stored in
|
||||
// decode order and the player reorders by timecode. Use PTS, not DTS —
|
||||
// DTS presents B-frames in decode order (visible judder) and breaks
|
||||
// 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;
|
||||
let mut has_picture = false;
|
||||
|
||||
// Scan for start codes in the elementary stream data.
|
||||
let mut pos = 0;
|
||||
@@ -176,6 +181,7 @@ impl CodecParser for Mpeg2Parser {
|
||||
PICTURE_CODE => {
|
||||
// Picture header: bytes after start code contain temporal_reference
|
||||
// (10 bits) + picture_coding_type (3 bits).
|
||||
has_picture = true;
|
||||
if sc + 5 < data.len() {
|
||||
let picture_coding_type = (data[sc + 5] >> 3) & 0x07;
|
||||
if picture_coding_type == PICTURE_TYPE_I {
|
||||
@@ -190,10 +196,28 @@ impl CodecParser for Mpeg2Parser {
|
||||
}
|
||||
}
|
||||
|
||||
// A PES that carried a sequence header but no picture start code is a
|
||||
// parameter-set-only access unit: it has no coded picture to emit.
|
||||
// Emitting it as a standalone keyframe would put bare sequence-header
|
||||
// bytes into frame data with no picture. The sequence header is
|
||||
// captured into codec_private above and is re-emitted in-band on the
|
||||
// next real picture's PES, so dropping the empty access unit loses
|
||||
// nothing. Mirrors how the H.264/HEVC parsers skip parameter-set-only
|
||||
// access units.
|
||||
//
|
||||
// Conservative: only drop when this PES actually contained a sequence
|
||||
// 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) {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
vec![Frame {
|
||||
pts_ns,
|
||||
keyframe,
|
||||
data: pes.data.clone(),
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
@@ -240,6 +264,21 @@ 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 {
|
||||
@@ -428,15 +467,15 @@ mod tests {
|
||||
assert!(has_ext, "codec_private should include sequence extension");
|
||||
}
|
||||
|
||||
// --- I-frame with sequence header = keyframe ---
|
||||
// --- sequence header + picture = keyframe ---
|
||||
|
||||
#[test]
|
||||
fn sequence_header_implies_keyframe() {
|
||||
fn sequence_header_with_picture_is_keyframe() {
|
||||
let mut parser = Mpeg2Parser::new();
|
||||
|
||||
let mut data = Vec::new();
|
||||
data.extend_from_slice(&make_seq_header(720, 480, 3, 4));
|
||||
// Even without an explicit picture header, a sequence header implies I-frame.
|
||||
data.extend_from_slice(&make_picture_header(PICTURE_TYPE_I));
|
||||
data.extend_from_slice(&[0xFF; 16]);
|
||||
|
||||
let pes = make_pes(data, Some(0));
|
||||
@@ -444,6 +483,43 @@ mod tests {
|
||||
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert!(frames[0].keyframe);
|
||||
// codecPrivate is still captured.
|
||||
assert!(parser.codec_private().is_some());
|
||||
}
|
||||
|
||||
// --- parameter-set-only PES (seq header, no picture) emits no frame ---
|
||||
|
||||
#[test]
|
||||
fn sequence_header_only_pes_emits_no_frame() {
|
||||
let mut parser = Mpeg2Parser::new();
|
||||
|
||||
// A PES carrying only a sequence header (+ extension), no picture.
|
||||
let mut data = Vec::new();
|
||||
data.extend_from_slice(&make_seq_header(1920, 1080, 3, 4));
|
||||
data.extend_from_slice(&[0x00, 0x00, 0x01, SEQ_EXT_CODE]);
|
||||
data.extend_from_slice(&[0x14, 0x8A, 0x00, 0x01, 0x00, 0x00]);
|
||||
|
||||
let pes = make_pes(data, Some(0));
|
||||
let frames = parser.parse(&pes);
|
||||
|
||||
// No coded picture → no frame emitted, but the sequence header is
|
||||
// still captured for codecPrivate.
|
||||
assert!(
|
||||
frames.is_empty(),
|
||||
"parameter-set-only PES should not emit a frame"
|
||||
);
|
||||
assert!(
|
||||
parser.codec_private().is_some(),
|
||||
"sequence header should still be captured into codec_private"
|
||||
);
|
||||
|
||||
// A following picture-bearing PES emits the real keyframe.
|
||||
let mut data2 = Vec::new();
|
||||
data2.extend_from_slice(&make_picture_header(PICTURE_TYPE_I));
|
||||
data2.extend_from_slice(&[0xFF; 16]);
|
||||
let frames2 = parser.parse(&make_pes(data2, Some(3600)));
|
||||
assert_eq!(frames2.len(), 1);
|
||||
assert!(frames2[0].keyframe);
|
||||
}
|
||||
|
||||
// --- PTS conversion ---
|
||||
|
||||
+169
-28
@@ -1,12 +1,39 @@
|
||||
//! HDMV PGS (Presentation Graphics Stream) subtitle parser.
|
||||
//!
|
||||
//! PGS segments: PCS, WDS, PDS, ODS, END.
|
||||
//! Each PES packet contains one or more segments.
|
||||
//! All segments are keyframes (no inter-segment dependencies).
|
||||
//! PGS segments: PCS, WDS, PDS, ODS, END. Each PES packet starts with
|
||||
//! one of those (segment_type byte at offset 0).
|
||||
//!
|
||||
//! Subtitle display lifecycle (BD spec):
|
||||
//! - A "display" PCS (number_of_composition_objects > 0) starts a
|
||||
//! visible subtitle. Its WDS/PDS/ODS follow.
|
||||
//! - A later "empty" PCS (number_of_composition_objects == 0) clears
|
||||
//! the screen.
|
||||
//!
|
||||
//! For Matroska output we collapse that pair into one block with
|
||||
//! `BlockDuration` set to (clear_pts - display_pts). Without a
|
||||
//! duration, hardware players linger on the last bitmap until the
|
||||
//! next subtitle replaces it — which can be many seconds, and on a
|
||||
//! disc where the final subtitle has no follower, until end of file.
|
||||
|
||||
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
|
||||
|
||||
pub struct PgsParser;
|
||||
const SEGMENT_PCS: u8 = 0x16;
|
||||
// Upper bound on a pending display set's accumulated bytes. Real PGS
|
||||
// display sets are small (a 1080p RLE bitmap plus palette is well under
|
||||
// 1 MB); a stream that keeps appending non-PCS segments without ever
|
||||
// emitting a PCS is malformed. Cap accumulation to bound memory and
|
||||
// drop further appends until the next PCS resyncs the parser. Mirrors
|
||||
// the MAX_*_BYTES / MAX_*_BUF caps in the DTS and AC-3 parsers.
|
||||
const MAX_PGS_PENDING_BYTES: usize = 4 * 1024 * 1024;
|
||||
// Offset within the PES payload at which number_of_composition_objects
|
||||
// lives in a PCS: 3-byte segment header + 10 bytes of PCS fields
|
||||
// (video_w/h, frame_rate, comp_num, comp_state, palette_update,
|
||||
// palette_id_ref) = 13.
|
||||
const PCS_NUM_OBJECTS_OFFSET: usize = 13;
|
||||
|
||||
pub struct PgsParser {
|
||||
pending: Option<(i64, Vec<u8>)>,
|
||||
}
|
||||
|
||||
impl Default for PgsParser {
|
||||
fn default() -> Self {
|
||||
@@ -16,7 +43,7 @@ impl Default for PgsParser {
|
||||
|
||||
impl PgsParser {
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
Self { pending: None }
|
||||
}
|
||||
}
|
||||
|
||||
@@ -26,11 +53,68 @@ impl CodecParser for PgsParser {
|
||||
return Vec::new();
|
||||
}
|
||||
let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0);
|
||||
vec![Frame {
|
||||
pts_ns,
|
||||
keyframe: true,
|
||||
data: pes.data.clone(),
|
||||
}]
|
||||
|
||||
let is_pcs = pes.data[0] == SEGMENT_PCS;
|
||||
let pcs_num_objects = if is_pcs && pes.data.len() > PCS_NUM_OBJECTS_OFFSET {
|
||||
Some(pes.data[PCS_NUM_OBJECTS_OFFSET])
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let mut out = Vec::new();
|
||||
match pcs_num_objects {
|
||||
// Clear/empty PCS — closes any pending display. Drop the
|
||||
// clear segment itself; BlockDuration covers the screen
|
||||
// wipe.
|
||||
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 {
|
||||
pts_ns: start_pts,
|
||||
keyframe: true,
|
||||
data,
|
||||
duration_ns: Some(duration),
|
||||
});
|
||||
}
|
||||
}
|
||||
// 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 {
|
||||
pts_ns: start_pts,
|
||||
keyframe: true,
|
||||
data,
|
||||
duration_ns: Some(duration),
|
||||
});
|
||||
}
|
||||
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.
|
||||
None => {
|
||||
if let Some((_, ref mut buf)) = self.pending {
|
||||
// Bound accumulation: a well-formed display set is small.
|
||||
// Past the cap, drop further appends (malformed stream);
|
||||
// the next PCS will take/replace `pending` and resync.
|
||||
if buf.len() + pes.data.len() <= MAX_PGS_PENDING_BYTES {
|
||||
buf.extend_from_slice(&pes.data);
|
||||
}
|
||||
} else {
|
||||
out.push(Frame {
|
||||
pts_ns,
|
||||
keyframe: true,
|
||||
data: pes.data.clone(),
|
||||
duration_ns: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
fn codec_private(&self) -> Option<Vec<u8>> {
|
||||
@@ -52,29 +136,86 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_basic_segment() {
|
||||
let mut parser = PgsParser::new();
|
||||
// PGS segment data (PCS = presentation composition segment)
|
||||
let data = vec![0x16, 0x00, 0x00, 0x11, 0x01, 0x02, 0x03];
|
||||
let pes = make_pes(data.clone(), Some(90000));
|
||||
let frames = parser.parse(&pes);
|
||||
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(frames[0].data, data);
|
||||
assert_eq!(frames[0].pts_ns, 1_000_000_000);
|
||||
// Minimum-viable PCS bytes: type 0x16, segment_length (2 bytes),
|
||||
// then 11 bytes of PCS fields ending in number_of_composition_objects.
|
||||
fn pcs_bytes(num_objects: u8) -> Vec<u8> {
|
||||
let mut v = vec![SEGMENT_PCS, 0x00, 0x0B];
|
||||
v.extend_from_slice(&[0x07, 0x80, 0x04, 0x38]); // 1920x1080
|
||||
v.push(0x10); // frame_rate
|
||||
v.extend_from_slice(&[0x00, 0x01]); // composition_number
|
||||
v.push(0x80); // composition_state = EpochStart
|
||||
v.push(0x00); // palette_update + reserved
|
||||
v.push(0x00); // palette_id_ref
|
||||
v.push(num_objects);
|
||||
v
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_keyframes() {
|
||||
fn display_then_clear_yields_duration() {
|
||||
let mut parser = PgsParser::new();
|
||||
for i in 0..3 {
|
||||
let data = vec![0x16, 0x00, i];
|
||||
let pes = make_pes(data, Some(90000 * i as i64));
|
||||
let frames = parser.parse(&pes);
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert!(frames[0].keyframe, "PGS segment should always be keyframe");
|
||||
|
||||
// Display PCS at PTS 90000 (= 1s)
|
||||
let display = pcs_bytes(1);
|
||||
let frames = parser.parse(&make_pes(display.clone(), Some(90000)));
|
||||
assert!(frames.is_empty(), "display PCS should be pending");
|
||||
|
||||
// Empty PCS at PTS 270000 (= 3s)
|
||||
let clear = pcs_bytes(0);
|
||||
let frames = parser.parse(&make_pes(clear, Some(270000)));
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(frames[0].pts_ns, 1_000_000_000);
|
||||
assert_eq!(frames[0].duration_ns, Some(2_000_000_000));
|
||||
assert_eq!(frames[0].data, display);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replace_without_clear_still_emits_prior_with_duration() {
|
||||
let mut parser = PgsParser::new();
|
||||
let _ = parser.parse(&make_pes(pcs_bytes(1), Some(90000)));
|
||||
let frames = parser.parse(&make_pes(pcs_bytes(1), Some(180000)));
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(frames[0].pts_ns, 1_000_000_000);
|
||||
assert_eq!(frames[0].duration_ns, Some(1_000_000_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_pcs_segment_appends_to_pending() {
|
||||
let mut parser = PgsParser::new();
|
||||
let _ = parser.parse(&make_pes(pcs_bytes(1), Some(90000)));
|
||||
// ODS-like segment (type 0x15)
|
||||
let frames = parser.parse(&make_pes(vec![0x15, 0x00, 0x02, 0xAA, 0xBB], Some(90000)));
|
||||
assert!(frames.is_empty());
|
||||
// Clear closes the set; data should include the appended bytes.
|
||||
let frames = parser.parse(&make_pes(pcs_bytes(0), Some(180000)));
|
||||
assert_eq!(frames.len(), 1);
|
||||
let data = &frames[0].data;
|
||||
assert!(data.windows(5).any(|w| w == [0x15, 0x00, 0x02, 0xAA, 0xBB]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_buffer_is_capped() {
|
||||
let mut parser = PgsParser::new();
|
||||
// Open a display set.
|
||||
let _ = parser.parse(&make_pes(pcs_bytes(1), Some(90000)));
|
||||
|
||||
// Flood with non-PCS segments far exceeding the cap.
|
||||
let chunk = vec![0x15u8; 256 * 1024]; // 256 KB ODS-like segment
|
||||
let floods = (MAX_PGS_PENDING_BYTES / chunk.len()) + 32;
|
||||
for _ in 0..floods {
|
||||
let frames = parser.parse(&make_pes(chunk.clone(), Some(90000)));
|
||||
assert!(frames.is_empty(), "non-PCS appends should not emit");
|
||||
}
|
||||
|
||||
// The pending buffer must not have grown without bound.
|
||||
let pending_len = parser.pending.as_ref().map(|(_, b)| b.len()).unwrap_or(0);
|
||||
assert!(
|
||||
pending_len <= MAX_PGS_PENDING_BYTES,
|
||||
"pending buffer {pending_len} exceeded cap {MAX_PGS_PENDING_BYTES}"
|
||||
);
|
||||
|
||||
// A following PCS still resyncs and emits the (capped) pending set.
|
||||
let frames = parser.parse(&make_pes(pcs_bytes(0), Some(180000)));
|
||||
assert_eq!(frames.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -16,6 +16,12 @@ use super::{CodecParser, Frame, PesPacket, pts_to_ns};
|
||||
/// Duration of one TrueHD access unit in nanoseconds (1/1200 second).
|
||||
const AU_DURATION_NS: i64 = 833_333;
|
||||
|
||||
/// Hard cap on the reassembly buffer. A valid TrueHD/MAT access unit is
|
||||
/// well under 32 KiB; if the buffer grows far past that without yielding a
|
||||
/// frame the stream is malformed, so we drop it and resync rather than grow
|
||||
/// without bound. Parity with the AC-3 / DTS / PGS caps.
|
||||
const MAX_TRUEHD_BUF: usize = 256 * 1024;
|
||||
|
||||
pub struct TrueHdParser {
|
||||
buf: Vec<u8>,
|
||||
next_pts_ns: i64,
|
||||
@@ -139,11 +145,18 @@ impl CodecParser for TrueHdParser {
|
||||
pts_ns: self.next_pts_ns,
|
||||
keyframe: is_major_sync,
|
||||
data: self.buf[..unit_bytes].to_vec(),
|
||||
duration_ns: None,
|
||||
});
|
||||
self.buf.drain(..unit_bytes);
|
||||
self.next_pts_ns += AU_DURATION_NS;
|
||||
}
|
||||
|
||||
// Bound memory on malformed input: a stream that never yields a
|
||||
// complete frame must not grow the buffer without limit.
|
||||
if self.buf.len() > MAX_TRUEHD_BUF {
|
||||
self.buf.clear();
|
||||
}
|
||||
|
||||
frames
|
||||
}
|
||||
|
||||
@@ -152,6 +165,52 @@ impl CodecParser for TrueHdParser {
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-bit channel counts for the TrueHD 8-channel and 6-channel presentation
|
||||
/// channel-assignment masks (per the MLP spec / FFmpeg `thd_channels`). Some
|
||||
/// bits denote a stereo pair (2), others a single channel (1).
|
||||
const THD_8CH: [u8; 13] = [2, 1, 1, 2, 2, 2, 2, 1, 1, 2, 2, 1, 1];
|
||||
const THD_6CH: [u8; 5] = [2, 1, 1, 2, 1];
|
||||
|
||||
/// Decode the true channel count from a TrueHD major-sync `format_info` word
|
||||
/// (the 32 bits immediately after the 0xF8726FBA sync). Returns the richest
|
||||
/// presentation's channel count — the 8-channel (e.g. 7.1) presentation when
|
||||
/// present, else the 6-channel (5.1) one. This is the real layout that the MPLS
|
||||
/// `audio_format` base field (often 5.1 even on a 7.1/Atmos track) understates.
|
||||
pub fn truehd_channels(format_info: u32) -> Option<u8> {
|
||||
let ch8 = (format_info & 0x1FFF) as u16; // 8ch_presentation_channel_assignment (13 bits)
|
||||
let ch6 = ((format_info >> 15) & 0x1F) as u16; // 6ch_presentation_channel_assignment (5 bits)
|
||||
let count = |mask: u16, tbl: &[u8]| -> u8 {
|
||||
tbl.iter()
|
||||
.enumerate()
|
||||
.filter(|(i, _)| mask & (1 << i) != 0)
|
||||
.map(|(_, &c)| c)
|
||||
.sum()
|
||||
};
|
||||
if ch8 != 0 {
|
||||
Some(count(ch8, &THD_8CH))
|
||||
} else if ch6 != 0 {
|
||||
Some(count(ch6, &THD_6CH))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Scan a demuxed TrueHD elementary-stream chunk for the first major sync and
|
||||
/// decode its true channel count. The stream may interleave AC-3; we scan for
|
||||
/// the major-sync word anywhere and read the following `format_info`.
|
||||
pub fn truehd_channels_from_stream(data: &[u8]) -> Option<u8> {
|
||||
let mut p = 0;
|
||||
while p + 8 <= data.len() {
|
||||
let w = u32::from_be_bytes([data[p], data[p + 1], data[p + 2], data[p + 3]]);
|
||||
if (w & 0xFFFF_FFFE) == 0xF872_6FBA {
|
||||
let fi = u32::from_be_bytes([data[p + 4], data[p + 5], data[p + 6], data[p + 7]]);
|
||||
return truehd_channels(fi);
|
||||
}
|
||||
p += 1;
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -246,4 +305,27 @@ mod tests {
|
||||
let parser = TrueHdParser::new();
|
||||
assert!(parser.codec_private().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truehd_channels_71_from_8ch_presentation() {
|
||||
// 8ch presentation assignment bits 0-4 (LR,C,LFE,LsRs,back-LR) = 2+1+1+2+2 = 8.
|
||||
let format_info = 0x1F; // low 13 bits = 0x1F
|
||||
assert_eq!(truehd_channels(format_info), Some(8));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truehd_channels_51_from_6ch_presentation() {
|
||||
// No 8ch presentation; 6ch bits 0-3 (LR,C,LFE,LsRs) = 2+1+1+2 = 6.
|
||||
let format_info = 0xF << 15; // 6ch field = 0xF, 8ch field = 0
|
||||
assert_eq!(truehd_channels(format_info), Some(6));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truehd_channels_scan_finds_major_sync() {
|
||||
// [junk][major sync 0xF8726FBA][format_info: 8ch=0x1F -> 7.1]
|
||||
let mut data = vec![0xAA, 0xBB];
|
||||
data.extend_from_slice(&0xF872_6FBAu32.to_be_bytes());
|
||||
data.extend_from_slice(&0x0000_001Fu32.to_be_bytes());
|
||||
assert_eq!(truehd_channels_from_stream(&data), Some(8));
|
||||
}
|
||||
}
|
||||
|
||||
+65
-10
@@ -41,9 +41,13 @@ impl CodecParser for Vc1Parser {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// Use DTS when available (monotonic for B-frame content), fall back to PTS
|
||||
let ts_ns = pes.dts.or(pes.pts).map(pts_to_ns).unwrap_or(0);
|
||||
// MKV block timecodes are PRESENTATION timestamps; frames are stored in
|
||||
// decode order and the player reorders by timecode. Use PTS, not DTS —
|
||||
// DTS presents B-frames in decode order (visible judder) and breaks
|
||||
// PTS-based seeking. Fall back to DTS only if PTS is absent.
|
||||
let ts_ns = pes.pts.or(pes.dts).map(pts_to_ns).unwrap_or(0);
|
||||
let mut has_seq_header = false;
|
||||
let mut has_entry_point = false;
|
||||
let mut frame_start: Option<usize> = None;
|
||||
|
||||
// Scan for start codes (00 00 01 XX)
|
||||
@@ -67,6 +71,7 @@ impl CodecParser for Vc1Parser {
|
||||
SC_ENTRY_POINT => {
|
||||
let end = find_next_sc(data, i + 4).unwrap_or(data.len());
|
||||
self.entry_point = Some(data[i..end].to_vec());
|
||||
has_entry_point = true;
|
||||
}
|
||||
SC_FRAME => {
|
||||
// Frame data starts at this start code
|
||||
@@ -85,17 +90,30 @@ impl CodecParser for Vc1Parser {
|
||||
// Keyframe = this PES contains a sequence header (I-frame indicator in BD)
|
||||
let keyframe = has_seq_header;
|
||||
|
||||
// Strip sequence header + entry point from frame data — those are in codecPrivate.
|
||||
// Only include data from the frame start code onwards.
|
||||
// Strip sequence header + entry point from frame data — those are in
|
||||
// codecPrivate, not coded-picture data. Only include data from the
|
||||
// frame start code onwards.
|
||||
let frame_data = match frame_start {
|
||||
Some(start) => &data[start..],
|
||||
None => data, // no frame start code found, pass through entire PES
|
||||
None => {
|
||||
// No frame start code. If this PES carried only parameter sets
|
||||
// (sequence header / entry point, captured above into
|
||||
// codecPrivate), there is no coded picture to emit — drop it
|
||||
// rather than passing parameter bytes through as a bogus
|
||||
// keyframe. Mirrors how the H.264/HEVC parsers skip
|
||||
// parameter-set-only access units.
|
||||
if has_seq_header || has_entry_point {
|
||||
return Vec::new();
|
||||
}
|
||||
data // genuine picture payload with no leading 0x0D — pass through
|
||||
}
|
||||
};
|
||||
|
||||
vec![Frame {
|
||||
pts_ns: ts_ns,
|
||||
keyframe,
|
||||
data: frame_data.to_vec(),
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
@@ -342,6 +360,42 @@ mod tests {
|
||||
assert_eq!(&frames[0].data[0..4], &[0x00, 0x00, 0x01, SC_FRAME]);
|
||||
}
|
||||
|
||||
// --- parameter-set-only PES (seq header + entry point, no frame SC) ---
|
||||
|
||||
#[test]
|
||||
fn param_set_only_pes_emits_no_frame() {
|
||||
let mut parser = Vc1Parser::new();
|
||||
|
||||
// Sequence header + entry point, but NO frame start code (0x0D).
|
||||
let mut data = Vec::new();
|
||||
data.extend_from_slice(&[0x00, 0x00, 0x01, SC_SEQUENCE_HEADER]);
|
||||
data.extend_from_slice(&[0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF]);
|
||||
data.extend_from_slice(&[0x00, 0x00, 0x01, SC_ENTRY_POINT]);
|
||||
data.extend_from_slice(&[0x11, 0x22, 0x33, 0x44]);
|
||||
|
||||
let pes = make_pes(data, Some(90000));
|
||||
let frames = parser.parse(&pes);
|
||||
|
||||
// No coded picture → no frame emitted (parameter bytes must not be
|
||||
// passed through as a bogus keyframe).
|
||||
assert!(
|
||||
frames.is_empty(),
|
||||
"parameter-set-only PES should not emit a frame"
|
||||
);
|
||||
// But codecPrivate is still captured.
|
||||
assert!(parser.seq_header.is_some());
|
||||
assert!(parser.entry_point.is_some());
|
||||
assert!(parser.codec_private().is_some());
|
||||
|
||||
// A following frame-bearing PES still emits its picture.
|
||||
let mut data2 = Vec::new();
|
||||
data2.extend_from_slice(&[0x00, 0x00, 0x01, SC_FRAME]);
|
||||
data2.extend_from_slice(&[0x55, 0x66, 0x77]);
|
||||
let frames2 = parser.parse(&make_pes(data2, Some(180000)));
|
||||
assert_eq!(frames2.len(), 1);
|
||||
assert_eq!(&frames2[0].data[0..4], &[0x00, 0x00, 0x01, SC_FRAME]);
|
||||
}
|
||||
|
||||
// --- empty PES ---
|
||||
|
||||
#[test]
|
||||
@@ -368,10 +422,10 @@ mod tests {
|
||||
assert_eq!(frames[0].pts_ns, 1_000_000_000);
|
||||
}
|
||||
|
||||
// --- DTS preferred over PTS ---
|
||||
// --- PTS (presentation) used for the MKV block timecode, not DTS ---
|
||||
|
||||
#[test]
|
||||
fn dts_preferred_over_pts() {
|
||||
fn pts_preferred_over_dts() {
|
||||
let mut parser = Vc1Parser::new();
|
||||
|
||||
let mut data = Vec::new();
|
||||
@@ -380,13 +434,14 @@ mod tests {
|
||||
|
||||
let pes = PesPacket {
|
||||
pid: 0x1011,
|
||||
pts: Some(180000),
|
||||
dts: Some(90000),
|
||||
pts: Some(180000), // presentation
|
||||
dts: Some(90000), // decode
|
||||
data,
|
||||
};
|
||||
let frames = parser.parse(&pes);
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(frames[0].pts_ns, 1_000_000_000);
|
||||
// PTS must be used — MKV block timecodes are presentation timestamps.
|
||||
assert_eq!(frames[0].pts_ns, 2_000_000_000);
|
||||
}
|
||||
|
||||
// --- find_next_sc utility ---
|
||||
|
||||
@@ -0,0 +1,350 @@
|
||||
//! `DemuxThread` — runs the read+decrypt+demux pipeline on a
|
||||
//! dedicated thread, feeding completed `PesPacket` batches to the
|
||||
//! caller via a bounded channel.
|
||||
//!
|
||||
//! ## Why a second worker thread
|
||||
//!
|
||||
//! With [`crate::sector::PrefetchedSectorSource`] alone, read+decrypt
|
||||
//! already runs on a producer thread; the *consumer* (main) thread
|
||||
//! still serialises `ts_demuxer.feed` (M2TS parsing) with the codec
|
||||
//! parsers. Profiling on the rip1 testbed showed feed at ~37 % and
|
||||
//! codec parse at ~44 % of consumer wall time — i.e. feed is heavy
|
||||
//! enough that pipelining it with parse pays for itself.
|
||||
//!
|
||||
//! Splitting them: feed runs in [`DemuxThread`]; the consumer thread
|
||||
//! receives `Vec<PesPacket>` batches and runs codec parse + frame
|
||||
//! emission only. Total throughput becomes `1/max(feed, parse)`
|
||||
//! instead of `1/(feed + parse)`.
|
||||
//!
|
||||
//! ## 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.
|
||||
|
||||
use crate::halt::Halt;
|
||||
use crate::sector::SectorSource;
|
||||
use crossbeam_channel::{Receiver, Sender, bounded};
|
||||
use std::thread::JoinHandle;
|
||||
|
||||
/// Output channel depth. Two batches in flight keeps the consumer
|
||||
/// (codec parser) busy without piling up demuxed bytes if it stalls.
|
||||
const DEMUX_CHANNEL_DEPTH: usize = 2;
|
||||
|
||||
/// One demuxed batch flowing from the demux thread to the consumer.
|
||||
pub enum DemuxBatch {
|
||||
/// Successfully demuxed PesPackets — non-empty.
|
||||
Ts(Vec<super::ts::PesPacket>),
|
||||
Ps(Vec<super::ps::PsPacket>),
|
||||
/// Underlying reader returned an error. Terminal.
|
||||
Err(std::io::Error),
|
||||
}
|
||||
|
||||
/// Spawned demux thread. Drop joins.
|
||||
///
|
||||
/// In zero-copy mode the thread also owns an opaque
|
||||
/// `producer_shell: Option<Box<dyn Send>>` — the join handle of the
|
||||
/// upstream producer (sector or byte prefetcher). Dropping the
|
||||
/// `DemuxThread` runs the shell's `Drop`, which joins the producer.
|
||||
/// `Box<dyn Send>` rather than a concrete type so the same demux
|
||||
/// worker can be wired behind either prefetcher kind.
|
||||
pub struct DemuxThread {
|
||||
handle: Option<JoinHandle<()>>,
|
||||
#[allow(dead_code)]
|
||||
producer_shell: Option<Box<dyn Send>>,
|
||||
}
|
||||
|
||||
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
|
||||
/// 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
|
||||
/// returns them to `recycle_tx` for the producer to re-fill.
|
||||
/// Eliminates the 16 MiB memcpy per batch that the SectorSource
|
||||
/// adapter incurred (and, with the producer-side recycling pool,
|
||||
/// also eliminates the per-batch heap alloc / cross-thread free
|
||||
/// that was costing 40 %+ of demux-thread time before).
|
||||
///
|
||||
/// `producer_shell` is an opaque handle whose only purpose is to
|
||||
/// outlive the demux thread and join the upstream producer when
|
||||
/// dropped. Both
|
||||
/// [`crate::sector::PrefetchedSectorSource::into_channels`] and
|
||||
/// [`crate::io::byte_prefetcher::BytePrefetcher::into_channels`]
|
||||
/// hand back a shell that fits — pass either.
|
||||
pub fn spawn_zero_copy<S: Send + 'static>(
|
||||
prefetch_rx: Receiver<std::io::Result<Vec<u8>>>,
|
||||
recycle_tx: Sender<Vec<u8>>,
|
||||
producer_shell: S,
|
||||
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 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_bytes: u64 = 0;
|
||||
loop {
|
||||
if halt.as_ref().map(|h| h.is_cancelled()).unwrap_or(false) {
|
||||
return;
|
||||
}
|
||||
let t0 = if prof {
|
||||
Some(std::time::Instant::now())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let buf = match prefetch_rx.recv() {
|
||||
Ok(Ok(b)) => b,
|
||||
Ok(Err(e)) => {
|
||||
let _ = tx.send(DemuxBatch::Err(e));
|
||||
return;
|
||||
}
|
||||
Err(_) => break, // producer done → EOF
|
||||
};
|
||||
let t1 = if prof {
|
||||
Some(std::time::Instant::now())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let n = buf.len();
|
||||
if let Some(ref mut d) = ts {
|
||||
let pkts = d.feed(&buf);
|
||||
let t2 = if prof {
|
||||
Some(std::time::Instant::now())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
// Recycle the buffer back to the producer
|
||||
// before pushing the demuxed packets. If the
|
||||
// recycle channel is closed the producer has
|
||||
// exited; we drop the buffer and continue.
|
||||
let _ = recycle_tx.send(buf);
|
||||
if !pkts.is_empty() && tx.send(DemuxBatch::Ts(pkts)).is_err() {
|
||||
return;
|
||||
}
|
||||
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_bytes += n as u64;
|
||||
let now = std::time::Instant::now();
|
||||
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={}%",
|
||||
el,
|
||||
mbps,
|
||||
prof_read_ns / 10_000 / el,
|
||||
prof_feed_ns / 10_000 / el,
|
||||
);
|
||||
prof_last_dump = now;
|
||||
prof_started = now;
|
||||
prof_read_ns = 0;
|
||||
prof_feed_ns = 0;
|
||||
prof_bytes = 0;
|
||||
}
|
||||
}
|
||||
} else if let Some(ref mut d) = ps {
|
||||
let pkts = d.feed(&buf);
|
||||
let _ = recycle_tx.send(buf);
|
||||
if !pkts.is_empty() && tx.send(DemuxBatch::Ps(pkts)).is_err() {
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
let _ = recycle_tx.send(buf);
|
||||
}
|
||||
}
|
||||
// Flush tail packets at EOF.
|
||||
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));
|
||||
}
|
||||
}
|
||||
})
|
||||
.expect("freemkv-demux thread spawn failed");
|
||||
|
||||
(
|
||||
Self {
|
||||
handle: Some(handle),
|
||||
producer_shell: Some(Box::new(producer_shell)),
|
||||
},
|
||||
rx,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for DemuxThread {
|
||||
fn drop(&mut self) {
|
||||
if let Some(h) = self.handle.take() {
|
||||
let _ = h.join();
|
||||
}
|
||||
}
|
||||
}
|
||||
+127
-9
@@ -1,7 +1,7 @@
|
||||
//! DiscStream — read any disc (physical drive or ISO file) → PES frames.
|
||||
//!
|
||||
//! One stream type for all disc sources. The source is a SectorSource —
|
||||
//! Drive (hardware) or IsoSectorReader (file). DiscStream doesn't care.
|
||||
//! Drive (hardware) or FileSectorSource (file). DiscStream doesn't care.
|
||||
//!
|
||||
//! Read-only. For disc→ISO (raw sector copy), use `Disc::copy()`.
|
||||
|
||||
@@ -149,7 +149,11 @@ pub struct DiscStream {
|
||||
// a percent without a separate API call.
|
||||
bytes_total_extents: u64,
|
||||
|
||||
// PES output
|
||||
// PES output — single-threaded inline demux + codec parse. The
|
||||
// pipeline-mode mux (3-stage threaded) lives in
|
||||
// [`super::pipelined_stream::PipelinedPesStream`]; this type is
|
||||
// the legacy in-thread path for live-disc reads where adaptive
|
||||
// batch retry on bad sectors lives in `fill_extents`.
|
||||
ts_demuxer: Option<super::ts::TsDemuxer>,
|
||||
ps_demuxer: Option<super::ps::PsDemuxer>,
|
||||
parsers: Vec<(u16, Box<dyn super::codec::CodecParser>)>,
|
||||
@@ -191,7 +195,8 @@ impl DiscStream {
|
||||
};
|
||||
pids.push(pid);
|
||||
pid_to_track.push((pid, idx));
|
||||
parsers.push((pid, super::codec::parser_for_codec(codec, None)));
|
||||
let is_dvd_ps = matches!(content_format, crate::disc::ContentFormat::MpegPs);
|
||||
parsers.push((pid, super::codec::parser_for_codec(codec, None, is_dvd_ps)));
|
||||
}
|
||||
|
||||
let mut ts_demuxer = None;
|
||||
@@ -403,6 +408,71 @@ 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).
|
||||
struct StageProf {
|
||||
started: std::time::Instant,
|
||||
last_dump: std::time::Instant,
|
||||
fill_ns: u128,
|
||||
feed_ns: u128,
|
||||
consume_ns: u128,
|
||||
bytes_in: u64,
|
||||
}
|
||||
|
||||
const PROFILE_INTERVAL: std::time::Duration = std::time::Duration::from_secs(5);
|
||||
|
||||
thread_local! {
|
||||
static STAGE_PROF: std::cell::RefCell<Option<StageProf>> = const { std::cell::RefCell::new(None) };
|
||||
}
|
||||
|
||||
fn prof_active() -> bool {
|
||||
std::env::var_os("FREEMKV_PROFILE").is_some()
|
||||
}
|
||||
|
||||
fn prof_tick(stage: &str, ns: u128, bytes: u64) {
|
||||
STAGE_PROF.with(|cell| {
|
||||
let mut slot = cell.borrow_mut();
|
||||
if slot.is_none() {
|
||||
if !prof_active() {
|
||||
return;
|
||||
}
|
||||
let now = std::time::Instant::now();
|
||||
*slot = Some(StageProf {
|
||||
started: now,
|
||||
last_dump: now,
|
||||
fill_ns: 0,
|
||||
feed_ns: 0,
|
||||
consume_ns: 0,
|
||||
bytes_in: 0,
|
||||
});
|
||||
}
|
||||
let p = slot.as_mut().unwrap();
|
||||
match stage {
|
||||
"fill" => p.fill_ns += ns,
|
||||
"feed" => p.feed_ns += ns,
|
||||
"consume" => p.consume_ns += ns,
|
||||
_ => {}
|
||||
}
|
||||
p.bytes_in += bytes;
|
||||
let now = std::time::Instant::now();
|
||||
if now.duration_since(p.last_dump) < PROFILE_INTERVAL {
|
||||
return;
|
||||
}
|
||||
let elapsed_ms = now.duration_since(p.started).as_millis().max(1);
|
||||
let fill_pct = p.fill_ns / 10_000 / elapsed_ms;
|
||||
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!(
|
||||
"[profile] elapsed={}ms in={}MB/s fill={}% feed={}% consume={}%",
|
||||
elapsed_ms, mbps, fill_pct, feed_pct, consume_pct,
|
||||
);
|
||||
p.last_dump = now;
|
||||
});
|
||||
}
|
||||
|
||||
impl crate::pes::Stream for DiscStream {
|
||||
fn read(&mut self) -> io::Result<Option<crate::pes::PesFrame>> {
|
||||
if let Some(frame) = self.pending_frames.pop_front() {
|
||||
@@ -414,6 +484,7 @@ impl crate::pes::Stream for DiscStream {
|
||||
}
|
||||
|
||||
loop {
|
||||
let t0 = std::time::Instant::now();
|
||||
if !self.fill_extents()? {
|
||||
self.eof = true;
|
||||
// Flush demuxer — last PES packet may still be in the assembler
|
||||
@@ -471,10 +542,25 @@ impl crate::pes::Stream for DiscStream {
|
||||
}
|
||||
}
|
||||
}
|
||||
// Drain any access unit a codec parser buffered past the last
|
||||
// PES (DTS-HD's final core+extension unit, assembled across
|
||||
// PES boundaries).
|
||||
let pid_to_track = &self.pid_to_track;
|
||||
let pending = &mut self.pending_frames;
|
||||
for (pid, parser) in self.parsers.iter_mut() {
|
||||
let Some(&(_, track)) = pid_to_track.iter().find(|(p, _)| p == pid) else {
|
||||
continue;
|
||||
};
|
||||
for frame in parser.flush() {
|
||||
pending.push_back(crate::pes::PesFrame::from_codec_frame(track, frame));
|
||||
}
|
||||
}
|
||||
return Ok(self.pending_frames.pop_front());
|
||||
}
|
||||
|
||||
let bytes = self.buf_valid;
|
||||
let t1 = std::time::Instant::now();
|
||||
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
|
||||
@@ -482,21 +568,42 @@ impl crate::pes::Stream for DiscStream {
|
||||
|
||||
if let Some(ref mut demuxer) = self.ts_demuxer {
|
||||
let packets = demuxer.feed(&self.read_buf[..bytes]);
|
||||
for pes in &packets {
|
||||
if let Some((_, track)) =
|
||||
self.pid_to_track.iter().find(|(pid, _)| *pid == pes.pid)
|
||||
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();
|
||||
for pes in packets {
|
||||
if let Some((_, track)) = self
|
||||
.pid_to_track
|
||||
.iter()
|
||||
.find(|(pid, _)| *pid == pes.pid)
|
||||
.copied()
|
||||
{
|
||||
if let Some((_, parser)) =
|
||||
if skip_parse {
|
||||
// Profiling escape hatch — bypass the codec
|
||||
// parser and pass the raw PES bytes straight
|
||||
// through as a single PesFrame. Lets us
|
||||
// attribute consumer-thread time to
|
||||
// "demux + framing" vs "codec parse".
|
||||
self.pending_frames.push_back(crate::pes::PesFrame {
|
||||
track,
|
||||
pts: pes.pts.map(super::codec::pts_to_ns).unwrap_or(0),
|
||||
keyframe: false,
|
||||
data: pes.data,
|
||||
duration_ns: None,
|
||||
});
|
||||
} else if let Some((_, parser)) =
|
||||
self.parsers.iter_mut().find(|(pid, _)| *pid == pes.pid)
|
||||
{
|
||||
for frame in parser.parse(pes) {
|
||||
for frame in parser.parse(&pes) {
|
||||
self.pending_frames.push_back(
|
||||
crate::pes::PesFrame::from_codec_frame(*track, frame),
|
||||
crate::pes::PesFrame::from_codec_frame(track, frame),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let t3 = std::time::Instant::now();
|
||||
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 {
|
||||
@@ -570,6 +677,13 @@ impl crate::pes::Stream for DiscStream {
|
||||
}
|
||||
|
||||
fn headers_ready(&self) -> bool {
|
||||
// FREEMKV_SKIP_PARSE bypasses codec parsers entirely for
|
||||
// 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() {
|
||||
return true;
|
||||
}
|
||||
for (idx, s) in self.title.streams.iter().enumerate() {
|
||||
if let crate::disc::Stream::Video(v) = s {
|
||||
if !v.secondary && self.codec_private(idx).is_none() {
|
||||
@@ -579,6 +693,10 @@ impl crate::pes::Stream for DiscStream {
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
fn errors(&self) -> u64 {
|
||||
self.errors
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
+59
-13
@@ -264,6 +264,13 @@ pub fn read_element_header(r: &mut impl Read) -> io::Result<(u32, u64, usize)> {
|
||||
|
||||
/// Read an unsigned integer value of `len` bytes.
|
||||
pub fn read_uint_val(r: &mut impl Read, len: usize) -> io::Result<u64> {
|
||||
// An EBML unsigned integer is at most 8 bytes. A malformed element
|
||||
// claiming `len > 8` would index past this stack buffer and panic
|
||||
// (DoS on untrusted input) — reject it at the source so every caller
|
||||
// is safe, not just the ones that pre-check.
|
||||
if len > 8 {
|
||||
return Err(crate::error::Error::MkvInvalid.into());
|
||||
}
|
||||
let mut buf = [0u8; 8];
|
||||
r.read_exact(&mut buf[..len])?;
|
||||
let mut val = 0u64;
|
||||
@@ -273,23 +280,33 @@ pub fn read_uint_val(r: &mut impl Read, len: usize) -> io::Result<u64> {
|
||||
Ok(val)
|
||||
}
|
||||
|
||||
/// Read a float value (4 or 8 bytes).
|
||||
/// 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.
|
||||
pub fn read_float_val(r: &mut impl Read, len: usize) -> io::Result<f64> {
|
||||
if len == 4 {
|
||||
let mut buf = [0u8; 4];
|
||||
r.read_exact(&mut buf)?;
|
||||
Ok(f32::from_be_bytes(buf) as f64)
|
||||
} else {
|
||||
let mut buf = [0u8; 8];
|
||||
r.read_exact(&mut buf)?;
|
||||
Ok(f64::from_be_bytes(buf))
|
||||
match len {
|
||||
0 => Ok(0.0),
|
||||
4 => {
|
||||
let mut buf = [0u8; 4];
|
||||
r.read_exact(&mut buf)?;
|
||||
Ok(f32::from_be_bytes(buf) as f64)
|
||||
}
|
||||
8 => {
|
||||
let mut buf = [0u8; 8];
|
||||
r.read_exact(&mut buf)?;
|
||||
Ok(f64::from_be_bytes(buf))
|
||||
}
|
||||
_ => Err(crate::error::Error::MkvInvalid.into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Read a UTF-8 string value of `len` bytes.
|
||||
pub fn read_string_val(r: &mut impl Read, len: usize) -> io::Result<String> {
|
||||
let mut buf = vec![0u8; len];
|
||||
r.read_exact(&mut buf)?;
|
||||
let mut buf = read_exact_bounded(r, len)?;
|
||||
// Strip trailing nulls
|
||||
while buf.last() == Some(&0) {
|
||||
buf.pop();
|
||||
@@ -299,8 +316,22 @@ pub fn read_string_val(r: &mut impl Read, len: usize) -> io::Result<String> {
|
||||
|
||||
/// Read binary data of `len` bytes.
|
||||
pub fn read_binary_val(r: &mut impl Read, len: usize) -> io::Result<Vec<u8>> {
|
||||
let mut buf = vec![0u8; len];
|
||||
r.read_exact(&mut buf)?;
|
||||
read_exact_bounded(r, len)
|
||||
}
|
||||
|
||||
/// Read exactly `len` bytes WITHOUT trusting `len` to size the allocation.
|
||||
///
|
||||
/// `vec![0u8; len]` on an attacker-controlled EBML size would allocate
|
||||
/// gigabytes before the read fails. Instead we cap the reader to `len`
|
||||
/// and grow the buffer as bytes actually arrive: a malformed element that
|
||||
/// claims a huge length but supplies few bytes allocates only what it
|
||||
/// delivers, then errors on the short read.
|
||||
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());
|
||||
}
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
@@ -321,6 +352,12 @@ pub const EBML_DOC_TYPE_READ_VERSION: u32 = 0x4285;
|
||||
// Segment
|
||||
pub const SEGMENT: u32 = 0x1853_8067;
|
||||
|
||||
// SeekHead
|
||||
pub const SEEK_HEAD: u32 = 0x114D_9B74;
|
||||
pub const SEEK: u32 = 0x4DBB;
|
||||
pub const SEEK_ID: u32 = 0x53AB;
|
||||
pub const SEEK_POSITION: u32 = 0x53AC;
|
||||
|
||||
// Segment Info
|
||||
pub const INFO: u32 = 0x1549_A966;
|
||||
pub const TIMESTAMP_SCALE: u32 = 0x2A_D7B1;
|
||||
@@ -356,6 +393,12 @@ pub const MATRIX_COEFFICIENTS: u32 = 0x55B1;
|
||||
pub const PRIMARIES: u32 = 0x55BB;
|
||||
pub const RANGE: u32 = 0x55B9;
|
||||
|
||||
// Dolby Vision — BlockAdditionMapping carries the DOVIDecoderConfigurationRecord
|
||||
// (dvcC) so players / mediainfo recognise the track as Dolby Vision.
|
||||
pub const BLOCK_ADDITION_MAPPING: u32 = 0x41E4;
|
||||
pub const BLOCK_ADD_ID_TYPE: u32 = 0x41E7;
|
||||
pub const BLOCK_ADD_ID_EXTRA_DATA: u32 = 0x41ED;
|
||||
|
||||
// Audio
|
||||
pub const AUDIO: u32 = 0xE1;
|
||||
pub const SAMPLING_FREQUENCY: u32 = 0xB5;
|
||||
@@ -366,6 +409,9 @@ pub const BIT_DEPTH: u32 = 0x6264;
|
||||
pub const CLUSTER: u32 = 0x1F43_B675;
|
||||
pub const CLUSTER_TIMESTAMP: u32 = 0xE7;
|
||||
pub const SIMPLE_BLOCK: u32 = 0xA3;
|
||||
pub const BLOCK_GROUP: u32 = 0xA0;
|
||||
pub const BLOCK: u32 = 0xA1;
|
||||
pub const BLOCK_DURATION: u32 = 0x9B;
|
||||
|
||||
// Cues
|
||||
pub const CUES: u32 = 0x1C53_BB6B;
|
||||
|
||||
@@ -0,0 +1,404 @@
|
||||
//! Fragmented MP4 muxer — **stub** for Phase 3.
|
||||
//!
|
||||
//! 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`].
|
||||
//!
|
||||
//! ## What's TODO (tracked in Phase 4 / v0.22.0 scope)
|
||||
//!
|
||||
//! - `moof` box: `mfhd` (sequence_number) + `traf` (`tfhd` + `tfdt`
|
||||
//! + `trun` with sample sizes, durations, flags, composition offsets).
|
||||
//! - `mdat` box: concatenated sample data.
|
||||
//! - 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.
|
||||
//! - Sample-flags computation (sync vs. delta, depends_on, etc.).
|
||||
//! - Edit lists / fragment_duration for accurate seeking.
|
||||
//!
|
||||
//! Reference: ISO/IEC 14496-12 §8 (Movie Fragments).
|
||||
|
||||
use std::io::{self, Write};
|
||||
|
||||
// Box type literals — four-character codes per ISO/IEC 14496-12 §4.2.
|
||||
|
||||
const FTYP: [u8; 4] = *b"ftyp";
|
||||
const MOOV: [u8; 4] = *b"moov";
|
||||
const MVHD: [u8; 4] = *b"mvhd";
|
||||
const TRAK: [u8; 4] = *b"trak";
|
||||
const TKHD: [u8; 4] = *b"tkhd";
|
||||
const MDIA: [u8; 4] = *b"mdia";
|
||||
const MDHD: [u8; 4] = *b"mdhd";
|
||||
const HDLR: [u8; 4] = *b"hdlr";
|
||||
const MINF: [u8; 4] = *b"minf";
|
||||
const VMHD: [u8; 4] = *b"vmhd";
|
||||
const DINF: [u8; 4] = *b"dinf";
|
||||
const DREF: [u8; 4] = *b"dref";
|
||||
const URL_: [u8; 4] = *b"url ";
|
||||
const STBL: [u8; 4] = *b"stbl";
|
||||
const STSD: [u8; 4] = *b"stsd";
|
||||
const STTS: [u8; 4] = *b"stts";
|
||||
const STSC: [u8; 4] = *b"stsc";
|
||||
const STSZ: [u8; 4] = *b"stsz";
|
||||
const STCO: [u8; 4] = *b"stco";
|
||||
const MVEX: [u8; 4] = *b"mvex";
|
||||
const TREX: [u8; 4] = *b"trex";
|
||||
|
||||
/// Default movie timescale — 90 kHz lines up with MPEG-TS PTS and the
|
||||
/// HEVC SPS `vui_time_scale` for film content, simplifying the math
|
||||
/// when fragment emission lands.
|
||||
const MOVIE_TIMESCALE: u32 = 90_000;
|
||||
/// Video track ID. fMP4 init segments conventionally use track_ID=1
|
||||
/// for the primary video track; a single-track DASH representation has
|
||||
/// no reason to deviate.
|
||||
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.
|
||||
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.
|
||||
#[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.
|
||||
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;
|
||||
}
|
||||
// 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)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the `ftyp` box. `major_brand = "iso6"`, `minor_version = 1`,
|
||||
/// compatible brands `iso6 dash msdh hvc1` — the same conservative set
|
||||
/// shaka-packager uses for HEVC-in-fMP4 outputs.
|
||||
fn build_ftyp() -> Vec<u8> {
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(b"iso6");
|
||||
body.extend_from_slice(&1u32.to_be_bytes());
|
||||
body.extend_from_slice(b"iso6");
|
||||
body.extend_from_slice(b"dash");
|
||||
body.extend_from_slice(b"msdh");
|
||||
body.extend_from_slice(b"hvc1");
|
||||
wrap_box(&FTYP, &body)
|
||||
}
|
||||
|
||||
/// Build the `moov` box — minimal skeleton. Single video trak, no
|
||||
/// hvcC inside stsd yet (TODO: full hvc1 sample entry).
|
||||
fn build_moov() -> Vec<u8> {
|
||||
let mvhd = build_mvhd();
|
||||
let trak = build_video_trak();
|
||||
let mvex = build_mvex();
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&mvhd);
|
||||
body.extend_from_slice(&trak);
|
||||
body.extend_from_slice(&mvex);
|
||||
wrap_box(&MOOV, &body)
|
||||
}
|
||||
|
||||
fn build_mvhd() -> Vec<u8> {
|
||||
// Version 0, 100 bytes total body. Fields per ISO/IEC 14496-12 §8.2.2.
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&[0, 0, 0, 0]); // version + flags
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // creation_time
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // modification_time
|
||||
body.extend_from_slice(&MOVIE_TIMESCALE.to_be_bytes());
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // duration = 0 (fragmented)
|
||||
body.extend_from_slice(&0x0001_0000u32.to_be_bytes()); // rate 1.0
|
||||
body.extend_from_slice(&0x0100u16.to_be_bytes()); // volume 1.0
|
||||
body.extend_from_slice(&[0u8; 2]); // reserved
|
||||
body.extend_from_slice(&[0u8; 8]); // reserved
|
||||
// 3x3 identity transformation matrix in 16.16 fixed point.
|
||||
for v in [0x1_0000u32, 0, 0, 0, 0x1_0000, 0, 0, 0, 0x4000_0000] {
|
||||
body.extend_from_slice(&v.to_be_bytes());
|
||||
}
|
||||
body.extend_from_slice(&[0u8; 24]); // pre_defined[6]
|
||||
body.extend_from_slice(&2u32.to_be_bytes()); // next_track_ID (1 reserved for video)
|
||||
wrap_box(&MVHD, &body)
|
||||
}
|
||||
|
||||
fn build_video_trak() -> Vec<u8> {
|
||||
let tkhd = build_tkhd();
|
||||
let mdia = build_mdia();
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&tkhd);
|
||||
body.extend_from_slice(&mdia);
|
||||
wrap_box(&TRAK, &body)
|
||||
}
|
||||
|
||||
fn build_tkhd() -> Vec<u8> {
|
||||
let mut body = Vec::new();
|
||||
// version=0 | flags=0x000007 (track_enabled | in_movie | in_preview)
|
||||
body.extend_from_slice(&[0, 0, 0, 7]);
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // creation_time
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // modification_time
|
||||
body.extend_from_slice(&VIDEO_TRACK_ID.to_be_bytes());
|
||||
body.extend_from_slice(&[0u8; 4]); // reserved
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // duration
|
||||
body.extend_from_slice(&[0u8; 8]); // reserved
|
||||
body.extend_from_slice(&0u16.to_be_bytes()); // layer
|
||||
body.extend_from_slice(&0u16.to_be_bytes()); // alternate_group
|
||||
body.extend_from_slice(&0u16.to_be_bytes()); // volume (video=0)
|
||||
body.extend_from_slice(&[0u8; 2]); // reserved
|
||||
// 3x3 identity matrix.
|
||||
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.
|
||||
body.extend_from_slice(&(1920u32 << 16).to_be_bytes());
|
||||
body.extend_from_slice(&(1080u32 << 16).to_be_bytes());
|
||||
wrap_box(&TKHD, &body)
|
||||
}
|
||||
|
||||
fn build_mdia() -> Vec<u8> {
|
||||
let mdhd = build_mdhd();
|
||||
let hdlr = build_hdlr_vide();
|
||||
let minf = build_minf();
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&mdhd);
|
||||
body.extend_from_slice(&hdlr);
|
||||
body.extend_from_slice(&minf);
|
||||
wrap_box(&MDIA, &body)
|
||||
}
|
||||
|
||||
fn build_mdhd() -> Vec<u8> {
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&[0, 0, 0, 0]); // version + flags
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // creation_time
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // modification_time
|
||||
body.extend_from_slice(&MOVIE_TIMESCALE.to_be_bytes());
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // duration
|
||||
// language: 'und' in 5-bit-per-char ISO 639-2 packed (bit 15 = 0).
|
||||
body.extend_from_slice(&[0x55, 0xC4]);
|
||||
body.extend_from_slice(&0u16.to_be_bytes()); // pre_defined
|
||||
wrap_box(&MDHD, &body)
|
||||
}
|
||||
|
||||
fn build_hdlr_vide() -> Vec<u8> {
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&[0, 0, 0, 0]); // version + flags
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // pre_defined
|
||||
body.extend_from_slice(b"vide");
|
||||
body.extend_from_slice(&[0u8; 12]); // reserved
|
||||
body.extend_from_slice(b"VideoHandler\0");
|
||||
wrap_box(&HDLR, &body)
|
||||
}
|
||||
|
||||
fn build_minf() -> Vec<u8> {
|
||||
let vmhd = build_vmhd();
|
||||
let dinf = build_dinf();
|
||||
let stbl = build_stbl();
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&vmhd);
|
||||
body.extend_from_slice(&dinf);
|
||||
body.extend_from_slice(&stbl);
|
||||
wrap_box(&MINF, &body)
|
||||
}
|
||||
|
||||
fn build_vmhd() -> Vec<u8> {
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&[0, 0, 0, 1]); // version + flags=1
|
||||
body.extend_from_slice(&0u16.to_be_bytes()); // graphicsmode
|
||||
body.extend_from_slice(&[0u8; 6]); // opcolor
|
||||
wrap_box(&VMHD, &body)
|
||||
}
|
||||
|
||||
fn build_dinf() -> Vec<u8> {
|
||||
let mut dref_body = Vec::new();
|
||||
dref_body.extend_from_slice(&[0, 0, 0, 0]);
|
||||
dref_body.extend_from_slice(&1u32.to_be_bytes()); // entry_count
|
||||
// url with flags=1 (self-contained) and zero name.
|
||||
let url_body = [0u8, 0, 0, 1];
|
||||
dref_body.extend_from_slice(&wrap_box(&URL_, &url_body));
|
||||
let dref = wrap_box(&DREF, &dref_body);
|
||||
wrap_box(&DINF, &dref)
|
||||
}
|
||||
|
||||
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.
|
||||
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
|
||||
let stsd = wrap_box(&STSD, &stsd_body);
|
||||
|
||||
// Empty stts/stsc/stsz/stco — fragmented init has no samples here.
|
||||
let stts = wrap_box(&STTS, &[0, 0, 0, 0, 0, 0, 0, 0]); // version+flags, count=0
|
||||
let stsc = wrap_box(&STSC, &[0, 0, 0, 0, 0, 0, 0, 0]);
|
||||
let stsz = wrap_box(
|
||||
&STSZ,
|
||||
&[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], // version+flags, sample_size=0, count=0
|
||||
);
|
||||
let stco = wrap_box(&STCO, &[0, 0, 0, 0, 0, 0, 0, 0]);
|
||||
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&stsd);
|
||||
body.extend_from_slice(&stts);
|
||||
body.extend_from_slice(&stsc);
|
||||
body.extend_from_slice(&stsz);
|
||||
body.extend_from_slice(&stco);
|
||||
wrap_box(&STBL, &body)
|
||||
}
|
||||
|
||||
fn build_mvex() -> Vec<u8> {
|
||||
// trex: track_ID=1, default_sample_description_index=1, others=0.
|
||||
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());
|
||||
trex_body.extend_from_slice(&1u32.to_be_bytes()); // default_sample_description_index
|
||||
trex_body.extend_from_slice(&0u32.to_be_bytes()); // default_sample_duration
|
||||
trex_body.extend_from_slice(&0u32.to_be_bytes()); // default_sample_size
|
||||
trex_body.extend_from_slice(&0u32.to_be_bytes()); // default_sample_flags
|
||||
let trex = wrap_box(&TREX, &trex_body);
|
||||
wrap_box(&MVEX, &trex)
|
||||
}
|
||||
|
||||
/// 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.
|
||||
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);
|
||||
out.extend_from_slice(&size.to_be_bytes());
|
||||
out.extend_from_slice(box_type);
|
||||
out.extend_from_slice(body);
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Decode the first box's size + type from `buf`.
|
||||
fn read_box_header(buf: &[u8]) -> (u32, [u8; 4]) {
|
||||
let size = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
|
||||
let bt = [buf[4], buf[5], buf[6], buf[7]];
|
||||
(size, bt)
|
||||
}
|
||||
|
||||
#[test]
|
||||
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.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
let (ftyp_size, ftyp_type) = read_box_header(&sink);
|
||||
assert_eq!(&ftyp_type, b"ftyp");
|
||||
assert!(ftyp_size >= 24, "ftyp too small: {ftyp_size}");
|
||||
|
||||
let (moov_size, moov_type) = read_box_header(&sink[ftyp_size as usize..]);
|
||||
assert_eq!(&moov_type, b"moov");
|
||||
assert!(moov_size > 100, "moov skeleton too small: {moov_size}");
|
||||
|
||||
// Stub guarantee: no media bytes after the init segment.
|
||||
let total = ftyp_size as usize + moov_size as usize;
|
||||
assert_eq!(sink.len(), total, "stub leaked media bytes past moov");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn moov_contains_trak_mvex() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = Fmp4Mux::new(&mut sink);
|
||||
mux.write_video(0, true, &[]).unwrap();
|
||||
mux.finish().unwrap();
|
||||
drop(sink);
|
||||
|
||||
// Re-emit into a fresh buffer for parsing.
|
||||
let mut buf: Vec<u8> = Vec::new();
|
||||
let mut mux2 = Fmp4Mux::new(&mut buf);
|
||||
mux2.write_video(0, true, &[]).unwrap();
|
||||
mux2.finish().unwrap();
|
||||
drop(mux2);
|
||||
|
||||
// Find moov payload start.
|
||||
let (ftyp_size, _) = read_box_header(&buf);
|
||||
let moov_start = ftyp_size as usize;
|
||||
let (moov_size, _) = read_box_header(&buf[moov_start..]);
|
||||
let moov_payload = &buf[moov_start + 8..moov_start + moov_size as usize];
|
||||
|
||||
// Scan for the trak and mvex four-CC anywhere in the moov payload.
|
||||
let has_trak = moov_payload.windows(4).any(|w| w == b"trak");
|
||||
let has_mvex = moov_payload.windows(4).any(|w| w == b"mvex");
|
||||
assert!(has_trak, "moov missing trak");
|
||||
assert!(has_mvex, "moov missing mvex");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,279 @@
|
||||
//! HEVC (H.265) elementary stream muxer — Annex B byte stream.
|
||||
//!
|
||||
//! Consumes [`PesFrame`](crate::pes::PesFrame)s for a single video track
|
||||
//! and writes them as a raw `.hevc` / `.h265` Annex B byte stream:
|
||||
//! `00 00 00 01 | NAL_unit | 00 00 00 01 | NAL_unit | …` with no
|
||||
//! container framing.
|
||||
//!
|
||||
//! On the first frame the muxer emits the codec_private's VPS, SPS, PPS
|
||||
//! (parsed from a `HEVCDecoderConfigurationRecord` in
|
||||
//! `length-prefixed-in-hvcC` form), then converts each PES frame's
|
||||
//! length-prefixed NAL units to Annex B and writes them.
|
||||
//!
|
||||
//! Sequential-only — no Cues, no backpatch. Target sink is any
|
||||
//! [`SequentialSink`](crate::io::sink::SequentialSink): file, socket,
|
||||
//! pipe, anything `Write + Send`.
|
||||
|
||||
use std::io::{self, Write};
|
||||
|
||||
/// Annex B 4-byte start code.
|
||||
pub(crate) const START_CODE: [u8; 4] = [0x00, 0x00, 0x00, 0x01];
|
||||
|
||||
/// HEVC NAL unit type bits live in `(byte0 >> 1) & 0x3F` in Annex B.
|
||||
/// We don't filter NAL types here — the muxer is format-only — but we
|
||||
/// keep the constant as documentation of the field layout.
|
||||
#[allow(dead_code)]
|
||||
const HEVC_NAL_TYPE_MASK: u8 = 0x3F;
|
||||
|
||||
/// Streaming HEVC Annex B muxer.
|
||||
///
|
||||
/// One instance per output stream. Tracks whether parameter sets have
|
||||
/// already been emitted so they're written exactly once at the head of
|
||||
/// the stream, mirroring the convention used by `ffmpeg -c:v copy -f
|
||||
/// hevc`.
|
||||
pub struct HevcMux<W: Write> {
|
||||
writer: W,
|
||||
/// `HEVCDecoderConfigurationRecord` payload (hvcC). Parsed lazily
|
||||
/// on the first `write_frame` so callers can set it after
|
||||
/// construction but before the first frame.
|
||||
codec_private: Option<Vec<u8>>,
|
||||
/// Set once VPS/SPS/PPS have been written to the stream. Subsequent
|
||||
/// frames write only their own NAL units.
|
||||
params_written: bool,
|
||||
}
|
||||
|
||||
impl<W: Write> HevcMux<W> {
|
||||
/// Construct over `writer`. The muxer does not impose any extra
|
||||
/// buffering of its own — the sink owns its write buffering policy
|
||||
/// (see [`LocalFileSink`](crate::io::sink::LocalFileSink) and
|
||||
/// [`SocketSink`](crate::io::sink::SocketSink)).
|
||||
pub fn new(writer: W) -> Self {
|
||||
Self {
|
||||
writer,
|
||||
codec_private: None,
|
||||
params_written: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Provide the `HEVCDecoderConfigurationRecord` (hvcC) so the muxer
|
||||
/// can prepend VPS/SPS/PPS Annex B NALs at stream start. Optional —
|
||||
/// if the PES frames already carry inline parameter sets (some
|
||||
/// upstream demuxers do this), skipping this call is fine.
|
||||
pub fn set_codec_private(&mut self, data: Vec<u8>) {
|
||||
self.codec_private = Some(data);
|
||||
}
|
||||
|
||||
/// Write one PES frame (= one access unit) as Annex B NAL units.
|
||||
///
|
||||
/// Input may be either:
|
||||
/// - 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.
|
||||
///
|
||||
/// `_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 {
|
||||
if let Some(cp) = &self.codec_private {
|
||||
if let Some(params) = hvcc_to_annex_b(cp) {
|
||||
self.writer.write_all(¶ms)?;
|
||||
}
|
||||
}
|
||||
self.params_written = true;
|
||||
}
|
||||
let annex_b = length_prefixed_to_annex_b(data);
|
||||
self.writer.write_all(&annex_b)
|
||||
}
|
||||
|
||||
/// Flush the underlying writer. No trailer NAL is needed — an Annex
|
||||
/// B stream ends whenever the file/socket ends.
|
||||
pub fn finish(&mut self) -> io::Result<()> {
|
||||
self.writer.flush()
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a `HEVCDecoderConfigurationRecord` (hvcC) into Annex B NAL
|
||||
/// units. Returns `Some(bytes)` if at least one NAL was extracted, else
|
||||
/// `None`.
|
||||
///
|
||||
/// Layout (per ISO/IEC 14496-15 §8.3.3.1.2):
|
||||
/// - 22-byte fixed header
|
||||
/// - byte 22 = `numOfArrays`
|
||||
/// - each array: `array_completeness:1 | reserved:1 | NAL_unit_type:6`,
|
||||
/// `numNalus:u16-BE`, then `numNalus` × `(nalUnitLength:u16-BE +
|
||||
/// NAL bytes)`.
|
||||
///
|
||||
/// 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>> {
|
||||
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;
|
||||
}
|
||||
offset += 1; // array_completeness + nal_type byte
|
||||
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(&START_CODE);
|
||||
out.extend_from_slice(&hvcc[offset..offset + nal_len]);
|
||||
offset += nal_len;
|
||||
}
|
||||
}
|
||||
if out.is_empty() { None } else { Some(out) }
|
||||
}
|
||||
|
||||
/// 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.
|
||||
pub(crate) fn length_prefixed_to_annex_b(data: &[u8]) -> Vec<u8> {
|
||||
let mut out = Vec::with_capacity(data.len() + (data.len() / 32));
|
||||
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() {
|
||||
// Mid-NAL truncation — fall through to the pass-through path
|
||||
// rather than emitting a half-NAL.
|
||||
return data.to_vec();
|
||||
}
|
||||
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();
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn length_prefixed_converts_to_annex_b() {
|
||||
// Two NALs: [3-byte payload AA BB CC] and [2-byte payload DD EE].
|
||||
let mut buf = Vec::new();
|
||||
buf.extend_from_slice(&3u32.to_be_bytes());
|
||||
buf.extend_from_slice(&[0xAA, 0xBB, 0xCC]);
|
||||
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
|
||||
];
|
||||
assert_eq!(&got[..], &want[..]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn already_annex_b_passes_through_when_no_lengths_match() {
|
||||
// A buffer < 4 bytes can't parse a length prefix at all →
|
||||
// pass-through path triggers.
|
||||
let raw = [0xAA, 0xBB, 0xCC];
|
||||
let got = length_prefixed_to_annex_b(&raw);
|
||||
assert_eq!(&got[..], &raw[..]);
|
||||
}
|
||||
|
||||
#[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.
|
||||
let mut raw = Vec::new();
|
||||
raw.extend_from_slice(&100u32.to_be_bytes());
|
||||
raw.extend_from_slice(&[0xAA, 0xBB, 0xCC]);
|
||||
let got = length_prefixed_to_annex_b(&raw);
|
||||
assert_eq!(&got[..], &raw[..]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hvcc_extracts_vps_sps_pps() {
|
||||
// Build a minimal-but-valid hvcC: 22-byte header, then 3 arrays
|
||||
// (VPS / SPS / PPS), each with 1 NAL of a 4-byte payload that
|
||||
// we can spot in the output.
|
||||
let mut hvcc = vec![0u8; 22];
|
||||
hvcc.push(3); // numOfArrays
|
||||
for (nal_type, payload) in [
|
||||
(32u8, [0x40, 0x01, 0x0C, 0x01]),
|
||||
(33, [0x42, 0x01, 0x01, 0x01]),
|
||||
(34, [0x44, 0x01, 0xC1, 0x72]),
|
||||
] {
|
||||
hvcc.push(nal_type & 0x3F);
|
||||
hvcc.extend_from_slice(&1u16.to_be_bytes()); // numNalus
|
||||
hvcc.extend_from_slice(&(payload.len() as u16).to_be_bytes());
|
||||
hvcc.extend_from_slice(&payload);
|
||||
}
|
||||
|
||||
let annex_b = hvcc_to_annex_b(&hvcc).expect("at least one NAL");
|
||||
// Three NALs × (4-byte start + 4-byte payload) = 24 bytes.
|
||||
assert_eq!(annex_b.len(), 24);
|
||||
assert_eq!(&annex_b[..4], &START_CODE);
|
||||
assert_eq!(&annex_b[8..12], &START_CODE);
|
||||
assert_eq!(&annex_b[16..20], &START_CODE);
|
||||
assert_eq!(annex_b[4], 0x40); // VPS first byte
|
||||
assert_eq!(annex_b[12], 0x42); // SPS first byte
|
||||
assert_eq!(annex_b[20], 0x44); // PPS first byte
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mux_writes_params_then_frames() {
|
||||
// Build hvcC with one SPS to verify params-once semantics.
|
||||
let mut hvcc = vec![0u8; 22];
|
||||
hvcc.push(1);
|
||||
hvcc.push(33);
|
||||
hvcc.extend_from_slice(&1u16.to_be_bytes());
|
||||
hvcc.extend_from_slice(&3u16.to_be_bytes());
|
||||
hvcc.extend_from_slice(&[0x42, 0x01, 0x01]);
|
||||
|
||||
let mut frame_data = Vec::new();
|
||||
frame_data.extend_from_slice(&2u32.to_be_bytes());
|
||||
frame_data.extend_from_slice(&[0xAA, 0xBB]);
|
||||
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = HevcMux::new(&mut sink);
|
||||
mux.set_codec_private(hvcc);
|
||||
mux.write_frame(0, &frame_data).unwrap();
|
||||
// Second frame — no SPS re-emission.
|
||||
mux.write_frame(40_000_000, &frame_data).unwrap();
|
||||
mux.finish().unwrap();
|
||||
|
||||
// SPS NAL (7 bytes) + 2× frame NAL (6 bytes) = 19 bytes.
|
||||
assert_eq!(sink.len(), 7 + 6 + 6);
|
||||
// Start codes at offsets 0 (SPS), 7 (frame1), 13 (frame2).
|
||||
assert_eq!(&sink[0..4], &START_CODE);
|
||||
assert_eq!(&sink[7..11], &START_CODE);
|
||||
assert_eq!(&sink[13..17], &START_CODE);
|
||||
assert_eq!(sink[4], 0x42);
|
||||
assert_eq!(sink[11], 0xAA);
|
||||
assert_eq!(sink[17], 0xAA);
|
||||
}
|
||||
}
|
||||
-101
@@ -1,101 +0,0 @@
|
||||
//! ISO sector reader — file-backed SectorSource for Blu-ray ISO images.
|
||||
//!
|
||||
//! An ISO is a flat image of 2048-byte sectors. Sector N starts at byte offset N * 2048.
|
||||
//! Used by DiscStream::open_iso() and Disc::scan_image().
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
use crate::sector::SectorSource;
|
||||
use std::fs::File;
|
||||
use std::io::{Read, Seek, SeekFrom};
|
||||
use std::path::Path;
|
||||
|
||||
const SECTOR_SIZE: u64 = 2048;
|
||||
|
||||
/// File-backed sector reader for ISO images.
|
||||
pub struct IsoSectorReader {
|
||||
file: File,
|
||||
capacity: u32,
|
||||
}
|
||||
|
||||
impl IsoSectorReader {
|
||||
pub fn open(path: &str) -> std::io::Result<Self> {
|
||||
let file = File::open(Path::new(path))?;
|
||||
let size = file.metadata()?.len();
|
||||
let sectors = size / SECTOR_SIZE;
|
||||
if sectors > u32::MAX as u64 {
|
||||
return Err(crate::error::Error::IsoTooLarge {
|
||||
path: path.to_string(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
let capacity = sectors as u32;
|
||||
Ok(Self { file, capacity })
|
||||
}
|
||||
|
||||
pub fn capacity_sectors(&self) -> u32 {
|
||||
self.capacity
|
||||
}
|
||||
}
|
||||
|
||||
impl SectorSource for IsoSectorReader {
|
||||
fn read_sectors(
|
||||
&mut self,
|
||||
lba: u32,
|
||||
count: u16,
|
||||
buf: &mut [u8],
|
||||
_recovery: bool,
|
||||
) -> Result<usize> {
|
||||
let bytes = count as usize * SECTOR_SIZE as usize;
|
||||
self.file
|
||||
.seek(SeekFrom::Start(lba as u64 * SECTOR_SIZE))
|
||||
.map_err(|e| Error::IoError { source: e })?;
|
||||
self.file
|
||||
.read_exact(&mut buf[..bytes])
|
||||
.map_err(|e| Error::IoError { source: e })?;
|
||||
Ok(bytes)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn iso_reader_read_sectors() {
|
||||
let mut data = vec![0u8; 4 * SECTOR_SIZE as usize];
|
||||
for i in 0..4u8 {
|
||||
let offset = i as usize * SECTOR_SIZE as usize;
|
||||
data[offset] = i + 1;
|
||||
data[offset + 2047] = i + 100;
|
||||
}
|
||||
|
||||
let dir = std::env::temp_dir().join("freemkv_test_iso_read");
|
||||
std::fs::write(&dir, &data).unwrap();
|
||||
|
||||
let mut reader = IsoSectorReader::open(dir.to_str().unwrap()).unwrap();
|
||||
assert_eq!(reader.capacity_sectors(), 4);
|
||||
|
||||
let mut buf = [0u8; 2048];
|
||||
reader.read_sectors(0, 1, &mut buf, true).unwrap();
|
||||
assert_eq!(buf[0], 1);
|
||||
assert_eq!(buf[2047], 100);
|
||||
|
||||
reader.read_sectors(2, 1, &mut buf, true).unwrap();
|
||||
assert_eq!(buf[0], 3);
|
||||
assert_eq!(buf[2047], 102);
|
||||
|
||||
std::fs::remove_file(&dir).ok();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn iso_reader_capacity() {
|
||||
let data = vec![0u8; 10 * SECTOR_SIZE as usize];
|
||||
let dir = std::env::temp_dir().join("freemkv_test_iso_cap");
|
||||
std::fs::write(&dir, &data).unwrap();
|
||||
|
||||
let reader = IsoSectorReader::open(dir.to_str().unwrap()).unwrap();
|
||||
assert_eq!(reader.capacity_sectors(), 10);
|
||||
|
||||
std::fs::remove_file(&dir).ok();
|
||||
}
|
||||
}
|
||||
+138
-226
@@ -1,57 +1,21 @@
|
||||
//! M2tsStream — BD transport stream with embedded metadata header.
|
||||
//! M2tsStream — BD transport stream write sink.
|
||||
//!
|
||||
//! Write: prepends FMKV metadata header, then muxes PES frames into BD-TS.
|
||||
//! Read: extracts metadata header (or scans PMT), then demuxes BD-TS into PES frames.
|
||||
//! Write: prepends FMKV metadata header, then muxes PES frames into
|
||||
//! BD-TS. The read direction lives on the pipeline highway —
|
||||
//! `m2ts://` URLs route through
|
||||
//! [`super::resolve::input`] → `build_m2ts_pipeline` →
|
||||
//! [`super::pipelined_stream::PipelinedPesStream`], so this type is
|
||||
//! write-only.
|
||||
|
||||
use super::{meta, ts};
|
||||
use super::meta;
|
||||
use crate::disc::{DiscTitle, Stream as DiscStream};
|
||||
use std::io::{self, Read, Write};
|
||||
use std::io::{self, Write};
|
||||
|
||||
type PesSetup = (
|
||||
Vec<u16>,
|
||||
Vec<(u16, Box<dyn super::codec::CodecParser>)>,
|
||||
Vec<(u16, usize)>,
|
||||
);
|
||||
|
||||
/// Size of initial scan buffer for PMT/stream detection.
|
||||
const SCAN_SIZE: usize = 1024 * 1024;
|
||||
|
||||
enum Mode {
|
||||
Write {
|
||||
muxer: super::tsmux::TsMuxer<Box<dyn Write + Send>>,
|
||||
},
|
||||
Read {
|
||||
reader: Box<dyn Read + Send>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Read as many bytes as possible into buf (multiple read calls if needed).
|
||||
/// Bounded by buf.len() — caller controls max bytes read.
|
||||
fn read_fill(r: &mut impl Read, buf: &mut [u8]) -> io::Result<usize> {
|
||||
let mut total = 0;
|
||||
while total < buf.len() {
|
||||
match r.read(&mut buf[total..]) {
|
||||
Ok(0) => break,
|
||||
Ok(n) => total += n,
|
||||
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
}
|
||||
Ok(total)
|
||||
}
|
||||
|
||||
/// BD transport stream with embedded metadata.
|
||||
/// BD transport stream write sink with embedded FMKV metadata
|
||||
/// header.
|
||||
pub struct M2tsStream {
|
||||
disc_title: DiscTitle,
|
||||
mode: Mode,
|
||||
// PES support
|
||||
demuxer: Option<ts::TsDemuxer>,
|
||||
parsers: Vec<(u16, Box<dyn super::codec::CodecParser>)>,
|
||||
pending_frames: std::collections::VecDeque<crate::pes::PesFrame>,
|
||||
pid_to_track: Vec<(u16, usize)>,
|
||||
pes_eof: bool,
|
||||
/// Codec private data per stream (from FMKV header).
|
||||
stored_codec_privates: Vec<Option<Vec<u8>>>,
|
||||
muxer: super::tsmux::TsMuxer<Box<dyn Write + Send>>,
|
||||
}
|
||||
|
||||
impl M2tsStream {
|
||||
@@ -81,205 +45,153 @@ impl M2tsStream {
|
||||
}
|
||||
Ok(Self {
|
||||
disc_title: title.clone(),
|
||||
mode: Mode::Write { muxer },
|
||||
demuxer: None,
|
||||
parsers: Vec::new(),
|
||||
pending_frames: std::collections::VecDeque::new(),
|
||||
pid_to_track: Vec::new(),
|
||||
pes_eof: false,
|
||||
stored_codec_privates: Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
fn setup_pes(streams: &[DiscStream]) -> PesSetup {
|
||||
let mut pids = Vec::new();
|
||||
let mut parsers: Vec<(u16, Box<dyn super::codec::CodecParser>)> = Vec::new();
|
||||
let mut pid_to_track = Vec::new();
|
||||
for (i, s) in streams.iter().enumerate() {
|
||||
let (pid, codec) = match s {
|
||||
DiscStream::Video(v) => (v.pid, v.codec),
|
||||
DiscStream::Audio(a) => (a.pid, a.codec),
|
||||
DiscStream::Subtitle(s) => (s.pid, s.codec),
|
||||
};
|
||||
pids.push(pid);
|
||||
pid_to_track.push((pid, i));
|
||||
parsers.push((pid, super::codec::parser_for_codec(codec, None)));
|
||||
}
|
||||
(pids, parsers, pid_to_track)
|
||||
}
|
||||
|
||||
/// Open an M2TS stream for reading. Takes any Read source — file, pipe, socket.
|
||||
///
|
||||
/// Tries FMKV metadata header first. Falls back to PMT scan of first 1 MB.
|
||||
pub fn open(mut reader: impl Read + Send + 'static) -> io::Result<Self> {
|
||||
// Read first chunk — enough for FMKV header or PMT scan
|
||||
let mut head = vec![0u8; SCAN_SIZE];
|
||||
let head_len = read_fill(&mut reader, &mut head)?;
|
||||
head.truncate(head_len);
|
||||
|
||||
// Try FMKV metadata header from the buffered head
|
||||
let mut cursor = io::Cursor::new(&head);
|
||||
if let Ok(Some(m)) = meta::read_header(&mut cursor) {
|
||||
let header_end = cursor.position() as usize;
|
||||
let title = m.to_title();
|
||||
let (pids, parsers, pid_to_track) = Self::setup_pes(&title.streams);
|
||||
// Chain: remaining head bytes + rest of reader
|
||||
let remaining_head = &head[header_end..];
|
||||
let chain: Box<dyn Read + Send> =
|
||||
Box::new(io::Cursor::new(remaining_head.to_vec()).chain(reader));
|
||||
return Ok(Self {
|
||||
disc_title: title.clone(),
|
||||
mode: Mode::Read { reader: chain },
|
||||
demuxer: if pids.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(ts::TsDemuxer::new(&pids))
|
||||
},
|
||||
parsers,
|
||||
pending_frames: std::collections::VecDeque::new(),
|
||||
pid_to_track,
|
||||
pes_eof: false,
|
||||
stored_codec_privates: title.codec_privates,
|
||||
});
|
||||
}
|
||||
|
||||
// No FMKV header — scan head for PMT
|
||||
let streams = ts::scan_streams(&head)
|
||||
.ok_or_else(|| -> io::Error { crate::error::Error::NoStreams.into() })?;
|
||||
|
||||
let (pids, parsers, pid_to_track) = Self::setup_pes(&streams);
|
||||
|
||||
// Chain: full head (it's all TS data) + rest of reader
|
||||
let chain: Box<dyn Read + Send> = Box::new(io::Cursor::new(head).chain(reader));
|
||||
|
||||
Ok(Self {
|
||||
disc_title: DiscTitle {
|
||||
duration_secs: 0.0, // unknown without seeking
|
||||
streams,
|
||||
..DiscTitle::empty()
|
||||
},
|
||||
mode: Mode::Read { reader: chain },
|
||||
demuxer: if pids.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(ts::TsDemuxer::new(&pids))
|
||||
},
|
||||
parsers,
|
||||
pending_frames: std::collections::VecDeque::new(),
|
||||
pid_to_track,
|
||||
pes_eof: false,
|
||||
stored_codec_privates: Vec::new(),
|
||||
muxer,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl crate::pes::Stream for M2tsStream {
|
||||
fn read(&mut self) -> io::Result<Option<crate::pes::PesFrame>> {
|
||||
if let Some(frame) = self.pending_frames.pop_front() {
|
||||
return Ok(Some(frame));
|
||||
}
|
||||
if self.pes_eof {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
loop {
|
||||
let reader = match &mut self.mode {
|
||||
Mode::Read { reader } => reader,
|
||||
_ => return Err(crate::error::Error::StreamWriteOnly.into()),
|
||||
};
|
||||
let mut buf = vec![0u8; 192 * 1024];
|
||||
let n = reader.read(&mut buf)?;
|
||||
if n == 0 {
|
||||
self.pes_eof = true;
|
||||
// Flush demuxer — last PES packet may still be in the assembler
|
||||
if let Some(ref mut demuxer) = self.demuxer {
|
||||
for pes in &demuxer.flush() {
|
||||
if let Some((_, track)) =
|
||||
self.pid_to_track.iter().find(|(pid, _)| *pid == pes.pid)
|
||||
{
|
||||
if let Some((_, parser)) =
|
||||
self.parsers.iter_mut().find(|(pid, _)| *pid == pes.pid)
|
||||
{
|
||||
for frame in parser.parse(pes) {
|
||||
self.pending_frames.push_back(
|
||||
crate::pes::PesFrame::from_codec_frame(*track, frame),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return Ok(self.pending_frames.pop_front());
|
||||
}
|
||||
|
||||
if let Some(ref mut demuxer) = self.demuxer {
|
||||
let packets = demuxer.feed(&buf[..n]);
|
||||
for pes in &packets {
|
||||
if let Some((_, track)) =
|
||||
self.pid_to_track.iter().find(|(pid, _)| *pid == pes.pid)
|
||||
{
|
||||
if let Some((_, parser)) =
|
||||
self.parsers.iter_mut().find(|(pid, _)| *pid == pes.pid)
|
||||
{
|
||||
for frame in parser.parse(pes) {
|
||||
self.pending_frames.push_back(
|
||||
crate::pes::PesFrame::from_codec_frame(*track, frame),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(frame) = self.pending_frames.pop_front() {
|
||||
return Ok(Some(frame));
|
||||
}
|
||||
}
|
||||
// Write-only sink. The m2ts:// read direction is served by
|
||||
// `super::resolve::build_m2ts_pipeline` →
|
||||
// `PipelinedPesStream`; routing through this type for reads
|
||||
// was removed when the highway became the only ingress.
|
||||
Err(crate::error::Error::StreamWriteOnly.into())
|
||||
}
|
||||
|
||||
fn write(&mut self, frame: &crate::pes::PesFrame) -> io::Result<()> {
|
||||
match &mut self.mode {
|
||||
Mode::Write { muxer } => muxer.write_frame(frame.track, frame.pts, &frame.data),
|
||||
Mode::Read { .. } => Err(crate::error::Error::StreamReadOnly.into()),
|
||||
}
|
||||
self.muxer
|
||||
.write_frame(frame.track, frame.pts, frame.keyframe, &frame.data)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> io::Result<()> {
|
||||
match &mut self.mode {
|
||||
Mode::Write { muxer } => muxer.finish(),
|
||||
Mode::Read { .. } => Ok(()),
|
||||
}
|
||||
self.muxer.finish()
|
||||
}
|
||||
|
||||
fn info(&self) -> &crate::disc::DiscTitle {
|
||||
&self.disc_title
|
||||
}
|
||||
|
||||
fn codec_private(&self, track: usize) -> Option<Vec<u8>> {
|
||||
// First check stored codec_privates from FMKV header
|
||||
if let Some(Some(cp)) = self.stored_codec_privates.get(track) {
|
||||
return Some(cp.clone());
|
||||
fn codec_private(&self, _track: usize) -> Option<Vec<u8>> {
|
||||
// Write side doesn't have parsers; codec_private flows in
|
||||
// via the title metadata at `create` time and gets baked
|
||||
// into the FMKV header. Nothing to surface back here.
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::disc::{
|
||||
Codec, ColorSpace, ContentFormat, DiscTitle, FrameRate, HdrFormat, Resolution,
|
||||
Stream as DiscStream, VideoStream,
|
||||
};
|
||||
use crate::pes::{PesFrame, Stream as PesStreamTrait};
|
||||
|
||||
const VIDEO_PID: u16 = 0x1011;
|
||||
|
||||
fn make_title() -> DiscTitle {
|
||||
DiscTitle {
|
||||
playlist: String::new(),
|
||||
playlist_id: 0,
|
||||
duration_secs: 0.0,
|
||||
size_bytes: 0,
|
||||
clips: Vec::new(),
|
||||
streams: vec![DiscStream::Video(VideoStream {
|
||||
pid: VIDEO_PID,
|
||||
codec: Codec::Hevc,
|
||||
resolution: Resolution::R1080p,
|
||||
frame_rate: FrameRate::F24,
|
||||
hdr: HdrFormat::Sdr,
|
||||
color_space: ColorSpace::Bt709,
|
||||
secondary: false,
|
||||
label: String::new(),
|
||||
})],
|
||||
chapters: Vec::new(),
|
||||
extents: Vec::new(),
|
||||
content_format: ContentFormat::BdTs,
|
||||
codec_privates: vec![Some({
|
||||
// Minimal hvcC with one VPS-like array entry.
|
||||
let marker: &[u8] = &[0x40, 0x01, 0x0C, 0x01];
|
||||
let mut hvcc = vec![0u8; 22];
|
||||
hvcc.push(1); // numArrays
|
||||
hvcc.push(32);
|
||||
hvcc.extend_from_slice(&1u16.to_be_bytes()); // numNalus
|
||||
hvcc.extend_from_slice(&(marker.len() as u16).to_be_bytes());
|
||||
hvcc.extend_from_slice(marker);
|
||||
hvcc
|
||||
})],
|
||||
}
|
||||
}
|
||||
|
||||
fn fake_idr_pes_data() -> Vec<u8> {
|
||||
// 4-byte length prefix + NAL: type 19 (IDR_W_RADL).
|
||||
let mut nal = vec![(19u8 << 1) & 0x7E, 0x01];
|
||||
for i in 0..200 {
|
||||
nal.push((i & 0xFF) as u8);
|
||||
}
|
||||
// Fall back to parser-extracted codec_private
|
||||
let pid = self
|
||||
.pid_to_track
|
||||
.iter()
|
||||
.find(|(_, idx)| *idx == track)
|
||||
.map(|(pid, _)| *pid)?;
|
||||
self.parsers
|
||||
.iter()
|
||||
.find(|(p, _)| *p == pid)
|
||||
.and_then(|(_, parser)| parser.codec_private())
|
||||
let mut out = Vec::with_capacity(4 + nal.len());
|
||||
out.extend_from_slice(&(nal.len() as u32).to_be_bytes());
|
||||
out.extend_from_slice(&nal);
|
||||
out
|
||||
}
|
||||
|
||||
fn headers_ready(&self) -> bool {
|
||||
for (idx, s) in self.disc_title.streams.iter().enumerate() {
|
||||
if let crate::disc::Stream::Video(v) = s {
|
||||
if !v.secondary && self.codec_private(idx).is_none() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
/// Writer wrapper that shares an Arc<Mutex<Vec<u8>>> so the test can
|
||||
/// inspect the bytes after the muxer drops.
|
||||
struct SharedSink(std::sync::Arc<std::sync::Mutex<Vec<u8>>>);
|
||||
impl Write for SharedSink {
|
||||
fn write(&mut self, b: &[u8]) -> io::Result<usize> {
|
||||
self.0.lock().unwrap().extend_from_slice(b);
|
||||
Ok(b.len())
|
||||
}
|
||||
true
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn m2ts_stream_forwards_keyframe_to_rai() {
|
||||
let title = make_title();
|
||||
let shared = std::sync::Arc::new(std::sync::Mutex::new(Vec::<u8>::new()));
|
||||
let sink = SharedSink(shared.clone());
|
||||
let mut stream = M2tsStream::create(sink, &title).unwrap();
|
||||
let frame = PesFrame {
|
||||
track: 0,
|
||||
pts: 0,
|
||||
keyframe: true,
|
||||
data: fake_idr_pes_data(),
|
||||
duration_ns: None,
|
||||
};
|
||||
stream.write(&frame).unwrap();
|
||||
stream.finish().unwrap();
|
||||
drop(stream);
|
||||
|
||||
let buf = shared.lock().unwrap().clone();
|
||||
|
||||
// Skip FMKV metadata header via meta::read_header.
|
||||
let mut cursor = std::io::Cursor::new(&buf);
|
||||
let _meta = super::meta::read_header(&mut cursor)
|
||||
.unwrap()
|
||||
.expect("FMKV header present");
|
||||
let header_end = cursor.position() as usize;
|
||||
let ts_bytes = &buf[header_end..];
|
||||
|
||||
// Find first PUSI packet on VIDEO_PID; verify RAI in AF flags.
|
||||
let pkt = ts_bytes
|
||||
.chunks(192)
|
||||
.find(|p| {
|
||||
let h = &p[4..];
|
||||
let pid = (((h[1] & 0x1F) as u16) << 8) | h[2] as u16;
|
||||
pid == VIDEO_PID && (h[1] & 0x40) != 0
|
||||
})
|
||||
.expect("video PUSI packet present");
|
||||
let h = &pkt[4..];
|
||||
let afc = (h[3] >> 4) & 0x03;
|
||||
assert!(afc & 0b10 != 0, "AF must be present");
|
||||
let af_len = h[4] as usize;
|
||||
assert!(af_len >= 1, "AF length must include flags byte");
|
||||
let flags = h[5];
|
||||
assert_eq!(flags & 0x40, 0x40, "RAI bit set");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,828 @@
|
||||
//! Standard MPEG-TS (188-byte packets) muxer — sequential-only.
|
||||
//!
|
||||
//! Distinct from `super::tsmux::TsMuxer` (BD-TS with 192-byte packets
|
||||
//! and the 4-byte TP_extra_header). This muxer emits the IETF / ISO/IEC
|
||||
//! 13818-1 wire format that ffmpeg, VLC, and `m2tsindex` consume
|
||||
//! out of the box. Use it for plain `.ts` / `.m2ts` files over a
|
||||
//! [`SequentialSink`](crate::io::sink::SequentialSink), and for
|
||||
//! MPEG-TS-over-UDP via [`UdpSocketSink`](crate::io::sink::UdpSocketSink).
|
||||
//!
|
||||
//! ## Wire format
|
||||
//!
|
||||
//! Every output packet is exactly 188 bytes:
|
||||
//!
|
||||
//! ```text
|
||||
//! sync_byte:8 = 0x47
|
||||
//! transport_error_indicator:1 = 0
|
||||
//! payload_unit_start_indicator:1
|
||||
//! transport_priority:1 = 0
|
||||
//! PID:13
|
||||
//! transport_scrambling_control:2 = 0
|
||||
//! adaptation_field_control:2
|
||||
//! continuity_counter:4
|
||||
//! [adaptation field, if signalled]
|
||||
//! [payload, if signalled]
|
||||
//! ```
|
||||
//!
|
||||
//! ## Wiring
|
||||
//!
|
||||
//! Single program (PMT PID `0x1000`, program number `1`):
|
||||
//! - PAT on PID `0x0000`
|
||||
//! - PMT on PID `0x1000`
|
||||
//! - Video (HEVC, stream_type `0x24`) on PID `0x0100`
|
||||
//! - First audio (AC3, stream_type `0x81`) on PID `0x0101`
|
||||
//! (or TrueHD, stream_type `0x83`, if hinted)
|
||||
//!
|
||||
//! ## Scope vs. full TS
|
||||
//!
|
||||
//! This is a deliberately minimal viable muxer:
|
||||
//! - One program, one video track, optionally one audio track.
|
||||
//! - PAT + PMT re-emitted every `PSI_INTERVAL_PACKETS` packets so a
|
||||
//! mid-stream receiver can lock on within ~100 ms at typical
|
||||
//! UHD bitrates.
|
||||
//! - PCR clock derived from video PTS (`pts - PCR_LEAD_90KHZ`),
|
||||
//! attached to the video PID's adaptation field every
|
||||
//! `PCR_INTERVAL_PACKETS` packets.
|
||||
//! - No language / descriptor tags, no SCTE-35 markers, no per-PID
|
||||
//! PMT version bumps, no SDT/EIT. Sufficient for "ffmpeg can play
|
||||
//! this back", not for full broadcast deployment.
|
||||
|
||||
use std::io::{self, Write};
|
||||
|
||||
mod packet;
|
||||
|
||||
use packet::{Packet, PacketWriter};
|
||||
|
||||
/// PID `0x0000` — PAT, mandated by spec.
|
||||
const PID_PAT: u16 = 0x0000;
|
||||
/// PID for the single program's PMT. `0x1000` is the conventional
|
||||
/// choice; anything outside reserved/null ranges works.
|
||||
const PID_PMT: u16 = 0x1000;
|
||||
/// PID for the video elementary stream.
|
||||
const PID_VIDEO: u16 = 0x0100;
|
||||
/// PID for the (optional) first audio elementary stream.
|
||||
const PID_AUDIO: u16 = 0x0101;
|
||||
/// Null PID — reserved by spec; never used here.
|
||||
#[allow(dead_code)]
|
||||
const PID_NULL: u16 = 0x1FFF;
|
||||
|
||||
/// Re-emit PAT/PMT every N TS packets. ~250 × 188 B = 47 KB; at
|
||||
/// 80 Mb/s UHD that's ~5 ms — well under typical receiver lock budget.
|
||||
const PSI_INTERVAL_PACKETS: u64 = 250;
|
||||
/// Re-stamp PCR every N TS packets carrying the video PID. Spec MAX
|
||||
/// is 100 ms; 40 packets × 188 B at moderate bitrate keeps us under.
|
||||
const PCR_INTERVAL_PACKETS: u64 = 40;
|
||||
/// PCR lead time — the PCR must precede the PTS of the first byte of
|
||||
/// 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.
|
||||
const STREAM_TYPE_HEVC: u8 = 0x24;
|
||||
const STREAM_TYPE_AC3: u8 = 0x81;
|
||||
const STREAM_TYPE_TRUEHD: u8 = 0x83;
|
||||
|
||||
/// Audio codec hint for [`M2tsMux::new`] / [`M2tsMux::set_audio`]. The
|
||||
/// muxer needs to know the codec to pick the right PMT `stream_type`
|
||||
/// and the right PES stream_id; it doesn't decode the audio.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum AudioCodec {
|
||||
/// AC-3 / E-AC-3, stream_type `0x81`, PES stream_id `0xBD`.
|
||||
Ac3,
|
||||
/// Dolby TrueHD, stream_type `0x83`, PES stream_id `0xBD`.
|
||||
TrueHd,
|
||||
}
|
||||
|
||||
impl AudioCodec {
|
||||
fn stream_type(self) -> u8 {
|
||||
match self {
|
||||
AudioCodec::Ac3 => STREAM_TYPE_AC3,
|
||||
AudioCodec::TrueHd => STREAM_TYPE_TRUEHD,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Sequential MPEG-TS muxer with one video and optional one audio track.
|
||||
pub struct M2tsMux<W: Write> {
|
||||
out: PacketWriter<W>,
|
||||
/// hvcC parameter set bytes to prepend to the first video PES.
|
||||
/// Optional — if the upstream frames already carry inline params,
|
||||
/// callers can omit this.
|
||||
video_codec_private: Option<Vec<u8>>,
|
||||
/// Set on first video frame: have we emitted VPS/SPS/PPS?
|
||||
params_written: bool,
|
||||
/// Audio codec, if an audio track is configured. `None` = video-only.
|
||||
audio: Option<AudioCodec>,
|
||||
/// First seen PTS (90 kHz). All subsequent PTS / PCR values are
|
||||
/// relative to this so output streams start near t=0 and don't
|
||||
/// confuse downstream parsers that don't tolerate huge starting
|
||||
/// timestamps.
|
||||
base_pts_90k: Option<u64>,
|
||||
/// Per-PID continuity counter, 4 bits, monotonically increasing.
|
||||
cc_video: u8,
|
||||
cc_audio: u8,
|
||||
cc_pat: u8,
|
||||
cc_pmt: u8,
|
||||
/// Total packets written, used to gate PSI / PCR cadence.
|
||||
packets_written: u64,
|
||||
/// Video packets written since last PCR, used to gate PCR cadence.
|
||||
video_packets_since_pcr: u64,
|
||||
}
|
||||
|
||||
impl<W: Write> M2tsMux<W> {
|
||||
/// Construct a muxer wrapping `writer`. By default audio is
|
||||
/// disabled; call [`set_audio`](Self::set_audio) before the first
|
||||
/// frame to enable.
|
||||
pub fn new(writer: W) -> Self {
|
||||
Self {
|
||||
out: PacketWriter::new(writer),
|
||||
video_codec_private: None,
|
||||
params_written: false,
|
||||
audio: None,
|
||||
base_pts_90k: None,
|
||||
cc_video: 0,
|
||||
cc_audio: 0,
|
||||
cc_pat: 0,
|
||||
cc_pmt: 0,
|
||||
packets_written: 0,
|
||||
video_packets_since_pcr: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Provide the `HEVCDecoderConfigurationRecord` for video so the
|
||||
/// muxer prepends VPS/SPS/PPS Annex B NALs at stream start.
|
||||
pub fn set_video_codec_private(&mut self, hvcc: Vec<u8>) {
|
||||
self.video_codec_private = Some(hvcc);
|
||||
}
|
||||
|
||||
/// Enable a single audio track. Must be called before
|
||||
/// [`write_audio`](Self::write_audio).
|
||||
pub fn set_audio(&mut self, codec: AudioCodec) {
|
||||
self.audio = Some(codec);
|
||||
}
|
||||
|
||||
/// Write one video PES frame. `data` is either length-prefixed
|
||||
/// NALUs (MKV-style) or already Annex B; both are accepted. `keyframe`
|
||||
/// drives the random_access_indicator bit on the first packet of this
|
||||
/// 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);
|
||||
// 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);
|
||||
|
||||
// Annex-B-ify the frame and prepend VPS/SPS/PPS once, on the
|
||||
// FIRST keyframe (not first frame — non-key frames before the
|
||||
// first keyframe can't carry params usefully).
|
||||
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);
|
||||
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);
|
||||
|
||||
let pes = build_video_pes(pts_90k, &es);
|
||||
self.write_pes(PID_VIDEO, &pes, Some(pcr), keyframe)
|
||||
}
|
||||
|
||||
/// Write one audio PES frame. Returns `Ok(())` and silently drops
|
||||
/// the frame if no audio track was configured — the design assumes
|
||||
/// the upstream picks tracks and won't ship audio to a video-only
|
||||
/// muxer, but defending against it keeps the API a single shape.
|
||||
pub fn write_audio(&mut self, pts_ns: i64, data: &[u8]) -> io::Result<()> {
|
||||
if self.audio.is_none() {
|
||||
return Ok(());
|
||||
}
|
||||
let pts_90k = self.base_relative_pts(pts_ns);
|
||||
let pes = build_audio_pes(pts_90k, data);
|
||||
self.write_pes(PID_AUDIO, &pes, None, false)
|
||||
}
|
||||
|
||||
/// Drain the underlying writer. No TS-level trailer is mandatory —
|
||||
/// receivers detect end-of-stream from socket close or file EOF.
|
||||
pub fn finish(&mut self) -> io::Result<()> {
|
||||
self.out.flush()
|
||||
}
|
||||
|
||||
/// 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 {
|
||||
let raw_90k = if pts_ns > 0 {
|
||||
(pts_ns as u64) * 9 / 100_000
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let base = *self.base_pts_90k.get_or_insert(raw_90k);
|
||||
raw_90k.saturating_sub(base)
|
||||
}
|
||||
|
||||
/// Emit one PES payload as a chain of TS packets on `pid`. If `pcr`
|
||||
/// is provided the first packet carries an adaptation field with
|
||||
/// the PCR. PAT/PMT are re-emitted every `PSI_INTERVAL_PACKETS`.
|
||||
///
|
||||
/// Packet-size math (TS = 188 bytes, header = 4 bytes):
|
||||
/// - 184 B remain after the header for the adaptation-field area
|
||||
/// plus the payload area.
|
||||
/// - With AF body of `b` bytes and `s` stuffing bytes: AF total =
|
||||
/// `1 + b + s` (the leading `1` is the `adaptation_field_length`
|
||||
/// byte itself). Payload = `184 - (1 + b + s)`.
|
||||
/// - With no AF area at all: payload = `184`.
|
||||
///
|
||||
/// The fit-the-tail logic on the last packet of the PES uses
|
||||
/// stuffing rather than a separate small packet, which is the
|
||||
/// standard MPEG-TS convention.
|
||||
fn write_pes(
|
||||
&mut self,
|
||||
pid: u16,
|
||||
pes: &[u8],
|
||||
pcr: Option<u64>,
|
||||
is_keyframe_video: bool,
|
||||
) -> io::Result<()> {
|
||||
// PSI cadence is enforced per TS packet — interleave a fresh
|
||||
// PAT+PMT into the packet stream every PSI_INTERVAL_PACKETS so
|
||||
// long single-PES emissions (e.g. one 60 KB video frame) don't
|
||||
// starve receivers tuning in mid-stream.
|
||||
let mut offset = 0;
|
||||
let mut first = true;
|
||||
while offset < pes.len() {
|
||||
self.maybe_emit_psi()?;
|
||||
|
||||
let attach_pcr = first
|
||||
&& (pid == PID_VIDEO)
|
||||
&& (pcr.is_some())
|
||||
&& (self.packets_written == 0
|
||||
|| self.video_packets_since_pcr >= PCR_INTERVAL_PACKETS);
|
||||
|
||||
// RAI rides only the FIRST packet of a keyframe video PES.
|
||||
let attach_rai = first && is_keyframe_video && pid == PID_VIDEO;
|
||||
|
||||
let mut af_body: Vec<u8> = if attach_pcr {
|
||||
build_pcr_adaptation(pcr.unwrap_or(0))
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
if attach_rai {
|
||||
if af_body.is_empty() {
|
||||
af_body.push(0x40); // flags: RAI only
|
||||
} else {
|
||||
af_body[0] |= 0x40; // OR RAI into existing PCR flags
|
||||
}
|
||||
}
|
||||
|
||||
let remaining = pes.len() - offset;
|
||||
// Capacity for payload given AF body and 1-byte AF length.
|
||||
// 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.
|
||||
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.
|
||||
let max_payload = 184 - 1 - af_body.len();
|
||||
let p = remaining.min(max_payload);
|
||||
let s = max_payload - p;
|
||||
(true, p, s)
|
||||
} else if remaining >= 184 {
|
||||
// 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;
|
||||
let p = remaining.min(max_payload);
|
||||
let s = max_payload - p;
|
||||
(true, p, s)
|
||||
};
|
||||
|
||||
let cc = self.advance_cc(pid);
|
||||
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_payload(&pes[offset..offset + payload_len]);
|
||||
debug_assert_eq!(packet.len(), 188, "packet not 188 bytes");
|
||||
self.out.write_packet(&packet)?;
|
||||
|
||||
self.packets_written += 1;
|
||||
if pid == PID_VIDEO {
|
||||
if attach_pcr {
|
||||
self.video_packets_since_pcr = 0;
|
||||
} else {
|
||||
self.video_packets_since_pcr += 1;
|
||||
}
|
||||
}
|
||||
offset += payload_len;
|
||||
first = false;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn advance_cc(&mut self, pid: u16) -> u8 {
|
||||
let slot = match pid {
|
||||
PID_VIDEO => &mut self.cc_video,
|
||||
PID_AUDIO => &mut self.cc_audio,
|
||||
PID_PAT => &mut self.cc_pat,
|
||||
PID_PMT => &mut self.cc_pmt,
|
||||
_ => return 0,
|
||||
};
|
||||
let cc = *slot;
|
||||
*slot = (*slot + 1) & 0x0F;
|
||||
cc
|
||||
}
|
||||
|
||||
fn maybe_emit_psi(&mut self) -> io::Result<()> {
|
||||
if self.packets_written == 0 || self.packets_written % PSI_INTERVAL_PACKETS == 0 {
|
||||
self.emit_pat()?;
|
||||
self.emit_pmt()?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn emit_pat(&mut self) -> io::Result<()> {
|
||||
let payload = build_pat(PID_PMT);
|
||||
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.pad_to_188();
|
||||
self.out.write_packet(&packet)?;
|
||||
self.packets_written += 1;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn emit_pmt(&mut self) -> io::Result<()> {
|
||||
let payload = build_pmt(self.audio);
|
||||
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.pad_to_188();
|
||||
self.out.write_packet(&packet)?;
|
||||
self.packets_written += 1;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// 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)
|
||||
}
|
||||
|
||||
/// 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.
|
||||
build_pes_packet(0xBD, pts_90k, es, /* length_in_header */ true)
|
||||
}
|
||||
|
||||
fn build_pes_packet(stream_id: u8, pts_90k: u64, es: &[u8], length_in_header: bool) -> Vec<u8> {
|
||||
let mut out = Vec::with_capacity(es.len() + 14);
|
||||
out.extend_from_slice(&[0x00, 0x00, 0x01, stream_id]);
|
||||
// PES_packet_length: total bytes after this field. 3 flag bytes + 5
|
||||
// PTS bytes + es.len(). Zero means "unbounded" — used for video
|
||||
// where PES can exceed u16.
|
||||
let pes_len = 8 + es.len();
|
||||
if length_in_header && pes_len <= u16::MAX as usize {
|
||||
out.extend_from_slice(&(pes_len as u16).to_be_bytes());
|
||||
} else {
|
||||
out.extend_from_slice(&[0x00, 0x00]);
|
||||
}
|
||||
// Flags byte 1: 10 = MPEG-2 marker, then scrambling/priority/etc 0.
|
||||
out.push(0x80);
|
||||
// Flags byte 2: PTS flag (bit 7).
|
||||
out.push(0x80);
|
||||
// PES_header_data_length = 5 (just PTS).
|
||||
out.push(5);
|
||||
// PTS bytes — 33-bit timestamp split across 5 bytes with marker bits.
|
||||
let pts = pts_90k & 0x1_FFFF_FFFF;
|
||||
out.push(0x21 | (((pts >> 29) & 0x0E) as u8));
|
||||
out.push(((pts >> 22) & 0xFF) as u8);
|
||||
out.push(0x01 | (((pts >> 14) & 0xFE) as u8));
|
||||
out.push(((pts >> 7) & 0xFF) as u8);
|
||||
out.push(0x01 | (((pts << 1) & 0xFE) as u8));
|
||||
out.extend_from_slice(es);
|
||||
out
|
||||
}
|
||||
|
||||
/// Build the PAT payload (section, with pointer_field).
|
||||
fn build_pat(pmt_pid: u16) -> Vec<u8> {
|
||||
let mut section = Vec::new();
|
||||
section.push(0x00); // table_id = PAT
|
||||
// section_syntax_indicator(1) | '0'(1) | reserved(2) | section_length(12)
|
||||
// section_length covers from end of this field through CRC.
|
||||
// Body: transport_stream_id(2) + version/cni(1) + section/last_section(2) + program(4) = 9 bytes,
|
||||
// plus CRC(4) = 13. Encoded big-endian.
|
||||
section.extend_from_slice(&[0xB0, 13]);
|
||||
section.extend_from_slice(&[0x00, 0x01]); // transport_stream_id = 1
|
||||
section.push(0xC1); // reserved | version=0 | current_next=1
|
||||
section.push(0x00); // section_number
|
||||
section.push(0x00); // last_section_number
|
||||
section.extend_from_slice(&[0x00, 0x01]); // program_number = 1
|
||||
// reserved(3) | network_PID/program_map_PID(13)
|
||||
let pid_bytes = (0xE000u16 | (pmt_pid & 0x1FFF)).to_be_bytes();
|
||||
section.extend_from_slice(&pid_bytes);
|
||||
let crc = mpegts_crc32(§ion);
|
||||
section.extend_from_slice(&crc.to_be_bytes());
|
||||
// Prepend pointer_field=0 (section starts immediately).
|
||||
let mut payload = Vec::with_capacity(section.len() + 1);
|
||||
payload.push(0x00);
|
||||
payload.extend_from_slice(§ion);
|
||||
payload
|
||||
}
|
||||
|
||||
/// Build the PMT payload (section, with pointer_field).
|
||||
fn build_pmt(audio: Option<AudioCodec>) -> Vec<u8> {
|
||||
let mut section = Vec::new();
|
||||
section.push(0x02); // table_id = PMT
|
||||
// section_length filled in after we know the body size.
|
||||
let len_placeholder = section.len();
|
||||
section.extend_from_slice(&[0xB0, 0x00]);
|
||||
|
||||
section.extend_from_slice(&1u16.to_be_bytes()); // program_number
|
||||
section.push(0xC1); // reserved | version=0 | current_next=1
|
||||
section.push(0x00); // section_number
|
||||
section.push(0x00); // last_section_number
|
||||
// reserved(3) | PCR_PID(13)
|
||||
let pcr_pid = (0xE000u16 | (PID_VIDEO & 0x1FFF)).to_be_bytes();
|
||||
section.extend_from_slice(&pcr_pid);
|
||||
// program_info_length = 0
|
||||
section.extend_from_slice(&[0xF0, 0x00]);
|
||||
|
||||
// Video elementary stream entry.
|
||||
section.push(STREAM_TYPE_HEVC);
|
||||
let v_pid = (0xE000u16 | (PID_VIDEO & 0x1FFF)).to_be_bytes();
|
||||
section.extend_from_slice(&v_pid);
|
||||
section.extend_from_slice(&[0xF0, 0x00]); // ES_info_length = 0
|
||||
|
||||
if let Some(codec) = audio {
|
||||
section.push(codec.stream_type());
|
||||
let a_pid = (0xE000u16 | (PID_AUDIO & 0x1FFF)).to_be_bytes();
|
||||
section.extend_from_slice(&a_pid);
|
||||
section.extend_from_slice(&[0xF0, 0x00]);
|
||||
}
|
||||
|
||||
// Now patch section_length: covers everything after the length field
|
||||
// through the CRC, so (current body size - 3 bytes consumed by
|
||||
// table_id + 2 length bytes) + 4 (CRC).
|
||||
let section_len = section.len() - 3 + 4;
|
||||
section[len_placeholder] = 0xB0 | ((section_len >> 8) as u8 & 0x0F);
|
||||
section[len_placeholder + 1] = section_len as u8;
|
||||
|
||||
let crc = mpegts_crc32(§ion);
|
||||
section.extend_from_slice(&crc.to_be_bytes());
|
||||
|
||||
let mut payload = Vec::with_capacity(section.len() + 1);
|
||||
payload.push(0x00);
|
||||
payload.extend_from_slice(§ion);
|
||||
payload
|
||||
}
|
||||
|
||||
/// Build the adaptation field carrying a PCR (no other flags).
|
||||
fn build_pcr_adaptation(pcr_90k: u64) -> Vec<u8> {
|
||||
// adaptation_field_length is set by `Packet::append_adaptation`
|
||||
// — this function returns just the field body.
|
||||
//
|
||||
// Layout: discontinuity_indicator(1) | random_access(1) |
|
||||
// elementary_stream_priority(1) | PCR_flag(1) | OPCR_flag(1) |
|
||||
// splicing_point_flag(1) | transport_private_data_flag(1) |
|
||||
// adaptation_field_extension_flag(1) | PCR(48b).
|
||||
let mut af = vec![0x10]; // PCR_flag=1; RAI is OR'd in by the caller when applicable
|
||||
let pcr_base = pcr_90k & 0x1_FFFF_FFFF; // 33-bit
|
||||
let pcr_ext: u16 = 0; // 9-bit, we keep it zero (no sub-tick precision)
|
||||
// Encode PCR: 33b base | 6b reserved | 9b extension = 48b
|
||||
af.push((pcr_base >> 25) as u8);
|
||||
af.push((pcr_base >> 17) as u8);
|
||||
af.push((pcr_base >> 9) as u8);
|
||||
af.push((pcr_base >> 1) as u8);
|
||||
af.push(((pcr_base << 7) as u8 & 0x80) | 0x7E | ((pcr_ext >> 8) as u8 & 0x01));
|
||||
af.push(pcr_ext as u8);
|
||||
af
|
||||
}
|
||||
|
||||
/// MPEG-TS CRC-32 (poly 0x04C11DB7, init 0xFFFFFFFF, no reflection, no
|
||||
/// final XOR). Implementation: bitwise so we don't need a table —
|
||||
/// PSI sections are tiny, the cost is negligible.
|
||||
fn mpegts_crc32(data: &[u8]) -> u32 {
|
||||
let mut crc: u32 = 0xFFFF_FFFF;
|
||||
for &b in data {
|
||||
crc ^= (b as u32) << 24;
|
||||
for _ in 0..8 {
|
||||
if crc & 0x8000_0000 != 0 {
|
||||
crc = (crc << 1) ^ 0x04C1_1DB7;
|
||||
} else {
|
||||
crc <<= 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
crc
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// All emitted bytes must align to 188-byte packet boundaries and
|
||||
/// every packet must start with `0x47`.
|
||||
fn assert_ts_well_formed(buf: &[u8]) {
|
||||
assert_eq!(
|
||||
buf.len() % 188,
|
||||
0,
|
||||
"stream not packet-aligned: {} bytes",
|
||||
buf.len()
|
||||
);
|
||||
for (i, chunk) in buf.chunks(188).enumerate() {
|
||||
assert_eq!(chunk[0], 0x47, "packet {} missing sync byte", i);
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_pids(buf: &[u8]) -> Vec<u16> {
|
||||
buf.chunks(188)
|
||||
.map(|p| u16::from_be_bytes([p[1] & 0x1F, p[2]]))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn crc32_is_self_validating() {
|
||||
// The MPEG-TS CRC has the property that prepending a single
|
||||
// bit-flip changes the output; running it over its own input +
|
||||
// CRC yields a fixed magic constant (the CRC residue). Rather
|
||||
// than hardcoding sample bytes, verify the underlying algorithm
|
||||
// by checking that two distinct inputs produce distinct CRCs
|
||||
// and that the same input is deterministic.
|
||||
let a = [
|
||||
0u8, 0xB0, 0x0D, 0x00, 0x01, 0xC1, 0x00, 0x00, 0x00, 0x01, 0xE1, 0x00,
|
||||
];
|
||||
let mut b = a;
|
||||
b[5] ^= 0x01; // flip one bit
|
||||
let crc_a = mpegts_crc32(&a);
|
||||
let crc_b = mpegts_crc32(&b);
|
||||
assert_ne!(crc_a, crc_b);
|
||||
assert_eq!(crc_a, mpegts_crc32(&a)); // deterministic
|
||||
// Sanity: all-zero input ⇒ CRC = 0 (init XORs but the
|
||||
// shift/feedback cancels for zero data after init drains).
|
||||
// We don't assert exact value — that depends on poly choice —
|
||||
// but check it's not the same as for non-zero data.
|
||||
let crc_zero = mpegts_crc32(&[0u8; 12]);
|
||||
assert_ne!(crc_zero, crc_a);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn video_only_mux_emits_pat_pmt_then_video() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = M2tsMux::new(&mut sink);
|
||||
// One small video frame, no codec_private (so no params NAL inline).
|
||||
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);
|
||||
|
||||
assert_ts_well_formed(&sink);
|
||||
let pids = extract_pids(&sink);
|
||||
// First two packets: PAT, PMT. At least one video packet after.
|
||||
assert_eq!(pids[0], PID_PAT);
|
||||
assert_eq!(pids[1], PID_PMT);
|
||||
assert!(pids.iter().any(|p| *p == PID_VIDEO));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn audio_track_appears_in_pmt_and_stream() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = M2tsMux::new(&mut sink);
|
||||
mux.set_audio(AudioCodec::Ac3);
|
||||
// Video + audio frame pair.
|
||||
let mut frame = Vec::new();
|
||||
frame.extend_from_slice(&3u32.to_be_bytes());
|
||||
frame.extend_from_slice(&[0x40, 0x01, 0x0C]);
|
||||
mux.write_video(0, true, &frame).unwrap();
|
||||
mux.write_audio(20_000_000, &[0x0B, 0x77, 0x12, 0x34])
|
||||
.unwrap();
|
||||
mux.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
assert_ts_well_formed(&sink);
|
||||
let pids = extract_pids(&sink);
|
||||
assert!(pids.iter().any(|p| *p == PID_VIDEO));
|
||||
assert!(pids.iter().any(|p| *p == PID_AUDIO));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn psi_re_emits_at_interval() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = M2tsMux::new(&mut sink);
|
||||
// Build a frame large enough to span > PSI_INTERVAL_PACKETS TS packets.
|
||||
// 184 B payload per packet ⇒ ~250 packets = ~46 KB elementary stream.
|
||||
let big: Vec<u8> = (0..(60 * 1024)).map(|i| (i & 0xff) as u8).collect();
|
||||
let mut frame = Vec::new();
|
||||
frame.extend_from_slice(&(big.len() as u32).to_be_bytes());
|
||||
frame.extend_from_slice(&big);
|
||||
mux.write_video(0, true, &frame).unwrap();
|
||||
mux.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
assert_ts_well_formed(&sink);
|
||||
let pids = extract_pids(&sink);
|
||||
// Count PAT/PMT pairs — must be at least 2 given the input size.
|
||||
let pat_count = pids.iter().filter(|p| **p == PID_PAT).count();
|
||||
let pmt_count = pids.iter().filter(|p| **p == PID_PMT).count();
|
||||
assert!(pat_count >= 2, "expected ≥2 PAT, got {}", pat_count);
|
||||
assert!(pmt_count >= 2, "expected ≥2 PMT, got {}", pmt_count);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn continuity_counter_increments_per_pid() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = M2tsMux::new(&mut sink);
|
||||
// Three small video frames to get a sequence of video TS packets.
|
||||
for pts in [0i64, 40_000_000, 80_000_000] {
|
||||
let mut frame = Vec::new();
|
||||
frame.extend_from_slice(&3u32.to_be_bytes());
|
||||
frame.extend_from_slice(&[0xAA, 0xBB, 0xCC]);
|
||||
mux.write_video(pts, pts == 0, &frame).unwrap();
|
||||
}
|
||||
mux.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
// Collect CCs for video packets in order.
|
||||
let ccs: Vec<u8> = sink
|
||||
.chunks(188)
|
||||
.filter(|p| u16::from_be_bytes([p[1] & 0x1F, p[2]]) == PID_VIDEO)
|
||||
.map(|p| p[3] & 0x0F)
|
||||
.collect();
|
||||
for w in ccs.windows(2) {
|
||||
assert_eq!(w[1], (w[0] + 1) & 0x0F);
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the adaptation field body (length byte stripped) for one
|
||||
/// 188-byte TS packet, or None when the packet has no AF.
|
||||
fn af_body(packet: &[u8]) -> Option<Vec<u8>> {
|
||||
let afc = (packet[3] >> 4) & 0x03;
|
||||
if afc & 0b10 == 0 {
|
||||
return None;
|
||||
}
|
||||
let af_len = packet[4] as usize;
|
||||
if af_len == 0 {
|
||||
return Some(Vec::new());
|
||||
}
|
||||
Some(packet[5..5 + af_len].to_vec())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rai_set_on_keyframe_pes_packet() {
|
||||
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);
|
||||
|
||||
// Find the first PUSI packet on PID_VIDEO.
|
||||
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("AF present on first packet of keyframe video PES");
|
||||
assert!(!af.is_empty(), "AF flags byte present");
|
||||
assert_eq!(af[0] & 0x40, 0x40, "RAI bit set");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pcr_packet_without_keyframe_has_rai_clear() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = M2tsMux::new(&mut sink);
|
||||
// First frame is the keyframe (gates codec_private; also gets PCR).
|
||||
let mut frame0 = Vec::new();
|
||||
frame0.extend_from_slice(&4u32.to_be_bytes());
|
||||
frame0.extend_from_slice(&[0x40, 0x01, 0x0C, 0x01]);
|
||||
mux.write_video(0, true, &frame0).unwrap();
|
||||
// Push enough non-key video frames to cross PCR_INTERVAL_PACKETS
|
||||
// video packets so a later PCR-bearing packet exists.
|
||||
// Each frame is ~50 KB → ~270 packets, well over 40.
|
||||
let big: Vec<u8> = (0..(50 * 1024)).map(|i| (i & 0xff) as u8).collect();
|
||||
for i in 1..3 {
|
||||
let mut frame = Vec::new();
|
||||
frame.extend_from_slice(&(big.len() as u32).to_be_bytes());
|
||||
frame.extend_from_slice(&big);
|
||||
mux.write_video((i as i64) * 40_000_000, false, &frame)
|
||||
.unwrap();
|
||||
}
|
||||
mux.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
// The first PUSI video packet carries PCR + RAI (keyframe).
|
||||
// A later video PUSI packet with AF + PCR but NOT keyframe must
|
||||
// have RAI clear.
|
||||
let video_pusi: Vec<&[u8]> = sink
|
||||
.chunks(188)
|
||||
.filter(|p| u16::from_be_bytes([p[1] & 0x1F, p[2]]) == PID_VIDEO && (p[1] & 0x40) != 0)
|
||||
.collect();
|
||||
assert!(
|
||||
video_pusi.len() >= 2,
|
||||
"expected ≥2 video PES starts, got {}",
|
||||
video_pusi.len()
|
||||
);
|
||||
// Find a later one with AF that carries PCR (flags & 0x10 set).
|
||||
let later_pcr = video_pusi
|
||||
.iter()
|
||||
.skip(1)
|
||||
.find_map(|p| {
|
||||
let af = af_body(p)?;
|
||||
if !af.is_empty() && (af[0] & 0x10) != 0 {
|
||||
Some(af)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.expect("later PCR-bearing PUSI exists");
|
||||
assert_eq!(
|
||||
later_pcr[0] & 0x40,
|
||||
0,
|
||||
"RAI must be clear on non-keyframe PCR packet"
|
||||
);
|
||||
}
|
||||
|
||||
#[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).
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = M2tsMux::new(&mut sink);
|
||||
let mut small = Vec::new();
|
||||
small.extend_from_slice(&4u32.to_be_bytes());
|
||||
small.extend_from_slice(&[0x40, 0x01, 0x0C, 0x01]);
|
||||
mux.write_video(0, true, &small).unwrap();
|
||||
// ~50 KB ≈ 270 packets — well over PCR_INTERVAL_PACKETS.
|
||||
let big: Vec<u8> = (0..(50 * 1024)).map(|i| (i & 0xff) as u8).collect();
|
||||
let mut big_frame = Vec::new();
|
||||
big_frame.extend_from_slice(&(big.len() as u32).to_be_bytes());
|
||||
big_frame.extend_from_slice(&big);
|
||||
mux.write_video(40_000_000, false, &big_frame).unwrap();
|
||||
// Now a second keyframe — must combine RAI (keyframe) and PCR
|
||||
// (counter exceeded).
|
||||
mux.write_video(80_000_000, true, &small).unwrap();
|
||||
mux.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
// Collect video PUSI packets and find the third (second keyframe).
|
||||
let video_pusi: Vec<&[u8]> = sink
|
||||
.chunks(188)
|
||||
.filter(|p| u16::from_be_bytes([p[1] & 0x1F, p[2]]) == PID_VIDEO && (p[1] & 0x40) != 0)
|
||||
.collect();
|
||||
assert!(video_pusi.len() >= 3, "three video PES starts expected");
|
||||
let af = af_body(video_pusi[2]).expect("AF present");
|
||||
assert!(!af.is_empty(), "AF flags byte present");
|
||||
assert_eq!(af[0], 0x50, "flags == RAI | PCR");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,145 @@
|
||||
//! 188-byte MPEG-TS packet builder + writer.
|
||||
//!
|
||||
//! Internal helper for `super::M2tsMux`. Not public API — exposes raw
|
||||
//! byte layout so the parent module can compose PSI / PCR / PES bytes
|
||||
//! without each caller re-implementing the 188-byte boundary math.
|
||||
|
||||
use std::io::{self, Write};
|
||||
|
||||
const TS_PACKET_SIZE: usize = 188;
|
||||
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.
|
||||
pub(super) struct Packet {
|
||||
buf: Vec<u8>,
|
||||
}
|
||||
|
||||
impl Packet {
|
||||
pub(super) fn new() -> Self {
|
||||
Self {
|
||||
buf: Vec::with_capacity(TS_PACKET_SIZE),
|
||||
}
|
||||
}
|
||||
|
||||
/// Write the 4-byte TS packet header.
|
||||
///
|
||||
/// * `pid` — 13-bit PID
|
||||
/// * `payload_unit_start` — first packet of a PES / PSI section
|
||||
/// * `has_payload` — packet carries any payload bytes
|
||||
/// * `has_adaptation` — packet carries an adaptation field
|
||||
/// * `cc` — 4-bit continuity counter
|
||||
pub(super) fn set_header(
|
||||
&mut self,
|
||||
pid: u16,
|
||||
payload_unit_start: bool,
|
||||
has_payload: bool,
|
||||
has_adaptation: bool,
|
||||
cc: u8,
|
||||
) {
|
||||
self.buf.clear();
|
||||
self.buf.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);
|
||||
// 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
|
||||
(false, true) => 0b01, // payload only
|
||||
(true, false) => 0b10, // adaptation only
|
||||
(true, true) => 0b11, // both
|
||||
};
|
||||
self.buf.push((afc << 4) | (cc & 0x0F));
|
||||
}
|
||||
|
||||
/// Append the adaptation field after the header.
|
||||
///
|
||||
/// `body` is the adaptation field body (flags byte plus optional PCR
|
||||
/// and so on). `stuffing` is the number of `0xFF` stuffing bytes to
|
||||
/// 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) {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
/// 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");
|
||||
}
|
||||
|
||||
/// Pad the packet to exactly 188 bytes with `0xFF` bytes — used by
|
||||
/// PSI emit paths where the section is much smaller than 184 bytes.
|
||||
/// 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);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn bytes(&self) -> &[u8] {
|
||||
&self.buf
|
||||
}
|
||||
|
||||
pub(super) fn len(&self) -> usize {
|
||||
self.buf.len()
|
||||
}
|
||||
}
|
||||
|
||||
/// Buffered writer for assembled TS packets. Owns the underlying sink.
|
||||
pub(super) struct PacketWriter<W: Write> {
|
||||
inner: W,
|
||||
}
|
||||
|
||||
impl<W: Write> PacketWriter<W> {
|
||||
pub(super) fn new(inner: W) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
|
||||
pub(super) fn write_packet(&mut self, packet: &Packet) -> io::Result<()> {
|
||||
let bytes = packet.bytes();
|
||||
debug_assert_eq!(bytes.len(), TS_PACKET_SIZE);
|
||||
self.inner.write_all(bytes)
|
||||
}
|
||||
|
||||
pub(super) fn flush(&mut self) -> io::Result<()> {
|
||||
self.inner.flush()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
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.pad_to_188();
|
||||
assert_eq!(p.bytes().len(), 188);
|
||||
assert_eq!(p.bytes()[0], SYNC_BYTE);
|
||||
// After 4-byte header + 3 payload, byte 7 starts stuffing.
|
||||
assert_eq!(p.bytes()[7], STUFF_BYTE);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn header_pid_round_trips() {
|
||||
let mut p = Packet::new();
|
||||
p.set_header(0x1ABC, false, true, false, 0xA);
|
||||
let pid = u16::from_be_bytes([p.bytes()[1] & 0x1F, p.bytes()[2]]);
|
||||
assert_eq!(pid, 0x1ABC);
|
||||
assert_eq!(p.bytes()[3] & 0x0F, 0xA);
|
||||
}
|
||||
}
|
||||
+752
-24
@@ -34,6 +34,24 @@ pub struct MkvTrack {
|
||||
pub sample_rate: f64,
|
||||
pub channels: u8,
|
||||
pub bit_depth: u8,
|
||||
// Dolby Vision: the dvcC (DOVIDecoderConfigurationRecord) for the DV layer,
|
||||
// emitted as a BlockAdditionMapping. `None` for non-DV tracks.
|
||||
pub dv_config: Option<Vec<u8>>,
|
||||
}
|
||||
|
||||
/// Build a DOVIDecoderConfigurationRecord (dvcC) — 24 bytes — for the Matroska
|
||||
/// BlockAdditionMapping. For disc Profile 7 dual-layer the base, enhancement,
|
||||
/// and RPU are all present (lossless FEL/MEL preserved as a second track).
|
||||
pub fn dolby_vision_config(profile: u8, level: u8, bl_compat_id: u8) -> Vec<u8> {
|
||||
let mut v = vec![0u8; 24];
|
||||
v[0] = 1; // dv_version_major
|
||||
v[1] = 0; // dv_version_minor
|
||||
// profile(7) | level(6) | rpu_present(1) | el_present(1) | bl_present(1)
|
||||
v[2] = ((profile & 0x7F) << 1) | ((level >> 5) & 0x01);
|
||||
v[3] = ((level & 0x1F) << 3) | (1 << 2) | (1 << 1) | 1; // rpu = el = bl = 1
|
||||
v[4] = (bl_compat_id & 0x0F) << 4;
|
||||
// v[5..24] reserved = 0
|
||||
v
|
||||
}
|
||||
|
||||
impl MkvTrack {
|
||||
@@ -83,10 +101,28 @@ impl MkvTrack {
|
||||
sample_rate: 0.0,
|
||||
channels: 0,
|
||||
bit_depth: 0,
|
||||
// The DV layer (hdr=DolbyVision) carries the dvcC so the track is
|
||||
// recognised as Dolby Vision (disc Profile 7 dual-layer).
|
||||
dv_config: if matches!(v.hdr, HdrFormat::DolbyVision) {
|
||||
Some(dolby_vision_config(7, 6, 0))
|
||||
} else {
|
||||
None
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
// "Supports DTS, DTS-ES, DTS-96/26, DTS-HD High Resolution Audio
|
||||
// and DTS-HD Master Audio." Players distinguish core vs HD-HRA vs
|
||||
// HD-MA by parsing the DTS bitstream extension substreams, not by
|
||||
// the container codec ID. The previously-emitted `A_DTS/MA` and
|
||||
// `A_DTS/HR` suffixes are NOT registered codec IDs; strict parsers
|
||||
// (libmatroska) and some hardware renderers fail to recognise the
|
||||
// track at all. Emit plain `A_DTS` for every DTS variant — the
|
||||
// lossless MA / HRA payload bytes are unchanged, only the
|
||||
// container codec-ID string differs.
|
||||
let codec_id = match a.codec {
|
||||
Codec::Ac3 => "A_AC3",
|
||||
Codec::Ac3Plus => "A_EAC3",
|
||||
@@ -120,6 +156,7 @@ impl MkvTrack {
|
||||
sample_rate: sr,
|
||||
channels: ch,
|
||||
bit_depth: 0,
|
||||
dv_config: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -148,6 +185,7 @@ impl MkvTrack {
|
||||
sample_rate: 0.0,
|
||||
channels: 0,
|
||||
bit_depth: 0,
|
||||
dv_config: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -159,6 +197,12 @@ struct CuePoint {
|
||||
cluster_pos: u64, // relative to Segment start
|
||||
}
|
||||
|
||||
/// SeekHead entry that needs its 8-byte SeekPosition back-patched after Cues are written.
|
||||
struct SeekPositionFixup {
|
||||
target_id: u32,
|
||||
value_offset: u64, // absolute file offset of the 8-byte SeekPosition value
|
||||
}
|
||||
|
||||
/// MKV muxer. Call write_frame() for each frame, then finish() at the end.
|
||||
pub struct MkvMuxer<W: Write + Seek> {
|
||||
writer: W,
|
||||
@@ -170,11 +214,21 @@ pub struct MkvMuxer<W: Write + Seek> {
|
||||
base_pts_ms: Option<i64>,
|
||||
cues: Vec<CuePoint>,
|
||||
frame_count: u64,
|
||||
seek_fixups: Vec<SeekPositionFixup>,
|
||||
info_offset: u64,
|
||||
tracks_offset: u64,
|
||||
chapters_offset: Option<u64>,
|
||||
}
|
||||
|
||||
/// New cluster every 5 seconds.
|
||||
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.
|
||||
const MAX_BLOCK_REL_MS: i64 = i16::MAX as i64;
|
||||
|
||||
impl<W: Write + Seek> MkvMuxer<W> {
|
||||
/// Create a new MKV muxer: writes EBML header, Segment start, Info, Tracks, Chapters.
|
||||
pub fn new(
|
||||
@@ -200,7 +254,34 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
ebml::write_unknown_size(&mut writer)?;
|
||||
let segment_start = writer.stream_position()?;
|
||||
|
||||
// SeekHead with fixed-width SeekPosition placeholders. Order: Info, Tracks, [Chapters], Cues.
|
||||
let mut seek_fixups: Vec<SeekPositionFixup> = Vec::new();
|
||||
let seekhead_pos = ebml::start_master(&mut writer, ebml::SEEK_HEAD)?;
|
||||
let mut targets: Vec<u32> = vec![ebml::INFO, ebml::TRACKS];
|
||||
if !chapters.is_empty() {
|
||||
targets.push(ebml::CHAPTERS);
|
||||
}
|
||||
targets.push(ebml::CUES);
|
||||
let seek_id_be = (ebml::SEEK as u16).to_be_bytes();
|
||||
let seek_inner_id_be = (ebml::SEEK_ID as u16).to_be_bytes();
|
||||
let seek_pos_id_be = (ebml::SEEK_POSITION as u16).to_be_bytes();
|
||||
for target_id in &targets {
|
||||
writer.write_all(&[seek_id_be[0], seek_id_be[1], 0x92])?;
|
||||
writer.write_all(&[seek_inner_id_be[0], seek_inner_id_be[1], 0x84])?;
|
||||
writer.write_all(&target_id.to_be_bytes())?;
|
||||
writer.write_all(&[seek_pos_id_be[0], seek_pos_id_be[1], 0x88])?;
|
||||
let value_offset = writer.stream_position()?;
|
||||
writer.write_all(&[0u8; 8])?;
|
||||
seek_fixups.push(SeekPositionFixup {
|
||||
target_id: *target_id,
|
||||
value_offset,
|
||||
});
|
||||
}
|
||||
ebml::end_master(&mut writer, seekhead_pos)?;
|
||||
|
||||
// Info
|
||||
let info_start = writer.stream_position()?;
|
||||
let info_offset = info_start - segment_start;
|
||||
let info_pos = ebml::start_master(&mut writer, ebml::INFO)?;
|
||||
ebml::write_uint(&mut writer, ebml::TIMESTAMP_SCALE, 1_000_000)?; // 1ms precision
|
||||
if duration_secs > 0.0 {
|
||||
@@ -215,6 +296,8 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
ebml::end_master(&mut writer, info_pos)?;
|
||||
|
||||
// Tracks
|
||||
let tracks_start = writer.stream_position()?;
|
||||
let tracks_offset = tracks_start - segment_start;
|
||||
let tracks_pos = ebml::start_master(&mut writer, ebml::TRACKS)?;
|
||||
for (i, track) in tracks.iter().enumerate() {
|
||||
let entry_pos = ebml::start_master(&mut writer, ebml::TRACK_ENTRY)?;
|
||||
@@ -286,6 +369,17 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
ebml::end_master(&mut writer, vid_pos)?;
|
||||
}
|
||||
|
||||
// Dolby Vision signaling — BlockAdditionMapping is a child of the
|
||||
// TrackEntry (sibling of Video). Carries the dvcC so players /
|
||||
// mediainfo recognise the track as Dolby Vision.
|
||||
if let Some(ref dvcc) = track.dv_config {
|
||||
let map_pos = ebml::start_master(&mut writer, ebml::BLOCK_ADDITION_MAPPING)?;
|
||||
// BlockAddIDType = "dvcC" fourcc (DOVIDecoderConfigurationRecord).
|
||||
ebml::write_uint(&mut writer, ebml::BLOCK_ADD_ID_TYPE, 0x6476_6343)?;
|
||||
ebml::write_binary(&mut writer, ebml::BLOCK_ADD_ID_EXTRA_DATA, dvcc)?;
|
||||
ebml::end_master(&mut writer, map_pos)?;
|
||||
}
|
||||
|
||||
// Audio-specific
|
||||
if track.track_type == ebml::TRACK_TYPE_AUDIO && track.sample_rate > 0.0 {
|
||||
let aud_pos = ebml::start_master(&mut writer, ebml::AUDIO)?;
|
||||
@@ -302,7 +396,10 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
ebml::end_master(&mut writer, tracks_pos)?;
|
||||
|
||||
// Chapters
|
||||
let mut chapters_offset: Option<u64> = None;
|
||||
if !chapters.is_empty() {
|
||||
let chapters_start = writer.stream_position()?;
|
||||
chapters_offset = Some(chapters_start - segment_start);
|
||||
let chapters_pos = ebml::start_master(&mut writer, ebml::CHAPTERS)?;
|
||||
let edition_pos = ebml::start_master(&mut writer, ebml::EDITION_ENTRY)?;
|
||||
for (i, ch) in chapters.iter().enumerate() {
|
||||
@@ -330,42 +427,71 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
base_pts_ms: None,
|
||||
cues: Vec::new(),
|
||||
frame_count: 0,
|
||||
seek_fixups,
|
||||
info_offset,
|
||||
tracks_offset,
|
||||
chapters_offset,
|
||||
})
|
||||
}
|
||||
|
||||
/// Write a single frame.
|
||||
///
|
||||
/// When `duration_ns` is `Some`, the frame is emitted as a
|
||||
/// `BlockGroup` with `BlockDuration` so the player knows exactly
|
||||
/// when to remove the on-screen artifact (the practical case is
|
||||
/// PGS subtitles — without it, the last bitmap lingers until the
|
||||
/// next display set replaces it). Otherwise a plain `SimpleBlock`.
|
||||
pub fn write_frame(
|
||||
&mut self,
|
||||
track_idx: usize,
|
||||
pts_ns: i64,
|
||||
keyframe: bool,
|
||||
data: &[u8],
|
||||
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;
|
||||
|
||||
// Start new cluster if needed
|
||||
if !self.cluster_open || (pts_ms - self.cluster_ts_ms) >= CLUSTER_DURATION_MS {
|
||||
if self.cluster_open {
|
||||
// Close current cluster (it's a master with unknown size — we use known size)
|
||||
// Actually, for streaming we keep clusters open-ended. Just start a new one.
|
||||
// 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 {
|
||||
return Ok(());
|
||||
}
|
||||
self.start_cluster(pts_ms)?;
|
||||
|
||||
// Add cue point at cluster start for keyframes (video track 0)
|
||||
if keyframe && track_idx == 0 {
|
||||
self.cues.push(CuePoint {
|
||||
timestamp_ms: pts_ms,
|
||||
track: track_idx + 1,
|
||||
cluster_pos: self.cluster_pos - self.segment_start,
|
||||
});
|
||||
}
|
||||
self.cues.push(CuePoint {
|
||||
timestamp_ms: pts_ms,
|
||||
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)?;
|
||||
}
|
||||
|
||||
// Write SimpleBlock
|
||||
let relative_ts = (pts_ms - self.cluster_ts_ms) as i16;
|
||||
self.write_simple_block(track_idx + 1, relative_ts, keyframe, data)?;
|
||||
match duration_ns {
|
||||
Some(dur_ns) => {
|
||||
let duration_ms = (dur_ns / 1_000_000).max(1);
|
||||
self.write_block_group(track_idx + 1, relative_ts, keyframe, data, duration_ms)?;
|
||||
}
|
||||
None => {
|
||||
self.write_simple_block(track_idx + 1, relative_ts, keyframe, data)?;
|
||||
}
|
||||
}
|
||||
self.frame_count += 1;
|
||||
|
||||
Ok(())
|
||||
@@ -377,6 +503,8 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
self.end_cluster()?;
|
||||
|
||||
// Write Cues
|
||||
let cues_start = self.writer.stream_position()?;
|
||||
let cues_offset = cues_start - self.segment_start;
|
||||
if !self.cues.is_empty() {
|
||||
let cues_pos = ebml::start_master(&mut self.writer, ebml::CUES)?;
|
||||
for cue in &self.cues {
|
||||
@@ -395,6 +523,21 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
ebml::end_master(&mut self.writer, cues_pos)?;
|
||||
}
|
||||
|
||||
// Back-patch SeekHead SeekPosition values now that all element offsets are known.
|
||||
for fixup in &self.seek_fixups {
|
||||
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::CUES => cues_offset,
|
||||
_ => 0,
|
||||
};
|
||||
self.writer
|
||||
.seek(std::io::SeekFrom::Start(fixup.value_offset))?;
|
||||
self.writer.write_all(&offset.to_be_bytes())?;
|
||||
}
|
||||
self.writer.seek(std::io::SeekFrom::End(0))?;
|
||||
|
||||
self.writer.flush()?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -447,6 +590,34 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_block_group(
|
||||
&mut self,
|
||||
track_num: usize,
|
||||
relative_ts: i16,
|
||||
keyframe: bool,
|
||||
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 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(&relative_ts.to_be_bytes())?;
|
||||
self.writer.write_all(&[flags])?;
|
||||
self.writer.write_all(data)?;
|
||||
ebml::write_uint(&mut self.writer, ebml::BLOCK_DURATION, duration_ms)?;
|
||||
ebml::end_master(&mut self.writer, bg_pos)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
@@ -499,6 +670,7 @@ mod tests {
|
||||
sample_rate: 0.0,
|
||||
channels: 0,
|
||||
bit_depth: 0,
|
||||
dv_config: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -523,9 +695,60 @@ mod tests {
|
||||
sample_rate: 48000.0,
|
||||
channels: 6,
|
||||
bit_depth: 0,
|
||||
dv_config: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn audio_stream(codec: Codec) -> AudioStream {
|
||||
use crate::disc::{AudioChannels, LabelPurpose, SampleRate};
|
||||
AudioStream {
|
||||
pid: 0x1100,
|
||||
codec,
|
||||
channels: AudioChannels::Surround51,
|
||||
language: "eng".into(),
|
||||
sample_rate: SampleRate::S48,
|
||||
secondary: false,
|
||||
purpose: LabelPurpose::Normal,
|
||||
label: String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dts_variants_map_to_registered_a_dts_codec_id() {
|
||||
// The Matroska codec-ID registry defines `A_DTS` for the whole DTS
|
||||
// family (core, DTS-HD HRA, DTS-HD MA). The `/MA` and `/HR` suffixes
|
||||
// are not registered and break strict parsers, so every DTS variant
|
||||
// must emit plain `A_DTS`.
|
||||
for codec in [Codec::Dts, Codec::DtsHdMa, Codec::DtsHdHr] {
|
||||
let track = MkvTrack::audio(&audio_stream(codec));
|
||||
assert_eq!(
|
||||
track.codec_id, "A_DTS",
|
||||
"{codec:?} must map to registered codec ID A_DTS, got {}",
|
||||
track.codec_id
|
||||
);
|
||||
}
|
||||
// Sanity: the non-DTS variants keep their distinct IDs.
|
||||
assert_eq!(MkvTrack::audio(&audio_stream(Codec::Ac3)).codec_id, "A_AC3");
|
||||
assert_eq!(
|
||||
MkvTrack::audio(&audio_stream(Codec::TrueHd)).codec_id,
|
||||
"A_TRUEHD"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dolby_vision_config_profile7() {
|
||||
// dvcC for disc Profile 7 dual-layer: version 1.0, profile 7, all of
|
||||
// bl/el/rpu present. 24 bytes.
|
||||
let c = dolby_vision_config(7, 6, 0);
|
||||
assert_eq!(c.len(), 24);
|
||||
assert_eq!(c[0], 1); // dv_version_major
|
||||
assert_eq!(c[1], 0); // dv_version_minor
|
||||
// profile in the top 7 bits of byte 2
|
||||
assert_eq!(c[2] >> 1, 7, "dv_profile must be 7");
|
||||
// rpu/el/bl present flags in byte 3 (low 3 bits after level)
|
||||
assert_eq!(c[3] & 0b0000_0111, 0b0000_0111, "rpu+el+bl all present");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mkv_writes_ebml_header() {
|
||||
let buf = Cursor::new(Vec::new());
|
||||
@@ -556,7 +779,7 @@ mod tests {
|
||||
let tracks = [make_video_track()];
|
||||
let mut muxer = MkvMuxer::new(buf, &tracks, None, 60.0, &[]).unwrap();
|
||||
muxer
|
||||
.write_frame(0, 0, true, &[0xDE, 0xAD, 0xBE, 0xEF])
|
||||
.write_frame(0, 0, true, &[0xDE, 0xAD, 0xBE, 0xEF], None)
|
||||
.unwrap();
|
||||
let data = muxer.writer.into_inner();
|
||||
assert!(
|
||||
@@ -594,7 +817,9 @@ mod tests {
|
||||
let writer = SharedWriter(shared.clone());
|
||||
let tracks = [make_video_track()];
|
||||
let mut muxer = MkvMuxer::new(writer, &tracks, Some("Cue Test"), 60.0, &[]).unwrap();
|
||||
muxer.write_frame(0, 0, true, &[0x01, 0x02, 0x03]).unwrap();
|
||||
muxer
|
||||
.write_frame(0, 0, true, &[0x01, 0x02, 0x03], None)
|
||||
.unwrap();
|
||||
muxer.finish().unwrap();
|
||||
|
||||
let data = shared.lock().unwrap().clone().into_inner();
|
||||
@@ -610,13 +835,17 @@ mod tests {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let mut muxer = MkvMuxer::new(buf, &tracks, Some("Multi"), 120.0, &[]).unwrap();
|
||||
// Write frames to both tracks
|
||||
muxer.write_frame(0, 0, true, &[0x00, 0x00, 0x01]).unwrap();
|
||||
muxer.write_frame(1, 0, false, &[0x0B, 0x77, 0x00]).unwrap();
|
||||
muxer
|
||||
.write_frame(0, 40_000_000, false, &[0x00, 0x00, 0x01])
|
||||
.write_frame(0, 0, true, &[0x00, 0x00, 0x01], None)
|
||||
.unwrap();
|
||||
muxer
|
||||
.write_frame(1, 32_000_000, false, &[0x0B, 0x77, 0x01])
|
||||
.write_frame(1, 0, false, &[0x0B, 0x77, 0x00], None)
|
||||
.unwrap();
|
||||
muxer
|
||||
.write_frame(0, 40_000_000, false, &[0x00, 0x00, 0x01], None)
|
||||
.unwrap();
|
||||
muxer
|
||||
.write_frame(1, 32_000_000, false, &[0x0B, 0x77, 0x01], None)
|
||||
.unwrap();
|
||||
// Should not panic
|
||||
let data = muxer.writer.into_inner();
|
||||
@@ -631,10 +860,12 @@ mod tests {
|
||||
|
||||
// Record position before first frame
|
||||
let pos_before_kf = muxer.writer.position();
|
||||
muxer.write_frame(0, 0, true, &[0xAA]).unwrap();
|
||||
muxer.write_frame(0, 0, true, &[0xAA], None).unwrap();
|
||||
let pos_after_kf = muxer.writer.position();
|
||||
|
||||
muxer.write_frame(0, 1_000_000, false, &[0xBB]).unwrap();
|
||||
muxer
|
||||
.write_frame(0, 1_000_000, false, &[0xBB], None)
|
||||
.unwrap();
|
||||
let pos_after_nkf = muxer.writer.position();
|
||||
|
||||
let data = muxer.writer.into_inner();
|
||||
@@ -812,4 +1043,501 @@ mod tests {
|
||||
"FlagForced element should not be present for non-forced subtitle"
|
||||
);
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Seekability tests: SeekHead, keyframe-aligned clusters, Cues
|
||||
// ============================================================
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
/// Writer that lets the test inspect the buffer after `finish()` consumes the muxer.
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
/// Build interleaved frames at 24 fps video (IDR every gop_secs) + 48 kHz audio (1024 samples per frame).
|
||||
fn frames_for(duration_secs: f64, gop_secs: f64) -> Vec<(usize, i64, bool, Vec<u8>)> {
|
||||
let video_interval_ns: i64 = 1_000_000_000 / 24;
|
||||
let audio_interval_ns: i64 = (1024i64 * 1_000_000_000) / 48_000;
|
||||
let gop_frames = (gop_secs * 24.0).round() as i64;
|
||||
|
||||
let mut out: Vec<(usize, i64, bool, Vec<u8>)> = Vec::new();
|
||||
let total_ns = (duration_secs * 1_000_000_000.0) as i64;
|
||||
|
||||
let mut vi: i64 = 0;
|
||||
loop {
|
||||
let pts = vi * video_interval_ns;
|
||||
if pts >= total_ns {
|
||||
break;
|
||||
}
|
||||
let keyframe = vi % gop_frames == 0;
|
||||
out.push((0, pts, keyframe, vec![0xAB; 64]));
|
||||
vi += 1;
|
||||
}
|
||||
|
||||
let mut ai: i64 = 0;
|
||||
loop {
|
||||
let pts = ai * audio_interval_ns;
|
||||
if pts >= total_ns {
|
||||
break;
|
||||
}
|
||||
out.push((1, pts, true, vec![0xCD; 32]));
|
||||
ai += 1;
|
||||
}
|
||||
|
||||
out.sort_by_key(|f| f.1);
|
||||
out
|
||||
}
|
||||
|
||||
/// Mux frames through a SharedWriter and return the final buffer.
|
||||
fn mux_to_bytes(
|
||||
tracks: &[MkvTrack],
|
||||
chapters: &[Chapter],
|
||||
frames: &[(usize, i64, bool, Vec<u8>)],
|
||||
) -> (Vec<u8>, u64) {
|
||||
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, chapters).unwrap();
|
||||
for (t, pts, kf, data) in frames {
|
||||
muxer.write_frame(*t, *pts, *kf, data, None).unwrap();
|
||||
}
|
||||
let frame_count = muxer.frame_count;
|
||||
muxer.finish().unwrap();
|
||||
let data = shared.lock().unwrap().clone().into_inner();
|
||||
(data, frame_count)
|
||||
}
|
||||
|
||||
/// Find the Segment header in the buffer and return (segment_id_pos, segment_start_pos).
|
||||
/// segment_start = position immediately after Segment's id + size bytes.
|
||||
fn locate_segment(data: &[u8]) -> (usize, usize) {
|
||||
let segment_id_pos = find_id(data, ebml::SEGMENT).expect("segment id not found");
|
||||
// Segment is written via write_id + write_unknown_size: 4 byte id + 8 byte size
|
||||
(segment_id_pos, segment_id_pos + 4 + 8)
|
||||
}
|
||||
|
||||
/// Walk Segment's top-level children. Returns Vec<(id, data_start_offset, data_size)>
|
||||
/// where data_start_offset is absolute file offset and data_size is the element body size.
|
||||
fn segment_children(data: &[u8]) -> Vec<(u32, usize, u64)> {
|
||||
let (_, seg_start) = locate_segment(data);
|
||||
let mut out = Vec::new();
|
||||
let mut cursor = Cursor::new(&data[seg_start..]);
|
||||
while (cursor.position() as usize) < data.len() - seg_start {
|
||||
let pos_before = cursor.position();
|
||||
let (id, size, hdr_len) = match ebml::read_element_header(&mut cursor) {
|
||||
Ok(v) => v,
|
||||
Err(_) => break,
|
||||
};
|
||||
let data_abs = seg_start + pos_before as usize + hdr_len;
|
||||
out.push((id, data_abs, size));
|
||||
// Skip the body to advance to the next element.
|
||||
cursor
|
||||
.seek(io::SeekFrom::Current(size as i64))
|
||||
.expect("seek past element body");
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Find every Cluster: returns Vec<(cluster_data_start_abs, cluster_data_size, cluster_timestamp_ms)>.
|
||||
fn find_clusters(data: &[u8]) -> Vec<(usize, u64, u64)> {
|
||||
let mut out = Vec::new();
|
||||
for (id, body_start, body_size) in segment_children(data) {
|
||||
if id == ebml::CLUSTER {
|
||||
let mut cursor = Cursor::new(&data[body_start..body_start + body_size as usize]);
|
||||
let (tid, tsize, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(
|
||||
tid,
|
||||
ebml::CLUSTER_TIMESTAMP,
|
||||
"cluster must start with timestamp"
|
||||
);
|
||||
let ts = ebml::read_uint_val(&mut cursor, tsize as usize).unwrap();
|
||||
out.push((body_start, body_size, ts));
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Parse the first SimpleBlock that appears in a cluster body slice.
|
||||
/// Returns (track_num, flags_byte). track_num decoded from VINT.
|
||||
fn first_simple_block(cluster_body: &[u8]) -> (u64, u8) {
|
||||
let mut cursor = Cursor::new(cluster_body);
|
||||
loop {
|
||||
let (id, size, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
if id == ebml::SIMPLE_BLOCK {
|
||||
let body_start = cursor.position() as usize;
|
||||
// Decode track VINT.
|
||||
let b0 = cluster_body[body_start];
|
||||
let (track_num, vint_len) = if b0 & 0x80 != 0 {
|
||||
((b0 & 0x7F) as u64, 1usize)
|
||||
} else if b0 & 0x40 != 0 {
|
||||
let b1 = cluster_body[body_start + 1];
|
||||
((((b0 & 0x3F) as u64) << 8) | b1 as u64, 2)
|
||||
} else {
|
||||
panic!("unsupported track vint width");
|
||||
};
|
||||
let flags = cluster_body[body_start + vint_len + 2];
|
||||
return (track_num, flags);
|
||||
}
|
||||
// Skip non-SimpleBlock child.
|
||||
cursor.seek(io::SeekFrom::Current(size as i64)).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse the Cues element body into Vec<(cue_time, cue_track, cue_cluster_position)>.
|
||||
fn parse_cues(data: &[u8]) -> Vec<(u64, u64, u64)> {
|
||||
let mut out = Vec::new();
|
||||
let (cues_id, cues_body_start, cues_body_size) = segment_children(data)
|
||||
.into_iter()
|
||||
.find(|(id, _, _)| *id == ebml::CUES)
|
||||
.expect("cues element not found");
|
||||
assert_eq!(cues_id, ebml::CUES);
|
||||
let cues_body = &data[cues_body_start..cues_body_start + cues_body_size as usize];
|
||||
let mut cursor = Cursor::new(cues_body);
|
||||
while (cursor.position() as usize) < cues_body.len() {
|
||||
let (id, size, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(id, ebml::CUE_POINT);
|
||||
let cp_end = cursor.position() + size;
|
||||
let mut cue_time = 0u64;
|
||||
let mut cue_track = 0u64;
|
||||
let mut cue_pos = 0u64;
|
||||
while cursor.position() < cp_end {
|
||||
let (sid, ssize, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
match sid {
|
||||
ebml::CUE_TIME => {
|
||||
cue_time = ebml::read_uint_val(&mut cursor, ssize as usize).unwrap();
|
||||
}
|
||||
ebml::CUE_TRACK_POSITIONS => {
|
||||
let ctp_end = cursor.position() + ssize;
|
||||
while cursor.position() < ctp_end {
|
||||
let (iid, isize_, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
match iid {
|
||||
ebml::CUE_TRACK => {
|
||||
cue_track =
|
||||
ebml::read_uint_val(&mut cursor, isize_ as usize).unwrap();
|
||||
}
|
||||
ebml::CUE_CLUSTER_POSITION => {
|
||||
cue_pos =
|
||||
ebml::read_uint_val(&mut cursor, isize_ as usize).unwrap();
|
||||
}
|
||||
_ => {
|
||||
cursor.seek(io::SeekFrom::Current(isize_ as i64)).unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
cursor.seek(io::SeekFrom::Current(ssize as i64)).unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
out.push((cue_time, cue_track, cue_pos));
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Parse the SeekHead body into Vec<(seek_id, seek_position)>.
|
||||
fn parse_seekhead(data: &[u8]) -> Vec<(u32, u64)> {
|
||||
let mut out = Vec::new();
|
||||
let (sh_id, sh_body_start, sh_body_size) = segment_children(data)
|
||||
.into_iter()
|
||||
.find(|(id, _, _)| *id == ebml::SEEK_HEAD)
|
||||
.expect("seekhead not found");
|
||||
assert_eq!(sh_id, ebml::SEEK_HEAD);
|
||||
let sh_body = &data[sh_body_start..sh_body_start + sh_body_size as usize];
|
||||
let mut cursor = Cursor::new(sh_body);
|
||||
while (cursor.position() as usize) < sh_body.len() {
|
||||
let (id, size, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(id, ebml::SEEK);
|
||||
let seek_end = cursor.position() + size;
|
||||
let mut seek_id_val: u32 = 0;
|
||||
let mut seek_pos_val: u64 = 0;
|
||||
while cursor.position() < seek_end {
|
||||
let (sid, ssize, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
match sid {
|
||||
ebml::SEEK_ID => {
|
||||
let raw = ebml::read_uint_val(&mut cursor, ssize as usize).unwrap();
|
||||
seek_id_val = raw as u32;
|
||||
}
|
||||
ebml::SEEK_POSITION => {
|
||||
seek_pos_val = ebml::read_uint_val(&mut cursor, ssize as usize).unwrap();
|
||||
}
|
||||
_ => {
|
||||
cursor.seek(io::SeekFrom::Current(ssize as i64)).unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
out.push((seek_id_val, seek_pos_val));
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cluster_starts_only_on_video_keyframe() {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let frames = frames_for(30.0, 1.0);
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames);
|
||||
let clusters = find_clusters(&data);
|
||||
assert!(!clusters.is_empty(), "expected at least one cluster");
|
||||
for (body_start, body_size, _ts) in clusters {
|
||||
let body = &data[body_start..body_start + body_size as usize];
|
||||
// Skip past the CLUSTER_TIMESTAMP element first.
|
||||
let mut cursor = Cursor::new(body);
|
||||
let (tid, tsize, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(tid, ebml::CLUSTER_TIMESTAMP);
|
||||
cursor.seek(io::SeekFrom::Current(tsize as i64)).unwrap();
|
||||
let after_ts = cursor.position() as usize;
|
||||
let (track_num, flags) = first_simple_block(&body[after_ts..]);
|
||||
assert_eq!(
|
||||
track_num, 1,
|
||||
"first block in cluster must be track 1 (video)"
|
||||
);
|
||||
assert_eq!(
|
||||
flags & 0x80,
|
||||
0x80,
|
||||
"first block in cluster must have keyframe flag set, got 0x{:02X}",
|
||||
flags
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cue_count_equals_cluster_count() {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let frames = frames_for(30.0, 1.0);
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames);
|
||||
let clusters = find_clusters(&data);
|
||||
let cues = parse_cues(&data);
|
||||
assert_eq!(
|
||||
clusters.len(),
|
||||
cues.len(),
|
||||
"cluster count {} != cue count {}",
|
||||
clusters.len(),
|
||||
cues.len()
|
||||
);
|
||||
// For 30s @ 5s min cluster duration with 1s GOP, expect 6 clusters / 6 cues.
|
||||
assert_eq!(
|
||||
clusters.len(),
|
||||
6,
|
||||
"expected 6 clusters for 30s @ 5s cluster duration"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cue_positions_resolve_to_clusters() {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let frames = frames_for(30.0, 1.0);
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames);
|
||||
let (_, seg_start) = locate_segment(&data);
|
||||
let cues = parse_cues(&data);
|
||||
assert!(!cues.is_empty());
|
||||
for (_time, _track, pos) in cues {
|
||||
let abs = seg_start + pos as usize;
|
||||
let mut cursor = Cursor::new(&data[abs..]);
|
||||
let (id, _size, _hdr_len) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(
|
||||
id,
|
||||
ebml::CLUSTER,
|
||||
"cue position 0x{:X} did not resolve to a cluster",
|
||||
pos
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cue_times_match_cluster_timestamps() {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let frames = frames_for(30.0, 1.0);
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames);
|
||||
let (_, seg_start) = locate_segment(&data);
|
||||
let cues = parse_cues(&data);
|
||||
for (time, _track, pos) in cues {
|
||||
let abs = seg_start + pos as usize;
|
||||
let mut cursor = Cursor::new(&data[abs..]);
|
||||
let (id, size, _hdr_len) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(id, ebml::CLUSTER);
|
||||
let body_start = abs + (cursor.position() as usize);
|
||||
let body = &data[body_start..body_start + size as usize];
|
||||
let mut bc = Cursor::new(body);
|
||||
let (tid, tsize, _) = ebml::read_element_header(&mut bc).unwrap();
|
||||
assert_eq!(tid, ebml::CLUSTER_TIMESTAMP);
|
||||
let cluster_ts = ebml::read_uint_val(&mut bc, tsize as usize).unwrap();
|
||||
assert_eq!(
|
||||
cluster_ts, time,
|
||||
"cluster timestamp {} != cue time {}",
|
||||
cluster_ts, time
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seekhead_is_first_child_of_segment() {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames_for(10.0, 1.0));
|
||||
let children = segment_children(&data);
|
||||
assert!(!children.is_empty());
|
||||
assert_eq!(
|
||||
children[0].0,
|
||||
ebml::SEEK_HEAD,
|
||||
"first child of segment must be SeekHead, got id 0x{:X}",
|
||||
children[0].0
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seekhead_points_to_real_elements() {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames_for(10.0, 1.0));
|
||||
let (_, seg_start) = locate_segment(&data);
|
||||
let entries = parse_seekhead(&data);
|
||||
let required = [ebml::INFO, ebml::TRACKS, ebml::CUES];
|
||||
for &want_id in &required {
|
||||
let entry = entries
|
||||
.iter()
|
||||
.find(|(id, _)| *id == want_id)
|
||||
.unwrap_or_else(|| panic!("seekhead missing entry for id 0x{:X}", want_id));
|
||||
let abs = seg_start + entry.1 as usize;
|
||||
let mut cursor = Cursor::new(&data[abs..]);
|
||||
let (got_id, _, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(
|
||||
got_id, want_id,
|
||||
"seekhead entry for 0x{:X} resolves to wrong id 0x{:X}",
|
||||
want_id, got_id
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seekhead_omits_chapters_when_empty() {
|
||||
let tracks = [make_video_track()];
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames_for(5.0, 1.0));
|
||||
let entries = parse_seekhead(&data);
|
||||
assert_eq!(
|
||||
entries.len(),
|
||||
3,
|
||||
"expected 3 seek entries (Info, Tracks, Cues), got {}",
|
||||
entries.len()
|
||||
);
|
||||
assert!(
|
||||
entries.iter().all(|(id, _)| *id != ebml::CHAPTERS),
|
||||
"seekhead should not contain Chapters entry when chapters are empty"
|
||||
);
|
||||
}
|
||||
|
||||
/// Collect every (cluster_ts_ms, block_relative_ts_i16, absolute_ms) for
|
||||
/// all SimpleBlocks across all clusters, so a test can assert that the
|
||||
/// reconstructed absolute timestamp (cluster_ts + relative_ts) is correct
|
||||
/// and that no relative_ts ever wrapped the i16 range.
|
||||
fn all_block_timestamps(data: &[u8]) -> Vec<(i64, i16, i64)> {
|
||||
let mut out = Vec::new();
|
||||
for (body_start, body_size, cluster_ts) in find_clusters(data) {
|
||||
let body = &data[body_start..body_start + body_size as usize];
|
||||
let mut cursor = Cursor::new(body);
|
||||
// Skip CLUSTER_TIMESTAMP.
|
||||
let (tid, tsize, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(tid, ebml::CLUSTER_TIMESTAMP);
|
||||
cursor.seek(io::SeekFrom::Current(tsize as i64)).unwrap();
|
||||
while (cursor.position() as usize) < body.len() {
|
||||
let (id, sz, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
if id == ebml::SIMPLE_BLOCK {
|
||||
let bstart = cursor.position() as usize;
|
||||
let b0 = body[bstart];
|
||||
let vint_len = if b0 & 0x80 != 0 { 1 } else { 2 };
|
||||
let ts_pos = bstart + vint_len;
|
||||
let rel = i16::from_be_bytes([body[ts_pos], body[ts_pos + 1]]);
|
||||
out.push((cluster_ts as i64, rel, cluster_ts as i64 + rel as i64));
|
||||
}
|
||||
cursor.seek(io::SeekFrom::Current(sz as i64)).unwrap();
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn long_audio_gap_forces_cluster_no_i16_overflow() {
|
||||
// Regression for the `(pts_ms - cluster_ts_ms) as i16` truncation:
|
||||
// a single video keyframe at t=0 opens one cluster, then a long
|
||||
// audio-only stretch (no further video keyframe) drifts well past
|
||||
// i16::MAX ms (~32.767 s). Without the overflow guard the audio
|
||||
// blocks past 32.767 s would write a wrapped (negative) relative
|
||||
// timestamp into the SimpleBlock. With the guard a fresh cluster is
|
||||
// forced so every relative_ts stays in range and reconstructs to the
|
||||
// true absolute timestamp.
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let mut frames: Vec<(usize, i64, bool, Vec<u8>)> = Vec::new();
|
||||
// One video keyframe at t=0 (opens the first cluster).
|
||||
frames.push((0, 0, true, vec![0xAB; 16]));
|
||||
// Audio frames every 100 ms out to 60 s — past the 32.767 s i16 limit
|
||||
// and past two i16 spans, with NO further video keyframe.
|
||||
let mut t_ms = 0i64;
|
||||
while t_ms <= 60_000 {
|
||||
frames.push((1, t_ms * 1_000_000, true, vec![0xCD; 16]));
|
||||
t_ms += 100;
|
||||
}
|
||||
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames);
|
||||
|
||||
let blocks = all_block_timestamps(&data);
|
||||
assert!(!blocks.is_empty());
|
||||
// Every block's relative timestamp must be within i16 range (it is by
|
||||
// type), AND must reconstruct to a non-negative, monotonic-ish
|
||||
// absolute timestamp matching the source — i.e. no silent wrap.
|
||||
for (cluster_ts, rel, abs) in &blocks {
|
||||
assert!(
|
||||
*rel as i64 >= 0 && (*rel as i64) <= MAX_BLOCK_REL_MS,
|
||||
"block relative_ts {rel} out of [0, i16::MAX] range \
|
||||
(cluster_ts={cluster_ts}, abs={abs}) — i16 overflow"
|
||||
);
|
||||
}
|
||||
// The latest audio frame is at 60_000 ms; its reconstructed absolute
|
||||
// timestamp must equal that, proving no truncation occurred.
|
||||
let max_abs = blocks.iter().map(|(_, _, abs)| *abs).max().unwrap();
|
||||
assert_eq!(max_abs, 60_000, "last block must reconstruct to 60_000 ms");
|
||||
// The overflow guard must have opened more than one cluster (the
|
||||
// single keyframe alone would otherwise yield exactly one).
|
||||
let clusters = find_clusters(&data);
|
||||
assert!(
|
||||
clusters.len() >= 2,
|
||||
"expected the i16 guard to force extra clusters, got {}",
|
||||
clusters.len()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pre_first_keyframe_frames_dropped() {
|
||||
let tracks = [make_video_track()];
|
||||
let frames = vec![
|
||||
(0usize, 0i64, false, vec![0x11; 16]),
|
||||
(0usize, 41_000_000i64, true, vec![0x22; 16]),
|
||||
];
|
||||
let (data, frame_count) = mux_to_bytes(&tracks, &[], &frames);
|
||||
assert_eq!(frame_count, 1, "muxer.frame_count must equal 1");
|
||||
let clusters = find_clusters(&data);
|
||||
assert_eq!(clusters.len(), 1, "expected exactly one cluster");
|
||||
let (body_start, body_size, _ts) = clusters[0];
|
||||
let body = &data[body_start..body_start + body_size as usize];
|
||||
let mut cursor = Cursor::new(body);
|
||||
// Skip CLUSTER_TIMESTAMP.
|
||||
let (tid, tsize, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(tid, ebml::CLUSTER_TIMESTAMP);
|
||||
cursor.seek(io::SeekFrom::Current(tsize as i64)).unwrap();
|
||||
let mut sb_count = 0;
|
||||
while (cursor.position() as usize) < body.len() {
|
||||
let (id, sz, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
if id == ebml::SIMPLE_BLOCK {
|
||||
sb_count += 1;
|
||||
}
|
||||
cursor.seek(io::SeekFrom::Current(sz as i64)).unwrap();
|
||||
}
|
||||
assert_eq!(sb_count, 1, "expected exactly one SimpleBlock in output");
|
||||
}
|
||||
}
|
||||
|
||||
+211
-19
@@ -14,6 +14,49 @@ fn skip_bytes(r: &mut impl Read, n: u64) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ── Sanity caps for untrusted EBML element sizes ──────────────
|
||||
//
|
||||
// Sizes come straight from the EBML stream (file or network) and are
|
||||
// otherwise cast to `usize` and used to allocate/read. An adversarial
|
||||
// or corrupt container can claim a multi-GB element and trigger an OOM
|
||||
// allocation, or claim an integer element wider than 8 bytes and panic
|
||||
// the fixed 8-byte reader. Every untrusted size is validated against
|
||||
// one of these caps before allocation.
|
||||
|
||||
/// Largest accepted SIMPLE_BLOCK payload. A block is a small vint track
|
||||
/// header + 2-byte rel-ts + 1-byte flags + one frame of elementary data.
|
||||
/// UHD HEVC keyframes run a few MB; 64 MiB is generously above any real
|
||||
/// single-frame block while still bounding a hostile allocation.
|
||||
const MAX_BLOCK_SIZE: u64 = 64 * 1024 * 1024;
|
||||
/// Largest accepted CODEC_PRIVATE payload. hvcC/avcC/setup blobs are a
|
||||
/// few KB in practice; 16 MiB is far above any legitimate value.
|
||||
const MAX_CODEC_PRIVATE: u64 = 16 * 1024 * 1024;
|
||||
/// Largest accepted string element (TITLE/CODEC_ID/LANGUAGE/TRACK_NAME).
|
||||
const MAX_STRING_LEN: u64 = 64 * 1024;
|
||||
/// EBML unsigned-int elements are at most 8 bytes wide.
|
||||
const MAX_UINT_LEN: u64 = 8;
|
||||
|
||||
/// Reject an untrusted element size that exceeds `cap` before it is used
|
||||
/// to allocate or read. Returns the size as `usize` when within bounds.
|
||||
fn checked_size(size: u64, cap: u64) -> io::Result<usize> {
|
||||
if size > cap {
|
||||
return Err(crate::error::Error::MkvInvalid.into());
|
||||
}
|
||||
Ok(size as usize)
|
||||
}
|
||||
|
||||
/// Read a bounded unsigned int. Guards against `size > 8` (which would
|
||||
/// otherwise index out of the fixed 8-byte buffer in `read_uint_val`)
|
||||
/// before delegating.
|
||||
fn read_uint_bounded(r: &mut impl Read, size: u64) -> io::Result<u64> {
|
||||
ebml::read_uint_val(r, checked_size(size, MAX_UINT_LEN)?)
|
||||
}
|
||||
|
||||
/// Read a bounded UTF-8 string element.
|
||||
fn read_string_bounded(r: &mut impl Read, size: u64) -> io::Result<String> {
|
||||
ebml::read_string_val(r, checked_size(size, MAX_STRING_LEN)?)
|
||||
}
|
||||
|
||||
use crate::disc::*;
|
||||
use std::io::{self, Read};
|
||||
|
||||
@@ -108,11 +151,12 @@ impl crate::pes::Stream for MkvStream {
|
||||
match id {
|
||||
ebml::CLUSTER => continue,
|
||||
ebml::CLUSTER_TIMESTAMP => {
|
||||
rs.cluster_ts_ms = ebml::read_uint_val(&mut rs.reader, size as usize)? as i64;
|
||||
rs.cluster_ts_ms = read_uint_bounded(&mut rs.reader, size)? as i64;
|
||||
continue;
|
||||
}
|
||||
ebml::SIMPLE_BLOCK => {
|
||||
let block = ebml::read_binary_val(&mut rs.reader, size as usize)?;
|
||||
let block =
|
||||
ebml::read_binary_val(&mut rs.reader, checked_size(size, MAX_BLOCK_SIZE)?)?;
|
||||
if block.len() < 4 {
|
||||
continue;
|
||||
}
|
||||
@@ -138,6 +182,7 @@ impl crate::pes::Stream for MkvStream {
|
||||
pts: pts_ms * 1_000_000, // ms → ns
|
||||
keyframe,
|
||||
data,
|
||||
duration_ns: None,
|
||||
}));
|
||||
}
|
||||
_ => {
|
||||
@@ -150,9 +195,13 @@ impl crate::pes::Stream for MkvStream {
|
||||
|
||||
fn write(&mut self, frame: &crate::pes::PesFrame) -> io::Result<()> {
|
||||
match &mut self.mode {
|
||||
Mode::Write { muxer: Some(m) } => {
|
||||
m.write_frame(frame.track, frame.pts, frame.keyframe, &frame.data)
|
||||
}
|
||||
Mode::Write { muxer: Some(m) } => m.write_frame(
|
||||
frame.track,
|
||||
frame.pts,
|
||||
frame.keyframe,
|
||||
&frame.data,
|
||||
frame.duration_ns,
|
||||
),
|
||||
Mode::Write { muxer: None } => Ok(()),
|
||||
Mode::Read(_) => Err(crate::error::Error::StreamReadOnly.into()),
|
||||
}
|
||||
@@ -193,7 +242,10 @@ impl crate::pes::Stream for MkvStream {
|
||||
/// Returns (DiscTitle, codec_privates: Vec<(track_number, codec_private_bytes)>)
|
||||
fn parse_mkv_header(r: &mut impl Read) -> MkvHeaderResult {
|
||||
let mut title = String::new();
|
||||
let mut duration_ms = 0.0f64;
|
||||
// EBML `DURATION` is a float expressed in TimestampScale ticks, not
|
||||
// milliseconds (Matroska spec). Named accordingly; converted to
|
||||
// seconds below as ticks * ts_scale_ns / 1e9.
|
||||
let mut duration_ticks = 0.0f64;
|
||||
let mut ts_scale: u64 = 1_000_000;
|
||||
let mut streams: Vec<crate::disc::Stream> = Vec::new();
|
||||
let mut codec_privates: Vec<(u16, Vec<u8>)> = Vec::new();
|
||||
@@ -230,9 +282,9 @@ fn parse_mkv_header(r: &mut impl Read) -> MkvHeaderResult {
|
||||
let (cid, cs, hlen) = ebml::read_element_header(r)?;
|
||||
remaining = remaining.saturating_sub(hlen as u64 + cs);
|
||||
match cid {
|
||||
ebml::TIMESTAMP_SCALE => ts_scale = ebml::read_uint_val(r, cs as usize)?,
|
||||
ebml::DURATION => duration_ms = ebml::read_float_val(r, cs as usize)?,
|
||||
ebml::TITLE => title = ebml::read_string_val(r, cs as usize)?,
|
||||
ebml::TIMESTAMP_SCALE => ts_scale = read_uint_bounded(r, cs)?,
|
||||
ebml::DURATION => duration_ticks = ebml::read_float_val(r, cs as usize)?,
|
||||
ebml::TITLE => title = read_string_bounded(r, cs)?,
|
||||
_ => {
|
||||
skip_bytes(r, cs)?;
|
||||
}
|
||||
@@ -269,7 +321,7 @@ fn parse_mkv_header(r: &mut impl Read) -> MkvHeaderResult {
|
||||
|
||||
let disc_title = DiscTitle {
|
||||
playlist: title,
|
||||
duration_secs: duration_ms * (ts_scale as f64) / 1_000_000_000.0,
|
||||
duration_secs: duration_ticks * (ts_scale as f64) / 1_000_000_000.0,
|
||||
streams,
|
||||
..DiscTitle::empty()
|
||||
};
|
||||
@@ -291,20 +343,25 @@ fn parse_track(
|
||||
let (cid, cs, hlen) = ebml::read_element_header(r)?;
|
||||
remaining = remaining.saturating_sub(hlen as u64 + cs);
|
||||
match cid {
|
||||
ebml::TRACK_NUMBER => tnum = ebml::read_uint_val(r, cs as usize)? as u16,
|
||||
ebml::TRACK_TYPE => ttype = ebml::read_uint_val(r, cs as usize)?,
|
||||
ebml::CODEC_ID => codec_id = ebml::read_string_val(r, cs as usize)?,
|
||||
ebml::CODEC_PRIVATE => codec_priv = Some(ebml::read_binary_val(r, cs as usize)?),
|
||||
ebml::LANGUAGE => lang = ebml::read_string_val(r, cs as usize)?,
|
||||
ebml::TRACK_NAME => name = ebml::read_string_val(r, cs as usize)?,
|
||||
ebml::FLAG_FORCED => forced = ebml::read_uint_val(r, cs as usize)? != 0,
|
||||
ebml::TRACK_NUMBER => tnum = read_uint_bounded(r, cs)? as u16,
|
||||
ebml::TRACK_TYPE => ttype = read_uint_bounded(r, cs)?,
|
||||
ebml::CODEC_ID => codec_id = read_string_bounded(r, cs)?,
|
||||
ebml::CODEC_PRIVATE => {
|
||||
codec_priv = Some(ebml::read_binary_val(
|
||||
r,
|
||||
checked_size(cs, MAX_CODEC_PRIVATE)?,
|
||||
)?)
|
||||
}
|
||||
ebml::LANGUAGE => lang = read_string_bounded(r, cs)?,
|
||||
ebml::TRACK_NAME => name = read_string_bounded(r, cs)?,
|
||||
ebml::FLAG_FORCED => forced = read_uint_bounded(r, cs)? != 0,
|
||||
ebml::VIDEO => {
|
||||
let mut vrem = cs;
|
||||
while vrem > 0 {
|
||||
let (vid, vs, vhlen) = ebml::read_element_header(r)?;
|
||||
vrem = vrem.saturating_sub(vhlen as u64 + vs);
|
||||
if vid == ebml::PIXEL_HEIGHT {
|
||||
ph = ebml::read_uint_val(r, vs as usize)? as u32;
|
||||
ph = read_uint_bounded(r, vs)? as u32;
|
||||
} else {
|
||||
skip_bytes(r, vs)?;
|
||||
}
|
||||
@@ -317,7 +374,7 @@ fn parse_track(
|
||||
arem = arem.saturating_sub(ahlen as u64 + as_);
|
||||
match aid {
|
||||
ebml::SAMPLING_FREQUENCY => sr = ebml::read_float_val(r, as_ as usize)?,
|
||||
ebml::CHANNELS => ch = ebml::read_uint_val(r, as_ as usize)? as u8,
|
||||
ebml::CHANNELS => ch = read_uint_bounded(r, as_)? as u8,
|
||||
_ => {
|
||||
skip_bytes(r, as_)?;
|
||||
}
|
||||
@@ -423,3 +480,138 @@ fn block_vint(d: &[u8]) -> (u64, usize) {
|
||||
}
|
||||
(0, 1) // Unsupported 5+ byte VINT — treat as track 0
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::pes::Stream as _;
|
||||
use std::io::Cursor;
|
||||
|
||||
// `From<Error> for io::Error` encodes the numeric code into the
|
||||
// Display string as "E{code}: ...". Check the prefix.
|
||||
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 checked_size_rejects_over_cap() {
|
||||
// Within cap → Ok with usize value.
|
||||
assert_eq!(checked_size(100, 256).unwrap(), 100);
|
||||
assert_eq!(checked_size(256, 256).unwrap(), 256);
|
||||
// Over cap → MkvInvalid, never a giant allocation.
|
||||
let e = checked_size(257, 256).unwrap_err();
|
||||
assert!(is_mkv_invalid(&e));
|
||||
// A hostile multi-GB block size is rejected as MkvInvalid.
|
||||
let e = checked_size(4 * 1024 * 1024 * 1024, MAX_BLOCK_SIZE).unwrap_err();
|
||||
assert!(is_mkv_invalid(&e));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_uint_bounded_rejects_oversized_int() {
|
||||
// size > 8 would index out of the fixed 8-byte buffer in
|
||||
// read_uint_val (panic / OOB). The guard turns it into a clean
|
||||
// MkvInvalid error instead.
|
||||
let mut data = Cursor::new(vec![0u8; 16]);
|
||||
let e = read_uint_bounded(&mut data, 9).unwrap_err();
|
||||
assert!(is_mkv_invalid(&e));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_uint_bounded_accepts_valid_width() {
|
||||
// 8 bytes is the max legal EBML uint width and must still work.
|
||||
let mut data = Cursor::new(vec![0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x02]);
|
||||
assert_eq!(read_uint_bounded(&mut data, 8).unwrap(), 0x0102);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_string_bounded_rejects_huge_string() {
|
||||
// Claimed string length far above the cap must not allocate.
|
||||
let mut data = Cursor::new(vec![0u8; 16]);
|
||||
let e = read_string_bounded(&mut data, MAX_STRING_LEN + 1).unwrap_err();
|
||||
assert!(is_mkv_invalid(&e));
|
||||
}
|
||||
|
||||
/// Build a minimal MKV (EBML header + Segment + Info + Tracks) so the
|
||||
/// reader is positioned in the cluster body, then append the given
|
||||
/// cluster bytes. Returns the full byte stream.
|
||||
fn minimal_mkv_with_cluster(cluster_body: &[u8]) -> Vec<u8> {
|
||||
let mut out = Vec::new();
|
||||
// EBML header (empty body).
|
||||
ebml::write_id(&mut out, ebml::EBML).unwrap();
|
||||
ebml::write_size(&mut out, 0).unwrap();
|
||||
// Segment (unknown size so the reader streams children).
|
||||
ebml::write_id(&mut out, ebml::SEGMENT).unwrap();
|
||||
ebml::write_unknown_size(&mut out).unwrap();
|
||||
// Empty Info.
|
||||
ebml::write_id(&mut out, ebml::INFO).unwrap();
|
||||
ebml::write_size(&mut out, 0).unwrap();
|
||||
// Empty Tracks.
|
||||
ebml::write_id(&mut out, ebml::TRACKS).unwrap();
|
||||
ebml::write_size(&mut out, 0).unwrap();
|
||||
out.extend_from_slice(cluster_body);
|
||||
out
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn simple_block_oversized_size_is_rejected() {
|
||||
// Cluster containing a SIMPLE_BLOCK that claims a 2 GiB payload.
|
||||
// The reader must reject it (MkvInvalid) rather than attempt a
|
||||
// multi-GB allocation. Header parse stops at CLUSTER, so the
|
||||
// SIMPLE_BLOCK is hit on the first read().
|
||||
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, 2 * 1024 * 1024 * 1024).unwrap();
|
||||
// No payload follows — but we must fail on the size check, before
|
||||
// any read of the body.
|
||||
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));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn well_formed_simple_block_round_trips() {
|
||||
// A small, well-formed SIMPLE_BLOCK must still parse into a frame.
|
||||
// We need at least one stream so the track index is in range, so
|
||||
// give Tracks one video TRACK_ENTRY (track number 1).
|
||||
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();
|
||||
|
||||
// Tracks → one TRACK_ENTRY (track number 1, type 1 = video).
|
||||
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();
|
||||
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);
|
||||
|
||||
// Cluster with a SIMPLE_BLOCK: track vint=0x81 (track 1),
|
||||
// rel_ts=0x0000, flags=0x80 (keyframe), then 4 bytes of data.
|
||||
ebml::write_id(&mut out, ebml::CLUSTER).unwrap();
|
||||
ebml::write_unknown_size(&mut out).unwrap();
|
||||
let block = [0x81u8, 0x00, 0x00, 0x80, 0xAA, 0xBB, 0xCC, 0xDD];
|
||||
ebml::write_id(&mut out, ebml::SIMPLE_BLOCK).unwrap();
|
||||
ebml::write_size(&mut out, block.len() as u64).unwrap();
|
||||
out.extend_from_slice(&block);
|
||||
|
||||
let mut stream = MkvStream::open(Cursor::new(out)).unwrap();
|
||||
let frame = stream.read().unwrap().expect("expected a frame");
|
||||
assert_eq!(frame.track, 0);
|
||||
assert!(frame.keyframe);
|
||||
assert_eq!(frame.data, vec![0xAA, 0xBB, 0xCC, 0xDD]);
|
||||
}
|
||||
}
|
||||
|
||||
+16
-3
@@ -6,7 +6,7 @@
|
||||
//! ```text
|
||||
//! let mut input = input("iso://Disc.iso", &opts)?;
|
||||
//! let title = input.info().clone();
|
||||
//! let mut output = output("mkv://Dune.mkv", &title)?;
|
||||
//! let mut output = output("mkv://Movie.mkv", &title)?;
|
||||
//! while let Ok(Some(frame)) = input.read() {
|
||||
//! output.write(&frame)?;
|
||||
//! }
|
||||
@@ -17,8 +17,9 @@
|
||||
|
||||
// Public modules — types here are intentionally part of the consumable API.
|
||||
pub mod codec;
|
||||
pub mod demux_thread;
|
||||
pub mod disc;
|
||||
pub mod iso;
|
||||
pub mod pipelined_stream;
|
||||
pub mod resolve;
|
||||
|
||||
// Internal modules — implementation details. Their *types* are re-exported
|
||||
@@ -33,6 +34,17 @@ pub(crate) mod m2ts;
|
||||
/// to round-trip codec_privates that don't fit inside the underlying format).
|
||||
/// Exposed for integration tests that exercise the wire format directly.
|
||||
pub mod meta;
|
||||
|
||||
// ── Phase 3 sequential 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;
|
||||
pub(crate) mod mkv;
|
||||
pub(crate) mod mkvstream;
|
||||
pub(crate) mod network;
|
||||
@@ -43,11 +55,12 @@ pub(crate) mod ts;
|
||||
pub(crate) mod tsmux;
|
||||
|
||||
pub use disc::DiscStream;
|
||||
pub use iso::IsoSectorReader;
|
||||
pub use m2ts::M2tsStream;
|
||||
pub use mkvstream::MkvStream;
|
||||
pub use network::NetworkStream;
|
||||
pub use null::NullStream;
|
||||
pub use pipelined_stream::PipelinedPesStream;
|
||||
pub use resolve::build_iso_pipeline;
|
||||
pub use resolve::{InputOptions, StreamUrl, input, output, parse_url};
|
||||
pub use stdio::StdioStream;
|
||||
|
||||
|
||||
@@ -186,6 +186,7 @@ mod tests {
|
||||
pts: 90000,
|
||||
keyframe: true,
|
||||
data: vec![0x47; 192],
|
||||
duration_ns: None,
|
||||
};
|
||||
pes::Stream::write(&mut writer, &frame).unwrap();
|
||||
pes::Stream::finish(&mut writer).unwrap();
|
||||
|
||||
@@ -47,6 +47,7 @@ mod tests {
|
||||
pts: 0,
|
||||
keyframe: true,
|
||||
data: vec![0x01, 0x02, 0x03],
|
||||
duration_ns: None,
|
||||
};
|
||||
sink.write(&frame).unwrap();
|
||||
let _ = sink.info();
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
//! `PipelinedPesStream` — the read-side of the freemkv mux
|
||||
//! highway.
|
||||
//!
|
||||
//! Given a [`crate::mux::demux_thread::DemuxThread`] (which has the
|
||||
//! producer + demux workers already spawned), a set of codec
|
||||
//! parsers, and the title metadata, this struct implements
|
||||
//! [`crate::pes::Stream`] by running codec parse on the caller's
|
||||
//! thread and emitting `PesFrame`s one at a time.
|
||||
//!
|
||||
//! The pipeline runs three threads in parallel:
|
||||
//!
|
||||
//! ```text
|
||||
//! Thread A: read + decrypt (PrefetchedSectorSource / BytePrefetcher)
|
||||
//! Thread B: M2TS demux (DemuxThread)
|
||||
//! Thread C: codec parse (this struct, on the caller's thread)
|
||||
//! ```
|
||||
//!
|
||||
//! Communication between A→B and B→C is via bounded channels with
|
||||
//! 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.
|
||||
|
||||
use super::codec::CodecParser;
|
||||
use super::demux_thread::{DemuxBatch, DemuxThread};
|
||||
use super::ts::PesPacket;
|
||||
use crate::disc::DiscTitle;
|
||||
use crate::pes::{PesFrame, Stream};
|
||||
use crossbeam_channel::Receiver;
|
||||
use std::io;
|
||||
|
||||
/// Stream impl that consumes pre-demuxed `PesPacket` batches from a
|
||||
/// [`DemuxThread`] and runs codec parse on the caller's thread.
|
||||
pub struct PipelinedPesStream {
|
||||
title: DiscTitle,
|
||||
parsers: Vec<(u16, Box<dyn CodecParser>)>,
|
||||
pid_to_track: Vec<(u16, usize)>,
|
||||
|
||||
demux_rx: Receiver<DemuxBatch>,
|
||||
/// Kept alive so dropping this stream joins the demux + producer
|
||||
/// workers deterministically. Never poked directly after spawn.
|
||||
#[allow(dead_code)]
|
||||
demux_thread: DemuxThread,
|
||||
|
||||
pending_frames: std::collections::VecDeque<PesFrame>,
|
||||
eof: bool,
|
||||
}
|
||||
|
||||
impl PipelinedPesStream {
|
||||
/// Wire up the stream. Caller has already spawned the
|
||||
/// `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(
|
||||
demux_thread: DemuxThread,
|
||||
demux_rx: Receiver<DemuxBatch>,
|
||||
title: DiscTitle,
|
||||
parsers: Vec<(u16, Box<dyn CodecParser>)>,
|
||||
pid_to_track: Vec<(u16, usize)>,
|
||||
) -> Self {
|
||||
Self {
|
||||
title,
|
||||
parsers,
|
||||
pid_to_track,
|
||||
demux_rx,
|
||||
demux_thread,
|
||||
pending_frames: std::collections::VecDeque::new(),
|
||||
eof: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Pull one batch of `PesPacket`s from the demux thread, run
|
||||
/// codec parse on each, enqueue resulting `PesFrame`s on
|
||||
/// `pending_frames`. Returns Ok(true) on success, Ok(false) on
|
||||
/// EOF (channel closed cleanly), Err on demuxer error.
|
||||
fn pump_one_batch(&mut self) -> io::Result<bool> {
|
||||
match self.demux_rx.recv() {
|
||||
Ok(DemuxBatch::Ts(packets)) => {
|
||||
self.consume_ts(packets);
|
||||
Ok(true)
|
||||
}
|
||||
Ok(DemuxBatch::Ps(packets)) => {
|
||||
self.consume_ps(packets);
|
||||
Ok(true)
|
||||
}
|
||||
Ok(DemuxBatch::Err(e)) => Err(e),
|
||||
Err(_) => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn consume_ts(&mut self, packets: Vec<PesPacket>) {
|
||||
let skip_parse = std::env::var_os("FREEMKV_SKIP_PARSE").is_some();
|
||||
for pes in packets {
|
||||
if let Some((_, track)) = self
|
||||
.pid_to_track
|
||||
.iter()
|
||||
.find(|(pid, _)| *pid == pes.pid)
|
||||
.copied()
|
||||
{
|
||||
if skip_parse {
|
||||
// Profiling escape hatch — bypass codec parser.
|
||||
self.pending_frames.push_back(PesFrame {
|
||||
track,
|
||||
pts: pes.pts.map(super::codec::pts_to_ns).unwrap_or(0),
|
||||
keyframe: false,
|
||||
data: pes.data,
|
||||
duration_ns: None,
|
||||
});
|
||||
} else if let Some((_, parser)) =
|
||||
self.parsers.iter_mut().find(|(pid, _)| *pid == pes.pid)
|
||||
{
|
||||
for frame in parser.parse(&pes) {
|
||||
self.pending_frames
|
||||
.push_back(PesFrame::from_codec_frame(track, frame));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn consume_ps(&mut self, packets: Vec<super::ps::PsPacket>) {
|
||||
for ps in packets {
|
||||
let track = match ps.stream_id {
|
||||
0xE0..=0xEF => 0,
|
||||
0xC0..=0xDF => 1,
|
||||
0xBD => ps
|
||||
.sub_stream_id
|
||||
.map(|s| (s & 0x1F) as usize + 1)
|
||||
.unwrap_or(1),
|
||||
_ => continue,
|
||||
};
|
||||
if track >= self.title.streams.len() {
|
||||
continue;
|
||||
}
|
||||
let pid = self
|
||||
.pid_to_track
|
||||
.iter()
|
||||
.find(|(_, idx)| *idx == track)
|
||||
.map(|(p, _)| *p)
|
||||
.unwrap_or(0);
|
||||
let pes = PesPacket {
|
||||
pid,
|
||||
pts: ps.pts.map(|p| p as i64),
|
||||
dts: ps.dts.map(|d| d as i64),
|
||||
data: ps.data,
|
||||
};
|
||||
if let Some((_, parser)) = self.parsers.iter_mut().find(|(p, _)| *p == pid) {
|
||||
for frame in parser.parse(&pes) {
|
||||
self.pending_frames
|
||||
.push_back(PesFrame::from_codec_frame(track, frame));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Stream for PipelinedPesStream {
|
||||
fn read(&mut self) -> io::Result<Option<PesFrame>> {
|
||||
if let Some(frame) = self.pending_frames.pop_front() {
|
||||
return Ok(Some(frame));
|
||||
}
|
||||
if self.eof {
|
||||
return Ok(None);
|
||||
}
|
||||
loop {
|
||||
match self.pump_one_batch()? {
|
||||
true => {
|
||||
if let Some(frame) = self.pending_frames.pop_front() {
|
||||
return Ok(Some(frame));
|
||||
}
|
||||
// Batch contained no trackable packets — pull again.
|
||||
}
|
||||
false => {
|
||||
self.eof = true;
|
||||
// Drain any access unit a parser buffered past the last
|
||||
// PES (e.g. DTS-HD's final core+extension unit).
|
||||
let pid_to_track = &self.pid_to_track;
|
||||
let pending = &mut self.pending_frames;
|
||||
for (pid, parser) in self.parsers.iter_mut() {
|
||||
let Some(&(_, track)) = pid_to_track.iter().find(|(p, _)| p == pid) else {
|
||||
continue;
|
||||
};
|
||||
for frame in parser.flush() {
|
||||
pending.push_back(PesFrame::from_codec_frame(track, frame));
|
||||
}
|
||||
}
|
||||
return Ok(self.pending_frames.pop_front());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn write(&mut self, _: &PesFrame) -> io::Result<()> {
|
||||
Err(crate::error::Error::StreamReadOnly.into())
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn info(&self) -> &DiscTitle {
|
||||
&self.title
|
||||
}
|
||||
|
||||
fn headers_ready(&self) -> bool {
|
||||
// Match the previous DiscStream semantics: video tracks need
|
||||
// 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() {
|
||||
return true;
|
||||
}
|
||||
for (idx, s) in self.title.streams.iter().enumerate() {
|
||||
if let crate::disc::Stream::Video(v) = s {
|
||||
if !v.secondary && self.codec_private(idx).is_none() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
fn codec_private(&self, track: usize) -> Option<Vec<u8>> {
|
||||
let pid = self
|
||||
.pid_to_track
|
||||
.iter()
|
||||
.find(|(_, idx)| *idx == track)
|
||||
.map(|(p, _)| *p)?;
|
||||
self.parsers
|
||||
.iter()
|
||||
.find(|(p, _)| *p == pid)
|
||||
.and_then(|(_, parser)| parser.codec_private())
|
||||
}
|
||||
}
|
||||
+285
-24
@@ -15,11 +15,13 @@
|
||||
//! Bare paths without a scheme are rejected.
|
||||
//! For disc→ISO (raw sector copy), use `Disc::copy()` instead.
|
||||
|
||||
use super::disc::DiscStream;
|
||||
use super::network::NetworkStream;
|
||||
use super::null::NullStream;
|
||||
use super::pipelined_stream::PipelinedPesStream;
|
||||
use super::stdio::StdioStream;
|
||||
use super::{M2tsStream, MkvStream};
|
||||
use crate::disc::{ContentFormat, DiscTitle};
|
||||
use crate::sector::SectorSource;
|
||||
use std::io;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
@@ -160,12 +162,32 @@ fn validate_network_addr(addr: &str) -> io::Result<()> {
|
||||
/// Options for opening an input stream.
|
||||
#[derive(Default)]
|
||||
pub struct InputOptions {
|
||||
pub keydb_path: Option<String>,
|
||||
/// 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])>,
|
||||
pub title_index: Option<usize>,
|
||||
/// Skip decryption — return raw encrypted bytes.
|
||||
pub raw: bool,
|
||||
}
|
||||
|
||||
/// Decide whether an ISO mux must abort for lack of a usable AACS key.
|
||||
///
|
||||
/// Returns `true` only when ALL hold: decryption is requested (`!raw`), the
|
||||
/// disc carries AACS state (`has_aacs` — AACS-encrypted, not CSS/unencrypted),
|
||||
/// and key resolution produced no usable key (`keys` is
|
||||
/// [`crate::decrypt::DecryptKeys::None`]). In that case muxing would emit
|
||||
/// undecryptable garbage, so the caller fails fast with [`Error::NoDiscKey`].
|
||||
///
|
||||
/// `--raw` (raw=true) always returns `false` — raw intentionally skips
|
||||
/// decryption and needs no key. A non-AACS disc (`has_aacs=false`) always
|
||||
/// returns `false`: unencrypted content has `None` keys legitimately, and CSS
|
||||
/// DVDs resolve to `DecryptKeys::Css{..}` (never `None`).
|
||||
fn aacs_key_missing(raw: bool, has_aacs: bool, keys: &crate::decrypt::DecryptKeys) -> bool {
|
||||
!raw && has_aacs && matches!(keys, crate::decrypt::DecryptKeys::None)
|
||||
}
|
||||
|
||||
/// Open a PES input stream (produces PES frames).
|
||||
pub fn input(url: &str, opts: &InputOptions) -> io::Result<Box<dyn crate::pes::Stream>> {
|
||||
let parsed = parse_url(url);
|
||||
@@ -179,16 +201,45 @@ pub fn input(url: &str, opts: &InputOptions) -> io::Result<Box<dyn crate::pes::S
|
||||
}
|
||||
StreamUrl::Iso { ref path } => {
|
||||
validate_file_path(path, "iso")?;
|
||||
let scan_opts = match &opts.keydb_path {
|
||||
Some(p) => crate::disc::ScanOptions {
|
||||
keydb_path: Some(p.into()),
|
||||
},
|
||||
None => crate::disc::ScanOptions::default(),
|
||||
};
|
||||
let mut reader = super::iso::IsoSectorReader::open(&path.to_string_lossy())?;
|
||||
// FileSectorSource is the sole file-backed sector source.
|
||||
// It carries the platform-tuned SEQUENTIAL fadvise hint
|
||||
// (so the kernel readahead window widens) and the periodic
|
||||
// DONTNEED page-cache eviction that bounds memory pressure
|
||||
// when the mux output is being written to the same disk.
|
||||
let mut reader = crate::io::file_sector_source::FileSectorSource::open(path)?;
|
||||
let capacity = reader.capacity_sectors();
|
||||
let disc = crate::disc::Disc::scan_image(&mut reader, capacity, &scan_opts)
|
||||
.map_err(|e| -> io::Error { e.into() })?;
|
||||
let mut disc = crate::disc::Disc::scan_image(
|
||||
&mut reader,
|
||||
capacity,
|
||||
&crate::disc::ScanOptions::default(),
|
||||
)
|
||||
.map_err(|e| -> io::Error { e.into() })?;
|
||||
// Apply the caller-resolved keys (lookup-free); decrypt_keys() then
|
||||
// yields them for the stream below. Propagate a failed application
|
||||
// rather than silently muxing an undecryptable stream.
|
||||
if !opts.unit_keys.is_empty() {
|
||||
// These UKs were already resolved AND validated by the caller
|
||||
// (the CLI's keydb loop), so no re-validation sample is needed.
|
||||
disc.decrypt_with(crate::disc::Key::Unit(opts.unit_keys.clone()), &[])
|
||||
.map_err(|e| -> io::Error { e.into() })?;
|
||||
}
|
||||
// No-key guard: if decryption is requested (not --raw) and the disc
|
||||
// is AACS-encrypted but key resolution yielded no usable key, FAIL
|
||||
// here — muxing an undecryptable stream produces ~100 MB of garbage
|
||||
// (encrypted m2ts → no TS syncs → demuxer emits nothing). A cheap
|
||||
// result-check on `decrypt_keys()`; no probe decryption needed.
|
||||
// CSS (DVD) decrypts from compiled keys (`decrypt_keys()` returns
|
||||
// `Css{..}`, never `None`), so this gate is AACS-only via `disc.aacs`.
|
||||
if aacs_key_missing(opts.raw, disc.aacs.is_some(), &disc.decrypt_keys()) {
|
||||
// Surface the disc hash (40-hex, no `0x` prefix) so the caller
|
||||
// can name the disc. Empty if scan didn't capture it.
|
||||
let disc_hash = disc
|
||||
.aacs
|
||||
.as_ref()
|
||||
.map(|a| a.disc_hash.trim_start_matches("0x").to_string())
|
||||
.unwrap_or_default();
|
||||
return Err(crate::error::Error::NoDiscKey { disc_hash }.into());
|
||||
}
|
||||
if disc.titles.is_empty() {
|
||||
return Err(crate::error::Error::NoStreams.into());
|
||||
}
|
||||
@@ -200,24 +251,50 @@ pub fn input(url: &str, opts: &InputOptions) -> io::Result<Box<dyn crate::pes::S
|
||||
}
|
||||
.into());
|
||||
}
|
||||
let title = disc.titles[idx].clone();
|
||||
// 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.
|
||||
let keys = disc.decrypt_keys();
|
||||
let format = disc.content_format;
|
||||
// ISO file: use large batch size (16 MB) — sequential read from fast storage, no bad sectors.
|
||||
// Physical drives need small batches for adaptive error handling and retry logic.
|
||||
const ISO_MUX_BATCH_SECTORS: u16 = 8192;
|
||||
let mut stream =
|
||||
DiscStream::new(Box::new(reader), title, keys, ISO_MUX_BATCH_SECTORS, format);
|
||||
if opts.raw {
|
||||
stream.set_raw();
|
||||
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]);
|
||||
}
|
||||
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
|
||||
// batches for adaptive error handling.
|
||||
const ISO_MUX_BATCH_SECTORS: u16 = 8192;
|
||||
|
||||
// Pass `DecryptKeys::None` to the decrypt decorator when
|
||||
// --raw is set — the read stack still flows through the
|
||||
// same producer+demux+parse pipeline, just without the
|
||||
// AACS / CSS step. Single highway for all ISO reads.
|
||||
let effective_keys = if opts.raw {
|
||||
crate::decrypt::DecryptKeys::None
|
||||
} else {
|
||||
keys
|
||||
};
|
||||
let stream = build_iso_pipeline(
|
||||
reader,
|
||||
title,
|
||||
effective_keys,
|
||||
ISO_MUX_BATCH_SECTORS,
|
||||
format,
|
||||
None,
|
||||
None,
|
||||
);
|
||||
Ok(Box::new(stream))
|
||||
}
|
||||
StreamUrl::M2ts { ref path } => {
|
||||
validate_file_path(path, "m2ts")?;
|
||||
let file = std::fs::File::open(path)?;
|
||||
let reader = std::io::BufReader::with_capacity(IO_BUF_SIZE, file);
|
||||
Ok(Box::new(M2tsStream::open(reader)?))
|
||||
let stream = build_m2ts_pipeline(reader)?;
|
||||
Ok(Box::new(stream))
|
||||
}
|
||||
StreamUrl::Mkv { ref path } => {
|
||||
validate_file_path(path, "mkv")?;
|
||||
@@ -250,11 +327,13 @@ pub fn output(
|
||||
// writeback) so a UHD-scale MKV mux to slow / network-attached
|
||||
// staging doesn't hit the dirty-page burst pathology that
|
||||
// sweep already side-steps. BufWriter sits on top to coalesce
|
||||
// mux's many small EBML element writes.
|
||||
// mux's many small EBML element writes. Pre-reserve the
|
||||
// target's worth of extents on Linux via fallocate(KEEP_SIZE)
|
||||
// to reduce extent fragmentation during the mux.
|
||||
let writer: Box<dyn super::WriteSeek + Send> =
|
||||
Box::new(std::io::BufWriter::with_capacity(
|
||||
IO_BUF_SIZE,
|
||||
crate::io::WritebackFile::create(path)?,
|
||||
crate::io::WritebackFile::create_with_size_hint(path, title.size_bytes)?,
|
||||
));
|
||||
Ok(Box::new(MkvStream::create(writer, title)?))
|
||||
}
|
||||
@@ -262,7 +341,7 @@ pub fn output(
|
||||
validate_file_path(path, "m2ts")?;
|
||||
let writer = std::io::BufWriter::with_capacity(
|
||||
IO_BUF_SIZE,
|
||||
crate::io::WritebackFile::create(path)?,
|
||||
crate::io::WritebackFile::create_with_size_hint(path, title.size_bytes)?,
|
||||
);
|
||||
Ok(Box::new(M2tsStream::create(writer, title)?))
|
||||
}
|
||||
@@ -279,3 +358,185 @@ pub fn output(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Demuxer-side state derived from a `DiscTitle`: the codec parser
|
||||
/// table (keyed by PID), the PID-to-track index map, and an initial
|
||||
/// `TsDemuxer` / `PsDemuxer` (whichever the content format calls
|
||||
/// for).
|
||||
type DemuxState = (
|
||||
Vec<(u16, Box<dyn super::codec::CodecParser>)>,
|
||||
Vec<(u16, usize)>,
|
||||
Option<super::ts::TsDemuxer>,
|
||||
Option<super::ps::PsDemuxer>,
|
||||
);
|
||||
|
||||
/// Build the title's codec parser table + initial `TsDemuxer` /
|
||||
/// `PsDemuxer`. Used by both the ISO and M2TS pipeline builders.
|
||||
fn build_demux_state(title: &DiscTitle, format: ContentFormat) -> DemuxState {
|
||||
let mut pids = Vec::new();
|
||||
let mut parsers = Vec::new();
|
||||
let mut pid_to_track = Vec::new();
|
||||
for (idx, s) in title.streams.iter().enumerate() {
|
||||
let (pid, codec) = match s {
|
||||
crate::disc::Stream::Video(v) => (v.pid, v.codec),
|
||||
crate::disc::Stream::Audio(a) => (a.pid, a.codec),
|
||||
crate::disc::Stream::Subtitle(s) => (s.pid, s.codec),
|
||||
};
|
||||
pids.push(pid);
|
||||
pid_to_track.push((pid, idx));
|
||||
let is_dvd_ps = matches!(format, ContentFormat::MpegPs);
|
||||
parsers.push((pid, super::codec::parser_for_codec(codec, None, is_dvd_ps)));
|
||||
}
|
||||
let (ts, ps) = match format {
|
||||
ContentFormat::MpegPs => (None, Some(super::ps::PsDemuxer::new())),
|
||||
ContentFormat::BdTs => {
|
||||
if pids.is_empty() {
|
||||
(None, None)
|
||||
} else {
|
||||
(Some(super::ts::TsDemuxer::new(&pids)), None)
|
||||
}
|
||||
}
|
||||
};
|
||||
(parsers, pid_to_track, ts, ps)
|
||||
}
|
||||
|
||||
/// Assemble the ISO mux pipeline (read+decrypt → demux → parse) for
|
||||
/// a `FileSectorSource`-backed reader. Returns the resulting
|
||||
/// `PipelinedPesStream`.
|
||||
pub fn build_iso_pipeline<S: SectorSource + Send + 'static>(
|
||||
reader: S,
|
||||
title: DiscTitle,
|
||||
keys: crate::decrypt::DecryptKeys,
|
||||
batch_sectors: u16,
|
||||
format: ContentFormat,
|
||||
halt: Option<crate::halt::Halt>,
|
||||
event_fn: Option<crate::sector::prefetched::EventFn>,
|
||||
) -> PipelinedPesStream {
|
||||
let extents = title.extents.clone();
|
||||
let decrypting =
|
||||
crate::sector::DecryptingSectorSource::new(Box::new(reader) as Box<dyn SectorSource>, keys);
|
||||
let prefetched = crate::sector::PrefetchedSectorSource::new_with_events(
|
||||
decrypting,
|
||||
extents,
|
||||
batch_sectors,
|
||||
halt.clone(),
|
||||
event_fn,
|
||||
);
|
||||
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)
|
||||
}
|
||||
|
||||
/// Assemble the M2TS file mux pipeline (read → demux → parse) for a
|
||||
/// byte-stream reader. Scans the head for FMKV header or PMT/PAT,
|
||||
/// rebuilds the title metadata, then wraps a chained reader (head +
|
||||
/// remainder) in a `BytePrefetcher` feeding the demux + parse
|
||||
/// threads.
|
||||
fn build_m2ts_pipeline<R: std::io::Read + Send + 'static>(
|
||||
mut reader: R,
|
||||
) -> io::Result<PipelinedPesStream> {
|
||||
use super::meta;
|
||||
use std::io::Read;
|
||||
|
||||
const M2TS_SCAN_BYTES: usize = 1024 * 1024;
|
||||
let mut head = vec![0u8; M2TS_SCAN_BYTES];
|
||||
let head_len = {
|
||||
let mut filled = 0;
|
||||
while filled < head.len() {
|
||||
match reader.read(&mut head[filled..])? {
|
||||
0 => break,
|
||||
n => filled += n,
|
||||
}
|
||||
}
|
||||
filled
|
||||
};
|
||||
head.truncate(head_len);
|
||||
|
||||
// Try FMKV metadata header first; fall back to PMT scan.
|
||||
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)
|
||||
};
|
||||
|
||||
// Chain: any un-consumed head bytes + the remainder of the
|
||||
// reader. The demuxer sees a contiguous M2TS byte stream.
|
||||
let remaining_head = head[head_consumed..].to_vec();
|
||||
let chained: Box<dyn Read + Send> = Box::new(io::Cursor::new(remaining_head).chain(reader));
|
||||
|
||||
let prefetcher = crate::io::byte_prefetcher::BytePrefetcher::new(
|
||||
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);
|
||||
Ok(PipelinedPesStream::new(
|
||||
demux_thread,
|
||||
demux_rx,
|
||||
title,
|
||||
parsers,
|
||||
pid_to_track,
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::aacs_key_missing;
|
||||
use crate::decrypt::DecryptKeys;
|
||||
|
||||
fn aacs_keys() -> DecryptKeys {
|
||||
DecryptKeys::Aacs {
|
||||
unit_keys: vec![(1, [0x11u8; 16])],
|
||||
read_data_key: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn css_keys() -> DecryptKeys {
|
||||
DecryptKeys::Css {
|
||||
title_key: [0u8; 5],
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encrypted_no_key_aborts() {
|
||||
// AACS disc, decryption requested, resolver yielded no key → abort.
|
||||
assert!(aacs_key_missing(false, true, &DecryptKeys::None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encrypted_with_key_proceeds() {
|
||||
// AACS disc with a usable key → proceed.
|
||||
assert!(!aacs_key_missing(false, true, &aacs_keys()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn not_encrypted_proceeds() {
|
||||
// No AACS state: unencrypted (None keys) and CSS (Css keys) both OK.
|
||||
assert!(!aacs_key_missing(false, false, &DecryptKeys::None));
|
||||
assert!(!aacs_key_missing(false, false, &css_keys()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_never_aborts() {
|
||||
// --raw skips decryption — must never hit the no-key abort, even on an
|
||||
// AACS disc with no key resolved.
|
||||
assert!(!aacs_key_missing(true, true, &DecryptKeys::None));
|
||||
assert!(!aacs_key_missing(true, true, &aacs_keys()));
|
||||
assert!(!aacs_key_missing(true, false, &DecryptKeys::None));
|
||||
}
|
||||
}
|
||||
|
||||
+164
-87
@@ -35,16 +35,35 @@ struct PesAssembler {
|
||||
pts: Option<i64>,
|
||||
dts: Option<i64>,
|
||||
active: bool,
|
||||
/// PES-header bytes still to be skipped on the next continuation
|
||||
/// packet(s). A PES header (9 + PES_header_data_length, up to 264
|
||||
/// bytes) can exceed a single 184-byte TS payload, spilling into the
|
||||
/// following continuation packet. Those spillover bytes are NOT
|
||||
/// elementary-stream data and must be skipped, or the PES start code
|
||||
/// (`00 00 01 …`) and timestamp bytes get injected into the ES — for
|
||||
/// HEVC/H264 that reads as a spurious start code / corrupt slice
|
||||
/// payload. Tracks how many header bytes remain across packets.
|
||||
header_remaining: usize,
|
||||
}
|
||||
|
||||
/// Initial capacity for a fresh PES buffer. Sized to cover the
|
||||
/// common BD-TS audio / subtitle PES outright (a few KB to ~16 KB).
|
||||
/// Video PES (typically 150–300 KB on UHD) will grow this via the
|
||||
/// standard Vec doubling, but the doublings hit the allocator's
|
||||
/// slab caches instead of the 64-page first-touch faults that the
|
||||
/// previous `Vec::with_capacity(256 * 1024)` triggered on every PES
|
||||
/// boundary.
|
||||
const PES_BUFFER_INIT_CAP: usize = 16 * 1024;
|
||||
|
||||
impl PesAssembler {
|
||||
fn new(pid: u16) -> Self {
|
||||
Self {
|
||||
pid,
|
||||
buffer: Vec::with_capacity(256 * 1024),
|
||||
buffer: Vec::with_capacity(PES_BUFFER_INIT_CAP),
|
||||
pts: None,
|
||||
dts: None,
|
||||
active: false,
|
||||
header_remaining: 0,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -55,7 +74,7 @@ impl PesAssembler {
|
||||
pid: self.pid,
|
||||
pts: self.pts,
|
||||
dts: self.dts,
|
||||
data: std::mem::replace(&mut self.buffer, Vec::with_capacity(256 * 1024)),
|
||||
data: std::mem::replace(&mut self.buffer, Vec::with_capacity(PES_BUFFER_INIT_CAP)),
|
||||
})
|
||||
} else {
|
||||
self.buffer.clear();
|
||||
@@ -124,97 +143,54 @@ impl TsDemuxer {
|
||||
}
|
||||
}
|
||||
|
||||
/// Feed a chunk of BD transport stream data. Handles non-192-byte-aligned input
|
||||
/// by buffering leftover bytes between calls. Returns completed PES packets.
|
||||
/// Feed a chunk of BD transport stream data. Handles non-192-byte-
|
||||
/// aligned input by buffering leftover bytes between calls. Returns
|
||||
/// completed PES packets.
|
||||
///
|
||||
/// 16 MiB ISO batches never divide evenly into 192-byte BD-TS
|
||||
/// packets, so every call after the first carries a ~64-byte
|
||||
/// remainder. The pre-0.24 implementation handled this by building
|
||||
/// a `combined` Vec containing remainder + the entire new input —
|
||||
/// a 16 MiB+ memcpy on every call. Now we splice exactly one
|
||||
/// boundary packet from a stack buffer, then process the rest of
|
||||
/// `data` in place. Zero-copy on the bulk path; one 192-byte copy
|
||||
/// on the boundary.
|
||||
pub fn feed(&mut self, data: &[u8]) -> Vec<PesPacket> {
|
||||
let mut completed = Vec::with_capacity(4);
|
||||
|
||||
// Prepend any remainder from previous call
|
||||
let mut combined: Vec<u8> = Vec::new();
|
||||
let work: &[u8] = if !self.remainder.is_empty() {
|
||||
combined.reserve(self.remainder.len() + data.len());
|
||||
combined.extend_from_slice(&self.remainder);
|
||||
combined.extend_from_slice(data);
|
||||
self.remainder.clear();
|
||||
&combined
|
||||
} else {
|
||||
data
|
||||
};
|
||||
|
||||
let mut offset = 0;
|
||||
|
||||
while offset + BD_TS_PACKET_SIZE <= work.len() {
|
||||
let packet = &work[offset..offset + BD_TS_PACKET_SIZE];
|
||||
offset += BD_TS_PACKET_SIZE;
|
||||
|
||||
// Skip 4-byte TP_extra_header, check sync byte
|
||||
if packet[4] != SYNC_BYTE {
|
||||
continue;
|
||||
}
|
||||
|
||||
let ts = &packet[4..]; // 188-byte standard TS packet
|
||||
|
||||
// Parse TS header
|
||||
let pid = (((ts[1] & 0x1F) as u16) << 8) | ts[2] as u16;
|
||||
let pusi = ts[1] & 0x40 != 0; // Payload Unit Start Indicator
|
||||
let adaptation = (ts[3] >> 4) & 0x03;
|
||||
|
||||
// Check if we're tracking this PID
|
||||
let idx = if (pid as usize) < self.pid_index.len() {
|
||||
self.pid_index[pid as usize]
|
||||
} else {
|
||||
-1
|
||||
};
|
||||
if idx < 0 {
|
||||
continue;
|
||||
}
|
||||
let asm = &mut self.assemblers[idx as usize];
|
||||
|
||||
// Find payload start (skip adaptation field if present)
|
||||
let payload_start = if adaptation == 0x03 || adaptation == 0x02 {
|
||||
// Adaptation field present
|
||||
let af_len = ts[4] as usize;
|
||||
if af_len > 183 {
|
||||
continue; // Malformed: AF length exceeds TS payload
|
||||
}
|
||||
5 + af_len
|
||||
} else {
|
||||
4
|
||||
};
|
||||
|
||||
if payload_start >= TS_PACKET_SIZE {
|
||||
continue;
|
||||
}
|
||||
|
||||
// No payload
|
||||
if adaptation == 0x02 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let payload = &ts[payload_start..];
|
||||
|
||||
if pusi {
|
||||
// New PES packet starts here — parse PES header
|
||||
let (pts, dts, pes_data_start) = parse_pes_header(payload);
|
||||
if let Some(prev) = asm.start(pts, dts) {
|
||||
completed.push(prev);
|
||||
}
|
||||
if pes_data_start < payload.len() {
|
||||
asm.push(&payload[pes_data_start..]);
|
||||
}
|
||||
} else {
|
||||
// Continuation of current PES packet
|
||||
asm.push(payload);
|
||||
// Boundary packet: if a partial packet was left from the last
|
||||
// call, complete it from the head of `data` without touching
|
||||
// the rest of `data`.
|
||||
if !self.remainder.is_empty() {
|
||||
let need = BD_TS_PACKET_SIZE - self.remainder.len();
|
||||
if data.len() < need {
|
||||
// Still not a full packet — accumulate and wait.
|
||||
self.remainder.extend_from_slice(data);
|
||||
return completed;
|
||||
}
|
||||
let mut boundary = [0u8; BD_TS_PACKET_SIZE];
|
||||
boundary[..self.remainder.len()].copy_from_slice(&self.remainder);
|
||||
boundary[self.remainder.len()..].copy_from_slice(&data[..need]);
|
||||
self.remainder.clear();
|
||||
self.process_packet(&boundary, &mut completed);
|
||||
offset = need;
|
||||
}
|
||||
|
||||
// Save leftover bytes for next call (cap at one packet to prevent unbounded growth)
|
||||
if offset < work.len() {
|
||||
let leftover = &work[offset..];
|
||||
// Aligned-packets fast path — reads directly out of `data`.
|
||||
while offset + BD_TS_PACKET_SIZE <= data.len() {
|
||||
let packet = &data[offset..offset + BD_TS_PACKET_SIZE];
|
||||
offset += BD_TS_PACKET_SIZE;
|
||||
self.process_packet(packet, &mut completed);
|
||||
}
|
||||
|
||||
// Save leftover bytes for next call (cap at one packet to
|
||||
// prevent unbounded growth on a desynchronised stream).
|
||||
if offset < data.len() {
|
||||
let leftover = &data[offset..];
|
||||
if leftover.len() < BD_TS_PACKET_SIZE {
|
||||
self.remainder.extend_from_slice(leftover);
|
||||
} else {
|
||||
// More than one full packet leftover — something is wrong, discard
|
||||
self.remainder.clear();
|
||||
}
|
||||
}
|
||||
@@ -222,6 +198,88 @@ impl TsDemuxer {
|
||||
completed
|
||||
}
|
||||
|
||||
/// Demux a single 192-byte BD-TS packet (4-byte TP_extra_header +
|
||||
/// 188-byte TS). Routes payload bytes into the per-PID
|
||||
/// `PesAssembler`; completed PES packets are pushed onto
|
||||
/// `completed` so the caller's allocation amortises across the
|
||||
/// batch.
|
||||
fn process_packet(&mut self, packet: &[u8], completed: &mut Vec<PesPacket>) {
|
||||
// Sync byte check skips malformed packets.
|
||||
if packet[4] != SYNC_BYTE {
|
||||
return;
|
||||
}
|
||||
let ts = &packet[4..]; // 188-byte standard TS packet
|
||||
|
||||
let pid = (((ts[1] & 0x1F) as u16) << 8) | ts[2] as u16;
|
||||
let pusi = ts[1] & 0x40 != 0; // Payload Unit Start Indicator
|
||||
let adaptation = (ts[3] >> 4) & 0x03;
|
||||
|
||||
let idx = if (pid as usize) < self.pid_index.len() {
|
||||
self.pid_index[pid as usize]
|
||||
} else {
|
||||
-1
|
||||
};
|
||||
if idx < 0 {
|
||||
return;
|
||||
}
|
||||
let asm = &mut self.assemblers[idx as usize];
|
||||
|
||||
let payload_start = if adaptation == 0x03 || adaptation == 0x02 {
|
||||
let af_len = ts[4] as usize;
|
||||
if af_len > 183 {
|
||||
return; // Malformed: AF length exceeds TS payload
|
||||
}
|
||||
5 + af_len
|
||||
} else {
|
||||
4
|
||||
};
|
||||
|
||||
if payload_start >= TS_PACKET_SIZE {
|
||||
return;
|
||||
}
|
||||
// adaptation == 0x02 → AF only, no payload.
|
||||
if adaptation == 0x02 {
|
||||
return;
|
||||
}
|
||||
|
||||
let payload = &ts[payload_start..];
|
||||
|
||||
if pusi {
|
||||
// `header_len` is the FULL (uncapped) PES-header length:
|
||||
// 0 = malformed (payload is not a PES start), else 6/9+N.
|
||||
let (pts, dts, header_len) = parse_pes_header(payload);
|
||||
if let Some(prev) = asm.start(pts, dts) {
|
||||
completed.push(prev);
|
||||
}
|
||||
if header_len == 0 {
|
||||
// PUSI packet whose payload is not a valid PES start. Do
|
||||
// NOT push it — those bytes are not elementary-stream data
|
||||
// and would inject a spurious start code / garbage.
|
||||
asm.header_remaining = 0;
|
||||
} else if header_len <= payload.len() {
|
||||
// Header fits in this packet (the common case).
|
||||
asm.header_remaining = 0;
|
||||
if header_len < payload.len() {
|
||||
asm.push(&payload[header_len..]);
|
||||
}
|
||||
} else {
|
||||
// Header spills past this packet — skip the remainder on
|
||||
// the following continuation packet(s).
|
||||
asm.header_remaining = header_len - payload.len();
|
||||
}
|
||||
} else if asm.header_remaining > 0 {
|
||||
// Continuation packet still inside a PES header that spanned
|
||||
// the boundary — consume header bytes before any ES data.
|
||||
let skip = asm.header_remaining.min(payload.len());
|
||||
asm.header_remaining -= skip;
|
||||
if skip < payload.len() {
|
||||
asm.push(&payload[skip..]);
|
||||
}
|
||||
} else {
|
||||
asm.push(payload);
|
||||
}
|
||||
}
|
||||
|
||||
/// Flush all assemblers, returning any remaining PES packets.
|
||||
pub fn flush(&mut self) -> Vec<PesPacket> {
|
||||
let mut completed = Vec::new();
|
||||
@@ -235,7 +293,14 @@ impl TsDemuxer {
|
||||
}
|
||||
|
||||
/// Parse a PES packet header, extracting PTS and DTS.
|
||||
/// Returns (pts, dts, offset_to_elementary_stream_data).
|
||||
///
|
||||
/// Returns `(pts, dts, header_len)` where `header_len` is the FULL,
|
||||
/// UNCAPPED PES-header length in bytes (`9 + PES_header_data_length`, or
|
||||
/// 6 for stream IDs without the standard extension). `0` signals the
|
||||
/// payload is not a valid PES start (malformed / too short). The caller
|
||||
/// must treat `header_len` as bytes-to-skip and carry any remainder past
|
||||
/// this packet's payload into the next continuation packet — the header
|
||||
/// can exceed one TS payload, and the spillover is header, not ES data.
|
||||
fn parse_pes_header(data: &[u8]) -> (Option<i64>, Option<i64>, usize) {
|
||||
// PES packet: 00 00 01 [stream_id] [length:2] [flags...]
|
||||
if data.len() < 9 || data[0] != 0x00 || data[1] != 0x00 || data[2] != 0x01 {
|
||||
@@ -263,7 +328,10 @@ fn parse_pes_header(data: &[u8]) -> (Option<i64>, Option<i64>, usize) {
|
||||
|
||||
let pts_dts_flags = (data[7] >> 6) & 0x03;
|
||||
let header_data_len = data[8] as usize;
|
||||
let data_start = (9 + header_data_len).min(data.len());
|
||||
// Full, uncapped header length. PTS/DTS (if present) live in the
|
||||
// first ~19 bytes, always within this packet's payload, so they parse
|
||||
// here; only the *skip* length may extend into the next packet.
|
||||
let header_len = 9 + header_data_len;
|
||||
|
||||
let mut pts = None;
|
||||
let mut dts = None;
|
||||
@@ -275,7 +343,7 @@ fn parse_pes_header(data: &[u8]) -> (Option<i64>, Option<i64>, usize) {
|
||||
dts = parse_timestamp(&data[14..19]);
|
||||
}
|
||||
|
||||
(pts, dts, data_start)
|
||||
(pts, dts, header_len)
|
||||
}
|
||||
|
||||
/// Parse a 5-byte PTS/DTS timestamp (33 bits in 90kHz).
|
||||
@@ -379,10 +447,19 @@ pub fn scan_streams(data: &[u8]) -> Option<Vec<crate::disc::Stream>> {
|
||||
|
||||
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;
|
||||
let end = pmt_start + 3 + section_len - 4;
|
||||
// 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];
|
||||
|
||||
+272
-35
@@ -42,28 +42,39 @@ impl<W: Write> TsMuxer<W> {
|
||||
/// 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.
|
||||
pub fn write_frame(&mut self, track: usize, pts_ns: i64, data: &[u8]) -> io::Result<()> {
|
||||
pub fn write_frame(
|
||||
&mut self,
|
||||
track: usize,
|
||||
pts_ns: i64,
|
||||
keyframe: bool,
|
||||
data: &[u8],
|
||||
) -> io::Result<()> {
|
||||
if track >= self.pids.len() {
|
||||
return Ok(()); // unknown track, skip
|
||||
}
|
||||
let base = *self.base_pts_ns.get_or_insert(pts_ns);
|
||||
let pts_ns = pts_ns - base;
|
||||
|
||||
let pid = self.pids[track];
|
||||
let is_video = (0x1011..=0x101F).contains(&pid);
|
||||
|
||||
// For video: convert length-prefixed NALUs to Annex B (start codes)
|
||||
// On first keyframe, prepend parameter sets from codec_private
|
||||
// Drop non-key video before any keyframe — decoder has no IDR or
|
||||
// parameter sets to anchor on.
|
||||
if is_video && !keyframe && !self.params_written[track] {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let base = *self.base_pts_ns.get_or_insert(pts_ns);
|
||||
let pts_ns = pts_ns - 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() {
|
||||
let mut annex_b = Vec::new();
|
||||
// Prepend codec_private parameter sets on first keyframe
|
||||
if !self.params_written[track] {
|
||||
if keyframe && !self.params_written[track] {
|
||||
if let Some(ref cp) = self.codec_privates[track] {
|
||||
if let Some(params) = hvcc_to_annex_b(cp) {
|
||||
annex_b.extend_from_slice(¶ms);
|
||||
self.params_written[track] = true;
|
||||
}
|
||||
}
|
||||
self.params_written[track] = true;
|
||||
}
|
||||
annex_b.extend_from_slice(&length_prefixed_to_annex_b(data));
|
||||
annex_b
|
||||
@@ -85,8 +96,24 @@ impl<W: Write> TsMuxer<W> {
|
||||
let mut first = true;
|
||||
while offset < pes_packet.len() {
|
||||
let remaining = pes_packet.len() - offset;
|
||||
let payload_len = remaining.min(TS_PAYLOAD);
|
||||
let need_stuffing = payload_len < TS_PAYLOAD;
|
||||
|
||||
// Invariant: TP_extra(4) + TS_header(4) + AF(af_bytes) + payload(payload_len) = 192,
|
||||
// i.e. af_bytes + payload_len = TS_PAYLOAD (184).
|
||||
// RAI on first packet of a keyframe video PES requires AF with flags=0x40.
|
||||
let want_rai = first && keyframe && is_video;
|
||||
|
||||
// Pick payload_len and af_bytes per case.
|
||||
let (af_bytes, payload_len): (usize, usize) = if want_rai {
|
||||
// Minimum AF = 2 bytes (length=1, flags=0x40). Payload caps at 182.
|
||||
let max_payload = TS_PAYLOAD - 2;
|
||||
let p = remaining.min(max_payload);
|
||||
(TS_PAYLOAD - p, p)
|
||||
} else if remaining >= TS_PAYLOAD {
|
||||
(0, TS_PAYLOAD) // no AF, full payload
|
||||
} else {
|
||||
// Stuffing-only AF, payload = remaining.
|
||||
(TS_PAYLOAD - remaining, remaining)
|
||||
};
|
||||
|
||||
// TP_extra_header (4 bytes — arrival time, set to 0)
|
||||
let tp_extra = [0u8; 4];
|
||||
@@ -102,38 +129,45 @@ impl<W: Write> TsMuxer<W> {
|
||||
ts_header[1] |= 0x40; // PUSI
|
||||
}
|
||||
ts_header[2] = pid as u8;
|
||||
ts_header[3] = 0x10 | cc; // no adaptation, has payload
|
||||
ts_header[3] = if af_bytes > 0 {
|
||||
0x30 | cc // AF + payload
|
||||
} else {
|
||||
0x10 | cc // payload only
|
||||
};
|
||||
|
||||
if need_stuffing {
|
||||
// Adaptation field for stuffing
|
||||
let stuff_len = TS_PAYLOAD - payload_len;
|
||||
ts_header[3] = 0x30 | cc; // adaptation + payload
|
||||
self.writer.write_all(&tp_extra)?;
|
||||
self.writer.write_all(&ts_header)?;
|
||||
|
||||
self.writer.write_all(&tp_extra)?;
|
||||
self.writer.write_all(&ts_header)?;
|
||||
|
||||
// Write adaptation field: length byte + flags byte + 0xFF padding
|
||||
// stuff_len == 1: AF length = 0 (just the length byte, no flags)
|
||||
// stuff_len >= 2: AF length = stuff_len-1, flags = 0, rest 0xFF
|
||||
if af_bytes > 0 {
|
||||
static STUFF_FF: [u8; 184] = [0xFF; 184];
|
||||
if stuff_len == 1 {
|
||||
self.writer.write_all(&[0u8])?; // adaptation_field_length = 0
|
||||
if want_rai {
|
||||
// RAI AF: length byte + flags(0x40) + (af_bytes - 2) stuffing.
|
||||
let af_len_field = (af_bytes - 1) as u8;
|
||||
self.writer.write_all(&[af_len_field])?;
|
||||
self.writer.write_all(&[0x40u8])?;
|
||||
let stuff = af_bytes - 2;
|
||||
if stuff > 0 {
|
||||
self.writer.write_all(&STUFF_FF[..stuff])?;
|
||||
}
|
||||
} else {
|
||||
self.writer.write_all(&[(stuff_len - 1) as u8])?; // AF length
|
||||
self.writer.write_all(&[0u8])?; // flags
|
||||
if stuff_len > 2 {
|
||||
self.writer.write_all(&STUFF_FF[..stuff_len - 2])?;
|
||||
// Stuffing-only AF.
|
||||
// af_bytes == 1: length=0, no flags.
|
||||
// af_bytes >= 2: length = af_bytes-1, flags=0, rest 0xFF.
|
||||
if af_bytes == 1 {
|
||||
self.writer.write_all(&[0u8])?;
|
||||
} else {
|
||||
self.writer.write_all(&[(af_bytes - 1) as u8])?;
|
||||
self.writer.write_all(&[0u8])?;
|
||||
if af_bytes > 2 {
|
||||
self.writer.write_all(&STUFF_FF[..af_bytes - 2])?;
|
||||
}
|
||||
}
|
||||
}
|
||||
self.writer
|
||||
.write_all(&pes_packet[offset..offset + payload_len])?;
|
||||
} else {
|
||||
self.writer.write_all(&tp_extra)?;
|
||||
self.writer.write_all(&ts_header)?;
|
||||
self.writer
|
||||
.write_all(&pes_packet[offset..offset + payload_len])?;
|
||||
}
|
||||
|
||||
self.writer
|
||||
.write_all(&pes_packet[offset..offset + payload_len])?;
|
||||
|
||||
offset += payload_len;
|
||||
first = false;
|
||||
}
|
||||
@@ -257,3 +291,206 @@ fn length_prefixed_to_annex_b(data: &[u8]) -> Vec<u8> {
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const BD_PACKET_SIZE: usize = 192;
|
||||
const VIDEO_PID: u16 = 0x1011;
|
||||
|
||||
/// Parsed BD-TS packet (192 bytes total: 4 TP_extra + 4 TS header + 184 body).
|
||||
struct TsPacket {
|
||||
pid: u16,
|
||||
pusi: bool,
|
||||
#[allow(dead_code)]
|
||||
cc: u8,
|
||||
/// Adaptation field body (length byte stripped) when present.
|
||||
af: Option<Vec<u8>>,
|
||||
/// Payload bytes (after AF, if any).
|
||||
payload: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Walk 192-byte BD-TS packets.
|
||||
fn parse_bd_ts(buf: &[u8]) -> Vec<TsPacket> {
|
||||
let mut out = Vec::new();
|
||||
for chunk in buf.chunks(BD_PACKET_SIZE) {
|
||||
if chunk.len() != BD_PACKET_SIZE {
|
||||
break;
|
||||
}
|
||||
// Skip TP_extra_header (4 bytes), parse TS header.
|
||||
let h = &chunk[4..];
|
||||
assert_eq!(h[0], 0x47, "bad sync byte");
|
||||
let pusi = (h[1] & 0x40) != 0;
|
||||
let pid = (((h[1] & 0x1F) as u16) << 8) | h[2] as u16;
|
||||
let afc = (h[3] >> 4) & 0x03;
|
||||
let cc = h[3] & 0x0F;
|
||||
let body = &h[4..]; // 184 bytes
|
||||
|
||||
let (af, payload) = match afc {
|
||||
0b01 => (None, body.to_vec()),
|
||||
0b11 => {
|
||||
let af_len = body[0] as usize;
|
||||
let af_body = body[1..1 + af_len].to_vec();
|
||||
let payload = body[1 + af_len..].to_vec();
|
||||
(Some(af_body), payload)
|
||||
}
|
||||
0b10 => {
|
||||
let af_len = body[0] as usize;
|
||||
(Some(body[1..1 + af_len].to_vec()), Vec::new())
|
||||
}
|
||||
_ => (None, Vec::new()),
|
||||
};
|
||||
|
||||
out.push(TsPacket {
|
||||
pid,
|
||||
pusi,
|
||||
cc,
|
||||
af,
|
||||
payload,
|
||||
});
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Build a fake HEVC NAL with a 4-byte length prefix.
|
||||
/// nal_type=19/20 are IDR; 1 is non-key (TRAIL_N/R).
|
||||
fn fake_hevc_nal(nal_type: u8, body_len: usize) -> Vec<u8> {
|
||||
let mut nal = Vec::with_capacity(2 + body_len);
|
||||
// 2-byte NAL header: forbidden_zero(1)=0 | nal_unit_type(6) | layer_id(6)=0 | tid_plus1(3)=1
|
||||
nal.push((nal_type & 0x3F) << 1);
|
||||
nal.push(0x01);
|
||||
for i in 0..body_len {
|
||||
nal.push((i & 0xFF) as u8);
|
||||
}
|
||||
let mut framed = Vec::with_capacity(4 + nal.len());
|
||||
framed.extend_from_slice(&(nal.len() as u32).to_be_bytes());
|
||||
framed.extend_from_slice(&nal);
|
||||
framed
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keyframe_param_threads_through() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
{
|
||||
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
|
||||
let idr = fake_hevc_nal(19, 100);
|
||||
mux.write_frame(0, 0, true, &idr).unwrap();
|
||||
let p = fake_hevc_nal(1, 80);
|
||||
mux.write_frame(0, 41_000_000, false, &p).unwrap();
|
||||
mux.finish().unwrap();
|
||||
}
|
||||
assert!(!sink.is_empty());
|
||||
let packets = parse_bd_ts(&sink);
|
||||
assert!(packets.iter().any(|p| p.pid == VIDEO_PID && p.pusi));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rai_set_on_first_packet_of_keyframe_pes() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
{
|
||||
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
|
||||
let idr = fake_hevc_nal(19, 200);
|
||||
mux.write_frame(0, 0, true, &idr).unwrap();
|
||||
mux.finish().unwrap();
|
||||
}
|
||||
let packets = parse_bd_ts(&sink);
|
||||
let first_pusi = packets
|
||||
.iter()
|
||||
.find(|p| p.pid == VIDEO_PID && p.pusi)
|
||||
.expect("video PUSI packet exists");
|
||||
let af = first_pusi.af.as_ref().expect("AF present on keyframe PES");
|
||||
assert!(!af.is_empty(), "AF body has flags byte");
|
||||
assert_eq!(af[0] & 0x40, 0x40, "RAI bit set");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rai_clear_on_non_keyframe_pes() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
{
|
||||
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
|
||||
let idr = fake_hevc_nal(19, 100);
|
||||
mux.write_frame(0, 0, true, &idr).unwrap();
|
||||
let p = fake_hevc_nal(1, 100);
|
||||
mux.write_frame(0, 41_000_000, false, &p).unwrap();
|
||||
mux.finish().unwrap();
|
||||
}
|
||||
let packets = parse_bd_ts(&sink);
|
||||
// Second PUSI packet on the video PID belongs to the non-key frame.
|
||||
let pusi_video: Vec<&TsPacket> = packets
|
||||
.iter()
|
||||
.filter(|p| p.pid == VIDEO_PID && p.pusi)
|
||||
.collect();
|
||||
assert!(pusi_video.len() >= 2, "two PUSI packets expected");
|
||||
let second = pusi_video[1];
|
||||
match &second.af {
|
||||
None => {}
|
||||
Some(af) if af.is_empty() => {} // length=0 case
|
||||
Some(af) => assert_eq!(af[0] & 0x40, 0, "RAI must be clear on non-key PES"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codec_private_prepended_only_on_first_keyframe() {
|
||||
// Build a minimal hvcC with one recognizable NAL.
|
||||
let marker: &[u8] = &[0xDE, 0xAD, 0xBE, 0xEF, 0xCA, 0xFE];
|
||||
let mut hvcc = vec![0u8; 22];
|
||||
hvcc.push(1); // numArrays
|
||||
hvcc.push(32); // VPS NAL type byte (high bits arbitrary)
|
||||
hvcc.extend_from_slice(&1u16.to_be_bytes()); // numNalus
|
||||
hvcc.extend_from_slice(&(marker.len() as u16).to_be_bytes());
|
||||
hvcc.extend_from_slice(marker);
|
||||
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
{
|
||||
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
|
||||
mux.set_codec_private(0, hvcc);
|
||||
// Non-IDR before any IDR: should be dropped.
|
||||
let p = fake_hevc_nal(1, 50);
|
||||
mux.write_frame(0, 0, false, &p).unwrap();
|
||||
// IDR: should carry codec_private NALs prepended.
|
||||
let idr = fake_hevc_nal(19, 50);
|
||||
mux.write_frame(0, 41_000_000, true, &idr).unwrap();
|
||||
mux.finish().unwrap();
|
||||
}
|
||||
let packets = parse_bd_ts(&sink);
|
||||
// Concatenate all video PID payloads in emission order.
|
||||
let video_bytes: Vec<u8> = packets
|
||||
.iter()
|
||||
.filter(|p| p.pid == VIDEO_PID)
|
||||
.flat_map(|p| p.payload.clone())
|
||||
.collect();
|
||||
// marker bytes must appear in the stream (codec_private was prepended).
|
||||
let pos_marker = video_bytes
|
||||
.windows(marker.len())
|
||||
.position(|w| w == marker)
|
||||
.expect("codec_private marker bytes present in TS payload");
|
||||
// Find IDR body byte (0x26 = (19<<1)). pos_idr must be AFTER marker.
|
||||
let idr_header = (19u8 << 1) & 0x7E;
|
||||
let pos_idr = video_bytes
|
||||
.iter()
|
||||
.position(|&b| b == idr_header)
|
||||
.expect("IDR NAL header present in TS payload");
|
||||
assert!(
|
||||
pos_marker < pos_idr,
|
||||
"codec_private must precede IDR in TS payload"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_key_before_first_keyframe_dropped() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
{
|
||||
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
|
||||
let p = fake_hevc_nal(1, 80);
|
||||
mux.write_frame(0, 0, false, &p).unwrap();
|
||||
mux.finish().unwrap();
|
||||
}
|
||||
// Nothing should be emitted for that PID.
|
||||
let packets = parse_bd_ts(&sink);
|
||||
assert!(
|
||||
!packets.iter().any(|p| p.pid == VIDEO_PID),
|
||||
"non-key before first keyframe must be dropped"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+17
@@ -17,6 +17,10 @@ pub struct PesFrame {
|
||||
pub keyframe: bool,
|
||||
/// Raw elementary stream data (NAL units, audio samples, etc).
|
||||
pub data: Vec<u8>,
|
||||
/// Optional duration in nanoseconds. In-memory only; not part of
|
||||
/// the on-wire serialization. Currently set by the PGS parser so
|
||||
/// the MKV muxer can emit `BlockDuration`.
|
||||
pub duration_ns: Option<u64>,
|
||||
}
|
||||
|
||||
impl PesFrame {
|
||||
@@ -64,6 +68,7 @@ impl PesFrame {
|
||||
pts,
|
||||
keyframe,
|
||||
data,
|
||||
duration_ns: None,
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -74,6 +79,7 @@ impl PesFrame {
|
||||
pts: frame.pts_ns,
|
||||
keyframe: frame.keyframe,
|
||||
data: frame.data,
|
||||
duration_ns: frame.duration_ns,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -115,6 +121,16 @@ pub trait Stream: Send {
|
||||
fn headers_ready(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// Cumulative count of read errors the stream skipped past (e.g.
|
||||
/// zero-filled bad sectors on a live drive). Default `0` for
|
||||
/// streams that don't have a notion of skip-on-error (file ISO,
|
||||
/// network, stdio, the pipeline highway, etc.); concrete impls
|
||||
/// with adaptive retry (`DiscStream` on the drive single-pass
|
||||
/// path) override.
|
||||
fn errors(&self) -> u64 {
|
||||
0
|
||||
}
|
||||
}
|
||||
|
||||
/// Wraps any output stream and counts bytes written.
|
||||
@@ -183,6 +199,7 @@ mod tests {
|
||||
pts,
|
||||
keyframe: track == 0 && pts == 0,
|
||||
data: vec![track as u8, (pts & 0xff) as u8, 0xAA],
|
||||
duration_ns: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
//! Linux `statfs64`-based filesystem type detection.
|
||||
//!
|
||||
//! Recognised local-FS magics: ext2/3/4, xfs, btrfs, tmpfs. NFS is the
|
||||
//! one network FS this layer cares about (the buffering decision keys
|
||||
//! off it). Anything else maps to [`FsType::Unknown`].
|
||||
|
||||
use std::ffi::CString;
|
||||
use std::os::unix::ffi::OsStrExt;
|
||||
use std::os::unix::io::RawFd;
|
||||
use std::path::Path;
|
||||
|
||||
use super::FsType;
|
||||
|
||||
// Magic numbers from `<linux/magic.h>`. Kept literal here so we don't
|
||||
// depend on libc exposing each one — only `NFS_SUPER_MAGIC` is
|
||||
// guaranteed to be present across libc / musl revisions.
|
||||
const EXT2_SUPER_MAGIC: i64 = 0xEF53;
|
||||
const XFS_SUPER_MAGIC: i64 = 0x5846_5342;
|
||||
const BTRFS_SUPER_MAGIC: i64 = 0x9123_683E;
|
||||
const TMPFS_MAGIC: i64 = 0x0102_1994;
|
||||
|
||||
/// Classify an `f_type` magic from `statfs`/`fstatfs`. Single source
|
||||
/// of truth for the magic comparisons used by both the path-based
|
||||
/// (`detect_impl`) and fd-based (`detect_fd_impl`) entry points.
|
||||
///
|
||||
/// Cast to `i64` because `statfs::f_type` is signed `__fsword_t` on
|
||||
/// glibc and unsigned `c_ulong` on musl; a portable comparison needs a
|
||||
/// common type. On glibc x86_64 both already are i64 — clippy flags
|
||||
/// the cast as unnecessary on that target only, but we need it for
|
||||
/// musl.
|
||||
#[allow(clippy::unnecessary_cast)]
|
||||
fn classify_f_type(f_type: i64) -> FsType {
|
||||
let nfs_magic = libc::NFS_SUPER_MAGIC as i64;
|
||||
if f_type == nfs_magic {
|
||||
return FsType::Nfs;
|
||||
}
|
||||
match f_type {
|
||||
EXT2_SUPER_MAGIC | XFS_SUPER_MAGIC | BTRFS_SUPER_MAGIC | TMPFS_MAGIC => FsType::Local,
|
||||
_ => FsType::Unknown,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn detect_impl(path: &Path) -> FsType {
|
||||
let cpath = match CString::new(path.as_os_str().as_bytes()) {
|
||||
Ok(c) => c,
|
||||
Err(_) => return FsType::Unknown,
|
||||
};
|
||||
// `statfs64` is repr(C); zeroing is the documented init pattern for
|
||||
// the kernel uapi struct.
|
||||
let mut buf: libc::statfs64 = unsafe { std::mem::zeroed() };
|
||||
let rc = unsafe { libc::statfs64(cpath.as_ptr(), &mut buf) };
|
||||
if rc != 0 {
|
||||
return FsType::Unknown;
|
||||
}
|
||||
#[allow(clippy::unnecessary_cast)]
|
||||
classify_f_type(buf.f_type as i64)
|
||||
}
|
||||
|
||||
/// fd-based variant of [`detect_impl`]. Same return semantics. Used
|
||||
/// by the writeback pipeline, which knows the open `File` but not its
|
||||
/// original path.
|
||||
pub(super) fn detect_fd_impl(fd: RawFd) -> FsType {
|
||||
let mut buf: libc::statfs = unsafe { std::mem::zeroed() };
|
||||
let rc = unsafe { libc::fstatfs(fd, &mut buf) };
|
||||
if rc != 0 {
|
||||
return FsType::Unknown;
|
||||
}
|
||||
#[allow(clippy::unnecessary_cast)]
|
||||
classify_f_type(buf.f_type as i64)
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
//! macOS `statfs`-based filesystem type detection.
|
||||
//!
|
||||
//! macOS exposes a textual `f_fstypename` field (e.g. `"apfs"`, `"hfs"`,
|
||||
//! `"nfs"`, `"smbfs"`) on its `statfs` struct, which is far more
|
||||
//! reliable than chasing magic numbers. Anything starting with `"nfs"`
|
||||
//! counts as NFS; anything else recognised maps to `Local`; otherwise
|
||||
//! `Unknown`.
|
||||
|
||||
use std::ffi::CString;
|
||||
use std::os::unix::ffi::OsStrExt;
|
||||
use std::path::Path;
|
||||
|
||||
use super::FsType;
|
||||
|
||||
pub(super) fn detect_impl(path: &Path) -> FsType {
|
||||
let cpath = match CString::new(path.as_os_str().as_bytes()) {
|
||||
Ok(c) => c,
|
||||
Err(_) => return FsType::Unknown,
|
||||
};
|
||||
let mut buf: libc::statfs = unsafe { std::mem::zeroed() };
|
||||
let rc = unsafe { libc::statfs(cpath.as_ptr(), &mut buf) };
|
||||
if rc != 0 {
|
||||
return FsType::Unknown;
|
||||
}
|
||||
// `f_fstypename` is a NUL-terminated C string of length MFSTYPENAMELEN.
|
||||
// SAFETY: libc guarantees the field is initialised to a NUL-terminated
|
||||
// string by a successful statfs.
|
||||
let name_ptr = buf.f_fstypename.as_ptr();
|
||||
let cstr = unsafe { std::ffi::CStr::from_ptr(name_ptr) };
|
||||
let name = cstr.to_bytes();
|
||||
if name.starts_with(b"nfs") {
|
||||
return FsType::Nfs;
|
||||
}
|
||||
// Recognised local types. SMB is not NFS but the buffering-policy
|
||||
// outcome on macOS is the same as for any other local FS — there's
|
||||
// no `WritebackFile` machinery to opt out of on this OS.
|
||||
match name {
|
||||
b"apfs" | b"hfs" | b"exfat" | b"msdos" | b"tmpfs" | b"smbfs" | b"webdav" => FsType::Local,
|
||||
_ => FsType::Unknown,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
//! Filesystem-type detection.
|
||||
//!
|
||||
//! The buffering architecture (Phase 2) selects different output sinks
|
||||
//! for local vs network filesystems: NFS gets the adaptive
|
||||
//! `WritebackFile` machinery on Linux; local disks get `LocalFileSink`
|
||||
//! and rely on the kernel's default writeback policy. This module
|
||||
//! provides the construction-site primitive that picks which one.
|
||||
//!
|
||||
//! Per the per-OS file-split convention, the actual `statfs` call lives
|
||||
//! in the matching platform file (`linux.rs`, `macos.rs`, `windows.rs`,
|
||||
//! `other.rs`); this `mod.rs` exposes only the cross-platform enum and
|
||||
//! the `detect` entry point.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
/// What kind of filesystem a path lives on, to the extent we can tell
|
||||
/// cheaply at construction time.
|
||||
///
|
||||
/// `Unknown` is the fail-open default: a misdetection here should not
|
||||
/// be load-bearing for correctness, only for the choice of sink (and
|
||||
/// hence buffering policy). Callers that need a binary local/non-local
|
||||
/// answer should treat `Unknown` as "probably local".
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum FsType {
|
||||
/// A local on-disk filesystem (ext4, xfs, btrfs, apfs, ntfs, …).
|
||||
Local,
|
||||
/// A network filesystem with NFS-like semantics. The current
|
||||
/// detector lumps SMB / UNC into this on Windows because the
|
||||
/// buffering policy outcome is the same.
|
||||
Nfs,
|
||||
/// `statfs` failed, the filesystem type is not on our recognised
|
||||
/// list, or we are on an OS without a real implementation.
|
||||
Unknown,
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
mod linux;
|
||||
#[cfg(target_os = "macos")]
|
||||
mod macos;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
|
||||
mod other;
|
||||
#[cfg(target_os = "windows")]
|
||||
mod windows;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
use linux::detect_impl;
|
||||
#[cfg(target_os = "macos")]
|
||||
use macos::detect_impl;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
|
||||
use other::detect_impl;
|
||||
#[cfg(target_os = "windows")]
|
||||
use windows::detect_impl;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
use linux::detect_fd_impl;
|
||||
|
||||
/// Best-effort classification of the filesystem under `path`.
|
||||
///
|
||||
/// Falls back to [`FsType::Unknown`] on any syscall error or unrecognised
|
||||
/// filesystem signature. Never panics. Never blocks beyond the cost of
|
||||
/// a single `statfs(2)` (Unix) or a string check (Windows).
|
||||
pub fn detect(path: &Path) -> FsType {
|
||||
detect_impl(path)
|
||||
}
|
||||
|
||||
/// fd-based classification. Same return semantics as [`detect`], but
|
||||
/// takes a `RawFd` so callers that only have an open file (notably
|
||||
/// [`crate::io::writeback::WritebackPipeline`]) don't have to
|
||||
/// round-trip through the path.
|
||||
///
|
||||
/// Only implemented on Linux; other platforms return
|
||||
/// [`FsType::Unknown`] (none of them have a writeback policy that
|
||||
/// keys off this classification today).
|
||||
#[cfg(target_os = "linux")]
|
||||
pub fn detect_fd(fd: std::os::unix::io::RawFd) -> FsType {
|
||||
detect_fd_impl(fd)
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
#[allow(dead_code)] // API parity with the linux impl; callers cfg-gate.
|
||||
pub fn detect_fd(_fd: i32) -> FsType {
|
||||
FsType::Unknown
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn detect_does_not_panic_on_missing_path() {
|
||||
// Non-existent path should fall through to Unknown, not panic.
|
||||
let p = std::path::Path::new("/this/path/should/not/exist/freemkv-test");
|
||||
let r = detect(p);
|
||||
// We don't assert == Unknown because Windows' heuristic looks at
|
||||
// the leading bytes and might still classify; just confirm the
|
||||
// call returns rather than panicking.
|
||||
let _ = r;
|
||||
}
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
#[test]
|
||||
fn macos_tmp_is_local() {
|
||||
// `/tmp` on macOS dev rigs is APFS via the symlink to
|
||||
// `/private/tmp`. Either way, never NFS.
|
||||
let r = detect(std::path::Path::new("/tmp"));
|
||||
assert!(
|
||||
matches!(r, FsType::Local | FsType::Unknown),
|
||||
"expected Local or Unknown for /tmp on macOS, got {r:?}",
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
#[test]
|
||||
fn linux_tmp_is_local_or_unknown() {
|
||||
// `/tmp` is tmpfs on most distros (which we recognise) but
|
||||
// could be ext4 on others. NFS would be unusual.
|
||||
let r = detect(std::path::Path::new("/tmp"));
|
||||
assert!(
|
||||
matches!(r, FsType::Local | FsType::Unknown),
|
||||
"expected Local or Unknown for /tmp on Linux, got {r:?}",
|
||||
);
|
||||
}
|
||||
|
||||
/// Real NFS exercise needs an actual NFS mount and isn't available
|
||||
/// in CI. Kept here as a manual probe.
|
||||
#[test]
|
||||
#[ignore = "needs an NFS mount path (e.g. /mnt/nfs/...) to validate live"]
|
||||
fn nfs_path_classified_as_nfs() {
|
||||
// Operator passes the mount as FREEMKV_NFS_PROBE; gated behind
|
||||
// `--ignored` because there's no portable NFS path.
|
||||
let p = std::env::var("FREEMKV_NFS_PROBE").expect("set FREEMKV_NFS_PROBE");
|
||||
assert_eq!(detect(std::path::Path::new(&p)), FsType::Nfs);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
//! Fallback fs-type detection for platforms without a specific impl.
|
||||
//!
|
||||
//! Always returns [`FsType::Unknown`]. Callers treat that as
|
||||
//! "probably local" and pick `LocalFileSink`.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use super::FsType;
|
||||
|
||||
pub(super) fn detect_impl(_path: &Path) -> FsType {
|
||||
FsType::Unknown
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
//! Windows filesystem-type detection.
|
||||
//!
|
||||
//! 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
|
||||
//! proper `GetVolumeInformation` query is a Phase 4 concern.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use super::FsType;
|
||||
|
||||
pub(super) fn detect_impl(path: &Path) -> FsType {
|
||||
// `Path::starts_with("\\\\")` won't match because component-wise
|
||||
// matching strips the prefix. Look at the raw OsStr instead.
|
||||
let s = path.as_os_str();
|
||||
// OsStr -> [u16] is the platform-correct way on Windows, but
|
||||
// checking the leading bytes via `to_string_lossy` is good enough
|
||||
// for a UNC prefix probe and works on any encoding.
|
||||
let lossy = s.to_string_lossy();
|
||||
if lossy.starts_with("\\\\") || lossy.starts_with("//") {
|
||||
return FsType::Nfs;
|
||||
}
|
||||
FsType::Local
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
//! Platform-specific drive initialization and disc probing.
|
||||
|
||||
pub mod fs_type;
|
||||
pub mod mt1959;
|
||||
|
||||
use crate::error::Result;
|
||||
@@ -17,4 +18,19 @@ pub(crate) trait PlatformDriver: Send {
|
||||
|
||||
/// True after successful init().
|
||||
fn is_ready(&self) -> bool;
|
||||
|
||||
/// True if the drive is currently in the extended-access state —
|
||||
/// per-drive runtime firmware uploaded AND the unlock response's
|
||||
/// marker bytes confirm the mode is live. When true:
|
||||
/// - host can issue the per-drive OEM CDBs in
|
||||
/// [`crate::profile::DriveProfile`]
|
||||
/// - VID retrieval works via the OEM CDB path (no cert-based
|
||||
/// mutual auth required)
|
||||
/// - SCSI READ_10 returns plaintext sectors (no bus encryption)
|
||||
///
|
||||
/// Default `false` — platforms without this mode always report
|
||||
/// inactive.
|
||||
fn is_unlocked(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
+188
-9
@@ -25,8 +25,16 @@ const SUB_CMD_INIT: u8 = 0x12;
|
||||
const SUB_CMD_PROBE: u8 = 0x14;
|
||||
const UNLOCK_RESPONSE_SIZE: u8 = 64;
|
||||
const VALIDATE_RESPONSE_SIZE: u8 = 4;
|
||||
/// Primary mode marker at bytes [12..16] of the unlock response — set
|
||||
/// by the platform firmware when the runtime image is loaded and the
|
||||
/// extended-access surface is live.
|
||||
const FIRMWARE_ACTIVE_OFFSET: usize = 12;
|
||||
const FIRMWARE_ACTIVE_SIG: [u8; 4] = [0x4D, 0x4D, 0x6B, 0x76];
|
||||
/// Secondary mode marker repeated through bytes [16..64] of the unlock
|
||||
/// response. Confirms the runtime firmware is the one driving the
|
||||
/// response, not a stale image's residual buffer.
|
||||
const FIRMWARE_MODE_OFFSET: usize = 16;
|
||||
const FIRMWARE_MODE_SIG: [u8; 4] = [0x4C, 0x62, 0x44, 0x72];
|
||||
|
||||
// ── Init address (per disc type) ──────────────────────────────────────
|
||||
const INIT_ADDR_BD: u16 = 0x0100;
|
||||
@@ -46,7 +54,18 @@ pub struct Mt1959 {
|
||||
pub(crate) profile: DriveProfile,
|
||||
pub(crate) mode: u8,
|
||||
pub(crate) buffer_id: u8,
|
||||
pub(crate) unlocked: bool,
|
||||
/// True after `run_init` has completed the unlock handshake (and any
|
||||
/// required firmware upload). Gates probe + downstream control
|
||||
/// commands; says nothing about whether the drive is in
|
||||
/// extended-access mode.
|
||||
pub(crate) init_complete: bool,
|
||||
/// True when the unlock response carried both the per-drive
|
||||
/// signature AND the primary mode marker at offset 12 AND the
|
||||
/// secondary mode marker at offset 16. When true the drive is in
|
||||
/// the extended-access state — host can issue the per-drive
|
||||
/// OEM CDBs and read sectors without the cert-based AACS bus
|
||||
/// encryption / mutual-auth gate.
|
||||
unlocked: bool,
|
||||
probed: bool,
|
||||
}
|
||||
|
||||
@@ -61,6 +80,7 @@ impl Mt1959 {
|
||||
profile,
|
||||
mode,
|
||||
buffer_id,
|
||||
init_complete: false,
|
||||
unlocked: false,
|
||||
probed: false,
|
||||
}
|
||||
@@ -141,7 +161,20 @@ impl Mt1959 {
|
||||
return Err(Error::UnlockFailed);
|
||||
}
|
||||
|
||||
self.unlocked = true;
|
||||
// Extended-access state is active when BOTH the per-drive
|
||||
// signature matched AND the response carries the secondary
|
||||
// marker at offset 16 (repeated through bytes 16..64) AND the
|
||||
// primary mode marker at [12..16] is present. The active-mode
|
||||
// marker at [12..16] is the primary gate; the [16..20] marker
|
||||
// is the redundant confirmation the firmware writes through
|
||||
// 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
|
||||
&& response[FIRMWARE_ACTIVE_OFFSET..FIRMWARE_ACTIVE_OFFSET + 4] == FIRMWARE_ACTIVE_SIG
|
||||
&& response[FIRMWARE_MODE_OFFSET..FIRMWARE_MODE_OFFSET + 4] == FIRMWARE_MODE_SIG;
|
||||
|
||||
self.init_complete = true;
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
@@ -177,11 +210,11 @@ impl Mt1959 {
|
||||
// ── Init (unlock + firmware) ───────────────────────────────────────
|
||||
|
||||
fn run_init(&mut self, scsi: &mut dyn ScsiTransport) -> Result<()> {
|
||||
let mut unlocked = false;
|
||||
let mut succeeded = false;
|
||||
for _attempt in 0..3 {
|
||||
match self.do_unlock(scsi) {
|
||||
Ok(_) => {
|
||||
unlocked = true;
|
||||
succeeded = true;
|
||||
break;
|
||||
}
|
||||
Err(Error::SignatureMismatch { .. }) => {
|
||||
@@ -202,7 +235,7 @@ impl Mt1959 {
|
||||
}
|
||||
}
|
||||
}
|
||||
if !unlocked {
|
||||
if !succeeded {
|
||||
return Err(Error::UnlockFailed);
|
||||
}
|
||||
Ok(())
|
||||
@@ -214,14 +247,14 @@ impl Mt1959 {
|
||||
/// per region. Two passes, then SET_CD_SPEED(max). After this the
|
||||
/// drive manages per-zone speeds internally.
|
||||
fn run_probe(&mut self, scsi: &mut dyn ScsiTransport) -> Result<()> {
|
||||
if !self.unlocked {
|
||||
if !self.init_complete {
|
||||
self.do_unlock(scsi)?;
|
||||
}
|
||||
|
||||
// Detect disc type from capacity to select probe mode.
|
||||
// BD: 3C 01 44 12 01 00 00 00 04 00 (init_addr = 0x0100)
|
||||
// UHD: 3C 01 44 12 02 00 00 00 04 00 (init_addr = 0x0200)
|
||||
// Verified from MakeMKV strace: BD and UHD use different init addresses.
|
||||
// Empirically verified via SCSI capture: BD and UHD use different init addresses.
|
||||
let cap_cdb = [
|
||||
SCSI_READ_CAPACITY,
|
||||
0x00,
|
||||
@@ -313,14 +346,14 @@ impl Mt1959 {
|
||||
|
||||
impl PlatformDriver for Mt1959 {
|
||||
fn init(&mut self, scsi: &mut dyn ScsiTransport) -> Result<()> {
|
||||
if self.unlocked {
|
||||
if self.init_complete {
|
||||
return Ok(());
|
||||
}
|
||||
self.run_init(scsi)
|
||||
}
|
||||
|
||||
fn probe_disc(&mut self, scsi: &mut dyn ScsiTransport) -> Result<()> {
|
||||
if !self.unlocked {
|
||||
if !self.init_complete {
|
||||
// Don't retry init here — if init() failed, probing can't work either.
|
||||
// Retrying causes repeated USB bus resets on BU40N.
|
||||
return Ok(());
|
||||
@@ -332,6 +365,152 @@ impl PlatformDriver for Mt1959 {
|
||||
}
|
||||
|
||||
fn is_ready(&self) -> bool {
|
||||
self.init_complete
|
||||
}
|
||||
|
||||
fn is_unlocked(&self) -> bool {
|
||||
self.unlocked
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::profile::{DriveProfile, Identity};
|
||||
use crate::scsi::{DataDirection, ScsiResult, ScsiTransport};
|
||||
|
||||
/// Minimal mock transport that returns a scripted response to the
|
||||
/// next `execute()` call. Only used for verifying that `do_unlock`
|
||||
/// classifies the response correctly — no general SCSI coverage.
|
||||
struct ScriptedTransport {
|
||||
response: Vec<u8>,
|
||||
}
|
||||
|
||||
impl ScsiTransport for ScriptedTransport {
|
||||
fn execute(
|
||||
&mut self,
|
||||
_cdb: &[u8],
|
||||
_dir: DataDirection,
|
||||
data: &mut [u8],
|
||||
_timeout_ms: u32,
|
||||
) -> Result<ScsiResult> {
|
||||
let n = self.response.len().min(data.len());
|
||||
data[..n].copy_from_slice(&self.response[..n]);
|
||||
Ok(ScsiResult {
|
||||
status: 0,
|
||||
bytes_transferred: n,
|
||||
sense: [0u8; 32],
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn fixture_profile(signature: [u8; 4]) -> DriveProfile {
|
||||
DriveProfile {
|
||||
identity: Identity {
|
||||
vendor_id: "TEST".into(),
|
||||
product_revision: String::new(),
|
||||
vendor_specific: String::new(),
|
||||
firmware_date: String::new(),
|
||||
},
|
||||
signature,
|
||||
firmware: Vec::new(),
|
||||
unlock_init_value: 0,
|
||||
unlock_response_size: 0,
|
||||
read_vid_cdb: None,
|
||||
read_disc_keys_cdb: None,
|
||||
drive_nominal_speed_cdb: None,
|
||||
set_speed_max_cdb: None,
|
||||
read10_raw_2sec_cdb: None,
|
||||
read10_raw_1sec_cdb: None,
|
||||
read_buffer_verify_cdb: None,
|
||||
write_buffer_cdb: None,
|
||||
read_buffer_unlock_cdb: None,
|
||||
speed_zone_table: None,
|
||||
speed_calc_table: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a synthetic 64-byte unlock response.
|
||||
///
|
||||
/// `mode_marker`: bytes [12..16]. Pass `FIRMWARE_ACTIVE_SIG` for the
|
||||
/// active-mode primary marker.
|
||||
/// `id_marker`: bytes [16..20] (and repeated through [20..64] in
|
||||
/// real responses; only [16..20] is checked).
|
||||
fn build_response(signature: [u8; 4], mode_marker: [u8; 4], id_marker: [u8; 4]) -> Vec<u8> {
|
||||
let mut r = vec![0u8; 64];
|
||||
r[0..4].copy_from_slice(&signature);
|
||||
// bytes [4..12] left as zeros (version + reserved per format)
|
||||
r[12..16].copy_from_slice(&mode_marker);
|
||||
// Real firmware repeats the secondary marker through [16..64];
|
||||
// the parser only checks [16..20], so we just write the marker
|
||||
// once.
|
||||
r[16..20].copy_from_slice(&id_marker);
|
||||
r
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn do_unlock_sets_unlocked_when_both_markers_present() {
|
||||
let sig = [0x99, 0x9E, 0xC3, 0x75];
|
||||
let response = build_response(sig, FIRMWARE_ACTIVE_SIG, FIRMWARE_MODE_SIG);
|
||||
let mut transport = ScriptedTransport { response };
|
||||
let mut mt = Mt1959::new(fixture_profile(sig), false);
|
||||
|
||||
let raw = mt.do_unlock(&mut transport).expect("unlock should succeed");
|
||||
assert_eq!(raw.len(), 64);
|
||||
assert!(mt.init_complete, "init_complete set after success");
|
||||
assert!(
|
||||
mt.is_unlocked(),
|
||||
"both markers present -> extended-access state"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn do_unlock_init_complete_but_not_unlocked_when_id_marker_missing() {
|
||||
// Primary mode marker present (so init passes) but the
|
||||
// secondary marker is replaced with zeros — drive isn't in
|
||||
// extended-access state.
|
||||
let sig = [0x99, 0x9E, 0xC3, 0x75];
|
||||
let response = build_response(sig, FIRMWARE_ACTIVE_SIG, [0u8; 4]);
|
||||
let mut transport = ScriptedTransport { response };
|
||||
let mut mt = Mt1959::new(fixture_profile(sig), false);
|
||||
|
||||
mt.do_unlock(&mut transport).expect("unlock should succeed");
|
||||
assert!(mt.init_complete);
|
||||
assert!(
|
||||
!mt.is_unlocked(),
|
||||
"missing secondary marker -> not in extended-access state"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn do_unlock_rejects_signature_mismatch() {
|
||||
let response = build_response(
|
||||
[0xAA, 0xBB, 0xCC, 0xDD],
|
||||
FIRMWARE_ACTIVE_SIG,
|
||||
FIRMWARE_MODE_SIG,
|
||||
);
|
||||
let mut transport = ScriptedTransport { response };
|
||||
let mut mt = Mt1959::new(fixture_profile([0x99, 0x9E, 0xC3, 0x75]), false);
|
||||
|
||||
let err = mt.do_unlock(&mut transport).unwrap_err();
|
||||
assert!(matches!(err, Error::SignatureMismatch { .. }));
|
||||
assert!(!mt.init_complete);
|
||||
assert!(!mt.is_unlocked());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn do_unlock_rejects_inactive_mode_marker() {
|
||||
// Signature matches but the primary marker at [12..16] is
|
||||
// missing -> drive is not in active mode; init_complete and the
|
||||
// unlocked flag must both stay false.
|
||||
let sig = [0x99, 0x9E, 0xC3, 0x75];
|
||||
let response = build_response(sig, [0u8; 4], FIRMWARE_MODE_SIG);
|
||||
let mut transport = ScriptedTransport { response };
|
||||
let mut mt = Mt1959::new(fixture_profile(sig), false);
|
||||
|
||||
let err = mt.do_unlock(&mut transport).unwrap_err();
|
||||
assert!(matches!(err, Error::UnlockFailed));
|
||||
assert!(!mt.init_complete);
|
||||
assert!(!mt.is_unlocked());
|
||||
}
|
||||
}
|
||||
|
||||
+108
@@ -35,6 +35,42 @@ pub struct DriveProfile {
|
||||
pub signature: [u8; 4],
|
||||
#[serde(default, deserialize_with = "deserialize_base64")]
|
||||
pub firmware: Vec<u8>,
|
||||
|
||||
// ── OEM-extended-access CDB templates ──────────────────────────────
|
||||
//
|
||||
// All optional — older profile blobs that pre-date the CDB capture
|
||||
// pipeline simply omit these fields and decode as `None`. Encoded
|
||||
// in the JSON as lowercase hex strings without separators
|
||||
// (e.g. `"3c014410e29100002400"` for a 10-byte CDB).
|
||||
#[serde(default)]
|
||||
pub unlock_init_value: u8,
|
||||
#[serde(default)]
|
||||
pub unlock_response_size: u8,
|
||||
|
||||
#[serde(default, deserialize_with = "deserialize_opt_hex_bytes_10")]
|
||||
pub read_vid_cdb: Option<[u8; 10]>,
|
||||
#[serde(default, deserialize_with = "deserialize_opt_hex_bytes_10")]
|
||||
pub read_disc_keys_cdb: Option<[u8; 10]>,
|
||||
#[serde(default, deserialize_with = "deserialize_opt_hex_bytes_12")]
|
||||
pub drive_nominal_speed_cdb: Option<[u8; 12]>,
|
||||
#[serde(default, deserialize_with = "deserialize_opt_hex_bytes_12")]
|
||||
pub set_speed_max_cdb: Option<[u8; 12]>,
|
||||
#[serde(default, deserialize_with = "deserialize_opt_hex_bytes_10")]
|
||||
pub read10_raw_2sec_cdb: Option<[u8; 10]>,
|
||||
#[serde(default, deserialize_with = "deserialize_opt_hex_bytes_10")]
|
||||
pub read10_raw_1sec_cdb: Option<[u8; 10]>,
|
||||
#[serde(default, deserialize_with = "deserialize_opt_hex_bytes_10")]
|
||||
pub read_buffer_verify_cdb: Option<[u8; 10]>,
|
||||
#[serde(default, deserialize_with = "deserialize_opt_hex_bytes_10")]
|
||||
pub write_buffer_cdb: Option<[u8; 10]>,
|
||||
#[serde(default, deserialize_with = "deserialize_opt_hex_bytes_10")]
|
||||
pub read_buffer_unlock_cdb: Option<[u8; 10]>,
|
||||
|
||||
// Per-drive identifier tables — variable-length hex strings.
|
||||
#[serde(default, deserialize_with = "deserialize_opt_hex_bytes")]
|
||||
pub speed_zone_table: Option<Vec<u8>>,
|
||||
#[serde(default, deserialize_with = "deserialize_opt_hex_bytes")]
|
||||
pub speed_calc_table: Option<Vec<u8>>,
|
||||
}
|
||||
|
||||
/// Chipset + variant — determined by which section the profile was found in.
|
||||
@@ -99,6 +135,78 @@ where
|
||||
.map_err(serde::de::Error::custom)
|
||||
}
|
||||
|
||||
// ── Fixed-length hex deserializers for CDB templates ────────────────────
|
||||
//
|
||||
// Profile JSON encodes CDBs as lowercase hex strings without separators.
|
||||
// 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)
|
||||
}
|
||||
|
||||
fn deserialize_opt_hex_bytes_10<'de, D>(
|
||||
deserializer: D,
|
||||
) -> std::result::Result<Option<[u8; 10]>, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
{
|
||||
let opt: Option<String> = Option::deserialize(deserializer)?;
|
||||
let Some(s) = opt else { return Ok(None) };
|
||||
if s.is_empty() {
|
||||
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);
|
||||
Ok(Some(out))
|
||||
}
|
||||
|
||||
fn deserialize_opt_hex_bytes_12<'de, D>(
|
||||
deserializer: D,
|
||||
) -> std::result::Result<Option<[u8; 12]>, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
{
|
||||
let opt: Option<String> = Option::deserialize(deserializer)?;
|
||||
let Some(s) = opt else { return Ok(None) };
|
||||
if s.is_empty() {
|
||||
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);
|
||||
Ok(Some(out))
|
||||
}
|
||||
|
||||
fn deserialize_opt_hex_bytes<'de, D>(
|
||||
deserializer: D,
|
||||
) -> std::result::Result<Option<Vec<u8>>, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
{
|
||||
let opt: Option<String> = Option::deserialize(deserializer)?;
|
||||
let Some(s) = opt else { return Ok(None) };
|
||||
if s.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
let bytes = parse_hex_bytes(&s).map_err(serde::de::Error::custom)?;
|
||||
Ok(Some(bytes))
|
||||
}
|
||||
|
||||
// ── Loading ────────────────────────────────────────────────────────────
|
||||
|
||||
const BUNDLED_PROFILES: &str = include_str!("../profiles.json");
|
||||
|
||||
+12
-81
@@ -1,83 +1,19 @@
|
||||
//! File-backed sector I/O — read and write 2048-byte sectors against
|
||||
//! an ISO image on disk.
|
||||
//! File-backed sector sink — write 2048-byte sectors to an ISO image
|
||||
//! on disk.
|
||||
//!
|
||||
//! [`FileSectorSource`] is the read side (open-only). [`FileSectorSink`]
|
||||
//! is the write side (create or open-rw); writes go through
|
||||
//! [`crate::io::WritebackFile`] so big sequential ISO writes share
|
||||
//! the same bounded-cache writeback pipeline used by sweep / patch /
|
||||
//! mux.
|
||||
//! The read-side counterpart ([`crate::io::file_sector_source::FileSectorSource`])
|
||||
//! lives under `io/` because its internals (read-ahead buffer, per-OS
|
||||
//! `fadvise`/`F_RDADVISE` hints) are I/O infrastructure rather than
|
||||
//! sector-trait business logic. Both types remain re-exported at
|
||||
//! [`crate::sector`] for ergonomic imports.
|
||||
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::io::{Read, Seek, SeekFrom, Write};
|
||||
use std::fs::OpenOptions;
|
||||
use std::io::{Seek, SeekFrom, Write};
|
||||
use std::path::Path;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
use super::{SectorSink, SectorSource};
|
||||
|
||||
/// SectorSource backed by a file (ISO image).
|
||||
///
|
||||
/// Seeks to `lba * 2048`, reads `count * 2048` bytes per call. The
|
||||
/// file is held directly: every `read_sectors` call performs an
|
||||
/// absolute seek, so a wrapping `BufReader` would have its buffer
|
||||
/// invalidated on every call (its internal cursor moves with the
|
||||
/// `Seek` impl) — pure overhead. Callers that benefit from buffered
|
||||
/// reads should compose their own `BufReader` at the `read_sectors`
|
||||
/// granularity they care about.
|
||||
pub struct FileSectorSource {
|
||||
file: File,
|
||||
capacity: u32,
|
||||
}
|
||||
|
||||
impl FileSectorSource {
|
||||
/// Open an existing ISO file for reading. Capacity is derived
|
||||
/// from `metadata().len() / 2048`. Returns
|
||||
/// [`Error::IsoTooLarge`] if the file would exceed the 32-bit
|
||||
/// LBA address space (~8 TB).
|
||||
pub fn open(path: &Path) -> std::io::Result<Self> {
|
||||
let file = File::open(path)?;
|
||||
let len = file.metadata()?.len();
|
||||
let sectors = len / 2048;
|
||||
if sectors > u32::MAX as u64 {
|
||||
return Err(Error::IsoTooLarge {
|
||||
path: path.to_string_lossy().into_owned(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
let capacity = sectors as u32;
|
||||
Ok(Self { file, capacity })
|
||||
}
|
||||
}
|
||||
|
||||
// Implement the legacy `SectorSource` trait. The blanket impl in
|
||||
// `super` produces the `SectorSource` impl automatically — no need
|
||||
// to write both, and writing both would conflict. This keeps the
|
||||
// 0.17 method-resolution path intact (callers with `SectorSource`
|
||||
// in scope can still write `fsr.read_sectors(..)` against a
|
||||
// `FileSectorSource`).
|
||||
impl SectorSource for FileSectorSource {
|
||||
fn capacity_sectors(&self) -> u32 {
|
||||
self.capacity
|
||||
}
|
||||
|
||||
fn read_sectors(
|
||||
&mut self,
|
||||
lba: u32,
|
||||
count: u16,
|
||||
buf: &mut [u8],
|
||||
_recovery: bool,
|
||||
) -> Result<usize> {
|
||||
let offset = lba as u64 * 2048;
|
||||
let bytes = count as usize * 2048;
|
||||
self.file
|
||||
.seek(SeekFrom::Start(offset))
|
||||
.map_err(|e| Error::IoError { source: e })?;
|
||||
self.file
|
||||
.read_exact(&mut buf[..bytes])
|
||||
.map_err(|e| Error::IoError { source: e })?;
|
||||
Ok(bytes)
|
||||
}
|
||||
}
|
||||
use super::SectorSink;
|
||||
|
||||
/// SectorSink backed by a file (ISO image).
|
||||
///
|
||||
@@ -146,13 +82,8 @@ impl SectorSink for FileSectorSink {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
// Bring the 0.18 trait into scope (not super::*: the super
|
||||
// module also re-exports the legacy `SectorSource`, and
|
||||
// having both `SectorSource::read_sectors` and
|
||||
// `SectorSource::read_sectors` visible would force every
|
||||
// call site to disambiguate). External consumers see the
|
||||
// same surface this test exercises.
|
||||
use super::{FileSectorSink, FileSectorSource};
|
||||
use super::FileSectorSink;
|
||||
use crate::io::file_sector_source::FileSectorSource;
|
||||
use crate::sector::{SectorSink, SectorSource};
|
||||
use tempfile::tempdir;
|
||||
|
||||
|
||||
+5
-2
@@ -6,7 +6,7 @@
|
||||
//! ISO images are opened for read OR write at construction time.
|
||||
//!
|
||||
//! - [`SectorSource`] is implemented by `Drive` (hardware) and
|
||||
//! [`FileSectorSource`] / `IsoSectorReader` (file-backed).
|
||||
//! [`FileSectorSource`] (file-backed).
|
||||
//! - [`SectorSink`] is implemented by [`FileSectorSink`]
|
||||
//! (ISO-backed) and sweep/patch consumer adapters.
|
||||
//! - [`DecryptingSectorSource`] is a decorator that wraps any
|
||||
@@ -15,6 +15,7 @@
|
||||
|
||||
pub mod decrypting;
|
||||
pub mod file;
|
||||
pub mod prefetched;
|
||||
|
||||
use crate::error::Result;
|
||||
|
||||
@@ -109,5 +110,7 @@ pub trait SectorSink: Send {
|
||||
fn finish(self: Box<Self>) -> Result<()>;
|
||||
}
|
||||
|
||||
pub use crate::io::file_sector_source::FileSectorSource;
|
||||
pub use decrypting::DecryptingSectorSource;
|
||||
pub use file::{FileSectorSink, FileSectorSource};
|
||||
pub use file::FileSectorSink;
|
||||
pub use prefetched::PrefetchedSectorSource;
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user