diff --git a/src/coder/decoder.rs b/src/coder/decoder.rs index e8a5911..44e2869 100644 --- a/src/coder/decoder.rs +++ b/src/coder/decoder.rs @@ -31,7 +31,7 @@ impl Decode for RgbDecoder { pub fn decode(image: RgbaImage) -> Result<(String, Vec), Box> { 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); diff --git a/src/coder/decoder/alpha_decoder.rs b/src/coder/decoder/alpha_decoder.rs index 36f8cb1..e27e9af 100644 --- a/src/coder/decoder/alpha_decoder.rs +++ b/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(); diff --git a/src/coder/decoder/header_decoder.rs b/src/coder/decoder/header_decoder.rs index f7816c6..a5d366c 100644 --- a/src/coder/decoder/header_decoder.rs +++ b/src/coder/decoder/header_decoder.rs @@ -1,21 +1,26 @@ -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) -> Header { +pub fn decode(buffer: &Vec) -> Result { 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 where I: Iterator, { 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())), } } @@ -23,19 +28,23 @@ fn decode_alpha() -> AlphaHeader { AlphaHeader {} } -fn decode_rgb<'a, I>(iter: &mut I) -> RgbHeader +fn decode_rgb<'a, I>(iter: &mut I) -> Result where I: Iterator, { - 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() { @@ -43,7 +52,7 @@ mod tests { 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())) + ); + } } diff --git a/src/coder/encoder.rs b/src/coder/encoder.rs index 2eea2c2..411fe90 100644 --- a/src/coder/encoder.rs +++ b/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()?; diff --git a/src/coder/encoder/header_encoder.rs b/src/coder/encoder/header_encoder.rs index adb2dac..64c66f9 100644 --- a/src/coder/encoder/header_encoder.rs +++ b/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) { +pub fn encode(header: Header, buffer: &mut Vec) -> 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, { + 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, { - *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() + )) + ); + } } diff --git a/src/coder/error.rs b/src/coder/error.rs index 3d8692e..d27afdc 100644 --- a/src/coder/error.rs +++ b/src/coder/error.rs @@ -1,4 +1,5 @@ use core::fmt; +use std::string::FromUtf8Error; #[derive(Debug, Clone, PartialEq)] pub struct EncodeError(pub String); @@ -6,6 +7,12 @@ 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 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"); + } +}