rpdfium_codec/jpeg/

jpegmodule.rs

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

//! DCTDecode filter — JPEG decompression using zune-jpeg.

use crate::error::DecodeError;

/// Maximum number of bytes to scan forward looking for SOI marker.
const SOI_SCAN_LIMIT: usize = 1024;

/// Byte offset of the dimension fields relative to the SOF marker start.
///
/// SOF marker layout: FF Cx (2 bytes marker) + length (2 bytes) + bpc (1 byte)
/// → dimensions start at offset 5 from the marker.
const SOF_MARKER_BYTE_OFFSET: usize = 5;

/// Known byte offsets (from SOI) where SOF markers appear in malformed JPEG
/// headers that set height to 0xFFFF.
///
/// Ported from PDFium's `kKnownBadHeaderWithInvalidHeightByteOffsetStarts`.
const KNOWN_BAD_HEADER_OFFSETS: &[usize] = &[94, 163];

/// Information extracted from JPEG markers before decoding.
#[derive(Debug, Clone, Default)]
pub struct JpegInfo {
    /// Image width in pixels.
    pub width: u32,
    /// Image height in pixels.
    pub height: u32,
    /// Number of color components (e.g. 1 for grayscale, 3 for RGB/YCbCr, 4 for CMYK).
    pub num_components: u8,
    /// Bits per component (typically 8).
    pub bits_per_component: u8,
    /// Whether an Adobe APP14 marker was found.
    pub has_adobe_marker: bool,
    /// Adobe color transform flag (0=Unknown, 1=YCbCr, 2=YCCK).
    pub adobe_color_transform: u8,
}

/// Probe JPEG data to extract image information without fully decoding.
///
/// Scans for SOF0/SOF1/SOF2 markers to extract dimensions and component info,
/// and for Adobe APP14 markers to detect color transform flags.
/// Returns `None` if no SOF marker is found.
pub fn probe_jpeg_info(input: &[u8]) -> Option<JpegInfo> {
    let data = find_soi(input)?;
    let mut info = JpegInfo::default();
    let mut found_sof = false;
    let mut pos = 2; // Skip SOI

    while pos + 1 < data.len() {
        if data[pos] != 0xFF {
            pos += 1;
            continue;
        }

        // Skip padding 0xFF bytes
        while pos + 1 < data.len() && data[pos + 1] == 0xFF {
            pos += 1;
        }
        if pos + 1 >= data.len() {
            break;
        }

        let marker = data[pos + 1];
        pos += 2;

        // SOS (Start of Scan) — stop scanning markers
        if marker == 0xDA {
            break;
        }

        // Markers without length
        if marker == 0x00 || marker == 0x01 || (0xD0..=0xD7).contains(&marker) {
            continue;
        }

        // Read marker length
        if pos + 2 > data.len() {
            break;
        }
        let length = u16::from_be_bytes([data[pos], data[pos + 1]]) as usize;
        if length < 2 || pos + length > data.len() {
            break;
        }

        match marker {
            // SOF0, SOF1, SOF2 (Baseline, Extended, Progressive)
            0xC0..=0xC2 => {
                if length >= 8 {
                    info.bits_per_component = data[pos + 2];
                    info.height = u16::from_be_bytes([data[pos + 3], data[pos + 4]]) as u32;
                    info.width = u16::from_be_bytes([data[pos + 5], data[pos + 6]]) as u32;
                    info.num_components = data[pos + 7];
                    found_sof = true;
                }
            }
            // APP14 (Adobe marker)
            0xEE => {
                // Adobe marker: length >= 14 (2 length + 5 "Adobe" + 7 data)
                if length >= 14 && pos + 2 + 5 <= data.len() && &data[pos + 2..pos + 7] == b"Adobe"
                {
                    info.has_adobe_marker = true;
                    // Color transform byte is at offset 11 from marker data start
                    // (2 length bytes + 5 "Adobe" + 2 version + 2 flags0)
                    if pos + 13 < data.len() {
                        info.adobe_color_transform = data[pos + 13];
                    }
                }
            }
            _ => {}
        }

        pos += length;
    }

    if found_sof { Some(info) } else { None }
}

