From 340ee18468622664d324ca9d43de306cc8a2fab7 Mon Sep 17 00:00:00 2001
From: David Kellum <dek-oss@gravitext.com>
Date: Tue, 1 Dec 2020 12:07:32 -0800
Subject: [PATCH] ValueDrain without RawLinks as potential fix of
 remove_entry_mult

Using RawLinks (to access links of other entries) is suspect for
ValueDrain, particularly in `remove_entry_mult` where `remove_found`
is used to remove the entire target entry (and adjust links
accordingly).

RawLinks was originally introduced to `ValueDrain` in
82d53dbdfdb1ffbeb0323200a0bbd30b5f895fa7.

Note that this introduces further duplication of code in
`remove_extra_value_2` which isn't ideal.
---
 src/header/map.rs | 102 +++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 96 insertions(+), 6 deletions(-)

diff --git a/src/header/map.rs b/src/header/map.rs
index ad6b66cd..31529353 100644
--- a/src/header/map.rs
+++ b/src/header/map.rs
@@ -206,7 +206,6 @@ pub struct ValueIterMut<'a, T> {
 /// An drain iterator of all values associated with a single header name.
 #[derive(Debug)]
 pub struct ValueDrain<'a, T> {
-    raw_links: RawLinks<T>,
     extra_values: *mut Vec<ExtraValue<T>>,
     first: Option<T>,
     next: Option<usize>,
@@ -1192,11 +1191,9 @@ impl<T> HeaderMap<T> {
             links = entry.links.take();
         }
 
-        let raw_links = self.raw_links();
         let extra_values = &mut self.extra_values as *mut _;
 
         ValueDrain {
-            raw_links,
             extra_values,
             first: Some(old),
             next: links.map(|l| l.next),
@@ -1703,6 +1700,101 @@ fn remove_extra_value<T>(mut raw_links: RawLinks<T>, extra_values: &mut Vec<Extr
     extra
 }
 
+/// Removes the `ExtraValue` at the given index.
+#[inline]
+fn remove_extra_value_2<T>(extra_values: &mut Vec<ExtraValue<T>>, idx: usize) -> ExtraValue<T> {
+    let prev;
+    let next;
+
+    {
+        debug_assert!(extra_values.len() > idx);
+        let extra = &extra_values[idx];
+        prev = extra.prev;
+        next = extra.next;
+    }
+
+    // First unlink the extra value
+    match (prev, next) {
+        (Link::Entry(prev), Link::Entry(next)) => {
+            debug_assert_eq!(prev, next);
+        }
+        (Link::Entry(prev), Link::Extra(next)) => {
+            debug_assert!(extra_values.len() > next);
+            extra_values[next].prev = Link::Entry(prev);
+        }
+        (Link::Extra(prev), Link::Entry(next)) => {
+            debug_assert!(extra_values.len() > prev);
+            extra_values[prev].next = Link::Entry(next);
+        }
+        (Link::Extra(prev), Link::Extra(next)) => {
+            debug_assert!(extra_values.len() > next);
+            debug_assert!(extra_values.len() > prev);
+
+            extra_values[prev].next = Link::Extra(next);
+            extra_values[next].prev = Link::Extra(prev);
+        }
+    }
+
+    // Remove the extra value
+    let mut extra = extra_values.swap_remove(idx);
+
+    // This is the index of the value that was moved (possibly `extra`)
+    let old_idx = extra_values.len();
+
+    // Update the links
+    if extra.prev == Link::Extra(old_idx) {
+        extra.prev = Link::Extra(idx);
+    }
+
+    if extra.next == Link::Extra(old_idx) {
+        extra.next = Link::Extra(idx);
+    }
+
+    // Check if another entry was displaced. If it was, then the links
+    // need to be fixed.
+    if idx != old_idx {
+        let next;
+        let prev;
+
+        {
+            debug_assert!(extra_values.len() > idx);
+            let moved = &extra_values[idx];
+            next = moved.next;
+            prev = moved.prev;
+        }
+
+        // An entry was moved, we have to the links
+        match prev {
+            Link::Entry(_) => {
+            }
+            Link::Extra(extra_idx) => {
+                debug_assert!(extra_values.len() > extra_idx);
+                extra_values[extra_idx].next = Link::Extra(idx);
+            }
+        }
+
+        match next {
+            Link::Entry(_) => {
+            }
+            Link::Extra(extra_idx) => {
+                debug_assert!(extra_values.len() > extra_idx);
+                extra_values[extra_idx].prev = Link::Extra(idx);
+            }
+        }
+    }
+
+    debug_assert!({
+        for v in &*extra_values {
+            assert!(v.next != Link::Extra(old_idx));
+            assert!(v.prev != Link::Extra(old_idx));
+        }
+
+        true
+    });
+
+    extra
+}
+
 impl<'a, T> IntoIterator for &'a HeaderMap<T> {
     type Item = (&'a HeaderName, &'a T);
     type IntoIter = Iter<'a, T>;
@@ -2937,10 +3029,8 @@ impl<'a, T> OccupiedEntry<'a, T> {
     /// returned.
     pub fn remove_entry_mult(self) -> (HeaderName, ValueDrain<'a, T>) {
         let entry = self.map.remove_found(self.probe, self.index);
-        let raw_links = self.map.raw_links();
         let extra_values = &mut self.map.extra_values as *mut _;
         let drain = ValueDrain {
-            raw_links,
             extra_values,
             first: Some(entry.value),
             next: entry.links.map(|l| l.next),
@@ -3039,7 +3129,7 @@ impl<'a, T> Iterator for ValueDrain<'a, T> {
         } else if let Some(next) = self.next {
             // Remove the extra value
             let extra = unsafe {
-                remove_extra_value(self.raw_links, &mut *self.extra_values, next)
+                remove_extra_value_2(&mut *self.extra_values, next)
             };
 
             match extra.next {