[ADT] Add ASAN Support for SmallVector

[njames93]

Summary

Mark addresses that are currently unused in a SmallVector as poisoned. This can help catch some iterator invalidation bugs when the SmallVector grows.

Running the ADTTest suite found an iterator invalidation bug in SmallVectorImpl<T>::insert(iterator, const T&) when the item to insert is contained inside the SmallVector.

[njames93]

llvm/include/llvm/ADT/SmallVector.h

594–595

Is this check the right behaviour, or is inserting an item from a SmallVector into itself considered UB?
If that would be the right behaviour then there is case to also add this check to append, push_back and resize.
If its classed as UB then the test case SmallVectorTest.PushPopTest needs addressing, specifically this line.

[njames93]

Undo fix for self referencing insertion

[njames93]

Undo SmallVectorTest fix

[RKSimon]

Adding some potential reviewers

[dblaikie]

llvm/include/llvm/ADT/SmallVector.h

222

is attribute unused the right one to use here?

(also LLVM doesn't usually use top-level const on local variables)

I'd probably add a (void)Size; down next to the poison call, like we'd use with variables only used in assertions

594–595

Not sure I understand the question - Was this comment related to a different line of code?

Ah, maybe the 608 "if (I <= EltPtr", etc bit?

I think that's the one @MaskRay pointed out on the self-insert vector review, the commit message for the commit that added this support seems to indicate that the at-the-time libc++ maintainer suggested this was valid for std::vector, and was implemented in a similar way there - but it would be good to check on the implementation strategy and be able to quote some words from the standard to back this up. Could you check those things?

llvm/lib/Support/Twine.cpp

54–55

Do existing LLVM unit tests fail under ASan due to this unpoison being missing? Or is this untested? If it's untested, it'd be good to add some test coverage for it (even if those tests only fail under asan and so don't fail for most/other uses)

Similarly for the SmallString change earlier.

llvm/unittests/ADT/SmallVectorTest.cpp

256–257

Is the intent to fix this in a later commit? Would that then make this change cause asan buildbots/builds to fail when they aren't currently? (that wouldn't be good/acceptable - so if that's the case, we'll need to decide on a fix before this is committed)

[njames93]

In D87237#2270189, @RKSimon wrote:

Adding some potential reviewers

Should point out this is current state is not intended for review, it was for context in a mailing list. Once some other issues are resolved I'll do this using the asan sequential container operations.

llvm/include/llvm/ADT/SmallString.h

public:
 
   /// Explicit conversion to StringRef.
   StringRef str() const { return StringRef(this->data(), this->size()); }
 
   // TODO: Make this const, if it's safe...
   const char* c_str() {
     this->push_back(0);
     this->pop_back();
+    // Unpoison the address of the null terminator.
+    __asan_unpoison_memory_region(this->end(), 1);
     return this->data();
   }
 
   /// Implicit conversion to StringRef.
   operator StringRef() const { return str(); }
 
   explicit operator std::string() const {
     return std::string(this->data(), this->size());

llvm/include/llvm/ADT/SmallVector.h

 template <typename T, bool = (is_trivially_copy_constructible<T>::value) &&
                              (is_trivially_move_constructible<T>::value) &&
                              std::is_trivially_destructible<T>::value>
 class SmallVectorTemplateBase : public SmallVectorTemplateCommon<T> {
 protected:
   SmallVectorTemplateBase(size_t Size) : SmallVectorTemplateCommon<T>(Size) {}
 
   static void destroy_range(T *S, T *E) {
+    const LLVM_ATTRIBUTE_UNUSED size_t Size = E - S;
     while (S != E) {
       --E;
       E->~T();
     }
+    __asan_poison_memory_region(S, Size * sizeof(T));
   }
 
   /// Move the range [I, E) into the uninitialized memory starting with "Dest",
   /// constructing elements as needed.
   template<typename It1, typename It2>
   static void uninitialized_move(It1 I, It1 E, It2 Dest) {
+    __asan_unpoison_memory_region(Dest, std::distance(I, E) * sizeof(T));
     std::uninitialized_copy(std::make_move_iterator(I),
                             std::make_move_iterator(E), Dest);
   }
 
   /// Copy the range [I, E) onto the uninitialized memory starting with "Dest",
   /// constructing elements as needed.
   template<typename It1, typename It2>
   static void uninitialized_copy(It1 I, It1 E, It2 Dest) {
+    __asan_unpoison_memory_region(Dest, std::distance(I, E) * sizeof(T));
     std::uninitialized_copy(I, E, Dest);
   }
 
   /// Grow the allocated memory (without initializing new elements), doubling
   /// the size of the allocated memory. Guarantees space for at least one more
   /// element, or MinSize more elements if specified.
   void grow(size_t MinSize = 0);
 
 public:
   void push_back(const T &Elt) {
     if (LLVM_UNLIKELY(this->size() >= this->capacity()))
       this->grow();
+    __asan_unpoison_memory_region(this->end(), sizeof(T));
     ::new ((void*) this->end()) T(Elt);
     this->set_size(this->size() + 1);
   }
 
   void push_back(T &&Elt) {
     if (LLVM_UNLIKELY(this->size() >= this->capacity()))
       this->grow();
+    __asan_unpoison_memory_region(this->end(), sizeof(T));
     ::new ((void*) this->end()) T(::std::move(Elt));
     this->set_size(this->size() + 1);
   }
 
   void pop_back() {
     this->set_size(this->size() - 1);
     this->end()->~T();
+    __asan_poison_memory_region(this->end(), sizeof(T));
   }
 };
 
 // Define this out-of-line to dissuade the C++ compiler from inlining it.
 template <typename T, bool TriviallyCopyable>
 void SmallVectorTemplateBase<T, TriviallyCopyable>::grow(size_t MinSize) {
   // Ensure we can fit the new capacity.
   // This is only going to be applicable when the capacity is 32 bit.

 /// T's. This allows using memcpy in place of copy/move construction and
 /// skipping destruction.
 template <typename T>
 class SmallVectorTemplateBase<T, true> : public SmallVectorTemplateCommon<T> {
 protected:
   SmallVectorTemplateBase(size_t Size) : SmallVectorTemplateCommon<T>(Size) {}
 
   // No need to do a destroy loop for POD's.
-  static void destroy_range(T *, T *) {}
+  static void destroy_range(T *B, T *E) {
+    __asan_poison_memory_region(B, (E - B) * sizeof(T));
+  }
 
   /// Move the range [I, E) onto the uninitialized memory
   /// starting with "Dest", constructing elements into it as needed.
   template<typename It1, typename It2>
   static void uninitialized_move(It1 I, It1 E, It2 Dest) {
     // Just do a copy.
     uninitialized_copy(I, E, Dest);
   }
 
   /// Copy the range [I, E) onto the uninitialized memory
   /// starting with "Dest", constructing elements into it as needed.
   template<typename It1, typename It2>
   static void uninitialized_copy(It1 I, It1 E, It2 Dest) {
     // Arbitrary iterator types; just use the basic implementation.
+    __asan_unpoison_memory_region(Dest, std::distance(I, E) * sizeof(T));
     std::uninitialized_copy(I, E, Dest);
   }
 
   /// Copy the range [I, E) onto the uninitialized memory
   /// starting with "Dest", constructing elements into it as needed.
   template <typename T1, typename T2>
   static void uninitialized_copy(
       T1 *I, T1 *E, T2 *Dest,
       std::enable_if_t<std::is_same<typename std::remove_const<T1>::type,
                                     T2>::value> * = nullptr) {
     // Use memcpy for PODs iterated by pointers (which includes SmallVector
     // iterators): std::uninitialized_copy optimizes to memmove, but we can
     // use memcpy here. Note that I and E are iterators and thus might be
     // invalid for memcpy if they are equal.
+    __asan_unpoison_memory_region(Dest, (E - I) * sizeof(T));
     if (I != E)
       memcpy(reinterpret_cast<void *>(Dest), I, (E - I) * sizeof(T));
   }
 
   /// Double the size of the allocated memory, guaranteeing space for at
   /// least one more element or MinSize if specified.
   void grow(size_t MinSize = 0) { this->grow_pod(MinSize, sizeof(T)); }
 
 public:
   void push_back(const T &Elt) {
     if (LLVM_UNLIKELY(this->size() >= this->capacity()))
       this->grow();
+    __asan_unpoison_memory_region(this->end(), sizeof(T));
     memcpy(reinterpret_cast<void *>(this->end()), &Elt, sizeof(T));
     this->set_size(this->size() + 1);
   }
 
-  void pop_back() { this->set_size(this->size() - 1); }
+  void pop_back() {
+    this->set_size(this->size() - 1);
+    __asan_poison_memory_region(this->end(), sizeof(T));
+  }
 };
 
 /// This class consists of common code factored out of the SmallVector class to
 /// reduce code duplication based on the SmallVector 'N' template parameter.
 template <typename T>
 class SmallVectorImpl : public SmallVectorTemplateBase<T> {
   using SuperClass = SmallVectorTemplateBase<T>;
 
public:
 
   void resize(size_type N) {
     if (N < this->size()) {
       this->destroy_range(this->begin()+N, this->end());
       this->set_size(N);
     } else if (N > this->size()) {
       if (this->capacity() < N)
         this->grow(N);
+      __asan_unpoison_memory_region(this->end(),
+                                    (N - this->size()) * sizeof(T));
       for (auto I = this->end(), E = this->begin() + N; I != E; ++I)
         new (&*I) T();
       this->set_size(N);
     }
   }
 
   void resize(size_type N, const T &NV) {
     if (N < this->size()) {
       this->destroy_range(this->begin()+N, this->end());
       this->set_size(N);
     } else if (N > this->size()) {
       if (this->capacity() < N)
         this->grow(N);
+      __asan_unpoison_memory_region(this->end(),
+                                    (N - this->size()) * sizeof(T));
       std::uninitialized_fill(this->end(), this->begin()+N, NV);
       this->set_size(N);
     }
   }
 
   void reserve(size_type N) {
     if (this->capacity() < N)
       this->grow(N);
void append(in_iter in_start, in_iter in_end) {
     this->set_size(this->size() + NumInputs);
   }
 
   /// Append \p NumInputs copies of \p Elt to the end.
   void append(size_type NumInputs, const T &Elt) {
     if (NumInputs > this->capacity() - this->size())
       this->grow(this->size()+NumInputs);
 
+    __asan_unpoison_memory_region(this->end(), NumInputs * sizeof(T));
     std::uninitialized_fill_n(this->end(), NumInputs, Elt);
     this->set_size(this->size() + NumInputs);
   }
 
   void append(std::initializer_list<T> IL) {
     append(IL.begin(), IL.end());
   }
 
   // FIXME: Consider assigning over existing elements, rather than clearing &
   // re-initializing them - for all assign(...) variants.
 
   void assign(size_type NumElts, const T &Elt) {
     clear();
     if (this->capacity() < NumElts)
       this->grow(NumElts);
+    __asan_unpoison_memory_region(this->begin(), NumElts * sizeof(T));
     this->set_size(NumElts);
     std::uninitialized_fill(this->begin(), this->end(), Elt);
   }
 
   template <typename in_iter,
             typename = std::enable_if_t<std::is_convertible<
                 typename std::iterator_traits<in_iter>::iterator_category,
                 std::input_iterator_tag>::value>>
iterator insert(iterator I, T &&Elt) {
     assert(I <= this->end() && "Inserting past the end of the vector.");
 
     if (this->size() >= this->capacity()) {
       size_t EltNo = I-this->begin();
       this->grow();
       I = this->begin()+EltNo;
     }
 
+    __asan_unpoison_memory_region(this->end(), sizeof(T));
     ::new ((void*) this->end()) T(::std::move(this->back()));
     // Push everything else over.
     std::move_backward(I, this->end()-1, this->end());
     this->set_size(this->size() + 1);
 
     // If we just moved the element we're inserting, be sure to update
     // the reference.
     T *EltPtr = &Elt;
if (I == this->end()) {  // Important special case for empty vector.
       this->push_back(Elt);
       return this->end()-1;
     }
 
     assert(I >= this->begin() && "Insertion iterator is out of bounds.");
     assert(I <= this->end() && "Inserting past the end of the vector.");
 
     if (this->size() >= this->capacity()) {
-      size_t EltNo = I-this->begin();
+      size_t EltNo = I - this->begin();
       this->grow();
-      I = this->begin()+EltNo;
+      I = this->begin() + EltNo;
     }
+
+    __asan_unpoison_memory_region(this->end(), sizeof(T));
     ::new ((void*) this->end()) T(std::move(this->back()));
     // Push everything else over.
     std::move_backward(I, this->end()-1, this->end());
     this->set_size(this->size() + 1);
 
     // If we just moved the element we're inserting, be sure to update
     // the reference.
     const T *EltPtr = &Elt;
if (size_t(this->end()-I) >= NumToInsert) {
       std::fill_n(I, NumToInsert, Elt);
       return I;
     }
 
     // Otherwise, we're inserting more elements than exist already, and we're
     // not inserting at the end.
 
     // Move over the elements that we're about to overwrite.
+    __asan_unpoison_memory_region(this->end(), NumToInsert * sizeof(T));
     T *OldEnd = this->end();
     this->set_size(this->size() + NumToInsert);
     size_t NumOverwritten = OldEnd-I;
     this->uninitialized_move(I, OldEnd, this->end()-NumOverwritten);
 
     // Replace the overwritten part.
     std::fill_n(I, NumOverwritten, Elt);
 
if (size_t(this->end()-I) >= NumToInsert) {
       return I;
     }
 
     // Otherwise, we're inserting more elements than exist already, and we're
     // not inserting at the end.
 
     // Move over the elements that we're about to overwrite.
     T *OldEnd = this->end();
+    __asan_unpoison_memory_region(this->end(), NumToInsert * sizeof(T));
     this->set_size(this->size() + NumToInsert);
     size_t NumOverwritten = OldEnd-I;
     this->uninitialized_move(I, OldEnd, this->end()-NumOverwritten);
 
     // Replace the overwritten part.
     for (T *J = I; NumOverwritten > 0; --NumOverwritten) {
       *J = *From;
       ++J; ++From;
     }
 
     // Insert the non-overwritten middle part.
     this->uninitialized_copy(From, To, OldEnd);
     return I;
   }
 
   void insert(iterator I, std::initializer_list<T> IL) {
     insert(I, IL.begin(), IL.end());
   }
 
   template <typename... ArgTypes> reference emplace_back(ArgTypes &&... Args) {
     if (LLVM_UNLIKELY(this->size() >= this->capacity()))
       this->grow();
+    __asan_unpoison_memory_region(this->end(), sizeof(T));
     ::new ((void *)this->end()) T(std::forward<ArgTypes>(Args)...);
     this->set_size(this->size() + 1);
     return this->back();
   }
 
   SmallVectorImpl &operator=(const SmallVectorImpl &RHS);
 
   SmallVectorImpl &operator=(SmallVectorImpl &&RHS);
if (this->capacity() < RHSSize) {
     this->destroy_range(this->begin(), this->end());
     this->set_size(0);
     CurSize = 0;
     this->grow(RHSSize);
   } else if (CurSize) {
     // Otherwise, use assignment for the already-constructed elements.
     std::copy(RHS.begin(), RHS.begin()+CurSize, this->begin());
   }
-
+  __asan_unpoison_memory_region(this->begin(), RHSSize * sizeof(T));
   // Copy construct the new elements in place.
   this->uninitialized_copy(RHS.begin()+CurSize, RHS.end(),
                            this->begin()+CurSize);
 
   // Set end.
   this->set_size(RHSSize);
   return *this;
 }
if (this->capacity() < RHSSize) {
     this->destroy_range(this->begin(), this->end());
     this->set_size(0);
     CurSize = 0;
     this->grow(RHSSize);
   } else if (CurSize) {
     // Otherwise, use assignment for the already-constructed elements.
     std::move(RHS.begin(), RHS.begin()+CurSize, this->begin());
   }
-
+  __asan_unpoison_memory_region(this->begin(), RHSSize * sizeof(T));
   // Move-construct the new elements in place.
   this->uninitialized_move(RHS.begin()+CurSize, RHS.end(),
                            this->begin()+CurSize);
 
   // Set end.
   this->set_size(RHSSize);
 
   RHS.clear();
   return *this;
 }
 
 /// Storage for the SmallVector elements.  This is specialized for the N=0 case
 /// to avoid allocating unnecessary storage.
 template <typename T, unsigned N>
 struct SmallVectorStorage {
+#if LLVM_ADDRESS_SANITIZER_BUILD
+  SmallVectorStorage() {
+    __asan_poison_memory_region(&InlineElts, sizeof(InlineElts));
+  }
+  ~SmallVectorStorage() {
+    // Although we don't want to be able to dereference pointers held to
+    // InlineElts, we need to unpoison the memory for the case when the memory
+    // is reused. This is especially important if the SmallVector lives on the
+    // stack.
+    __asan_unpoison_memory_region(&InlineElts, sizeof(InlineElts));
+  }
+#endif
   AlignedCharArrayUnion<T> InlineElts[N];
 };
 
 /// We need the storage to be properly aligned even for small-size of 0 so that
 /// the pointer math in \a SmallVectorTemplateCommon::getFirstEl() is
 /// well-defined.
 template <typename T> struct alignas(alignof(T)) SmallVectorStorage<T, 0> {};
 

llvm/lib/Support/Twine.cpp

if (isUnary()) {
     }
     default:
       break;
     }
   }
   toVector(Out);
   Out.push_back(0);
   Out.pop_back();
+  // Unpoison the address of the null terminator.
+  __asan_unpoison_memory_region(Out.end(), 1);
   return StringRef(Out.data(), Out.size());
 }
 
 void Twine::printOneChild(raw_ostream &OS, Child Ptr,
                           NodeKind Kind) const {
   switch (Kind) {
   case Twine::NullKind: break;
   case Twine::EmptyKind: break;

llvm/unittests/ADT/SmallVectorTest.cpp

TYPED_TEST(SmallVectorTest, PushPopTest) {
   EXPECT_FALSE(this->theVector.begin() == this->theVector.end());
   EXPECT_FALSE(this->theVector.empty());
 
   // Push another element
   this->theVector.push_back(Constructable(2));
   this->assertValuesInOrder(this->theVector, 2u, 1, 2);
 
   // Insert at beginning
+  // TODO: Fix self referencing insertion. Right now this causes asan to fail.
   this->theVector.insert(this->theVector.begin(), this->theVector[1]);
   this->assertValuesInOrder(this->theVector, 3u, 2, 1, 2);
 
   // Pop one element
   this->theVector.pop_back();
   this->assertValuesInOrder(this->theVector, 2u, 2, 1);
 
   // Pop remaining elements
   this->theVector.pop_back();