sources: stateful one-key-at-a-time providers, UK-first keydb, shared resolve loop

Each source implements next_key (a cursor over its candidates) instead of
returning them all at once. The keydb hands its per-disc candidates out
UK-first (UK > VK > MK > DK) so a stale/wrong VUK never pre-empts a good UK in
the same entry; online and mapfile are one-shot. MultiSource composes sources
in the caller's chosen order and resolve_and_apply drives the
next_key -> decrypt_with loop, stopping at the first key that decrypts.
read_sample_units moves here so the CLI and autorip share one content sampler.

src/keydb.rs

@@ -16,17 +16,24 @@
 use std::path::PathBuf;
 
 use libfreemkv::aacs::{HostCert, KeyDb};
-use libfreemkv::{DiscInputs, Key, KeySource, Result};
+use libfreemkv::{DiscInputs, Key, KeySource};
 
 /// A [`KeySource`] backed by a local `keydb.cfg` file.
 pub struct KeydbSource {
     path: PathBuf,
+    /// Lazily-built candidate list (UK ▸ VK ▸ MK ▸ DK ▸ …) plus its cursor —
+    /// the keydb owns the order and hands one candidate per `next_key`. `None`
+    /// until the first `next_key` parses the file.
+    cursor: Option<std::vec::IntoIter<Key>>,
 }
 
 impl KeydbSource {
     /// A keydb source reading the given `keydb.cfg` path.
     pub fn new(path: impl Into<PathBuf>) -> Self {
-        Self { path: path.into() }
+        Self {
+            path: path.into(),
+            cursor: None,
+        }
     }
 
     /// The host certificate(s) in this keydb — the second kind of data the one
@@ -42,17 +49,28 @@ impl KeydbSource {
 
     /// Build the ordered candidate list from a parsed keydb. Pure (no I/O), so
     /// it is unit-testable without a file on disk.
+    ///
+    /// Order = cheapest + most authoritative first: **UK ▸ VK ▸ MK ▸ DK**. The
+    /// UK is the final per-CPS-unit content key — zero derivation, directly
+    /// usable — so it is tried first; the VUK needs one derivation step, an MK
+    /// two, and the device-key pool the full MKB walk (AACS-1.0-only, slowest),
+    /// so it is the last-resort fallback. Trying the UK first is also what lets a
+    /// stale/wrong per-disc VUK be skipped in favour of a good UK in the SAME
+    /// entry (`decrypt_with` rejects the VUK; the loop falls through to the UK).
     fn candidates_from(db: &KeyDb, inputs: &DiscInputs) -> Vec<Key> {
         let mut out = Vec::new();
 
         // Per-disc hit (most specific). find_disc normalizes the hash form.
         if let Some(entry) = db.find_disc(&inputs.disc_hash) {
-            if let Some(vuk) = entry.vuk {
-                out.push(Key::Volume(vuk));
-            }
+            // UK first — terminal content key, no derivation.
             if !entry.unit_keys.is_empty() {
                 out.push(Key::Unit(entry.unit_keys.clone()));
             }
+            // VK next — one step (decrypt Unit_Key_RO.inf).
+            if let Some(vuk) = entry.vuk {
+                out.push(Key::Volume(vuk));
+            }
+            // MK — two steps (derive the VUK, then the unit keys).
             if let Some(mk) = entry.media_key {
                 out.push(Key::Media(vec![mk]));
             }
@@ -80,15 +98,19 @@ impl KeydbSource {
 }
 
 impl KeySource for KeydbSource {
-    fn resolve(&self, inputs: &DiscInputs) -> Result<Vec<Key>> {
-        // A missing keydb is not an error — another source may have the key.
-        // (Parse/format problems surface as an empty/partial keydb, same as the
-        // library's own loader; this source never fails the whole resolve.)
-        let db = match KeyDb::load(&self.path) {
-            Ok(db) => db,
-            Err(_) => return Ok(Vec::new()),
+    fn next_key(&mut self, inputs: &DiscInputs) -> Option<Key> {
+        // On the first ask, parse the keydb once and build the ordered candidate
+        // list; later asks just advance the cursor. A missing/unreadable keydb
+        // is not an error — it simply yields no candidates (another source may
+        // have the key), the same as the library's own loader.
+        if self.cursor.is_none() {
+            let cands = match KeyDb::load(&self.path) {
+                Ok(db) => Self::candidates_from(&db, inputs),
+                Err(_) => Vec::new(),
             };
-        Ok(Self::candidates_from(&db, inputs))
+            self.cursor = Some(cands.into_iter());
+        }
+        self.cursor.as_mut().and_then(Iterator::next)
     }
 }
 
@@ -150,6 +172,32 @@ mod tests {
         );
     }
 
+    #[test]
+    fn per_disc_uk_ranks_before_vuk() {
+        // An entry with BOTH a UK and a VUK (the Being There shape) must hand the
+        // terminal UK out first, so a stale/wrong VUK never pre-empts a good UK.
+        let mut entries = HashMap::new();
+        let mut e = entry_with_vuk("0xaabb", [0x11u8; 16]);
+        e.unit_keys = vec![(1, [0x22u8; 16])];
+        entries.insert("0xaabb".into(), e);
+        let db = KeyDb {
+            device_keys: Vec::new(),
+            processing_keys: Vec::new(),
+            host_certs: Vec::new(),
+            disc_entries: entries,
+        };
+
+        let cands = KeydbSource::candidates_from(&db, &inputs("0xaabb"));
+        assert!(
+            matches!(cands.first(), Some(Key::Unit(_))),
+            "the terminal UK must be the first candidate"
+        );
+        assert!(
+            matches!(cands.get(1), Some(Key::Volume(v)) if *v == [0x11u8; 16]),
+            "the VUK follows the UK"
+        );
+    }
+
     #[test]
     fn no_disc_hit_offers_only_universal_material() {
         let db = KeyDb {

src/lib.rs

@@ -12,9 +12,13 @@
 //! their own config — the local-vs-online policy is just which impls they plug
 //! in — then resolve and hand the resulting key to `Disc::decrypt_with`.
 //!
-//! Sources are dumb: they enumerate the raw material they hold as candidate
-//! keys and do NO derivation or validation. The caller tries the candidates in
-//! order and keeps the first that decrypts a sample ([`resolve_first`]).
+//! Sources are dumb and stateful: each hands its candidate keys out one at a
+//! time via [`KeySource::next_key`], in its own best order, and reports
+//! exhaustion. Compose several with [`MultiSource`] in the caller's chosen
+//! order; [`resolve_and_apply`] drives the loop — handing each key to
+//! `Disc::decrypt_with` (which validates against the disc's content samples) and
+//! stopping at the first that decrypts, or reporting a genuine "no key" when
+//! every source is spent.
 
 mod keydb;
 mod mapfile;
@@ -28,33 +32,122 @@ pub use online::OnlineSource;
 // for the source-side types.
 pub use libfreemkv::{DiscInputs, Key, KeySource};
 
-use libfreemkv::Result;
+use libfreemkv::{Disc, DiscTitle, SectorSource};
 
-/// Try each source's candidate keys in order and return the first that the
-/// `accept` predicate approves — the *validate-before-return* policy.
-///
-/// `accept` is the caller's validation (typically: clone the disc, apply the
-/// key with `Disc::decrypt_with`, decrypt a sample sector, and check it looks
-/// like cleartext). It lives with the caller because only the caller can read
-/// disc content. A stale or wrong candidate is rejected and the next is tried,
-/// so a wrong keydb entry transparently falls through to the next source.
-///
-/// `Ok(None)` means no source offered a candidate the validator accepted; an
-/// `Err` from any source's `resolve` is propagated.
-pub fn resolve_first<F>(
-    sources: &[&dyn KeySource],
+/// An ordered composition of key sources, driven as one. `next_key` exhausts
+/// the first source (one candidate per call), then the next, … then `None`.
+/// **The caller supplies the list AND the order** — local-first `[Keydb,
+/// Online]`, online-first `[Online, Keydb]`, resume `[Mapfile, Keydb]`, etc. —
+/// so the "which sources, in what order" policy lives entirely with the
+/// application, not the library. `MultiSource` is itself a [`KeySource`], so it
+/// nests and composes.
+pub struct MultiSource {
+    sources: Vec<Box<dyn KeySource>>,
+    idx: usize,
+}
+
+impl MultiSource {
+    /// Compose the given sources, tried in the order supplied.
+    pub fn new(sources: Vec<Box<dyn KeySource>>) -> Self {
+        Self { sources, idx: 0 }
+    }
+}
+
+impl KeySource for MultiSource {
+    fn next_key(&mut self, inputs: &DiscInputs) -> Option<Key> {
+        while self.idx < self.sources.len() {
+            if let Some(key) = self.sources[self.idx].next_key(inputs) {
+                return Some(key);
+            }
+            self.idx += 1; // this source is spent — advance to the next
+        }
+        None
+    }
+
+    fn needs_samples(&self) -> bool {
+        self.sources.iter().any(|s| s.needs_samples())
+    }
+
+    fn errored(&self) -> bool {
+        self.sources.iter().any(|s| s.errored())
+    }
+}
+
+/// Drive `sources` until one key decrypts `disc`. Loops `next_key` and hands
+/// each candidate to [`Disc::decrypt_with`] (which validates it against
+/// `inputs.samples` and only mutates the disc on success), returning `true` at
+/// the first key that decrypts and `false` once every source is exhausted — the
+/// genuine "no key for this disc". THE shared key-resolution loop: every
+/// application (the `freemkv` CLI, autorip) uses it instead of re-rolling the
+/// candidate/retry logic, so the "no key" verdict is identical everywhere.
+pub fn resolve_and_apply(
+    sources: &mut dyn KeySource,
     inputs: &DiscInputs,
-    mut accept: F,
-) -> Result<Option<Key>>
-where
-    F: FnMut(&Key) -> bool,
-{
-    for src in sources {
-        for key in src.resolve(inputs)? {
-            if accept(&key) {
-                return Ok(Some(key));
+    disc: &mut Disc,
+) -> bool {
+    while let Some(key) = sources.next_key(inputs) {
+        if disc.decrypt_with(key, &inputs.samples).is_ok() {
+            return true;
+        }
+    }
+    false
+}
+
+/// Read up to `n` ENCRYPTED 6144-byte aligned units from `title`'s body, raw (no
+/// decrypt) — the content samples a caller hands to [`resolve_and_apply`] (for
+/// `Disc::decrypt_with` to validate a key against) and that a sample-needing
+/// source (an online key service) byte-validates against.
+///
+/// "Encrypted" is decided by `libfreemkv::aacs::is_aacs_scrambled` — the SAME
+/// predicate the library's decrypt gate and a key service use — so all sides
+/// agree. A clip opens with clear navigation units (PAT/PMT, menus); only the
+/// feature body is scrambled, and a clear unit proves nothing, so this collects
+/// only scrambled ones, sampling the largest extent at its midpoint forward.
+pub fn read_sample_units(
+    reader: &mut dyn SectorSource,
+    title: &DiscTitle,
+    n: usize,
+) -> Vec<Vec<u8>> {
+    const UNIT_LEN: usize = 6144;
+    const UNIT_SECTORS: u32 = 3; // 6144 / 2048
+    const CHUNK_UNITS: u32 = 15; // 45 sectors/read — under the drive transfer cap
+    const MAX_CHUNKS_PER_EXTENT: u32 = 4; // ~60 units scanned at each extent's midpoint
+
+    let mut out: Vec<Vec<u8>> = Vec::new();
+    for ext in &title.extents {
+        let total_units = ext.sector_count / UNIT_SECTORS;
+        if total_units == 0 {
+            continue;
+        }
+        let mut unit = total_units / 2; // midpoint (past the clear nav at the head)
+        for _ in 0..MAX_CHUNKS_PER_EXTENT {
+            if unit >= total_units {
+                break;
+            }
+            let units_this = CHUNK_UNITS.min(total_units - unit);
+            let lba = ext.start_lba + unit * UNIT_SECTORS;
+            let count = (units_this * 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.
+            if reader.read_sectors(lba, count, &mut buf, false).is_err() {
+                break;
+            }
+            for i in 0..units_this as usize {
+                let o = i * UNIT_LEN;
+                if o + UNIT_LEN > buf.len() {
+                    break;
+                }
+                let u = &buf[o..o + UNIT_LEN];
+                if libfreemkv::aacs::is_aacs_scrambled(u) {
+                    out.push(u.to_vec());
+                    if out.len() >= n {
+                        return out;
                     }
                 }
             }
-    Ok(None)
+            unit += units_this;
+        }
+    }
+    out
 }

src/mapfile.rs

@@ -12,31 +12,35 @@
 use std::path::PathBuf;
 
 use libfreemkv::disc::mapfile::Mapfile;
-use libfreemkv::{DiscInputs, Key, KeySource, Result};
+use libfreemkv::{DiscInputs, Key, KeySource};
 
 /// A [`KeySource`] backed by a rip mapfile's persisted unit keys.
 pub struct MapfileSource {
     path: PathBuf,
+    /// The mapfile holds exactly one (terminal) UK set, so it is read once —
+    /// this flips true after the first `next_key`.
+    asked: bool,
 }
 
 impl MapfileSource {
     /// A mapfile source reading the given `*.mapfile` path.
     pub fn new(path: impl Into<PathBuf>) -> Self {
-        Self { path: path.into() }
+        Self {
+            path: path.into(),
+            asked: false,
+        }
     }
 }
 
 impl KeySource for MapfileSource {
-    fn resolve(&self, _inputs: &DiscInputs) -> Result<Vec<Key>> {
+    fn next_key(&mut self, _inputs: &DiscInputs) -> Option<Key> {
+        if self.asked {
+            return None;
+        }
+        self.asked = true;
         // A missing/unreadable/keyless mapfile simply offers nothing.
-        let Ok(map) = Mapfile::load(&self.path) else {
-            return Ok(Vec::new());
-        };
+        let map = Mapfile::load(&self.path).ok()?;
         let uks = map.unit_keys();
-        if uks.is_empty() {
-            Ok(Vec::new())
-        } else {
-            Ok(vec![Key::Unit(uks.to_vec())])
-        }
+        (!uks.is_empty()).then(|| Key::Unit(uks.to_vec()))
     }
 }

src/online.rs

@@ -3,7 +3,7 @@
 use std::time::Duration;
 
 use base64::Engine;
-use libfreemkv::{DiscInputs, Key, KeySource, Result};
+use libfreemkv::{DiscInputs, Key, KeySource};
 
 const MAX_MKB_BYTES: usize = 10 * 1024 * 1024;
 const TIMEOUT_SECS: u64 = 180;
@@ -11,6 +11,14 @@ const TIMEOUT_SECS: u64 = 180;
 pub struct OnlineSource {
     base_url: String,
     secret: String,
+    /// The key service pre-validates server-side and returns a single UK, so it
+    /// is asked **at most once** — this flips true after the first `next_key`,
+    /// and every later ask returns `None` without re-hitting the network.
+    asked: bool,
+    /// Set when the round-trip itself failed (network down, bad response) — as
+    /// opposed to the service simply having no key. Lets the caller report
+    /// "key service unreachable" distinctly from "no key for this disc".
+    errored: bool,
 }
 
 impl OnlineSource {
@@ -18,14 +26,16 @@ impl OnlineSource {
         Self {
             base_url: base_url.into(),
             secret: secret.into(),
-        }
+            asked: false,
+            errored: false,
         }
     }
 
-impl KeySource for OnlineSource {
-    fn resolve(&self, inputs: &DiscInputs) -> Result<Vec<Key>> {
+    /// The single server-resolved UK for this disc, or `None`. Runs exactly the
+    /// one network round-trip; `next_key` gates it to one call per session.
+    fn query(&mut self, inputs: &DiscInputs) -> Option<Key> {
         if self.base_url.is_empty() || inputs.mkb.len() > MAX_MKB_BYTES {
-            return Ok(Vec::new());
+            return None;
         }
         let b64 = base64::engine::general_purpose::STANDARD;
         let mut body = serde_json::json!({
@@ -50,21 +60,43 @@ impl KeySource for OnlineSource {
         }
         let resp = match req.send_json(body) {
             Ok(r) => r,
-            Err(_) => return Ok(Vec::new()),
+            Err(_) => {
+                self.errored = true;
+                return None;
+            }
         };
         let json: serde_json::Value = match resp.into_json() {
             Ok(j) => j,
-            Err(_) => return Ok(Vec::new()),
-        };
-        match json.get("UK").and_then(|u| u.as_str()).and_then(parse_uk) {
-            Some(uk) => Ok(vec![Key::Unit(vec![(1, uk)])]),
-            None => Ok(Vec::new()),
+            Err(_) => {
+                self.errored = true;
+                return None;
             }
+        };
+        json.get("UK")
+            .and_then(|u| u.as_str())
+            .and_then(parse_uk)
+            .map(|uk| Key::Unit(vec![(1, uk)]))
+    }
+}
+
+impl KeySource for OnlineSource {
+    fn next_key(&mut self, inputs: &DiscInputs) -> Option<Key> {
+        // One shot: the service pre-validates and returns a single UK, so a
+        // second ask has nothing new to offer — don't re-hit the network.
+        if self.asked {
+            return None;
+        }
+        self.asked = true;
+        self.query(inputs)
     }
 
     fn needs_samples(&self) -> bool {
         true
     }
+
+    fn errored(&self) -> bool {
+        self.errored
+    }
 }
 
 fn parse_uk(hex: &str) -> Option<[u8; 16]> {