/// Upstream-aligned alias for [`probe_jpeg_info()`](probe_jpeg_info). Corresponds to `JpegModule::LoadInfo()`.
#[inline]
pub fn load_info(input: &[u8]) -> Option<JpegInfo> {
    probe_jpeg_info(input)
}

/// Find the SOI (Start of Image) marker in the input.
///
/// If input starts with 0xFF 0xD8, returns it directly.
/// Otherwise scans forward up to `SOI_SCAN_LIMIT` bytes.
fn find_soi(input: &[u8]) -> Option<&[u8]> {
    if input.len() < 2 {
        return None;
    }

    if input[0] == 0xFF && input[1] == 0xD8 {
        return Some(input);
    }

    // Scan forward for SOI marker
    let limit = input.len().min(SOI_SCAN_LIMIT);
    for i in 0..limit.saturating_sub(1) {
        if input[i] == 0xFF && input[i + 1] == 0xD8 {
            return Some(&input[i..]);
        }
    }

    None
}

/// Check whether a SOF marker exists at the given byte offset within JPEG data.
///
/// SOF markers are 0xFF followed by 0xC0–0xC2 (baseline/extended/progressive).
fn is_sof_at_offset(data: &[u8], dimension_offset: usize) -> bool {
    if dimension_offset < SOF_MARKER_BYTE_OFFSET {
        return false;
    }
    let marker_offset = dimension_offset - SOF_MARKER_BYTE_OFFSET;
    if marker_offset + 1 >= data.len() {
        return false;
    }
    data[marker_offset] == 0xFF && (0xC0..=0xC2).contains(&data[marker_offset + 1])
}

/// Check if JPEG data has a known bad header with height set to 0xFFFF
/// at a specific dimension offset.
///
/// Matches upstream PDFium's `HasKnownBadHeaderWithInvalidHeight()`.
fn has_invalid_height_at_offset(data: &[u8], dimension_offset: usize, expected_width: u32) -> bool {
    if dimension_offset + 3 >= data.len() {
        return false;
    }
    if !is_sof_at_offset(data, dimension_offset) {
        return false;
    }

    let width_hi = ((expected_width >> 8) & 0xFF) as u8;
    let width_lo = (expected_width & 0xFF) as u8;

    // Height high byte, height low byte, width high byte, width low byte
    data[dimension_offset] == 0xFF
        && data[dimension_offset + 1] == 0xFF
        && data[dimension_offset + 2] == width_hi
        && data[dimension_offset + 3] == width_lo
}

/// Detect and attempt to fix JPEG data with a known bad header where height is 0xFFFF.
///
/// Some malformed JPEG encoders set the image height to 0xFFFF in the SOF marker.
/// This function checks known byte offsets for this pattern and, if found, patches
/// the height to 0 (which per JPEG spec means "determine from data/DNL marker").
///
/// Returns `Some(patched_data)` if the known bad pattern was found and patched,
/// `None` if the data doesn't match the known bad pattern.
pub fn try_patch_invalid_height(input: &[u8]) -> Option<Vec<u8>> {
    let data = find_soi(input)?;
    let soi_offset = input.len() - data.len();

    let info = probe_jpeg_info(data)?;
    if info.height != 0xFFFF {
        return None;
    }

    for &offset in KNOWN_BAD_HEADER_OFFSETS {
        if has_invalid_height_at_offset(data, offset, info.width as u32) {
            let mut patched = input.to_vec();
            let abs_offset = soi_offset + offset;
            if abs_offset + 1 < patched.len() {
                // Patch height to 0 (JPEG spec: height determined from data)
                patched[abs_offset] = 0x00;
                patched[abs_offset + 1] = 0x00;
                return Some(patched);
            }
        }
    }

    None
}

