Rollup merge of #136579 - bjorn3:fix_thinvec_ext_ub, r=BoxyUwU

Fix UB in ThinVec::flat_map_in_place

`thin_vec.as_ptr()` goes through the `Deref` impl of `ThinVec`, which will not allow access to any memory as we did call `set_len(0)` first.

Found in the process of investigating https://github.com/rust-lang/rust/issues/135870.
This commit is contained in:
Matthias Krüger 2025-02-27 08:56:36 +01:00 committed by GitHub
commit 18c47ad639
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,4 +1,4 @@
use std::ptr;
use std::{mem, ptr};
use smallvec::{Array, SmallVec};
use thin_vec::ThinVec;
@ -13,39 +13,44 @@ pub trait FlatMapInPlace<T>: Sized {
// The implementation of this method is syntactically identical for all the
// different vector types.
macro_rules! flat_map_in_place {
() => {
($vec:ident $( where T: $bound:path)?) => {
fn flat_map_in_place<F, I>(&mut self, mut f: F)
where
F: FnMut(T) -> I,
I: IntoIterator<Item = T>,
{
struct LeakGuard<'a, T $(: $bound)?>(&'a mut $vec<T>);
impl<'a, T $(: $bound)?> Drop for LeakGuard<'a, T> {
fn drop(&mut self) {
unsafe {
self.0.set_len(0); // make sure we just leak elements in case of panic
}
}
}
let this = LeakGuard(self);
let mut read_i = 0;
let mut write_i = 0;
unsafe {
let mut old_len = self.len();
self.set_len(0); // make sure we just leak elements in case of panic
while read_i < old_len {
while read_i < this.0.len() {
// move the read_i'th item out of the vector and map it
// to an iterator
let e = ptr::read(self.as_ptr().add(read_i));
let e = ptr::read(this.0.as_ptr().add(read_i));
let iter = f(e).into_iter();
read_i += 1;
for e in iter {
if write_i < read_i {
ptr::write(self.as_mut_ptr().add(write_i), e);
ptr::write(this.0.as_mut_ptr().add(write_i), e);
write_i += 1;
} else {
// If this is reached we ran out of space
// in the middle of the vector.
// However, the vector is in a valid state here,
// so we just do a somewhat inefficient insert.
self.set_len(old_len);
self.insert(write_i, e);
old_len = self.len();
self.set_len(0);
this.0.insert(write_i, e);
read_i += 1;
write_i += 1;
@ -54,20 +59,23 @@ macro_rules! flat_map_in_place {
}
// write_i tracks the number of actually written new items.
self.set_len(write_i);
this.0.set_len(write_i);
// The ThinVec is in a sane state again. Prevent the LeakGuard from leaking the data.
mem::forget(this);
}
}
};
}
impl<T> FlatMapInPlace<T> for Vec<T> {
flat_map_in_place!();
flat_map_in_place!(Vec);
}
impl<T, A: Array<Item = T>> FlatMapInPlace<T> for SmallVec<A> {
flat_map_in_place!();
flat_map_in_place!(SmallVec where T: Array);
}
impl<T> FlatMapInPlace<T> for ThinVec<T> {
flat_map_in_place!();
flat_map_in_place!(ThinVec);
}