libfreemkv/src/keysource.rs

//! Key sources — the layer that hands libfreemkv a disc's terminal Unit Keys.
//!
//! libfreemkv performs NO key lookup. An application resolves a disc's keys
//! through one or more [`KeySource`]s, each an adapter over a backing store (a
//! keydb file, a key server, the mapfile cache). A source's job is to return the
//! disc's terminal **Unit Keys** ([`crate::aacs::types::UnitKey`]). It knows what
//! material it holds (a DK / MK / VUK / pre-decrypted UK) and what it must fetch
//! from the disc (VID, MKB, encrypted title keys, content samples) to get there;
//! it orchestrates the derivation by calling libfreemkv's own derivation
//! primitives ([`crate::aacs::derive::derive_media_key_from_dk`] /
//! [`crate::aacs::derive::derive_vuk`] / [`crate::aacs::derive::decrypt_unit_key`])
//! through the [`ResolveCtx`] handed to it.
//!
//! libfreemkv still OWNS the crypto: the boil-down primitives and the AES live
//! here. A source owns only PATH ORCHESTRATION — deciding which primitive to
//! call with what input for the material it happens to hold. Source
//! implementations are published in the companion `freemkv-keysources` crate,
//! keeping key *policy* (which store, which order, online vs local) out of the
//! library.

use crate::aacs::types::HostCert;
use crate::aacs::types::{UnitKey, Vid};
use crate::disc::Key;
use crate::error::Error;

/// Minimum encrypted-content unit samples a single online key request must carry.
///
/// The key service identifies a key by which of the submitted units it decrypts,
/// so too few samples — especially on FMTS, where a segment interleaves several
/// variants at the unit level — can return a key that matches an incidental unit
/// rather than the one asked about (a false positive). This many distinct units
/// make the request unambiguous.
///
/// Canonical here (the base crate) so BOTH consumers agree on one value: the
/// online source in `freemkv-keysources` (which refuses to send an under-sampled
/// request) re-exports it, and libfreemkv's own FMTS forensic query
/// ([`crate::mux`]) sizes its per-segment batch by it. Layering forbids the
/// reverse import (keysources depends on libfreemkv, not vice versa), so the
/// value lives at the lower layer both share.
pub const MIN_SAMPLE_UNITS: usize = 8;

/// A set of encrypted content-unit samples PROVEN to carry at least
/// [`MIN_SAMPLE_UNITS`] units — the online `/decode` request's proof-of-ownership.
///
/// "Parse, don't validate": the only constructor, [`DecodeSampleSet::new`], returns
/// `None` for an under-sized slice, so an online key request simply *cannot be built*
/// from too few samples. The runtime `len() < MIN_SAMPLE_UNITS` check that used to
/// live at the request site (and was silently forgotten by an under-sampling caller,
/// reading as "key service down") becomes a compile-time obligation: a request builder
/// that takes `&DecodeSampleSet` can never receive an unchecked `Vec`.
///
/// The *count* enforced here is a runtime property of the disc (how many encrypted
/// units it yields); the *requested* count is a caller-side compile-time constant that
/// callers pin to `MIN_SAMPLE_UNITS` (see e.g. autorip's `SAMPLE_UNITS`). Together the
/// two make under-sampling unrepresentable at the request boundary.
///
/// The wrapped samples are on-disc AACS ciphertext — the same bytes the sibling
/// [`DiscInputs::samples`] redacts as key MATERIAL — so [`Debug`] is hand-written
/// and redacting; see the impl below.
#[derive(Clone)]
pub struct DecodeSampleSet(Vec<Vec<u8>>);

impl std::fmt::Debug for DecodeSampleSet {
    /// Prints the SHAPE only. A derived `Debug` dumped every wrapped sample
    /// verbatim: a `DecodeSampleSet` carries at least [`MIN_SAMPLE_UNITS`]
    /// 6144-byte aligned units (≥ 49 KiB, in practice multi-MB) of AACS
    /// ciphertext plus each unit's clear 16-byte derivation seed, so one
    /// `tracing::debug!("{set:?}")` on a failed `/decode` request — or an
    /// `assert_eq!` whose panic message formats it — wrote all of it to the log
    /// that gets attached to a bug report. Same policy and same shape as
    /// [`DiscInputs`]'s impl below.
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DecodeSampleSet")
            .field("units", &"<redacted>")
            .field("units_len", &self.0.len())
            .finish()
    }
}

impl DecodeSampleSet {
    /// Wrap `units` iff it carries at least [`MIN_SAMPLE_UNITS`] samples; `None`
    /// otherwise (the caller then skips the online source rather than sending an
    /// ambiguous request). This is the sole way to obtain a `DecodeSampleSet`.
    pub fn new(units: Vec<Vec<u8>>) -> Option<Self> {
        (units.len() >= MIN_SAMPLE_UNITS).then_some(Self(units))
    }

    /// The proven-sufficient samples. Guaranteed `>= MIN_SAMPLE_UNITS` in length.
    pub fn units(&self) -> &[Vec<u8>] {
        &self.0
    }

    /// Number of samples — always `>= MIN_SAMPLE_UNITS`.
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Always `false` (a `DecodeSampleSet` never holds fewer than `MIN_SAMPLE_UNITS`);
    /// provided so the type satisfies the usual `len`/`is_empty` pairing.
    pub fn is_empty(&self) -> bool {
        false
    }
}

/// The public AACS inputs a key source needs to look a disc up. Captured at
/// scan; carries no DERIVED secrets (no media key, VUK or plaintext unit key) —
/// only the disc identity and the on-disc AACS structures a source or key server
/// may key on. The on-disc structures are nonetheless key MATERIAL (the encrypted
/// title keys live in `unit_key_ro`), so [`Debug`] is hand-written and redacting;
/// see the impl below.
#[derive(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],
    /// AACS major version (1 = V10 / BD AACS 1.0, 2 = V20+ / UHD). Drives the
    /// `Unit_Key_RO.inf` parse stride (48-byte V10 vs 64-byte V20/V21) when a
    /// source returns a VUK to derive unit keys from. Defaults to 2.
    pub version: u8,
    /// 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>>,
    /// The disc's human title — the UDF/ISO volume identifier (e.g.
    /// `TITLE_2024`), falling back to the BDMV `<di:name>` when present.
    /// `None` when not captured. Identity only, no secret; a key service may
    /// record it (keyed by `disc_hash`) to build a hash→title catalog. Not used
    /// in any AACS derivation.
    pub volume_label: Option<String>,
}

/// Redacting `Debug`, per the policy `aacs::types` documents (and which
/// `aacs::types::Vid` already applies to this very Volume ID). `DiscInputs` is
/// public and returned by [`crate::Disc::inputs`], so a consumer's
/// `tracing::debug!("{inputs:?}")` used to print the Volume ID, the whole
/// `Unit_Key_RO.inf` (the encrypted title keys), the entire MKB and every
/// ciphertext sample verbatim into a log that ends up attached to a bug report.
/// Only non-secret identity and shape (presence, lengths) is printed.
impl std::fmt::Debug for DiscInputs {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DiscInputs")
            .field("disc_hash", &self.disc_hash)
            .field("volume_id", &"<redacted>")
            .field("version", &self.version)
            .field("mkb", &"<redacted>")
            .field("mkb_len", &self.mkb.len())
            .field("unit_key_ro", &"<redacted>")
            .field("unit_key_ro_len", &self.unit_key_ro.len())
            .field("samples", &"<redacted>")
            .field("samples_len", &self.samples.len())
            .field("volume_label", &self.volume_label)
            .finish()
    }
}

