use std::io::Write;
use std::time::Instant;
use flate2::write::ZlibEncoder;
use flate2::Compression;
use lopdf::{Dictionary, Document, LoadOptions, Object, Stream};
const MIB: usize = 1024 * 1024;
fn flate_bomb(target: usize) -> Vec<u8> {
let mut encoder = ZlibEncoder::new(Vec::new(), Compression::best());
let zeros = [0u8; 64 * 1024];
let mut remaining = target;
while remaining > 0 {
let n = remaining.min(zeros.len());
encoder.write_all(&zeros[..n]).unwrap();
remaining -= n;
}
encoder.finish().unwrap()
}
fn zlib_compress(data: &[u8]) -> Vec<u8> {
let mut encoder = ZlibEncoder::new(Vec::new(), Compression::best());
encoder.write_all(data).unwrap();
encoder.finish().unwrap()
}
fn main() {
println!("== lopdf decompression-bomb guard ==\n");
let logical_output = 256 * MIB;
let bomb = flate_bomb(logical_output);
println!(
"Built a FlateDecode bomb: {} bytes compressed -> {} bytes decompressed ({}x amplification).",
bomb.len(),
logical_output,
logical_output / bomb.len().max(1)
);
let mut dict = Dictionary::new();
dict.set("Filter", "FlateDecode");
let stream = Stream::new(dict, bomb.clone());
let limit = 8 * MIB;
let start = Instant::now();
match stream.decompressed_content_with_limit(limit) {
Ok(data) => println!(" UNEXPECTED: decompressed {} bytes", data.len()),
Err(e) => println!(
" decompressed_content_with_limit({} MiB) rejected it in {:?}: {e}",
limit / MIB,
start.elapsed()
),
}
println!(
" (The default decompressed_content() has NO limit and would allocate all {} MiB.)\n",
logical_output / MIB
);
let inner = flate_bomb(256 * MIB);
let outer = zlib_compress(&inner);
let mut dict = Dictionary::new();
dict.set("Filter", vec![Object::from("FlateDecode"), Object::from("FlateDecode")]);
let chained = Stream::new(dict, outer.clone());
println!(
"Built a chained [FlateDecode FlateDecode] bomb: {} bytes -> ~256 MiB after two layers.",
outer.len()
);
match chained.decompressed_content_with_limit(limit) {
Ok(data) => println!(" UNEXPECTED: decompressed {} bytes", data.len()),
Err(e) => println!(" rejected at a filter layer with an {} MiB limit: {e}\n", limit / MIB),
}
let pdf = xref_stream_bomb_pdf(&bomb);
println!("Loading a PDF whose cross-reference stream is the bomb...");
match Document::load_mem(&pdf) {
Ok(_) => println!(" load_mem (no limit): parsed after decoding the full output"),
Err(e) => println!(" load_mem (no limit): {e}"),
}
match Document::load_mem_with_options(&pdf, LoadOptions::with_max_decompressed_size(limit)) {
Ok(_) => println!(" UNEXPECTED: guarded load succeeded"),
Err(e) => println!(" load_mem_with_options(max_decompressed_size = {} MiB): {e}", limit / MIB),
}
}
fn xref_stream_bomb_pdf(bomb: &[u8]) -> Vec<u8> {
let mut pdf = Vec::new();
pdf.extend_from_slice(b"%PDF-1.5\n");
let obj_offset = pdf.len();
pdf.extend_from_slice(b"1 0 obj\n");
pdf.extend_from_slice(
format!(
"<< /Type /XRef /Size 1 /W [1 1 1] /Root 1 0 R /Filter /FlateDecode /Length {} >>\n",
bomb.len()
)
.as_bytes(),
);
pdf.extend_from_slice(b"stream\n");
pdf.extend_from_slice(bomb);
pdf.extend_from_slice(b"\nendstream\nendobj\n");
pdf.extend_from_slice(format!("startxref\n{obj_offset}\n%%EOF").as_bytes());
pdf
}