pdfrum-page 0.1.0

//! JPEG 2000 decoding, behind a thin entry point.
//!
//! # The colorspace override table
//!
//! A JPX image's colour space comes from **both** the PDF dictionary and the
//! codestream, and reconciling them is a table rather than a rule. The
//! interesting rows are the failures: a `/DeviceGray` image whose codestream
//! says RGB does not fall back to anything — **the whole load fails**. And a
//! three-component space whose codestream has four channels and says sRGB
//! takes the alpha-dropping path, a special case added for PDFs generated by
//! iOS.
//!
//! # `/SMaskInData`
//!
//! Only the value **1** is distinguished. It un-premultiplies the colour
//! against white and captures the alpha into a side buffer that becomes a
//! synthetic grayscale mask. Value 2, nominally "premultiplied", is *not*
//! distinguished and takes the same drop-alpha path as 0.

use crate::color::ColorSpace;
use crate::error::Error;
use crate::image::RequestedSize;
use pdfrum_common::Limits;

/// A decoded JPEG 2000 image.
#[derive(Debug, Clone, PartialEq)]
pub struct JpxImage {
    /// Width in pixels.
    pub width: u32,
    /// Height in pixels.
    pub height: u32,
    /// Components per pixel after the conversion action was applied.
    pub components: u8,
    /// Interleaved eight-bit samples, `width * height * components` bytes.
    pub data: Vec<u8>,
    /// What the conversion table did to the dictionary's colour space.
    pub space_override: SpaceOverride,
    /// The alpha channel, when `/SMaskInData 1` captured one.
    pub alpha: Option<Vec<u8>>,
}

/// What a [`JpxAction`] does to the image dictionary's `/ColorSpace`.
///
/// Three distinct outcomes, which is why this is an enum rather than a
/// nested `Option`: the dictionary's space can be kept, replaced, or
/// **removed**. Removing it is not the same as keeping it — the RGB actions
/// hand back samples that are already device RGB, so keeping the original
/// space would convert them a second time.
#[derive(Debug, Clone, PartialEq)]
pub enum SpaceOverride {
    /// The dictionary's `/ColorSpace` stands.
    Keep,
    /// The dictionary's `/ColorSpace` is dropped: the samples are already in
    /// the device's own space.
    Clear,
    /// The dictionary's `/ColorSpace` is replaced by this one.
    Replace(ColorSpace),
}

/// What the conversion table decided to do with the decoded channels.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JpxAction {
    /// Take one channel as grey.
    UseGray,
    /// Take three channels as RGB.
    UseRgb,
    /// Take four channels as CMYK.
    UseCmyk,
    /// Take the first three of four or more channels, dropping the rest.
    ConvertArgbToRgb,
    /// Take the channels as palette indices.
    UseIndexed,
    /// Leave the channels as they are.
    DoNothing,
    /// Refuse the image outright.
    Fail,
}

/// The codestream's own idea of its colour space, reduced to what the
/// conversion table asks about.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JpxColorSpace {
    /// Greyscale.
    Gray,
    /// sRGB.
    Srgb,
    /// CMYK.
    Cmyk,
    /// The codestream did not say, or said something unrecognised.
    Unspecified,
}

/// Whether a codestream space is the expected one, or simply unstated.
///
/// "Unspecified" counts as a match for every expectation, which is what lets
/// a raw codestream with no colour boxes be used at all.
fn matches_or_unspecified(actual: JpxColorSpace, expected: JpxColorSpace) -> bool {
    actual == expected || actual == JpxColorSpace::Unspecified
}

