Add unit test JpegStreamReader::read_header_from_too_small_input_buff…

…er_fails

src/decoder.rs

@@ -1,7 +1,7 @@
 // Copyright (c) Team CharLS.
 // SPDX-License-Identifier: BSD-3-Clause
 
-use std::io::{Read, self};
+use std::io::{Read};
 
 #[warn(unused_variables)]
 
@@ -19,7 +19,7 @@ pub struct Decoder<R: Read> {
 
 
 impl<R: Read> Decoder<R> {
-    pub fn new(mut r: R) -> Decoder<R> {
+    pub fn new(r: R) -> Decoder<R> {
         let width = 0;
         let height = 0;
         let bits_per_sample = 0;

src/jpeg_marker_code.rs

@@ -1,6 +1,7 @@
 // Copyright (c) Team CharLS.
 // SPDX-License-Identifier: BSD-3-Clause
 
+use std::convert::TryFrom;
 
 #[derive(Debug, Eq, PartialEq)]
 pub enum JpegMarkerCode {
@@ -46,3 +47,16 @@ pub enum JpegMarkerCode {
     ApplicationData15 = 0xEF, // APP15: Application data 15.
     Comment = 0xFE             // COM:   Comment block.
 }
+
+impl TryFrom<u8> for JpegMarkerCode {
+    type Error = ();
+
+    fn try_from(v: u8) -> Result<Self, Self::Error> {
+        match v {
+            x if x == JpegMarkerCode::StartOfImage as u8 => Ok(JpegMarkerCode::StartOfImage),
+            x if x == JpegMarkerCode::EndOfImage as u8 => Ok(JpegMarkerCode::EndOfImage),
+            x if x == JpegMarkerCode::StartOfScan as u8 => Ok(JpegMarkerCode::StartOfScan),
+            _ => Err(()),
+        }
+    }
+}

src/jpeg_stream_reader.rs

@@ -3,7 +3,7 @@
 
 //mod jpeg_marker_code;
 
-use std::io::{Read, self};
+use std::io::Read;
 
 use crate::jpeg_marker_code::JpegMarkerCode;
 use crate::decoding_error::DecodingError;
@@ -40,7 +40,7 @@ pub struct JpegStreamReader<R: Read> {
 
 
 impl<R: Read> JpegStreamReader<R> {
-    pub fn new(mut r: R) -> JpegStreamReader<R> {
+    pub fn new(r: R) -> JpegStreamReader<R> {
         let width = 0;
         let height = 0;
         let bits_per_sample = 0;
@@ -58,13 +58,28 @@ impl<R: Read> JpegStreamReader<R> {
         }
     }
 
-    pub fn read_next_marker_code(&mut self) -> JpegMarkerCode {
-        JpegMarkerCode::StartOfImage
+    pub fn read_next_marker_code(&mut self) -> Result<JpegMarkerCode, DecodingError> {
+        let mut value = self.read_u8()?;
+        if value != 255 {
+            return Err(DecodingError::StartOfImageMarkerNotFound);
+        }
+
+        // Read all preceding 0xFF fill values until a non 0xFF value has been found. (see ISO/IEC 10918-1, B.1.1.2)
+        while value == 255 {
+            value = self.read_u8()?;
+        }
+
+        let r = JpegMarkerCode::try_from(value);
+        if r.is_err() {
+            return Err(DecodingError::StartOfImageMarkerNotFound);
+        }
+
+        return Ok(r.unwrap())
     }
 
     pub fn read_header(&mut self) -> Result<(), DecodingError> {
         if self.state == ReaderState::BeforeStartOfImage {
-            if self.read_next_marker_code() != JpegMarkerCode::StartOfImage {
+            if self.read_next_marker_code()? != JpegMarkerCode::StartOfImage {
                 return Err(DecodingError::StartOfImageMarkerNotFound);
             }
 
@@ -73,5 +88,29 @@ impl<R: Read> JpegStreamReader<R> {
 
         Ok(())
     }
+
+    fn read_u8(&mut self) -> Result<u8, DecodingError> {
+        let mut buf = [0; 1];
+        let result = self.reader.read_exact(&mut buf);
+        if result.is_err() {
+            return Err(DecodingError::UnknownError);
+        }
+
+        Ok(buf[0])
+    }
 }
 
+#[cfg(test)]
+mod tests {
+    use std::io::Write;
+    use super::*;
+
+    #[test]
+    fn read_header_from_too_small_input_buffer_fails() {
+        let mut buffer = Vec::new();
+        buffer.write_all(&[1]).unwrap();
+
+        let mut reader = JpegStreamReader::new(buffer.as_slice());
+        assert!(reader.read_header().is_err());
+    }
+}