Make encrypt_unit report a refused slice, and expand its key once

Two defects in the encrypt_unit promoted to public API last round, both found by
round 2 auditing that new code.

It returned silently without encrypting when the slice was shorter than
ALIGNED_UNIT_LEN. Its own contract requires the caller to set the container's
encrypted flag BEFORE calling — the header is the key seed — so a silent no-op
leaves a unit advertised as encrypted while still carrying plaintext, with
nothing for an authoring caller to check. It now returns bool and is
#[must_use], so ignoring the refusal is a compile-time warning; every call site
was updated to assert on it.

bool rather than Result deliberately: a wrong buffer length is a programming
error at a library boundary, not a disc condition, and a new Error variant would
mean a new numeric code plus its rendering in another repo.

It also drove CBC from the single-block aes_ecb_encrypt, rebuilding the AES key
schedule for each of the 383 blocks in a unit — an order of magnitude slower
than its inverse, which expands the key once via aes_cbc_decrypt. The missing
counterpart aes_cbc_encrypt now exists alongside it, and encrypt_unit calls it,
so the two directions are symmetric in structure as well as in result. For an
authoring caller encrypting a 90 GB image that removes ~5.6 billion redundant
key expansions.

New test pins the boundary: ALIGNED_UNIT_LEN - 1 returns false and leaves the
buffer byte-identical, ALIGNED_UNIT_LEN succeeds. The existing round-trip and
padding-asymmetry tests still pass, so the CBC rewrite is provably the same
transform.

src/aacs/content.rs

@@ -6,7 +6,11 @@ use aes::Aes128;
 #[cfg(test)]
 use aes::cipher::{KeyInit, generic_array::GenericArray};
 
-use super::crypto::{AACS_IV, aes_cbc_decrypt, aes_ecb_encrypt};
+use super::crypto::{aes_cbc_decrypt, aes_cbc_encrypt, aes_ecb_encrypt};
+// Only this module's test fixtures build CBC ciphertext by hand now — the
+// production paths get the IV from `aes_cbc_encrypt` / `aes_cbc_decrypt`.
+#[cfg(test)]
+use super::crypto::AACS_IV;
 
 // ── AACS constants ──────────────────────────────────────────────────────────
 