/// How a decoded codestream's channels are reconciled with the dictionary.
///
/// `space` is the PDF dictionary's colour space, `None` when it stated none.
#[must_use]
pub fn conversion_action(
    space: Option<&ColorSpace>,
    codestream: JpxColorSpace,
    channels: u8,
) -> JpxAction {
    let Some(space) = space else {
        // With no PDF colour space the codestream decides alone.
        return match codestream {
            JpxColorSpace::Unspecified => {
                if channels == 3 {
                    JpxAction::UseRgb
                } else {
                    JpxAction::DoNothing
                }
            }
            JpxColorSpace::Srgb => {
                if channels > 3 {
                    JpxAction::ConvertArgbToRgb
                } else {
                    JpxAction::UseRgb
                }
            }
            JpxColorSpace::Gray => JpxAction::UseGray,
            JpxColorSpace::Cmyk => JpxAction::UseCmyk,
        };
    };
    match space {
        ColorSpace::DeviceGray => {
            if matches_or_unspecified(codestream, JpxColorSpace::Gray) {
                JpxAction::UseGray
            } else {
                // No fallback: the whole load fails.
                JpxAction::Fail
            }
        }
        ColorSpace::DeviceRgb => {
            if !matches_or_unspecified(codestream, JpxColorSpace::Srgb) {
                JpxAction::Fail
            } else if channels > 3 {
                JpxAction::ConvertArgbToRgb
            } else {
                JpxAction::UseRgb
            }
        }
        ColorSpace::DeviceCmyk => {
            if matches_or_unspecified(codestream, JpxColorSpace::Cmyk) {
                JpxAction::UseCmyk
            } else {
                JpxAction::Fail
            }
        }
        ColorSpace::Indexed(_) if space.n_components() == 1 => JpxAction::UseIndexed,
        // The iOS special case: a three-component space whose codestream has
        // a fourth channel and claims sRGB drops the alpha.
        other
            if other.n_components() == 3 && channels == 4 && codestream == JpxColorSpace::Srgb =>
        {
            JpxAction::ConvertArgbToRgb
        }
        _ => JpxAction::DoNothing,
    }
}

impl JpxAction {
    /// What this action does to the dictionary's `/ColorSpace`.
    #[must_use]
    pub fn space_override(self) -> SpaceOverride {
        match self {
            Self::UseGray => SpaceOverride::Replace(ColorSpace::DeviceGray),
            Self::UseCmyk => SpaceOverride::Replace(ColorSpace::DeviceCmyk),
            Self::UseRgb | Self::ConvertArgbToRgb => SpaceOverride::Clear,
            Self::DoNothing | Self::UseIndexed | Self::Fail => SpaceOverride::Keep,
        }
    }

    /// How many components the samples carry after this action.
    #[must_use]
    pub fn components(self, channels: u8) -> u8 {
        match self {
            Self::UseGray | Self::UseIndexed => 1,
            Self::UseRgb | Self::ConvertArgbToRgb => 3,
            Self::UseCmyk => 4,
            Self::DoNothing => channels,
            Self::Fail => 0,
        }
    }
}

/// Whether every component in the codestream agrees on subsampling and depth.
///
/// `CJPX_Decoder::Decode` walks the components and gives up on the first that
/// disagrees with the one before it:
///
/// ```text
/// if (components[i].dx != components[i - 1].dx ||
///     components[i].dy != components[i - 1].dy ||
///     components[i].prec != components[i - 1].prec) {
///   return false;
/// }
/// ```
///
/// The whole image is then refused — `LoadJpxBitmap` returns null and the page
/// draws nothing at all. It is a real gate rather than a paranoia check:
/// `bug_557223` is a 904-byte codestream claiming a 707×6131 three-component
/// image whose components declare subsampling 3×7, 1×7, 1×7 at precisions 1,
/// 2 and 3. PDFium prints "has an empty bitmap" and paints white; a decoder
/// that presses on invents a picture that is not in the file.
///
/// The fields live in the `SIZ` marker segment, which follows `SOC` at the
/// head of the codestream, so this reads them directly rather than going
/// through the decoder — `hayro-jpeg2000` exposes a component's depth only
/// after a decode and its subsampling not at all.
///
/// A codestream this cannot find or parse is left alone: the gate exists to
/// reject a specific disagreement, not to second-guess the decoder.
fn components_agree(data: &[u8]) -> bool {
    // `SOC` immediately followed by `SIZ`, which is the only place the pair
    // may appear: a raw codestream opens with it, and a JP2 file's `jp2c` box
    // contains it.
    let Some(soc) = data.windows(4).position(|w| w == [0xFF, 0x4F, 0xFF, 0x51]) else {
        return true;
    };
    // `Lsiz`(2) `Rsiz`(2) then eight 4-byte grid fields, then `Csiz`(2).
    let header = soc + 4;
    let Some(csiz_at) = header.checked_add(2 + 2 + 32) else {
        return true;
    };
    let Some(count) = data
        .get(csiz_at..)
        .and_then(<[u8]>::first_chunk::<2>)
        .map(|b| usize::from(u16::from_be_bytes(*b)))
    else {
        return true;
    };
    // Three bytes per component: `Ssiz` (depth, biased by one), `XRsiz`,
    // `YRsiz`.
    let first = csiz_at + 2;
    let Some(fields) = count
        .checked_mul(3)
        .and_then(|len| data.get(first..first.checked_add(len)?))
    else {
        return true;
    };
    fields
        .as_chunks::<3>()
        .0
        .iter()
        .all(|c| fields.first_chunk::<3>() == Some(c))
}

