Fix double free in IntoIter

declanvk · declanvk · commit 1c4664957b53 · 2024-08-27T17:28:12.000-07:00
**Description**
 - The previous implementation of `deallocate_leaves` was faulty
   because it would read all the way through the linked list of leaf
   nodes even if some leaves had been popped off the back by the
   iterator.
    - Fixing this by changing the `deallocate_leaves` take an
      iterator with start and end, instead of just start pointer.
 - Added some missing safety docs
 - Modified the `IntoIter` fuzz tests to iterate on the front and the
   back
 - Added new unit test to ensure drop count was correct

**Testing Done**
Ran fuzzer and full test script
diff --git a/fuzz/fuzz_targets/fuzz_tree_map_api.rs b/fuzz/fuzz_targets/fuzz_tree_map_api.rs
@@ -44,7 +44,7 @@ enum Action {
     EntryRef(EntryAction, Box<[u8]>),
     Fuzzy(Box<[u8]>),
     Prefix(Box<[u8]>),
-    IntoIter(usize),
+    IntoIter { take_front: usize, take_back: usize },
 }
 
 libfuzzer_sys::fuzz_target!(|actions: Vec<Action>| {
@@ -221,10 +221,23 @@ libfuzzer_sys::fuzz_target!(|actions: Vec<Action>| {
                 let v: Vec<_> = tree.prefix(&key).collect();
                 std::hint::black_box(v);
             },
-            Action::IntoIter(take) => {
+            Action::IntoIter {
+                take_front,
+                take_back,
+            } => {
                 let tree = mem::replace(&mut tree, TreeMap::<_, u32>::new());
-                let num_entries = tree.into_iter().take(take).count();
-                assert!(num_entries <= take);
+
+                let original_size = tree.len();
+                let mut it = tree.into_iter();
+                let front_count = it.by_ref().take(take_front).count();
+                let back_count = it.rev().take(take_back).count();
+
+                assert_eq!(
+                    front_count + back_count,
+                    original_size.min(take_front.saturating_add(take_back))
+                );
+                assert!(front_count <= take_front);
+                assert!(back_count <= take_back);
             },
         }
     }
diff --git a/src/collections/map/iterators/into_iter.rs b/src/collections/map/iterators/into_iter.rs
@@ -1,4 +1,4 @@
-use std::mem::ManuallyDrop;
+use std::mem::{self, ManuallyDrop};
 
 use crate::{deallocate_leaves, deallocate_tree_non_leaves, NodePtr, RawIterator, TreeMap};
 
@@ -17,13 +17,11 @@ pub struct IntoIter<K, V, const PREFIX_LEN: usize> {
 
 impl<K, V, const PREFIX_LEN: usize> Drop for IntoIter<K, V, PREFIX_LEN> {
     fn drop(&mut self) {
-        if let Some(next) = unsafe { self.inner.next() } {
-            // SAFETY: TODO
-            unsafe { deallocate_leaves(next) }
-        }
+        // SAFETY: The `deallocate_tree_non_leaves` function is called earlier on the
+        // trie (which we have unique access to), so the leaves will still be allocated.
+        unsafe { deallocate_leaves(mem::replace(&mut self.inner, RawIterator::empty())) }
 
-        // Just to be safe, clear the iterator
-        self.inner = RawIterator::empty();
+        // Just to be safe, clear the iterator size
         self.size = 0;
     }
 }
@@ -35,6 +33,8 @@ impl<K, V, const PREFIX_LEN: usize> IntoIter<K, V, PREFIX_LEN> {
         let tree = ManuallyDrop::new(tree);
 
         if let Some(state) = &tree.state {
+            // SAFETY: By construction (and maintained on insert/delete), the `min_leaf` is
+            // always before or equal to `max_leaf` in the leaf node order.
             let inner = unsafe { RawIterator::new(state.min_leaf, state.max_leaf) };
 
             // SAFETY: Since this function takes an owned `TreeMap`, we can assume there is
@@ -58,11 +58,13 @@ impl<K, V, const PREFIX_LEN: usize> Iterator for IntoIter<K, V, PREFIX_LEN> {
     type Item = (K, V);
 
     fn next(&mut self) -> Option<Self::Item> {
-        // SAFETY: TODO
+        // SAFETY: By construction of the `IntoIter` we know that we have ownership over
+        // the trie, and there will be no other concurrent access (read or write).
         let leaf_ptr = unsafe { self.inner.next() };
 
         leaf_ptr.map(|leaf_ptr| {
-            // SAFETY: TODO
+            // SAFETY: This function is only called once for a given `leaf_ptr` since the
+            // iterator will never repeat
             unsafe { NodePtr::deallocate_node_ptr(leaf_ptr) }.into_entry()
         })
     }
@@ -74,20 +76,24 @@ impl<K, V, const PREFIX_LEN: usize> Iterator for IntoIter<K, V, PREFIX_LEN> {
 
 impl<K, V, const PREFIX_LEN: usize> DoubleEndedIterator for IntoIter<K, V, PREFIX_LEN> {
     fn next_back(&mut self) -> Option<Self::Item> {
-        // SAFETY: TODO
+        // SAFETY: By construction of the `IntoIter` we know that we have ownership over
+        // the trie, and there will be no other concurrent access (read or write).
         let leaf_ptr = unsafe { self.inner.next_back() };
 
         leaf_ptr.map(|leaf_ptr| {
-            // SAFETY: TODO
+            // SAFETY: This function is only called once for a given `leaf_ptr` since the
+            // iterator will never repeat
             unsafe { NodePtr::deallocate_node_ptr(leaf_ptr) }.into_entry()
         })
     }
 }
 
-/// An owning iterator over the keys of a `TreeMap`.
+impl<K, V, const PREFIX_LEN: usize> ExactSizeIterator for IntoIter<K, V, PREFIX_LEN> {}
+
+/// An owning iterator over the keys of a [`TreeMap`].
 ///
 /// This `struct` is created by the [`crate::TreeMap::into_keys`] method on
-/// `TreeMap`. See its documentation for more.
+/// [`TreeMap`]. See its documentation for more.
 pub struct IntoKeys<K, V, const PREFIX_LEN: usize>(IntoIter<K, V, PREFIX_LEN>);
 
 impl<K, V, const PREFIX_LEN: usize> IntoKeys<K, V, PREFIX_LEN> {
@@ -102,6 +108,10 @@ impl<K, V, const PREFIX_LEN: usize> Iterator for IntoKeys<K, V, PREFIX_LEN> {
     fn next(&mut self) -> Option<Self::Item> {
         Some(self.0.next()?.0)
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.0.size_hint()
+    }
 }
 
 impl<K, V, const PREFIX_LEN: usize> DoubleEndedIterator for IntoKeys<K, V, PREFIX_LEN> {
@@ -110,9 +120,11 @@ impl<K, V, const PREFIX_LEN: usize> DoubleEndedIterator for IntoKeys<K, V, PREFI
     }
 }
 
-/// An owning iterator over the values of a `TreeMap`.
+impl<K, V, const PREFIX_LEN: usize> ExactSizeIterator for IntoKeys<K, V, PREFIX_LEN> {}
+
+/// An owning iterator over the values of a [`TreeMap`].
 ///
-/// This `struct` is created by the [`into_values`] method on `TreeMap`.
+/// This `struct` is created by the [`into_values`] method on [`TreeMap`].
 /// See its documentation for more.
 ///
 /// [`into_values`]: crate::TreeMap::into_values
@@ -130,10 +142,153 @@ impl<K, V, const PREFIX_LEN: usize> Iterator for IntoValues<K, V, PREFIX_LEN> {
     fn next(&mut self) -> Option<Self::Item> {
         Some(self.0.next()?.1)
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.0.size_hint()
+    }
 }
 
 impl<K, V, const PREFIX_LEN: usize> DoubleEndedIterator for IntoValues<K, V, PREFIX_LEN> {
     fn next_back(&mut self) -> Option<Self::Item> {
         Some(self.0.next_back()?.1)
     }
 }
+
+impl<K, V, const PREFIX_LEN: usize> ExactSizeIterator for IntoValues<K, V, PREFIX_LEN> {}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    };
+
+    use crate::{AsBytes, NoPrefixesBytes, OrderedBytes};
+
+    use super::*;
+
+    #[derive(Debug)]
+    struct DropCounter<T>(Arc<AtomicUsize>, T);
+
+    impl<T> DropCounter<T> {
+        fn new(counter: &Arc<AtomicUsize>, value: T) -> Self {
+            counter.fetch_add(1, Ordering::Relaxed);
+            DropCounter(Arc::clone(counter), value)
+        }
+    }
+
+    impl<T> Drop for DropCounter<T> {
+        fn drop(&mut self) {
+            self.0.fetch_sub(1, Ordering::Relaxed);
+        }
+    }
+
+    impl<T: Ord> Ord for DropCounter<T> {
+        fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+            self.1.cmp(&other.1)
+        }
+    }
+
+    impl<T: PartialOrd> PartialOrd for DropCounter<T> {
+        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+            self.1.partial_cmp(&other.1)
+        }
+    }
+
+    impl<T: Eq> Eq for DropCounter<T> {}
+
+    impl<T: PartialEq> PartialEq for DropCounter<T> {
+        fn eq(&self, other: &Self) -> bool {
+            self.1 == other.1
+        }
+    }
+
+    impl<T: AsBytes> AsBytes for DropCounter<T> {
+        fn as_bytes(&self) -> &[u8] {
+            self.1.as_bytes()
+        }
+    }
+
+    unsafe impl<T: NoPrefixesBytes> NoPrefixesBytes for DropCounter<T> {}
+    unsafe impl<T: OrderedBytes + Ord> OrderedBytes for DropCounter<T> {}
+
+    fn setup_will_deallocate_unconsumed_iter_values(
+        drop_counter: &Arc<AtomicUsize>,
+    ) -> TreeMap<DropCounter<[u8; 3]>, usize> {
+        fn swap<A, B>((a, b): (A, B)) -> (B, A) {
+            (b, a)
+        }
+
+        assert_eq!(drop_counter.load(Ordering::Relaxed), 0);
+
+        [
+            [0u8, 0u8, 0u8],
+            [0, 0, u8::MAX],
+            [0, u8::MAX, 0],
+            [0, u8::MAX, u8::MAX],
+            [u8::MAX, 0, 0],
+            [u8::MAX, 0, u8::MAX],
+            [u8::MAX, u8::MAX, 0],
+            [u8::MAX, u8::MAX, u8::MAX],
+        ]
+        .into_iter()
+        .map(|arr| DropCounter::new(&drop_counter, arr))
+        .enumerate()
+        .map(swap)
+        .collect()
+    }
+
+    fn check_will_deallocate_unconsumed_iter_values(
+        drop_counter: &Arc<AtomicUsize>,
+        mut iter: impl Iterator + DoubleEndedIterator + ExactSizeIterator,
+    ) {
+        assert_eq!(drop_counter.load(Ordering::Relaxed), 8);
+
+        assert_eq!(iter.len(), 8);
+
+        assert_eq!(drop_counter.load(Ordering::Relaxed), 8);
+
+        let _ = iter.next().unwrap();
+        let _ = iter.next_back().unwrap();
+
+        assert_eq!(drop_counter.load(Ordering::Relaxed), 6);
+
+        drop(iter);
+
+        assert_eq!(drop_counter.load(Ordering::Relaxed), 0);
+    }
+
+    #[test]
+    fn into_iter_will_deallocate_unconsumed_iter_values() {
+        let drop_counter = Arc::new(AtomicUsize::new(0));
+
+        let tree = setup_will_deallocate_unconsumed_iter_values(&drop_counter);
+
+        check_will_deallocate_unconsumed_iter_values(&drop_counter, tree.into_iter());
+    }
+
+    #[test]
+    fn into_keys_will_deallocate_unconsumed_iter_values() {
+        let drop_counter = Arc::new(AtomicUsize::new(0));
+
+        let tree = setup_will_deallocate_unconsumed_iter_values(&drop_counter);
+
+        check_will_deallocate_unconsumed_iter_values(&drop_counter, tree.into_keys());
+    }
+
+    #[test]
+    fn into_values_will_deallocate_unconsumed_iter_values() {
+        let drop_counter = Arc::new(AtomicUsize::new(0));
+
+        let tree = setup_will_deallocate_unconsumed_iter_values(&drop_counter);
+
+        check_will_deallocate_unconsumed_iter_values(&drop_counter, tree.into_values());
+    }
+
+    #[test]
+    fn empty_tree_empty_iterator() {
+        let tree: TreeMap<u8, usize> = TreeMap::new();
+        let mut it = tree.into_iter();
+        assert_eq!(it.next(), None);
+    }
+}
diff --git a/src/nodes/iterator.rs b/src/nodes/iterator.rs
@@ -9,6 +9,17 @@ struct RawIteratorInner<K, V, const PREFIX_LEN: usize> {
     end: LeafPtr<K, V, PREFIX_LEN>,
 }
 
+impl<K, V, const PREFIX_LEN: usize> Copy for RawIteratorInner<K, V, PREFIX_LEN> {}
+
+impl<K, V, const PREFIX_LEN: usize> Clone for RawIteratorInner<K, V, PREFIX_LEN> {
+    fn clone(&self) -> Self {
+        Self {
+            start: self.start.clone(),
+            end: self.end.clone(),
+        }
+    }
+}
+
 impl<K, V, const PREFIX_LEN: usize> fmt::Debug for RawIteratorInner<K, V, PREFIX_LEN> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("RawIteratorInner")
@@ -27,6 +38,16 @@ pub struct RawIterator<K, V, const PREFIX_LEN: usize> {
     state: Option<RawIteratorInner<K, V, PREFIX_LEN>>,
 }
 
+impl<K, V, const PREFIX_LEN: usize> Copy for RawIterator<K, V, PREFIX_LEN> {}
+
+impl<K, V, const PREFIX_LEN: usize> Clone for RawIterator<K, V, PREFIX_LEN> {
+    fn clone(&self) -> Self {
+        Self {
+            state: self.state.clone(),
+        }
+    }
+}
+
 impl<K, V, const PREFIX_LEN: usize> fmt::Debug for RawIterator<K, V, PREFIX_LEN> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("RawIterator")
diff --git a/src/nodes/operations/deallocate.rs b/src/nodes/operations/deallocate.rs
@@ -1,6 +1,7 @@
-use crate::{ConcreteNodePtr, InnerNode, LeafNode, NodePtr, OpaqueNodePtr};
+use crate::{ConcreteNodePtr, InnerNode, LeafNode, NodePtr, OpaqueNodePtr, RawIterator};
 
-/// Deallocate all the leaf nodes in the linked list starting from `start`.
+/// Deallocate all the leaf nodes in the linked list starting that are within
+/// the given iterator.
 ///
 /// # Safety
 ///  - This function must only be called once for this `start` node and all
@@ -9,19 +10,12 @@ use crate::{ConcreteNodePtr, InnerNode, LeafNode, NodePtr, OpaqueNodePtr};
 ///  - This function should not be called concurrently with any read of the
 ///    tree, otherwise it could result in a use-after-free.
 pub unsafe fn deallocate_leaves<K, V, const PREFIX_LEN: usize>(
-    start: NodePtr<PREFIX_LEN, LeafNode<K, V, PREFIX_LEN>>,
+    mut leaf_range: RawIterator<K, V, PREFIX_LEN>,
 ) {
-    let mut current = start;
-
-    loop {
+    // SAFETY: Covered by function safety doc
+    while let Some(leaf_ptr) = unsafe { leaf_range.next() } {
         // SAFETY: Covered by function safety doc
-        let leaf = unsafe { NodePtr::deallocate_node_ptr(current) };
-
-        if let Some(next) = leaf.next {
-            current = next;
-        } else {
-            break;
-        }
+        let _ = unsafe { NodePtr::deallocate_node_ptr(leaf_ptr) };
     }
 }
 
