use crate::color::{ColorSpace, Family};
use crate::error::Error;
use crate::image::dict::is_allowed_bits_per_component;
use zune_jpeg::zune_core::bytestream::ZCursor;
use zune_jpeg::zune_core::colorspace::ColorSpace as ZColorSpace;
use zune_jpeg::zune_core::options::DecoderOptions;
const VALID_COMPONENTS: [u8; 3] = [1, 3, 4];
#[derive(Debug, Clone, PartialEq)]
pub struct DctImage {
pub width: u32,
pub height: u32,
pub components: u8,
pub bpc: u32,
pub data: Vec<u8>,
}
#[must_use]
#[allow(
dead_code,
reason = "unwired: no decode path asks for reduced resolution yet"
)]
pub fn scale_denominator(levels: u8) -> u32 {
1u32 << levels.min(3)
}
#[must_use]
#[allow(
dead_code,
reason = "unwired: no decode path asks for reduced resolution yet"
)]
pub fn scaled_size(dimension: u32, denominator: u32) -> u32 {
if denominator == 0 {
return dimension;
}
dimension.div_ceil(denominator)
}
#[must_use]
#[allow(
dead_code,
reason = "unwired: no decode path asks for reduced resolution yet"
)]
pub fn allows_reduced_resolution(width: u32, height: u32, max_h: u32, max_v: u32) -> bool {
let h_mcu = max_h.saturating_mul(8);
let v_mcu = max_v.saturating_mul(8);
h_mcu != 0 && v_mcu != 0 && width.is_multiple_of(h_mcu) && height.is_multiple_of(v_mcu)
}
#[must_use]
pub fn component_mismatch_allowed(space: Option<&ColorSpace>, components: u8) -> bool {
let Some(space) = space else {
return true;
};
let cs_components = u8::try_from(space.n_components()).unwrap_or(u8::MAX);
match space.family() {
Family::DeviceGray | Family::DeviceRgb | Family::DeviceCmyk => {
let min = match space.family() {
Family::DeviceGray => 1,
Family::DeviceRgb => 3,
_ => 4,
};
cs_components >= min && components >= min
}
Family::Lab => components == 3 && cs_components >= 3,
Family::IccBased => {
crate::color::is_valid_icc_components(i64::from(components))
&& crate::color::is_valid_icc_components(i64::from(cs_components))
&& cs_components >= components
}
_ => cs_components == components,
}
}
fn output_space(channels: u8, input: Option<ZColorSpace>) -> Option<ZColorSpace> {
if channels != 4 {
return None;
}
Some(match input {
Some(ZColorSpace::YCCK) => ZColorSpace::YCCK,
_ => ZColorSpace::CMYK,
})
}
const SCALEBITS: i32 = 16;
const ONE_HALF: i32 = 1 << (SCALEBITS - 1);
const CR_R: i32 = 91_881;
const CB_B: i32 = 116_130;
const CR_G: i32 = 46_802;
const CB_G: i32 = 22_554;
const MAX_SAMPLE: i32 = 255;
const CENTER_SAMPLE: i32 = 128;
fn ycck_to_cmyk(samples: &mut [u8]) {
for [c, m, y_channel, _k] in samples.as_chunks_mut::<4>().0 {
let (y, cb, cr) = (
i32::from(*c),
i32::from(*m) - CENTER_SAMPLE,
i32::from(*y_channel) - CENTER_SAMPLE,
);
let red = y + ((CR_R * cr + ONE_HALF) >> SCALEBITS);
let green = y + ((-CB_G * cb - CR_G * cr + ONE_HALF) >> SCALEBITS);
let blue = y + ((CB_B * cb + ONE_HALF) >> SCALEBITS);
*c = clamp_sample(MAX_SAMPLE - red);
*m = clamp_sample(MAX_SAMPLE - green);
*y_channel = clamp_sample(MAX_SAMPLE - blue);
}
}
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
fn clamp_sample(value: i32) -> u8 {
value.clamp(0, MAX_SAMPLE) as u8
}
const KNOWN_BAD_HEIGHT_OFFSETS: [usize; 2] = [94, 163];
const SOF_MARKER_BACK_OFFSET: usize = 5;
fn has_known_bad_height(data: &[u8], offset: usize, declared: (u32, u32)) -> bool {
let (width, height) = declared;
if width == 0 || width > JPEG_MAX_DIMENSION || height == 0 || height > JPEG_MAX_DIMENSION {
return false;
}
let Some(marker_at) = offset.checked_sub(SOF_MARKER_BACK_OFFSET) else {
return false;
};
if !matches!(
data.get(marker_at..marker_at + 2),
Some(&[0xff, sof]) if (0xc0..=0xcf).contains(&sof)
) {
return false;
}
let expected = big_endian(width);
matches!(
data.get(offset..offset + 4),
Some(&[0xff, 0xff, high, low]) if [high, low] == expected
)
}
const JPEG_MAX_DIMENSION: u32 = 65500;
#[allow(clippy::cast_possible_truncation)]
fn big_endian(dimension: u32) -> [u8; 2] {
[((dimension >> 8) & 0xff) as u8, (dimension & 0xff) as u8]
}
pub fn decode_dct(data: &[u8], declared: (u32, u32)) -> Result<DctImage, Error> {
let patched: Option<Vec<u8>> = probe(data)
.is_none()
.then(|| {
KNOWN_BAD_HEIGHT_OFFSETS
.into_iter()
.find(|&offset| has_known_bad_height(data, offset, declared))
.map(|offset| {
let mut copy = data.to_vec();
if let Some(slot) = copy.get_mut(offset..offset + 2) {
slot.copy_from_slice(&big_endian(declared.1));
}
copy
})
})
.flatten();
let data = patched.as_deref().unwrap_or(data);
let (channels, input) = read_header(data).map_or((0, None), |(_, _, c, space)| (c, space));
let pinned = output_space(channels, input);
let options = pinned.map(|space| DecoderOptions::default().jpeg_set_out_colorspace(space));
let mut decoder = match options {
Some(options) => zune_jpeg::JpegDecoder::new_with_options(ZCursor::new(data), options),
None => zune_jpeg::JpegDecoder::new(ZCursor::new(data)),
};
let mut pixels = decoder
.decode()
.map_err(|_| Error::CodecRejected { codec: "DCT" })?;
let info = decoder
.info()
.ok_or(Error::CodecRejected { codec: "DCT" })?;
if pinned == Some(ZColorSpace::YCCK) {
ycck_to_cmyk(&mut pixels);
}
let components = u8::try_from(
pixels.len() / usize::from(info.width).max(1) / usize::from(info.height).max(1),
)
.unwrap_or(0);
if !VALID_COMPONENTS.contains(&components) || !is_allowed_bits_per_component(8) {
return Err(Error::CodecRejected { codec: "DCT" });
}
Ok(DctImage {
width: u32::from(info.width),
height: u32::from(info.height),
components,
bpc: 8,
data: pixels,
})
}
fn read_header(data: &[u8]) -> Option<(u32, u32, u8, Option<ZColorSpace>)> {
let mut decoder = zune_jpeg::JpegDecoder::new(ZCursor::new(data));
decoder.decode_headers().ok()?;
let info = decoder.info()?;
Some((
u32::from(info.width),
u32::from(info.height),
info.components,
decoder.input_colorspace(),
))
}
#[must_use]
pub fn probe(data: &[u8]) -> Option<(u32, u32, u8)> {
read_header(data).map(|(width, height, components, _)| (width, height, components))
}
#[cfg(test)]
mod tests {
#![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::{
ZColorSpace, allows_reduced_resolution, component_mismatch_allowed, decode_dct,
has_known_bad_height, output_space, probe, scale_denominator, scaled_size, ycck_to_cmyk,
};
use crate::color::{ColorSpace, Indexed};
#[test]
fn the_scale_denominator_caps_at_one_eighth() {
assert_eq!(scale_denominator(0), 1);
assert_eq!(scale_denominator(1), 2);
assert_eq!(scale_denominator(3), 8);
assert_eq!(scale_denominator(4), 8);
assert_eq!(scale_denominator(255), 8);
}
#[test]
fn scaling_rounds_up_unlike_the_jpx_path() {
assert_eq!(scaled_size(400, 8), 50);
assert_eq!(scaled_size(401, 8), 51);
assert_eq!(scaled_size(400, 1), 400);
assert_eq!(scaled_size(400, 0), 400);
}
#[test]
fn non_mcu_aligned_images_refuse_reduced_resolution() {
assert!(allows_reduced_resolution(400, 400, 2, 2));
assert!(!allows_reduced_resolution(408, 408, 2, 2));
assert!(allows_reduced_resolution(408, 408, 1, 1));
assert!(!allows_reduced_resolution(0, 0, 0, 0));
}
#[test]
fn device_families_need_only_enough_components() {
assert!(component_mismatch_allowed(Some(&ColorSpace::DeviceRgb), 3));
assert!(component_mismatch_allowed(Some(&ColorSpace::DeviceRgb), 4));
assert!(!component_mismatch_allowed(Some(&ColorSpace::DeviceRgb), 1));
assert!(component_mismatch_allowed(Some(&ColorSpace::DeviceGray), 1));
}
#[test]
fn other_families_need_an_exact_match() {
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!(component_mismatch_allowed(Some(&indexed), 1));
assert!(!component_mismatch_allowed(Some(&indexed), 3));
}
#[test]
fn lab_needs_exactly_three_components() {
let lab = ColorSpace::Lab(Box::new(crate::color::Lab {
white_point: [0.9505, 1.0, 1.089],
black_point: [0.0; 3],
ranges: [-100.0, 100.0, -100.0, 100.0],
}));
assert!(component_mismatch_allowed(Some(&lab), 3));
assert!(!component_mismatch_allowed(Some(&lab), 1));
assert!(!component_mismatch_allowed(Some(&lab), 4));
}
#[test]
fn with_no_colour_space_any_component_count_is_accepted() {
assert!(component_mismatch_allowed(None, 1));
assert!(component_mismatch_allowed(None, 4));
}
#[test]
fn garbage_is_rejected_rather_than_panicked_on() {
for data in [&b""[..], b"\xFF\xD8", b"not a jpeg", &[0u8; 64]] {
assert!(
decode_dct(data, (0, 0)).is_err(),
"{data:?} should be rejected"
);
assert!(probe(data).is_none());
}
}
fn known_bad_header() -> Vec<u8> {
let mut data = vec![0u8; 200];
data[158] = 0xff;
data[159] = 0xc2;
data[163] = 0xff;
data[164] = 0xff;
data[165] = 0x02;
data[166] = 0x64;
data
}
#[test]
fn a_sof_height_of_ffff_beside_the_declared_width_is_the_known_bad_header() {
let data = known_bad_header();
assert!(has_known_bad_height(&data, 163, (612, 792)));
assert!(!has_known_bad_height(&data, 94, (612, 792)));
}
#[test]
fn the_repair_declines_every_way_the_evidence_can_fall_short() {
let data = known_bad_header();
assert!(!has_known_bad_height(&data, 163, (613, 792)));
assert!(!has_known_bad_height(&data, 163, (612, 0)));
assert!(!has_known_bad_height(&data, 163, (612, 65501)));
let mut not_sof = data.clone();
not_sof[159] = 0xd8;
assert!(!has_known_bad_height(¬_sof, 163, (612, 792)));
let mut sane = data;
sane[163] = 0x03;
sane[164] = 0x18;
assert!(!has_known_bad_height(&sane, 163, (612, 792)));
assert!(!has_known_bad_height(&[0xff, 0xc2], 163, (612, 792)));
}
#[test]
fn a_four_channel_jpeg_stays_four_channels() {
assert_eq!(
output_space(4, Some(ZColorSpace::CMYK)),
Some(ZColorSpace::CMYK)
);
assert_eq!(output_space(1, Some(ZColorSpace::Luma)), None);
assert_eq!(output_space(3, Some(ZColorSpace::YCbCr)), None);
assert_eq!(output_space(2, None), None);
assert_eq!(output_space(0, None), None);
}
#[test]
fn a_ycck_jpeg_is_asked_for_its_own_space_not_cmyk() {
assert_eq!(
output_space(4, Some(ZColorSpace::YCCK)),
Some(ZColorSpace::YCCK)
);
assert_eq!(output_space(4, None), Some(ZColorSpace::CMYK));
}
#[test]
fn ycck_leaves_a_neutral_pixel_black_and_passes_k_through() {
let mut samples = [255, 128, 128, 42];
ycck_to_cmyk(&mut samples);
assert_eq!(samples, [0, 0, 0, 42]);
let mut dark = [0, 128, 128, 7];
ycck_to_cmyk(&mut dark);
assert_eq!(dark, [255, 255, 255, 7]);
}
#[test]
fn ycck_matches_libjpegs_fixed_point_arithmetic_exactly() {
let reference = |y: i32, cb: i32, cr: i32| -> [i32; 3] {
let (x_b, x_r) = (cb - 128, cr - 128);
let red_from_cr = (91881 * x_r + 32768) >> 16;
let blue_from_cb = (116130 * x_b + 32768) >> 16;
let green_chroma_red_term = -46802 * x_r;
let green_chroma_blue_term = -22554 * x_b + 32768;
[
255 - (y + red_from_cr),
255 - (y + ((green_chroma_blue_term + green_chroma_red_term) >> 16)),
255 - (y + blue_from_cb),
]
};
for y in (0u8..=255).step_by(17) {
for cb in (0u8..=255).step_by(23) {
for cr in (0u8..=255).step_by(29) {
let mut got = [y, cb, cr, 200];
ycck_to_cmyk(&mut got);
let want = reference(i32::from(y), i32::from(cb), i32::from(cr));
for (channel, expected) in want.into_iter().enumerate() {
assert_eq!(
i32::from(got[channel]),
expected.clamp(0, 255),
"channel {channel} at y={y} cb={cb} cr={cr}"
);
}
assert_eq!(got[3], 200, "K must pass through");
}
}
}
}
#[test]
fn a_ycck_codestream_decodes_to_four_channels_rather_than_being_dropped() {
let Some(data) = ycck_codestream() else {
return;
};
let image = decode_dct(&data, (429, 542)).expect("a YCCK JPEG must decode");
assert_eq!((image.width, image.height), (429, 542));
assert_eq!(image.components, 4);
assert_eq!(image.data.len(), 429 * 542 * 4);
let mut expected = [242, 111, 130, 0];
ycck_to_cmyk(&mut expected);
assert_eq!(&image.data[..4], &expected);
assert_ne!(&image.data[..3], &[242, 111, 130]);
}
fn oracle_checkout() -> std::path::PathBuf {
std::env::var_os("PDFRUM_ORACLE_CHECKOUT").map_or_else(
|| std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../../pdfium-c++"),
std::path::PathBuf::from,
)
}
fn ycck_codestream() -> Option<Vec<u8>> {
let pdf = std::fs::read(oracle_checkout().join("testing/corpus/fx/other/1.pdf")).ok()?;
let start = pdf
.windows(2)
.enumerate()
.filter(|(_, pair)| *pair == b"\xff\xd8")
.map(|(at, _)| at)
.next_back()?;
let end = pdf
.windows(9)
.enumerate()
.find(|(at, window)| *at > start && *window == b"endstream")
.map(|(at, _)| at)?;
Some(pdf.get(start..end)?.to_vec())
}
}