/// Decode a JPEG 2000 codestream or JP2 file.
///
/// `space` is the PDF dictionary's colour space, and `smask_in_data` its
/// `/SMaskInData`. `target` is a hint: JPEG 2000 stores a pyramid of
/// resolution levels, so a reduced request decodes fewer packets rather than
/// decoding everything and shrinking it. The returned dimensions are the
/// decoder's answer, which may be larger than the request — it clamps to the
/// levels the codestream carries and refuses to reduce a palettized image at
/// all.
///
/// # Errors
///
/// [`Error::CodecRejected`] when the codestream will not decode, the
/// conversion table refuses the space combination, or the components disagree
/// on subsampling or depth, and [`Error::ImageTooLarge`] when the result
/// exceeds the byte budget.
pub fn decode_jpx(
    data: &[u8],
    space: Option<&ColorSpace>,
    smask_in_data: i64,
    target: RequestedSize,
    limits: &Limits,
) -> Result<JpxImage, Error> {
    if !components_agree(data) {
        return Err(Error::CodecRejected { codec: "JPX" });
    }
    let settings = hayro_jpeg2000::DecodeSettings {
        // An `Indexed` PDF space wants the raw indices, not the palette's
        // colours: the palette lives in the PDF, not the codestream.
        resolve_palette_indices: !matches!(space, Some(ColorSpace::Indexed(_))),
        strict: false,
        target_resolution: match target {
            // No samples never reaches a codec: the build returns first.
            RequestedSize::Full | RequestedSize::NoSamples => None,
            RequestedSize::Reduced { width, height } => {
                // A zero on either axis would make the decoder's own
                // `checked_div` fall through to zero levels; asking for it is
                // meaningless, so it is spelt as no request at all.
                (width != 0 && height != 0).then_some((width, height))
            }
        },
    };
    let image = hayro_jpeg2000::Image::new(data, &settings)
        .map_err(|_| Error::CodecRejected { codec: "JPX" })?;

    let codestream_space = match image.color_space() {
        hayro_jpeg2000::ColorSpace::Gray => JpxColorSpace::Gray,
        hayro_jpeg2000::ColorSpace::RGB => JpxColorSpace::Srgb,
        hayro_jpeg2000::ColorSpace::CMYK => JpxColorSpace::Cmyk,
        // An embedded ICC profile is deliberately discarded, so such an
        // image is treated as having said nothing.
        hayro_jpeg2000::ColorSpace::Icc { .. } | hayro_jpeg2000::ColorSpace::Unknown { .. } => {
            JpxColorSpace::Unspecified
        }
    };
    let channels = image.color_space().num_channels() + u8::from(image.has_alpha());

    let action = conversion_action(space, codestream_space, channels);
    if action == JpxAction::Fail {
        return Err(Error::CodecRejected { codec: "JPX" });
    }

    // The decoder's own dimensions, already reduced if it honoured the hint.
    // Shifting them again here — which is what this did while the hint was
    // never sent — would halve an image that had already been halved, and read
    // the top-left quarter of the samples as if it were the whole picture.
    let width = image.width();
    let height = image.height();
    if width == 0 || height == 0 {
        return Err(Error::CodecRejected { codec: "JPX" });
    }

    let mut context = hayro_jpeg2000::DecoderContext::default();
    let decoded = image
        .decode(&mut context)
        .map_err(|_| Error::CodecRejected { codec: "JPX" })?;
    let samples = decoded.data_u8();
    let source_channels = usize::from(channels).max(1);
    let out_components = action.components(channels);

    let pixels = usize::try_from(width)
        .ok()
        .and_then(|w| w.checked_mul(usize::try_from(height).ok()?))
        .ok_or(Error::ImageTooLarge)?;
    let out_len = pixels
        .checked_mul(usize::from(out_components).max(1))
        .ok_or(Error::ImageTooLarge)?;
    if out_len > limits.max_decoded_stream_len {
        return Err(Error::ImageTooLarge);
    }

    let mut out = vec![0u8; out_len];
    let mut alpha =
        (smask_in_data == 1 && action == JpxAction::ConvertArgbToRgb).then(|| vec![0u8; pixels]);
    let keep = usize::from(out_components).max(1);
    for i in 0..pixels {
        let src = i * source_channels;
        // `/SMaskInData 1` un-premultiplies against white; every other value,
        // including 2, simply drops the extra channel.
        let a = alpha
            .is_some()
            .then(|| samples.get(src + 3).copied().unwrap_or(255));
        for c in 0..keep {
            let v = samples.get(src + c).copied().unwrap_or(0);
            let v = match a {
                Some(a) => {
                    let na = u32::from(255 - a);
                    #[expect(
                        clippy::cast_possible_truncation,
                        reason = "the weighted average of two bytes stays within a byte"
                    )]
                    let blended = ((u32::from(v) * u32::from(a) + 255 * na) / 255) as u8;
                    blended
                }
                None => v,
            };
            if let Some(slot) = out.get_mut(i * keep + c) {
                *slot = v;
            }
        }
        if let (Some(buffer), Some(a)) = (alpha.as_mut(), a)
            && let Some(slot) = buffer.get_mut(i)
        {
            *slot = a;
        }
    }

    Ok(JpxImage {
        width,
        height,
        components: out_components,
        data: out,
        space_override: action.space_override(),
        alpha,
    })
}

