Skip to content

Commit d9725a5

Browse files
authored
Unrolled build for rust-lang#136579
Rollup merge of rust-lang#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 rust-lang#135870.
2 parents e6059f5 + 3477297 commit d9725a5

File tree

1 file changed

+25
-17
lines changed

1 file changed

+25
-17
lines changed
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use std::ptr;
1+
use std::{mem, ptr};
22

33
use smallvec::{Array, SmallVec};
44
use thin_vec::ThinVec;
@@ -13,39 +13,44 @@ pub trait FlatMapInPlace<T>: Sized {
1313
// The implementation of this method is syntactically identical for all the
1414
// different vector types.
1515
macro_rules! flat_map_in_place {
16-
() => {
16+
($vec:ident $( where T: $bound:path)?) => {
1717
fn flat_map_in_place<F, I>(&mut self, mut f: F)
1818
where
1919
F: FnMut(T) -> I,
2020
I: IntoIterator<Item = T>,
2121
{
22+
struct LeakGuard<'a, T $(: $bound)?>(&'a mut $vec<T>);
23+
24+
impl<'a, T $(: $bound)?> Drop for LeakGuard<'a, T> {
25+
fn drop(&mut self) {
26+
unsafe {
27+
self.0.set_len(0); // make sure we just leak elements in case of panic
28+
}
29+
}
30+
}
31+
32+
let this = LeakGuard(self);
33+
2234
let mut read_i = 0;
2335
let mut write_i = 0;
2436
unsafe {
25-
let mut old_len = self.len();
26-
self.set_len(0); // make sure we just leak elements in case of panic
27-
28-
while read_i < old_len {
37+
while read_i < this.0.len() {
2938
// move the read_i'th item out of the vector and map it
3039
// to an iterator
31-
let e = ptr::read(self.as_ptr().add(read_i));
40+
let e = ptr::read(this.0.as_ptr().add(read_i));
3241
let iter = f(e).into_iter();
3342
read_i += 1;
3443

3544
for e in iter {
3645
if write_i < read_i {
37-
ptr::write(self.as_mut_ptr().add(write_i), e);
46+
ptr::write(this.0.as_mut_ptr().add(write_i), e);
3847
write_i += 1;
3948
} else {
4049
// If this is reached we ran out of space
4150
// in the middle of the vector.
4251
// However, the vector is in a valid state here,
4352
// so we just do a somewhat inefficient insert.
44-
self.set_len(old_len);
45-
self.insert(write_i, e);
46-
47-
old_len = self.len();
48-
self.set_len(0);
53+
this.0.insert(write_i, e);
4954

5055
read_i += 1;
5156
write_i += 1;
@@ -54,20 +59,23 @@ macro_rules! flat_map_in_place {
5459
}
5560

5661
// write_i tracks the number of actually written new items.
57-
self.set_len(write_i);
62+
this.0.set_len(write_i);
63+
64+
// The ThinVec is in a sane state again. Prevent the LeakGuard from leaking the data.
65+
mem::forget(this);
5866
}
5967
}
6068
};
6169
}
6270

6371
impl<T> FlatMapInPlace<T> for Vec<T> {
64-
flat_map_in_place!();
72+
flat_map_in_place!(Vec);
6573
}
6674

6775
impl<T, A: Array<Item = T>> FlatMapInPlace<T> for SmallVec<A> {
68-
flat_map_in_place!();
76+
flat_map_in_place!(SmallVec where T: Array);
6977
}
7078

7179
impl<T> FlatMapInPlace<T> for ThinVec<T> {
72-
flat_map_in_place!();
80+
flat_map_in_place!(ThinVec);
7381
}

0 commit comments

Comments
 (0)