rpdfium_codec/flate/

flatemodule.rs

// Derived from PDFium's flate/flatemodule.cpp
// Original: Copyright 2014 The PDFium Authors
// Licensed under BSD-3-Clause / Apache-2.0
// See pdfium-upstream/LICENSE for the original license.

//! FlateDecode filter — zlib/deflate decompression.

use flate2::{Decompress, FlushDecompress, Status};

use super::predictor;
use crate::error::DecodeError;

/// Maximum decompressed stream size from security limits.
const MAX_DECOMPRESSED_SIZE: u64 = rpdfium_core::fx_system::DEFAULT_MAX_DECOMPRESSED_STREAM_SIZE;

/// Decode Flate (zlib) compressed data, optionally applying predictor processing.
///
/// Parameters:
/// - `predictor`: Predictor type (1=none, 2=TIFF, 10-15=PNG). Defaults to 1.
/// - `columns`: Number of columns per row. Defaults to 1.
/// - `colors`: Number of color components per sample. Defaults to 1.
/// - `bits_per_component`: Bits per color component. Defaults to 8.
pub fn decode(
    input: &[u8],
    predictor: Option<i32>,
    columns: Option<i32>,
    colors: Option<i32>,
    bits_per_component: Option<i32>,
) -> Result<Vec<u8>, DecodeError> {
    let decompressed = inflate(input)?;

    let predictor_val = predictor.unwrap_or(1);
    if predictor_val <= 1 {
        return Ok(decompressed);
    }

    let columns = columns.unwrap_or(1);
    let colors = colors.unwrap_or(1);
    let bpc = bits_per_component.unwrap_or(8);

    predictor::apply_predictor(&decompressed, predictor_val, columns, colors, bpc)
}

/// Encode data using zlib/deflate compression.
///
/// Mirrors PDFium's `FlateModule::Encode()`.
pub fn encode(input: &[u8]) -> Vec<u8> {
    use flate2::Compression;
    use flate2::write::ZlibEncoder;
    use std::io::Write;

    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
    encoder.write_all(input).unwrap_or_default();
    encoder.finish().unwrap_or_default()
}