#[cfg(test)]
mod tests {
    // Test fixtures quote the oracle's own vectors, compare floats exactly
    // where the behaviour being pinned is exact, and index arrays whose
    // length the fixture itself fixes.
    #![allow(
        clippy::unreadable_literal,
        clippy::float_cmp,
        clippy::indexing_slicing,
        clippy::cast_precision_loss,
        clippy::cast_possible_truncation,
        reason = "test fixtures quote oracle vectors verbatim and compare exactly"
    )]

    use super::{
        JpxAction, JpxColorSpace, RequestedSize, SpaceOverride, components_agree,
        conversion_action, decode_jpx,
    };
    use crate::color::{ColorSpace, Indexed};
    use pdfrum_common::Limits;

    #[test]
    fn device_gray_needs_a_grey_or_silent_codestream() {
        let gray = ColorSpace::DeviceGray;
        assert_eq!(
            conversion_action(Some(&gray), JpxColorSpace::Gray, 1),
            JpxAction::UseGray
        );
        assert_eq!(
            conversion_action(Some(&gray), JpxColorSpace::Unspecified, 1),
            JpxAction::UseGray
        );
        // An RGB codestream under a gray dictionary fails outright.
        assert_eq!(
            conversion_action(Some(&gray), JpxColorSpace::Srgb, 3),
            JpxAction::Fail
        );
    }

    #[test]
    fn device_rgb_drops_a_fourth_channel() {
        let rgb = ColorSpace::DeviceRgb;
        assert_eq!(
            conversion_action(Some(&rgb), JpxColorSpace::Srgb, 3),
            JpxAction::UseRgb
        );
        assert_eq!(
            conversion_action(Some(&rgb), JpxColorSpace::Srgb, 4),
            JpxAction::ConvertArgbToRgb
        );
        assert_eq!(
            conversion_action(Some(&rgb), JpxColorSpace::Cmyk, 4),
            JpxAction::Fail
        );
    }

    #[test]
    fn device_cmyk_needs_a_cmyk_or_silent_codestream() {
        let cmyk = ColorSpace::DeviceCmyk;
        assert_eq!(
            conversion_action(Some(&cmyk), JpxColorSpace::Cmyk, 4),
            JpxAction::UseCmyk
        );
        assert_eq!(
            conversion_action(Some(&cmyk), JpxColorSpace::Gray, 1),
            JpxAction::Fail
        );
    }

    #[test]
    fn an_indexed_space_takes_the_raw_indices() {
        let indexed = ColorSpace::Indexed(Box::new(Indexed {
            base: Box::new(ColorSpace::DeviceRgb),
            max_index: 3,
            lookup: Box::from(&[0u8; 12][..]),
            component_ranges: Box::from(&[(0.0f32, 1.0f32); 3][..]),
        }));
        assert_eq!(
            conversion_action(Some(&indexed), JpxColorSpace::Srgb, 1),
            JpxAction::UseIndexed
        );
    }

    #[test]
    fn the_ios_special_case_drops_alpha_for_any_three_component_space() {
        // A `CalRGB` space, four channels, sRGB codestream.
        let cal = ColorSpace::CalRgb(Box::new(crate::color::CalRgb {
            white_point: [0.9505, 1.0, 1.089],
            black_point: [0.0; 3],
            gamma: None,
            matrix: None,
        }));
        assert_eq!(
            conversion_action(Some(&cal), JpxColorSpace::Srgb, 4),
            JpxAction::ConvertArgbToRgb
        );
        // Three channels instead takes the do-nothing path.
        assert_eq!(
            conversion_action(Some(&cal), JpxColorSpace::Srgb, 3),
            JpxAction::DoNothing
        );
    }

    #[test]
    fn without_a_pdf_space_the_codestream_decides_alone() {
        assert_eq!(
            conversion_action(None, JpxColorSpace::Unspecified, 3),
            JpxAction::UseRgb
        );
        assert_eq!(
            conversion_action(None, JpxColorSpace::Unspecified, 2),
            JpxAction::DoNothing
        );
        assert_eq!(
            conversion_action(None, JpxColorSpace::Srgb, 4),
            JpxAction::ConvertArgbToRgb
        );
        assert_eq!(
            conversion_action(None, JpxColorSpace::Gray, 1),
            JpxAction::UseGray
        );
        assert_eq!(
            conversion_action(None, JpxColorSpace::Cmyk, 4),
            JpxAction::UseCmyk
        );
    }

    #[test]
    fn the_rgb_actions_reset_the_space_rather_than_replacing_it() {
        assert_eq!(JpxAction::UseRgb.space_override(), SpaceOverride::Clear);
        assert_eq!(
            JpxAction::ConvertArgbToRgb.space_override(),
            SpaceOverride::Clear
        );
        assert_eq!(
            JpxAction::UseGray.space_override(),
            SpaceOverride::Replace(ColorSpace::DeviceGray)
        );
        assert_eq!(JpxAction::UseIndexed.space_override(), SpaceOverride::Keep);
        assert_eq!(JpxAction::DoNothing.space_override(), SpaceOverride::Keep);
    }

    /// A `SOC`+`SIZ` head with `count` components described by `fields`.
    fn siz(count: u16, fields: &[[u8; 3]]) -> Vec<u8> {
        let mut out = vec![0xFF, 0x4F, 0xFF, 0x51];
        // `Lsiz`, `Rsiz`, then the eight grid words the parser skips.
        out.extend_from_slice(&[0, 47, 0, 0]);
        out.extend_from_slice(&[0u8; 32]);
        out.extend_from_slice(&count.to_be_bytes());
        for f in fields {
            out.extend_from_slice(f);
        }
        out
    }

    #[test]
    fn components_that_disagree_on_subsampling_or_depth_are_refused() {
        // `bug_557223`'s shape: three components at 3x7, 1x7, 1x7 and
        // precisions 1, 2, 3. `Ssiz` is the depth biased by one.
        let bad = siz(3, &[[0, 3, 7], [1, 1, 7], [2, 1, 7]]);
        assert!(!components_agree(&bad));
        // Subsampling alone is enough.
        assert!(!components_agree(&siz(2, &[[7, 1, 1], [7, 2, 1]])));
        // So is depth alone.
        assert!(!components_agree(&siz(2, &[[7, 1, 1], [6, 1, 1]])));
        // Agreement passes, at any component count including one and zero.
        assert!(components_agree(&siz(
            3,
            &[[7, 1, 1], [7, 1, 1], [7, 1, 1]]
        )));
        assert!(components_agree(&siz(1, &[[7, 2, 2]])));
        assert!(components_agree(&siz(0, &[])));
    }

    #[test]
    fn a_codestream_the_gate_cannot_read_is_left_to_the_decoder() {
        // No `SOC`+`SIZ` pair, a truncated header, or a component table that
        // runs off the end: the gate declines to judge rather than rejecting.
        assert!(components_agree(b""));
        assert!(components_agree(b"not a codestream at all"));
        assert!(components_agree(&[0xFF, 0x4F, 0xFF, 0x51]));
        let mut short = siz(4, &[[7, 1, 1]]);
        short.truncate(short.len() - 1);
        assert!(components_agree(&short));
    }

    #[test]
    fn garbage_is_rejected_rather_than_panicked_on() {
        let limits = Limits::default();
        for data in [&b""[..], b"\x00\x00", b"not jpeg2000", &[0xFFu8; 32]] {
            assert!(decode_jpx(data, None, 0, RequestedSize::Full, &limits).is_err());
        }
    }
}