Merge pull request #279 from image-rs/seeking-multiple-times

Seeking multiple times

src/decoder/cycles.rs

@@ -0,0 +1,169 @@
+use crate::{tags::IfdPointer, TiffError, TiffFormatError};
+use std::collections::HashMap;
+
+/// The IFD structure of a TIFF file should be limited to a forest of entries, and a tree when we
+/// only consider having a primary IFD and its children. There is up to one primary child of a node
+/// that is the `next` entry in an IFD itself. Since we offer iteration over this chain we want to
+/// detect when a new edge would introduce a cycle, such that callers can be guaranteed to
+/// terminate iteration at some point even in malicious images.
+///
+/// However, we are not necessarily visiting the IFDs in their topological order.
+///
+/// Since we only consider one child, we run union find to assign each pointer to a known chain.
+/// Then, when we insert a new edge we check if they belong to the same component.
+#[derive(Default, Debug)]
+pub struct IfdCycles {
+    /// The root of each component union.
+    component_union: HashMap<ComponentId, ComponentId>,
+    links: HashMap<IfdPointer, u64>,
+    chains: HashMap<IfdPointer, ComponentId>,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+struct ComponentId(u64);
+
+impl IfdCycles {
+    pub fn new() -> Self {
+        IfdCycles::default()
+    }
+
+    pub fn insert_next(
+        &mut self,
+        from: IfdPointer,
+        to: Option<IfdPointer>,
+    ) -> Result<bool, TiffError> {
+        let to_offset = to.map_or(0, |p| p.0);
+
+        // For some reason we got
+        match self.links.get(&from) {
+            Some(existing) if *existing == to_offset => return Ok(false),
+            // We got here twice with two different reads of the same IFD. That is .. unusual and
+            // we interpret it as a cycle.
+            Some(_) => return Err(TiffError::FormatError(TiffFormatError::CycleInOffsets)),
+            None => self.links.insert(from, to_offset),
+        };
+
+        self.ensure_node(from);
+
+        if let Some(to) = to {
+            self.ensure_node(to);
+
+            let parent = self.nominal_component(from);
+            let child = self.nominal_component(to);
+
+            if parent == child {
+                return Err(TiffError::FormatError(TiffFormatError::CycleInOffsets));
+            }
+
+            self.component_union.insert(child, parent);
+        }
+
+        Ok(true)
+    }
+
+    fn ensure_node(&mut self, ifd: IfdPointer) {
+        self.chains.entry(ifd).or_insert_with(|| {
+            let id = ComponentId(self.component_union.len() as u64);
+            self.component_union.insert(id, id);
+            id
+        });
+    }
+
+    fn nominal_component(&mut self, node: IfdPointer) -> ComponentId {
+        let id = self.chains[&node];
+
+        let nomimal = {
+            let mut iter = id;
+
+            loop {
+                let parent = self.component_union[&iter];
+
+                if parent == iter {
+                    break parent;
+                }
+
+                iter = parent;
+            }
+        };
+
+        if nomimal != id {
+            // Compress.
+            let mut iter = id;
+
+            loop {
+                let parent = self.component_union[&iter];
+
+                if parent == iter {
+                    break;
+                }
+
+                self.component_union.insert(iter, nomimal);
+                iter = parent;
+            }
+        }
+
+        nomimal
+    }
+}
+
+#[test]
+fn cycles_are_detected() {
+    let mut cycles = IfdCycles::new();
+
+    cycles
+        .insert_next(IfdPointer(0x20), Some(IfdPointer(0x800)))
+        .expect("non-existing link is valid");
+
+    cycles
+        .insert_next(IfdPointer(0x800), Some(IfdPointer(0x20)))
+        .expect_err("cycle must be detected");
+}
+
+#[test]
+fn reflective_cycle() {
+    let mut cycles = IfdCycles::new();
+
+    cycles
+        .insert_next(IfdPointer(0x20), Some(IfdPointer(0x20)))
+        .expect_err("self-referential cycle must be detected");
+}
+
+#[test]
+fn late_cycle() {
+    let mut cycles = IfdCycles::new();
+
+    cycles
+        .insert_next(IfdPointer(0x20), Some(IfdPointer(0x40)))
+        .expect("non-existing link is valid");
+
+    cycles
+        .insert_next(IfdPointer(0x60), Some(IfdPointer(0x80)))
+        .expect("non-existing link is valid");
+    cycles
+        .insert_next(IfdPointer(0x80), Some(IfdPointer(0x20)))
+        .expect("non-existing link is valid");
+
+    cycles
+        .insert_next(IfdPointer(0x40), Some(IfdPointer(0x60)))
+        .expect_err("non-existing link is valid");
+}
+
+#[test]
+fn odd_cycle() {
+    let mut cycles = IfdCycles::new();
+
+    cycles
+        .insert_next(IfdPointer(0x20), Some(IfdPointer(0x40)))
+        .expect("non-existing link is valid");
+
+    cycles
+        .insert_next(IfdPointer(0x60), Some(IfdPointer(0x80)))
+        .expect("non-existing link is valid");
+    cycles
+        .insert_next(IfdPointer(0x80), Some(IfdPointer(0x20)))
+        .expect("non-existing link is valid");
+
+    cycles
+        .insert_next(IfdPointer(0x40), Some(IfdPointer(0x80)))
+        .expect_err("non-existing link is valid");
+}

src/decoder/mod.rs

@@ -1,4 +1,4 @@
-use std::collections::{BTreeMap, HashSet};
+use std::collections::BTreeMap;
 use std::io::{self, Read, Seek};
 
 use crate::tags::{
@@ -14,6 +14,7 @@ use half::f16;
 use self::image::Image;
 use self::stream::{ByteOrder, EndianReader};
 
+mod cycles;
 pub mod ifd;
 mod image;
 mod stream;
@@ -263,9 +264,12 @@ where
     reader: EndianReader<R>,
     bigtiff: bool,
     limits: Limits,
+    current_ifd: Option<IfdPointer>,
     next_ifd: Option<IfdPointer>,
+    /// The IFDs we visited already in this chain of IFDs.
     ifd_offsets: Vec<IfdPointer>,
-    seen_ifds: HashSet<IfdPointer>,
+    /// Map from the ifd into the `ifd_offsets` ordered list.
+    seen_ifds: cycles::IfdCycles,
     image: Image,
 }
 
@@ -495,9 +499,8 @@ impl<R: Read + Seek> Decoder<R> {
         }
         .map(IfdPointer);
 
-        let mut seen_ifds = HashSet::new();
-        seen_ifds.insert(*next_ifd.as_ref().unwrap());
-        let ifd_offsets = vec![*next_ifd.as_ref().unwrap()];
+        let current_ifd = *next_ifd.as_ref().unwrap();
+        let ifd_offsets = vec![current_ifd];
 
         let mut decoder = Decoder {
             reader,
@@ -505,7 +508,8 @@ impl<R: Read + Seek> Decoder<R> {
             limits: Default::default(),
             next_ifd,
             ifd_offsets,
-            seen_ifds,
+            current_ifd: None,
+            seen_ifds: cycles::IfdCycles::new(),
             image: Image {
                 ifd: None,
                 width: 0,
@@ -542,6 +546,11 @@ impl<R: Read + Seek> Decoder<R> {
         self.image().colortype()
     }
 
+    /// The offset of the directory representing the current image.
+    pub fn ifd_pointer(&mut self) -> Option<IfdPointer> {
+        self.current_ifd
+    }
+
     fn image(&self) -> &Image {
         &self.image
     }
@@ -552,6 +561,8 @@ impl<R: Read + Seek> Decoder<R> {
         if ifd_index >= self.ifd_offsets.len() {
             // We possibly need to load in the next IFD
             if self.next_ifd.is_none() {
+                self.current_ifd = None;
+
                 return Err(TiffError::FormatError(
                     TiffFormatError::ImageFileDirectoryNotFound,
                 ));
@@ -574,6 +585,8 @@ impl<R: Read + Seek> Decoder<R> {
         // If the index is within the list of ifds then we can load the selected image/IFD
         if let Some(ifd_offset) = self.ifd_offsets.get(ifd_index) {
             let ifd = Self::read_ifd(&mut self.reader, self.bigtiff, *ifd_offset)?;
+            self.next_ifd = ifd.next();
+            self.current_ifd = Some(*ifd_offset);
             self.image = Image::from_reader(&mut self.reader, ifd, &self.limits, self.bigtiff)?;
 
             Ok(())
@@ -591,20 +604,20 @@ impl<R: Read + Seek> Decoder<R> {
             ));
         }
 
-        let ifd = Self::read_ifd(
-            &mut self.reader,
-            self.bigtiff,
-            self.next_ifd.take().unwrap(),
-        )?;
+        let ifd_offset = self.next_ifd.take().unwrap();
+        let ifd = Self::read_ifd(&mut self.reader, self.bigtiff, ifd_offset)?;
 
-        if let Some(next) = ifd.next() {
-            if !self.seen_ifds.insert(next) {
-                return Err(TiffError::FormatError(TiffFormatError::CycleInOffsets));
-            }
-            self.next_ifd = Some(next);
-            self.ifd_offsets.push(next);
+        // Ensure this walk does not get us into a cycle.
+        self.seen_ifds.insert_next(ifd_offset, ifd.next())?;
+
+        // Extend the list of known IFD offsets in this chain, if needed.
+        if self.ifd_offsets.last().copied() == self.current_ifd {
+            self.ifd_offsets.push(ifd_offset);
         }
 
+        self.current_ifd = Some(ifd_offset);
+        self.next_ifd = ifd.next();
+
         Ok(ifd)
     }
 
@@ -622,7 +635,7 @@ impl<R: Read + Seek> Decoder<R> {
         self.next_ifd.is_some()
     }
 
-    /// Returns the byte_order
+    /// Returns the byte_order of the file.
     pub fn byte_order(&self) -> ByteOrder {
         self.reader.byte_order
     }

src/directory.rs

@@ -158,4 +158,24 @@ mod tests {
             u16::MAX.into()
         );
     }
+
+    #[test]
+    fn iteration_order() {
+        let mut dir = Directory::empty(IfdPointer(0));
+        assert_eq!(dir.len(), 0);
+
+        let fake_entry = Entry::new_u64(crate::tags::Type::BYTE, 0, [0; 8]);
+        dir.extend((0..32).map(|i| {
+            let tag = Tag::Unknown(i);
+            let entry = fake_entry.clone();
+            (tag, entry)
+        }));
+
+        let iter_order: Vec<u16> = dir.iter().map(|(tag, _e)| tag.to_u16()).collect();
+        assert_eq!(
+            iter_order,
+            (0..32).collect::<Vec<_>>(),
+            "Tags must be in ascending order according to the specification"
+        );
+    }
 }

tests/COPYRIGHT

@@ -0,0 +1,4 @@
+subsubifd.tif:
+  url: https://data.kitware.com/api/v1/file/hashsum/sha512/372ca32735c8a8fdbe2286dabead9e779da63f10ba81eda84625e5273f76d74ca1a47a978f67e9c00c12f7f72009c7b2c07a641e643bb0c463812f4ae7f15d6e/download
+  credit: https://github.com/DigitalSlideArchive/tifftools/commit/4d00c70dbce828262f86c1431f7c66b1748965ac
+  license: CC0

tests/subsubifds.rs

@@ -0,0 +1,41 @@
+use tiff::decoder::Decoder;
+
+// This file has the following IFD structure:
+//
+// ```
+// IFD -> IFD -> IFD
+//        |-> IFD
+//        |-> IFD -> IFD
+//        |          \-> IFD -> IFD
+//        \-> IFD
+// ```
+//
+// where the main chain of IFD has 3 images, the second of which has several sub-IFDs and some have
+// more nested IFDs.
+const TEST_IMAGE_SUBIFD: &str = "./tests/subsubifds.tif";
+use std::fs::File;
+
+#[test]
+fn decode_seek_chain() {
+    let file = File::open(TEST_IMAGE_SUBIFD).expect("Cannot open test image");
+    let mut decoder = Decoder::new(file).expect("Invalid format to create decoder");
+
+    // Remember the offsets in the first iteration.
+    let offset0 = decoder.ifd_pointer().expect("First IFD pointer not found");
+    decoder.next_image().unwrap();
+    let offset1 = decoder.ifd_pointer().expect("Second IFD pointer not found");
+    decoder.next_image().unwrap();
+    let offset2 = decoder.ifd_pointer().expect("Third IFD pointer not found");
+    assert!(!decoder.more_images());
+
+    // Ensure we can seek, to do this again.
+    decoder
+        .seek_to_image(0)
+        .expect("Failed to seek to first image");
+    assert_eq!(Some(offset0), decoder.ifd_pointer());
+    decoder.next_image().unwrap();
+    assert_eq!(Some(offset1), decoder.ifd_pointer());
+    decoder.next_image().unwrap();
+    assert_eq!(Some(offset2), decoder.ifd_pointer());
+    assert!(!decoder.more_images());
+}