@@ -347,9 +351,17 @@ pub fn decrypt_unit(unit: &mut [u8], unit_key: &[u8; 16]) {
 /// packet encrypts to ciphertext that is not all-zero, so it is not mistaken for
 /// padding on the way back and the round trip is still exact. Authoring that wants
 /// true source-zero padding leaves those packets unencrypted instead.
-pub fn encrypt_unit(unit: &mut [u8], unit_key: &[u8; 16]) {
+///
+/// Returns `false` — encrypting nothing — when `unit` is shorter than
+/// [`ALIGNED_UNIT_LEN`]. That case MUST be checked: the caller has already set the
+/// container's encrypted flag by then (this function's contract requires it), so
+/// ignoring the result leaves a unit marked encrypted while still carrying
+/// plaintext, which is the worst possible outcome for an authoring tool.
+#[must_use = "returns false when the slice is too short to encrypt, leaving \
+              plaintext behind a flag that already says 'encrypted'"]
+pub fn encrypt_unit(unit: &mut [u8], unit_key: &[u8; 16]) -> bool {
     if unit.len() < ALIGNED_UNIT_LEN {
-        return;
+        return false;
     }
     let mut header = [0u8; 16];
     header.copy_from_slice(&unit[..16]);
@@ -358,19 +370,11 @@ pub fn encrypt_unit(unit: &mut [u8], unit_key: &[u8; 16]) {
     for i in 0..16 {
         k[i] = derived[i] ^ header[i];
     }
-    // CBC-encrypt bytes 16.. under the fixed AACS IV (forward of `aes_cbc_decrypt`).
+    // CBC-encrypt bytes 16.. under the fixed AACS IV — the exact forward of the
-    let mut prev = AACS_IV;
+    // `aes_cbc_decrypt` call in `decrypt_unit`, and one key expansion for the whole
-    let num_blocks = (ALIGNED_UNIT_LEN - 16) / 16;
+    // unit rather than one per 16-byte block.
-    for i in 0..num_blocks {
+    aes_cbc_encrypt(&k, &mut unit[16..ALIGNED_UNIT_LEN]);
-        let off = 16 + i * 16;
+    true
-        let mut block = [0u8; 16];
-        for j in 0..16 {
-            block[j] = unit[off + j] ^ prev[j];
-        }
-        let enc = aes_ecb_encrypt(&k, &block);
-        unit[off..off + 16].copy_from_slice(&enc);
-        prev.copy_from_slice(&enc);
-    }
 }
 
 /// Remove bus encryption from an aligned unit (AACS 2.0 / UHD).
@@ -412,7 +416,10 @@ mod tests {
         clear[0] |= 0xC0;
 
         let mut unit = clear.clone();
-        encrypt_unit(&mut unit, &key);
+        assert!(
+            encrypt_unit(&mut unit, &key),
+            "a full-length unit must encrypt"
+        );
         assert_ne!(
             unit[16..],
             clear[16..],
@@ -433,6 +440,31 @@ mod tests {
     /// to non-zero bytes, so it is not mistaken for padding and still round-trips.
     /// This is the one place the two functions are deliberately not symmetric, so
     /// the claim is worth pinning rather than asserting in prose alone.
+    /// A slice too short to encrypt must SAY so. The caller has already set the
+    /// container's encrypted flag by the time it calls this (the contract requires
+    /// flag-before-crypto, since the header is the key seed), so a silent no-op
+    /// leaves a unit advertised as encrypted while still carrying plaintext.
+    #[test]
+    fn encrypt_unit_reports_a_slice_too_short_to_encrypt() {
+        let key = [0x11u8; 16];
+        let mut short = vec![0u8; ALIGNED_UNIT_LEN - 1];
+        short[0] |= 0xC0; // the caller already flagged it encrypted
+        let before = short.clone();
+        assert!(
+            !encrypt_unit(&mut short, &key),
+            "a short slice must report false, not silently succeed"
+        );
+        assert_eq!(
+            short, before,
+            "a refused encrypt must leave the buffer untouched"
+        );
+
+        // Exactly ALIGNED_UNIT_LEN is the boundary and must succeed.
+        let mut exact = vec![0u8; ALIGNED_UNIT_LEN];
+        exact[0] |= 0xC0;
+        assert!(encrypt_unit(&mut exact, &key), "a full unit must encrypt");
+    }
+
     #[test]
     fn encrypt_unit_round_trips_all_zero_plaintext_packets() {
         let key = [0xA5u8; 16];
@@ -444,7 +476,10 @@ mod tests {
         }
 
         let mut unit = clear.clone();
-        encrypt_unit(&mut unit, &key);
+        assert!(
+            encrypt_unit(&mut unit, &key),
+            "a full-length unit must encrypt"
+        );
         // No later packet may encipher to all-zero, or decrypt would treat it as
         // source padding and the asymmetry would bite.
         for p in 1..ALIGNED_UNIT_LEN / BD_SOURCE_PACKET_BYTES {
@@ -732,7 +767,10 @@ mod tests {
         // Delegate to the module-scope `pub(crate)` helper (the single encrypt
         // implementation, shared with the mux `driver.rs` decrypt test).
         unit[0] |= 0xC0;
-        super::encrypt_unit(unit, unit_key);
+        assert!(
+            super::encrypt_unit(unit, unit_key),
+            "a full-length unit must encrypt"
+        );
     }
 
     /// Build a clear aligned unit with TS sync bytes at offset 4 + k*192.

src/aacs/crypto.rs

@@ -40,6 +40,35 @@ pub(crate) fn aes_ecb_decrypt(key: &[u8; 16], data: &[u8; 16]) -> [u8; 16] {
 /// Precondition: `data.len()` is a multiple of 16. Any trailing partial
 /// block is silently ignored; all callers pass aligned regions (6128 and
 /// 2032 bytes), and the assert documents/enforces that contract.
+/// AES-128-CBC encrypt in place under the fixed [`AACS_IV`] — the forward
+/// direction of [`aes_cbc_decrypt`], and its exact inverse.
+///
+/// Constructs the cipher ONCE for the whole slice. Driving this from the
+/// single-block [`aes_ecb_encrypt`] instead rebuilds the AES key schedule per
+/// 16-byte block, which for a 6144-byte aligned unit is 383 redundant key
+/// expansions.
+pub(crate) fn aes_cbc_encrypt(key: &[u8; 16], data: &mut [u8]) {
+    debug_assert!(
+        data.len() % 16 == 0,
+        "aes_cbc_encrypt requires a block-aligned slice"
+    );
+    let cipher = Aes128::new(GenericArray::from_slice(key));
+    let num_blocks = data.len() / 16;
+    let mut prev = AACS_IV;
+    // Forward order: each block is XORed with the PRECEDING ciphertext block.
+    for i in 0..num_blocks {
+        let offset = i * 16;
+        let mut block = [0u8; 16];
+        for j in 0..16 {
+            block[j] = data[offset + j] ^ prev[j];
+        }
+        let mut ga = GenericArray::clone_from_slice(&block);
+        cipher.encrypt_block(&mut ga);
+        data[offset..offset + 16].copy_from_slice(&ga);
+        prev.copy_from_slice(&ga);
+    }
+}
+
 pub(crate) fn aes_cbc_decrypt(key: &[u8; 16], data: &mut [u8]) {
     debug_assert!(
         data.len() % 16 == 0,

src/decrypt.rs

@@ -1065,7 +1065,10 @@ mod tests {
     /// seed, so the flag must be set before the crypto runs).
     fn aacs_encrypt_unit_for_test(unit: &mut [u8], unit_key: &[u8; 16]) {
         unit[0] |= 0xC0;
-        aacs::content::encrypt_unit(unit, unit_key);
+        assert!(
+            aacs::content::encrypt_unit(unit, unit_key),
+            "a full-length unit must encrypt"
+        );
     }
 
     /// Build a clear aligned unit with TS sync bytes placed at the BD-TS stride

src/disc/extract.rs

@@ -1057,7 +1057,10 @@ mod tests {
     /// `tag` distinguishes two units' payloads.
     fn encrypt_aacs_unit(unit_key: &[u8; 16], tag: u8) -> Vec<u8> {
         let mut unit = clear_aacs_unit(tag);
-        crate::aacs::content::encrypt_unit(&mut unit, unit_key);
+        assert!(
+            crate::aacs::content::encrypt_unit(&mut unit, unit_key),
+            "a full-length unit must encrypt"
+        );
         unit
     }
 

src/mux/driver.rs

@@ -1555,7 +1555,10 @@ mod tests {
         unit[..192].copy_from_slice(&pkt);
         // Flag encrypted BEFORE encrypting: bytes 0..16 are the key seed.
         unit[0] |= 0xC0;
-        crate::aacs::content::encrypt_unit(&mut unit, unit_key);
+        assert!(
+            crate::aacs::content::encrypt_unit(&mut unit, unit_key),
+            "a full-length unit must encrypt"
+        );
         unit
     }
 

src/mux/resolve.rs

@@ -2642,7 +2642,10 @@ mod tests {
         }
         // Flag encrypted BEFORE encrypting: bytes 0..16 are the key seed.
         u[0] |= 0xC0;
-        crate::aacs::content::encrypt_unit(&mut u, key);
+        assert!(
+            crate::aacs::content::encrypt_unit(&mut u, key),
+            "a full-length unit must encrypt"
+        );
         u
     }
 

src/sector/decrypting.rs

@@ -759,7 +759,10 @@ mod tests {
     /// encrypt it under `unit_key` so `aacs::content::decrypt_unit` recovers it.
     fn encrypt_aacs_unit(unit_key: &[u8; 16]) -> Vec<u8> {
         let mut unit = clear_aacs_unit();
-        crate::aacs::content::encrypt_unit(&mut unit, unit_key);
+        assert!(
+            crate::aacs::content::encrypt_unit(&mut unit, unit_key),
+            "a full-length unit must encrypt"
+        );
         unit
     }