forked from llvm/llvm-project
-
Notifications
You must be signed in to change notification settings - Fork 18
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[mlir] Fix use-after-free bugs in {RankedTensorType|VectorType}::Buil…
…der (llvm#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
- Loading branch information
Showing
3 changed files
with
192 additions
and
44 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,82 @@ | ||
| //===- ADTExtras.h - Extra ADTs for use in MLIR -----------------*- C++ -*-===// | ||
| // | ||
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. | ||
| // See https://llvm.org/LICENSE.txt for license information. | ||
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | ||
| // | ||
| //===----------------------------------------------------------------------===// | ||
|
|
||
| #ifndef MLIR_SUPPORT_ADTEXTRAS_H | ||
| #define MLIR_SUPPORT_ADTEXTRAS_H | ||
|
|
||
| #include "llvm/ADT/ArrayRef.h" | ||
| #include "llvm/ADT/SmallVector.h" | ||
|
|
||
| namespace mlir { | ||
|
|
||
| //===----------------------------------------------------------------------===// | ||
| // CopyOnWriteArrayRef<T> | ||
| //===----------------------------------------------------------------------===// | ||
|
|
||
| // A wrapper around an ArrayRef<T> that copies to a SmallVector<T> on | ||
| // modification. This is for use in the mlir::<Type>::Builders. | ||
| template <typename T> | ||
| class CopyOnWriteArrayRef { | ||
| public: | ||
| CopyOnWriteArrayRef(ArrayRef<T> array) : nonOwning(array){}; | ||
|
|
||
| CopyOnWriteArrayRef &operator=(ArrayRef<T> array) { | ||
| nonOwning = array; | ||
| owningStorage = {}; | ||
| return *this; | ||
| } | ||
|
|
||
| void insert(size_t index, T value) { | ||
| SmallVector<T> &vector = ensureCopy(); | ||
| vector.insert(vector.begin() + index, value); | ||
| } | ||
|
|
||
| void erase(size_t index) { | ||
| // Note: A copy can be avoided when just dropping the front/back dims. | ||
| if (isNonOwning() && index == 0) { | ||
| nonOwning = nonOwning.drop_front(); | ||
| } else if (isNonOwning() && index == size() - 1) { | ||
| nonOwning = nonOwning.drop_back(); | ||
| } else { | ||
| SmallVector<T> &vector = ensureCopy(); | ||
| vector.erase(vector.begin() + index); | ||
| } | ||
| } | ||
|
|
||
| void set(size_t index, T value) { ensureCopy()[index] = value; } | ||
|
|
||
| size_t size() const { return ArrayRef<T>(*this).size(); } | ||
|
|
||
| bool empty() const { return ArrayRef<T>(*this).empty(); } | ||
|
|
||
| operator ArrayRef<T>() const { | ||
| return nonOwning.empty() ? ArrayRef<T>(owningStorage) : nonOwning; | ||
| } | ||
|
|
||
| private: | ||
| bool isNonOwning() const { return !nonOwning.empty(); } | ||
|
|
||
| SmallVector<T> &ensureCopy() { | ||
| // Empty non-owning storage signals the array has been copied to the owning | ||
| // storage (or both are empty). Note: `nonOwning` should never reference | ||
| // `owningStorage`. This can lead to dangling references if the | ||
| // CopyOnWriteArrayRef<T> is copied. | ||
| if (isNonOwning()) { | ||
| owningStorage = SmallVector<T>(nonOwning); | ||
| nonOwning = {}; | ||
| } | ||
| return owningStorage; | ||
| } | ||
|
|
||
| ArrayRef<T> nonOwning; | ||
| SmallVector<T> owningStorage; | ||
| }; | ||
|
|
||
| } // namespace mlir | ||
|
|
||
| #endif |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters