use std::collections::{HashMap, HashSet};
use crate::SOURCE_FORMAT_PDF;
#[cfg(feature = "rans")]
use crate::container::observation::{
ObservationIndex, ObservationSelector, OpEntry, SECTION_OP_TABLE, SECTION_PDF_SELECTORS,
SELECTOR_OBJECT, SELECTOR_REVISION, SELECTOR_STREAM,
};
use crate::container::{Descriptor, ObjectSource, UNIVERSE};
use crate::dra::op::{DEFLATE_SOURCE_CHANNEL, DEFLATE_SOURCE_OBJECT};
use crate::dra::{Op, Program};
use crate::encode::candidates::{Candidate, CandidateKind};
use crate::error::Result;
use crate::integrity::sha256;
use crate::limits::Limits;
use crate::codec::deflate::{ReplayPlan, try_replay, try_replay_detailed};
use super::cos::FilterClass;
use super::physical::{PdfPhysical, scan};
pub(super) type SpanPlans = HashMap<(u64, u64), ReplayPlan>;
pub(super) fn collect_plans(
input: &[u8],
limits: Limits,
) -> Result<Option<(PdfPhysical, SpanPlans)>> {
let physical = match scan(input, limits) {
Ok(p) => p,
Err(_) => return Ok(None),
};
if physical.header.is_none() || physical.objects.is_empty() || physical.eofs.is_empty() {
return Ok(None);
}
let mut by_span: SpanPlans = HashMap::new();
for stream in &physical.streams {
if stream.filter != FilterClass::FlateDecode {
continue;
}
if by_span.len() as u64 >= u64::from(limits.max_object_count) {
break;
}
let start = stream.data_start as usize;
let end = match start.checked_add(stream.data_len as usize) {
Some(e) if e <= input.len() => e,
_ => continue,
};
if let Some(plan) = try_replay(&input[start..end], limits) {
by_span.insert((stream.data_start, stream.data_len), plan);
}
}
if by_span.is_empty() {
return Ok(None);
}
Ok(Some((physical, by_span)))
}
pub fn propose_pdf_deflate_replay(input: &[u8], limits: Limits) -> Result<Option<Candidate>> {
let Some((physical, by_span)) = collect_plans(input, limits)? else {
return Ok(None);
};
if physical.spans.len() as u64 > u64::from(limits.max_graph_ops) {
return Ok(None);
}
let mut objects: Vec<Vec<u8>> = Vec::new();
let mut index: HashMap<Vec<u8>, u32> = HashMap::new();
let mut ops: Vec<Op> = Vec::with_capacity(physical.spans.len());
let mut replayed = 0usize;
for span in &physical.spans {
let key = (span.start, span.len);
if let Some(plan) = by_span.get(&key) {
let plaintext_id = intern(&mut objects, &mut index, plan.plaintext.clone());
let corrections_id = intern(&mut objects, &mut index, plan.corrections.clone());
if objects.len() as u64 > u64::from(limits.max_object_count) {
return Ok(None);
}
ops.push(Op::Inline {
bytes: plan.header.to_vec(),
});
ops.push(Op::DeflateReplay {
replay_codec: crate::dra::op::REPLAY_DEFLATE_PREFLATE_0_7_6,
source_kind: DEFLATE_SOURCE_OBJECT,
source_id: plaintext_id,
corrections_object: corrections_id,
declared_output_len: plan.raw_len,
});
ops.push(Op::Inline {
bytes: plan.adler.to_vec(),
});
replayed += 1;
} else {
let start = span.start as usize;
let end = start + span.len as usize;
ops.push(Op::Inline {
bytes: input[start..end].to_vec(),
});
}
if ops.len() as u64 > u64::from(limits.max_graph_ops) {
return Ok(None);
}
}
let format_basis = format!(
"pdf-deflate-replay;streams={};replayed={replayed};objects={}",
physical.streams.len(),
objects.len()
);
let descriptor = Descriptor {
universe: UNIVERSE.to_string(),
source_format: SOURCE_FORMAT_PDF,
format_basis,
models: vec![],
channels: vec![],
objects: objects.into_iter().map(ObjectSource::Inline).collect(),
program: Program::new(ops),
observation_index: None,
seek_directory: false,
source_sha256: sha256(input),
source_len: input.len() as u64,
};
Ok(Some(Candidate {
kind: CandidateKind::PdfDeflateReplay,
descriptor,
}))
}
#[cfg(feature = "rans")]
pub fn propose_pdf_deflate_replay_rans(input: &[u8], limits: Limits) -> Result<Option<Candidate>> {
Ok(
build_pdf_deflate_replay_rans(input, limits)?.map(|(_, descriptor)| Candidate {
kind: CandidateKind::PdfDeflateReplayRans,
descriptor,
}),
)
}
#[cfg(feature = "rans")]
fn build_pdf_deflate_replay_rans(
input: &[u8],
limits: Limits,
) -> Result<Option<(PdfPhysical, Descriptor)>> {
use crate::entropy::{
CODER_ORDER0_BYTE_RANS, CODER_VERSION_1, EntropyChannelDescriptor, EntropyModel,
encode_channel,
};
let Some((physical, by_span)) = collect_plans(input, limits)? else {
return Ok(None);
};
if physical.spans.len() as u64 > u64::from(limits.max_graph_ops) {
return Ok(None);
}
let mut models: Vec<EntropyModel> = Vec::new();
let mut channels: Vec<EntropyChannelDescriptor> = Vec::new();
let mut channel_id: HashMap<Vec<u8>, u32> = HashMap::new();
let mut objects: Vec<Vec<u8>> = Vec::new();
let mut object_id: HashMap<Vec<u8>, u32> = HashMap::new();
let mut ops: Vec<Op> = Vec::with_capacity(physical.spans.len());
let mut replayed = 0usize;
for span in &physical.spans {
let key = (span.start, span.len);
if let Some(plan) = by_span.get(&key) {
let cid = if let Some(&id) = channel_id.get(&plan.plaintext) {
id
} else {
if channels.len() as u64 >= u64::from(limits.max_channel_count)
|| models.len() as u64 >= u64::from(limits.max_model_count)
{
return Ok(None);
}
let mut counts = [0u64; crate::entropy::ALPHABET];
for &b in &plan.plaintext {
counts[b as usize] += 1;
}
let model = EntropyModel::from_counts(&counts, 12)?;
let scale_bits = model.scale_bits;
let capsule = encode_channel(&model, &plan.plaintext)?;
let id = channels.len() as u32;
models.push(model);
channels.push(EntropyChannelDescriptor {
coder: CODER_ORDER0_BYTE_RANS,
coder_version: CODER_VERSION_1,
scale_bits,
lane_count: 1,
model_id: id,
symbol_count: capsule.symbol_count,
decoded_length: capsule.decoded_length,
initial_state: capsule.initial_state,
payload: capsule.payload,
});
channel_id.insert(plan.plaintext.clone(), id);
id
};
let corr = intern(&mut objects, &mut object_id, plan.corrections.clone());
if objects.len() as u64 > u64::from(limits.max_object_count) {
return Ok(None);
}
ops.push(Op::Inline {
bytes: plan.header.to_vec(),
});
ops.push(Op::DeflateReplay {
replay_codec: crate::dra::op::REPLAY_DEFLATE_PREFLATE_0_7_6,
source_kind: DEFLATE_SOURCE_CHANNEL,
source_id: cid,
corrections_object: corr,
declared_output_len: plan.raw_len,
});
ops.push(Op::Inline {
bytes: plan.adler.to_vec(),
});
replayed += 1;
} else {
let start = span.start as usize;
let end = start + span.len as usize;
ops.push(Op::Inline {
bytes: input[start..end].to_vec(),
});
}
if ops.len() as u64 > u64::from(limits.max_graph_ops) {
return Ok(None);
}
}
let format_basis = format!(
"pdf-deflate-replay-rans;streams={};replayed={replayed};channels={};objects={}",
physical.streams.len(),
channels.len(),
objects.len()
);
let descriptor = Descriptor {
universe: UNIVERSE.to_string(),
source_format: SOURCE_FORMAT_PDF,
format_basis,
models,
channels,
objects: objects.into_iter().map(ObjectSource::Inline).collect(),
program: Program::new(ops),
observation_index: None,
seek_directory: false,
source_sha256: sha256(input),
source_len: input.len() as u64,
};
Ok(Some((physical, descriptor)))
}
#[cfg(all(feature = "deflate-replay", feature = "rans"))]
pub fn propose_pdf_deflate_replay_rans_indexed(
input: &[u8],
limits: Limits,
) -> Result<Option<Candidate>> {
let Some((physical, mut descriptor)) = build_pdf_deflate_replay_rans(input, limits)? else {
return Ok(None);
};
let Some(index) = build_observation_index(&physical, &descriptor, limits)? else {
return Ok(None);
};
descriptor.observation_index = Some(index);
descriptor.seek_directory = true;
Ok(Some(Candidate {
kind: CandidateKind::PdfDeflateReplayRansIndexed,
descriptor,
}))
}
#[cfg(feature = "rans")]
fn build_observation_index(
physical: &PdfPhysical,
descriptor: &Descriptor,
limits: Limits,
) -> Result<Option<ObservationIndex>> {
let object_lens: Vec<u64> = descriptor.objects.iter().map(|o| o.len()).collect();
let channel_lens: Vec<u64> = descriptor
.channels
.iter()
.map(|c| c.decoded_length)
.collect();
let per_op = match descriptor
.program
.analyze_ops(&object_lens, &channel_lens, limits)
{
Ok(v) => v,
Err(_) => return Ok(None),
};
if per_op.len() != descriptor.program.ops.len() {
return Ok(None);
}
let mut ops: Vec<OpEntry> = Vec::with_capacity(per_op.len());
for (i, len) in per_op.iter().enumerate() {
let Ok(out_len) = u32::try_from(*len) else {
return Ok(None);
};
let (dep_kind, dep_id) = primary_dependency(&descriptor.program.ops[i]);
ops.push(OpEntry {
out_len,
dep_kind,
dep_id,
});
}
let raw_count = physical
.objects
.len()
.saturating_add(physical.streams.len())
.saturating_add(physical.revisions.len());
if raw_count as u64 > u64::from(limits.max_index_selectors) {
return Ok(None);
}
let mut selectors: Vec<ObservationSelector> = Vec::with_capacity(raw_count);
let mut seen: HashSet<(u8, u32, u32)> = HashSet::new();
for o in &physical.objects {
let (Ok(number), Ok(generation)) = (u32::try_from(o.number), u32::try_from(o.generation))
else {
return Ok(None);
};
let len = o.end.saturating_sub(o.start);
if len == 0
|| o.end > descriptor.source_len
|| !seen.insert((SELECTOR_OBJECT, number, generation))
{
continue;
}
selectors.push(ObservationSelector {
kind: SELECTOR_OBJECT,
number,
generation,
out_off: o.start,
out_len: len,
});
}
for s in &physical.streams {
let (Ok(number), Ok(generation)) = (u32::try_from(s.object), u32::try_from(s.generation))
else {
return Ok(None);
};
let end = s.data_start.saturating_add(s.data_len);
if s.data_len == 0
|| end > descriptor.source_len
|| !seen.insert((SELECTOR_STREAM, number, generation))
{
continue;
}
selectors.push(ObservationSelector {
kind: SELECTOR_STREAM,
number,
generation,
out_off: s.data_start,
out_len: s.data_len,
});
}
for r in &physical.revisions {
let len = r.end.saturating_sub(r.start);
if len == 0
|| r.end > descriptor.source_len
|| !seen.insert((SELECTOR_REVISION, r.index, 0))
{
continue;
}
selectors.push(ObservationSelector {
kind: SELECTOR_REVISION,
number: r.index,
generation: 0,
out_off: r.start,
out_len: len,
});
}
selectors.sort_by_key(|s| (s.kind, s.number, s.generation, s.out_off));
let index = ObservationIndex {
section_flags: SECTION_OP_TABLE | SECTION_PDF_SELECTORS,
ops,
selectors,
digests: Vec::new(),
};
if index.encode()?.len() as u64 > u64::from(limits.max_record_len) {
return Ok(None);
}
Ok(Some(index))
}
#[cfg(feature = "rans")]
fn primary_dependency(op: &Op) -> (u8, u32) {
crate::container::observation::primary_dependency(op)
}
fn intern(objects: &mut Vec<Vec<u8>>, index: &mut HashMap<Vec<u8>, u32>, bytes: Vec<u8>) -> u32 {
if let Some(&id) = index.get(&bytes) {
return id;
}
let id = objects.len() as u32;
index.insert(bytes.clone(), id);
objects.push(bytes);
id
}
pub const RANS_SCALE_BITS: u8 = 12;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StreamStats {
pub object: u64,
pub generation: u64,
pub compressed_bytes: u64,
pub replayed: bool,
pub decline_reason: Option<&'static str>,
pub plaintext_bytes: Option<u64>,
pub correction_bytes: Option<u64>,
pub rans_plaintext_bytes: Option<u64>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct DeflateSummary {
pub flate_streams: u64,
pub replayed: u64,
pub declined: u64,
pub compressed_bytes: u64,
pub plaintext_bytes: u64,
pub correction_bytes: u64,
pub rans_plaintext_bytes: u64,
pub replayed_rans_dedup_bytes: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeflateStats {
pub is_pdf: bool,
pub streams: Vec<StreamStats>,
pub summary: DeflateSummary,
}
pub fn deflate_stats(input: &[u8], limits: Limits) -> Result<DeflateStats> {
if !super::detect(input, limits) {
return Ok(DeflateStats {
is_pdf: false,
streams: Vec::new(),
summary: DeflateSummary::default(),
});
}
let physical = scan(input, limits)?;
let mut streams: Vec<StreamStats> = Vec::new();
let mut summary = DeflateSummary::default();
let mut unique_plaintexts: HashSet<Vec<u8>> = HashSet::new();
let mut unique_corrections: HashSet<Vec<u8>> = HashSet::new();
for stream in &physical.streams {
if stream.filter != FilterClass::FlateDecode {
continue;
}
summary.flate_streams += 1;
summary.compressed_bytes += stream.data_len;
let mut st = StreamStats {
object: stream.object,
generation: stream.generation,
compressed_bytes: stream.data_len,
replayed: false,
decline_reason: None,
plaintext_bytes: None,
correction_bytes: None,
rans_plaintext_bytes: None,
};
let start = stream.data_start as usize;
let end = start
.checked_add(stream.data_len as usize)
.filter(|&e| e <= input.len());
match end {
None => st.decline_reason = Some("span_out_of_range"),
Some(end) => match try_replay_detailed(&input[start..end], limits) {
Ok(plan) => {
let plaintext_bytes = plan.plaintext.len() as u64;
let correction_bytes = plan.corrections.len() as u64;
st.replayed = true;
st.plaintext_bytes = Some(plaintext_bytes);
st.correction_bytes = Some(correction_bytes);
st.rans_plaintext_bytes = rans_plaintext_bytes(&plan.plaintext);
summary.replayed += 1;
summary.plaintext_bytes += plaintext_bytes;
summary.correction_bytes += correction_bytes;
summary.rans_plaintext_bytes += st.rans_plaintext_bytes.unwrap_or(0);
if unique_corrections.insert(plan.corrections.clone()) {
summary.replayed_rans_dedup_bytes += correction_bytes;
}
if unique_plaintexts.insert(plan.plaintext.clone()) {
summary.replayed_rans_dedup_bytes += st.rans_plaintext_bytes.unwrap_or(0);
}
}
Err(reason) => st.decline_reason = Some(reason.name()),
},
}
if !st.replayed {
summary.declined += 1;
}
streams.push(st);
}
Ok(DeflateStats {
is_pdf: true,
streams,
summary,
})
}
#[cfg(feature = "rans")]
fn rans_plaintext_bytes(plaintext: &[u8]) -> Option<u64> {
use crate::entropy::codec::WIRE_HEADER_LEN;
use crate::entropy::{ALPHABET, EntropyModel, encode_channel};
let mut counts = [0u64; ALPHABET];
for &b in plaintext {
counts[b as usize] += 1;
}
let model = EntropyModel::from_counts(&counts, RANS_SCALE_BITS).ok()?;
let model_bytes = model.encode().ok()?.len() as u64;
let capsule = encode_channel(&model, plaintext).ok()?;
Some(model_bytes + WIRE_HEADER_LEN as u64 + capsule.payload.len() as u64)
}
#[cfg(not(feature = "rans"))]
fn rans_plaintext_bytes(_plaintext: &[u8]) -> Option<u64> {
None
}