Made codecs object-safe and deal with bytes only

3 files changed, 55 insertions, 92 deletions

diff --git a/src/codec.rs b/src/codec.rs
index 8b38a9b..11fa76e 100644
--- a/src/codec.rs
+++ b/src/codec.rs
@@ -1,26 +1,28 @@
+use thiserror::Error;
+
 /// Codecs enable the concealment of payload data inside the data of a carrier.
 pub trait Codec {
-    /// Data type representing the carrier.
-    type Carrier;
-
-    /// Data type representing the payload.
-    type Payload;
-
-    /// Data type representing encoder output/decoder input (usually the same as the carrier).
-    type Output;
-
-    /// Type of errors produced by this codec.
-    type Error;
-
     /// Embeds payload data inside carrier, returning the result.
-    fn encode<C, P>(&self, carrier: C, payload: P) -> Result<Self::Output, Self::Error>
-    where
-        C: Into<Self::Carrier>,
-        P: Into<Self::Payload>;
+    fn encode(&self, carrier: &[u8], payload: &[u8]) -> Result<Vec<u8>, CodecError>;
 
     /// Extracts payload data from an encoded carrier, returning the carrier with data removed and the
     /// payload data.
-    fn decode<E>(&self, encoded: E) -> Result<(Self::Carrier, Self::Payload), Self::Error>
-    where
-        E: Into<Self::Output>;
+    fn decode(&self, encoded: &[u8]) -> Result<(Vec<u8>, Vec<u8>), CodecError>;
+}
+
+/// Errors produced by a codec
+#[derive(Debug, Error)]
+pub enum CodecError {
+    /// Variant used when data is determined not to be encoded. Note that a codec may have no way
+    /// of knowing this, so this may not be returned even if the data was not encoded
+    #[error("Data was not encoded with this codec")]
+    DataNotEncoded,
+
+    /// Variant used when data is invalid in some way. Allows a message string for further context
+    #[error("Provided data invalid: {0}")]
+    DataInvalid(String),
+
+    /// Variant used when some dependency, such as a file load, fails
+    #[error("Error occured in dependency: {0}")]
+    DependencyError(String),
 }
diff --git a/src/jpeg/segment.rs b/src/jpeg/segment.rs
index b20ba3b..cd1a651 100644
--- a/src/jpeg/segment.rs
+++ b/src/jpeg/segment.rs
@@ -1,9 +1,8 @@
 use std::{mem::size_of, usize};
 
 use img_parts::jpeg::{markers, Jpeg, JpegSegment};
-use thiserror::Error;
 
-use crate::codec::Codec;
+use crate::{codec::Codec, CodecError};
 
 /// Codec for storing payload data in JPEG comment (COM) segments. Can store an arbitrary amount of
 /// data, as long as the number of comment segments does not exceed u64::MAX.
