use crate::aacs;
use crate::css;
use rayon::prelude::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, RwLock};
const PARALLEL_MIN_UNITS: usize = 8;
pub const MAX_THREADS: usize = 64;
static DECRYPT_THREADS: AtomicUsize = AtomicUsize::new(0);
static DECRYPT_POOL: RwLock<Option<Arc<rayon::ThreadPool>>> = RwLock::new(None);
pub fn set_decrypt_threads(n: usize) {
let clamped = n.min(MAX_THREADS);
DECRYPT_THREADS.store(clamped, Ordering::Relaxed);
if let Ok(mut guard) = DECRYPT_POOL.write() {
*guard = None;
}
}
fn decrypt_pool() -> Option<Arc<rayon::ThreadPool>> {
{
let guard = DECRYPT_POOL.read().unwrap_or_else(|e| e.into_inner());
if let Some(pool) = guard.as_ref() {
return Some(Arc::clone(pool));
}
}
let mut guard = DECRYPT_POOL.write().unwrap_or_else(|e| e.into_inner());
if let Some(pool) = guard.as_ref() {
return Some(Arc::clone(pool));
}
let n = decrypt_threads();
let pool = rayon::ThreadPoolBuilder::new()
.num_threads(n)
.thread_name(|i| format!("freemkv-decrypt-{i}"))
.build()
.ok()
.map(Arc::new)?;
*guard = Some(Arc::clone(&pool));
Some(pool)
}
pub fn decrypt_threads() -> usize {
let explicit = DECRYPT_THREADS.load(Ordering::Relaxed);
if explicit > 0 {
return explicit;
}
let env = std::env::var("FREEMKV_THREADS")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.unwrap_or(0);
if env > 0 {
return env.min(MAX_THREADS);
}
let cores = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(2);
cores.clamp(1, MAX_THREADS)
}
#[derive(Clone)]
pub enum DecryptKeys {
None,
Aacs {
unit_keys: Vec<(u32, [u8; 16])>,
read_data_key: Option<[u8; 16]>,
},
Css { title_key: [u8; 5] },
}
impl DecryptKeys {
pub fn is_encrypted(&self) -> bool {
!matches!(self, DecryptKeys::None)
}
}
pub fn decrypt_sectors(
buf: &mut [u8],
keys: &DecryptKeys,
unit_key_idx: usize,
) -> Result<(), crate::error::Error> {
match keys {
DecryptKeys::None => {}
DecryptKeys::Aacs {
unit_keys,
read_data_key,
} => {
let uk = match unit_keys.get(unit_key_idx) {
Some((_, k)) => *k,
None => {
return Err(crate::error::Error::DecryptFailed);
}
};
let rdk: Option<[u8; 16]> = *read_data_key;
let unit_len = aacs::ALIGNED_UNIT_LEN;
let partial_len = buf.len() % unit_len;
if partial_len != 0 {
let partial = &buf[buf.len() - partial_len..];
let packets = aacs::ts_packet_total(partial);
if packets > 0 && aacs::ts_sync_count(partial) <= packets / 2 {
return Err(crate::error::Error::DecryptFailed);
}
}
let nthreads = decrypt_threads();
let nunits = buf.len() / unit_len;
let decrypt_one = |chunk: &mut [u8]| {
if chunk.len() == unit_len && aacs::is_aacs_scrambled(chunk) {
let original: Vec<u8> = chunk.to_vec();
if !aacs::decrypt_unit_full(chunk, &uk, rdk.as_ref()) {
chunk.copy_from_slice(&original);
}
}
};
if nthreads <= 1 || nunits < PARALLEL_MIN_UNITS {
for chunk in buf.chunks_mut(unit_len) {
decrypt_one(chunk);
}
} else {
match decrypt_pool() {
Some(pool) => {
let chunks: Vec<&mut [u8]> = buf.chunks_mut(unit_len).collect();
pool.install(|| {
chunks.into_par_iter().for_each(|chunk| {
decrypt_one(chunk);
});
});
}
None => {
for chunk in buf.chunks_mut(unit_len) {
decrypt_one(chunk);
}
}
}
}
}
DecryptKeys::Css { title_key } => {
for chunk in buf.chunks_mut(2048) {
css::lfsr::descramble_sector(title_key, chunk);
}
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn nav_file_unit_survives_decrypt_attempt() {
let mut unit = vec![0u8; aacs::ALIGNED_UNIT_LEN];
unit[0] = b'M';
unit[1] = b'P';
unit[2] = b'L';
unit[3] = b'S';
for (i, b) in unit.iter_mut().enumerate().skip(4) {
*b = (i as u8).wrapping_mul(31);
}
let snapshot = unit.clone();
let keys = DecryptKeys::Aacs {
unit_keys: vec![(0, [0xAB; 16])],
read_data_key: None,
};
decrypt_sectors(&mut unit, &keys, 0).unwrap();
assert_eq!(
unit, snapshot,
"non-m2ts unit must be restored after failed decrypt"
);
}
fn clear_ts_region(len: usize) -> Vec<u8> {
let mut v: Vec<u8> = (0..len).map(|i| (i as u8).wrapping_mul(31)).collect();
let mut off = 4;
while off < len {
v[off] = 0x47;
off += 192;
}
v
}
fn scrambled_region(len: usize) -> Vec<u8> {
let mut v: Vec<u8> = (0..len).map(|i| (i as u8).wrapping_mul(31)).collect();
let mut off = 4;
while off < len {
v[off] = 0xA5;
off += 192;
}
v
}
#[test]
fn aacs_clear_trailing_partial_is_tolerated_unchanged() {
let keys = DecryptKeys::Aacs {
unit_keys: vec![(0, [0xAB; 16])],
read_data_key: None,
};
let unit = scrambled_region(aacs::ALIGNED_UNIT_LEN);
let tail = clear_ts_region(2048);
let mut buf = unit;
buf.extend_from_slice(&tail);
decrypt_sectors(&mut buf, &keys, 0).expect("clear trailing partial is Ok");
assert_eq!(
&buf[aacs::ALIGNED_UNIT_LEN..],
&tail[..],
"clear trailing partial unit must be left unchanged"
);
}
#[test]
fn aacs_scrambled_trailing_partial_is_rejected() {
let keys = DecryptKeys::Aacs {
unit_keys: vec![(0, [0xAB; 16])],
read_data_key: None,
};
let unit = clear_ts_region(aacs::ALIGNED_UNIT_LEN);
let tail = scrambled_region(4096);
let mut buf = unit;
buf.extend_from_slice(&tail);
let err = decrypt_sectors(&mut buf, &keys, 0)
.expect_err("scrambled trailing partial must be rejected");
assert_eq!(
err.code(),
crate::error::Error::DecryptFailed.code(),
"scrambled trailing partial must fail with DecryptFailed"
);
}
#[test]
fn aacs_empty_buffer_is_ok() {
let keys = DecryptKeys::Aacs {
unit_keys: vec![(0, [0xAB; 16])],
read_data_key: None,
};
let mut buf: Vec<u8> = Vec::new();
assert!(decrypt_sectors(&mut buf, &keys, 0).is_ok());
}
#[test]
fn aacs_exact_multiple_unchanged() {
let keys = DecryptKeys::Aacs {
unit_keys: vec![(0, [0xAB; 16])],
read_data_key: None,
};
let mut buf = clear_ts_region(aacs::ALIGNED_UNIT_LEN * 2);
let snapshot = buf.clone();
decrypt_sectors(&mut buf, &keys, 0).expect("exact-multiple buffer is Ok");
assert_eq!(
buf, snapshot,
"clear exact-multiple buffer must be left unchanged"
);
}
#[test]
fn none_keys_is_noop() {
let mut buf: Vec<u8> = (0..4096u32).map(|i| (i % 256) as u8).collect();
let snapshot = buf.clone();
decrypt_sectors(&mut buf, &DecryptKeys::None, 0).expect("None is always Ok");
assert_eq!(buf, snapshot, "None must not touch the buffer");
}
#[test]
fn is_encrypted_matches_variant() {
assert!(!DecryptKeys::None.is_encrypted());
assert!(DecryptKeys::Css { title_key: [0; 5] }.is_encrypted());
assert!(
DecryptKeys::Aacs {
unit_keys: vec![(0, [0; 16])],
read_data_key: None,
}
.is_encrypted()
);
}
fn make_css_sector(title_key: &[u8; 5], seed: &[u8; 5], body_fill: u8) -> (Vec<u8>, Vec<u8>) {
let mut sector = vec![body_fill; 2048];
sector[0x14] = 0x30; sector[0x54..0x59].copy_from_slice(seed);
let plaintext = sector.clone();
css::lfsr::descramble_sector(title_key, &mut sector);
sector[0x14] = 0x30;
(sector, plaintext)
}
#[test]
fn css_descrambles_with_title_key() {
let title_key = [0x42, 0x13, 0x37, 0xBE, 0xEF];
let seed = [0xDE, 0xAD, 0xBE, 0xEF, 0x42];
let (mut sector, plaintext) = make_css_sector(&title_key, &seed, 0xA5);
let keys = DecryptKeys::Css { title_key };
decrypt_sectors(&mut sector, &keys, 0).expect("CSS decrypt is Ok");
assert_eq!(
§or[0x80..2048],
&plaintext[0x80..2048],
"CSS body must round-trip to plaintext"
);
assert_eq!(
sector[0x14] & 0x30,
0,
"scramble flag cleared after CSS decrypt"
);
}
#[test]
fn css_processes_every_sector_in_buffer() {
let title_key = [0x01, 0x02, 0x03, 0x04, 0x05];
let (s0, p0) = make_css_sector(&title_key, &[0x11, 0x22, 0x33, 0x44, 0x55], 0x3C);
let (s1, p1) = make_css_sector(&title_key, &[0x66, 0x77, 0x88, 0x99, 0xAA], 0xC3);
let mut buf = s0;
buf.extend_from_slice(&s1);
let keys = DecryptKeys::Css { title_key };
decrypt_sectors(&mut buf, &keys, 0).expect("CSS multi-sector decrypt is Ok");
assert_eq!(
&buf[0x80..2048],
&p0[0x80..2048],
"sector 0 body must round-trip"
);
assert_eq!(
&buf[2048 + 0x80..4096],
&p1[0x80..2048],
"sector 1 body must round-trip (loop must reach the 2nd sector)"
);
}
#[test]
fn css_leaves_clear_sector_unchanged() {
let title_key = [0x01, 0x02, 0x03, 0x04, 0x05];
let mut sector = vec![0x77u8; 2048];
sector[0x14] = 0x00; let snapshot = sector.clone();
let keys = DecryptKeys::Css { title_key };
decrypt_sectors(&mut sector, &keys, 0).unwrap();
assert_eq!(sector, snapshot, "clear CSS sector must be left untouched");
}
#[test]
fn css_empty_buffer_is_ok() {
let mut buf: Vec<u8> = Vec::new();
let keys = DecryptKeys::Css { title_key: [0; 5] };
assert!(decrypt_sectors(&mut buf, &keys, 0).is_ok());
}
#[test]
fn aacs_out_of_range_unit_key_idx_errors() {
let keys = DecryptKeys::Aacs {
unit_keys: vec![(0, [0xAB; 16])],
read_data_key: None,
};
let mut buf = clear_ts_region(aacs::ALIGNED_UNIT_LEN);
let err = decrypt_sectors(&mut buf, &keys, 5)
.expect_err("unit_key_idx 5 is out of range for a 1-key list");
assert_eq!(
err.code(),
crate::error::Error::DecryptFailed.code(),
"out-of-range unit key index must be DecryptFailed"
);
}
#[test]
fn aacs_empty_unit_keys_errors() {
let keys = DecryptKeys::Aacs {
unit_keys: vec![],
read_data_key: None,
};
let mut buf = clear_ts_region(aacs::ALIGNED_UNIT_LEN);
let err = decrypt_sectors(&mut buf, &keys, 0).expect_err("empty unit_keys must error");
assert_eq!(err.code(), crate::error::Error::DecryptFailed.code());
}
#[test]
fn decrypt_threads_within_valid_pool_range() {
let n = decrypt_threads();
assert!(n >= 1, "decrypt thread count must be at least 1, got {n}");
assert!(
n <= MAX_THREADS,
"decrypt thread count must not exceed MAX_THREADS ({MAX_THREADS}), got {n}"
);
}
}