rpdfium_codec/

streaming.rs

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

//! Streaming (incremental) decode API for large image decompression.
//!
//! The [`StreamingDecoder`] trait allows feeding input data in chunks
//! and reading decompressed output incrementally, which is useful for
//! large inline images and progressive rendering.

use crate::error::DecodeError;

/// A streaming (incremental) decoder for PDF stream data.
///
/// Callers feed input chunks via [`feed()`](StreamingDecoder::feed),
/// then read available output via [`read_output()`](StreamingDecoder::read_output).
/// Call [`finish()`](StreamingDecoder::finish) after all input is provided.
pub trait StreamingDecoder {
    /// Feed a chunk of compressed input data.
    fn feed(&mut self, input: &[u8]) -> Result<(), DecodeError>;

    /// Signal that all input has been provided and perform final flush.
    fn finish(&mut self) -> Result<(), DecodeError>;

    /// Returns the number of decompressed bytes available to read.
    fn output_available(&self) -> usize;

    /// Read decompressed output into the provided buffer.
    /// Returns the number of bytes actually written to `buf`.
    fn read_output(&mut self, buf: &mut [u8]) -> usize;
}

/// Types of filters that support streaming decode.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamingDecoderType {
    /// FlateDecode (zlib/deflate) — true incremental streaming.
    Flate,
    /// DCTDecode (JPEG) — accumulates input, decodes on finish().
    Dct,
}

/// Create a streaming decoder for the given filter type.
pub fn create_streaming_decoder(filter: StreamingDecoderType) -> Box<dyn StreamingDecoder> {
    match filter {
        StreamingDecoderType::Flate => Box::new(FlateStreamingDecoder::new()),
        StreamingDecoderType::Dct => Box::new(DctStreamingDecoder::new()),
    }
}

// ---------------------------------------------------------------------------
// FlateStreamingDecoder — true incremental zlib decompression
// ---------------------------------------------------------------------------

/// Streaming decoder for FlateDecode (zlib/deflate).
///
/// Decompresses data incrementally as chunks are fed in via [`feed()`](StreamingDecoder::feed).
pub struct FlateStreamingDecoder {
    decompressor: flate2::Decompress,
    output_buffer: Vec<u8>,
    finished: bool,
}

impl Default for FlateStreamingDecoder {
    fn default() -> Self {
        Self::new()
    }
}

impl FlateStreamingDecoder {
    /// Create a new streaming Flate decoder.
    pub fn new() -> Self {
        Self {
            decompressor: flate2::Decompress::new(true), // zlib header
            output_buffer: Vec::new(),
            finished: false,
        }
    }
}

impl StreamingDecoder for FlateStreamingDecoder {
    fn feed(&mut self, input: &[u8]) -> Result<(), DecodeError> {
        if self.finished {
            return Err(DecodeError::InvalidInput("decoder already finished".into()));
        }
        let mut buf = [0u8; 32 * 1024];
        let mut offset = 0;
        while offset < input.len() {
            let before_in = self.decompressor.total_in();
            let before_out = self.decompressor.total_out();
            let status = self
                .decompressor
                .decompress(&input[offset..], &mut buf, flate2::FlushDecompress::None)
                .map_err(|e| DecodeError::Flate(e.to_string()))?;
            let consumed = (self.decompressor.total_in() - before_in) as usize;
            let produced = (self.decompressor.total_out() - before_out) as usize;
            self.output_buffer.extend_from_slice(&buf[..produced]);
            offset += consumed;
            if matches!(status, flate2::Status::StreamEnd) {
                self.finished = true;
                break;
            }
            if consumed == 0 && produced == 0 {
                break;
            }
        }
        Ok(())
    }

    fn finish(&mut self) -> Result<(), DecodeError> {
        if !self.finished {
            let mut buf = [0u8; 4096];
            let before_out = self.decompressor.total_out();
            let _ = self
                .decompressor
                .decompress(&[], &mut buf, flate2::FlushDecompress::Finish);
            let produced = (self.decompressor.total_out() - before_out) as usize;
            self.output_buffer.extend_from_slice(&buf[..produced]);
            self.finished = true;
        }
        Ok(())
    }

    fn output_available(&self) -> usize {
        self.output_buffer.len()
    }

    fn read_output(&mut self, buf: &mut [u8]) -> usize {
        let len = buf.len().min(self.output_buffer.len());
        buf[..len].copy_from_slice(&self.output_buffer[..len]);
        self.output_buffer.drain(..len);
        len
    }
}

// ---------------------------------------------------------------------------
// DctStreamingDecoder — accumulates input, decodes on finish()
// ---------------------------------------------------------------------------

/// Streaming decoder for DCTDecode (JPEG).
///
/// Since JPEG requires the full input to decode, this accumulates bytes
/// and performs the actual decode in [`finish()`](StreamingDecoder::finish).
pub struct DctStreamingDecoder {
    accumulated: Vec<u8>,
    output_buffer: Vec<u8>,
    finished: bool,
}

impl Default for DctStreamingDecoder {
    fn default() -> Self {
        Self::new()
    }
}

impl DctStreamingDecoder {
    /// Create a new streaming DCT decoder.
    pub fn new() -> Self {
        Self {
            accumulated: Vec::new(),
            output_buffer: Vec::new(),
            finished: false,
        }
    }
}

impl StreamingDecoder for DctStreamingDecoder {
    fn feed(&mut self, input: &[u8]) -> Result<(), DecodeError> {
        if self.finished {
            return Err(DecodeError::InvalidInput("decoder already finished".into()));
        }
        self.accumulated.extend_from_slice(input);
        Ok(())
    }

