fix: correct HeaderMap iterators' len and size_hint for multi-value (#3971)

* test: add headermap iterator tests

confirmed new tests fail as expected

* fix: correct HeaderMap iterators' len and size_hint for multi-value
headers

fixes #3969

Author: robjtede

actix-http/CHANGES.md

@@ -3,6 +3,7 @@
 ## Unreleased
 
 - Encode the HTTP/1 `Connection: Upgrade` header in Camel-Case when camel-case header formatting is enabled.[#3953]
+- Fix `HeaderMap` iterators' `len()` and `size_hint()` implementations for multi-value headers.
 
 [#3953]: https://github.com/actix/actix-web/pull/3953
 

actix-http/src/header/map.rs

@@ -537,7 +537,7 @@ impl HeaderMap {
     /// assert!(pairs.contains(&(&header::SET_COOKIE, &HeaderValue::from_static("two=2"))));
     /// ```
     pub fn iter(&self) -> Iter<'_> {
-        Iter::new(self.inner.iter())
+        Iter::new(self.inner.iter(), self.len())
     }
 
     /// An iterator over all contained header names.
@@ -626,7 +626,8 @@ impl HeaderMap {
     /// assert!(map.is_empty());
     /// ```
     pub fn drain(&mut self) -> Drain<'_> {
-        Drain::new(self.inner.drain())
+        let len = self.len();
+        Drain::new(self.inner.drain(), len)
     }
 }
 
@@ -638,7 +639,8 @@ impl IntoIterator for HeaderMap {
 
     #[inline]
     fn into_iter(self) -> Self::IntoIter {
-        IntoIter::new(self.inner.into_iter())
+        let len = self.len();
+        IntoIter::new(self.inner.into_iter(), len)
     }
 }
 
@@ -648,7 +650,7 @@ impl<'a> IntoIterator for &'a HeaderMap {
 
     #[inline]
     fn into_iter(self) -> Self::IntoIter {
-        Iter::new(self.inner.iter())
+        Iter::new(self.inner.iter(), self.len())
     }
 }
 
@@ -760,14 +762,16 @@ pub struct Iter<'a> {
     inner: hash_map::Iter<'a, HeaderName, Value>,
     multi_inner: Option<(&'a HeaderName, &'a SmallVec<[HeaderValue; 4]>)>,
     multi_idx: usize,
+    remaining: usize,
 }
 
 impl<'a> Iter<'a> {
-    fn new(iter: hash_map::Iter<'a, HeaderName, Value>) -> Self {
+    fn new(iter: hash_map::Iter<'a, HeaderName, Value>, remaining: usize) -> Self {
         Self {
             inner: iter,
             multi_idx: 0,
             multi_inner: None,
+            remaining,
         }
     }
 }
@@ -781,6 +785,7 @@ impl<'a> Iterator for Iter<'a> {
             match vals.get(self.multi_idx) {
                 Some(val) => {
                     self.multi_idx += 1;
+                    self.remaining -= 1;
                     return Some((name, val));
                 }
                 None => {
@@ -800,9 +805,7 @@ impl<'a> Iterator for Iter<'a> {
 
     #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        // take inner lower bound
-        // make no attempt at an upper bound
-        (self.inner.size_hint().0, None)
+        (self.remaining, Some(self.remaining))
     }
 }
 
@@ -818,14 +821,16 @@ pub struct Drain<'a> {
     inner: hash_map::Drain<'a, HeaderName, Value>,
     multi_inner: Option<(Option<HeaderName>, SmallVec<[HeaderValue; 4]>)>,
     multi_idx: usize,
+    remaining: usize,
 }
 
 impl<'a> Drain<'a> {
-    fn new(iter: hash_map::Drain<'a, HeaderName, Value>) -> Self {
+    fn new(iter: hash_map::Drain<'a, HeaderName, Value>, remaining: usize) -> Self {
         Self {
             inner: iter,
             multi_inner: None,
             multi_idx: 0,
+            remaining,
         }
     }
 }
@@ -838,6 +843,7 @@ impl Iterator for Drain<'_> {
         if let Some((ref mut name, ref mut vals)) = self.multi_inner {
             if !vals.is_empty() {
                 // OPTIMIZE: array removals
+                self.remaining -= 1;
                 return Some((name.take(), vals.remove(0)));
             } else {
                 // no more items in value iterator; reset state
@@ -855,9 +861,7 @@ impl Iterator for Drain<'_> {
 
     #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        // take inner lower bound
-        // make no attempt at an upper bound
-        (self.inner.size_hint().0, None)
+        (self.remaining, Some(self.remaining))
     }
 }
 
@@ -872,13 +876,15 @@ impl iter::FusedIterator for Drain<'_> {}
 pub struct IntoIter {
     inner: hash_map::IntoIter<HeaderName, Value>,
     multi_inner: Option<(HeaderName, smallvec::IntoIter<[HeaderValue; 4]>)>,
+    remaining: usize,
 }
 
 impl IntoIter {
-    fn new(inner: hash_map::IntoIter<HeaderName, Value>) -> Self {
+    fn new(inner: hash_map::IntoIter<HeaderName, Value>, remaining: usize) -> Self {
         Self {
             inner,
             multi_inner: None,
+            remaining,
         }
     }
 }
@@ -891,6 +897,7 @@ impl Iterator for IntoIter {
         if let Some((ref name, ref mut vals)) = self.multi_inner {
             match vals.next() {
                 Some(val) => {
+                    self.remaining -= 1;
                     return Some((name.clone(), val));
                 }
                 None => {
@@ -909,9 +916,7 @@ impl Iterator for IntoIter {
 
     #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        // take inner lower bound
-        // make no attempt at an upper bound
-        (self.inner.size_hint().0, None)
+        (self.remaining, Some(self.remaining))
     }
 }
 
@@ -1160,6 +1165,40 @@ mod tests {
         assert!(vals.next().is_none());
     }
 
+    #[test]
+    fn iter_len_counts_values() {
+        let mut map = HeaderMap::new();
+        map.append(header::SET_COOKIE, HeaderValue::from_static("a=1"));
+        map.append(header::SET_COOKIE, HeaderValue::from_static("b=2"));
+        map.append(header::SET_COOKIE, HeaderValue::from_static("c=3"));
+
+        assert_eq!(map.iter().count(), 3);
+        assert_eq!(map.iter().len(), 3);
+    }
+
+    #[test]
+    fn into_iter_len_counts_values() {
+        let mut map = HeaderMap::new();
+        map.append(header::SET_COOKIE, HeaderValue::from_static("a=1"));
+        map.append(header::SET_COOKIE, HeaderValue::from_static("b=2"));
+        map.append(header::SET_COOKIE, HeaderValue::from_static("c=3"));
+
+        assert_eq!(map.clone().into_iter().count(), 3);
+        assert_eq!(map.into_iter().len(), 3);
+    }
+
+    #[test]
+    fn drain_len_counts_values() {
+        let mut map = HeaderMap::new();
+        map.append(header::SET_COOKIE, HeaderValue::from_static("a=1"));
+        map.append(header::SET_COOKIE, HeaderValue::from_static("b=2"));
+        map.append(header::SET_COOKIE, HeaderValue::from_static("c=3"));
+
+        let mut drained = map.clone();
+        assert_eq!(map.drain().count(), 3);
+        assert_eq!(drained.drain().len(), 3);
+    }
+
     fn owned_pair<'a>((name, val): (&'a HeaderName, &'a HeaderValue)) -> (HeaderName, HeaderValue) {
         (name.clone(), val.clone())
     }