/// A lazy view of a disc's AACS material, handed to [`KeySource::get_unit_keys`] so a
/// source can drive the derivation chain without holding the disc reader.
///
/// "Lazy" by contract: each accessor returns only what the source asks for, so a
/// source that already holds terminal Unit Keys never touches the MKB or
/// samples. (Today the backing [`DiscInputsCtx`] is eagerly populated from a
/// scan-time [`DiscInputs`]; the trait keeps the lazy signature so a future
/// implementation can fetch on demand without a source-API break.)
pub trait ResolveCtx {
    /// SHA-1 of `Unit_Key_RO.inf`, `0x`-prefixed hex — the per-disc lookup key.
    fn disc_hash(&self) -> &str;
    /// The disc's human title (UDF/ISO volume identifier), when captured.
    fn title(&self) -> Option<&str>;
    /// Volume ID, or `None` when no authenticated handshake ran (the all-zero
    /// sentinel) — VID-dependent derivation (`MK → VUK`) is then impossible.
    fn vid(&self) -> Option<Vid>;
    /// Raw MKB bytes (may be empty when not captured).
    fn mkb(&self) -> Result<&[u8], Error>;
    /// The disc's encrypted title keys, parsed from `Unit_Key_RO.inf` the same
    /// way the library's resolver parses them ([`crate::aacs::inf::parse_unit_key_ro`]),
    /// in on-disc order. Feed straight into [`crate::aacs::derive::decrypt_unit_key`].
    fn enc_title_keys(&self) -> Result<&[[u8; 16]], Error>;
    /// Up to `n` encrypted on-disc content sample units, for a source that
    /// validates a candidate server-side against real ciphertext.
    fn samples(&self, n: usize) -> Result<Vec<Vec<u8>>, Error>;
    /// Raw `Unit_Key_RO.inf` bytes, verbatim. Most sources derive locally from
    /// the parsed [`Self::enc_title_keys`]; a source that forwards the on-disc
    /// structure to a server doing its OWN derivation (an online key service)
    /// needs the unparsed blob. Empty when not captured. Defaults to empty so
    /// existing/foreign `ResolveCtx` impls keep compiling unchanged.
    fn unit_key_ro(&self) -> &[u8] {
        &[]
    }
}

/// [`ResolveCtx`] over a scan-time [`DiscInputs`].
///
/// Pre-parses the encrypted title keys at construction (so `enc_title_keys` can
/// hand back a borrowed slice) at the version-appropriate `Unit_Key_RO.inf`
/// stride — `version_u8` is the disc's AACS major (1 → 48-byte V10 stride, else
/// 64-byte V20/V21 stride), matching the library resolver's dispatch.
pub struct DiscInputsCtx<'a> {
    inner: &'a DiscInputs,
    enc_keys: Vec<[u8; 16]>,
}

impl<'a> DiscInputsCtx<'a> {
    /// Build a context over `inputs`, parsing the encrypted title keys at the
    /// stride for the disc's own AACS major (`inputs.version`: 1 → 48-byte V10
    /// stride, else 64-byte V20/V21) — the single source of truth, no separate
    /// version argument to drift from it.
    ///
    /// A present-but-malformed `unit_key_ro` (truncated / wrong magic / wrong
    /// stride) parses to an empty key set, so a later [`Self::enc_title_keys`]
    /// returns `Ok(&[])` indistinguishably from a disc that legitimately has no
    /// title keys — the parse failure is swallowed here, not surfaced as an
    /// error.
    pub fn new(inputs: &'a DiscInputs) -> Self {
        use crate::aacs::inf::parse_unit_key_ro;
        use crate::aacs::mkb::AacsVersion;
        let enc_keys = if inputs.unit_key_ro.is_empty() {
            Vec::new()
        } else {
            parse_unit_key_ro(&inputs.unit_key_ro, AacsVersion::from_major(inputs.version))
                .map(|f| f.encrypted_keys.into_iter().map(|(_, k)| k).collect())
                .unwrap_or_default()
        };
        Self {
            inner: inputs,
            enc_keys,
        }
    }
}

impl ResolveCtx for DiscInputsCtx<'_> {
    fn disc_hash(&self) -> &str {
        &self.inner.disc_hash
    }
    fn title(&self) -> Option<&str> {
        self.inner.volume_label.as_deref()
    }
    fn vid(&self) -> Option<Vid> {
        if self.inner.volume_id == [0u8; 16] {
            None
        } else {
            Some(Vid(self.inner.volume_id))
        }
    }
    fn mkb(&self) -> Result<&[u8], Error> {
        Ok(&self.inner.mkb)
    }
    fn enc_title_keys(&self) -> Result<&[[u8; 16]], Error> {
        Ok(&self.enc_keys)
    }
    fn samples(&self, n: usize) -> Result<Vec<Vec<u8>>, Error> {
        Ok(self.inner.samples.iter().take(n).cloned().collect())
    }
    fn unit_key_ro(&self) -> &[u8] {
        &self.inner.unit_key_ro
    }
}

/// A key source: an adapter over a backing store that resolves a disc's terminal
/// Unit Keys.
///
/// Dumb about *policy*, smart about *its own material*: given a [`ResolveCtx`] a
/// source looks the disc up in its store and, from whatever level of material it
/// holds, orchestrates the derivation down to Unit Keys using the library's
/// boil-down crypto primitives — never re-implementing AES. A source that holds
/// pre-decrypted Unit Keys returns them directly; one that holds a VUK calls
/// [`crate::aacs::derive::decrypt_unit_key`]; one that holds device keys calls
/// [`crate::aacs::derive::derive_media_key_from_dk`] → [`crate::aacs::derive::derive_vuk`] → `decrypt_unit_key`.
///
/// Returning an empty `Vec` means "no key for this disc from this source"; an
/// `Err` means the source itself failed (I/O, parse, network). The caller
/// ([`resolve_and_apply`]) tries each source in order and validates the returned
/// keys against real ciphertext before committing them, so a wrong key from one
/// source transparently falls through to the next.
///
/// Two explicit resolve operations, one per key kind — never one overloaded call
/// whose meaning depends on how many keys came back:
/// * [`get_unit_keys`](Self::get_unit_keys) — the disc's base per-CPS-unit Unit
///   Keys (index space = CPS-unit number). The common path for every disc.
/// * [`get_fmts_indexes`](Self::get_fmts_indexes) — the AACS 2.1 forensic index
///   keys (index space = forensic index 1..N). Defaults to empty: a source with
///   no forensic material opts out, and only an FMTS disc ever asks.
///
/// What each source must do to answer is the source's own business: a keydb keys
/// on `disc_hash` and reads no samples; the online source submits the ctx's
/// content samples (a base batch for `get_unit_keys`, an index-1 anchor batch for
/// `get_fmts_indexes`) to the key service.
pub trait KeySource {
    /// Resolve this disc's base per-CPS-unit Unit Keys from this source. An empty
    /// `Vec` is a genuine "no key here"; `Err` is a source failure.
    fn get_unit_keys(&self, ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error>;

    /// Resolve this disc's AACS 2.1 forensic index keys — the per-index keys the
    /// base Unit Key cannot open (see [`crate::aacs::segment`]) — ordered by
    /// forensic index (element `i` carries `UnitKey.idx == i`, forensic index
    /// `i + 1`). The source hands back the COMPLETE set it holds; the caller
    /// trusts any non-empty result as all of them and never assumes a fixed count.
    /// Defaults to empty: a source with no forensic material (a plain keydb, the
    /// mapfile) opts out, and only an FMTS disc's mux ever calls this.
    fn get_fmts_indexes(&self, _ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
        Ok(Vec::new())
    }

    /// The AACS host certificate(s) this source can supply for the live-drive
    /// SCSI mutual-auth handshake (the OEM/AACS baseline route). `mkb` is the
    /// disc's MKB generation when known, so a source MAY return only certs whose
    /// generation matches (the default ignores it). A host cert unlocks the
    /// authenticated bus so the drive reports the Volume ID and bus key; it is
    /// **perishable** (revocable on a drive's HRL), so it is served by a source,
    /// never compiled in. A source holding no cert returns the empty vec.
    fn host_certs(&self, _mkb: Option<u32>) -> Vec<HostCert> {
        Vec::new()
    }

    /// A short, stable identifier for this source kind (`"keydb"`, `"online"`,
    /// `"mapfile"`, …). For logging which source produced a key, and for
    /// composition/ordering. A format string, not user-facing English.
    fn label(&self) -> &'static str {
        "source"
    }
}

/// Drive `sources` until one resolves Unit Keys that decrypt `disc`. Returns
/// `true` at the first source whose keys validate and commit, `false` once every
/// source is exhausted (the genuine "no key for this disc"). Thin wrapper over
/// [`resolve_and_apply_traced`] that discards the trace.
pub fn resolve_and_apply(
    sources: &[Box<dyn KeySource>],
    inputs: &DiscInputs,
    disc: &mut crate::Disc,
) -> bool {
    resolve_and_apply_traced(sources, inputs, disc).0
}

