Unverified Commit b0b8e83e authored by Benjamin Maxwell's avatar Benjamin Maxwell Committed by GitHub
Browse files

[mlir] Fix use-after-free bugs in {RankedTensorType|VectorType}::Builder (#68969)

Previously, these would set their ArrayRef members to reference their
storage SmallVectors after a copy-on-write (COW) operation. This leads
to a use-after-free if the builder is copied and the original destroyed
(as the new builder would still reference the old SmallVector).

This could easily accidentally occur in code like (annotated):
```c++
// 1. `VectorType::Builder(type)` constructs a new temporary builder
// 2. `.dropDim(0)` updates the temporary builder by reference, and returns a `VectorType::Builder&`
//    - Modifying the shape is a COW operation, so `storage` is used, and `shape` updated the reference it
// 3. Assigning the reference to `auto` copies the builder (via the default C++ copy ctor)
//    -  There's no special handling for `shape` and `storage`, so the new shape points to the old builder's `storage`
auto newType = VectorType::Builder(type).dropDim(0);
// 4. When this line is reached the original temporary builder is destroyed
//    - Actually constructing the vector type is now a use-after-free
VectorType newVectorType = VectorType(newType);
```

This is fixed with these changes by using `CopyOnWriteArrayRef<T>`,
which implements the same functionality, but ensures no
dangling references are possible if it's copied. 

---

The VectorType::Builder also set the ArrayRef<bool> scalableDims member
to a temporary SmallVector when the provided scalableDims are empty.
This again leads to a use-after-free, and is unnecessary as
VectorType::get already handles being passed an empty scalableDims
array.

These bugs were in-part caught by UBSAN, see:
https://lab.llvm.org/buildbot/#/builders/5/builds/37355
parent 28e4f973
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment