From eacacdc4b15d92b949e2081bcf0017198143fbd2 Mon Sep 17 00:00:00 2001 From: Tony Kan Date: Fri, 13 Mar 2026 13:19:52 -0700 Subject: [PATCH] fix(iter): Panic on contract violation in `Filter`/`FilterMap::next_chunk` Co-authored-by: Kyle Ehrlich <4015993+KyleDavidE@users.noreply.github.com> --- library/core/src/iter/adapters/filter.rs | 14 +++++++-- library/core/src/iter/adapters/filter_map.rs | 14 +++++++-- library/core/src/iter/traits/iterator.rs | 8 +++++ .../coretests/tests/iter/adapters/filter.rs | 30 +++++++++++++++++++ .../tests/iter/adapters/filter_map.rs | 30 +++++++++++++++++++ 5 files changed, 90 insertions(+), 6 deletions(-) diff --git a/library/core/src/iter/adapters/filter.rs b/library/core/src/iter/adapters/filter.rs index b22419ccf080a..19b9b8489e1ad 100644 --- a/library/core/src/iter/adapters/filter.rs +++ b/library/core/src/iter/adapters/filter.rs @@ -41,12 +41,20 @@ where let mut array: [MaybeUninit; 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(()) } @@ -54,11 +62,11 @@ 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(()) => { - // 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) }) } } diff --git a/library/core/src/iter/adapters/filter_map.rs b/library/core/src/iter/adapters/filter_map.rs index 24ec6b1741ce1..86b25453b3388 100644 --- a/library/core/src/iter/adapters/filter_map.rs +++ b/library/core/src/iter/adapters/filter_map.rs @@ -70,6 +70,11 @@ where ) -> Result<[Self::Item; N], array::IntoIter> { let mut array: [MaybeUninit; 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], 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 = @@ -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) }) } } diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs index 2e82f45771823..09322ed7a54ab 100644 --- a/library/core/src/iter/traits/iterator.rs +++ b/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: diff --git a/library/coretests/tests/iter/adapters/filter.rs b/library/coretests/tests/iter/adapters/filter.rs index 167851e33336e..0b4cddc79fc1f 100644 --- a/library/coretests/tests/iter/adapters/filter.rs +++ b/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 { + unreachable!() + } + fn try_for_each(&mut self, mut f: F) -> R + where + F: FnMut(usize) -> R, + R: std::ops::Try, + { + for i in 0..self.1 { + let _ = f(self.0 + i); + } + R::from_output(()) + } + } + let _ = OverEager(0, 8).filter(|_| true).next_chunk::<4>(); +} diff --git a/library/coretests/tests/iter/adapters/filter_map.rs b/library/coretests/tests/iter/adapters/filter_map.rs index 46738eda63f3d..07f2611847960 100644 --- a/library/coretests/tests/iter/adapters/filter_map.rs +++ b/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 { + unreachable!() + } + fn try_for_each(&mut self, mut f: F) -> R + where + F: FnMut(usize) -> R, + R: std::ops::Try, + { + for i in 0..self.1 { + let _ = f(self.0 + i); + } + R::from_output(()) + } + } + let _ = OverEager(0, 8).filter_map(Some).next_chunk::<4>(); +}