fix(android): bound archive extraction and decoder resources

senbei-engine/src/android/extract/pipeline.rs

@@ -3,7 +3,9 @@ use std::fs::File;
 use std::path::{Path, PathBuf};
 
 use memmap2::MmapOptions;
-use senbei_crypto::android::{Module9bConfig, decode_container};
+use senbei_crypto::android::{Module9bConfig, decode_container_with_limit};
+
+use crate::android::{ByteBudget, EntryBudget, ResourceLimits};
 use serde_json::to_vec_pretty;
 
 use super::super::common;
@@ -61,6 +63,13 @@ struct ArtifactSpec<'a> {
     classification: &'a str,
 }
 
+struct DecodeBudget {
+    limits: ResourceLimits,
+    bytes: ByteBudget,
+    records: EntryBudget,
+    config: Module9bConfig,
+}
+
 struct Extractor {
     output_dir: PathBuf,
     streams: Vec<StreamReport>,
@@ -70,6 +79,20 @@ struct Extractor {
 }
 
 pub fn extract_stage2(options: &ExtractOptions) -> Result<ExtractionReport> {
+    extract_stage2_with_limits(options, &ResourceLimits::default())
+}
+
+/// [`extract_stage2`] with explicit ceilings for decoded records and nesting.
+///
+/// `limits.max_entries` counts decoded stage-2 records. `limits.max_depth`
+/// counts nested interpreter streams, with the bootstrap stream at depth 0.
+/// `limits.max_entry_bytes` bounds each decoded container and its writers, and
+/// `limits.max_total_bytes` sums those decoded payloads for this call. Neither
+/// budget is a peak resident-memory measurement.
+pub fn extract_stage2_with_limits(
+    options: &ExtractOptions,
+    limits: &ResourceLimits,
+) -> Result<ExtractionReport> {
     let input_path = absolute(&options.input)?;
     let output_dir = absolute(&options.output_dir)?;
     if !input_path.is_file() {
@@ -98,6 +121,10 @@ pub fn extract_stage2(options: &ExtractOptions) -> Result<ExtractionReport> {
         &input_path,
         options.outer_size,
         options.cipher_constant,
+        // Stage 1 plaintext is the first thing this task retains. Its copy
+        // is capped by both the single-entry ceiling and whatever remains of
+        // the task total, which is the whole total at this point.
+        Some(limits.max_entry_bytes.min(limits.max_total_bytes)),
     )?;
     if let Some(stage2_output) = &options.stage2_output {
         write_atomic(&absolute(stage2_output)?, &stage1.plaintext)?;
@@ -119,13 +146,22 @@ pub fn extract_stage2(options: &ExtractOptions) -> Result<ExtractionReport> {
         registry: BTreeMap::new(),
         seen_streams: HashSet::new(),
     };
+    let mut budget = DecodeBudget {
+        limits: *limits,
+        bytes: ByteBudget::new(limits),
+        records: EntryBudget::new(limits),
+        config: core_config,
+    };
+    budget
+        .bytes
+        .charge(u64::try_from(stage1.plaintext.len()).unwrap_or(u64::MAX))?;
     extractor.extract_stream(
         bootstrap,
         0xe2,
         0,
         None,
         Some(stage1.remaining_file_offset),
-        core_config,
+        &mut budget,
     )?;
 
     let module_registry = extractor
@@ -170,8 +206,9 @@ impl Extractor {
         depth: usize,
         parent: Option<StreamParent>,
         source_file_offset: Option<usize>,
-        config: Module9bConfig,
+        budget: &mut DecodeBudget,
     ) -> Result<()> {
+        budget.limits.check_depth(depth)?;
         let digest = sha256(stream);
         if !self.seen_streams.insert((stream_id, digest.clone())) {
             return Ok(());
@@ -200,10 +237,15 @@ impl Extractor {
                     "decoded_interpreter"
                 },
                 (depth != 0).then_some(stream_id),
-                &config,
+                &budget.config,
             ),
-            records: Vec::with_capacity(records.len()),
+            records: Vec::new(),
         };
+        budget.records.charge(records.len())?;
+        stream_report
+            .records
+            .try_reserve_exact(records.len())
+            .map_err(|_| Error::Invalid("stage-2 record report reservation failed".to_owned()))?;
         let mut direct_records = Vec::new();
         let mut modules_at_level = BTreeSet::new();
 
@@ -213,16 +255,24 @@ impl Extractor {
             let mut metadata_data = None;
 
             if !record.direct() && record.image_size != 0 {
+                // Charge the declared size before the decoder reserves it.
+                // A metadata-only record is charged below, on the same rule.
+                budget.bytes.charge(u64::from(record.image_size))?;
                 let image_source = record_tail(stream, record.image_offset)?;
-                let image = decode_container(image_source, &config, record.image_size as usize)
-                    .map_err(|source| Error::RecordDecode {
-                        depth,
-                        stream_id,
-                        record_index: record.index,
-                        command_id: record.command_id,
-                        part: "image decode",
-                        source,
-                    })?;
+                let image = decode_container_with_limit(
+                    image_source,
+                    &budget.config,
+                    record.image_size as usize,
+                    budget.limits.output_limit(),
+                )
+                .map_err(|source| Error::RecordDecode {
+                    depth,
+                    stream_id,
+                    record_index: record.index,
+                    command_id: record.command_id,
+                    part: "image decode",
+                    source,
+                })?;
                 let classification = if record.metadata_size != 0 {
                     "module_image"
                 } else {
@@ -243,17 +293,22 @@ impl Extractor {
                 image_data = Some((image, artifact));
             }
             if record.metadata_size != 0 {
+                budget.bytes.charge(u64::from(record.metadata_size))?;
                 let metadata_source = record_tail(stream, record.metadata_offset)?;
-                let metadata =
-                    decode_container(metadata_source, &config, record.metadata_size as usize)
-                        .map_err(|source| Error::RecordDecode {
-                            depth,
-                            stream_id,
-                            record_index: record.index,
-                            command_id: record.command_id,
-                            part: "metadata decode",
-                            source,
-                        })?;
+                let metadata = decode_container_with_limit(
+                    metadata_source,
+                    &budget.config,
+                    record.metadata_size as usize,
+                    budget.limits.output_limit(),
+                )
+                .map_err(|source| Error::RecordDecode {
+                    depth,
+                    stream_id,
+                    record_index: record.index,
+                    command_id: record.command_id,
+                    part: "metadata decode",
+                    source,
+                })?;
                 let artifact = self.write_artifact(
                     &record,
                     depth,
@@ -339,14 +394,17 @@ impl Extractor {
                         source,
                     }
                 })?;
+            budget.limits.check_depth(depth + 1)?;
+            let parent_config = std::mem::replace(&mut budget.config, interpreter_config);
             self.extract_stream(
                 child_data,
                 next_stream_id,
                 depth + 1,
                 Some(parent),
                 None,
-                interpreter_config,
+                budget,
             )?;
+            budget.config = parent_config;
         }
         Ok(())
     }
@@ -506,3 +564,197 @@ fn absolute(path: &Path) -> Result<PathBuf> {
 fn sha256(data: &[u8]) -> String {
     common::sha256(data)
 }
+
+#[cfg(test)]
+mod limit_tests {
+    use super::*;
+    use crate::android::LimitExceeded;
+
+    fn config() -> Module9bConfig {
+        Module9bConfig {
+            header_seed: 0,
+            container_seed: 0,
+            aes_key: [0; 32],
+            skip_aes: true,
+            schedule_offset: 0,
+        }
+    }
+
+    fn tight(entries: usize, depth: usize) -> ResourceLimits {
+        ResourceLimits {
+            max_entry_bytes: 1024,
+            max_total_bytes: 1024,
+            max_entries: entries,
+            max_depth: depth,
+        }
+    }
+
+    /// One direct record whose payload is `child`.
+    ///
+    /// The command is `stream_id + 0x10`, which is the id a nested stream uses
+    /// for its interpreter. The payload begins immediately after the single
+    /// descriptor, so parsing stops at that boundary.
+    fn direct_stream(stream_id: u32, child: &[u8]) -> Vec<u8> {
+        let key = stream_id.wrapping_mul(0x9d32_3cd7);
+        let shift = stream_id & 7;
+        let base = (key >> shift)
+            .wrapping_add(0x5e72_7d74)
+            .wrapping_add(key.wrapping_shl(stream_id & 0xb))
+            .wrapping_add(0xf71e_3005);
+        // record_state = plain1 + base, and plain1 is chosen so that state is
+        // zero. That keeps the record cipher's index mask zero as well.
+        let plain1 = 0_u32.wrapping_sub(base);
+        let cipher0 =
+            unmultiply(1_u32 ^ 0xebe8_1dba_u32.wrapping_add(base)).wrapping_sub(0xcbf0_c1d8);
+        let cipher1 = unmultiply(plain1 ^ 0xebe8_1dba_u32.wrapping_mul(5).wrapping_add(base))
+            .wrapping_sub(cipher0);
+        let mut stream = Vec::new();
+        stream.extend_from_slice(&cipher0.to_le_bytes());
+        stream.extend_from_slice(&cipher1.to_le_bytes());
+        // `command - 0x10` is the interpreter id a nested stream looks up.
+        // Using `stream_id + 0x10` makes a lone record direct, because that
+        // interpreter was not decoded from this same stream.
+        let command = stream_id.wrapping_add(0x10);
+        let words = [
+            command,
+            super::super::stream::DIRECT_FLAG,
+            (8 + super::super::stream::RECORD_SIZE) as u32,
+            child.len() as u32,
+            0,
+            0,
+            command,
+            0,
+            0,
+        ];
+        stream.extend_from_slice(&encrypt_record(&words));
+        stream.extend_from_slice(child);
+        stream
+    }
+
+    fn encrypt_record(words: &[u32]) -> Vec<u8> {
+        // Inverse of `decrypt_record` at record state zero: replay its
+        // feedback from each ciphertext word, then undo the field multiply.
+        let mut feedback = 0xf02f_7685_u32;
+        let mut accumulator = 0x7993_4cf6_u32;
+        let mut raw = vec![0_u8; super::super::stream::RECORD_SIZE];
+        for (word_index, plain) in words.iter().copied().enumerate() {
+            feedback = feedback.wrapping_mul(feedback);
+            let mixed = plain
+                .wrapping_sub(accumulator)
+                .wrapping_add(0 >> ((word_index * 4 + 3) & 5));
+            let cipher = unmultiply(mixed) ^ (feedback >> 3);
+            raw[word_index * 4..word_index * 4 + 4].copy_from_slice(&cipher.to_le_bytes());
+            accumulator = accumulator.wrapping_add(0xe64d_33d8);
+            feedback = cipher;
+        }
+        raw
+    }
+
+    fn unmultiply(encoded: u32) -> u32 {
+        let mut multiplier = 0xdc49_cd06_u32;
+        let mut value = encoded;
+        let mut result = 0_u32;
+        while multiplier != 0 {
+            if multiplier & 1 != 0 {
+                result ^= value;
+            }
+            let carry = value >> 31;
+            value = value.wrapping_shl(1);
+            if carry != 0 {
+                value ^= 0x5793_57eb;
+            }
+            multiplier >>= 1;
+        }
+        result
+    }
+
+    fn extractor() -> Extractor {
+        Extractor {
+            output_dir: PathBuf::from("."),
+            streams: Vec::new(),
+            artifacts: Vec::new(),
+            registry: BTreeMap::new(),
+            seen_streams: HashSet::new(),
+        }
+    }
+
+    fn budget(entries: usize, depth: usize) -> DecodeBudget {
+        let limits = tight(entries, depth);
+        DecodeBudget {
+            limits,
+            bytes: ByteBudget::new(&limits),
+            records: EntryBudget::new(&limits),
+            config: config(),
+        }
+    }
+
+    #[test]
+    fn record_budget_rejects_one_past_the_ceiling_and_allows_the_boundary() {
+        let stream = direct_stream(0xe2, &[0]);
+        let error = extractor()
+            .extract_stream(&stream, 0xe2, 0, None, None, &mut budget(0, 32))
+            .expect_err("one record is already past a zero-record ceiling");
+        assert!(
+            matches!(
+                error,
+                Error::Limit(LimitExceeded::Entries {
+                    requested: 1,
+                    limit: 0
+                })
+            ),
+            "unexpected error: {error}"
+        );
+
+        let directory = tempfile::tempdir().expect("temporary output directory");
+        let mut at_ceiling = extractor();
+        at_ceiling.output_dir = directory.path().to_path_buf();
+        at_ceiling
+            .extract_stream(&stream, 0xe2, 0, None, None, &mut budget(1, 0))
+            .expect("one record at depth zero is the boundary");
+    }
+
+    #[test]
+    fn nesting_one_past_the_ceiling_is_rejected_before_the_child_is_decoded() {
+        let stream = direct_stream(0xe2, &[0]);
+        let directory = tempfile::tempdir().expect("temporary output directory");
+        let mut at_ceiling = extractor();
+        at_ceiling.output_dir = directory.path().to_path_buf();
+        at_ceiling
+            .extract_stream(&stream, 0xe2, 32, None, None, &mut budget(1, 32))
+            .expect("depth 32 is the nesting boundary");
+
+        let error = extractor()
+            .extract_stream(&stream, 0xe2, 33, None, None, &mut budget(1, 32))
+            .expect_err("depth 33 is rejected before the stream is parsed");
+        match &error {
+            Error::Limit(limit) => assert_eq!(
+                limit,
+                &LimitExceeded::Depth {
+                    requested: 33,
+                    limit: 32
+                }
+            ),
+            other => panic!("unexpected error: {other:?}"),
+        }
+    }
+
+    #[test]
+    fn decoded_payloads_share_one_task_budget() {
+        let stream = direct_stream(0xe2, &[0]);
+        let directory = tempfile::tempdir().expect("temporary output directory");
+        let mut budget = budget(1, 0);
+        budget.limits.max_total_bytes = 1;
+        budget.bytes = ByteBudget::new(&budget.limits);
+        // A direct record has no decoded container, so it does not charge the
+        // byte budget. Seed the only allowed byte to prove the budget is the
+        // one the caller passed in and is not reset per stream.
+        budget.bytes.charge(1).unwrap();
+        let mut extractor = extractor();
+        extractor.output_dir = directory.path().to_path_buf();
+        extractor
+            .extract_stream(&stream, 0xe2, 0, None, None, &mut budget)
+            .unwrap();
+        assert_eq!(budget.bytes.used(), 1);
+        assert!(budget.bytes.charge(1024).is_err());
+    }
+}