fix(iter): Fix Unsoundness in Filter/FilterMap::next_chunk

library/core/src/iter/adapters/filter.rs

@@ -41,24 +41,32 @@ where
         let mut array: [MaybeUninit<I::Item>; N] = [const { MaybeUninit::uninit() }; N];
         let mut initialized = 0;
 
+        if N == 0 {
+            // SAFETY: the array is empty, vacuously initialized.
+            return Ok(unsafe { MaybeUninit::array_assume_init(array) });
+        }
+
         let result = self.iter.try_for_each(|element| {
             let idx = initialized;
+            if idx >= N {
+                panic!("closure called after `ControlFlow::Break` was returned");
+            }
             // branchless index update combined with unconditionally copying the value even when
             // it is filtered reduces branching and dependencies in the loop.
             initialized = idx + (self.predicate)(&element) as usize;
-            // SAFETY: Loop conditions ensure the index is in bounds.
+            // SAFETY: idx < N is guaranteed by the check above.
             unsafe { array.get_unchecked_mut(idx) }.write(element);
 
             if initialized < N { ControlFlow::Continue(()) } else { ControlFlow::Break(()) }
         });
 
         match result {
             ControlFlow::Break(()) => {
-                // SAFETY: The loop above is only explicitly broken when the array has been fully initialized
+                // SAFETY: The loop breaks only when initialized == N.
                 Ok(unsafe { MaybeUninit::array_assume_init(array) })
             }
             ControlFlow::Continue(()) => {
-                // SAFETY: The range is in bounds since the loop breaks when reaching N elements.
+                // SAFETY: The range is in bounds: initialized <= N.
                 Err(unsafe { array::IntoIter::new_unchecked(array, 0..initialized) })
             }
         }

library/core/src/iter/adapters/filter_map.rs

@@ -70,6 +70,11 @@ where
     ) -> Result<[Self::Item; N], array::IntoIter<Self::Item, N>> {
         let mut array: [MaybeUninit<Self::Item>; N] = [const { MaybeUninit::uninit() }; N];
 
+        if N == 0 {
+            // SAFETY: the array is empty, vacuously initialized.
+            return Ok(unsafe { MaybeUninit::array_assume_init(array) });
+        }
+
         struct Guard<'a, T> {
             array: &'a mut [MaybeUninit<T>],
             initialized: usize,
@@ -91,10 +96,13 @@ where
 
         let result = self.iter.try_for_each(|element| {
             let idx = guard.initialized;
+            if idx >= N {
+                panic!("closure called after `ControlFlow::Break` was returned");
+            }
             let val = (self.f)(element);
             guard.initialized = idx + val.is_some() as usize;
 
-            // SAFETY: Loop conditions ensure the index is in bounds.
+            // SAFETY: idx < N is guaranteed by the check above.
 
             unsafe {
                 let opt_payload_at: *const MaybeUninit<B> =
@@ -111,12 +119,12 @@ where
 
         match result {
             ControlFlow::Break(()) => {
-                // SAFETY: The loop above is only explicitly broken when the array has been fully initialized
+                // SAFETY: The loop breaks only when initialized == N.
                 Ok(unsafe { MaybeUninit::array_assume_init(array) })
             }
             ControlFlow::Continue(()) => {
                 let initialized = guard.initialized;
-                // SAFETY: The range is in bounds since the loop breaks when reaching N elements.
+                // SAFETY: The range is in bounds: initialized <= N.
                 Err(unsafe { array::IntoIter::new_unchecked(array, 0..initialized) })
             }
         }

library/core/src/iter/traits/iterator.rs

@@ -81,6 +81,14 @@ pub const trait Iterator {
     /// If there are not enough elements to fill the array then `Err` is returned
     /// containing an iterator over the remaining elements.
     ///
+    /// # Panics
+    ///
+    /// May panic if the underlying iterator's [`try_for_each`] calls the closure
+    /// again after it returned [`ControlFlow::Break`].
+    ///
+    /// [`try_for_each`]: Iterator::try_for_each
+    /// [`ControlFlow::Break`]: core::ops::ControlFlow::Break
+    ///
     /// # Examples
     ///
     /// Basic usage:

library/coretests/tests/iter/adapters/filter.rs

@@ -63,3 +63,33 @@ fn test_next_chunk_does_not_leak() {
         assert_eq!(Rc::strong_count(w), 1);
     }
 }
+
+#[test]
+fn test_next_chunk_zero_len() {
+    let result = (0usize..).filter(|_| true).next_chunk::<0>();
+    assert!(result.is_ok());
+    assert_eq!(result.unwrap(), [] as [usize; 0]);
+}
+
+#[test]
+#[should_panic(expected = "closure called after `ControlFlow::Break` was returned")]
+fn test_next_chunk_malicious_try_for_each() {
+    struct OverEager(usize, usize);
+    impl Iterator for OverEager {
+        type Item = usize;
+        fn next(&mut self) -> Option<usize> {
+            unreachable!()
+        }
+        fn try_for_each<F, R>(&mut self, mut f: F) -> R
+        where
+            F: FnMut(usize) -> R,
+            R: std::ops::Try<Output = ()>,
+        {
+            for i in 0..self.1 {
+                let _ = f(self.0 + i);
+            }
+            R::from_output(())
+        }
+    }
+    let _ = OverEager(0, 8).filter(|_| true).next_chunk::<4>();
+}

library/coretests/tests/iter/adapters/filter_map.rs

@@ -48,3 +48,33 @@ fn test_double_ended_filter_map() {
     assert_eq!(it.next().unwrap(), 4);
     assert_eq!(it.next_back(), None);
 }
+
+#[test]
+fn test_next_chunk_zero_len() {
+    let result = (0usize..).filter_map(Some).next_chunk::<0>();
+    assert!(result.is_ok());
+    assert_eq!(result.unwrap(), [] as [usize; 0]);
+}
+
+#[test]
+#[should_panic(expected = "closure called after `ControlFlow::Break` was returned")]
+fn test_next_chunk_malicious_try_for_each() {
+    struct OverEager(usize, usize);
+    impl Iterator for OverEager {
+        type Item = usize;
+        fn next(&mut self) -> Option<usize> {
+            unreachable!()
+        }
+        fn try_for_each<F, R>(&mut self, mut f: F) -> R
+        where
+            F: FnMut(usize) -> R,
+            R: std::ops::Try<Output = ()>,
+        {
+            for i in 0..self.1 {
+                let _ = f(self.0 + i);
+            }
+            R::from_output(())
+        }
+    }
+    let _ = OverEager(0, 8).filter_map(Some).next_chunk::<4>();
+}