/// Decode JPEG data with automatic recovery for known bad headers.
///
/// If normal decoding fails and the JPEG has height=0xFFFF at a known offset,
/// this patches the height and retries. Ported from upstream PDFium's
/// malformed JPEG handling logic.
pub fn decode_with_recovery(input: &[u8]) -> Result<Vec<u8>, DecodeError> {
    match decode(input) {
        Ok(data) => Ok(data),
        Err(original_err) => {
            if let Some(patched) = try_patch_invalid_height(input) {
                decode(&patched)
            } else {
                Err(original_err)
            }
        }
    }
}

/// Decode JPEG (DCT) compressed data.
///
/// Returns raw pixel data (RGB or grayscale depending on the JPEG).
/// If the input doesn't start with an SOI marker, scans forward up to
/// 1024 bytes to find one.
pub fn decode(input: &[u8]) -> Result<Vec<u8>, DecodeError> {
    let data = find_soi(input).unwrap_or(input);
    let cursor = std::io::Cursor::new(data);
    let mut decoder = zune_jpeg::JpegDecoder::new(cursor);
    decoder
        .decode()
        .map_err(|e| DecodeError::Jpeg(format!("{e}")))
}

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

    #[test]
    fn test_decode_invalid_jpeg() {
        let result = decode(b"not a jpeg");
        assert!(result.is_err());
    }

    #[test]
    fn test_decode_empty_input() {
        let result = decode(b"");
        assert!(result.is_err());
    }

    #[test]
    fn test_decode_truncated_header() {
        // Start of JPEG marker but truncated
        let result = decode(&[0xFF, 0xD8, 0xFF]);
        assert!(result.is_err());
    }

    #[test]
    fn test_decode_minimal_valid_jpeg() {
        // A minimal 1x1 white JPEG (generated programmatically is complex,
        // so we just test that valid-looking but incomplete data fails gracefully)
        let partial_jpeg = [
            0xFF, 0xD8, // SOI
            0xFF, 0xE0, // APP0
            0x00, 0x10, // length
            b'J', b'F', b'I', b'F', 0x00, // JFIF identifier
            0x01, 0x01, // version
            0x00, // units
            0x00, 0x01, // X density
            0x00, 0x01, // Y density
            0x00, 0x00, // thumbnail
        ];
        let result = decode(&partial_jpeg);
        // This is incomplete so it should error
        assert!(result.is_err());
    }

    /// Build a minimal JPEG-like byte sequence with SOF0 marker for probe tests.
    fn build_jpeg_with_sof(width: u16, height: u16, components: u8, bpc: u8) -> Vec<u8> {
        let mut data = Vec::new();
        // SOI
        data.extend_from_slice(&[0xFF, 0xD8]);
        // SOF0
        data.extend_from_slice(&[0xFF, 0xC0]);
        // Length = 8 + 3 * components
        let sof_len = 8 + 3 * components as u16;
        data.extend_from_slice(&sof_len.to_be_bytes());
        data.push(bpc);
        data.extend_from_slice(&height.to_be_bytes());
        data.extend_from_slice(&width.to_be_bytes());
        data.push(components);
        // Component specs (3 bytes each): id, sampling, qtable
        for i in 0..components {
            data.extend_from_slice(&[i + 1, 0x11, 0]);
        }
        data
    }

    #[test]
    fn test_probe_jpeg_info_basic() {
        let jpeg = build_jpeg_with_sof(640, 480, 3, 8);
        let info = probe_jpeg_info(&jpeg).unwrap();
        assert_eq!(info.width, 640);
        assert_eq!(info.height, 480);
        assert_eq!(info.num_components, 3);
        assert_eq!(info.bits_per_component, 8);
        assert!(!info.has_adobe_marker);
    }

    #[test]
    fn test_probe_jpeg_info_grayscale() {
        let jpeg = build_jpeg_with_sof(100, 200, 1, 8);
        let info = probe_jpeg_info(&jpeg).unwrap();
        assert_eq!(info.width, 100);
        assert_eq!(info.height, 200);
        assert_eq!(info.num_components, 1);
    }

    #[test]
    fn test_probe_jpeg_info_with_adobe_marker() {
        let mut jpeg = build_jpeg_with_sof(320, 240, 4, 8);
        // Append APP14 Adobe marker
        jpeg.extend_from_slice(&[0xFF, 0xEE]);
        // Length = 14 (2 length + 5 Adobe + 2 version + 2 flags0 + 2 flags1 + 1 transform)
        jpeg.extend_from_slice(&[0x00, 0x0E]);
        // "Adobe"
        jpeg.extend_from_slice(b"Adobe");
        // Version (2 bytes)
        jpeg.extend_from_slice(&[0x00, 0x64]);
        // Flags0 (2 bytes)
        jpeg.extend_from_slice(&[0x00, 0x00]);
        // Flags1 (2 bytes)
        jpeg.extend_from_slice(&[0x00, 0x00]);
        // Color transform: 2 (YCCK)
        jpeg.push(0x02);

        let info = probe_jpeg_info(&jpeg).unwrap();
        assert_eq!(info.width, 320);
        assert_eq!(info.num_components, 4);
        assert!(info.has_adobe_marker);
        assert_eq!(info.adobe_color_transform, 2);
    }

    #[test]
    fn test_probe_jpeg_info_no_sof_returns_none() {
        // Just an SOI with no SOF
        let data = [0xFF, 0xD8];
        assert!(probe_jpeg_info(&data).is_none());
    }

    #[test]
    fn test_probe_jpeg_info_empty_returns_none() {
        assert!(probe_jpeg_info(&[]).is_none());
    }

    #[test]
    fn test_find_soi_at_start() {
        let data = [0xFF, 0xD8, 0xFF, 0xC0];
        let result = find_soi(&data);
        assert!(result.is_some());
        assert_eq!(result.unwrap().len(), 4);
    }

    #[test]
    fn test_find_soi_with_offset() {
        // Garbage bytes followed by SOI
        let mut data = vec![0x00, 0x00, 0x00, 0x42];
        data.extend_from_slice(&[0xFF, 0xD8, 0xFF, 0xC0]);
        let result = find_soi(&data);
        assert!(result.is_some());
        assert_eq!(result.unwrap().len(), 4); // From SOI to end
    }

    #[test]
    fn test_find_soi_not_found() {
        let data = [0x00; 100];
        assert!(find_soi(&data).is_none());
    }

    #[test]
    fn test_find_soi_beyond_limit() {
        // SOI marker at position > SOI_SCAN_LIMIT should not be found
        let mut data = vec![0x00; SOI_SCAN_LIMIT + 10];
        data[SOI_SCAN_LIMIT + 5] = 0xFF;
        data[SOI_SCAN_LIMIT + 6] = 0xD8;
        assert!(find_soi(&data).is_none());
    }

    #[test]
    fn test_decode_with_leading_garbage_still_fails_gracefully() {
        // SOI is there but JPEG is incomplete
        let mut data = vec![0x00, 0x00, 0x00];
        data.extend_from_slice(&[0xFF, 0xD8, 0xFF]);
        let result = decode(&data);
        assert!(result.is_err());
    }

    // -----------------------------------------------------------------------
    // Known-bad-header tests
    // -----------------------------------------------------------------------

    /// Build JPEG data with SOF0 at a specific byte offset from SOI, with given height.
    fn build_jpeg_with_sof_at_offset(
        offset: usize,
        width: u16,
        height: u16,
        components: u8,
    ) -> Vec<u8> {
        let mut data = Vec::new();
        // SOI
        data.extend_from_slice(&[0xFF, 0xD8]);

        // Pad with APP0 markers to reach the desired offset for SOF
        // SOF marker position = offset - SOF_MARKER_BYTE_OFFSET (which is 5)
        let sof_marker_pos = offset - SOF_MARKER_BYTE_OFFSET;
        let current_pos = 2; // after SOI
        if sof_marker_pos > current_pos {
            // Fill with an APP0 segment of appropriate size
            let fill_len = sof_marker_pos - current_pos;
            if fill_len >= 4 {
                data.extend_from_slice(&[0xFF, 0xE0]); // APP0 marker
                let seg_len = (fill_len - 2) as u16;
                data.extend_from_slice(&seg_len.to_be_bytes());
                data.extend(vec![0x00; fill_len - 4]);
            } else {
                data.extend(vec![0x00; fill_len]);
            }
        }

        // SOF0 marker
        data.extend_from_slice(&[0xFF, 0xC0]);
        let sof_len = 8 + 3 * components as u16;
        data.extend_from_slice(&sof_len.to_be_bytes());
        data.push(8); // bpc
        data.extend_from_slice(&height.to_be_bytes());
        data.extend_from_slice(&width.to_be_bytes());
        data.push(components);
        for i in 0..components {
            data.extend_from_slice(&[i + 1, 0x11, 0]);
        }
        data
    }

    #[test]
    fn test_detect_known_bad_header_at_offset_94() {
        let data = build_jpeg_with_sof_at_offset(94, 640, 0xFFFF, 3);
        let info = probe_jpeg_info(&data).unwrap();
        assert_eq!(info.height, 0xFFFF);
        assert_eq!(info.width, 640);

        let patched = try_patch_invalid_height(&data);
        assert!(patched.is_some());

        // Verify height was patched to 0
        let patched = patched.unwrap();
        let patched_info = probe_jpeg_info(&patched).unwrap();
        assert_eq!(patched_info.height, 0);
        assert_eq!(patched_info.width, 640);
    }

    #[test]
    fn test_detect_known_bad_header_at_offset_163() {
        let data = build_jpeg_with_sof_at_offset(163, 320, 0xFFFF, 3);
        let info = probe_jpeg_info(&data).unwrap();
        assert_eq!(info.height, 0xFFFF);

        let patched = try_patch_invalid_height(&data);
        assert!(patched.is_some());
        let patched_info = probe_jpeg_info(&patched.unwrap()).unwrap();
        assert_eq!(patched_info.height, 0);
    }

    #[test]
    fn test_no_patch_for_normal_height() {
        let data = build_jpeg_with_sof(640, 480, 3, 8);
        assert!(try_patch_invalid_height(&data).is_none());
    }

    #[test]
    fn test_no_patch_for_unknown_offset() {
        // SOF at offset 50 (not a known bad offset)
        let data = build_jpeg_with_sof_at_offset(50, 640, 0xFFFF, 3);
        let info = probe_jpeg_info(&data).unwrap();
        assert_eq!(info.height, 0xFFFF);
        // Offset 50 is not in KNOWN_BAD_HEADER_OFFSETS
        assert!(try_patch_invalid_height(&data).is_none());
    }

    #[test]
    fn test_decode_with_recovery_on_invalid_input() {
        // Recovery still fails if data is not valid JPEG
        let result = decode_with_recovery(b"not a jpeg");
        assert!(result.is_err());
    }

    #[test]
    fn test_is_sof_at_offset_basic() {
        // SOF0 at position 0: FF C0
        let data = [0xFF, 0xC0, 0x00, 0x08, 0x08, 0x00, 0x01, 0x00, 0x01, 0x01];
        assert!(is_sof_at_offset(&data, SOF_MARKER_BYTE_OFFSET));
    }

    #[test]
    fn test_is_sof_at_offset_sof2() {
        // SOF2 (progressive)
        let data = [0xFF, 0xC2, 0x00, 0x08, 0x08, 0x00, 0x01, 0x00, 0x01, 0x01];
        assert!(is_sof_at_offset(&data, SOF_MARKER_BYTE_OFFSET));
    }

    #[test]
    fn test_is_sof_at_offset_not_sof() {
        // Not a SOF marker
        let data = [0xFF, 0xE0, 0x00, 0x08, 0x08, 0x00, 0x01, 0x00, 0x01, 0x01];
        assert!(!is_sof_at_offset(&data, SOF_MARKER_BYTE_OFFSET));
    }
}