From 0e1bda128a4277296cb1f620cdf12d35482823f4 Mon Sep 17 00:00:00 2001 From: pommicket Date: Sun, 3 Sep 2023 15:06:45 -0400 Subject: initial commit --- src/lib.rs | 328 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 328 insertions(+) create mode 100644 src/lib.rs (limited to 'src') diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..a801f9a --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,328 @@ +//#![cfg_attr(not(feature = "std"), no_std)] + +use core::cmp::min; +use core::fmt::{self, Debug, Display}; + +pub trait IOError: Sized + Display + Debug { +} +impl IOError for T { +} + +#[cfg(feature = "std")] +impl IOError for std::io::Error { + fn unexpected_eof() -> Self { + std::io::ErrorKind::UnexpectedEof.into() + } +} + +#[derive(Debug)] +#[non_exhaustive] +pub enum Error { + IO(I), + NotPng, + BadIhdr, + UnrecognizedChunk([u8; 4]), + BadBlockType, + TooMuchData, + UnexpectedEob, + BadZlibHeader, +} + +impl From for Error { + fn from(value: I) -> Self { + Self::IO(value) + } +} + +impl Display for Error { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::IO(e) => write!(f, "{e}"), + Self::NotPng => write!(f, "not a png file"), + Self::BadIhdr => write!(f, "bad IHDR chunk"), + Self::UnrecognizedChunk([a, b, c, d]) => write!(f, "unrecognized chunk type: {a} {b} {c} {d}"), + Self::BadBlockType => write!(f, "bad DEFLATE block type"), + Self::TooMuchData => write!(f, "decompressed data is larger than it should be"), + Self::UnexpectedEob => write!(f, "unexpected end of block"), + Self::BadZlibHeader => write!(f, "bad zlib header"), + } + } +} + +#[cfg(feature = "std")] +impl std::error::Error for Error {} + +pub trait Read { + type Error: IOError; + fn read(&mut self, buf: &mut [u8]) -> Result<(), Self::Error>; + fn skip_bytes(&mut self, count: usize) -> Result<(), Self::Error> { + let mut count = count; + const BUF_LEN: usize = 128; + let mut buf = [0; BUF_LEN]; + while count > 0 { + let c = min(BUF_LEN, count); + self.read(&mut buf[..c])?; + count -= c; + } + Ok(()) + } +} + +#[derive(Debug)] +pub struct UnexpectedEofError; + +impl core::fmt::Display for UnexpectedEofError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "unexpected EOF") + } +} + +impl<'a> Read for &'a [u8] { + type Error = UnexpectedEofError; + fn read(&mut self, buf: &mut [u8]) -> Result<(), Self::Error> { + if self.len() < buf.len() { + return Err(UnexpectedEofError); + } + buf.copy_from_slice(&self[..buf.len()]); + *self = &self[buf.len()..]; + Ok(()) + } + // TODO: skip_bytes implementation +} + +struct BlockReader<'a, R: Read> { + inner: &'a mut R, + bytes_left: usize +} + +impl<'a, R: Read> BlockReader<'_, R> { + fn read(&mut self, buf: &mut [u8]) -> Result<(), Error> { + if buf.len() > self.bytes_left { + return Err(Error::UnexpectedEob); + } + Ok(self.inner.read(buf)?) + } + + fn read_partial(&mut self, buf: &mut [u8]) -> Result> { + let count = min(self.bytes_left, buf.len()); + self.read(&mut buf[..count])?; + Ok(count) + } +} + +#[derive(Debug)] +pub struct ImageHeader { + width: u32, + height: u32, + bit_depth: BitDepth, + color_type: ColorType, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum BitDepth { + One, + Two, + Four, + Eight, + Sixteen, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum ColorType { + Gray, + GrayAlpha, + Rgb, + Rgba, + Indexed, +} + +impl BitDepth { + fn from_byte(x: u8) -> Option { + Some(match x { + 1 => Self::One, + 2 => Self::Two, + 4 => Self::Four, + 8 => Self::Eight, + 16 => Self::Sixteen, + _ => return None + }) + } +} + +impl ColorType { + fn from_byte(x: u8) -> Option { + Some(match x { + 0 => Self::Gray, + 2 => Self::Rgb, + 3 => Self::Indexed, + 4 => Self::GrayAlpha, + 6 => Self::Rgba, + _ => return None, + }) + } +} + +impl ImageHeader { + pub fn width(&self) -> u32 { + self.width + } + pub fn height(&self) -> u32 { + self.height + } + pub fn bit_depth(&self) -> BitDepth { + self.bit_depth + } + pub fn color_type(&self) -> ColorType { + self.color_type + } +} + +/// number of bits to read in each [`Read::read`] call. +/// +/// don't change this to something bigger than `u32`, since we don't want to overread past the zlib checksum. +type ReadBits = u32; +/// number of bits to store in the [`BitReader`] buffer. +type Bits = u64; + +struct BitReader<'a, R: Read> { + inner: BlockReader<'a, R>, + bits: Bits, + bits_left: u8, +} + +impl<'a, R: Read> From> for BitReader<'a, R> { + fn from(inner: BlockReader<'a, R>) -> Self { + Self { + inner, + bits: 0, + bits_left: 0, + } + } +} + +impl BitReader<'_, R> { + fn peek_bits(&mut self, count: u8) -> Result> { + debug_assert!(count > 0 && u32::from(count) <= 31); + if self.bits_left < count { + // read more bits + let mut new_bits = [0; ReadBits::BITS as usize / 8]; + self.inner.read_partial(&mut new_bits)?; + let new_bits = Bits::from(ReadBits::from_le_bytes(new_bits)); + self.bits |= new_bits << self.bits_left; + self.bits_left += ReadBits::BITS as u8; + } + Ok((self.bits as u32) & ((1 << count) - 1)) + } + + fn read_bits(&mut self, count: u8) -> Result> { + let bits = self.peek_bits(count)?; + self.bits_left -= count; + self.bits >>= count; + Ok(bits) + } +} + + +pub fn read_png_header(reader: &mut R) -> Result> { + let mut signature = [0; 8]; + reader.read(&mut signature)?; + if signature != [137, 80, 78, 71, 13, 10, 26, 10] { + return Err(Error::NotPng); + } + let mut ihdr = [0; 25]; + reader.read(&mut ihdr)?; + let ihdr_len = (u32::from_be_bytes([ihdr[0], ihdr[1], ihdr[2], ihdr[3]]) + 12) as usize; + if &ihdr[4..8] != b"IHDR" || ihdr_len < ihdr.len() { + return Err(Error::BadIhdr); + } + reader.skip_bytes(ihdr_len - ihdr.len())?; + + let width = u32::from_be_bytes([ihdr[8], ihdr[9], ihdr[10], ihdr[11]]); + let height = u32::from_be_bytes([ihdr[12], ihdr[13], ihdr[14], ihdr[15]]); + let bit_depth = BitDepth::from_byte(ihdr[16]).ok_or(Error::BadIhdr)?; + let color_type = ColorType::from_byte(ihdr[17]).ok_or(Error::BadIhdr)?; + let compression = ihdr[18]; + let filter = ihdr[19]; + if compression != 0 || filter != 0 { + return Err(Error::BadIhdr); + } + + Ok(ImageHeader { width, height, bit_depth, color_type }) +} + +struct ByteWriter<'a> { + slice: &'a mut [u8], +} + +fn read_idat(mut reader: BlockReader<'_, R>, writer: &mut ByteWriter) -> Result<(), Error> { + let mut zlib_header = [0; 2]; + reader.read(&mut zlib_header)?; + + let mut reader = BitReader::from(reader); + loop { + let bfinal = reader.read_bits(1)?; + let btype = reader.read_bits(2)?; + if btype == 0 { + // uncompressed block + let len = reader.read_bits(16)? as usize; + reader.read_bits(16)?; // nlen + if writer.slice.len() < len { + return Err(Error::TooMuchData); + } + reader.inner.read(&mut writer.slice[..len])?; + } else if btype == 1 || btype == 2 { + // compressed block + todo!("{btype}") + } else { + // 0b11 is not a valid block type + return Err(Error::BadBlockType); + } + + if bfinal != 0 { + break; + } + } + + // skip checksum + let mut trash = [0; 4]; + reader.inner.read(&mut trash[..4 - usize::from(reader.bits_left) / 8])?; + + Ok(()) +} + +pub fn read_png( + reader: &mut R, + header: Option, + buf: &mut [u8], +) -> Result<(), Error> { + let _header = match header { + None => read_png_header(reader)?, + Some(h) => h, + }; + let mut writer = ByteWriter { slice: buf }; + loop { + let mut chunk_header = [0; 8]; + reader.read(&mut chunk_header[..])?; + let chunk_len = u32::from_be_bytes([chunk_header[0], chunk_header[1], chunk_header[2], chunk_header[3]]) as usize; + let chunk_type = [chunk_header[4], chunk_header[5], chunk_header[6], chunk_header[7]]; + if &chunk_type == b"IEND" { + break; + } else if &chunk_type == b"IDAT" { + read_idat(BlockReader { + inner: reader, + bytes_left: chunk_len + }, &mut writer)?; + } else if &chunk_type == b"PLTE" { + todo!(); + } else if chunk_type[0].is_ascii_lowercase() { + // non-essential chunk + } else { + return Err(Error::UnrecognizedChunk(chunk_type)); + } + + reader.skip_bytes(chunk_len + 4)?; + } + + Ok(()) +} + -- cgit v1.2.3