From 01334d31d02741b53f4e42e92c035d2f2df1025c Mon Sep 17 00:00:00 2001 From: Jason Ross Date: Wed, 24 Jun 2026 19:57:52 -0500 Subject: [PATCH] Fix panic-on-malformed-input defects found in code review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Harden the decryptors against corrupt/edge-case input so a single bad file errors cleanly instead of panicking (and, in a batch, aborting the whole run): * ADEPT EPUB: the content-decrypt guard was `len < 16` but then sliced off the 16-byte pad block before reading the trailing pad length; a one-block (16-byte) ciphertext left an empty slice and panicked on `.last().unwrap()`. Guard is now `len <= 16`. * MOBI: validate the section table before indexing it — require >= 2 sections and reject a record count that exceeds the section count, preventing out-of-bounds panics on `section_offsets[1]` / `section_bounds(i)`. * MOBI: `trailing_size` now uses saturating/checked subtraction, and the caller clamps the trailing size to the record length, so a corrupt trailing-size encoding can't underflow. * MOBI: `normalize_pids` slices PIDs by characters, not bytes, so a non-ASCII `--pid` can't panic on a non-char-boundary. * CLI: wrap the per-file decrypt in `catch_unwind` so any future panic is contained to that file rather than aborting a multi-file run. Adds 5 regression tests (one per fix). 46 lib tests pass; clippy + fmt clean. Co-Authored-By: Claude Opus 4.8 --- crates/drmlibre-cli/src/remove.rs | 13 ++++- crates/drmlibre-core/src/adept/epub.rs | 13 ++++- crates/drmlibre-core/src/kindle/mobi.rs | 78 +++++++++++++++++++++++-- 3 files changed, 97 insertions(+), 7 deletions(-) diff --git a/crates/drmlibre-cli/src/remove.rs b/crates/drmlibre-cli/src/remove.rs index 781445d..4ee5c20 100644 --- a/crates/drmlibre-cli/src/remove.rs +++ b/crates/drmlibre-cli/src/remove.rs @@ -124,7 +124,18 @@ fn process_one( let tmp = temp_sibling(output); let _ = std::fs::remove_file(&tmp); - match decrypt(input, &tmp, keys, opts) { + // Contain any panic in the engine so one malformed file can't abort a whole + // multi-file run; turn it into a per-file error instead. + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + decrypt(input, &tmp, keys, opts) + })) + .unwrap_or_else(|_| { + Err(Error::Decrypt( + "internal error while processing this file (skipped)".into(), + )) + }); + + match result { Ok(Outcome::Decrypted) => { std::fs::rename(&tmp, output)?; Ok(()) diff --git a/crates/drmlibre-core/src/adept/epub.rs b/crates/drmlibre-core/src/adept/epub.rs index a35cafd..493aa8b 100644 --- a/crates/drmlibre-core/src/adept/epub.rs +++ b/crates/drmlibre-core/src/adept/epub.rs @@ -146,7 +146,9 @@ pub fn decrypt_epub(input: &Path, output: &Path, key: UserKey, opts: &Options) - /// raw-inflate. Mirrors `ineptepub.py::Decryptor.decrypt`. fn decrypt_content(bookkey: &[u8], data: &[u8], decompress: bool) -> Result> { let pt = crypto::aes128_cbc_decrypt_nopad(bookkey, &[0u8; 16], data)?; - if pt.len() < 16 { + // The first 16 bytes are a throwaway pad block, so there must be at least + // one more block after them before we can read the trailing pad length. + if pt.len() <= 16 { return Err(Error::Decrypt("ciphertext too short".into())); } let pt = &pt[16..]; @@ -333,6 +335,15 @@ mod tests { assert!(decrypt_epub(&input, &output, UserKey::Rsa(&other_der), &opts).is_err()); } + #[test] + fn decrypt_content_rejects_single_block_without_panic() { + // A one-block ciphertext decrypts to 16 bytes; after dropping the 16-byte + // pad block nothing is left, so this must error rather than panic. + let bookkey = [0u8; 16]; + let ct = [0xABu8; 16]; + assert!(decrypt_content(&bookkey, &ct, false).is_err()); + } + #[test] fn drm_free_epub_reports_already_free() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/drmlibre-core/src/kindle/mobi.rs b/crates/drmlibre-core/src/kindle/mobi.rs index b234252..a13d943 100644 --- a/crates/drmlibre-core/src/kindle/mobi.rs +++ b/crates/drmlibre-core/src/kindle/mobi.rs @@ -186,6 +186,20 @@ impl MobiBook { ))); } + // Guard the section table against a corrupt header before we index it: + // we need at least sections 0 and 1, and every text record (1..=records) + // must have a corresponding section. + if self.section_offsets.len() < 2 { + return Err(Error::Decrypt( + "MOBI has too few sections to decrypt".into(), + )); + } + if self.records >= self.section_offsets.len() { + return Err(Error::Decrypt( + "MOBI record count exceeds section count (corrupt)".into(), + )); + } + // Normalize the supplied PIDs to 8-character form (mobidedrm goodpids). let goodpids = normalize_pids(pidlist); @@ -235,7 +249,7 @@ impl MobiBook { for i in 1..=self.records { let (start, end) = self.section_bounds(i); let sec = &self.data[start..end]; - let extra = trailing_size(sec, self.extra_data_flags); + let extra = trailing_size(sec, self.extra_data_flags).min(sec.len()); let body = &sec[..sec.len() - extra]; out.extend_from_slice(&pc1(&found_key, body, true)?); if extra > 0 { @@ -270,8 +284,9 @@ impl MobiBook { fn normalize_pids(pidlist: &[String]) -> Vec { let mut good = Vec::new(); for pid in pidlist { - match pid.len() { - 10 => good.push(pid[..8].to_string()), + // Count/slice by characters, not bytes: a non-ASCII --pid must not panic. + match pid.chars().count() { + 10 => good.push(pid.chars().take(8).collect()), // drop the 2-char checksum 8 => good.push(pid.clone()), _ => tracing::debug!("ignoring PID {pid} with wrong length"), } @@ -342,12 +357,16 @@ fn trailing_size(ptr: &[u8], flags: u16) -> usize { let mut testflags = flags >> 1; while testflags != 0 { if testflags & 1 != 0 { - num += entry(ptr, ptr.len() - num); + // `num` may already exceed the record on corrupt input; saturate so + // the size param can't underflow. + num += entry(ptr, ptr.len().saturating_sub(num)); } testflags >>= 1; } if flags & 1 != 0 { - num += (ptr[ptr.len() - num - 1] & 0x3) as usize + 1; + if let Some(&b) = ptr.len().checked_sub(num + 1).and_then(|i| ptr.get(i)) { + num += (b & 0x3) as usize + 1; + } } num } @@ -483,4 +502,53 @@ mod tests { let out = book.process(&[]).unwrap(); assert_eq!(out, data); } + + /// A single-section BOOKMOBI (num_sections = 1) claiming encryption. + fn build_one_section_mobi(crypto_type: u16) -> Vec { + let mut sec0 = vec![0u8; 0x100]; + sec0[0..2].copy_from_slice(&1u16.to_be_bytes()); // compression + sec0[0xC..0xE].copy_from_slice(&crypto_type.to_be_bytes()); + sec0[0x14..0x18].copy_from_slice(&0xC8u32.to_be_bytes()); // mobi_length + sec0[0x68..0x6C].copy_from_slice(&6u32.to_be_bytes()); // mobi_version + + let num_sections = 1u16; + let header_len = 78 + num_sections as usize * 8; // 86 + let mut data = vec![0u8; header_len]; + data[0x3C..0x44].copy_from_slice(b"BOOKMOBI"); + data[76..78].copy_from_slice(&num_sections.to_be_bytes()); + data[78..82].copy_from_slice(&(header_len as u32).to_be_bytes()); + data.extend_from_slice(&sec0); + data + } + + #[test] + fn single_section_mobi_errors_not_panics() { + let data = build_one_section_mobi(2); + let book = MobiBook::from_bytes(data).unwrap(); + assert!(book.process(&[]).is_err()); + } + + #[test] + fn corrupt_record_count_errors_not_panics() { + let bookkey = [0xC3u8; 16]; + let mut data = build_encrypted_mobi(&bookkey, "12345678", b"hi"); + // Patch the record count (sec0 + 8; sec0 starts at 94) to exceed sections. + data[94 + 8..94 + 10].copy_from_slice(&0xFFFFu16.to_be_bytes()); + let book = MobiBook::from_bytes(data).unwrap(); + assert!(book.process(&["12345678".to_string()]).is_err()); + } + + #[test] + fn trailing_size_does_not_underflow_on_corrupt_input() { + // Encoded sizes far exceed the record length; must return, not panic. + let _ = trailing_size(&[0x7F, 0x7F], 0b110); // two size entries + let _ = trailing_size(&[0x7F], 0b11); // size entry + multibyte flag + } + + #[test] + fn normalize_pids_handles_non_ascii() { + // A 10-byte PID whose 8th byte is mid-character must not panic. + let pids = vec!["XXXXXXXéX".to_string()]; + assert!(normalize_pids(&pids).is_empty()); + } }