xberg-native-pdf 1.1.3

//! FlateDecode (zlib/deflate) implementation.
//!
//! This is the most common PDF compression filter, used in ~90% of PDFs.
//! Uses the flate2 crate for zlib decompression.

#![forbid(unsafe_code)]

use crate::decoders::StreamDecoder;
use crate::error::{Error, Result};
use flate2::read::{DeflateDecoder, ZlibDecoder};
use std::io::Read;

/// Default cap for [`FlateDecoder`]: 256 MB per stream.
///
/// Prevents zip-bomb / flate-bomb attacks where a tiny compressed stream
/// expands to an arbitrarily large output, exhausting virtual memory and
/// triggering an allocator abort (SIGABRT / exit 134).
///
/// 256 MB accommodates A4 @ 600 DPI RGB (~99 MB) with headroom.
///
/// Override via:
/// - `XBERG_NATIVE_PDF_MAX_DECOMPRESS_MB` environment variable (e.g. `64` for
///   64 MB), or the pre-rename `PDF_OXIDE_MAX_DECOMPRESS_MB` if the new name
///   is unset — this is a security control, so an upgrade must never silently
///   revert a deployment to the 256 MB default because only the old name was
///   ever set.
/// - [`FlateDecoder::with_limit`] for programmatic control
pub const DEFAULT_MAX_DECOMPRESSED_BYTES: u64 = 256 * 1024 * 1024;

fn effective_limit_from_str(val: Option<&str>) -> u64 {
    val.and_then(|v| v.parse::<u64>().ok())
        .map(|mb| mb * 1024 * 1024)
        .unwrap_or(DEFAULT_MAX_DECOMPRESSED_BYTES)
}

/// Resolve the limit from the two possible environment values, new name
/// first. Pulled apart from [`effective_limit`] so the fallback precedence
/// is testable without mutating real process environment variables.
fn effective_limit_from_env_pair(new_val: Option<&str>, old_val: Option<&str>) -> u64 {
    effective_limit_from_str(new_val.or(old_val))
}

/// Read the decompression limit from the environment, falling back to the
/// compile-time default.
///
/// Three names are accepted, newest first: `XBERG_NATIVE_PDF_MAX_DECOMPRESS_MB`, then the
/// pre-vendoring `XBERG_PDF_OXIDE_MAX_DECOMPRESS_MB`, then upstream's original
/// `PDF_OXIDE_MAX_DECOMPRESS_MB`. Each rename kept the older spellings working because this
/// is a deployment knob someone has already set in a Dockerfile or a systemd unit; dropping a
/// name silently reverts their limit to the compile-time default with no error. ~keep
fn effective_limit() -> u64 {
    let new_val = std::env::var("XBERG_NATIVE_PDF_MAX_DECOMPRESS_MB").ok();
    let previous_val = std::env::var("XBERG_PDF_OXIDE_MAX_DECOMPRESS_MB").ok();
    let old_val = std::env::var("PDF_OXIDE_MAX_DECOMPRESS_MB").ok();
    effective_limit_from_env_pair(new_val.or(previous_val).as_deref(), old_val.as_deref())
}

/// Heuristic validator for partial-recovery output from a failing decompress.
///
/// Returns `true` if the decoded bytes look like a plausible PDF stream — any
/// of: a content-stream operator (BT/ET/Tj/TJ/Tm/Td), a common PDF object
/// marker (<</>>/stream/obj/endobj), or a `%PDF-` prefix. The set is kept
/// intentionally broad because FlateDecode also wraps object streams, xref
/// streams, font programs, and image data, none of which carry text operators.
///
/// The guard exists to distinguish genuine partially-decoded output from
/// `deflate` "success" on a misaligned input — the latter produces short runs
/// of pseudo-random bytes (128 bytes of `P\xffj!}` repeating) that contain
/// none of these markers.
///
/// Conservative fallback: if the decoded bytes are mostly ASCII (printable
/// plus whitespace), the output is also treated as plausible, because stream
/// contents in the wild include ASCII-only data (hex-encoded images, small
/// object streams) that do not hit any specific marker.
fn looks_like_real_stream(output: &[u8]) -> bool {
    if output.is_empty() {
        return false;
    }
    const MARKERS: &[&[u8]] = &[b"BT", b"ET", b"Tj", b"TJ", b"Tm", b"Td", b"stream", b"endobj", b"%PDF-"];
    for m in MARKERS {
        if output.windows(m.len()).any(|w| w == *m) {
            return true;
        }
    }
    let printable = output
        .iter()
        .filter(|&&b| (0x20..=0x7E).contains(&b) || b == b'\t' || b == b'\n' || b == b'\r')
        .count();
    printable * 5 >= output.len() * 4
}

