Harden for multi-scan JPEG

dll/src/lib.rs

@@ -924,8 +924,8 @@ mod tests {
         "iphonecrop2",
         "iphoneprogressive",
         "iphoneprogressive2",
-        "progressive_late_dht", // image has huffman tables that come very late which causes a verification failure 
-        "out_of_order_dqt",     // image with quanatization table dqt that comes after image definition SOF
+        "progressive_late_dht",
+        "out_of_order_dqt",
         "narrowrst",
         "nofsync",
         "slrcity",
@@ -935,8 +935,8 @@ mod tests {
         "trailingrst",
         "trailingrst2",
         "trunc",
-        "eof_and_trailingrst",    // the lepton format has a wrongly set unexpected eof and trailing rst
-        "eof_and_trailinghdrdata" // the lepton format has a wrongly set unexpected eof and trailing header data
+        "eof_and_trailingrst",
+        "eof_and_trailinghdrdata"
     )]
         file: &str,
     ) {

images/README.md

@@ -0,0 +1,10 @@
+# Images
+
+- `dc_init_non_interleaved` is a progressive JPEG with DC initial scans that do not contain all components. Associated Lepton file produced by the original implementation.
+- `eof_and_trailinghdrdata` the lepton format has a wrongly set unexpected EOF and trailing header data.
+- `eof_and_trailingrst` the Lepton format has a wrongly set unexpected EOF and trailing RST.
+- `non-interleaved-multiple-restart` is a sequential JPEG with non-interleaved scans and chroma subsampling. The restart interval is set to one MCU row, meaning that there are multiple restart intervals defined on the JPEG. Associated Lepton file produced by the original implementation.
+- `out_of_order_dqt` has a quantization table that comes after the image definition SOF.
+- `progressive_late_dht`  has Huffman tables that come very late which causes a verification failure.
+- `scan_order_reversed` is a sequential, non-interleaved JPEG where the first scan does not contain the first component. Associated Lepton file produced by the original implementation.
+- `truncbad` the Lepton format is truncated and invalid.

lib/src/consts.rs

@@ -85,8 +85,8 @@ pub const MAX_THREADS_SUPPORTED_BY_LEPTON_FORMAT: usize = 16; // Number of threa
 pub const EOI: [u8; 2] = [0xFF, jpeg_code::EOI]; // EOI segment
 pub const SOI: [u8; 2] = [0xFF, jpeg_code::SOI]; // SOI segment
 pub const LEPTON_FILE_HEADER: [u8; 2] = [0xcf, 0x84]; // the tau symbol for a tau lepton in utf-8
-pub const LEPTON_HEADER_BASELINE_JPEG_TYPE: [u8; 1] = [b'Z'];
-pub const LEPTON_HEADER_PROGRESSIVE_JPEG_TYPE: [u8; 1] = [b'X'];
+pub const LEPTON_HEADER_SEQUENTIAL_SINGLE_SCAN_JPEG_TYPE: [u8; 1] = [b'Z'];
+pub const LEPTON_HEADER_MULTI_SCAN_JPEG_TYPE: [u8; 1] = [b'X'];
 pub const LEPTON_HEADER_MARKER: [u8; 3] = *b"HDR";
 pub const LEPTON_HEADER_PAD_MARKER: [u8; 3] = *b"P0D";
 pub const LEPTON_HEADER_JPG_RESTARTS_MARKER: [u8; 3] = *b"CRS";

lib/src/jpeg/bit_reader.rs

@@ -245,12 +245,33 @@ impl<R: BufRead> BitReader<R> {
         // can recode it again just by remembering the pad bit.
         self.undo_read_ahead();
 
+        if self.inner.fill_buf()?.is_empty() {
+            self.eof = true;
+            return Ok(());
+        }
+
         let mut h = [0u8; 2];
-        self.inner.read_exact(&mut h)?;
-        if h[0] != 0xff || h[1] != (jpeg_code::RST0 + (self.cpos as u8 & 7)) {
+        self.inner.read_exact(&mut h[0..1])?;
+        if h[0] != 0xff {
+            return err_exit_code(
+                ExitCode::InvalidResetCode,
+                format!(
+                    "reset code does not start with 0xff but with 0x{:x} found in stream",
+                    h[0]
+                ),
+            );
+        }
+
+        if self.inner.fill_buf()?.is_empty() {
+            self.eof = true;
+            return Ok(());
+        }
+
+        self.inner.read_exact(&mut h[1..2])?;
+        if h[1] != (jpeg_code::RST0 + (self.cpos as u8 & 7)) {
             return err_exit_code(
                 ExitCode::InvalidResetCode,
-                format!("invalid reset code {0:x} {1:x} found in stream", h[0], h[1]),
+                format!("invalid reset code 0xff {:x} found in stream", h[1]),
             );
         }
 

lib/src/jpeg/jpeg_header.rs

@@ -33,11 +33,11 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
 
 use std::fmt::Debug;
-use std::io::{Cursor, Read, Write};
+use std::io::{BufRead, Cursor, Read, Write};
 use std::num::NonZeroU32;
 
 use crate::LeptonError;
-use crate::consts::JpegType;
+use crate::consts::{EOI, JpegType};
 use crate::enabled_features::EnabledFeatures;
 use crate::helpers::*;
 use crate::lepton_error::{AddContext, ExitCode, Result, err_exit_code};
@@ -88,17 +88,20 @@ impl RestartSegmentCodingInfo {
 pub struct ReconstructionInfo {
     /// the maximum component in a truncated progressive image.
     ///
-    /// This is meant to be used for progressive images but is not yet implemented.
+    /// This is meant to be used for images that are sequential with multiple scans or are
+    /// progressive, but was never implemented in the original codebase, nor is it implemented here.
     pub max_cmp: u32,
 
     /// the maximum band in a truncated progressive image
     ///
-    /// This is meant to be used for progressive images but is not yet implemented.
+    /// This is meant to be used for images that are sequential with multiple scans or are
+    /// progressive, but was never implemented in the original codebase, nor is it implemented here.
     pub max_bpos: u32,
 
     /// The maximum bit in a truncated progressive image.
     ///
-    /// This is meant to be used for progressive images but is not yet implemented.
+    /// This is meant to be used for images that are sequential with multiple scans or are
+    /// progressive, but was never implemented in the original codebase, nor is it implemented here.
     pub max_sah: u8,
 
     /// the maximum dpos in a truncated image
@@ -137,7 +140,7 @@ pub struct ReconstructionInfo {
     pub garbage_data: Vec<u8>,
 }
 
-pub fn parse_jpeg_header<R: Read>(
+pub fn parse_jpeg_header<R: Read + BufRead>(
     reader: &mut R,
     enabled_features: &EnabledFeatures,
     jpeg_header: &mut JpegHeader,
@@ -152,19 +155,9 @@ pub fn parse_jpeg_header<R: Read>(
 
     let mut mirror = Mirror::new(reader, &mut output_cursor);
 
-    if jpeg_header.parse(&mut mirror, enabled_features).context()? {
-        // append the header if it was not the end of file marker
-        rinfo.raw_jpeg_header.append(&mut output);
-        return Ok(true);
-    } else {
-        // if the output was more than 2 bytes then was a trailing header, so keep that around as well,
-        // but we don't want the EOI since that goes into the garbage data.
-        if output.len() > 2 {
-            rinfo.raw_jpeg_header.extend(&output[0..output.len() - 2]);
-        }
-
-        return Ok(false);
-    }
+    let hit_sos = jpeg_header.parse(&mut mirror, enabled_features).context()?;
+    rinfo.raw_jpeg_header.append(&mut output);
+    return Ok(hit_sos);
 }
 
 // internal utility we use to collect the header that we read for later
@@ -193,6 +186,16 @@ impl<R: Read, W: Write> Read for Mirror<'_, R, W> {
     }
 }
 
+impl<R: BufRead, W: Write> BufRead for Mirror<'_, R, W> {
+    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
+        self.read.fill_buf()
+    }
+
+    fn consume(&mut self, amt: usize) {
+        self.read.consume(amt);
+    }
+}
+
 #[derive(Copy, Clone, Debug)]
 pub(crate) struct HuffCodes {
     pub c_val: [u16; 256],
@@ -546,10 +549,13 @@ impl Default for JpegHeader {
 }
 
 impl JpegHeader {
-    /// true if this image is a single scan, which can be partitioned and decode
+    /// True if this image is sequential with a single scan, which can be partitioned and decoded
     /// completely independently by separate threads. If this is not the case, then
     /// we need to decode the entire image in memory and then encode the JPEG sequentially.
-    pub fn is_single_scan(&self) -> bool {
+    /// In theory progressive images with a single scan (only first DC scan) can get the same
+    /// treatment, though such images practically do not exist. The original codebase does not
+    /// give such images this treatment, and doing so would likely cause incompatibility.
+    pub fn is_sequential_single_scan(&self) -> bool {
         assert!(self.jpeg_type != JpegType::Unknown);
 
         self.jpeg_type == JpegType::Sequential && self.cmpc == self.cs_cmpc
@@ -579,7 +585,7 @@ impl JpegHeader {
     /// until we hit either an SOS (image data) or EOI (end of image).
     ///
     /// Returns false if we hit EOI, true if we have an image to process.
-    pub fn parse<R: Read>(
+    pub fn parse<R: Read + BufRead>(
         &mut self,
         reader: &mut R,
         enabled_features: &EnabledFeatures,
@@ -695,13 +701,24 @@ impl JpegHeader {
     }
 
     // returns true we should continue parsing headers or false if we hit SOS and should stop
-    fn parse_next_segment<R: Read>(
+    fn parse_next_segment<R: Read + BufRead>(
         &mut self,
         reader: &mut R,
         enabled_features: &EnabledFeatures,
     ) -> Result<ParseSegmentResult> {
         let mut header = [0u8; 4];
 
+        if reader.fill_buf()?.starts_with(&EOI) {
+            // don't consume eoi
+            return Ok(ParseSegmentResult::EOI);
+        }
+
+        if reader.fill_buf()?.len() < 2 {
+            // if it's less than two remaining and not EOI,
+            // it's garbage and don't consume it
+            return Ok(ParseSegmentResult::EOI);
+        }
+
         if reader.read(&mut header[0..1]).context()? == 0 {
             // didn't get an EOI
             return Ok(ParseSegmentResult::EOI);

lib/src/jpeg/jpeg_position_state.rs

@@ -4,7 +4,7 @@
  *  This software incorporates material from third parties. See NOTICE.txt for details.
  *--------------------------------------------------------------------------------------------*/
 
-use crate::consts::{JpegDecodeStatus, JpegType};
+use crate::consts::JpegDecodeStatus;
 use crate::lepton_error::{AddContext, ExitCode, err_exit_code};
 use crate::{LeptonError, Result};
 
@@ -86,6 +86,11 @@ impl JpegPositionState {
         self.prev_eobrun = 0;
     }
 
+    pub fn trim_mcu(&mut self, jf: &JpegHeader) {
+        // After updating dpos, update the current MCU to be a fraction of that.
+        self.mcu = self.dpos / (jf.cmp_info[self.cmp].sfv * jf.cmp_info[self.cmp].sfh);
+    }
+
     /// calculates next position (non interleaved)
     fn next_mcu_pos_noninterleaved(&mut self, jf: &JpegHeader) -> JpegDecodeStatus {
         // increment position
@@ -103,11 +108,6 @@ impl JpegPositionState {
             self.dpos = cmp_info.bc;
         }
 
-        // now we've updated dpos, update the current MCU to be a fraction of that
-        if jf.jpeg_type == JpegType::Sequential {
-            self.mcu = self.dpos / (cmp_info.sfv * cmp_info.sfh);
-        }
-
         // check position
         if self.dpos >= cmp_info.bc {
             return JpegDecodeStatus::ScanCompleted;