mirror of
https://github.com/Momoko-Ayase/Senbei.git
synced 2026-11-04 10:31:29 -05:00
761 lines
27 KiB
Rust
761 lines
27 KiB
Rust
use std::collections::{BTreeMap, BTreeSet, HashSet};
|
|
use std::fs::File;
|
|
use std::path::{Path, PathBuf};
|
|
|
|
use memmap2::MmapOptions;
|
|
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;
|
|
use super::error::{Error, Result, invalid};
|
|
use super::report::{
|
|
ArtifactReport, DecoderReport, ExtractionReport, ModuleRegistryEntry, RecordReport,
|
|
Stage1Report, StreamParent, StreamReport,
|
|
};
|
|
use super::stage1::{
|
|
DEFAULT_CIPHER_CONSTANT, DEFAULT_OUTER_SIZE, SHT_LOUSER, Stage1Result, inspect,
|
|
};
|
|
use super::stream::{DIRECT_FLAG, Record, parse_record_stream};
|
|
|
|
/// Inputs and output locations for one complete static Stage 2 extraction.
|
|
#[derive(Debug, Clone)]
|
|
pub struct ExtractOptions {
|
|
pub input: PathBuf,
|
|
pub output_dir: PathBuf,
|
|
pub stage2_output: Option<PathBuf>,
|
|
pub outer_size: usize,
|
|
pub cipher_constant: u32,
|
|
}
|
|
|
|
impl ExtractOptions {
|
|
#[must_use]
|
|
pub fn with_defaults(input: PathBuf, output_dir: PathBuf) -> Self {
|
|
Self {
|
|
input,
|
|
output_dir,
|
|
stage2_output: None,
|
|
outer_size: DEFAULT_OUTER_SIZE,
|
|
cipher_constant: DEFAULT_CIPHER_CONSTANT,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct LoadedModule {
|
|
image: Vec<u8>,
|
|
metadata: Option<Vec<u8>>,
|
|
image_path: String,
|
|
metadata_path: Option<String>,
|
|
sha256: String,
|
|
depth: usize,
|
|
record_index: usize,
|
|
command_id: u32,
|
|
init_offset: u32,
|
|
entry_offset: u32,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy)]
|
|
struct ArtifactSpec<'a> {
|
|
suffix: &'a str,
|
|
kind: &'a str,
|
|
classification: &'a str,
|
|
}
|
|
|
|
struct DecodeBudget {
|
|
limits: ResourceLimits,
|
|
bytes: ByteBudget,
|
|
records: EntryBudget,
|
|
config: Module9bConfig,
|
|
}
|
|
|
|
struct Extractor {
|
|
output_dir: PathBuf,
|
|
streams: Vec<StreamReport>,
|
|
artifacts: Vec<ArtifactReport>,
|
|
registry: BTreeMap<u32, LoadedModule>,
|
|
seen_streams: HashSet<(u32, String)>,
|
|
}
|
|
|
|
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() {
|
|
return invalid(format!(
|
|
"protected ELF does not exist: {}",
|
|
input_path.display()
|
|
));
|
|
}
|
|
if let Some(stage2_output) = &options.stage2_output {
|
|
let stage2_output = absolute(stage2_output)?;
|
|
if stage2_output == input_path {
|
|
return invalid("refusing to overwrite the protected ELF with Stage 2 output");
|
|
}
|
|
}
|
|
std::fs::create_dir_all(&output_dir)
|
|
.map_err(|source| Error::io("create Stage 2 output directory", &output_dir, source))?;
|
|
|
|
let file = File::open(&input_path)
|
|
.map_err(|source| Error::io("open protected ELF", &input_path, source))?;
|
|
// SAFETY: the mapping is read-only, the file remains open for the mapping
|
|
// lifetime, and extraction never mutates or truncates the source.
|
|
let source = unsafe { MmapOptions::new().map(&file) }
|
|
.map_err(|source| Error::io("map protected ELF", &input_path, source))?;
|
|
let stage1 = inspect(
|
|
&source,
|
|
&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)?;
|
|
}
|
|
|
|
let core_config =
|
|
Module9bConfig::parse_embedded(&stage1.plaintext).map_err(Error::EmbeddedConfig)?;
|
|
let bootstrap_end = stage1
|
|
.remaining_file_offset
|
|
.checked_add(stage1.remaining_size)
|
|
.ok_or_else(|| Error::Invalid("Stage 2 bootstrap range overflow".to_owned()))?;
|
|
let bootstrap = source
|
|
.get(stage1.remaining_file_offset..bootstrap_end)
|
|
.ok_or_else(|| Error::Invalid("Stage 2 bootstrap range is outside the ELF".to_owned()))?;
|
|
let mut extractor = Extractor {
|
|
output_dir: output_dir.clone(),
|
|
streams: Vec::new(),
|
|
artifacts: Vec::new(),
|
|
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),
|
|
&mut budget,
|
|
)?;
|
|
|
|
let module_registry = extractor
|
|
.registry
|
|
.values()
|
|
.map(|module| ModuleRegistryEntry {
|
|
command_id: module.command_id,
|
|
size: module.image.len(),
|
|
sha256: module.sha256.clone(),
|
|
depth: module.depth,
|
|
record_index: module.record_index,
|
|
image_path: module.image_path.clone(),
|
|
metadata_path: module.metadata_path.clone(),
|
|
init_offset: module.init_offset,
|
|
entry_offset: module.entry_offset,
|
|
classification: if module.metadata.is_some() {
|
|
"module_image".to_owned()
|
|
} else {
|
|
"decoded_data".to_owned()
|
|
},
|
|
})
|
|
.collect::<Vec<_>>();
|
|
let report = ExtractionReport {
|
|
format_version: 4,
|
|
protected_elf: input_path.display().to_string(),
|
|
output_dir: output_dir.display().to_string(),
|
|
stage1: stage1_report(&stage1, options.outer_size),
|
|
streams: extractor.streams,
|
|
artifacts: extractor.artifacts,
|
|
errors: Vec::new(),
|
|
module_registry,
|
|
};
|
|
write_json_atomic(&output_dir.join("index.json"), &report)?;
|
|
Ok(report)
|
|
}
|
|
|
|
impl Extractor {
|
|
fn extract_stream(
|
|
&mut self,
|
|
stream: &[u8],
|
|
stream_id: u32,
|
|
depth: usize,
|
|
parent: Option<StreamParent>,
|
|
source_file_offset: Option<usize>,
|
|
budget: &mut DecodeBudget,
|
|
) -> Result<()> {
|
|
budget.limits.check_depth(depth)?;
|
|
let digest = sha256(stream);
|
|
if !self.seen_streams.insert((stream_id, digest.clone())) {
|
|
return Ok(());
|
|
}
|
|
let (header, records, table_size) =
|
|
parse_record_stream(stream, stream_id).map_err(|source| {
|
|
Error::Invalid(format!(
|
|
"depth {depth} stream 0x{stream_id:02X} record table: {source}"
|
|
))
|
|
})?;
|
|
let mut stream_report = StreamReport {
|
|
depth,
|
|
stream_id,
|
|
parent,
|
|
source_file_offset,
|
|
available_size: stream.len(),
|
|
descriptor_table_size: table_size,
|
|
encrypted_header_words: header.encrypted_words,
|
|
decrypted_header_words: header.decrypted_words,
|
|
record_state: header.record_state,
|
|
sha256: digest,
|
|
decoder: decoder_report(
|
|
if depth == 0 {
|
|
"embedded_stage2"
|
|
} else {
|
|
"decoded_interpreter"
|
|
},
|
|
(depth != 0).then_some(stream_id),
|
|
&budget.config,
|
|
),
|
|
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();
|
|
|
|
for record in records {
|
|
let mut result = record_report(record);
|
|
let mut image_data = None;
|
|
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_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 {
|
|
"decoded_data"
|
|
};
|
|
let artifact = self.write_artifact(
|
|
&record,
|
|
depth,
|
|
stream_id,
|
|
ArtifactSpec {
|
|
suffix: "module.bin",
|
|
kind: "decoded_container",
|
|
classification,
|
|
},
|
|
&image,
|
|
)?;
|
|
result.image = Some(artifact.clone());
|
|
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_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,
|
|
stream_id,
|
|
ArtifactSpec {
|
|
suffix: "metadata.bin",
|
|
kind: "decoded_metadata",
|
|
classification: "decoded_metadata",
|
|
},
|
|
&metadata,
|
|
)?;
|
|
result.metadata = Some(artifact.clone());
|
|
metadata_data = Some((metadata, artifact));
|
|
}
|
|
if let Some((image, image_artifact)) = image_data {
|
|
let (metadata, metadata_path) = if let Some((data, artifact)) = metadata_data {
|
|
(Some(data), Some(artifact.path))
|
|
} else {
|
|
(None, None)
|
|
};
|
|
self.register_module(LoadedModule {
|
|
sha256: image_artifact.sha256.clone(),
|
|
image_path: image_artifact.path.clone(),
|
|
metadata_path,
|
|
image,
|
|
metadata,
|
|
depth,
|
|
record_index: record.index,
|
|
command_id: record.command_id,
|
|
init_offset: record.init_offset,
|
|
entry_offset: record.entry_offset,
|
|
})?;
|
|
modules_at_level.insert(record.command_id);
|
|
}
|
|
if record.direct() && record.image_size != 0 {
|
|
direct_records.push((record, stream_report.records.len()));
|
|
}
|
|
stream_report.records.push(result);
|
|
}
|
|
|
|
let mut children = Vec::new();
|
|
for (record, report_index) in direct_records {
|
|
let next_stream_id = record.command_id.wrapping_sub(0x10);
|
|
if modules_at_level.contains(&next_stream_id) {
|
|
stream_report.records[report_index].nested_stream_id = Some(next_stream_id);
|
|
children.push((record, next_stream_id));
|
|
continue;
|
|
}
|
|
let direct_data = record_slice(stream, record.image_offset, record.image_size)?;
|
|
let artifact = self.write_artifact(
|
|
&record,
|
|
depth,
|
|
stream_id,
|
|
ArtifactSpec {
|
|
suffix: "direct.bin",
|
|
kind: "direct",
|
|
classification: "direct_data",
|
|
},
|
|
direct_data,
|
|
)?;
|
|
stream_report.records[report_index].image = Some(artifact);
|
|
}
|
|
|
|
self.streams.push(stream_report);
|
|
for (record, next_stream_id) in children {
|
|
let child_data = record_slice(stream, record.image_offset, record.image_size)?;
|
|
let parent = StreamParent {
|
|
stream_id,
|
|
record_index: record.index,
|
|
command_id: record.command_id,
|
|
};
|
|
let interpreter = self.registry.get(&next_stream_id).ok_or_else(|| {
|
|
Error::Invalid(format!(
|
|
"depth {depth} stream 0x{stream_id:02X} child 0x{next_stream_id:02X} has no interpreter module"
|
|
))
|
|
})?;
|
|
let interpreter_config =
|
|
Module9bConfig::parse(&interpreter.image).map_err(|source| {
|
|
Error::InterpreterConfig {
|
|
depth: depth + 1,
|
|
stream_id: next_stream_id,
|
|
interpreter_id: next_stream_id,
|
|
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,
|
|
budget,
|
|
)?;
|
|
budget.config = parent_config;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn register_module(&mut self, module: LoadedModule) -> Result<()> {
|
|
if let Some(previous) = self.registry.get(&module.command_id) {
|
|
if previous.sha256 != module.sha256 {
|
|
return invalid(format!(
|
|
"module 0x{:02X} produced conflicting images: {} and {}",
|
|
module.command_id, previous.sha256, module.sha256
|
|
));
|
|
}
|
|
return Ok(());
|
|
}
|
|
self.registry.insert(module.command_id, module);
|
|
Ok(())
|
|
}
|
|
|
|
fn write_artifact(
|
|
&mut self,
|
|
record: &Record,
|
|
depth: usize,
|
|
stream_id: u32,
|
|
spec: ArtifactSpec<'_>,
|
|
data: &[u8],
|
|
) -> Result<ArtifactReport> {
|
|
let digest = sha256(data);
|
|
let filename = format!(
|
|
"d{depth:02}_s{stream_id:02X}_r{:03}_id{:08X}_{}.{}",
|
|
record.index,
|
|
record.command_id,
|
|
&digest[..12],
|
|
spec.suffix
|
|
);
|
|
let path = self.output_dir.join(filename);
|
|
write_atomic(&path, data)?;
|
|
let artifact = ArtifactReport {
|
|
kind: spec.kind.to_owned(),
|
|
path: path
|
|
.file_name()
|
|
.ok_or_else(|| Error::Invalid("artifact path has no file name".to_owned()))?
|
|
.to_string_lossy()
|
|
.into_owned(),
|
|
size: data.len(),
|
|
sha256: digest,
|
|
depth,
|
|
stream_id,
|
|
record_index: Some(record.index),
|
|
command_id: Some(record.command_id),
|
|
classification: spec.classification.to_owned(),
|
|
};
|
|
self.artifacts.push(artifact.clone());
|
|
Ok(artifact)
|
|
}
|
|
}
|
|
|
|
fn record_report(record: Record) -> RecordReport {
|
|
RecordReport {
|
|
index: record.index,
|
|
command_id: record.command_id,
|
|
flags: record.flags,
|
|
image_offset: record.image_offset,
|
|
image_size: record.image_size,
|
|
metadata_offset: record.metadata_offset,
|
|
metadata_size: record.metadata_size,
|
|
id_copy: record.id_copy,
|
|
entry_offset: record.entry_offset,
|
|
init_offset: record.init_offset,
|
|
direct: record.flags & DIRECT_FLAG != 0,
|
|
extraction_status: "complete".to_owned(),
|
|
image: None,
|
|
metadata: None,
|
|
nested_stream_id: None,
|
|
}
|
|
}
|
|
|
|
fn decoder_report(
|
|
kind: &str,
|
|
interpreter_id: Option<u32>,
|
|
config: &Module9bConfig,
|
|
) -> DecoderReport {
|
|
DecoderReport {
|
|
kind: kind.to_owned(),
|
|
interpreter_id,
|
|
header_seed: config.header_seed,
|
|
container_seed: config.container_seed,
|
|
schedule_offset: config.schedule_offset,
|
|
aes_key_sha256: sha256(&config.aes_key),
|
|
skip_aes: config.skip_aes,
|
|
}
|
|
}
|
|
|
|
fn record_slice(stream: &[u8], offset: u32, size: u32) -> Result<&[u8]> {
|
|
let offset = usize::try_from(offset)
|
|
.map_err(|_| Error::Invalid("record payload offset exceeds usize".to_owned()))?;
|
|
let size = usize::try_from(size)
|
|
.map_err(|_| Error::Invalid("record payload size exceeds usize".to_owned()))?;
|
|
let end = offset
|
|
.checked_add(size)
|
|
.ok_or_else(|| Error::Invalid("record payload range overflows usize".to_owned()))?;
|
|
stream.get(offset..end).ok_or_else(|| {
|
|
Error::Invalid(format!(
|
|
"record payload range 0x{offset:x}..0x{end:x} exceeds stream 0x{:x}",
|
|
stream.len()
|
|
))
|
|
})
|
|
}
|
|
|
|
fn record_tail(stream: &[u8], offset: u32) -> Result<&[u8]> {
|
|
let offset = usize::try_from(offset)
|
|
.map_err(|_| Error::Invalid("record container offset exceeds usize".to_owned()))?;
|
|
stream.get(offset..).ok_or_else(|| {
|
|
Error::Invalid(format!(
|
|
"record container offset 0x{offset:x} exceeds stream 0x{:x}",
|
|
stream.len()
|
|
))
|
|
})
|
|
}
|
|
|
|
fn stage1_report(stage1: &Stage1Result, outer_size: usize) -> Stage1Report {
|
|
Stage1Report {
|
|
section_index: stage1.section_index,
|
|
section_type: SHT_LOUSER,
|
|
section_offset: stage1.section_offset,
|
|
section_size: stage1.section_size,
|
|
outer_size,
|
|
header_offset: stage1.header_offset,
|
|
header_key: stage1.header.key,
|
|
payload_offset: stage1.header.payload_offset,
|
|
payload_size: stage1.header.payload_size,
|
|
payload_key: stage1.header.payload_key,
|
|
entry_offset: stage1.header.entry_offset,
|
|
protect_size: stage1.header.protect_size,
|
|
stage2_file_offset: stage1.payload_file_offset,
|
|
stage2_size: stage1.plaintext.len(),
|
|
stage2_sha256: sha256(&stage1.plaintext),
|
|
remaining_file_offset: stage1.remaining_file_offset,
|
|
remaining_size: stage1.remaining_size,
|
|
}
|
|
}
|
|
|
|
fn write_json_atomic(path: &Path, value: &impl serde::Serialize) -> Result<()> {
|
|
let mut bytes = to_vec_pretty(value)?;
|
|
bytes.push(b'\n');
|
|
write_atomic(path, &bytes)
|
|
}
|
|
|
|
fn write_atomic(path: &Path, data: &[u8]) -> Result<()> {
|
|
common::write_atomic(path, data)
|
|
.map_err(|source| Error::io("write temporary output", path, source))
|
|
}
|
|
|
|
fn absolute(path: &Path) -> Result<PathBuf> {
|
|
common::absolute(path).map_err(|source| Error::io("query current directory", path, source))
|
|
}
|
|
|
|
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());
|
|
}
|
|
}
|