/// Raw zlib/deflate decompression with size limit enforcement.
fn inflate(input: &[u8]) -> Result<Vec<u8>, DecodeError> {
    let mut decompressor = Decompress::new(true); // zlib header
    let mut output = Vec::with_capacity(input.len().saturating_mul(2).min(1 << 20));
    let mut buf = [0u8; 32 * 1024]; // 32KB working buffer

    loop {
        let in_before = decompressor.total_in();
        let out_before = decompressor.total_out();

        let status = decompressor
            .decompress(
                &input[in_before as usize..],
                &mut buf,
                FlushDecompress::None,
            )
            .map_err(|e| DecodeError::Flate(e.to_string()))?;

        let bytes_written = (decompressor.total_out() - out_before) as usize;
        output.extend_from_slice(&buf[..bytes_written]);

        if output.len() as u64 > MAX_DECOMPRESSED_SIZE {
            return Err(DecodeError::OutputTooLarge {
                limit: MAX_DECOMPRESSED_SIZE,
            });
        }

        match status {
            Status::StreamEnd => break,
            Status::Ok | Status::BufError => {
                let in_after = decompressor.total_in();
                let out_after = decompressor.total_out();
                // No progress — we're stuck
                if in_after == in_before && out_after == out_before {
                    break;
                }
            }
        }
    }

    Ok(output)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_encode_decode_roundtrip() {
        let original = b"Hello, World! FlateEncode/FlateDecode round-trip.";
        let compressed = encode(original);
        assert!(!compressed.is_empty());
        let result = decode(&compressed, None, None, None, None).unwrap();
        assert_eq!(result, original);
    }

    #[test]
    fn test_encode_empty() {
        let compressed = encode(b"");
        // zlib header is emitted even for empty input
        assert!(!compressed.is_empty());
        let result = decode(&compressed, None, None, None, None).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_encode_all_bytes() {
        let original: Vec<u8> = (0u8..=255).collect();
        let compressed = encode(&original);
        // All 256 distinct bytes should compress without loss
        assert_eq!(
            decode(&compressed, None, None, None, None).unwrap(),
            original
        );
    }

    #[test]
    fn test_encode_large_data() {
        let original: Vec<u8> = (0..10_000).map(|i| (i % 256) as u8).collect();
        let compressed = encode(&original);
        // Repetitive data should compress
        assert!(compressed.len() < original.len());
        assert_eq!(
            decode(&compressed, None, None, None, None).unwrap(),
            original
        );
    }

    #[test]
    fn test_decode_invalid_data() {
        let result = decode(b"not valid zlib data", None, None, None, None);
        assert!(result.is_err());
    }

    #[test]
    fn test_decode_with_png_sub_predictor() {
        // 2 columns, 1 color, 8bpc -> row_bytes=2
        // Original row: [10, 20]; Sub-encoded: filter_byte=1, [10, 10] (20-10=10)
        let encoded_data = vec![1u8, 10, 10];
        let compressed = encode(&encoded_data);
        let result = decode(&compressed, Some(15), Some(2), Some(1), Some(8)).unwrap();
        assert_eq!(result, vec![10, 20]);
    }

    #[test]
    fn test_decode_with_png_up_predictor() {
        // 2 rows, 2 columns, 1 color, 8bpc
        // Row 0: filter=0 (None), data=[10, 20]
        // Row 1: filter=2 (Up), data=[5, 5] → result [15, 25]
        let encoded_data = vec![0u8, 10, 20, 2, 5, 5];
        let compressed = encode(&encoded_data);
        let result = decode(&compressed, Some(15), Some(2), Some(1), Some(8)).unwrap();
        assert_eq!(result, vec![10, 20, 15, 25]);
    }

    #[test]
    fn test_decode_with_tiff_predictor() {
        // 3 columns, 1 color, 8bpc; TIFF pred 2: [10, 15, 18]
        let encoded_data = vec![10u8, 5, 3];
        let compressed = encode(&encoded_data);
        let result = decode(&compressed, Some(2), Some(3), Some(1), Some(8)).unwrap();
        assert_eq!(result, vec![10, 15, 18]);
    }

    #[test]
    fn test_decode_no_predictor_explicit() {
        let original = b"test data";
        let compressed = encode(original);
        let result = decode(&compressed, Some(1), None, None, None).unwrap();
        assert_eq!(result, original);
    }

    #[test]
    fn test_flate_roundtrip_pseudo_random_data() {
        // Pseudo-random data with low compressibility
        let data: Vec<u8> = (0..1000).map(|i| ((i * 7 + 13) % 256) as u8).collect();
        let compressed = encode(&data);
        let decompressed = decode(&compressed, None, None, None, None).unwrap();
        assert_eq!(decompressed, data);
    }

    #[test]
    fn test_flate_roundtrip_all_same_bytes() {
        // Highly compressible: 10,000 identical bytes
        let data = vec![0xAA; 10_000];
        let compressed = encode(&data);
        assert!(
            compressed.len() < 100,
            "identical bytes should compress well"
        );
        let decompressed = decode(&compressed, None, None, None, None).unwrap();
        assert_eq!(decompressed, data);
    }

    #[test]
    fn test_flate_roundtrip_single_byte() {
        let data = vec![42u8];
        let compressed = encode(&data);
        let decompressed = decode(&compressed, None, None, None, None).unwrap();
        assert_eq!(decompressed, data);
    }

    #[test]
    fn test_flate_roundtrip_alternating_pattern() {
        // Alternating bytes pattern
        let data: Vec<u8> = (0..2000)
            .map(|i| if i % 2 == 0 { 0x00 } else { 0xFF })
            .collect();
        let compressed = encode(&data);
        let decompressed = decode(&compressed, None, None, None, None).unwrap();
        assert_eq!(decompressed, data);
    }
}