/// Returns `Err` if `output` reached the decompression cap, indicating that the
/// stream was truncated rather than fully decoded.
#[inline]
fn check_limit(output: &[u8], limit: u64) -> Result<()> {
    if output.len() as u64 >= limit {
        return Err(Error::Decode(format!(
            "FlateDecode output reached the {} MB safety limit; \
             stream may be a flate bomb or an unusually large image",
            limit / (1024 * 1024)
        )));
    }
    Ok(())
}

/// FlateDecode filter implementation.
///
/// Decompresses data using the zlib/deflate algorithm. The decompression cap
/// defaults to `DEFAULT_MAX_DECOMPRESSED_BYTES` and can be overridden with
/// [`FlateDecoder::with_limit`].
pub struct FlateDecoder {
    /// Maximum number of decompressed bytes accepted per stream.
    pub max_decompressed_bytes: u64,
}

impl Default for FlateDecoder {
    fn default() -> Self {
        Self {
            max_decompressed_bytes: effective_limit(),
        }
    }
}

impl FlateDecoder {
    /// Creates a decoder that rejects any stream decompressing to more than
    /// `limit` bytes. Use this to tighten or relax the default 256 MB cap.
    pub fn with_limit(limit: u64) -> Self {
        Self {
            max_decompressed_bytes: limit,
        }
    }
}