/// Like [`resolve_and_apply`] but also returns a structured
/// [`crate::aacs::trace::ResolutionTrace`] recording, per source, what happened — for
/// applications to render. ZERO English; the trace is typed enums only.
///
/// One-shot per source: each source's [`KeySource::get_unit_keys`] is called exactly
/// once with a [`DiscInputsCtx`] over `inputs`. Non-empty Unit Keys are mapped
/// to terminal [`Key::Unit`]s and applied via [`crate::Disc::decrypt_with`],
/// which validates them against `inputs.samples` and only mutates the disc on
/// success — so a wrong/partial key set is rejected and the loop continues.
///
/// CPS-unit numbering: a source returns Unit Keys carrying the POSITIONAL index
/// from [`crate::aacs::derive::decrypt_unit_key`]; the library's canonical CPS-unit number is
/// `position + 1` (matching [`crate::aacs::inf::parse_unit_key_ro`]'s `(i + 1)`), so
/// the committed `AacsState.unit_keys` is byte-identical to the library-resolved
/// path.
///
/// The NUMBER itself is not what descramble indexes by — but the ORDER is
/// load-bearing, so a source must return its keys in CPS-unit order. Trial
/// decrypt-and-check was deliberately deleted (see
/// [`crate::decrypt::AacsKeyMap`]: decryption is driven by the disc's CPS-unit /
/// FMTS-segment structure, "never by trial-decrypt-and-check per unit"), and
/// `decrypt_sectors_mapped` indexes the committed pool POSITIONALLY —
/// `unit_keys[key_idx].1`, where `key_idx` is a POSITION in the Vec a source
/// returned, recorded by `resolve_mux_key_map_cached` / `resolve_fmts_key_map`.
/// Return the same keys in a different order and every `AacsKeyMap` points at the
/// wrong key: the whole title decrypts under a neighbour's key, or a forensic
/// range trips the `is_clean` net into `DecryptFailed`. (The doc used to say the
/// number "is cosmetic for descramble (the decrypt path strips it and tries every
/// key)", which is what the DELETED trial-decrypt path did; the only place that
/// still tries every key is `Disc::decrypt_with`'s sample VALIDATION, which does
/// not descramble content.)
pub fn resolve_and_apply_traced(
    sources: &[Box<dyn KeySource>],
    inputs: &DiscInputs,
    disc: &mut crate::Disc,
) -> (bool, crate::aacs::trace::ResolutionTrace) {
    use crate::aacs::trace::{KeyNode, KeyOutcome, KeyStep};

    let mut trace = crate::aacs::trace::ResolutionTrace::new();

    // The FIRST source failure seen, if any. A source that returns `Err` did not
    // answer "no key for this disc" — it could not answer at all — and that
    // reason is stamped onto `disc.aacs_error` below so the decrypt gate reports
    // THAT instead of the generic `NoDiscKey`. First-wins (not last) so the
    // ordered sources' most-preferred failure is the one the operator is told
    // about, matching the first-valid-wins rule for successes.
    let mut source_failure: Option<crate::error::Error> = None;

    // The ctx parses Unit_Key_RO.inf at the stride for `inputs.version` (the
    // disc's own AACS major), so the stride is the disc's single source of truth.
    let ctx = DiscInputsCtx::new(inputs);

    for source in sources {
        // `who` is the source's own stable identifier — no enum to map back to.
        let who = source.label().to_string();
        match source.get_unit_keys(&ctx) {
            Ok(uks) if !uks.is_empty() => {
                // Positional index → canonical CPS-unit number (position + 1).
                let unit_keys: Vec<(u32, [u8; 16])> = uks
                    .iter()
                    .map(|uk| (uk.idx.saturating_add(1), uk.key))
                    .collect();
                if disc
                    .decrypt_with(Key::Unit(unit_keys), &inputs.samples)
                    .is_ok()
                {
                    trace.keys.push(KeyStep {
                        who,
                        path: vec![KeyNode::FoundUnitKeys, KeyNode::DerivedUnitKeys],
                        outcome: KeyOutcome::Resolved,
                    });
                    return (true, trace);
                }
                // Keys produced but rejected by validation — record and continue.
                trace.keys.push(KeyStep {
                    who,
                    path: vec![KeyNode::FoundUnitKeys],
                    outcome: KeyOutcome::NoKey,
                });
            }
            // The source ANSWERED and holds nothing for this disc. This — and
            // only this — is `NoEntry`: the claim "I looked, it is not there".
            Ok(_) => {
                trace.keys.push(KeyStep {
                    who,
                    path: vec![KeyNode::NoEntry],
                    outcome: KeyOutcome::NoKey,
                });
            }
            // The source could NOT answer — it was unreachable, it errored, or it
            // refused. Nothing is known about whether a key exists, so the path
            // is EMPTY: recording `NoEntry` here is exactly the conflation that
            // made a seven-hour run of HTTP 502s render as
            // `key: online > no entry > NO KEY` + `E7022 No key source has a
            // decryption key for this disc`, and sent operators hunting for a VUK
            // that was never missing.
            //
            // The reason itself rides out on `disc.aacs_error` (below), the
            // channel `Disc::ensure_decryptable_keys` already reads for the
            // E7017-vs-E7022 split — so the decrypt gate raises the SOURCE's code
            // (`KeyServiceUnavailable` / `KeyServiceUnauthorized` /
            // `KeyServiceRateLimited`) instead of the generic `NoDiscKey`.
            //
            // `KeyOutcome` deliberately gains no variant: it is matched
            // exhaustively by every front-end's trace renderer (freemkv's
            // `pipe::render_resolution_trace`, autorip's
            // `keysource::render_resolution_trace`), and this fix must not turn
            // into a breaking change across four repos to say something the error
            // code already says precisely.
            Err(e) => {
                if source_failure.is_none() {
                    source_failure = Some(e);
                }
                trace.keys.push(KeyStep {
                    who,
                    path: Vec::new(),
                    outcome: KeyOutcome::NoKey,
                });
            }
        }
    }
    // Nothing resolved. If a source FAILED rather than answered, stamp that
    // reason onto the disc so the decrypt gate can report it — but never clobber
    // a reason the scan already captured (e.g. `AacsVidUnavailable`), which is
    // closer to the disc itself than a source outage is.
    if let Some(e) = source_failure
        && disc.aacs_error.is_none()
    {
        disc.aacs_error = Some(e);
    }
    (false, trace)
}

/// THE single key-fetch: drive `sources` in order and return the first non-empty
/// Unit Key set. This is exactly what both paths do — only the samples differ:
/// * at disc open, `ctx` carries reachable-content samples → resolves the
///   up-front CPS units (the common one),
/// * in the read, on a decrypt miss, `ctx` carries the FAILING unit's ciphertext
///   → resolves the CPS unit that wasn't sampled up front.
///
/// Same sources, same call; there is no separate "fetch". Unlike
/// [`resolve_and_apply`] this does not validate/commit to a disc — the read's
/// decorator re-decrypts with the returned keys, which is the validation.
pub fn fetch_unit_keys(sources: &[Box<dyn KeySource>], ctx: &dyn ResolveCtx) -> Vec<UnitKey> {
    drive_unit_keys(sources, ctx).keys
}

/// Whether a driver run resolved keys, and — when it did NOT — whether the miss
/// was a genuine "no source holds this key" (`errored == false`) or at least one
/// source FAILED (`errored == true`, e.g. a network source was unreachable). The
/// distinction gates negative-result memoization: an empty-because-absent result
/// is safe to cache, an empty-because-a-source-was-down result is transient and
/// must NOT be cached (the key may resolve once the source recovers).
struct FetchOutcome {
    keys: Vec<UnitKey>,
    errored: bool,
}

/// [`fetch_unit_keys`] plus the error signal: drive `sources` in order, return the
/// first source's non-empty Unit Keys, and flag whether any source that failed to
/// answer did so with an `Err` (a source failure) rather than an empty `Ok`
/// (genuine absence — see [`KeySource::get_unit_keys`]).
fn drive_unit_keys(sources: &[Box<dyn KeySource>], ctx: &dyn ResolveCtx) -> FetchOutcome {
    let mut errored = false;
    for source in sources {
        match source.get_unit_keys(ctx) {
            Ok(uks) if !uks.is_empty() => {
                return FetchOutcome {
                    keys: uks,
                    errored: false,
                };
            }
            Ok(_) => {}
            Err(_) => errored = true,
        }
    }
    FetchOutcome {
        keys: Vec::new(),
        errored,
    }
}

