From b34af1fa7499050d2328f7ea139f9174135fab4b Mon Sep 17 00:00:00 2001
From: MattJackson <1085847+MattJackson@users.noreply.github.com>
Date: Thu, 4 Jun 2026 20:28:37 -0700
Subject: [PATCH] mux: wire TrueHD channel probe into iso:// mux setup
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

correct_truehd_channels() probes the first DECRYPTED access units of the
chosen title (TrueHD PIDs, bounded 8 MiB read of the first extent) and
sets AudioStream.channels from the MLP major sync — fixing the MPLS
audio_format understatement (5.1 declared on a 7.1/Atmos TrueHD track).
Regenerates the basic codec label for the corrected count; richer
editorial labels are left untouched. Wired in resolve.rs input() for
iso:// after decrypt_with (the m2ts is only decryptable post-key), using
a fresh reader so the mux reader is undisturbed. generate_audio_label
made pub(crate).
---
 src/disc/mod.rs    | 73 ++++++++++++++++++++++++++++++++++++++++++++++
 src/labels/mod.rs  |  2 +-
 src/mux/resolve.rs |  9 +++++-
 3 files changed, 82 insertions(+), 2 deletions(-)

diff --git a/src/disc/mod.rs b/src/disc/mod.rs
index 9d3d26a..dfb73a5 100644
--- a/src/disc/mod.rs
+++ b/src/disc/mod.rs
@@ -375,6 +375,79 @@ pub struct Extent {
     pub sector_count: u32,
 }
 
+/// Correct the channel count of a title's TrueHD audio streams by probing the
+/// first decrypted access units. The MPLS `audio_format` field declares the
+/// BASE layout (often 5.1) even for a 7.1/Atmos TrueHD track; the real count is
+/// in the MLP major sync. `reader` must yield DECRYPTED sectors (the m2ts is
+/// AACS-encrypted, so this can only run at mux time, not scan). Reads a bounded
+/// window of the title's first extent. Also regenerates the stream's codec
+/// label when it was the basic descriptor for the (now corrected) count —
+/// richer editorial labels (e.g. "Dolby Atmos") are left untouched.
+pub(crate) fn correct_truehd_channels(reader: &mut dyn SectorSource, title: &mut DiscTitle) {
+    use crate::mux::codec::truehd::truehd_channels_from_stream;
+
+    let pids: Vec<u16> = title
+        .streams
+        .iter()
+        .filter_map(|s| match s {
+            Stream::Audio(a) if matches!(a.codec, Codec::TrueHd) => Some(a.pid),
+            _ => None,
+        })
+        .collect();
+    if pids.is_empty() {
+        return;
+    }
+    let Some(ext) = title.extents.first() else {
+        return;
+    };
+    // Bounded probe: up to 8 MiB from the start of the title — enough for the
+    // first interleaved TrueHD major sync of each stream.
+    const PROBE_SECTORS: u32 = 4096;
+    let n = ext.sector_count.min(PROBE_SECTORS) as u16;
+    if n == 0 {
+        return;
+    }
+    let mut buf = vec![0u8; n as usize * 2048];
+    if reader
+        .read_sectors(ext.start_lba, n, &mut buf, true)
+        .is_err()
+    {
+        return;
+    }
+
+    let mut demux = crate::mux::ts::TsDemuxer::new(&pids);
+    let mut payloads: std::collections::HashMap<u16, Vec<u8>> = std::collections::HashMap::new();
+    for pes in demux.feed(&buf).into_iter().chain(demux.flush()) {
+        payloads
+            .entry(pes.pid)
+            .or_default()
+            .extend_from_slice(&pes.data);
+    }
+
+    for s in title.streams.iter_mut() {
+        let Stream::Audio(a) = s else { continue };
+        if !matches!(a.codec, Codec::TrueHd) {
+            continue;
+        }
+        let Some(payload) = payloads.get(&a.pid) else {
+            continue;
+        };
+        let Some(count) = truehd_channels_from_stream(payload) else {
+            continue;
+        };
+        let new_ch = AudioChannels::from_count(count);
+        if new_ch == AudioChannels::Unknown || new_ch == a.channels {
+            continue;
+        }
+        let was_basic =
+            a.label == crate::labels::generate_audio_label(&a.codec, &a.channels, a.secondary);
+        a.channels = new_ch;
+        if was_basic {
+            a.label = crate::labels::generate_audio_label(&a.codec, &new_ch, a.secondary);
+        }
+    }
+}
+
 /// Calculate how many bytes of bad/unreadable data fall within a title's extents.
 /// `pub(crate)` so autorip can use it for main-movie lost_ms computation.
 pub fn bytes_bad_in_title(title: &DiscTitle, bad_ranges: &[(u64, u64)]) -> u64 {
diff --git a/src/labels/mod.rs b/src/labels/mod.rs
index 72d3ff7..6ccdde1 100644
--- a/src/labels/mod.rs
+++ b/src/labels/mod.rs
@@ -412,7 +412,7 @@ fn codec_hint_adds_detail(hint: &str) -> bool {
     h.contains("atmos") || h.contains("dts:x") || h.contains("dts-x") || h.contains("dtsx")
 }
 
-fn generate_audio_label(
+pub(crate) fn generate_audio_label(
     codec: &crate::disc::Codec,
     channels: &crate::disc::AudioChannels,
     _secondary: bool,
diff --git a/src/mux/resolve.rs b/src/mux/resolve.rs
index 1ae7009..e71fa8a 100644
--- a/src/mux/resolve.rs
+++ b/src/mux/resolve.rs
@@ -216,8 +216,15 @@ 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();
+            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