Refactor cycle and IFD tracking

The previous strategy could not keep track of all cycles such as
reflexive ones at the very start. Additionally, it would fail to be able
to iterate after seeking since it was overly cautious and verified the
next IFD position before allowing the read of a current IFD.

The new implementation is furthermore proof against having multiple
independent chains of IFD which may occur when iterating SubIFDs for
thumbnails. (Note that we do not verify the tree-structure, only a
freedom from cycles).

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();
+        let current_ifd = *next_ifd.as_ref().unwrap();
-        seen_ifds.insert(*next_ifd.as_ref().unwrap());
+        let ifd_offsets = vec![current_ifd];
-        let ifd_offsets = vec![*next_ifd.as_ref().unwrap()];
 
         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(
+        let ifd_offset = self.next_ifd.take().unwrap();
-            &mut self.reader,
+        let ifd = Self::read_ifd(&mut self.reader, self.bigtiff, ifd_offset)?;
-            self.bigtiff,
-            self.next_ifd.take().unwrap(),
-        )?;
 
-        if let Some(next) = ifd.next() {
+        // Ensure this walk does not get us into a cycle.
-            if !self.seen_ifds.insert(next) {
+        self.seen_ifds.insert_next(ifd_offset, ifd.next())?;
-                return Err(TiffError::FormatError(TiffFormatError::CycleInOffsets));
+
-            }
+        // Extend the list of known IFD offsets in this chain, if needed.
-            self.next_ifd = Some(next);
+        if self.ifd_offsets.last().copied() == self.current_ifd {
-            self.ifd_offsets.push(next);
+            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
     }

tests/subsubifds.rs

@@ -1,4 +1,4 @@
-use tiff::{decoder::Decoder, tags::IfdPointer};
+use tiff::decoder::Decoder;
 
 // This file has the following IFD structure:
 //
@@ -20,14 +20,21 @@ 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());
 }