    fn finish(&mut self) -> Result<(), DecodeError> {
        if !self.finished {
            self.output_buffer = crate::jpeg::decode(&self.accumulated)?;
            self.accumulated = Vec::new();
            self.finished = true;
        }
        Ok(())
    }

    fn output_available(&self) -> usize {
        self.output_buffer.len()
    }

    fn read_output(&mut self, buf: &mut [u8]) -> usize {
        let len = buf.len().min(self.output_buffer.len());
        buf[..len].copy_from_slice(&self.output_buffer[..len]);
        self.output_buffer.drain(..len);
        len
    }
}

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

    fn zlib_compress(data: &[u8]) -> Vec<u8> {
        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
        encoder.write_all(data).unwrap();
        encoder.finish().unwrap()
    }

    #[test]
    fn test_flate_streaming_multi_chunk() {
        let original = b"Hello, World!";
        let compressed = zlib_compress(original);
        let mut decoder = FlateStreamingDecoder::new();

        // Feed in 3 chunks
        let chunk_size = compressed.len().div_ceil(3);
        for chunk in compressed.chunks(chunk_size) {
            decoder.feed(chunk).unwrap();
        }
        decoder.finish().unwrap();

        let mut output = vec![0u8; 256];
        let n = decoder.read_output(&mut output);
        assert_eq!(&output[..n], original);
        assert_eq!(decoder.output_available(), 0);
    }

    #[test]
    fn test_flate_streaming_single_byte() {
        let original = b"Streaming byte by byte test";
        let compressed = zlib_compress(original);
        let mut decoder = FlateStreamingDecoder::new();

        // Feed one byte at a time
        for &byte in &compressed {
            decoder.feed(&[byte]).unwrap();
        }
        decoder.finish().unwrap();

        let mut output = vec![0u8; 256];
        let n = decoder.read_output(&mut output);
        assert_eq!(&output[..n], original);
    }

    #[test]
    fn test_flate_streaming_matches_batch() {
        let original: Vec<u8> = (0..10_000).map(|i| (i % 256) as u8).collect();
        let compressed = zlib_compress(&original);

        // Batch decode
        let batch_result = crate::flate::decode(&compressed, None, None, None, None).unwrap();

        // Streaming decode
        let mut decoder = FlateStreamingDecoder::new();
        decoder.feed(&compressed).unwrap();
        decoder.finish().unwrap();

        let mut streaming_result = vec![0u8; decoder.output_available()];
        let n = decoder.read_output(&mut streaming_result);
        streaming_result.truncate(n);

        assert_eq!(streaming_result, batch_result);
        assert_eq!(streaming_result, original);
    }

    #[test]
    fn test_dct_streaming_accumulate() {
        let mut decoder = DctStreamingDecoder::new();
        // Feed some invalid JPEG data in 2 chunks
        decoder.feed(&[0xFF, 0xD8]).unwrap();
        decoder.feed(&[0xFF, 0xE0]).unwrap();
        // finish() should fail because the data is not a valid JPEG
        let result = decoder.finish();
        assert!(result.is_err());
    }

    #[test]
    fn test_output_available_accuracy() {
        let original = b"Test output availability tracking";
        let compressed = zlib_compress(original);
        let mut decoder = FlateStreamingDecoder::new();

        decoder.feed(&compressed).unwrap();
        decoder.finish().unwrap();

        let available = decoder.output_available();
        assert_eq!(available, original.len());

        let mut output = vec![0u8; available];
        let n = decoder.read_output(&mut output);
        assert_eq!(n, available);
        assert_eq!(decoder.output_available(), 0);
        assert_eq!(&output, original);
    }

    #[test]
    fn test_read_output_small_buffer() {
        let original = b"Drain me in small reads";
        let compressed = zlib_compress(original);
        let mut decoder = FlateStreamingDecoder::new();

        decoder.feed(&compressed).unwrap();
        decoder.finish().unwrap();

        let mut collected = Vec::new();
        let mut buf = [0u8; 4];
        loop {
            let n = decoder.read_output(&mut buf);
            if n == 0 {
                break;
            }
            collected.extend_from_slice(&buf[..n]);
        }
        assert_eq!(collected, original);
        assert_eq!(decoder.output_available(), 0);
    }

    #[test]
    fn test_feed_after_finish_errors() {
        let compressed = zlib_compress(b"data");
        let mut decoder = FlateStreamingDecoder::new();
        decoder.feed(&compressed).unwrap();
        decoder.finish().unwrap();

        let result = decoder.feed(b"more data");
        assert!(result.is_err());
    }

    #[test]
    fn test_dct_feed_after_finish_errors() {
        let mut decoder = DctStreamingDecoder::new();
        // finish with empty data — dct::decode on empty input will error
        let _ = decoder.finish();
        // Even after finish, feed should error because finished is set
        // (finish sets finished=true even if decode fails, but let's
        // use a successful-ish flow)
        let mut decoder2 = DctStreamingDecoder::new();
        decoder2.accumulated = Vec::new();
        // Force finish state
        decoder2.finished = true;
        let result = decoder2.feed(b"more");
        assert!(result.is_err());
    }

    #[test]
    fn test_create_streaming_decoder_flate() {
        let mut decoder = create_streaming_decoder(StreamingDecoderType::Flate);
        let compressed = zlib_compress(b"factory test");
        decoder.feed(&compressed).unwrap();
        decoder.finish().unwrap();
        let mut out = vec![0u8; 64];
        let n = decoder.read_output(&mut out);
        assert_eq!(&out[..n], b"factory test");
    }

    #[test]
    fn test_create_streaming_decoder_dct() {
        let mut decoder = create_streaming_decoder(StreamingDecoderType::Dct);
        decoder.feed(&[0xFF, 0xD8]).unwrap();
        // Invalid JPEG, finish should error
        assert!(decoder.finish().is_err());
    }
}