Error handling

src/coder/decoder.rs

@@ -31,7 +31,7 @@ impl Decode for RgbDecoder {
 
 pub fn decode(image: RgbaImage) -> Result<(String, Vec<u8>), Box<dyn Error>> {
     let mut buffer = image.into_raw();
-    let header = header_decoder::decode(&buffer);
+    let header = header_decoder::decode(&buffer)?;
     let buffer = buffer.split_off(header.size() * 4);
 
     let decoder = create_decoder(&header, buffer);

src/coder/decoder/alpha_decoder.rs

@@ -15,7 +15,7 @@ impl AlphaDecoder {
         let file_name_length = self.decode_length();
 
         self.validate_data_available(file_name_length, "filename")?;
-        let file_name = String::from_utf8(self.decode_data(file_name_length)).unwrap();
+        let file_name = String::from_utf8(self.decode_data(file_name_length))?;
 
         self.validate_data_available(4, "data length")?;
         let data_length = self.decode_length();

src/coder/decoder/header_decoder.rs

@@ -1,49 +1,58 @@
-use crate::coder::header::{AlgHeader, AlphaHeader, Header, RgbHeader, ALPHA_MODE, RGB_MODE};
+use crate::coder::{
+    error::HeaderDecodeError,
+    header::{AlgHeader, AlphaHeader, Header, RgbHeader, ALPHA_MODE, RGB_MODE},
+};
 
-pub fn decode(buffer: &Vec<u8>) -> Header {
+pub fn decode(buffer: &Vec<u8>) -> Result<Header, HeaderDecodeError> {
     let mut iter = buffer.iter().skip(3).step_by(4);
-    let mode = *iter.next().unwrap();
+    let mode = *iter.next().ok_or(HeaderDecodeError(
+        "Header decode error: Not enough data to decode mode.".to_string(),
+    ))?;
 
-    let alg_header = decode_alg_header(mode, &mut iter);
-    Header::new(mode, alg_header)
+    let alg_header = decode_alg_header(mode, &mut iter)?;
+    Ok(Header::new(mode, alg_header))
 }
 
-fn decode_alg_header<'a, I>(mode: u8, iter: &mut I) -> AlgHeader
+fn decode_alg_header<'a, I>(mode: u8, iter: &mut I) -> Result<AlgHeader, HeaderDecodeError>
 where
     I: Iterator<Item = &'a u8>,
 {
     match mode {
-        ALPHA_MODE => AlgHeader::Alpha(decode_alpha()),
-        RGB_MODE => AlgHeader::Rgb(decode_rgb(iter)),
-        _ => panic!("Unknown mode. Should never happen."),
+        ALPHA_MODE => Ok(AlgHeader::Alpha(decode_alpha())),
+        RGB_MODE => Ok(AlgHeader::Rgb(decode_rgb(iter)?)),
+        _ => Err(HeaderDecodeError("Unknown mode in header.".to_string())),
     }
 }
 
 fn decode_alpha() -> AlphaHeader {
     AlphaHeader {}
 }
 
-fn decode_rgb<'a, I>(iter: &mut I) -> RgbHeader
+fn decode_rgb<'a, I>(iter: &mut I) -> Result<RgbHeader, HeaderDecodeError>
 where
     I: Iterator<Item = &'a u8>,
 {
-    RgbHeader {
-        bits_per_channel: *iter.next().unwrap(),
-    }
+    let bits_per_channel = *iter.next().ok_or(HeaderDecodeError(
+        "Header decode error: Not enough data to decode bits per channel.".to_string(),
+    ))?;
+
+    Ok(RgbHeader { bits_per_channel })
 }
 
 #[cfg(test)]
 mod tests {
-
-    use crate::coder::header::{Header, ALPHA_MODE, RGB_MODE};
+    use crate::coder::{
+        error::HeaderDecodeError,
+        header::{Header, ALPHA_MODE, RGB_MODE},
+    };
 
     #[test]
     fn decode_alpha() {
         let mut buffer = vec![0; 4];
         let mut iter = buffer.iter_mut().skip(3).step_by(4);
         *iter.next().unwrap() = ALPHA_MODE;
 
-        let decoded = super::decode(&buffer);
+        let decoded = super::decode(&buffer).unwrap();
         assert_eq!(decoded, Header::new_alpha());
     }
 
@@ -55,7 +64,45 @@ mod tests {
         *iter.next().unwrap() = RGB_MODE;
         *iter.next().unwrap() = bits_per_channel;
 
-        let decoded = super::decode(&buffer);
+        let decoded = super::decode(&buffer).unwrap();
         assert_eq!(decoded, Header::new_rgb(bits_per_channel));
     }
+
+    #[test]
+    fn decode_error_missing_mode_data() {
+        let buffer = Vec::new();
+        let decoded = super::decode(&buffer);
+        assert_eq!(
+            decoded,
+            Err(HeaderDecodeError(
+                "Header decode error: Not enough data to decode mode.".to_string()
+            ))
+        );
+    }
+
+    #[test]
+    fn decode_error_missing_bits_per_channel_data() {
+        let buffer = vec![0, 0, 0, RGB_MODE];
+        let decoded = super::decode(&buffer);
+        assert_eq!(
+            decoded,
+            Err(HeaderDecodeError(
+                "Header decode error: Not enough data to decode bits per channel.".to_string()
+            ))
+        );
+    }
+
+    #[test]
+    fn decode_error_unknown_mode() {
+        let unknown_mode = 4;
+        let mut buffer = vec![0; 4];
+        let mut iter = buffer.iter_mut().skip(3).step_by(4);
+        *iter.next().unwrap() = unknown_mode;
+
+        let decoded = super::decode(&buffer);
+        assert_eq!(
+            decoded,
+            Err(HeaderDecodeError("Unknown mode in header.".to_string()))
+        );
+    }
 }

src/coder/encoder.rs

@@ -24,7 +24,7 @@ pub fn encode(
 
     let header = create_header(algorithm);
     let data_buffer = image_data.split_off(header.size() * 4);
-    header_encoder::encode(header.clone(), &mut image_data);
+    header_encoder::encode(header.clone(), &mut image_data)?;
 
     let encoder = create_encoder(algorithm, data_buffer, data, secret_filename);
     let mut encoded = encoder.run()?;

src/coder/encoder/header_encoder.rs

@@ -1,38 +1,54 @@
-use crate::coder::header::{AlgHeader, Header, RgbHeader};
+use crate::coder::{
+    error::HeaderEncodeError,
+    header::{AlgHeader, Header, RgbHeader},
+};
 
-pub fn encode(header: Header, buffer: &mut Vec<u8>) {
+pub fn encode(header: Header, buffer: &mut Vec<u8>) -> Result<(), HeaderEncodeError> {
     let mut iter = buffer.iter_mut().skip(3).step_by(4);
-    *iter.next().unwrap() = header.mode;
+    let mode_byte = iter.next().ok_or(HeaderEncodeError(
+        "Not enough to encode header mode.".to_string(),
+    ))?;
+    *mode_byte = header.mode;
 
     match header.alg_header {
-        AlgHeader::Alpha(_) => encode_alpha(&mut iter),
-        AlgHeader::Rgb(alg_header) => encode_rgb(&mut iter, &alg_header),
+        AlgHeader::Alpha(_) => encode_alpha(&mut iter)?,
+        AlgHeader::Rgb(alg_header) => encode_rgb(&mut iter, &alg_header)?,
     }
+
+    Ok(())
 }
 
-fn encode_alpha<'a, I>(_iter: &mut I)
+fn encode_alpha<'a, I>(_iter: &mut I) -> Result<(), HeaderEncodeError>
 where
     I: Iterator<Item = &'a mut u8>,
 {
+    Ok(())
 }
 
-fn encode_rgb<'a, I>(iter: &mut I, header: &RgbHeader)
+fn encode_rgb<'a, I>(iter: &mut I, header: &RgbHeader) -> Result<(), HeaderEncodeError>
 where
     I: Iterator<Item = &'a mut u8>,
 {
-    *iter.next().unwrap() = header.bits_per_channel;
+    let bits_per_channel_byte = iter.next().ok_or(HeaderEncodeError(
+        "Not enough to encode header bits per channel.".to_string(),
+    ))?;
+    *bits_per_channel_byte = header.bits_per_channel;
+    Ok(())
 }
 
 #[cfg(test)]
 mod tests {
 
-    use crate::coder::header::{Header, ALPHA_MODE, RGB_MODE};
+    use crate::coder::{
+        error::HeaderEncodeError,
+        header::{Header, ALPHA_MODE, RGB_MODE},
+    };
 
     #[test]
     fn encode_alpha() {
         let header = Header::new_alpha();
         let mut buffer = vec![0; 10];
-        super::encode(header, &mut buffer);
+        assert_eq!(Ok(()), super::encode(header, &mut buffer));
         assert_eq!(buffer, vec![0, 0, 0, ALPHA_MODE, 0, 0, 0, 0, 0, 0]);
     }
 
@@ -41,10 +57,37 @@ mod tests {
         let bits_per_channel = 4;
         let header = Header::new_rgb(bits_per_channel);
         let mut buffer = vec![0; 10];
-        super::encode(header, &mut buffer);
+        assert_eq!(Ok(()), super::encode(header, &mut buffer));
         assert_eq!(
             buffer,
             vec![0, 0, 0, RGB_MODE, 0, 0, 0, bits_per_channel, 0, 0]
         );
     }
+
+    #[test]
+    fn encode_error_not_enough_data_for_mode() {
+        let header = Header::new_alpha();
+        let mut buffer = vec![0; 1];
+        let encoded = super::encode(header, &mut buffer);
+        assert_eq!(
+            encoded,
+            Err(HeaderEncodeError(
+                "Not enough to encode header mode.".to_string()
+            ))
+        );
+    }
+
+    #[test]
+    fn encode_error_not_enough_data_for_bits_per_channel() {
+        let bits_per_channel = 1;
+        let header = Header::new_rgb(bits_per_channel);
+        let mut buffer = vec![0; 4];
+        let encoded = super::encode(header, &mut buffer);
+        assert_eq!(
+            encoded,
+            Err(HeaderEncodeError(
+                "Not enough to encode header bits per channel.".to_string()
+            ))
+        );
+    }
 }

src/coder/error.rs

@@ -1,11 +1,18 @@
 use core::fmt;
+use std::string::FromUtf8Error;
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct EncodeError(pub String);
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct DecodeError(pub String);
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct HeaderEncodeError(pub String);
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct HeaderDecodeError(pub String);
+
 impl fmt::Display for EncodeError {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(f, "Encode error: {}", self.0)
@@ -18,5 +25,58 @@ impl fmt::Display for DecodeError {
     }
 }
 
+impl fmt::Display for HeaderEncodeError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "Header encode error: {}", self.0)
+    }
+}
+
+impl fmt::Display for HeaderDecodeError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "Header decode error: {}", self.0)
+    }
+}
+
 impl std::error::Error for EncodeError {}
 impl std::error::Error for DecodeError {}
+impl std::error::Error for HeaderEncodeError {}
+impl std::error::Error for HeaderDecodeError {}
+
+impl From<FromUtf8Error> for DecodeError {
+    fn from(value: FromUtf8Error) -> Self {
+        DecodeError(value.to_string())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use super::DecodeError;
+    use super::EncodeError;
+    use super::HeaderDecodeError;
+    use super::HeaderEncodeError;
+
+    #[test]
+    fn display_encode_error() {
+        let error = EncodeError("some failure".to_string());
+        assert_eq!(error.to_string(), "Encode error: some failure");
+    }
+
+    #[test]
+    fn display_decode_error() {
+        let error = DecodeError("some failure".to_string());
+        assert_eq!(error.to_string(), "Decode error: some failure");
+    }
+
+    #[test]
+    fn display_header_encode_error() {
+        let error = HeaderEncodeError("some failure".to_string());
+        assert_eq!(error.to_string(), "Header encode error: some failure");
+    }
+
+    #[test]
+    fn display_header_decode_error() {
+        let error = HeaderDecodeError("some failure".to_string());
+        assert_eq!(error.to_string(), "Header decode error: some failure");
+    }
+}