impl StreamDecoder for FlateDecoder {
    fn decode(&self, input: &[u8]) -> Result<Vec<u8>> {
        let mut decoder = ZlibDecoder::new(input).take(self.max_decompressed_bytes);
        let mut output = Vec::new();

        match decoder.read_to_end(&mut output) {
            Ok(_) => {
                check_limit(&output, self.max_decompressed_bytes)?;
                Ok(output)
            }
            Err(e) => {
                // Partial recovery: return only if output *looks like* a
                // plausible stream. The pre-fix behaviour accepted any
                // non-empty buffer, which let strategies 2 and 3 return
                // misaligned-deflate garbage (`P\xffj!}` × 16) that the text
                // extractor then emitted as zero bytes of output. ~keep
                if !output.is_empty() && looks_like_real_stream(&output) {
                    check_limit(&output, self.max_decompressed_bytes)?;
                    tracing::warn!(
                        filter = "FlateDecode",
                        bytes = output.len(),
                        error = %e,
                        "partial recovery: extracted bytes before corruption"
                    );
                    return Ok(output);
                }

                // Strategy 2: Try raw deflate (no zlib wrapper)
                // Some PDFs have corrupt zlib headers but valid deflate data ~keep
                tracing::debug!(filter = "FlateDecode", "zlib decode failed, trying raw deflate");
                output.clear();
                let mut deflate_decoder = DeflateDecoder::new(input).take(self.max_decompressed_bytes);

                match deflate_decoder.read_to_end(&mut output) {
                    Ok(_) => {
                        check_limit(&output, self.max_decompressed_bytes)?;
                        tracing::warn!(
                            filter = "FlateDecode",
                            bytes = output.len(),
                            "recovered via raw deflate fallback after corrupt zlib header"
                        );
                        Ok(output)
                    }
                    Err(deflate_err) => {
                        if !output.is_empty() && looks_like_real_stream(&output) {
                            check_limit(&output, self.max_decompressed_bytes)?;
                            tracing::warn!(
                                filter = "FlateDecode",
                                bytes = output.len(),
                                "raw deflate partial recovery: extracted bytes before error"
                            );
                            return Ok(output);
                        }

                        if input.len() > 2 {
                            tracing::debug!(
                                filter = "FlateDecode",
                                "trying deflate after skipping potential corrupt zlib header"
                            );
                            output.clear();
                            let mut deflate_decoder =
                                DeflateDecoder::new(&input[2..]).take(self.max_decompressed_bytes);

                            match deflate_decoder.read_to_end(&mut output) {
                                Ok(_) => {
                                    check_limit(&output, self.max_decompressed_bytes)?;
                                    tracing::warn!(
                                        filter = "FlateDecode",
                                        bytes = output.len(),
                                        "recovered by skipping corrupt zlib header bytes"
                                    );
                                    return Ok(output);
                                }
                                Err(_) => {
                                    if !output.is_empty() && looks_like_real_stream(&output) {
                                        check_limit(&output, self.max_decompressed_bytes)?;
                                        tracing::warn!(
                                            filter = "FlateDecode",
                                            bytes = output.len(),
                                            "header-skip partial recovery"
                                        );
                                        return Ok(output);
                                    }
                                }
                            }
                        }

                        // Strategy 4: Try fixing corrupt zlib header byte
                        // If first byte has invalid compression method, replace with 0x78 (standard deflate)
                        // ~keep
                        if input.len() >= 2 {
                            let first_byte = input[0];
                            let compression_method = first_byte & 0x0F;
                            if compression_method != 8 {
                                tracing::debug!(
                                    filter = "FlateDecode",
                                    compression_method,
                                    header_byte = format_args!("0x{:02x}", first_byte),
                                    "invalid compression method in header byte, trying corrected header"
                                );
                                let mut corrected = input.to_vec();
                                // Replace CM bits (0-3) with 8 (deflate), keep CINFO bits (4-7)
                                // ~keep
                                corrected[0] = (first_byte & 0xF0) | 0x08;

                                output.clear();
                                let mut decoder = ZlibDecoder::new(&corrected[..]).take(self.max_decompressed_bytes);
                                match decoder.read_to_end(&mut output) {
                                    Ok(_) if !output.is_empty() => {
                                        check_limit(&output, self.max_decompressed_bytes)?;
                                        tracing::warn!(
                                            filter = "FlateDecode",
                                            bytes = output.len(),
                                            "recovered via corrected zlib header byte"
                                        );
                                        return Ok(output);
                                    }
                                    Err(_) if !output.is_empty() && looks_like_real_stream(&output) => {
                                        check_limit(&output, self.max_decompressed_bytes)?;
                                        tracing::warn!(
                                            filter = "FlateDecode",
                                            bytes = output.len(),
                                            "header-correction partial recovery"
                                        );
                                        return Ok(output);
                                    }
                                    _ => {
                                        tracing::debug!(filter = "FlateDecode", "header correction failed");
                                    }
                                }
                            }
                        }

                        // Strategy 5: Brute-force scan for valid deflate data
                        // Try starting deflate decompression from offsets 0-20
                        // BUT validate the output contains valid PDF operators ~keep
                        tracing::debug!(filter = "FlateDecode", "trying brute-force scan for valid deflate data");
                        let max_offset = std::cmp::min(20, input.len());
                        for offset in 0..max_offset {
                            if offset == 0 || offset == 2 {
                                continue;
                            }

                            output.clear();
                            let mut deflate_decoder =
                                DeflateDecoder::new(&input[offset..]).take(self.max_decompressed_bytes);

                            match deflate_decoder.read_to_end(&mut output) {
                                Ok(_) if !output.is_empty() => {
                                    check_limit(&output, self.max_decompressed_bytes)?;
                                    let decoded_str = String::from_utf8_lossy(&output);
                                    let has_pdf_operators = decoded_str.contains("BT")
                                        || decoded_str.contains("ET")
                                        || decoded_str.contains("Tj")
                                        || decoded_str.contains("TJ")
                                        || decoded_str.contains("Tm")
                                        || decoded_str.contains("Td");

                                    if has_pdf_operators {
                                        tracing::warn!(
                                            filter = "FlateDecode",
                                            offset,
                                            bytes = output.len(),
                                            "brute-force deflate recovery succeeded (validated PDF content)"
                                        );
                                        return Ok(output);
                                    } else {
                                        tracing::trace!(
                                            filter = "FlateDecode",
                                            offset,
                                            bytes = output.len(),
                                            "brute-force offset produced no valid PDF operators, trying next offset"
                                        );
                                        continue;
                                    }
                                }
                                Err(_) if !output.is_empty() => {
                                    check_limit(&output, self.max_decompressed_bytes)?;
                                    let decoded_str = String::from_utf8_lossy(&output);
                                    let has_pdf_operators = decoded_str.contains("BT")
                                        || decoded_str.contains("ET")
                                        || decoded_str.contains("Tj")
                                        || decoded_str.contains("TJ")
                                        || decoded_str.contains("Tm")
                                        || decoded_str.contains("Td");

                                    if has_pdf_operators {
                                        tracing::warn!(
                                            filter = "FlateDecode",
                                            offset,
                                            bytes = output.len(),
                                            "brute-force partial recovery (validated PDF content)"
                                        );
                                        return Ok(output);
                                    } else {
                                        tracing::trace!(
                                            filter = "FlateDecode",
                                            offset,
                                            "brute-force partial recovery has no valid PDF operators, \
                                             trying next offset"
                                        );
                                        continue;
                                    }
                                }
                                _ => continue,
                            }
                        }

                        // SPEC COMPLIANCE FIX: Removed strategies 8-9 that violated PDF spec
                        //
                        // Previous strategies 8-9 would return raw uncompressed data for streams
                        // labeled as /FlateDecode. This violates PDF Spec ISO 32000-1:2008,
                        // Section 7.3.8.2 which states that if a stream has /Filter /FlateDecode,
                        // it MUST be compressed with the FlateDecode algorithm.
                        //
                        // Returning raw data creates security risks:
                        // 1. Malicious PDFs could bypass compression validation
                        // 2. Type confusion attacks (treating compressed data as raw)
                        // 3. Inconsistent behavior across PDF processors
                        //
                        // Correct behavior: If all decompression strategies fail, return an error.
                        // The stream is either corrupted or malicious, and should not be processed.
                        // ~keep

                        Err(Error::Decode(format!(
                            "FlateDecode decompression failed: stream is labeled as compressed but all decompression attempts failed. \
                            This violates PDF Spec ISO 32000-1:2008, Section 7.3.8.2. \
                            Zlib error: {}, Deflate error: {}. Compressed size: {} bytes.",
                            e,
                            deflate_err,
                            input.len()
                        )))
                    }
                }
            }
        }
    }

