From 16d1838e5bca2e90ca7cf8584a786b84fc409708 Mon Sep 17 00:00:00 2001 From: Silas Bartha Date: Mon, 27 May 2024 11:18:36 -0400 Subject: Initial Commit --- .gitignore | 1 + Cargo.toml | 16 ++++++ src/codec.rs | 14 +++++ src/jpg/mod.rs | 2 + src/jpg/segment.rs | 66 ++++++++++++++++++++++ src/lib.rs | 7 +++ src/lossless/champleve.rs | 139 ++++++++++++++++++++++++++++++++++++++++++++++ src/lossless/mod.rs | 2 + 8 files changed, 247 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.toml create mode 100644 src/codec.rs create mode 100644 src/jpg/mod.rs create mode 100644 src/jpg/segment.rs create mode 100644 src/lib.rs create mode 100644 src/lossless/champleve.rs create mode 100644 src/lossless/mod.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ea8c4bf --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/target diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..9a471d6 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "occule" +version = "0.1.0" +edition = "2021" + +[features] +default = ["jpg", "png", "lossless"] +jpg = [] +png = [] +lossless = [] + +[dependencies] +colored = "2.1.0" +image = "0.24" +img-parts = "0.3.0" +thiserror = "1.0.61" diff --git a/src/codec.rs b/src/codec.rs new file mode 100644 index 0000000..03cdf15 --- /dev/null +++ b/src/codec.rs @@ -0,0 +1,14 @@ +pub trait Codec { + type Carrier; + type Payload; + type Output; + type Error; + + fn encode( + &self, + carrier: impl Into, + payload: impl Into, + ) -> Result; + + fn decode(&self, encoded: impl Into) -> Result<(Self::Carrier, Self::Payload), Self::Error>; +} diff --git a/src/jpg/mod.rs b/src/jpg/mod.rs new file mode 100644 index 0000000..3d6bdc3 --- /dev/null +++ b/src/jpg/mod.rs @@ -0,0 +1,2 @@ +mod segment; +pub use segment::*; diff --git a/src/jpg/segment.rs b/src/jpg/segment.rs new file mode 100644 index 0000000..f54d0e2 --- /dev/null +++ b/src/jpg/segment.rs @@ -0,0 +1,66 @@ +use std::{mem::size_of, usize}; + +use img_parts::jpeg::{markers, Jpeg, JpegSegment}; +use thiserror::Error; + +use crate::codec::Codec; + +#[derive(Debug, PartialEq, Eq)] +pub struct JpegSegmentCodec { + pub start_index: usize, +} + +impl Codec for JpegSegmentCodec { + type Carrier = Vec; + type Payload = Vec; + type Output = Self::Carrier; + type Error = JpegSegmentError; + + fn encode(&self, carrier: impl Into, payload: impl Into) -> Result { + let mut jpeg = match Jpeg::from_bytes(carrier.into().into()) { + Ok(image) => image, + Err(err) => return Err(JpegSegmentError::ParseFailed { inner: err }) + }; + let mut payload_bytes: Self::Carrier = payload.into(); + let segment_count = ((payload_bytes.len() + size_of::()) as u64).div_ceil((u16::MAX as usize - size_of::()) as u64); + payload_bytes.splice(0..0, segment_count.to_le_bytes()); + for (index, payload_chunk) in payload_bytes.chunks(u16::MAX as usize - size_of::()).enumerate() { + let segment = JpegSegment::new_with_contents(markers::COM, payload_chunk.to_vec().into()); + jpeg.segments_mut().insert(self.start_index + index, segment); + } + Ok(jpeg.encoder().bytes().to_vec()) + } + + fn decode(&self, encoded: impl Into) -> Result<(Self::Carrier, Self::Payload), Self::Error> { + let mut jpeg = match Jpeg::from_bytes(encoded.into().into()) { + Ok(image) => image, + Err(err) => return Err(JpegSegmentError::ParseFailed { inner: err }) + }; + let segment = jpeg.segments_mut().remove(self.start_index); + let segment_bytes = segment.contents(); + let segment_count = u64::from_le_bytes(segment_bytes[0..size_of::()].try_into().unwrap()) as usize; + let mut payload_vec: Vec = Vec::with_capacity((u16::MAX as usize - size_of::()) * segment_count); + payload_vec.extend(segment_bytes[size_of::()..].to_vec()); + + for _ in 0..segment_count-1 { + let segment = jpeg.segments_mut().remove(self.start_index); + payload_vec.extend(segment.contents()); + } + + Ok((jpeg.encoder().bytes().to_vec(), payload_vec)) + } +} + +impl Default for JpegSegmentCodec { + fn default() -> Self { + Self { + start_index: 3, + } + } +} + +#[derive(Error, Debug)] +pub enum JpegSegmentError { + #[error("Failed to parse JPEG data: {inner:?}")] + ParseFailed { inner: img_parts::Error } +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..a1ed909 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,7 @@ +pub mod codec; + +#[cfg(feature = "jpg")] +pub mod jpg; + +#[cfg(feature = "lossless")] +pub mod lossless; diff --git a/src/lossless/champleve.rs b/src/lossless/champleve.rs new file mode 100644 index 0000000..948408c --- /dev/null +++ b/src/lossless/champleve.rs @@ -0,0 +1,139 @@ +use std::cmp::Ordering; + +use image::{ColorType, DynamicImage, GenericImageView, Pixel}; +use thiserror::Error; + +use crate::codec::Codec; + +#[derive(Debug)] +pub struct ChampleveCodec; + +impl Codec for ChampleveCodec { + type Carrier = DynamicImage; + type Payload = Vec; + type Output = Self::Carrier; + type Error = ChampleveError; + + fn encode(&self, carrier: impl Into, payload: impl Into) -> Result { + let mut image: DynamicImage = carrier.into(); + let payload: Vec = payload.into(); + + if image.pixels().count() < payload.len() { + return Err(ChampleveError::PayloadTooBig); + } + + let mut payload_iter = payload.iter(); + + match image { + DynamicImage::ImageRgba8(ref mut image) => { + for pixel in image.pixels_mut() { + if let Some(payload_byte) = payload_iter.next() { + encode_pixel(pixel, *payload_byte, false); + } else { + encode_pixel(pixel, 0, true); + } + } + }, + DynamicImage::ImageRgb8(ref mut image) => { + for pixel in image.pixels_mut() { + if let Some(payload_byte) = payload_iter.next() { + encode_pixel(pixel, *payload_byte, false); + } else { + encode_pixel(pixel, 0, true); + } + } + }, + _ => return Err(ChampleveError::UnsupportedFormat { format: image.color() }) + } + + Ok(image) + } + + fn decode(&self, carrier: impl Into) -> Result<(Self::Carrier, Self::Payload), ChampleveError> { + let mut image: DynamicImage = carrier.into(); + let mut payload: Vec = Vec::new(); + + match image { + DynamicImage::ImageRgba8(ref mut image) => { + for pixel in image.pixels_mut() { + if let Some(payload_byte) = decode_pixel(pixel) { + payload.push(payload_byte); + } else { + break; + } + } + }, + DynamicImage::ImageRgb8(ref mut image) => { + for pixel in image.pixels_mut() { + if let Some(payload_byte) = decode_pixel(pixel) { + payload.push(payload_byte); + } else { + break; + } + } + }, + _ => return Err(ChampleveError::UnsupportedFormat { format: image.color() }) + } + + Ok((image, payload)) + } +} + +fn encode_pixel>(pixel: &mut P, payload_byte: u8, end_of_data: bool) { + let mut bits_remaining: i32 = 8; + for channel in pixel.channels_mut() { + *channel &= 0b11111000; + bits_remaining -= 3; + if bits_remaining <= -3 { + break; + } + + let mask = match bits_remaining.cmp(&0) { + Ordering::Less => payload_byte << -bits_remaining, + _ => payload_byte >> bits_remaining, + } & 0b00000111; + + *channel |= mask; + } + + // Add end-of-data marker to final bit if necessary + if end_of_data { + *pixel.channels_mut().last_mut().unwrap() |= 1; + } +} + +fn decode_pixel>(pixel: &mut P) -> Option { + + // Final bit as end-of-data marker + if pixel.channels().last().unwrap() & 1 == 1 { + return None; + } + + let mut bits_remaining: i32 = 8; + let mut payload_byte: u8 = 0; + for channel in pixel.channels_mut() { + bits_remaining -= 3; + if bits_remaining <= -3 { + break; + } + + let channel_bits = *channel & 0b00000111; + *channel &= 0b11111000; + let mask = match bits_remaining.cmp(&0) { + Ordering::Less => channel_bits >> -bits_remaining, + _ => channel_bits << bits_remaining, + }; + payload_byte |= mask; + } + Some(payload_byte) +} + +#[derive(Error, Debug)] +pub enum ChampleveError { + #[error("Payload is too big for the carrier. Choose a smaller payload or an image with greater pixel dimensions.")] + PayloadTooBig, + #[error("Specified image format ({format:?}) is unsupported.")] + UnsupportedFormat { + format: ColorType + }, +} diff --git a/src/lossless/mod.rs b/src/lossless/mod.rs new file mode 100644 index 0000000..a9ae225 --- /dev/null +++ b/src/lossless/mod.rs @@ -0,0 +1,2 @@ +mod champleve; +pub use champleve::*; -- cgit v1.2.3