stegoeggo 0.2.3

//! JPEG Header Parsing
//!
//! Parses JPEG file headers to extract quantization tables, Huffman tables,
//! and other metadata needed for transcoding.
#![allow(dead_code)] // JPEG spec reference types (color spaces, coding processes, lookup methods)

use super::{Result, TranscoderError};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JpegCodingProcess {
    SequentialDCT,
    ProgressiveDCT,
    Lossless,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[allow(clippy::upper_case_acronyms)]
pub enum JpegColorSpace {
    Grayscale,
    YCbCr,
    RGB,
    CMYK,
    YCCK,
}

#[derive(Debug, Clone, Copy)]
pub struct QuantizationTable {
    pub table_id: u8,
    pub precision: u8, // 0 = 8-bit, 1 = 16-bit
    pub values: [u16; 64],
}

impl QuantizationTable {
    pub fn get(&self, index: usize) -> u16 {
        if index < 64 {
            self.values[index]
        } else {
            0
        }
    }

    pub fn scaled(&self, scale: f32) -> [u16; 64] {
        let mut result = [0u16; 64];
        for (i, &val) in self.values.iter().enumerate() {
            result[i] = ((val as f32 * scale).round() as u16).max(1);
        }
        result
    }
}

#[derive(Debug, Clone)]
pub struct HuffmanTable {
    pub table_class: u8, // 0 = DC, 1 = AC
    pub table_id: u8,
    pub counts: [u16; 16],
    pub values: Vec<u8>,
}

impl HuffmanTable {
    pub fn is_dc(&self) -> bool {
        self.table_class == 0
    }

    pub fn is_ac(&self) -> bool {
        self.table_class == 1
    }
}

#[derive(Debug, Clone, Copy)]
pub struct ScanComponent {
    pub component_id: u8,
    pub h_sampling: u8,
    pub v_sampling: u8,
    pub quant_table_id: u8,
    pub dc_table_id: u8,
    pub ac_table_id: u8,
}

#[derive(Debug, Clone)]
pub struct JpegHeader {
    pub width: u16,
    pub height: u16,
    pub precision: u8,
    pub coding_process: JpegCodingProcess,
    pub color_space: JpegColorSpace,

    pub quantization_tables: [Option<QuantizationTable>; 4],
    pub huffman_tables_dc: Vec<Option<HuffmanTable>>,
    pub huffman_tables_ac: Vec<Option<HuffmanTable>>,

    pub components: Vec<ScanComponent>,

    pub app0_marker: Option<Vec<u8>>,
    pub app1_markers: Vec<Vec<u8>>,
    pub com_markers: Vec<Vec<u8>>,

    pub restart_interval: u16,

    pub is_progressive: bool,
}

impl Default for JpegHeader {
    fn default() -> Self {
        Self {
            width: 0,
            height: 0,
            precision: 8,
            coding_process: JpegCodingProcess::SequentialDCT,
            color_space: JpegColorSpace::YCbCr,
            quantization_tables: [None; 4],
            huffman_tables_dc: vec![None, None, None, None],
            huffman_tables_ac: vec![None, None, None, None],
            components: Vec::new(),
            app0_marker: None,
            app1_markers: Vec::new(),
            com_markers: Vec::new(),
            restart_interval: 0,
            is_progressive: false,
        }
    }
}

impl JpegHeader {
    pub fn parse(data: &[u8]) -> Result<Self> {
        Self::parse_with_limits(data, &crate::ResourceLimits::default())
    }

    pub fn parse_with_limits(data: &[u8], limits: &crate::ResourceLimits) -> Result<Self> {
        if data.len() < 2 {
            return Err(TranscoderError::InvalidFormat("Input too short".into()));
        }

        if data[0] != 0xFF || data[1] != 0xD8 {
            return Err(TranscoderError::InvalidFormat("No SOI marker found".into()));
        }

        if data.len() <= 2 {
            return Err(TranscoderError::InvalidFormat(
                "Truncated JPEG header".into(),
            ));
        }

        let mut header = JpegHeader::default();

        // Parse the outer JPEG stream sequentially from the leading SOI.
        // This avoids false positives from marker-like byte sequences inside
        // metadata payloads such as COM, APP1, or APP13 segments.
        let mut pos = 2;
        let mut largest_width = 0usize;
        let mut largest_height = 0usize;
        let mut segment_count: usize = 0;

        while pos < data.len() {
            // Find next marker - skip to 0xFF
            while pos < data.len() && data[pos] != 0xFF {
                pos += 1;
            }

            if pos >= data.len() {
                break;
            }

            if pos + 1 >= data.len() {
                return Err(TranscoderError::InvalidFormat(format!(
                    "Truncated marker at byte offset {}",
                    pos
                )));
            }

            let marker = data[pos + 1];

            // Handle stuffed zeros (0xFF 0x00)
            if marker == 0x00 {
                pos += 2;
                continue;
            }

            // RST markers (RST0-RST7: 0xFF 0xD0-0xFF 0xD7)
            if (0xD0..=0xD7).contains(&marker) {
                pos += 2;
                continue;
            }

            // Skip standalone 0xFF that's not a marker (shouldn't happen in valid JPEG, but handle it)
            if marker == 0xFF {
                pos += 1;
                continue;
            }

            //EOI - End of Image
            if marker == 0xD9 {
                break;
            }

            segment_count += 1;
            if segment_count > limits.max_jpeg_segments() {
                return Err(TranscoderError::InvalidFormat(format!(
                    "JPEG segment count {} exceeds limit {}",
                    segment_count,
                    limits.max_jpeg_segments()
                )));
            }

            // Get segment length
            if pos + 3 >= data.len() {
                return Err(TranscoderError::InvalidFormat(format!(
                    "Truncated segment length at byte offset {}",
                    pos
                )));
            }

            let segment_len = ((data[pos + 2] as usize) << 8) | (data[pos + 3] as usize);

            if segment_len > limits.max_jpeg_segment_bytes() {
                return Err(TranscoderError::InvalidFormat(format!(
                    "JPEG segment size {} exceeds limit {}",
                    segment_len,
                    limits.max_jpeg_segment_bytes()
                )));
            }

            let segment_data_start = pos + 4;
            let segment_data_end = (pos + 2 + segment_len)
                .min(data.len())
                .max(segment_data_start);

            let segment_data = &data[segment_data_start..segment_data_end];

            match marker {
                // APP0 (JFIF)
                0xE0 => {
                    header.app0_marker = Some(segment_data.to_vec());
                }
                // APP1 (EXIF, etc)
                0xE1 => {
                    header.app1_markers.push(segment_data.to_vec());
                }
                // DQT - Define Quantization Table
                0xDB => {
                    header.parse_dqt(segment_data)?;
                }
                // SOF0 - Start of Frame (baseline)
                0xC0 => {
                    // Parse all SOFs - we'll fix dimensions later
                    header.parse_sof(segment_data)?;
                    if (header.width as usize) * (header.height as usize)
                        > largest_width * largest_height
                    {
                        largest_width = header.width as usize;
                        largest_height = header.height as usize;
                    }
                }
                // SOF1 - Start of Frame (extended)
                0xC1 => {
                    header.parse_sof(segment_data)?;
                    if (header.width as usize) * (header.height as usize)
                        > largest_width * largest_height
                    {
                        largest_width = header.width as usize;
                        largest_height = header.height as usize;
                    }
                }
                // SOF2 - Start of Frame (progressive)
                0xC2 => {
                    header.parse_sof(segment_data)?;
                    header.is_progressive = true;
                    header.coding_process = JpegCodingProcess::ProgressiveDCT;
                    if (header.width as usize) * (header.height as usize)
                        > largest_width * largest_height
                    {
                        largest_width = header.width as usize;
                        largest_height = header.height as usize;
                    }
                }
                // DHT - Define Huffman Table
                0xC4 => {
                    header.parse_dht(segment_data)?;
                }
                // SOS - Start of Scan
                0xDA => {
                    header.parse_sos(segment_data);
                    // Stop parsing header - rest is scan data
                    break;
                }
                // COM - Comment
                0xFE => {
                    header.com_markers.push(segment_data.to_vec());
                }
                0xDD if segment_data.len() >= 2 => {
                    header.restart_interval =
                        ((segment_data[0] as u16) << 8) | (segment_data[1] as u16);
                }
                _ => {
                    // Unknown marker - skip
                }
            }

            pos += 2 + segment_len;
        }

        // Fix dimensions to use the largest SOF
        if largest_width > 0
            && (header.width as usize, header.height as usize) != (largest_width, largest_height)
        {
            header.width = largest_width as u16;
            header.height = largest_height as u16;
        }

        Ok(header)
    }

    fn parse_dqt(&mut self, data: &[u8]) -> Result<()> {
        let mut pos = 0;
        while pos + 64 < data.len() {
            let table_info = data[pos];
            let table_id = table_info & 0x0F;
            if table_id >= 4 {
                return Err(TranscoderError::InvalidFormat(format!(
                    "DQT table_id {} out of range (0-3)",
                    table_id
                )));
            }
            let precision = if (table_info & 0xF0) != 0 { 16 } else { 8 };

            let mut values = [0u16; 64];

            if precision == 8 {
                for i in 0..64 {
                    values[i] = data[pos + 1 + i] as u16;
                }
                pos += 65;
            } else {
                // 16-bit precision needs 128 bytes of table data
                if pos + 128 >= data.len() {
                    return Err(TranscoderError::InvalidFormat(
                        "Truncated 16-bit DQT segment".into(),
                    ));
                }
                for i in 0..64 {
                    values[i] =
                        ((data[pos + 1 + i * 2] as u16) << 8) | (data[pos + 2 + i * 2] as u16);
                }
                pos += 129;
            }

            let table = QuantizationTable {
                table_id,
                precision,
                values,
            };

            self.quantization_tables[table_id as usize] = Some(table);
        }
        Ok(())
    }

    fn parse_sof(&mut self, data: &[u8]) -> Result<()> {
        if data.len() < 6 {
            return Err(TranscoderError::InvalidFormat(
                "SOF segment too short".into(),
            ));
        }

        self.precision = data[0];
        self.height = ((data[1] as u16) << 8) | (data[2] as u16);
        self.width = ((data[3] as u16) << 8) | (data[4] as u16);

        let num_components = data[5] as usize;

        if data.len() < 6 + num_components * 3 {
            return Err(TranscoderError::InvalidFormat(
                "SOF segment too short for components".into(),
            ));
        }

        self.components.clear();
        for i in 0..num_components {
            let offset = 6 + i * 3;
            let component_id = data[offset];
            let sampling = data[offset + 1];
            let quant_table_id = data[offset + 2];

            let h_sampling = (sampling >> 4) & 0x0F;
            let v_sampling = sampling & 0x0F;

            if h_sampling == 0 || v_sampling == 0 {
                return Err(TranscoderError::InvalidFormat(format!(
                    "Component {} has zero sampling factor (h={}, v={}); \
                     JPEG spec requires 1-4",
                    component_id, h_sampling, v_sampling
                )));
            }
            if h_sampling > 4 || v_sampling > 4 {
                return Err(TranscoderError::InvalidFormat(format!(
                    "Component {} sampling factor exceeds JPEG maximum (h={}, v={})",
                    component_id, h_sampling, v_sampling
                )));
            }

            self.components.push(ScanComponent {
                component_id,
                h_sampling,
                v_sampling,
                quant_table_id,
                dc_table_id: 0,
                ac_table_id: 0,
            });
        }

        Ok(())
    }

    fn parse_dht(&mut self, data: &[u8]) -> Result<()> {
        let mut pos = 0;
        while pos + 17 < data.len() {
            let table_info = data[pos];
            let table_class = (table_info >> 4) & 0x0F;
            let table_id = table_info & 0x0F;

            if table_id >= 4 {
                return Err(TranscoderError::InvalidFormat(format!(
                    "DHT table_id {} out of range (0-3)",
                    table_id
                )));
            }

            let mut counts = [0u16; 16];
            let mut total = 0u16;
            for i in 0..16 {
                counts[i] = data[pos + 1 + i] as u16;
                total += counts[i];
            }

            let values_start = pos + 17;
            let values_end = values_start + total as usize;
            if values_end > data.len() {
                return Err(TranscoderError::InvalidFormat(
                    "Truncated DHT segment: not enough value bytes".into(),
                ));
            }
            let values = data[values_start..values_end].to_vec();

            let table = HuffmanTable {
                table_class,
                table_id,
                counts,
                values,
            };

            if table_class == 0 {
                self.huffman_tables_dc[table_id as usize] = Some(table);
            } else {
                self.huffman_tables_ac[table_id as usize] = Some(table);
            }

            pos = values_end;
        }

        Ok(())
    }

    fn parse_sos(&mut self, data: &[u8]) {
        if data.len() < 3 {
            return;
        }

        let num_components = data[0] as usize;
        if data.len() < 4 + num_components * 2 {
            return;
        }

        for i in 0..num_components {
            let component_id = data[1 + i * 2];
            let table_info = data[2 + i * 2];

            // Update components with their Huffman table assignments.
            // The 4-bit table IDs are clamped to 0-3 because the entropy
            // decoder/encoder arrays are sized for the JPEG spec maximum of
            // 4 Huffman tables per class. Malformed SOS segments with IDs > 3
            // would otherwise cause out-of-bounds array indexing.
            if let Some(comp) = self
                .components
                .iter_mut()
                .find(|c| c.component_id == component_id)
            {
                comp.dc_table_id = ((table_info >> 4) & 0x0F).min(3);
                comp.ac_table_id = (table_info & 0x0F).min(3);
            }
        }
    }

    pub fn get_quantization_table(&self, id: u8) -> Option<&QuantizationTable> {
        self.quantization_tables
            .get(id as usize)
            .and_then(|t| t.as_ref())
    }

    pub fn get_dc_huffman_table(&self, id: u8) -> Option<&HuffmanTable> {
        self.huffman_tables_dc
            .get(id as usize)
            .and_then(|t| t.as_ref())
    }

    pub fn get_ac_huffman_table(&self, id: u8) -> Option<&HuffmanTable> {
        self.huffman_tables_ac
            .get(id as usize)
            .and_then(|t| t.as_ref())
    }
}

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

    #[test]
    fn parse_empty_input_returns_error() {
        let result = JpegHeader::parse(&[]);
        assert!(result.is_err());
    }

    #[test]
    fn parse_single_byte_returns_error() {
        let result = JpegHeader::parse(&[0xFF]);
        assert!(result.is_err());
    }

    #[test]
    fn parse_tiny_data_returns_error() {
        let result = JpegHeader::parse(&[0xFF, 0xD8]);
        assert!(result.is_err());
    }

    #[test]
    fn com_markers_are_preserved_through_parse() {
        let mut data = vec![0xFF, 0xD8];

        let com_payload = b"Test comment";
        let com_len = (com_payload.len() + 2) as u16;
        data.extend_from_slice(&[0xFF, 0xFE]);
        data.extend_from_slice(&com_len.to_be_bytes());
        data.extend_from_slice(com_payload);

        data.extend_from_slice(&[0xFF, 0xD9]);

        let header = JpegHeader::parse(&data).unwrap();
        assert_eq!(header.com_markers.len(), 1);
        assert_eq!(header.com_markers[0], com_payload);
    }

    #[test]
    fn parse_ignores_marker_like_bytes_inside_metadata_payloads() {
        use crate::protected::metadata_trap::MetadataTrapProtector;
        use crate::types::{ImageOutputFormat, ProtectionContext};
        use image::DynamicImage;

        let img = DynamicImage::ImageRgb8(image::ImageBuffer::from_fn(32, 32, |x, y| {
            image::Rgb([(x * 3) as u8, (y * 5) as u8, ((x + y) * 7) as u8])
        }));

        let jpeg = crate::util::image::encode_image(&img, image::ImageFormat::Jpeg).unwrap();
        let ctx = ProtectionContext::new(0.5, 42).with_format(ImageOutputFormat::Jpeg);
        let injected = MetadataTrapProtector::new()
            .inject_bytes(&jpeg, &ctx)
            .unwrap();

        let header = JpegHeader::parse(&injected).unwrap();
        assert!(header.width > 0);
        assert!(header.height > 0);
        assert!(!header.quantization_tables.iter().all(|t| t.is_none()));
    }

    /// Regression test for a divide-by-zero panic discovered by the fuzz harness
    /// in `fuzz/fuzz_targets/pipeline_bytes.rs`. A SOF segment with a component
    /// whose sampling-factor nibble is 0 must be rejected at parse time, not
    /// panic later in the entropy decoder.
    #[test]
    fn parse_rejects_zero_sampling_factor() {
        // SOF0 + 1 component with h_sampling=0 (high nibble)
        let data: &[u8] = &[
            0xFF, 0xD8, // SOI
            0xFF, 0xC0, // SOF0
            0x00, 0x11, // segment length 17
            0x08, // precision 8
            0x00, 0x10, // height 16
            0x00, 0x10, // width 16
            0x01, // 1 component
            0x01, 0x00, 0x00, // id=1, sampling=0x00 (h=0, v=0), quant=0
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        ];
        let result = JpegHeader::parse(data);
        assert!(
            result.is_err(),
            "SOF with zero sampling factor must be rejected"
        );
    }

    /// Regression test for SOS table IDs > 3 causing OOB panic in entropy
    /// decoder. The 4-bit field can encode 0-15 but only 0-3 are valid per
    /// the JPEG spec. Clamping prevents the panic.
    #[test]
    fn parse_sos_clamps_table_ids_to_valid_range() {
        let data: &[u8] = &[
            0xFF, 0xD8, // SOI
            0xFF, 0xC0, // SOF0
            0x00, 0x11, // segment length 17
            0x08, // precision 8
            0x00, 0x10, // height 16
            0x00, 0x10, // width 16
            0x01, // 1 component
            0x01, 0x11, 0x00, // id=1, sampling=0x11 (h=1, v=1), quant=0
            // DHT - define DC table 0
            0xFF, 0xC4, 0x00, 0x1B, 0x00, // length=27, class=DC, id=0
            0x00, 0x01, 0x05, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, // counts
            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, // values
            // DHT - define AC table 0
            0xFF, 0xC4, 0x00, 0xB5, 0x10, // length=181, class=AC, id=0
            0x00, 0x02, 0x01, 0x03, 0x03, 0x02, 0x04, 0x03, 0x05, 0x05, 0x04, 0x04, 0x00, 0x00,
            0x01, 0x7D, // counts
            0x01, 0x02, 0x03, 0x00, 0x04, 0x11, 0x05, 0x12, 0x21, 0x31, 0x41, 0x06, 0x13, 0x51,
            0x61, 0x07, 0x22, 0x71, 0x14, 0x32, 0x81, 0x91, 0xA1, 0x08, 0x23, 0x42, 0xB1, 0xC1,
            0x15, 0x52, 0xD1, 0xF0, 0x24, 0x33, 0x62, 0x72, 0x82, 0x09, 0x0A, 0x16, 0x17, 0x18,
            0x19, 0x1A, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2A, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39,
            0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49, 0x4A, 0x53, 0x54, 0x55, 0x56, 0x57,
            0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6A, 0x73, 0x74, 0x75,
            0x76, 0x77, 0x78, 0x79, 0x7A, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8A, 0x92,
            0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6, 0xA7,
            0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3,
            0xC4, 0xC5, 0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8,
            0xD9, 0xDA, 0xE1, 0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF1, 0xF2,
            0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8, 0xF9, 0xFA, // values
            // SOS with table IDs = 0x55 (dc=5, ac=5) - should be clamped to 3
            0xFF, 0xDA, 0x00, 0x08, // SOS, length=8
            0x01, // 1 component
            0x01, 0x55, // component_id=1, table_info=0x55 (dc=5, ac=5)
            0x00, 0x3F, 0x00, // spectral selection, approx
        ];
        let header = JpegHeader::parse(data).unwrap();
        assert_eq!(header.components.len(), 1);
        assert_eq!(
            header.components[0].dc_table_id, 3,
            "dc_table_id should be clamped from 5 to 3"
        );
        assert_eq!(
            header.components[0].ac_table_id, 3,
            "ac_table_id should be clamped from 5 to 3"
        );
    }

    #[test]
    fn parse_rejects_oversized_sampling_factor() {
        // SOF0 + 1 component with h_sampling=5 (above JPEG max of 4)
        let data: &[u8] = &[
            0xFF, 0xD8, 0xFF, 0xC0, 0x00, 0x11, 0x08, 0x00, 0x10, 0x00, 0x10, 0x01, 0x01, 0x50,
            0x00, // id=1, sampling=0x50 (h=5, v=0)
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        ];
        let result = JpegHeader::parse(data);
        assert!(
            result.is_err(),
            "SOF with sampling factor > 4 must be rejected"
        );
    }
}