    fn name(&self) -> &str {
        "FlateDecode"
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use flate2::Compression;
    use flate2::write::ZlibEncoder;
    use std::io::Write;

    // When a strategy's partial recovery emits garbage (valid-looking
    // deflate bits that decode to pseudo-random bytes on a misaligned input),
    // the decoder must not accept it. Strategy 5 already validates via PDF
    // content-stream operators; strategies 1–4 now validate via
    // `looks_like_real_stream`. This test pins that guard. ~keep
    #[test]
    fn looks_like_real_stream_rejects_repeating_garbage() {
        // Actual symptom before the fix: 128 bytes of
        // `P\xffj!}\xef\xbd\xbd\xef\xbd\xbd...` high-bit-heavy repetition. ~keep
        let garbage =
            b"P\xffj!}\xef\xbd\xbd\xef\xbd\xbd\xef\xbd\xbd\xef\xbd\xbd\xef\xbd\xbd\xef\xbd\xbd\xef\xbd\xbd\xef\xbd\xbd"
                .repeat(4);
        assert!(
            !looks_like_real_stream(&garbage),
            "misaligned-deflate garbage must be rejected as a partial recovery"
        );
    }

    #[test]
    fn looks_like_real_stream_accepts_content_stream_operators() {
        let real = b"BT /F1 12 Tf 100 700 Td (hello) Tj ET";
        assert!(looks_like_real_stream(real));
    }

    #[test]
    fn looks_like_real_stream_accepts_ascii_only_object_stream() {
        // Object-stream-like payload: ASCII, no content-stream operators. ~keep
        let object_stream = b"1 0 obj\n<< /Length 42 >>\nstream\nhello world\nendstream\nendobj\n";
        assert!(looks_like_real_stream(object_stream));
    }

    #[test]
    fn looks_like_real_stream_rejects_empty() {
        assert!(!looks_like_real_stream(&[]));
    }

    #[test]
    fn test_flate_decode_simple() {
        let decoder = FlateDecoder::default();

        let original = b"Hello, FlateDecode!";
        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
        encoder.write_all(original).unwrap();
        let compressed = encoder.finish().unwrap();

        let decoded = decoder.decode(&compressed).unwrap();
        assert_eq!(decoded, original);
    }

    #[test]
    fn test_flate_decode_empty() {
        let decoder = FlateDecoder::default();

        let original = b"";
        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
        encoder.write_all(original).unwrap();
        let compressed = encoder.finish().unwrap();

        let decoded = decoder.decode(&compressed).unwrap();
        assert_eq!(decoded, original);
    }

    #[test]
    fn test_flate_decode_large_data() {
        let decoder = FlateDecoder::default();

        let original = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ".repeat(1000);
        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
        encoder.write_all(&original).unwrap();
        let compressed = encoder.finish().unwrap();

        let decoded = decoder.decode(&compressed).unwrap();
        assert_eq!(decoded, original);
    }

    #[test]
    fn test_flate_decode_invalid_data() {
        let decoder = FlateDecoder::default();

        // Invalid zlib data - should fail decompression
        // SPEC COMPLIANCE: We now correctly reject invalid compressed data
        // instead of returning it as raw data (which violated PDF spec) ~keep
        let invalid = b"This is not zlib compressed data";
        let result = decoder.decode(invalid);
        assert!(result.is_err());

        if let Err(e) = result {
            let error_msg = format!("{}", e);
            assert!(error_msg.contains("FlateDecode decompression failed"));
        }
    }

    #[test]
    fn test_flate_decoder_name() {
        let decoder = FlateDecoder::default();
        assert_eq!(decoder.name(), "FlateDecode");
    }

    #[test]
    fn test_flate_bomb_rejected() {
        let large = vec![0u8; DEFAULT_MAX_DECOMPRESSED_BYTES as usize];
        let result = check_limit(&large, DEFAULT_MAX_DECOMPRESSED_BYTES);
        assert!(result.is_err());
        let msg = format!("{}", result.unwrap_err());
        assert!(msg.contains("safety limit"));
    }

    #[test]
    fn test_check_limit_below_threshold() {
        let small = vec![0u8; 1024];
        assert!(check_limit(&small, DEFAULT_MAX_DECOMPRESSED_BYTES).is_ok());
    }

    #[test]
    fn test_custom_limit_accepts_data_within_limit() {
        let original = b"x".repeat(512);
        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
        encoder.write_all(&original).unwrap();
        let compressed = encoder.finish().unwrap();

        let decoder = FlateDecoder::with_limit(1024);
        let decoded = decoder.decode(&compressed).unwrap();
        assert_eq!(decoded, original);
    }

    #[test]
    fn test_custom_limit_rejects_data_over_limit() {
        let original = b"x".repeat(100);
        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
        encoder.write_all(&original).unwrap();
        let compressed = encoder.finish().unwrap();

        let decoder = FlateDecoder::with_limit(10);
        let result = decoder.decode(&compressed);
        assert!(result.is_err(), "expected rejection when output exceeds custom limit");
    }

    #[test]
    fn test_bomb_error_does_not_expose_internal_symbol_name() {
        let large = vec![0u8; DEFAULT_MAX_DECOMPRESSED_BYTES as usize];
        let result = check_limit(&large, DEFAULT_MAX_DECOMPRESSED_BYTES);
        assert!(result.is_err());
        let msg = format!("{}", result.unwrap_err());
        assert!(
            !msg.contains("MAX_DECOMPRESSED_BYTES"),
            "error message must not reference internal symbol names: {msg}"
        );
    }

    #[test]
    fn test_effective_limit_env_variable() {
        assert_eq!(effective_limit_from_str(None), DEFAULT_MAX_DECOMPRESSED_BYTES);
        assert_eq!(effective_limit_from_str(Some("64")), 64 * 1024 * 1024);
        assert_eq!(
            effective_limit_from_str(Some("not_a_number")),
            DEFAULT_MAX_DECOMPRESSED_BYTES
        );
    }

    #[test]
    fn test_effective_limit_prefers_new_env_var_and_falls_back_to_old() {
        assert_eq!(effective_limit_from_env_pair(Some("64"), Some("128")), 64 * 1024 * 1024);
        // New name unset: the pre-rename name still works, so an upgrade
        // never silently reverts a deployment to the compile-time default. ~keep
        assert_eq!(effective_limit_from_env_pair(None, Some("64")), 64 * 1024 * 1024);
        assert_eq!(
            effective_limit_from_env_pair(None, None),
            DEFAULT_MAX_DECOMPRESSED_BYTES
        );
    }
}