1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

use byteorder::{BigEndian, ReadBytesExt};
use colors::{BitDepth, ColorType};
use crc::crc32;
use error::PngError;
use std::collections::HashSet;
use std::io::Cursor;
use PngResult;

#[derive(Debug, Clone, Copy)]
/// Headers from the IHDR chunk of the image
pub struct IhdrData {
    /// The width of the image in pixels
    pub width: u32,
    /// The height of the image in pixels
    pub height: u32,
    /// The color type of the image
    pub color_type: ColorType,
    /// The bit depth of the image
    pub bit_depth: BitDepth,
    /// The compression method used for this image (0 for DEFLATE)
    pub compression: u8,
    /// The filter mode used for this image (currently only 0 is valid)
    pub filter: u8,
    /// The interlacing mode of the image (0 = None, 1 = Adam7)
    pub interlaced: u8,
}

#[derive(Debug, PartialEq, Clone)]
/// Options to use for performing operations on headers (such as stripping)
pub enum Headers {
    /// None
    None,
    /// Remove specific chunks
    Strip(Vec<String>),
    /// Headers that won't affect rendering (all but cHRM, gAMA, iCCP, sBIT, sRGB, bKGD, hIST, pHYs, sPLT)
    Safe,
    /// Remove all non-critical chunks except these
    Keep(HashSet<String>),
    /// All non-critical headers
    All,
}

#[inline]
pub fn file_header_is_valid(bytes: &[u8]) -> bool {
    let expected_header: [u8; 8] = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];

    *bytes == expected_header
}

#[derive(Debug, Clone, Copy)]
pub struct RawHeader<'a> {
    pub name: &'a [u8],
    pub data: &'a [u8],
}

pub fn parse_next_header<'a>(
    byte_data: &'a [u8],
    byte_offset: &mut usize,
    fix_errors: bool,
) -> PngResult<Option<RawHeader<'a>>> {
    let mut rdr = Cursor::new(
        byte_data
            .get(*byte_offset..*byte_offset + 4)
            .ok_or(PngError::TruncatedData)?,
    );
    let length = rdr.read_u32::<BigEndian>().unwrap();
    *byte_offset += 4;

    let header_start = *byte_offset;
    let chunk_name = byte_data
        .get(header_start..header_start + 4)
        .ok_or(PngError::TruncatedData)?;
    if chunk_name == b"IEND" {
        // End of data
        return Ok(None);
    }
    *byte_offset += 4;

    let data = byte_data
        .get(*byte_offset..*byte_offset + length as usize)
        .ok_or(PngError::TruncatedData)?;
    *byte_offset += length as usize;
    let mut rdr = Cursor::new(
        byte_data
            .get(*byte_offset..*byte_offset + 4)
            .ok_or(PngError::TruncatedData)?,
    );
    let crc = rdr.read_u32::<BigEndian>().unwrap();
    *byte_offset += 4;

    let header_bytes = byte_data
        .get(header_start..header_start + 4 + length as usize)
        .ok_or(PngError::TruncatedData)?;
    if !fix_errors && crc32::checksum_ieee(header_bytes) != crc {
        return Err(PngError::new(&format!(
            "CRC Mismatch in {} header; May be recoverable by using --fix",
            String::from_utf8_lossy(chunk_name)
        )));
    }

    Ok(Some(RawHeader {
        name: chunk_name,
        data,
    }))
}

pub fn parse_ihdr_header(byte_data: &[u8]) -> PngResult<IhdrData> {
    let mut rdr = Cursor::new(&byte_data[0..8]);
    Ok(IhdrData {
        color_type: match byte_data[9] {
            0 => ColorType::Grayscale,
            2 => ColorType::RGB,
            3 => ColorType::Indexed,
            4 => ColorType::GrayscaleAlpha,
            6 => ColorType::RGBA,
            _ => return Err(PngError::new("Unexpected color type in header")),
        },
        bit_depth: match byte_data[8] {
            1 => BitDepth::One,
            2 => BitDepth::Two,
            4 => BitDepth::Four,
            8 => BitDepth::Eight,
            16 => BitDepth::Sixteen,
            _ => return Err(PngError::new("Unexpected bit depth in header")),
        },
        width: rdr.read_u32::<BigEndian>().unwrap(),
        height: rdr.read_u32::<BigEndian>().unwrap(),
        compression: byte_data[10],
        filter: byte_data[11],
        interlaced: byte_data[12],
    })
}