use std::fmt;
use std::io::{self, Read};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct InflateLimits {
pub grace_size: u64,
pub max_expansion: u64,
pub max_size: u64,
}
impl InflateLimits {
pub const POI: Self = Self {
grace_size: 100 * 1024,
max_expansion: 100,
max_size: 0xFFFF_FFFF,
};
pub fn allows(&self, compressed: u64, inflated: u64) -> bool {
inflated <= self.max_size
&& (inflated <= self.grace_size
|| inflated <= compressed.saturating_mul(self.max_expansion))
}
pub fn check(&self, compressed: u64, inflated: u64) -> io::Result<()> {
if self.allows(compressed, inflated) {
Ok(())
} else {
Err(io::Error::new(
io::ErrorKind::InvalidData,
InflationLimit {
compressed,
inflated,
},
))
}
}
}
impl Default for InflateLimits {
fn default() -> Self {
Self::POI
}
}
#[derive(Debug)]
pub struct InflateGuard<R> {
inner: R,
compressed: u64,
inflated: u64,
limits: InflateLimits,
}
impl<R> InflateGuard<R> {
pub fn new(inner: R, compressed_size: u64) -> Self {
Self::with_limits(inner, compressed_size, InflateLimits::POI)
}
pub fn with_limits(inner: R, compressed_size: u64, limits: InflateLimits) -> Self {
Self {
inner,
compressed: compressed_size,
inflated: 0,
limits,
}
}
}
impl<R: Read> Read for InflateGuard<R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let n = self.inner.read(buf)?;
self.inflated += n as u64;
self.limits.check(self.compressed, self.inflated)?;
Ok(n)
}
}
#[derive(Debug)]
pub struct InflationLimit {
compressed: u64,
inflated: u64,
}
impl fmt::Display for InflationLimit {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"{} compressed bytes expanded past {} bytes: more than a document's data expands \
(a decompression bomb)",
self.compressed, self.inflated
)
}
}
impl std::error::Error for InflationLimit {}
pub fn is_inflation_limit(error: &io::Error) -> bool {
error
.get_ref()
.is_some_and(|inner| inner.is::<InflationLimit>())
}
#[cfg(test)]
mod tests {
use super::*;
fn read_all(data: impl Read, compressed: u64) -> io::Result<usize> {
InflateGuard::new(data, compressed).read_to_end(&mut Vec::new())
}
#[test]
fn small_output_is_never_refused() {
let grace = InflateLimits::POI.grace_size;
assert_eq!(
read_all(io::repeat(0).take(grace), 1).unwrap(),
grace as usize
);
}
#[test]
fn other_limits_move_both_bounds() {
let limits = InflateLimits {
grace_size: 1 << 20,
max_expansion: 600,
max_size: 8 << 20,
};
let read = |len: u64, compressed: u64| {
InflateGuard::with_limits(io::repeat(0).take(len), compressed, limits)
.read_to_end(&mut Vec::new())
};
assert!(read(5 << 20, (5 << 20) / 500).is_ok());
assert!(is_inflation_limit(
&read(5 << 20, (5 << 20) / 1000).unwrap_err()
));
assert!(read(1 << 20, 1).is_ok());
assert!(is_inflation_limit(&read(9 << 20, 9 << 20).unwrap_err()));
}
#[test]
fn check_applies_the_rule_without_a_reader() {
let poi = InflateLimits::default();
assert!(poi.check(1, poi.grace_size).is_ok());
let err = poi.check(1, poi.grace_size + 1).unwrap_err();
assert!(is_inflation_limit(&err));
assert!(poi.allows(u64::MAX, poi.max_size));
assert!(!poi.allows(u64::MAX, poi.max_size + 1));
}
#[test]
fn output_past_a_hundredfold_is_a_bomb() {
let err = read_all(io::repeat(0).take(10 << 20), 10 * 1024).unwrap_err();
assert!(is_inflation_limit(&err), "{err}");
assert!(err.to_string().contains("decompression bomb"), "{err}");
}
#[test]
fn ordinary_compression_passes() {
assert_eq!(
read_all(io::repeat(b'x').take(4 << 20), (4 << 20) / 20).unwrap(),
4 << 20
);
}
#[test]
fn an_io_error_underneath_is_not_the_limit() {
struct Failing;
impl Read for Failing {
fn read(&mut self, _: &mut [u8]) -> io::Result<usize> {
Err(io::Error::other("disk"))
}
}
let err = read_all(Failing, 10).unwrap_err();
assert!(!is_inflation_limit(&err));
}
}