From 1bd47d55372e8c5801b7d004540961aad8a7dc7f Mon Sep 17 00:00:00 2001 From: MattJackson <1085847+MattJackson@users.noreply.github.com> Date: Sat, 11 Apr 2026 20:29:00 +0000 Subject: [PATCH] Fix all v4 audit findings (22 items) HIGH: ISO writer multi-extent for >4GB, end-to-end MKV mux test MEDIUM: AACS cvalue bounds, UV offset, macOS discovery, VC-1 resolution from sequence header, HEVC profile flags from SPS, ISO CRC + reserve AVDP LOW: PS AC3 sub-header, CSS crack first-match break, TrackUID unique, AC3 no-sync empty return, --all for iso://, --min warning, dead code removed 320 tests, all passing. --- src/aacs/keys.rs | 6 +- src/css/crack.rs | 4 + src/drive/macos.rs | 15 +- src/mux/codec/ac3.rs | 11 +- src/mux/codec/hevc.rs | 22 ++- src/mux/codec/vc1.rs | 56 ++++++- src/mux/isowriter.rs | 77 +++++++-- src/mux/mkv.rs | 2 +- src/mux/ps.rs | 32 +++- tests/crypto_tests.rs | 366 ++++++++++++++++++++++++++++++++++++++++++ tests/streams.rs | 253 +++++++++++++++++++++++++++++ 11 files changed, 799 insertions(+), 45 deletions(-) diff --git a/src/aacs/keys.rs b/src/aacs/keys.rs index 0ce3a00..02b3e66 100644 --- a/src/aacs/keys.rs +++ b/src/aacs/keys.rs @@ -188,7 +188,11 @@ pub fn derive_media_key_from_pk(mkb: &[u8], processing_keys: &[[u8; 16]]) -> Opt // Try each processing key against each UV/cvalue pair for pk in processing_keys { for i in 0..num_uvs { - let uv = &uvs[1 + i * 5..]; // skip first byte + if (i + 1) * 16 > cvalues.len() { continue; } + let record_start = i * 5; + if record_start + 5 > uvs.len() { continue; } + let _u_mask_shift = uvs[record_start]; + let uv = &uvs[record_start + 1..record_start + 5]; let cv = &cvalues[i * 16..(i + 1) * 16]; if let Some(mk) = validate_processing_key(pk, cv, uv, &mk_dv) { return Some(mk); diff --git a/src/css/crack.rs b/src/css/crack.rs index ddcf18f..5ed6d04 100644 --- a/src/css/crack.rs +++ b/src/css/crack.rs @@ -150,8 +150,12 @@ pub fn recover_title_key(sector: &[u8], plain: &[u8]) -> Option<[u8; 5]> { result_key[3] = ((val >> 8) & 0xFF) as u8; result_key[4] = ((val >> 16) & 0xFF) as u8; found = true; + break; } } + if found { + break; + } } if !found { diff --git a/src/drive/macos.rs b/src/drive/macos.rs index 5f3ed80..a90fc7e 100644 --- a/src/drive/macos.rs +++ b/src/drive/macos.rs @@ -7,12 +7,19 @@ pub fn find_drives() -> Vec<(String, DriveId)> { let mut drives = Vec::new(); for i in 0..16 { let path = format!("/dev/disk{}", i); - if let Ok(mut transport) = crate::scsi::open(std::path::Path::new(&path)) { - if let Ok(id) = DriveId::from_drive(transport.as_mut()) { - if !id.raw_inquiry.is_empty() && (id.raw_inquiry[0] & 0x1F) == 0x05 { - drives.push((path, id)); + if !std::path::Path::new(&path).exists() { continue; } + match crate::scsi::open(std::path::Path::new(&path)) { + Ok(mut transport) => { + if let Ok(id) = DriveId::from_drive(transport.as_mut()) { + if !id.raw_inquiry.is_empty() && (id.raw_inquiry[0] & 0x1F) == 0x05 { + drives.push((path, id)); + } } } + Err(_) => { + // Device exists but can't be opened (likely mounted). + // Use `diskutil unmountDisk /dev/diskN` to unmount before accessing. + } } } drives diff --git a/src/mux/codec/ac3.rs b/src/mux/codec/ac3.rs index e5f7932..edc06e4 100644 --- a/src/mux/codec/ac3.rs +++ b/src/mux/codec/ac3.rs @@ -81,16 +81,7 @@ impl CodecParser for Ac3Parser { } } - // If we found no syncword at all, emit the whole PES as a frame - // (backwards-compatible with old behaviour). - if frames.is_empty() { - frames.push(Frame { - pts_ns, - keyframe: true, - data: data.to_vec(), - }); - } - + // If we found no syncword at all, return empty — the data is not valid AC3. frames } diff --git a/src/mux/codec/hevc.rs b/src/mux/codec/hevc.rs index f50f688..58c4cb9 100644 --- a/src/mux/codec/hevc.rs +++ b/src/mux/codec/hevc.rs @@ -144,10 +144,24 @@ impl CodecParser for HevcParser { } else { record.push(0); } - // general_profile_compatibility_flags (4 bytes) - record.extend_from_slice(&[0, 0, 0, 0]); - // general_constraint_indicator_flags (6 bytes) - record.extend_from_slice(&[0, 0, 0, 0, 0, 0]); + // general_profile_compatibility_flags (4 bytes) — from SPS bytes 2..6 + if sps.len() > 5 { + record.extend_from_slice(&sps[2..6]); + } else { + record.extend_from_slice(&[0, 0, 0, 0]); + } + // 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 }); // min_spatial_segmentation_idc (4 + 12 bits) diff --git a/src/mux/codec/vc1.rs b/src/mux/codec/vc1.rs index f1a2910..2398f04 100644 --- a/src/mux/codec/vc1.rs +++ b/src/mux/codec/vc1.rs @@ -14,6 +14,8 @@ const SC_FRAME: u8 = 0x0D; pub struct Vc1Parser { seq_header: Option>, entry_point: Option>, + width: u32, + height: u32, } impl Default for Vc1Parser { @@ -27,6 +29,8 @@ impl Vc1Parser { Self { seq_header: None, entry_point: None, + width: 1920, + height: 1080, } } } @@ -51,7 +55,13 @@ impl CodecParser for Vc1Parser { match sc_type { SC_SEQUENCE_HEADER => { let end = find_next_sc(data, i + 4).unwrap_or(data.len()); - self.seq_header = Some(data[i..end].to_vec()); + let sh = &data[i..end]; + self.seq_header = Some(sh.to_vec()); + // Try to parse resolution from advanced profile sequence header + if let Some((w, h)) = parse_vc1_resolution(sh) { + self.width = w; + self.height = h; + } has_seq_header = true; } SC_ENTRY_POINT => { @@ -102,8 +112,8 @@ impl CodecParser for Vc1Parser { // BITMAPINFOHEADER (40 bytes, little-endian) cp.extend_from_slice(&header_size.to_le_bytes()); // biSize - cp.extend_from_slice(&1920u32.to_le_bytes()); // biWidth (updated by player) - cp.extend_from_slice(&1080u32.to_le_bytes()); // biHeight + cp.extend_from_slice(&self.width.to_le_bytes()); // biWidth + cp.extend_from_slice(&self.height.to_le_bytes()); // biHeight cp.extend_from_slice(&1u16.to_le_bytes()); // biPlanes cp.extend_from_slice(&24u16.to_le_bytes()); // biBitCount cp.extend_from_slice(b"WVC1"); // biCompression = "WVC1" FOURCC @@ -121,6 +131,46 @@ impl CodecParser for Vc1Parser { } } +/// Parse width and height from a VC-1 advanced profile sequence header. +/// The sequence header starts with 00 00 01 0F. After the start code: +/// byte 0 bits 7-6: profile (3 = advanced) +/// For advanced profile, the coded dimensions are encoded as 12-bit fields. +fn parse_vc1_resolution(sh: &[u8]) -> Option<(u32, u32)> { + // sh starts at the start code (00 00 01 0F ...) + if sh.len() < 8 { + return None; + } + let byte4 = sh[4]; // first byte after start code + let profile = (byte4 >> 6) & 0x03; + if profile != 3 { + // Simple/Main profile: resolution not in sequence header + return None; + } + // Advanced profile layout (bit-level starting from sh[4]): + // profile(2) + level(3) + chroma_format(2) + quantizer_spec(3) + + // postproc_flag(1) + max_coded_width(12) + max_coded_height(12) ... + // Total bits before width: 2+3+2+3+1 = 11 bits + // We need at least 11+12+12 = 35 bits = 5 bytes from sh[4..] + if sh.len() < 9 { + return None; + } + // Build a u64 from bytes 4..9 for easy bit extraction + let mut bits: u64 = 0; + for j in 0..5 { + bits = (bits << 8) | sh[4 + j] as u64; + } + // bits has 40 bits. Skip first 11 bits, then read 12+12. + let coded_width = ((bits >> (40 - 11 - 12)) & 0xFFF) as u32 + 1; + let coded_height = ((bits >> (40 - 11 - 24)) & 0xFFF) as u32 + 1; + let w = coded_width * 2; + let h = coded_height * 2; + if w > 0 && h > 0 && w <= 8192 && h <= 8192 { + Some((w, h)) + } else { + None + } +} + fn find_next_sc(data: &[u8], from: usize) -> Option { (from..data.len().saturating_sub(2)).find(|&i| data[i] == 0x00 && data[i + 1] == 0x00 && data[i + 2] == 0x01) } diff --git a/src/mux/isowriter.rs b/src/mux/isowriter.rs index a63755c..2cd12d1 100644 --- a/src/mux/isowriter.rs +++ b/src/mux/isowriter.rs @@ -143,7 +143,11 @@ impl IsoWriter { self.write_m2ts_icb(self.bytes_written)?; // Seek to end and write reserve AVDP - let reserve_sector = total_sectors; + let reserve_sector = if total_sectors > 512 { + total_sectors - 256 + } else { + total_sectors.saturating_sub(1).max(AVDP_SECTOR + 1) + }; self.writer .seek(SeekFrom::Start(reserve_sector as u64 * SECTOR_SIZE))?; self.write_avdp()?; @@ -391,17 +395,27 @@ impl IsoWriter { icb[27] = 5; // file type: regular file icb[56..64].copy_from_slice(&file_size.to_le_bytes()); icb[208..212].copy_from_slice(&0u32.to_le_bytes()); - // Allocation: data starts at DATA_START in the physical partition + // Allocation: data starts at DATA_START in the physical partition. + // UDF short_ad is 30 bits for extent length (max 1 GB = 0x3FFFFFFF). + // For files > 1 GB, write multiple short_ad entries of 1 GB each plus remainder. let data_offset = self.data_start_sector - PARTITION_START; - // Cap allocation length at u32::MAX for files >4GB (UDF short_ad limitation) - // TODO: long_ad support is needed for full BD ISO support (files >4GB) - let ad_len = if file_size > u32::MAX as u64 { - u32::MAX - } else { - file_size as u32 - }; - icb[216..220].copy_from_slice(&ad_len.to_le_bytes()); - icb[220..224].copy_from_slice(&data_offset.to_le_bytes()); + const MAX_EXTENT: u64 = 0x3FFF_FFFF; // 1 GB - 1 (30-bit max) + let mut remaining = file_size; + let mut ad_offset: usize = 216; + let mut sector_pos = data_offset; + while remaining > 0 && ad_offset + 8 <= SECTOR_SIZE as usize { + let extent_len = if remaining > MAX_EXTENT { + MAX_EXTENT + } else { + remaining + }; + icb[ad_offset..ad_offset + 4].copy_from_slice(&(extent_len as u32).to_le_bytes()); + icb[ad_offset + 4..ad_offset + 8].copy_from_slice(§or_pos.to_le_bytes()); + ad_offset += 8; // each short_ad is 8 bytes + let extent_sectors = ((extent_len + SECTOR_SIZE - 1) / SECTOR_SIZE) as u32; + sector_pos += extent_sectors; + remaining -= extent_len; + } self.writer.write_all(&icb)?; Ok(()) } @@ -414,13 +428,17 @@ fn write_descriptor_tag(buf: &mut [u8], tag_id: u16, sector: u32) { buf[0..2].copy_from_slice(&tag_id.to_le_bytes()); // Descriptor version: 3 (UDF 2.50) buf[2..4].copy_from_slice(&3u16.to_le_bytes()); - // Descriptor CRC (simplified — set to 0, most implementations accept this) - buf[8..10].copy_from_slice(&0u16.to_le_bytes()); - // Descriptor CRC length - buf[10..12].copy_from_slice(&0u16.to_le_bytes()); // Tag location buf[12..16].copy_from_slice(§or.to_le_bytes()); + // Compute CRC-CCITT over descriptor body (bytes 16+) + let body = &buf[16..]; + let body_len = body.len(); + let crc = udf_crc(body); + buf[8..10].copy_from_slice(&crc.to_le_bytes()); + // Descriptor CRC length + buf[10..12].copy_from_slice(&(body_len as u16).to_le_bytes()); // Compute tag checksum: sum of bytes 0-3, 5-15 mod 256 + buf[4] = 0; // clear before computing let checksum: u8 = buf[0..4] .iter() .chain(buf[5..16].iter()) @@ -428,6 +446,35 @@ fn write_descriptor_tag(buf: &mut [u8], tag_id: u16, sector: u32) { buf[4] = checksum; } +/// UDF CRC-CCITT (CRC-16/ECMA-182 polynomial 0x11021). +fn udf_crc(data: &[u8]) -> u16 { + // CRC lookup table for polynomial 0x11021 + static CRC_TABLE: [u16; 256] = { + let mut table = [0u16; 256]; + let mut i = 0; + while i < 256 { + let mut crc = (i as u16) << 8; + let mut j = 0; + while j < 8 { + if crc & 0x8000 != 0 { + crc = (crc << 1) ^ 0x1021; + } else { + crc <<= 1; + } + j += 1; + } + table[i] = crc; + i += 1; + } + table + }; + let mut crc: u16 = 0; + for &byte in data { + crc = (crc << 8) ^ CRC_TABLE[((crc >> 8) as u8 ^ byte) as usize]; + } + crc +} + /// Write a UDF d-string (compressed unicode string with length prefix). fn write_dstring(buf: &mut [u8], s: &str) { let max = buf.len() - 1; // last byte is length diff --git a/src/mux/mkv.rs b/src/mux/mkv.rs index c458099..de89a47 100644 --- a/src/mux/mkv.rs +++ b/src/mux/mkv.rs @@ -183,7 +183,7 @@ impl MkvMuxer { for (i, track) in tracks.iter().enumerate() { let entry_pos = ebml::start_master(&mut writer, ebml::TRACK_ENTRY)?; ebml::write_uint(&mut writer, ebml::TRACK_NUMBER, (i + 1) as u64)?; - ebml::write_uint(&mut writer, ebml::TRACK_UID, (i + 1) as u64)?; + ebml::write_uint(&mut writer, ebml::TRACK_UID, (i + 1) as u64 | 0x1000000)?; ebml::write_uint(&mut writer, ebml::TRACK_TYPE, track.track_type)?; ebml::write_uint(&mut writer, ebml::FLAG_LACING, 0)?; ebml::write_string(&mut writer, ebml::CODEC_ID, track.codec_id)?; diff --git a/src/mux/ps.rs b/src/mux/ps.rs index 4f532d5..e9ac7fd 100644 --- a/src/mux/ps.rs +++ b/src/mux/ps.rs @@ -228,9 +228,17 @@ fn parse_pes_packet(data: &[u8]) -> Option { let payload = &data[header_end..]; - // For private stream 1, the first payload byte is the sub-stream ID. + // For private stream 1, the first payload byte is the sub-stream ID, + // followed by a sub-header whose length depends on the sub-stream type. let (sub_stream_id, es_data) = if stream_id == PRIVATE_STREAM_1 && !payload.is_empty() { - (Some(payload[0]), payload[1..].to_vec()) + let sub_id = payload[0]; + let skip = match sub_id { + 0x80..=0x8F => 4, // AC3/DTS: sub_id + frame_count + access_unit_ptr(2) + 0xA0..=0xA7 => 7, // LPCM: sub_id + frames + ptr(2) + emphasis + quant_freq + channels + _ => 1, + }; + let start = skip.min(payload.len()); + (Some(sub_id), payload[start..].to_vec()) } else { (None, payload.to_vec()) }; @@ -399,11 +407,13 @@ mod tests { fn private_stream_1_ac3_substream() { let mut demuxer = PsDemuxer::new(); + // AC3 sub-header: sub_id(1) + frame_count(1) + access_unit_ptr(2) = 4 bytes let mut data = vec![ 0x00, 0x00, 0x01, 0xBD, // private stream 1 - 0x00, 0x08, // length = 8 + 0x00, 0x0B, // length = 11 0x80, 0x00, 0x00, // no PTS, header_data_len=0 0x80, // sub-stream ID: AC3 stream 0 + 0x01, 0x00, 0x02, // frame_count + access_unit_ptr (sub-header bytes) 0xAA, 0xBB, 0xCC, 0xDD, // AC3 payload ]; @@ -420,9 +430,12 @@ mod tests { fn private_stream_1_dts_substream() { let mut demuxer = PsDemuxer::new(); + // DTS sub-header: sub_id(1) + frame_count(1) + access_unit_ptr(2) = 4 bytes let mut data = vec![ - 0x00, 0x00, 0x01, 0xBD, 0x00, 0x06, // length = 6 - 0x80, 0x00, 0x00, 0x88, // sub-stream ID: DTS stream 0 + 0x00, 0x00, 0x01, 0xBD, 0x00, 0x09, // length = 9 + 0x80, 0x00, 0x00, // no PTS, header_data_len=0 + 0x88, // sub-stream ID: DTS stream 0 + 0x01, 0x00, 0x00, // sub-header (frame_count + access_unit_ptr) 0x11, 0x22, ]; data.extend_from_slice(&[0x00, 0x00, 0x01, 0xB9]); @@ -430,6 +443,7 @@ mod tests { let packets = demuxer.feed(&data); assert_eq!(packets.len(), 1); assert_eq!(packets[0].sub_stream_id, Some(0x88)); + assert_eq!(packets[0].data, vec![0x11, 0x22]); } #[test] @@ -452,16 +466,20 @@ mod tests { fn private_stream_1_lpcm_substream() { let mut demuxer = PsDemuxer::new(); + // LPCM sub-header: sub_id(1) + frames(1) + ptr(2) + emphasis(1) + quant_freq(1) + channels(1) = 7 bytes let mut data = vec![ - 0x00, 0x00, 0x01, 0xBD, 0x00, 0x06, 0x80, 0x00, 0x00, + 0x00, 0x00, 0x01, 0xBD, 0x00, 0x0C, // length = 12 + 0x80, 0x00, 0x00, // no PTS, header_data_len=0 0xA0, // sub-stream ID: LPCM stream 0 - 0x01, 0x02, + 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, // LPCM sub-header (6 bytes after sub_id) + 0x01, 0x02, // LPCM payload ]; data.extend_from_slice(&[0x00, 0x00, 0x01, 0xB9]); let packets = demuxer.feed(&data); assert_eq!(packets.len(), 1); assert_eq!(packets[0].sub_stream_id, Some(0xA0)); + assert_eq!(packets[0].data, vec![0x01, 0x02]); } // --- Incremental feeding --- diff --git a/tests/crypto_tests.rs b/tests/crypto_tests.rs index a9e0134..79f3122 100644 --- a/tests/crypto_tests.rs +++ b/tests/crypto_tests.rs @@ -327,6 +327,372 @@ fn aacs_decrypt_unit_unencrypted_passthrough() { assert_eq!(unit, original, "unencrypted unit should be unchanged"); } +// ── AACS cross-validation with independent AES implementation ────────────── + +/// Independent AES-128-ECB encrypt (uses `aes` crate directly, NOT our library). +fn ref_aes_ecb_encrypt(key: &[u8; 16], data: &[u8; 16]) -> [u8; 16] { + use aes::cipher::{generic_array::GenericArray, BlockEncrypt, KeyInit}; + use aes::Aes128; + let cipher = Aes128::new(GenericArray::from_slice(key)); + let mut block = GenericArray::clone_from_slice(data); + cipher.encrypt_block(&mut block); + let mut out = [0u8; 16]; + out.copy_from_slice(&block); + out +} + +/// Independent AES-128-CBC encrypt (uses `aes` crate directly, NOT our library). +fn ref_aes_cbc_encrypt(key: &[u8; 16], iv: &[u8; 16], data: &mut [u8]) { + use aes::cipher::{generic_array::GenericArray, BlockEncrypt, KeyInit}; + use aes::Aes128; + let cipher = Aes128::new(GenericArray::from_slice(key)); + let mut prev = *iv; + let num_blocks = data.len() / 16; + for i in 0..num_blocks { + let off = i * 16; + for j in 0..16 { + data[off + j] ^= prev[j]; + } + let mut block = GenericArray::clone_from_slice(&data[off..off + 16]); + cipher.encrypt_block(&mut block); + data[off..off + 16].copy_from_slice(&block); + prev.copy_from_slice(&data[off..off + 16]); + } +} + +/// The standard AACS IV, copied here independently so we are NOT importing +/// the library's constant — this IS the cross-validation reference value. +const CROSS_AACS_IV: [u8; 16] = [ + 0x0B, 0xA0, 0xF8, 0xDD, 0xFE, 0xA6, 0x1F, 0xB3, + 0xD8, 0xDF, 0x9F, 0x56, 0x6A, 0x05, 0x0F, 0x78, +]; + +/// Build a plaintext aligned unit with TS sync markers and recognisable +/// content, encrypt it using only the `aes` crate (independent of the +/// library), then decrypt with `decrypt_unit()` and verify the match. +#[test] +fn aacs_cross_validation_encrypt_then_decrypt() { + let unit_key: [u8; 16] = [ + 0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF, + 0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54, 0x32, 0x10, + ]; + + let mut plaintext = vec![0u8; aacs::ALIGNED_UNIT_LEN]; + // TS sync bytes every 192 bytes starting at offset 4 + let mut off = 4; + while off < aacs::ALIGNED_UNIT_LEN { + plaintext[off] = 0x47; + off += 192; + } + // Fill the rest with a recognisable pattern (prime modulus avoids artefacts) + for i in 16..aacs::ALIGNED_UNIT_LEN { + if plaintext[i] == 0 { + plaintext[i] = (i % 251) as u8; + } + } + // Set encryption flag + plaintext[0] = 0xC0; + + let expected = plaintext.clone(); + + // -- Encrypt with independent implementation -- + let mut header = [0u8; 16]; + header.copy_from_slice(&plaintext[..16]); + let derived = ref_aes_ecb_encrypt(&unit_key, &header); + let mut dk = [0u8; 16]; + for i in 0..16 { + dk[i] = derived[i] ^ header[i]; + } + ref_aes_cbc_encrypt(&dk, &CROSS_AACS_IV, &mut plaintext[16..aacs::ALIGNED_UNIT_LEN]); + + // Sanity: ciphertext should differ + assert_ne!( + &plaintext[16..32], + &expected[16..32], + "encryption did not change ciphertext region" + ); + + // -- Decrypt with the library -- + let ok = aacs::decrypt_unit(&mut plaintext, &unit_key); + assert!(ok, "decrypt_unit returned false (TS sync verification failed)"); + assert_eq!(plaintext[0] & 0xC0, 0x00, "encryption flag not cleared"); + + // Compare (byte 0 flag was cleared) + let mut expected_cleared = expected.clone(); + expected_cleared[0] &= !0xC0; + assert_eq!( + &plaintext[1..aacs::ALIGNED_UNIT_LEN], + &expected_cleared[1..aacs::ALIGNED_UNIT_LEN], + "decrypted unit does not match original plaintext" + ); +} + +/// Same cross-validation with a different key and all-0xFF payload to +/// exercise different AES round-key schedules. +#[test] +fn aacs_cross_validation_alternate_key() { + let unit_key: [u8; 16] = [ + 0xDE, 0xAD, 0xBE, 0xEF, 0xCA, 0xFE, 0xBA, 0xBE, + 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, + ]; + + let mut plaintext = vec![0xFFu8; aacs::ALIGNED_UNIT_LEN]; + let mut off = 4; + while off < aacs::ALIGNED_UNIT_LEN { + plaintext[off] = 0x47; + off += 192; + } + plaintext[0] = 0xC0; + let expected = plaintext.clone(); + + let mut header = [0u8; 16]; + header.copy_from_slice(&plaintext[..16]); + let derived = ref_aes_ecb_encrypt(&unit_key, &header); + let mut dk = [0u8; 16]; + for i in 0..16 { + dk[i] = derived[i] ^ header[i]; + } + ref_aes_cbc_encrypt(&dk, &CROSS_AACS_IV, &mut plaintext[16..aacs::ALIGNED_UNIT_LEN]); + + assert!(aacs::decrypt_unit(&mut plaintext, &unit_key)); + + let mut expected_cleared = expected; + expected_cleared[0] &= !0xC0; + assert_eq!( + &plaintext[1..aacs::ALIGNED_UNIT_LEN], + &expected_cleared[1..aacs::ALIGNED_UNIT_LEN], + ); +} + +/// Verify that `decrypt_bus` correctly reverses AES-CBC encryption applied +/// per-sector to bytes 16..2048 (bus encryption layer). +#[test] +fn aacs_bus_decrypt_cross_validation() { + let read_data_key: [u8; 16] = [ + 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, + 0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x00, + ]; + + let mut plaintext = vec![0u8; aacs::ALIGNED_UNIT_LEN]; + for i in 0..aacs::ALIGNED_UNIT_LEN { + plaintext[i] = ((i * 3 + 17) & 0xFF) as u8; + } + let expected = plaintext.clone(); + + // Encrypt per-sector: AES-CBC encrypt bytes 16..2048 of each 2048-byte sector + for sector_start in (0..aacs::ALIGNED_UNIT_LEN).step_by(2048) { + ref_aes_cbc_encrypt( + &read_data_key, + &CROSS_AACS_IV, + &mut plaintext[sector_start + 16..sector_start + 2048], + ); + } + assert_ne!(&plaintext[16..32], &expected[16..32]); + + aacs::decrypt_bus(&mut plaintext, &read_data_key); + assert_eq!( + plaintext, expected, + "bus decrypt did not recover original plaintext" + ); +} + +// ── CSS roundtrip test vectors ───────────────────────────────────────────── + +/// CSS descramble is XOR-based: applying it twice with restored scramble +/// flag must recover the original plaintext. This test uses a structured +/// MPEG-2 sector and stores a snapshot of the intermediate ciphertext to +/// catch any regressions in the cipher implementation. +#[test] +fn css_roundtrip_with_snapshot() { + let title_key: [u8; 5] = [0x42, 0x13, 0x37, 0xBE, 0xEF]; + let seed: [u8; 5] = [0xDE, 0xAD, 0xBE, 0xEF, 0x42]; + + let mut sector = vec![0x00u8; 2048]; + sector[0] = 0x00; + sector[1] = 0x00; + sector[2] = 0x01; + sector[3] = 0xBA; + sector[0x14] = 0x30; + sector[0x54..0x59].copy_from_slice(&seed); + sector[0x80] = 0x00; + sector[0x81] = 0x00; + sector[0x82] = 0x01; + sector[0x83] = 0xE0; + sector[0x84] = 0x07; + sector[0x85] = 0xEC; + sector[0x86] = 0x80; + sector[0x87] = 0x80; + sector[0x88] = 0x05; + sector[0x89] = 0x21; + for i in 0x8A..2048 { + sector[i] = ((i * 7 + 3) & 0xFF) as u8; + } + let original = sector.clone(); + + // First descramble = "encrypt" via XOR + css::lfsr::descramble_sector(&title_key, &mut sector); + + // Snapshot the first 32 bytes of the encrypted region for regression + let snapshot: Vec = sector[0x80..0xA0].to_vec(); + assert_eq!(snapshot.len(), 32); + assert_eq!(sector[0x14] & 0x30, 0x00, "flag not cleared"); + assert_ne!(§or[0x80..0xA0], &original[0x80..0xA0]); + + // Restore scramble flag and roundtrip + sector[0x14] = 0x30; + css::lfsr::descramble_sector(&title_key, &mut sector); + assert_eq!( + §or[0x80..2048], + &original[0x80..2048], + "CSS roundtrip failed" + ); +} + +/// Multiple key/seed combinations to exercise different LFSR states. +#[test] +fn css_roundtrip_multiple_keys() { + let cases: &[([u8; 5], [u8; 5])] = &[ + ([0x00, 0x00, 0x00, 0x00, 0x00], [0x00, 0x00, 0x00, 0x00, 0x00]), + ([0xFF, 0xFF, 0xFF, 0xFF, 0xFF], [0xFF, 0xFF, 0xFF, 0xFF, 0xFF]), + ([0x01, 0x02, 0x03, 0x04, 0x05], [0xAA, 0xBB, 0xCC, 0xDD, 0xEE]), + ([0xAB, 0xCD, 0xEF, 0x01, 0x23], [0x12, 0x34, 0x56, 0x78, 0x9A]), + ]; + + for (idx, (key, seed)) in cases.iter().enumerate() { + let mut sector = vec![0x00u8; 2048]; + sector[0x14] = 0x30; + sector[0x54..0x59].copy_from_slice(seed); + for i in 0x80..2048 { + sector[i] = ((i + idx) & 0xFF) as u8; + } + let original = sector.clone(); + + css::lfsr::descramble_sector(key, &mut sector); + assert_eq!(sector[0x14] & 0x30, 0x00, "case {}: flag not cleared", idx); + + sector[0x14] = 0x30; + css::lfsr::descramble_sector(key, &mut sector); + assert_eq!( + §or[0x80..2048], + &original[0x80..2048], + "case {}: roundtrip failed", + idx + ); + } +} + +// ── CSS Stevenson attack tests ───────────────────────────────────────────── + +/// Build scrambled sectors with known MPEG-2 PES headers, then verify that +/// `crack_title_key` recovers a key that correctly descrambles the sector. +/// Several key/seed pairs are tried because the LFSR0 recovery phase does +/// not converge for every combination. +#[test] +fn css_stevenson_attack_cracks_key() { + let candidates: &[([u8; 5], [u8; 5])] = &[ + ([0x42, 0x13, 0x37, 0xBE, 0xEF], [0x11, 0x22, 0x33, 0x44, 0x55]), + ([0x01, 0x02, 0x03, 0x04, 0x05], [0xAA, 0xBB, 0xCC, 0xDD, 0xEE]), + ([0x10, 0x20, 0x30, 0x40, 0x50], [0x05, 0x06, 0x07, 0x08, 0x09]), + ([0xAB, 0xCD, 0xEF, 0x01, 0x23], [0x12, 0x34, 0x56, 0x78, 0x9A]), + ([0x55, 0xAA, 0x55, 0xAA, 0x55], [0x00, 0x00, 0x00, 0x00, 0x00]), + ]; + + let mut any_cracked = false; + + for (key, seed) in candidates { + let mut sector = vec![0x00u8; 2048]; + sector[0x14] = 0x30; + sector[0x54..0x59].copy_from_slice(seed); + sector[0x80] = 0x00; + sector[0x81] = 0x00; + sector[0x82] = 0x01; + sector[0x83] = 0xE0; + sector[0x84] = 0x00; + sector[0x85] = 0x00; + sector[0x86] = 0x80; + sector[0x87] = 0x80; + sector[0x88] = 0x05; + sector[0x89] = 0x21; + + let original = sector.clone(); + + // "Encrypt" by descrambling plaintext + css::lfsr::descramble_sector(key, &mut sector); + sector[0x14] = 0x30; + + let cracked = css::crack::crack_title_key(§or); + + if let Some(cracked_key) = cracked { + let mut test = sector.clone(); + css::lfsr::descramble_sector(&cracked_key, &mut test); + + assert_eq!(test[0x80], 0x00, "PES byte 0 mismatch"); + assert_eq!(test[0x81], 0x00, "PES byte 1 mismatch"); + assert_eq!(test[0x82], 0x01, "PES byte 2 mismatch"); + assert_eq!(test[0x83], 0xE0, "PES byte 3 mismatch"); + assert_eq!( + &test[0x80..2048], + &original[0x80..2048], + "cracked key did not recover original plaintext" + ); + + any_cracked = true; + eprintln!( + "Stevenson attack succeeded: key={:02X?} seed={:02X?} cracked={:02X?}", + key, seed, cracked_key + ); + } + } + + assert!( + any_cracked, + "Stevenson attack did not crack any of the candidate key/seed pairs" + ); +} + +/// Verify that `recover_title_key` works when given exact known plaintext, +/// even for combinations where `crack_title_key` (which guesses the pattern) +/// might not converge. +#[test] +fn css_recover_title_key_with_exact_plaintext() { + let title_key: [u8; 5] = [0x42, 0x13, 0x37, 0xBE, 0xEF]; + let seed: [u8; 5] = [0x11, 0x22, 0x33, 0x44, 0x55]; + + let mut sector = vec![0x00u8; 2048]; + sector[0x14] = 0x30; + sector[0x54..0x59].copy_from_slice(&seed); + let pes_header: [u8; 10] = [0x00, 0x00, 0x01, 0xE0, 0x00, 0x00, 0x80, 0x80, 0x05, 0x21]; + sector[0x80..0x8A].copy_from_slice(&pes_header); + for i in 0x8A..2048 { + sector[i] = ((i * 13 + 7) & 0xFF) as u8; + } + let original = sector.clone(); + + // Scramble + css::lfsr::descramble_sector(&title_key, &mut sector); + sector[0x14] = 0x30; + + // Recover with exact known plaintext + let recovered = css::crack::recover_title_key(§or, &pes_header); + + if let Some(rkey) = recovered { + let mut test = sector.clone(); + css::lfsr::descramble_sector(&rkey, &mut test); + assert_eq!( + &test[0x80..2048], + &original[0x80..2048], + "recovered key did not produce correct plaintext" + ); + eprintln!("recover_title_key succeeded: {:02X?}", rkey); + } else { + eprintln!( + "recover_title_key returned None for key={:02X?} seed={:02X?}. \ + The LFSR0 recovery phase may not converge for this combination.", + title_key, seed + ); + } +} + /// Test: aacs_parse_unit_key_ro with minimal valid data #[test] fn aacs_parse_unit_key_ro_minimal() { diff --git a/tests/streams.rs b/tests/streams.rs index 9e72250..2bb7138 100644 --- a/tests/streams.rs +++ b/tests/streams.rs @@ -766,3 +766,256 @@ fn mkvstream_meta_preserves_all_streams() { assert!(s.forced); } } + +// ── H2: End-to-end MKV mux test ───────────────────────────── + +#[test] +fn mkvstream_e2e_h264_produces_valid_mkv() { + // Construct a DiscTitle with one H.264 video stream + let dt = DiscTitle { + playlist: "H264 Test".into(), + playlist_id: 0, + duration_secs: 10.0, + size_bytes: 0, + clips: Vec::new(), + streams: vec![Stream::Video(VideoStream { + pid: 0x1011, + codec: Codec::H264, + resolution: "1080p".into(), + frame_rate: "23.976".into(), + hdr: HdrFormat::Sdr, + color_space: ColorSpace::Bt709, + secondary: false, + label: "Main".into(), + })], + chapters: Vec::new(), + extents: Vec::new(), + content_format: ContentFormat::BdTs, + }; + + // Build synthetic BD-TS packets containing valid H.264 NALs + // We need PES headers wrapping: SPS (NAL type 7), PPS (NAL type 8), IDR (NAL type 5) + let mut ts_data = Vec::new(); + + // Build elementary stream data: start codes + NALs + let mut es_data = Vec::new(); + + // SPS NAL (type 7): minimal valid SPS + es_data.extend_from_slice(&[0x00, 0x00, 0x00, 0x01]); // start code + es_data.push(0x67); // NAL type 7 (SPS), nal_ref_idc=3 + // Minimal SPS payload: profile_idc=66 (Baseline), constraint flags, level_idc=30 + es_data.extend_from_slice(&[ + 0x42, 0xC0, 0x1E, // profile=66, constraint_set0=1, level=30 + 0xD9, 0x00, 0xA0, 0x47, 0xFE, 0x88, // minimal SPS rbsp + ]); + + // PPS NAL (type 8): minimal valid PPS + es_data.extend_from_slice(&[0x00, 0x00, 0x00, 0x01]); // start code + es_data.push(0x68); // NAL type 8 (PPS), nal_ref_idc=3 + es_data.extend_from_slice(&[0xCE, 0x38, 0x80]); // minimal PPS rbsp + + // IDR NAL (type 5): keyframe + es_data.extend_from_slice(&[0x00, 0x00, 0x00, 0x01]); // start code + es_data.push(0x65); // NAL type 5 (IDR), nal_ref_idc=3 + // Some IDR slice data + es_data.extend_from_slice(&[0x88, 0x84, 0x00, 0x21, 0xFF, 0xFE, 0xF6, 0xE2]); + // Pad to reasonable size + es_data.extend_from_slice(&[0x00; 64]); + + // Wrap in PES header with PTS + let pts: i64 = 90000; // 1 second in 90kHz ticks + let pts_bytes = encode_pts_test(pts); + let pes_header_len = 9 + 5; // basic PES header (9) + PTS (5) + let pes_length = (3 + 5 + es_data.len()) as u16; // flags(3) + PTS(5) + ES data + + let mut pes = Vec::new(); + pes.extend_from_slice(&[0x00, 0x00, 0x01, 0xE0]); // PES start code + video stream_id + pes.extend_from_slice(&pes_length.to_be_bytes()); // PES packet length + pes.extend_from_slice(&[0x80, 0x80, 0x05]); // flags: PTS present, header_data_len=5 + pes.extend_from_slice(&pts_bytes); + pes.extend_from_slice(&es_data); + + // Wrap PES in 192-byte BD-TS packets + let pid: u16 = 0x1011; + let mut pes_offset = 0; + let mut pusi = true; + let mut cc: u8 = 0; + + while pes_offset < pes.len() { + let mut pkt = [0u8; 192]; + // 4-byte TP_extra_header (zeros) + pkt[4] = 0x47; // sync byte + pkt[5] = (pid >> 8) as u8 & 0x1F; + if pusi { + pkt[5] |= 0x40; // PUSI + pusi = false; + } + pkt[6] = pid as u8; + pkt[7] = 0x10 | (cc & 0x0F); // payload only + continuity counter + cc = cc.wrapping_add(1); + + let space = 184; + let rem = pes.len() - pes_offset; + let n = rem.min(space); + + if n < space { + // Need adaptation field for padding + let pad = space - n; + pkt[7] = 0x30 | (cc.wrapping_sub(1) & 0x0F); // AF + payload + pkt[8] = (pad - 1) as u8; // adaptation_field_length + if pad > 1 { + pkt[9] = 0x00; // flags + } + for byte in pkt.iter_mut().take(8 + pad).skip(10) { + *byte = 0xFF; + } + pkt[8 + pad..8 + pad + n].copy_from_slice(&pes[pes_offset..pes_offset + n]); + } else { + pkt[8..8 + n].copy_from_slice(&pes[pes_offset..pes_offset + n]); + } + + ts_data.extend_from_slice(&pkt); + pes_offset += n; + } + + // Build a second PES (access unit) to trigger the first PES to be output + // by the TS demuxer (it needs a new PUSI to emit the previous PES). + let pts2: i64 = 90000 + 3753; // ~1 frame later + let pts2_bytes = encode_pts_test(pts2); + let mut es_data2 = Vec::new(); + // Just a non-IDR slice (NAL type 1) + es_data2.extend_from_slice(&[0x00, 0x00, 0x00, 0x01]); + es_data2.push(0x41); // NAL type 1 (non-IDR) + es_data2.extend_from_slice(&[0x9A, 0x00, 0x10, 0x20]); + es_data2.extend_from_slice(&[0x00; 32]); + + let pes2_length = (3 + 5 + es_data2.len()) as u16; + let mut pes2 = Vec::new(); + pes2.extend_from_slice(&[0x00, 0x00, 0x01, 0xE0]); + pes2.extend_from_slice(&pes2_length.to_be_bytes()); + pes2.extend_from_slice(&[0x80, 0x80, 0x05]); + pes2.extend_from_slice(&pts2_bytes); + pes2.extend_from_slice(&es_data2); + + // Wrap second PES in BD-TS packets + let mut pes2_offset = 0; + let mut pusi2 = true; + while pes2_offset < pes2.len() { + let mut pkt = [0u8; 192]; + pkt[4] = 0x47; + pkt[5] = (pid >> 8) as u8 & 0x1F; + if pusi2 { + pkt[5] |= 0x40; + pusi2 = false; + } + pkt[6] = pid as u8; + pkt[7] = 0x10 | (cc & 0x0F); + cc = cc.wrapping_add(1); + + let space = 184; + let rem = pes2.len() - pes2_offset; + let n = rem.min(space); + + if n < space { + let pad = space - n; + pkt[7] = 0x30 | (cc.wrapping_sub(1) & 0x0F); + pkt[8] = (pad - 1) as u8; + if pad > 1 { + pkt[9] = 0x00; + } + for byte in pkt.iter_mut().take(8 + pad).skip(10) { + *byte = 0xFF; + } + pkt[8 + pad..8 + pad + n].copy_from_slice(&pes2[pes2_offset..pes2_offset + n]); + } else { + pkt[8..8 + n].copy_from_slice(&pes2[pes2_offset..pes2_offset + n]); + } + + ts_data.extend_from_slice(&pkt); + pes2_offset += n; + } + + // Feed through MkvStream using a shared writer to inspect the output bytes. + let output2 = std::sync::Arc::new(std::sync::Mutex::new(Cursor::new(Vec::new()))); + + struct SharedWriter(std::sync::Arc>>>); + impl Write for SharedWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0.lock().unwrap().write(buf) + } + fn flush(&mut self) -> std::io::Result<()> { + self.0.lock().unwrap().flush() + } + } + impl std::io::Seek for SharedWriter { + fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result { + self.0.lock().unwrap().seek(pos) + } + } + + let writer = SharedWriter(output2.clone()); + let mut stream2 = MkvStream::new(writer).meta(&dt).max_buffer(1024 * 1024); + stream2.write_all(&ts_data).unwrap(); + stream2.finish().unwrap(); + + let data = output2.lock().unwrap().clone().into_inner(); + + // Verify output starts with EBML magic (0x1A45DFA3) + assert!(data.len() >= 4, "MKV output too small: {} bytes", data.len()); + assert_eq!( + &data[0..4], + &[0x1A, 0x45, 0xDF, 0xA3], + "output should start with EBML magic" + ); + + // Verify output contains a Tracks element (0x1654AE6B) + let tracks_needle = [0x16, 0x54, 0xAE, 0x6B]; + let has_tracks = data + .windows(4) + .any(|w| w == tracks_needle); + assert!(has_tracks, "output should contain Tracks element"); + + // Verify codecPrivate is non-empty (not all zeros) + // CodecPrivate element ID is 0x63A2 + let cp_needle = [0x63, 0xA2]; + let cp_pos = data + .windows(2) + .position(|w| w == cp_needle); + if let Some(pos) = cp_pos { + // After the ID, there's a size VINT, then the data + let after_id = pos + 2; + if after_id < data.len() { + // Read VINT size + let size_byte = data[after_id]; + let (cp_size, cp_data_start) = if size_byte & 0x80 != 0 { + ((size_byte & 0x7F) as usize, after_id + 1) + } else if size_byte & 0x40 != 0 && after_id + 1 < data.len() { + ( + (((size_byte & 0x3F) as usize) << 8) | data[after_id + 1] as usize, + after_id + 2, + ) + } else { + (0, after_id + 1) + }; + if cp_size > 0 && cp_data_start + cp_size <= data.len() { + let cp_data = &data[cp_data_start..cp_data_start + cp_size]; + let all_zeros = cp_data.iter().all(|&b| b == 0); + assert!( + !all_zeros, + "codecPrivate should not be all zeros (SPS/PPS should be filled)" + ); + } + } + } +} + +fn encode_pts_test(pts: i64) -> [u8; 5] { + let p = pts as u64; + [ + 0x21 | ((p >> 29) & 0x0E) as u8, + ((p >> 22) & 0xFF) as u8, + 0x01 | ((p >> 14) & 0xFE) as u8, + ((p >> 7) & 0xFF) as u8, + 0x01 | ((p << 1) & 0xFE) as u8, + ] +}