bjoernager/rust
bjoernager/rust
1
Fork 0
Code Activity

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

Browse source 
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
1 changed files with 25 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

42
compiler/rustc_data_structures/src/flat_map_in_place.rs
View file

@ -1,4 +1,4 @@
use std::ptr; use std::{mem, ptr};
use smallvec::{Array, SmallVec}; use smallvec::{Array, SmallVec};
use thin_vec::ThinVec; use thin_vec::ThinVec;
@ -13,39 +13,44 @@ pub trait FlatMapInPlace<T>: Sized {
// The implementation of this method is syntactically identical for all the // The implementation of this method is syntactically identical for all the
// different vector types. // different vector types.
macro_rules! flat_map_in_place { macro_rules! flat_map_in_place {
() => { ($vec:ident $( where T: $bound:path)?) => {
fn flat_map_in_place<F, I>(&mut self, mut f: F) fn flat_map_in_place<F, I>(&mut self, mut f: F)
where where
F: FnMut(T) -> I, F: FnMut(T) -> I,
I: IntoIterator<Item = T>, 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 read_i = 0;
let mut write_i = 0; let mut write_i = 0;
unsafe { unsafe {
let mut old_len = self.len(); while read_i < this.0.len() {
self.set_len(0); // make sure we just leak elements in case of panic
while read_i < old_len {
// move the read_i'th item out of the vector and map it // move the read_i'th item out of the vector and map it
// to an iterator // 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(); let iter = f(e).into_iter();
read_i += 1; read_i += 1;
for e in iter { for e in iter {
if write_i < read_i { 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; write_i += 1;
} else { } else {
// If this is reached we ran out of space // If this is reached we ran out of space
// in the middle of the vector. // in the middle of the vector.
// However, the vector is in a valid state here, // However, the vector is in a valid state here,
// so we just do a somewhat inefficient insert. // so we just do a somewhat inefficient insert.
self.set_len(old_len); this.0.insert(write_i, e);
self.insert(write_i, e);
old_len = self.len();
self.set_len(0);
read_i += 1; read_i += 1;
write_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. // 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> { 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> { 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> { impl<T> FlatMapInPlace<T> for ThinVec<T> {
flat_map_in_place!(); flat_map_in_place!(ThinVec);
} }