/// The forensic counterpart to [`fetch_unit_keys`]: drive `sources` in order and
/// return the first source's non-empty AACS 2.1 forensic index set. `ctx` carries
/// the index-1 anchor batch (the mux, which owns disc geometry, gathers it and
/// injects it as the ctx's samples); a source that needs no samples (a keydb
/// keying on `disc_hash`) ignores them. Whatever the winning source returns —
/// ≥ 1 key — is trusted as the COMPLETE ordered set; no fixed count is assumed.
pub fn fetch_fmts_indexes(sources: &[Box<dyn KeySource>], ctx: &dyn ResolveCtx) -> Vec<UnitKey> {
    drive_fmts_indexes(sources, ctx).keys
}

/// [`fetch_fmts_indexes`] plus the error signal (see [`drive_unit_keys`]): the
/// forensic counterpart that flags whether any source `Err`ed during the miss.
fn drive_fmts_indexes(sources: &[Box<dyn KeySource>], ctx: &dyn ResolveCtx) -> FetchOutcome {
    let mut errored = false;
    for source in sources {
        match source.get_fmts_indexes(ctx) {
            Ok(uks) if !uks.is_empty() => {
                return FetchOutcome {
                    keys: uks,
                    errored: false,
                };
            }
            Ok(_) => {}
            Err(_) => errored = true,
        }
    }
    FetchOutcome {
        keys: Vec::new(),
        errored,
    }
}

/// Build the read-time [`crate::sector::KeyFetch`] from the disc's public AACS
/// inputs and a way to (re)build the application's key sources. The returned
/// resolver has the two explicit operations the mux and recovery decorator call:
/// [`unit_keys`](crate::sector::KeyFetch::unit_keys) drives [`fetch_unit_keys`]
/// (base per-CPS-unit keys), [`fmts_indexes`](crate::sector::KeyFetch::fmts_indexes)
/// drives [`fetch_fmts_indexes`] (the AACS 2.1 forensic set). Each is handed the
/// caller's sample batch as the ctx's `samples`, so a source pulls whatever
/// material it needs.
///
/// One builder, used by every read path (sweep / patch / mux) and every consumer
/// (CLI, autorip) — neither application contains the fetch logic, only its
/// key-source config. Cheap to clone; build once, clone into each read path.
/// `make_sources` is invoked per fetch (the cold path) so the resolver stays
/// `Send + Sync` without requiring `KeySource: Send`.
pub fn key_fetch(
    inputs: DiscInputs,
    make_sources: std::sync::Arc<dyn Fn() -> Vec<Box<dyn KeySource>> + Send + Sync>,
) -> crate::sector::KeyFetch {
    // One driver behind both operations: rebuild the sources, inject `samples`
    // as the ctx's content samples, run `drive` (the per-kind fetch), map the
    // resolved UnitKeys to raw keys. Memoized by the fingerprint of the sample
    // batch: the resolved keys are disc-level (a clip's index / CPS keys are
    // identical for every title that references it), so the first batch resolves
    // over the network and every repeat is answered from the cache with no
    // request. A GENUINELY-empty reply (every source ran and none held the key)
    // is cached too — the key the service lacks for a batch won't appear on a
    // re-ask, so re-hitting the network buys nothing. But an empty reply caused
    // by a source FAILURE (network down, source unreachable) is NOT cached: that
    // is a transient miss, and caching it would permanently drop a unit that
    // could be recovered once the source recovers — the `errored` flag on
    // `FetchOutcome` draws exactly that line. Each operation gets its OWN cache:
    // a base batch and a forensic anchor never collide, and the same bytes could
    // legitimately resolve differently per op.
    // The per-kind driver: `drive_unit_keys` or `drive_fmts_indexes`.
    type FetchDriver = fn(&[Box<dyn KeySource>], &dyn ResolveCtx) -> FetchOutcome;
    fn make_op(
        inputs: DiscInputs,
        make_sources: std::sync::Arc<dyn Fn() -> Vec<Box<dyn KeySource>> + Send + Sync>,
        drive: FetchDriver,
    ) -> crate::sector::KeyFetchFn {
        let cache: std::sync::Arc<std::sync::Mutex<std::collections::HashMap<u64, Vec<[u8; 16]>>>> =
            std::sync::Arc::new(std::sync::Mutex::new(std::collections::HashMap::new()));
        std::sync::Arc::new(move |samples: &[Vec<u8>]| -> Vec<[u8; 16]> {
            let fp = {
                use std::hash::{Hash, Hasher};
                let mut h = std::collections::hash_map::DefaultHasher::new();
                samples.len().hash(&mut h);
                for s in samples {
                    s.hash(&mut h);
                }
                h.finish()
            };
            if let Some(hit) = cache.lock().unwrap_or_else(|e| e.into_inner()).get(&fp) {
                return hit.clone();
            }
            let sources = make_sources();
            let mut di = inputs.clone();
            di.samples = samples.to_vec();
            // Parse Unit_Key_RO.inf at the disc's OWN stride (carried on `inputs`):
            // an online /decode reply that returns a VUK (not a terminal UK) then
            // derives unit keys from `enc_title_keys`, which a V10 disc parses at
            // the 48-byte stride — hardcoding the V20 stride here corrupted them.
            let ctx = DiscInputsCtx::new(&di);
            let outcome = drive(&sources, &ctx);
            let keys: Vec<[u8; 16]> = outcome.keys.into_iter().map(|u| u.key).collect();
            // Memoize a positive result always; memoize a NEGATIVE (empty) result
            // only when it is a genuine absence, never when a source errored — a
            // transient outage must not permanently poison this fingerprint.
            if !keys.is_empty() || !outcome.errored {
                cache
                    .lock()
                    .unwrap_or_else(|e| e.into_inner())
                    .insert(fp, keys.clone());
            }
            keys
        })
    }
    let unit = make_op(inputs.clone(), make_sources.clone(), drive_unit_keys);
    let fmts = make_op(inputs, make_sources, drive_fmts_indexes);
    crate::sector::KeyFetch::new(unit, fmts)
}