@@ -14,21 +13,13 @@ pub struct JpegSegmentCodec {
 }
 
 impl Codec for JpegSegmentCodec {
-    type Carrier = Vec<u8>;
-    type Payload = Vec<u8>;
-    type Output = Self::Carrier;
-    type Error = JpegSegmentError;
-
-    fn encode<C, P>(&self, carrier: C, payload: P) -> Result<Self::Output, Self::Error>
-    where
-        C: Into<Self::Carrier>,
-        P: Into<Self::Payload>,
+    fn encode(&self, carrier: &[u8], payload: &[u8]) -> Result<Vec<u8>, CodecError>
     {
-        let mut jpeg = match Jpeg::from_bytes(carrier.into().into()) {
-            Ok(image) => image,
-            Err(err) => return Err(JpegSegmentError::ParseFailed { inner: err })
+        let mut jpeg = match Jpeg::from_bytes(carrier.to_vec().into()) {
+            Ok(v) => v,
+            Err(e) => return Err(CodecError::DependencyError(e.to_string()))
         };
-        let mut payload_bytes: Self::Carrier = payload.into();
+        let mut payload_bytes = payload.to_vec();
         let segment_count = ((payload_bytes.len() + size_of::<u64>()) as u64).div_ceil((u16::MAX as usize - size_of::<u16>()) 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::<u16>()).enumerate() {
@@ -38,13 +29,11 @@ impl Codec for JpegSegmentCodec {
         Ok(jpeg.encoder().bytes().to_vec())
     }
 
-    fn decode<E>(&self, encoded: E) -> Result<(Self::Carrier, Self::Payload), Self::Error>
-    where
-        E: Into<Self::Output>,
+    fn decode(&self, encoded: &[u8]) -> Result<(Vec<u8>, Vec<u8>), CodecError>
     {
-        let mut jpeg = match Jpeg::from_bytes(encoded.into().into()) {
-            Ok(image) => image,
-            Err(err) => return Err(JpegSegmentError::ParseFailed { inner: err })
+        let mut jpeg = match Jpeg::from_bytes(encoded.to_vec().into()) {
+            Ok(v) => v,
+            Err(e) => return Err(CodecError::DependencyError(e.to_string()))
         };
         let segment = jpeg.segments_mut().remove(self.start_index);
         let segment_bytes = segment.contents();
@@ -68,14 +57,3 @@ impl Default for JpegSegmentCodec {
         }
     }
 }
-
-/// Errors thrown by the JPEG segment codec.
-#[derive(Error, Debug)]
-pub enum JpegSegmentError {
-    /// Parsing JPEG data failed.
-    #[error("Failed to parse JPEG data: {inner:?}")]
-    ParseFailed { 
-        /// Error thrown by parser.
-        inner: img_parts::Error,
-    }
-}
diff --git a/src/lossless/lsb.rs b/src/lossless/lsb.rs
index 59dcc0b..f8f1113 100644
--- a/src/lossless/lsb.rs
+++ b/src/lossless/lsb.rs
@@ -1,33 +1,25 @@
-use std::cmp::Ordering;
+use std::{cmp::Ordering, io::{BufWriter, Cursor}};
 
-use image::{ColorType, DynamicImage, GenericImageView, Pixel};
-use thiserror::Error;
+use image::{DynamicImage, GenericImageView, Pixel};
 
-use crate::codec::Codec;
+use crate::{codec::Codec, CodecError};
 
 /// Least-significant bit (LSB) steganography encodes data in the least-significant bits of colors
 /// in an image. This implementation reduces the colors in the carrier (irreversibly) in order to
 /// allow a byte of data to fit in each pixel of the image. 3 bits of data are encoded per pixel,
 /// and the 9th bit is used to signal the end of data.
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub struct LsbCodec;
 
 impl Codec for LsbCodec {
-    type Carrier = DynamicImage;
-    type Payload = Vec<u8>;
-    type Output = Self::Carrier;
-    type Error = LsbError;
-
-    fn encode<C, P>(&self, carrier: C, payload: P) -> Result<Self::Output, Self::Error>
-    where
-        C: Into<Self::Carrier>,
-        P: Into<Self::Payload>,
+    fn encode(&self, carrier: &[u8], payload: &[u8]) -> Result<Vec<u8>, CodecError>
     {
-        let mut image: DynamicImage = carrier.into();
-        let payload: Vec<u8> = payload.into();
+        let image_format = image::guess_format(carrier.into()).unwrap();
+        let mut image: DynamicImage = image::load_from_memory(carrier.into()).unwrap();
+        let payload: &[u8] = payload.into();
 
         if image.pixels().count() < payload.len() {
-            return Err(LsbError::PayloadTooBig);
+            return Err(CodecError::DataInvalid("Payload Too Big for Carrier".into()));
         }
 
         let mut payload_iter = payload.iter();
@@ -51,17 +43,20 @@ impl Codec for LsbCodec {
                     }
                 }
             },
-            _ => return Err(LsbError::UnsupportedFormat { format: image.color() })
+            _ => return Err(CodecError::DataInvalid("Unsupported Image Color Format".into()))
         }
 
-        Ok(image)
+        let mut buf = BufWriter::new(Cursor::new(Vec::<u8>::new()));
+        if let Err(e) = image.write_to(&mut buf, image_format) {
+            return Err(CodecError::DependencyError(e.to_string()))
+        }
+        Ok(buf.into_inner().unwrap().into_inner())
     }
 
-    fn decode<E>(&self, carrier: E) -> Result<(Self::Carrier, Self::Payload), LsbError>
-    where
-        E: Into<Self::Output>,
+    fn decode(&self, carrier: &[u8]) -> Result<(Vec<u8>, Vec<u8>), CodecError>
     {
-        let mut image: DynamicImage = carrier.into();
+        let image_format = image::guess_format(carrier.into()).unwrap();
+        let mut image: DynamicImage = image::load_from_memory(carrier.into()).unwrap();
         let mut payload: Vec<u8> = Vec::new();
 
         match image {
@@ -83,10 +78,14 @@ impl Codec for LsbCodec {
                     }
                 }
             },
-            _ => return Err(LsbError::UnsupportedFormat { format: image.color() })
+            _ => return Err(CodecError::DataInvalid("Unsupported Image Color Format".into()))
         }
         
-        Ok((image, payload))
+        let mut buf = BufWriter::new(Cursor::new(Vec::<u8>::new()));
+        if let Err(e) = image.write_to(&mut buf, image_format) {
+            return Err(CodecError::DependencyError(e.to_string()))
+        }
+        Ok((buf.into_inner().unwrap().into_inner(), payload))
     }
 }
 
@@ -138,19 +137,3 @@ fn decode_pixel<P: Pixel<Subpixel = u8>>(pixel: &mut P) -> Option<u8> {
     }
     Some(payload_byte)
 }
-
-/// Errors thrown by the LSB Codec.
-#[derive(Error, Debug)]
-pub enum LsbError {
-
-    /// Error thrown when payload is too big for the carrier.
-    #[error("Payload is too big for the carrier. Choose a smaller payload or an image with greater pixel dimensions.")]
-    PayloadTooBig,
-
-    /// Error thrown when pixel format is unsupported.
-    #[error("Specified image format ({format:?}) is unsupported.")]
-    UnsupportedFormat {
-        /// Provided (invalid) format.
-        format: ColorType
-    },
-}