/// Read up to `n` ENCRYPTED 6144-byte aligned units from `title`'s body, raw (no
/// decrypt) — the content samples that populate [`DiscInputs::samples`] for a
/// key server to validate a candidate against, and that [`resolve_and_apply`]
/// hands to [`crate::Disc::decrypt_with`].
///
/// Lives in the library, not a key-source crate: reading the disc and carving
/// AACS units is decryption *mechanism* (unit geometry anchored at each extent's
/// `start_lba`), which the library owns. A key source is *handed* these bytes
/// via `DiscInputs.samples`; it never reads the disc itself.
///
/// "Encrypted" is decided by [`crate::aacs::content::aacs_unit_encrypted`] — the
/// AACS Copy Permission Indicator (CPI) in the top 2 bits of byte 0, the
/// spec-correct signal (`buf[0] & 0xc0`). NOT the `is_clean` TS-sync
/// heuristic: a unit lacking clean TS syncs does not imply encryption (an FMTS
/// variant frame or an odd clear unit can lack syncs yet be unencrypted), and a clear
/// unit sent to a key server yields nothing to validate against — the "0
/// encrypted units" rejection. A clip opens with clear navigation units (PAT/PMT,
/// menus) whose CPI is clear; only CPI-flagged content units are collected —
/// probing several points spread across EACH extent so a title whose encrypted
/// body starts late (or whose midpoint lands in clear nav) still yields samples.
/// CPI is read at each extent's `start_lba` (clip-file-anchored), so byte 0 is a
/// real unit start and the flag is meaningful.
pub fn read_encrypted_units(
    reader: &mut dyn crate::sector::SectorSource,
    title: &crate::disc::DiscTitle,
    n: usize,
) -> Vec<Vec<u8>> {
    use crate::aacs::content::{ALIGNED_UNIT_LEN, ALIGNED_UNIT_SECTORS, aacs_unit_encrypted};
    const CHUNK_UNITS: u32 = 15; // 45 sectors/read — under the drive transfer cap
    // Probe several evenly-spaced points across EACH extent rather than only the
    // midpoint-and-forward: a title whose encrypted feature starts late, or whose
    // midpoint lands in a clear nav stretch, must STILL yield scrambled samples.
    // Empty samples make `Disc::decrypt_with` skip wrong-key validation, so a
    // real encrypted title returning nothing here is a silent wrong-key hazard.
    const PROBES_PER_EXTENT: u32 = 8;

    let mut out: Vec<Vec<u8>> = Vec::new();
    for ext in &title.extents {
        let total_units = ext.sector_count / ALIGNED_UNIT_SECTORS;
        if total_units == 0 {
            continue;
        }
        for p in 1..=PROBES_PER_EXTENT {
            // Probe at p/(P+1) of the extent — spreads P points across it while
            // skipping the clear nav at the very head.
            let unit = ((total_units as u64 * p as u64) / (PROBES_PER_EXTENT as u64 + 1)) as u32;
            if unit >= total_units {
                continue;
            }
            let units_this = CHUNK_UNITS.min(total_units - unit);
            // Saturate: start_lba comes from attacker-controlled UDF/MPLS
            // extents; a malformed extent near u32::MAX would otherwise panic
            // (debug) or wrap to a wrong LBA (release). An over-capacity LBA then
            // fails cleanly via the read_sectors().is_err() skip below.
            let lba = ext
                .start_lba
                .saturating_add(unit.saturating_mul(ALIGNED_UNIT_SECTORS));
            let count = (units_this * ALIGNED_UNIT_SECTORS) as u16;
            let mut buf = vec![0u8; count as usize * 2048];
            // `false` = no recovery retries; the reader is the raw drive/file
            // (no decrypt decorator), so these are the on-disc encrypted bytes. A
            // read error at one probe skips THAT probe only — it must not abandon
            // the rest of the extent (the old `break` blinded the sampler on a
            // single transient miss).
            if reader.read_sectors(lba, count, &mut buf, false).is_err() {
                continue;
            }
            for i in 0..units_this as usize {
                let o = i * ALIGNED_UNIT_LEN;
                if o + ALIGNED_UNIT_LEN > buf.len() {
                    break;
                }
                let u = &buf[o..o + ALIGNED_UNIT_LEN];
                if aacs_unit_encrypted(u, title.content_format) {
                    out.push(u.to_vec());
                    if out.len() >= n {
                        return out;
                    }
                }
            }
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::aacs::types::UnitKey;
    use std::sync::{Arc, Mutex};

    fn units(n: usize) -> Vec<Vec<u8>> {
        (0..n).map(|i| vec![i as u8; 4]).collect()
    }

    // ── DecodeSampleSet: the online request can't be built under-sized ─────────

    /// Fewer than MIN_SAMPLE_UNITS → no set. Mutation: accepting a short slice
    /// resurrects the exact autorip bug (a 4-sample request silently skipped /
    /// read as "service down").
    #[test]
    fn decode_sample_set_rejects_under_min() {
        for n in 0..MIN_SAMPLE_UNITS {
            assert!(
                DecodeSampleSet::new(units(n)).is_none(),
                "{n} samples (< {MIN_SAMPLE_UNITS}) must not build a DecodeSampleSet"
            );
        }
    }

    /// Exactly the minimum, and above it, construct — and expose all samples.
    #[test]
    fn decode_sample_set_accepts_min_and_above() {
        let exact = DecodeSampleSet::new(units(MIN_SAMPLE_UNITS)).expect("min builds");
        assert_eq!(exact.len(), MIN_SAMPLE_UNITS);
        assert_eq!(exact.units().len(), MIN_SAMPLE_UNITS);
        assert!(!exact.is_empty());

        let more = DecodeSampleSet::new(units(MIN_SAMPLE_UNITS + 5)).expect("above min builds");
        assert_eq!(more.len(), MIN_SAMPLE_UNITS + 5);
    }

    /// The wrapped units round-trip byte-for-byte (the request carries exactly what
    /// was gathered — no reordering/truncation).
    #[test]
    fn decode_sample_set_preserves_units() {
        let raw = units(MIN_SAMPLE_UNITS);
        let set = DecodeSampleSet::new(raw.clone()).unwrap();
        assert_eq!(set.units(), raw.as_slice());
    }

    // ── KeySource default-method behaviour ────────────────────────────────────

    /// KeySource::host_certs() defaults to empty regardless of the MKB argument.
    /// Spec: a source holding no cert returns the empty vec; the `mkb` param is
    /// forward-looking and the default ignores it.
    /// Mutation: a default returning a non-empty vec would inject phantom certs
    ///           into the OEM handshake.
    #[test]
    fn key_source_host_certs_defaults_to_empty() {
        struct MinimalSource;
        impl KeySource for MinimalSource {
            fn get_unit_keys(&self, _ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
                Ok(Vec::new())
            }
        }
        let s = MinimalSource;
        assert!(s.host_certs(None).is_empty());
        assert!(s.host_certs(Some(68)).is_empty());
    }

    /// DiscInputsCtx maps DiscInputs faithfully: zero VID → None, non-zero VID →
    /// Some; title from volume_label; samples truncate to n; enc_title_keys
    /// parses Unit_Key_RO.inf at the version stride.
    #[test]
    fn disc_inputs_ctx_maps_fields() {
        // Build a minimal V10 Unit_Key_RO.inf with one key (stride 48):
        // uk_pos = 32, num_uk = 1, key at uk_pos + 48 = 80.
        let mut uk_ro = vec![0u8; 96];
        let uk_pos = 32usize;
        uk_ro[0..4].copy_from_slice(&(uk_pos as u32).to_be_bytes());
        uk_ro[uk_pos] = 0x00;
        uk_ro[uk_pos + 1] = 0x01; // num_unit_keys = 1
        let key_bytes = [0x7Eu8; 16];
        uk_ro[80..96].copy_from_slice(&key_bytes);

        let inputs = DiscInputs {
            disc_hash: "0xABC".into(),
            volume_id: [0u8; 16],
            version: crate::aacs::mkb::AACS_MAJOR_BD,
            mkb: vec![1, 2, 3],
            unit_key_ro: uk_ro,
            samples: vec![vec![9u8; 4], vec![8u8; 4], vec![7u8; 4]],
            volume_label: Some("TITLE_X".into()),
        };

        // Zero VID → None.
        let ctx = DiscInputsCtx::new(&inputs);
        assert_eq!(ctx.disc_hash(), "0xABC");
        assert_eq!(ctx.title(), Some("TITLE_X"));
        assert!(ctx.vid().is_none(), "all-zero VID is the no-VID sentinel");
        assert_eq!(ctx.mkb().unwrap(), &[1, 2, 3]);
        assert_eq!(ctx.enc_title_keys().unwrap(), &[key_bytes]);
        assert_eq!(ctx.samples(2).unwrap().len(), 2, "samples truncates to n");

        // Non-zero VID → Some(vid).
        let mut inputs2 = inputs.clone();
        inputs2.volume_id = [0x42u8; 16];
        let ctx2 = DiscInputsCtx::new(&inputs2);
        assert_eq!(ctx2.vid(), Some(Vid([0x42u8; 16])));
    }

    /// `resolve_and_apply_traced` records each step's `who` as the source's own
    /// `label()`, carried verbatim — no enum round-trip. A source with a custom
    /// label surfaces it as-is in the trace.
    #[test]
    fn trace_who_is_the_source_label_verbatim() {
        struct LabeledSource(&'static str);
        impl KeySource for LabeledSource {
            fn get_unit_keys(&self, _ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
                Ok(Vec::new())
            }
            fn label(&self) -> &'static str {
                self.0
            }
        }
        let mut disc = crate::Disc {
            volume_id: String::new(),
            meta_title: None,
            format: crate::DiscFormat::BluRay,
            capacity_sectors: 0,
            capacity_bytes: 0,
            layers: 1,
            titles: Vec::new(),
            region: crate::disc::DiscRegion::Free,
            aacs: None,
            css: None,
            encrypted: false,
            aacs_error: None,
            css_error: None,
            content_format: crate::ContentFormat::BdTs,
        };
        let inputs = DiscInputs {
            disc_hash: "0x00".into(),
            volume_id: [0u8; 16],
            version: crate::aacs::mkb::AACS_MAJOR_UHD,
            mkb: Vec::new(),
            unit_key_ro: Vec::new(),
            samples: Vec::new(),
            volume_label: None,
        };
        let sources: Vec<Box<dyn KeySource>> = vec![
            Box::new(LabeledSource("keydb")),
            Box::new(LabeledSource("my-custom-source")),
        ];
        let (_ok, trace) = resolve_and_apply_traced(&sources, &inputs, &mut disc);
        let whos: Vec<&str> = trace.keys.iter().map(|s| s.who.as_str()).collect();
        assert_eq!(whos, vec!["keydb", "my-custom-source"]);
    }

    // ── fetch_unit_keys / key_fetch (the one shared fetch path) ───────────────

    fn empty_inputs() -> DiscInputs {
        DiscInputs {
            disc_hash: String::new(),
            volume_id: [0u8; 16],
            version: crate::aacs::mkb::AACS_MAJOR_UHD,
            mkb: Vec::new(),
            unit_key_ro: Vec::new(),
            samples: Vec::new(),
            volume_label: None,
        }
    }

    struct EmptySource;
    impl KeySource for EmptySource {
        fn get_unit_keys(&self, _ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
            Ok(Vec::new())
        }
    }
    struct ErroringSource;
    impl KeySource for ErroringSource {
        fn get_unit_keys(&self, _ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
            Err(Error::AacsNoKeys)
        }
    }
    struct HasKey([u8; 16]);
    impl KeySource for HasKey {
        fn get_unit_keys(&self, _ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
            Ok(vec![UnitKey::new(0, self.0)])
        }
    }

    /// `fetch_unit_keys` returns the FIRST source's non-empty keys, skipping a
    /// source that returns empty or errors; empty when no source answers.
    #[test]
    fn fetch_unit_keys_first_nonempty_skips_empty_and_errors() {
        let inputs = empty_inputs();
        let ctx = DiscInputsCtx::new(&inputs);
        let key = [0xABu8; 16];

        let sources: Vec<Box<dyn KeySource>> = vec![
            Box::new(EmptySource),
            Box::new(ErroringSource),
            Box::new(HasKey(key)),
        ];
        let got = fetch_unit_keys(&sources, &ctx);
        assert_eq!(got.len(), 1, "the first source that answers wins");
        assert_eq!(got[0].key, key);

        let none: Vec<Box<dyn KeySource>> = vec![Box::new(EmptySource), Box::new(ErroringSource)];
        assert!(
            fetch_unit_keys(&none, &ctx).is_empty(),
            "no source answers ⇒ empty"
        );
    }

    /// `key_fetch` builds a closure that runs the sources with the GIVEN failing
    /// samples and returns their keys — the exact bytes are forwarded to the
    /// source, and `make_sources` is invoked per call.
    #[test]
    fn key_fetch_closure_forwards_samples_and_returns_keys() {
        let key = [0x5au8; 16];
        let seen: Arc<Mutex<Vec<Vec<u8>>>> = Arc::new(Mutex::new(Vec::new()));
        let builds = Arc::new(Mutex::new(0usize));

        struct Probe {
            key: [u8; 16],
            seen: Arc<Mutex<Vec<Vec<u8>>>>,
        }
        impl KeySource for Probe {
            fn get_unit_keys(&self, ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
                if let Ok(s) = ctx.samples(8) {
                    self.seen.lock().unwrap().extend(s);
                }
                Ok(vec![UnitKey::new(0, self.key)])
            }
        }

        let seen_c = Arc::clone(&seen);
        let builds_c = Arc::clone(&builds);
        let make: Arc<dyn Fn() -> Vec<Box<dyn KeySource>> + Send + Sync> = Arc::new(move || {
            *builds_c.lock().unwrap() += 1;
            vec![Box::new(Probe {
                key,
                seen: Arc::clone(&seen_c),
            }) as Box<dyn KeySource>]
        });

        let cb = key_fetch(empty_inputs(), make);
        let samples = vec![vec![0xEEu8; crate::aacs::content::ALIGNED_UNIT_LEN]];
        let got = cb.unit_keys(&samples);
        assert_eq!(
            got,
            vec![key],
            "the source's key flows back through the closure"
        );
        assert_eq!(
            seen.lock().unwrap().len(),
            1,
            "the failing ciphertext sample is forwarded to the source"
        );
        assert_eq!(*builds.lock().unwrap(), 1, "make_sources invoked per fetch");
    }

    /// `key_fetch` memoizes each operation by the fingerprint of the sample batch:
    /// identical samples reuse the cached keys (no rebuild), different samples miss,
    /// the two operations keep independent caches, and even an empty reply is cached.
    #[test]
    fn key_fetch_memoizes_per_op_by_sample_fingerprint() {
        let builds = Arc::new(Mutex::new(0usize));
        let builds_c = Arc::clone(&builds);
        let key = [0x11u8; 16];
        let make: Arc<dyn Fn() -> Vec<Box<dyn KeySource>> + Send + Sync> = Arc::new(move || {
            *builds_c.lock().unwrap() += 1;
            vec![Box::new(HasKey(key)) as Box<dyn KeySource>]
        });
        let cb = key_fetch(empty_inputs(), make);
        let a = vec![vec![0xAAu8; 8]];
        let b = vec![vec![0xBBu8; 8]];

        // First resolve for `a` builds sources; the identical repeat is cached.
        assert_eq!(cb.unit_keys(&a), vec![key]);
        assert_eq!(cb.unit_keys(&a), vec![key]);
        assert_eq!(
            *builds.lock().unwrap(),
            1,
            "identical samples reuse the cache"
        );

        // A different sample batch is a cache miss → one more build.
        assert_eq!(cb.unit_keys(&b), vec![key]);
        assert_eq!(
            *builds.lock().unwrap(),
            2,
            "different samples miss the cache"
        );

        // The forensic op has its OWN cache (HasKey has no forensic keys → empty),
        // so `a` builds once more here; its empty reply is then cached too.
        assert!(cb.fmts_indexes(&a).is_empty());
        assert_eq!(
            *builds.lock().unwrap(),
            3,
            "unit/fmts caches are independent"
        );
        assert!(cb.fmts_indexes(&a).is_empty());
        assert_eq!(
            *builds.lock().unwrap(),
            3,
            "an empty reply is cached, not re-asked"
        );
    }

    /// A transient source outage must NOT be memoized as a permanent "no key":
    /// a fingerprint whose first fetch failed because the source errored must be
    /// re-asked, and once the source recovers the key resolves. Regression guard
    /// for the negative-result memoization fix — caching the errored empty would
    /// permanently drop a recoverable unit for the rest of the op.
    #[test]
    fn errored_empty_is_not_cached_and_retries_when_source_recovers() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let key = [0x77u8; 16];
        // Shared across every `make_sources()` rebuild: call 0 errors (source
        // down), every later call succeeds (source recovered).
        let calls = Arc::new(AtomicUsize::new(0));

        struct Flaky {
            calls: Arc<AtomicUsize>,
            key: [u8; 16],
        }
        impl KeySource for Flaky {
            fn get_unit_keys(&self, _ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
                if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
                    Err(Error::AacsNoKeys) // first attempt: source unreachable
                } else {
                    Ok(vec![UnitKey::new(0, self.key)])
                }
            }
        }

        let calls_c = Arc::clone(&calls);
        let make: Arc<dyn Fn() -> Vec<Box<dyn KeySource>> + Send + Sync> = Arc::new(move || {
            vec![Box::new(Flaky {
                calls: Arc::clone(&calls_c),
                key,
            }) as Box<dyn KeySource>]
        });

        let cb = key_fetch(empty_inputs(), make);
        let samples = vec![vec![0xCDu8; 8]];

        // First fetch: the source errors → empty, but the miss must NOT be cached.
        assert!(
            cb.unit_keys(&samples).is_empty(),
            "source down → empty this time"
        );
        // Second fetch, SAME samples: not blocked by a cached empty → the now-
        // recovered source resolves the key.
        assert_eq!(
            cb.unit_keys(&samples),
            vec![key],
            "recovered source resolves — errored empty was not memoized"
        );
    }

    /// A GENUINE absence (a source that runs and returns an empty `Ok`) is still
    /// memoized — the benefit the fix preserves. A source counting its calls must
    /// be asked exactly once for a fingerprint whose first (clean) reply was empty.
    #[test]
    fn genuine_empty_is_still_memoized() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let calls = Arc::new(AtomicUsize::new(0));

        struct AlwaysEmpty {
            calls: Arc<AtomicUsize>,
        }
        impl KeySource for AlwaysEmpty {
            fn get_unit_keys(&self, _ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
                self.calls.fetch_add(1, Ordering::SeqCst);
                Ok(Vec::new()) // ran fine, genuinely holds no key
            }
        }

        let calls_c = Arc::clone(&calls);
        let make: Arc<dyn Fn() -> Vec<Box<dyn KeySource>> + Send + Sync> = Arc::new(move || {
            vec![Box::new(AlwaysEmpty {
                calls: Arc::clone(&calls_c),
            }) as Box<dyn KeySource>]
        });

        let cb = key_fetch(empty_inputs(), make);
        let samples = vec![vec![0xEFu8; 8]];

        assert!(cb.unit_keys(&samples).is_empty());
        assert!(cb.unit_keys(&samples).is_empty());
        assert_eq!(
            calls.load(Ordering::SeqCst),
            1,
            "a clean empty reply is cached — the source is asked only once"
        );
    }

    /// The two `KeyFetch` operations route to the two DISTINCT trait methods:
    /// `unit_keys` drives `get_unit_keys`, `fmts_indexes` drives
    /// `get_fmts_indexes`. A source that returns different keys per method proves
    /// the seam no longer collapses "1 base key" and "the forensic set" into one
    /// overloaded call — the operation, not the return length, decides which.
    #[test]
    fn key_fetch_routes_unit_and_fmts_to_distinct_source_methods() {
        const BASE: [u8; 16] = [0xB0; 16];
        const F1: [u8; 16] = [0xF1; 16];
        const F2: [u8; 16] = [0xF2; 16];

        struct TwoOp;
        impl KeySource for TwoOp {
            fn get_unit_keys(&self, _ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
                Ok(vec![UnitKey::new(0, BASE)])
            }
            fn get_fmts_indexes(&self, _ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
                Ok(vec![UnitKey::new(0, F1), UnitKey::new(1, F2)])
            }
        }

        let make: Arc<dyn Fn() -> Vec<Box<dyn KeySource>> + Send + Sync> =
            Arc::new(|| vec![Box::new(TwoOp) as Box<dyn KeySource>]);
        let cb = key_fetch(empty_inputs(), make);
        let samples = vec![vec![0x01u8; 4]];

        assert_eq!(
            cb.unit_keys(&samples),
            vec![BASE],
            "unit_keys resolves the base Unit Key via get_unit_keys"
        );
        assert_eq!(
            cb.fmts_indexes(&samples),
            vec![F1, F2],
            "fmts_indexes resolves the forensic set (any length) via get_fmts_indexes"
        );
    }

    /// `KeyFetch::unit_only` serves base keys but NEVER a forensic set — the
    /// contract the sweep/patch recovery decorator relies on (it resolves CPS
    /// units only). Its `fmts_indexes` is unconditionally empty.
    #[test]
    fn key_fetch_unit_only_never_serves_forensic() {
        let f = crate::sector::KeyFetch::unit_only(std::sync::Arc::new(|_| vec![[0xAA; 16]]));
        assert_eq!(f.unit_keys(&[vec![0u8; 4]]), vec![[0xAA; 16]]);
        assert!(
            f.fmts_indexes(&[vec![0u8; 4]]).is_empty(),
            "unit_only resolver yields no forensic keys"
        );
    }

    /// `get_fmts_indexes` defaults to empty, so a base-only source (a keydb) opts
    /// out of the forensic path without implementing it. `fetch_fmts_indexes` then
    /// falls through to the next source, exactly like the unit-key driver.
    #[test]
    fn fetch_fmts_indexes_skips_default_optout_source() {
        struct BaseOnly; // uses the default (empty) get_fmts_indexes
        impl KeySource for BaseOnly {
            fn get_unit_keys(&self, _ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
                Ok(vec![UnitKey::new(0, [0x11; 16])])
            }
        }
        struct Forensic;
        impl KeySource for Forensic {
            fn get_unit_keys(&self, _ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
                Ok(Vec::new())
            }
            fn get_fmts_indexes(&self, _ctx: &dyn ResolveCtx) -> Result<Vec<UnitKey>, Error> {
                Ok(vec![UnitKey::new(0, [0x77; 16])])
            }
        }
        let inputs = empty_inputs();
        let ctx = DiscInputsCtx::new(&inputs);
        let sources: Vec<Box<dyn KeySource>> = vec![Box::new(BaseOnly), Box::new(Forensic)];
        let got = fetch_fmts_indexes(&sources, &ctx);
        assert_eq!(got.len(), 1);
        assert_eq!(got[0].key, [0x77; 16], "the base-only source is skipped");
    }

    /// #4 regression: encrypted content NOT at the extent midpoint (a late-
    /// starting feature, or a midpoint landing in clear nav) must still be
    /// sampled — empty samples make `decrypt_with` skip wrong-key validation.
    /// The old midpoint-and-forward sampler returned empty; the probe-spread
    /// finds the early scrambled band.
    #[test]
    fn read_encrypted_units_finds_scrambled_content_off_the_midpoint() {
        use crate::aacs::content::{ALIGNED_UNIT_LEN, ALIGNED_UNIT_SECTORS, aacs_unit_encrypted};
        use crate::error::Result;
        use crate::sector::SectorSource;

        // Units in the FIRST SIXTH of the extent are scrambled (0xFF → no TS
        // sync); everything else (incl. the midpoint) is clear (0x47 syncs).
        struct BandSource {
            ext_start: u32,
            total_units: u32,
        }
        impl SectorSource for BandSource {
            fn capacity_sectors(&self) -> u32 {
                self.ext_start + self.total_units * ALIGNED_UNIT_SECTORS + 64
            }
            fn read_sectors(
                &mut self,
                lba: u32,
                count: u16,
                buf: &mut [u8],
                _r: bool,
            ) -> Result<usize> {
                let bytes = count as usize * 2048;
                for (i, chunk) in buf[..bytes].chunks_mut(ALIGNED_UNIT_LEN).enumerate() {
                    if chunk.len() < ALIGNED_UNIT_LEN {
                        break;
                    }
                    let abs_unit = (lba - self.ext_start) / ALIGNED_UNIT_SECTORS + i as u32;
                    if abs_unit < self.total_units / 6 {
                        chunk.fill(0xFF); // scrambled: no TS sync
                    } else {
                        chunk.fill(0);
                        let mut o = 4;
                        while o < ALIGNED_UNIT_LEN {
                            chunk[o] = 0x47; // clear TS syncs
                            o += 192;
                        }
                    }
                }
                Ok(bytes)
            }
        }

        let total_units = 600u32;
        let ext_start = 1000u32;
        let mut src = BandSource {
            ext_start,
            total_units,
        };
        let title = crate::disc::DiscTitle {
            playlist: String::new(),
            playlist_id: 0,
            duration_secs: 0.0,
            size_bytes: 0,
            clips: Vec::new(),
            streams: Vec::new(),
            chapters: Vec::new(),
            extents: vec![crate::disc::Extent {
                start_lba: ext_start,
                sector_count: total_units * ALIGNED_UNIT_SECTORS,
            }],
            content_format: crate::disc::ContentFormat::BdTs,
            codec_privates: Vec::new(),
        };

        let samples = read_encrypted_units(&mut src, &title, 4);
        assert!(
            !samples.is_empty(),
            "the probe-spread must sample the early scrambled band the midpoint misses"
        );
        for s in &samples {
            assert!(
                aacs_unit_encrypted(s, crate::disc::ContentFormat::BdTs),
                "every sample is a CPI-flagged encrypted unit (byte0 & 0xC0 != 0)"
            );
        }
    }

    /// DISCRIMINATING: selection is by the AACS CPI (byte 0), NOT the
    /// TS-sync clarity heuristic. Half the units lack TS syncs but are
    /// CPI-CLEAR (`byte0 & 0xC0 == 0`) — genuinely UNencrypted units that merely
    /// lack TS syncs; the old sampler collected these and the key server rejected
    /// the POST as "0 encrypted units". `read_encrypted_units` must skip them and
    /// return ONLY CPI-flagged units. A regression to selecting by TS-sync clarity
    /// would collect the CPI-clear units too and fail the `& 0xC0` assertion.
    #[test]
    fn read_encrypted_units_selects_by_cpi_not_ts_sync() {
        use crate::aacs::content::{ALIGNED_UNIT_LEN, ALIGNED_UNIT_SECTORS, aacs_unit_encrypted};
        use crate::error::Result;
        use crate::sector::SectorSource;

        // Even units: CPI-clear (byte0 & 0xC0 == 0) AND sync-destroyed (no 0x47).
        // Odd units:  CPI-set (byte0 = 0xC0) with a scrambled body.
        // Neither has clean TS syncs, so `is_clean` is FALSE for BOTH;
        // `aacs_unit_encrypted` flags only the odd units.
        struct MixSource {
            ext_start: u32,
            total_units: u32,
        }
        impl SectorSource for MixSource {
            fn capacity_sectors(&self) -> u32 {
                self.ext_start + self.total_units * ALIGNED_UNIT_SECTORS + 64
            }
            fn read_sectors(
                &mut self,
                lba: u32,
                count: u16,
                buf: &mut [u8],
                _r: bool,
            ) -> Result<usize> {
                let bytes = count as usize * 2048;
                for (i, chunk) in buf[..bytes].chunks_mut(ALIGNED_UNIT_LEN).enumerate() {
                    if chunk.len() < ALIGNED_UNIT_LEN {
                        break;
                    }
                    let abs = (lba - self.ext_start) / ALIGNED_UNIT_SECTORS + i as u32;
                    if abs.is_multiple_of(2) {
                        chunk.fill(0x11); // CPI-clear (0x11 & 0xC0 == 0), no TS sync
                    } else {
                        chunk.fill(0xAB); // scrambled body (no TS sync)
                        chunk[0] = 0xC0; // CPI set -> encrypted
                    }
                }
                Ok(bytes)
            }
        }

        let total_units = 400u32;
        let ext_start = 500u32;
        let mut src = MixSource {
            ext_start,
            total_units,
        };
        let title = crate::disc::DiscTitle {
            playlist: String::new(),
            playlist_id: 0,
            duration_secs: 0.0,
            size_bytes: 0,
            clips: Vec::new(),
            streams: Vec::new(),
            chapters: Vec::new(),
            extents: vec![crate::disc::Extent {
                start_lba: ext_start,
                sector_count: total_units * ALIGNED_UNIT_SECTORS,
            }],
            content_format: crate::disc::ContentFormat::BdTs,
            codec_privates: Vec::new(),
        };

        let samples = read_encrypted_units(&mut src, &title, 8);
        assert!(
            !samples.is_empty(),
            "the CPI-flagged (odd) units must still be collected"
        );
        for s in &samples {
            assert!(
                aacs_unit_encrypted(s, crate::disc::ContentFormat::BdTs),
                "only CPI-flagged units are selected"
            );
            assert_eq!(
                s[0] & 0xC0,
                0xC0,
                "a CPI-clear sync-destroyed unit must never be sampled"
            );
        }
    }

    /// Audit #5 — a DISCRIMINATING test for the version→stride fix. A 2-key
    /// `Unit_Key_RO.inf` whose SECOND key sits at the V20 (64-byte) offset; a V10
    /// (48-byte) parse reads a DIFFERENT region. Confirms `DiscInputsCtx` parses
    /// at the stride for `inputs.version` — a swapped `from_major` branch or a
    /// hardcoded stride (the exact bug 1.2.0 fixes) would fail this, where the
    /// prior single-key fixtures passed regardless of stride.
    #[test]
    fn disc_inputs_ctx_parses_unit_keys_at_the_version_stride() {
        use crate::aacs::mkb::{AACS_MAJOR_BD, AACS_MAJOR_UHD};
        const UK_POS: usize = 64;
        let mut inf = vec![0u8; 200];
        inf[0..4].copy_from_slice(&(UK_POS as u32).to_be_bytes()); // uk_pos
        inf[UK_POS..UK_POS + 2].copy_from_slice(&2u16.to_be_bytes()); // num_uk = 2
        let key0_at = UK_POS + 48; // first key — same for both strides
        let key1_v10_at = key0_at + 48; // second key if parsed at V10 stride
        let key1_v20_at = key0_at + 64; // second key if parsed at V20 stride
        inf[key0_at..key0_at + 16].fill(0xA0);
        inf[key1_v10_at..key1_v10_at + 16].fill(0x10);
        inf[key1_v20_at..key1_v20_at + 16].fill(0x20);

        let base = DiscInputs {
            disc_hash: String::new(),
            volume_id: [0u8; 16],
            version: AACS_MAJOR_UHD,
            mkb: Vec::new(),
            unit_key_ro: inf,
            samples: Vec::new(),
            volume_label: None,
        };
        let k20 = DiscInputsCtx::new(&base).enc_title_keys().unwrap().to_vec();
        let v10_inputs = DiscInputs {
            version: AACS_MAJOR_BD,
            ..base.clone()
        };
        let k10 = DiscInputsCtx::new(&v10_inputs)
            .enc_title_keys()
            .unwrap()
            .to_vec();

        assert_eq!(k20.len(), 2);
        assert_eq!(k10.len(), 2);
        assert_eq!(k20[0], [0xA0; 16], "first key is at +48 for both strides");
        assert_eq!(k10[0], [0xA0; 16]);
        assert_eq!(k20[1], [0x20; 16], "V20 reads the 2nd key at +64");
        assert_eq!(k10[1], [0x10; 16], "V10 reads the 2nd key at +48");
        assert_ne!(k20[1], k10[1], "the parse stride follows inputs.version");
    }

    /// `DiscInputs` is public and returned by `Disc::inputs`, so any consumer's
    /// `tracing::debug!("{inputs:?}")` prints it. A derived `Debug` printed the
    /// Volume ID (the value `aacs::types::Vid` deliberately renders as
    /// `Vid(<redacted>)`), the whole `Unit_Key_RO.inf` (the encrypted title keys),
    /// the entire MKB and every ciphertext sample verbatim. Sentinel byte
    /// 0xD5 = decimal 213, matching `aacs::types::redaction_tests`. Mutation
    /// guard: restoring `#[derive(Debug)]` fails this.
    #[test]
    fn disc_inputs_debug_is_redacted() {
        let inputs = DiscInputs {
            disc_hash: "0xAA".into(),
            volume_id: [0xD5; 16],
            version: 2,
            mkb: vec![0xD5; 64],
            unit_key_ro: vec![0xD5; 48],
            samples: vec![vec![0xD5; 6144]],
            volume_label: Some("TITLE_2024".into()),
        };
        let dbg = format!("{inputs:?}");
        assert!(
            !dbg.contains("213"),
            "DiscInputs Debug leaked key material (decimal 213): {dbg}"
        );
        assert!(
            dbg.contains("redacted"),
            "DiscInputs Debug missing redaction marker: {dbg}"
        );
        // Non-secret identity and shape stay printable for diagnostics.
        assert!(dbg.contains("0xAA"), "{dbg}");
        assert!(dbg.contains("mkb_len: 64"), "{dbg}");
        assert!(dbg.contains("unit_key_ro_len: 48"), "{dbg}");
        assert!(dbg.contains("samples_len: 1"), "{dbg}");
        assert!(dbg.contains("TITLE_2024"), "{dbg}");
    }

    /// `DecodeSampleSet` is public and wraps the SAME on-disc ciphertext the
    /// sibling `DiscInputs` redacts, so a derived `Debug` dumped ≥ MIN_SAMPLE_UNITS
    /// × 6144 bytes of verbatim AACS ciphertext (plus every unit's clear 16-byte
    /// derivation seed) into any log that formatted it. Sentinel byte 0xD5 =
    /// decimal 213, matching `aacs::types::redaction_tests` and the
    /// `DiscInputs` test above. Mutation guard: restoring `#[derive(Debug)]`
    /// fails this.
    #[test]
    fn decode_sample_set_debug_is_redacted() {
        let set = DecodeSampleSet::new(vec![vec![0xD5; 6144]; MIN_SAMPLE_UNITS])
            .expect("MIN_SAMPLE_UNITS units is a valid set");
        let dbg = format!("{set:?}");
        assert!(
            !dbg.contains("213"),
            "DecodeSampleSet Debug leaked ciphertext (decimal 213): {dbg}"
        );
        assert!(
            dbg.contains("redacted"),
            "DecodeSampleSet Debug missing redaction marker: {dbg}"
        );
        // Non-secret shape stays printable for diagnostics.
        assert!(dbg.contains("units_len: 8"), "{dbg}");
    }
}