Fix use after free in BinaryStream library.

This was reported by the ASAN bot, and it turned out to be
a fairly fundamental problem with the design of VarStreamArray
and the way it passes context information to the extractor.

The fix was cumbersome, and I'm not entirely pleased with it,
so I plan to revisit this design in the future when I'm not
pressed to get the bots green again.  For now, this fixes
the issue by storing the context information by value instead
of by reference, and introduces some impossibly-confusing
template magic to make things "work".

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@301999 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm/DebugInfo/CodeView/CVRecord.h

@@ -53,7 +53,7 @@ struct VarStreamArrayExtractor<codeview::CVRecord<Kind>> {
   typedef void ContextType;
 
   static Error extract(BinaryStreamRef Stream, uint32_t &Len,
-                       codeview::CVRecord<Kind> &Item, void *Ctx) {
+                       codeview::CVRecord<Kind> &Item) {
     using namespace codeview;
     const RecordPrefix *Prefix = nullptr;
     BinaryStreamReader Reader(Stream);

include/llvm/DebugInfo/CodeView/ModuleDebugFileChecksumFragment.h

@@ -35,7 +35,7 @@ public:
   typedef void ContextType;
 
   static Error extract(BinaryStreamRef Stream, uint32_t &Len,
-                       codeview::FileChecksumEntry &Item, void *Ctx);
+                       codeview::FileChecksumEntry &Item);
 };
 }
 
@@ -55,8 +55,8 @@ public:
 
   Error initialize(BinaryStreamReader Reader);
 
-  Iterator begin() const { return Checksums.begin(); }
-  Iterator end() const { return Checksums.end(); }
+  Iterator begin() { return Checksums.begin(); }
+  Iterator end() { return Checksums.end(); }
 
   const FileChecksumArray &getArray() const { return Checksums; }
 

include/llvm/DebugInfo/CodeView/ModuleDebugFragmentRecord.h

@@ -57,8 +57,6 @@ private:
   ModuleDebugFragment &Frag;
 };
 
-typedef VarStreamArray<ModuleDebugFragmentRecord> ModuleDebugFragmentArray;
-
 } // namespace codeview
 
 template <>
@@ -66,13 +64,17 @@ struct VarStreamArrayExtractor<codeview::ModuleDebugFragmentRecord> {
   typedef void ContextType;
 
   static Error extract(BinaryStreamRef Stream, uint32_t &Length,
-                       codeview::ModuleDebugFragmentRecord &Info, void *Ctx) {
+                       codeview::ModuleDebugFragmentRecord &Info) {
     if (auto EC = codeview::ModuleDebugFragmentRecord::initialize(Stream, Info))
       return EC;
     Length = Info.getRecordLength();
     return Error::success();
   }
 };
+
+namespace codeview {
+typedef VarStreamArray<ModuleDebugFragmentRecord> ModuleDebugFragmentArray;
+}
 } // namespace llvm
 
 #endif // LLVM_DEBUGINFO_CODEVIEW_MODULEDEBUGFRAGMENTRECORD_H

include/llvm/DebugInfo/CodeView/ModuleDebugInlineeLinesFragment.h

@@ -42,11 +42,10 @@ struct InlineeSourceLine {
 }
 
 template <> struct VarStreamArrayExtractor<codeview::InlineeSourceLine> {
-  typedef codeview::ModuleDebugInlineeLineFragmentRef ContextType;
+  typedef bool ContextType;
 
   static Error extract(BinaryStreamRef Stream, uint32_t &Len,
-                       codeview::InlineeSourceLine &Item,
-                       ContextType *Fragment);
+                       codeview::InlineeSourceLine &Item, bool HasExtraFiles);
 };
 
 namespace codeview {

include/llvm/DebugInfo/CodeView/ModuleDebugLineFragment.h

@@ -61,10 +61,10 @@ struct LineColumnEntry {
 
 class LineColumnExtractor {
 public:
-  typedef const LineFragmentHeader ContextType;
+  typedef const LineFragmentHeader *ContextType;
 
   static Error extract(BinaryStreamRef Stream, uint32_t &Len,
-                       LineColumnEntry &Item, const LineFragmentHeader *Header);
+                       LineColumnEntry &Item, const LineFragmentHeader *Ctx);
 };
 
 class ModuleDebugLineFragmentRef final : public ModuleDebugFragmentRef {

include/llvm/DebugInfo/PDB/Native/DbiModuleDescriptor.h

@@ -66,7 +66,7 @@ struct ModuleInfoEx {
 template <> struct VarStreamArrayExtractor<pdb::DbiModuleDescriptor> {
   typedef void ContextType;
   static Error extract(BinaryStreamRef Stream, uint32_t &Length,
-                       pdb::DbiModuleDescriptor &Info, void *Ctx) {
+                       pdb::DbiModuleDescriptor &Info) {
     if (auto EC = pdb::DbiModuleDescriptor::initialize(Stream, Info))
       return EC;
     Length = Info.getRecordLength();

include/llvm/Support/BinaryStreamArray.h

@@ -42,99 +42,34 @@ namespace llvm {
 /// having to specify a second template argument to VarStreamArray (documented
 /// below).
 template <typename T> struct VarStreamArrayExtractor {
-  typedef void Context;
+  struct ContextType {};
 
   // Method intentionally deleted.  You must provide an explicit specialization
-  // with the following method implemented.
+  // with one of the following two methods implemented.
+  static Error extract(BinaryStreamRef Stream, uint32_t &Len, T &Item) = delete;
+
   static Error extract(BinaryStreamRef Stream, uint32_t &Len, T &Item,
-                       Context *Ctx) = delete;
+                       const ContextType &Ctx) = delete;
 };
 
-/// VarStreamArray represents an array of variable length records backed by a
-/// stream.  This could be a contiguous sequence of bytes in memory, it could
-/// be a file on disk, or it could be a PDB stream where bytes are stored as
-/// discontiguous blocks in a file.  Usually it is desirable to treat arrays
-/// as contiguous blocks of memory, but doing so with large PDB files, for
-/// example, could mean allocating huge amounts of memory just to allow
-/// re-ordering of stream data to be contiguous before iterating over it.  By
-/// abstracting this out, we need not duplicate this memory, and we can
-/// iterate over arrays in arbitrarily formatted streams.  Elements are parsed
-/// lazily on iteration, so there is no upfront cost associated with building
-/// or copying a VarStreamArray, no matter how large it may be.
-///
-/// You create a VarStreamArray by specifying a ValueType and an Extractor type.
-/// If you do not specify an Extractor type, you are expected to specialize
-/// VarStreamArrayExtractor<T> for your ValueType.
-///
-/// The default extractor type is stateless, but by specializing
-/// VarStreamArrayExtractor or defining your own custom extractor type and
-/// adding the appropriate ContextType typedef to the class, you can pass a
-/// context field during construction of the VarStreamArray that will be
-/// passed to each call to extract.
-///
-template <typename ValueType, typename ExtractorType>
-class VarStreamArrayIterator;
-
-template <typename ValueType,
-          typename ExtractorType = VarStreamArrayExtractor<ValueType>>
-class VarStreamArray {
-public:
-  typedef typename ExtractorType::ContextType ContextType;
-  typedef VarStreamArrayIterator<ValueType, ExtractorType> Iterator;
-  friend Iterator;
-
-  VarStreamArray() = default;
-
-  explicit VarStreamArray(BinaryStreamRef Stream,
-                          ContextType *Context = nullptr)
-      : Stream(Stream), Context(Context) {}
-
-  VarStreamArray(const VarStreamArray<ValueType, ExtractorType> &Other)
-      : Stream(Other.Stream), Context(Other.Context) {}
-
-  Iterator begin(bool *HadError = nullptr) const {
-    if (empty())
-      return end();
-
-    return Iterator(*this, Context, HadError);
-  }
-
-  Iterator end() const { return Iterator(); }
-
-  bool empty() const { return Stream.getLength() == 0; }
-
-  /// \brief given an offset into the array's underlying stream, return an
-  /// iterator to the record at that offset.  This is considered unsafe
-  /// since the behavior is undefined if \p Offset does not refer to the
-  /// beginning of a valid record.
-  Iterator at(uint32_t Offset) const {
-    return Iterator(*this, Context, Stream.drop_front(Offset), nullptr);
-  }
-
-  BinaryStreamRef getUnderlyingStream() const { return Stream; }
-
-private:
-  BinaryStreamRef Stream;
-  ContextType *Context = nullptr;
-};
-
-template <typename ValueType, typename ExtractorType>
+template <typename ArrayType, typename Value, typename Extractor,
+          typename WrappedCtx>
 class VarStreamArrayIterator
     : public iterator_facade_base<
-          VarStreamArrayIterator<ValueType, ExtractorType>,
-          std::forward_iterator_tag, ValueType> {
-  typedef typename ExtractorType::ContextType ContextType;
-  typedef VarStreamArrayIterator<ValueType, ExtractorType> IterType;
-  typedef VarStreamArray<ValueType, ExtractorType> ArrayType;
+          VarStreamArrayIterator<ArrayType, Value, Extractor, WrappedCtx>,
+          std::forward_iterator_tag, Value> {
+  typedef VarStreamArrayIterator<ArrayType, Value, Extractor, WrappedCtx>
+      IterType;
 
 public:
-  VarStreamArrayIterator(const ArrayType &Array, ContextType *Context,
+  VarStreamArrayIterator() = default;
+  VarStreamArrayIterator(const ArrayType &Array, const WrappedCtx &Ctx,
                          BinaryStreamRef Stream, bool *HadError = nullptr)
-      : IterRef(Stream), Context(Context), Array(&Array), HadError(HadError) {
+      : IterRef(Stream), Ctx(&Ctx), Array(&Array), HadError(HadError) {
     if (IterRef.getLength() == 0)
       moveToEnd();
     else {
-      auto EC = ExtractorType::extract(IterRef, ThisLen, ThisValue, Context);
+      auto EC = Ctx.template invoke<Extractor>(IterRef, ThisLen, ThisValue);
       if (EC) {
         consumeError(std::move(EC));
         markError();
@@ -142,11 +77,13 @@ public:
     }
   }
 
-  VarStreamArrayIterator(const ArrayType &Array, ContextType *Context,
+  VarStreamArrayIterator(const ArrayType &Array, const WrappedCtx &Ctx,
                          bool *HadError = nullptr)
-      : VarStreamArrayIterator(Array, Context, Array.Stream, HadError) {}
+      : VarStreamArrayIterator(Array, Ctx, Array.Stream, HadError) {}
+
+  VarStreamArrayIterator(const WrappedCtx &Ctx) : Ctx(&Ctx) {}
+  VarStreamArrayIterator(const VarStreamArrayIterator &Other) = default;
 
-  VarStreamArrayIterator() = default;
   ~VarStreamArrayIterator() = default;
 
   bool operator==(const IterType &R) const {
@@ -164,12 +101,12 @@ public:
     return false;
   }
 
-  const ValueType &operator*() const {
+  const Value &operator*() const {
     assert(Array && !HasError);
     return ThisValue;
   }
 
-  ValueType &operator*() {
+  Value &operator*() {
     assert(Array && !HasError);
     return ThisValue;
   }
@@ -185,7 +122,7 @@ public:
         moveToEnd();
       } else {
         // There is some data after the current record.
-        auto EC = ExtractorType::extract(IterRef, ThisLen, ThisValue, Context);
+        auto EC = Ctx->template invoke<Extractor>(IterRef, ThisLen, ThisValue);
         if (EC) {
           consumeError(std::move(EC));
           markError();
@@ -210,15 +147,134 @@ private:
       *HadError = true;
   }
 
-  ValueType ThisValue;
+  Value ThisValue;
   BinaryStreamRef IterRef;
-  ContextType *Context{nullptr};
+  const WrappedCtx *Ctx{nullptr};
   const ArrayType *Array{nullptr};
   uint32_t ThisLen{0};
   bool HasError{false};
   bool *HadError{nullptr};
 };
 
+template <typename T, typename Context> struct ContextWrapper {
+  ContextWrapper() = default;
+
+  explicit ContextWrapper(Context &&Ctx) : Ctx(Ctx) {}
+
+  template <typename Extractor>
+  Error invoke(BinaryStreamRef Stream, uint32_t &Len, T &Item) const {
+    return Extractor::extract(Stream, Len, Item, Ctx);
+  }
+
+  Context Ctx;
+};
+
+template <typename T> struct ContextWrapper<T, void> {
+  ContextWrapper() = default;
+
+  template <typename Extractor>
+  Error invoke(BinaryStreamRef Stream, uint32_t &Len, T &Item) const {
+    return Extractor::extract(Stream, Len, Item);
+  }
+};
+
+/// VarStreamArray represents an array of variable length records backed by a
+/// stream.  This could be a contiguous sequence of bytes in memory, it could
+/// be a file on disk, or it could be a PDB stream where bytes are stored as
+/// discontiguous blocks in a file.  Usually it is desirable to treat arrays
+/// as contiguous blocks of memory, but doing so with large PDB files, for
+/// example, could mean allocating huge amounts of memory just to allow
+/// re-ordering of stream data to be contiguous before iterating over it.  By
+/// abstracting this out, we need not duplicate this memory, and we can
+/// iterate over arrays in arbitrarily formatted streams.  Elements are parsed
+/// lazily on iteration, so there is no upfront cost associated with building
+/// or copying a VarStreamArray, no matter how large it may be.
+///
+/// You create a VarStreamArray by specifying a ValueType and an Extractor type.
+/// If you do not specify an Extractor type, you are expected to specialize
+/// VarStreamArrayExtractor<T> for your ValueType.
+///
+/// The default extractor type is stateless, but by specializing
+/// VarStreamArrayExtractor or defining your own custom extractor type and
+/// adding the appropriate ContextType typedef to the class, you can pass a
+/// context field during construction of the VarStreamArray that will be
+/// passed to each call to extract.
+///
+template <typename Value, typename Extractor, typename WrappedCtx>
+class VarStreamArrayBase {
+  typedef VarStreamArrayBase<Value, Extractor, WrappedCtx> MyType;
+
+public:
+  typedef VarStreamArrayIterator<MyType, Value, Extractor, WrappedCtx> Iterator;
+  friend Iterator;
+
+  VarStreamArrayBase() = default;
+
+  VarStreamArrayBase(BinaryStreamRef Stream, const WrappedCtx &Ctx)
+      : Stream(Stream), Ctx(Ctx) {}
+
+  VarStreamArrayBase(const MyType &Other)
+      : Stream(Other.Stream), Ctx(Other.Ctx) {}
+
+  Iterator begin(bool *HadError = nullptr) const {
+    if (empty())
+      return end();
+
+    return Iterator(*this, Ctx, Stream, HadError);
+  }
+
+  Iterator end() const { return Iterator(Ctx); }
+
+  bool empty() const { return Stream.getLength() == 0; }
+
+  /// \brief given an offset into the array's underlying stream, return an
+  /// iterator to the record at that offset.  This is considered unsafe
+  /// since the behavior is undefined if \p Offset does not refer to the
+  /// beginning of a valid record.
+  Iterator at(uint32_t Offset) const {
+    return Iterator(*this, Ctx, Stream.drop_front(Offset), nullptr);
+  }
+
+  BinaryStreamRef getUnderlyingStream() const { return Stream; }
+
+private:
+  BinaryStreamRef Stream;
+  WrappedCtx Ctx;
+};
+
+template <typename Value, typename Extractor, typename Context>
+class VarStreamArrayImpl
+    : public VarStreamArrayBase<Value, Extractor,
+                                ContextWrapper<Value, Context>> {
+  typedef ContextWrapper<Value, Context> WrappedContext;
+  typedef VarStreamArrayImpl<Value, Extractor, Context> MyType;
+  typedef VarStreamArrayBase<Value, Extractor, WrappedContext> BaseType;
+
+public:
+  typedef Context ContextType;
+
+  VarStreamArrayImpl() = default;
+  VarStreamArrayImpl(BinaryStreamRef Stream, Context &&Ctx)
+      : BaseType(Stream, WrappedContext(std::forward<Context>(Ctx))) {}
+};
+
+template <typename Value, typename Extractor>
+class VarStreamArrayImpl<Value, Extractor, void>
+    : public VarStreamArrayBase<Value, Extractor, ContextWrapper<Value, void>> {
+  typedef ContextWrapper<Value, void> WrappedContext;
+  typedef VarStreamArrayImpl<Value, Extractor, void> MyType;
+  typedef VarStreamArrayBase<Value, Extractor, WrappedContext> BaseType;
+
+public:
+  VarStreamArrayImpl() = default;
+  VarStreamArrayImpl(BinaryStreamRef Stream)
+      : BaseType(Stream, WrappedContext()) {}
+};
+
+template <typename Value, typename Extractor = VarStreamArrayExtractor<Value>>
+using VarStreamArray =
+    VarStreamArrayImpl<Value, Extractor, typename Extractor::ContextType>;
+
 template <typename T> class FixedStreamArrayIterator;
 
 /// FixedStreamArray is similar to VarStreamArray, except with each record

include/llvm/Support/BinaryStreamReader.h

@@ -173,13 +173,29 @@ public:
   /// \returns a success error code if the data was successfully read, otherwise
   /// returns an appropriate error code.
   template <typename T, typename U>
-  Error
-  readArray(VarStreamArray<T, U> &Array, uint32_t Size,
-            typename VarStreamArray<T, U>::ContextType *Context = nullptr) {
+  Error readArray(VarStreamArray<T, U> &Array, uint32_t Size) {
     BinaryStreamRef S;
     if (auto EC = readStreamRef(S, Size))
       return EC;
-    Array = VarStreamArray<T, U>(S, Context);
+    Array = VarStreamArray<T, U>(S);
+    return Error::success();
+  }
+
+  /// Read a VarStreamArray of size \p Size bytes and store the result into
+  /// \p Array.  Updates the stream's offset to point after the newly read
+  /// array.  Never causes a copy (although iterating the elements of the
+  /// VarStreamArray may, depending upon the implementation of the underlying
+  /// stream).
+  ///
+  /// \returns a success error code if the data was successfully read, otherwise
+  /// returns an appropriate error code.
+  template <typename T, typename U, typename ContextType>
+  Error readArray(VarStreamArray<T, U> &Array, uint32_t Size,
+                  ContextType &&Context) {
+    BinaryStreamRef S;
+    if (auto EC = readStreamRef(S, Size))
+      return EC;
+    Array = VarStreamArray<T, U>(S, std::move(Context));
     return Error::success();
   }
 

lib/DebugInfo/CodeView/ModuleDebugFileChecksumFragment.cpp

@@ -25,7 +25,7 @@ struct FileChecksumEntryHeader {
 };
 
 Error llvm::VarStreamArrayExtractor<FileChecksumEntry>::extract(
-    BinaryStreamRef Stream, uint32_t &Len, FileChecksumEntry &Item, void *Ctx) {
+    BinaryStreamRef Stream, uint32_t &Len, FileChecksumEntry &Item) {
   BinaryStreamReader Reader(Stream);
 
   const FileChecksumEntryHeader *Header;

lib/DebugInfo/CodeView/ModuleDebugInlineeLinesFragment.cpp

@@ -17,13 +17,13 @@ using namespace llvm::codeview;
 
 Error VarStreamArrayExtractor<InlineeSourceLine>::extract(
     BinaryStreamRef Stream, uint32_t &Len, InlineeSourceLine &Item,
-    ContextType *Fragment) {
+    bool HasExtraFiles) {
   BinaryStreamReader Reader(Stream);
 
   if (auto EC = Reader.readObject(Item.Header))
     return EC;
 
-  if (Fragment->hasExtraFiles()) {
+  if (HasExtraFiles) {
     uint32_t ExtraFileCount;
     if (auto EC = Reader.readInteger(ExtraFileCount))
       return EC;
@@ -42,7 +42,8 @@ Error ModuleDebugInlineeLineFragmentRef::initialize(BinaryStreamReader Reader) {
   if (auto EC = Reader.readEnum(Signature))
     return EC;
 
-  if (auto EC = Reader.readArray(Lines, Reader.bytesRemaining(), this))
+  if (auto EC =
+          Reader.readArray(Lines, Reader.bytesRemaining(), hasExtraFiles()))
     return EC;
 
   assert(Reader.bytesRemaining() == 0);

lib/DebugInfo/PDB/Native/DbiModuleDescriptorBuilder.cpp

@@ -89,6 +89,14 @@ uint32_t DbiModuleDescriptorBuilder::calculateSerializedLength() const {
   return alignTo(L + M + O, sizeof(uint32_t));
 }
 
+template <typename T> struct Foo {
+  explicit Foo(T &&Answer) : Answer(Answer) {}
+
+  T Answer;
+};
+
+template <typename T> Foo<T> makeFoo(T &&t) { return Foo<T>(std::move(t)); }
+
 void DbiModuleDescriptorBuilder::finalize() {
   Layout.FileNameOffs = 0; // TODO: Fix this
   Layout.Flags = 0;        // TODO: Fix this

unittests/Support/BinaryStreamTest.cpp

@@ -358,14 +358,14 @@ TEST_F(BinaryStreamTest, VarStreamArray) {
 
   struct StringExtractor {
   public:
-    typedef uint32_t ContextType;
+    typedef uint32_t &ContextType;
     static Error extract(BinaryStreamRef Stream, uint32_t &Len, StringRef &Item,
-                         uint32_t *Index) {
-      if (*Index == 0)
+                         uint32_t &Index) {
+      if (Index == 0)
         Len = strlen("1. Test");
-      else if (*Index == 1)
+      else if (Index == 1)
         Len = strlen("2. Longer Test");
-      else if (*Index == 2)
+      else if (Index == 2)
         Len = strlen("3. Really Long Test");
       else
         Len = strlen("4. Super Extra Longest Test Of All");
@@ -374,14 +374,14 @@ TEST_F(BinaryStreamTest, VarStreamArray) {
         return EC;
       Item =
           StringRef(reinterpret_cast<const char *>(Bytes.data()), Bytes.size());
-      ++(*Index);
+      ++Index;
       return Error::success();
     }
   };
 
   for (auto &Stream : Streams) {
     uint32_t Context = 0;
-    VarStreamArray<StringRef, StringExtractor> Array(*Stream.Input, &Context);
+    VarStreamArray<StringRef, StringExtractor> Array(*Stream.Input, Context);
     auto Iter = Array.begin();
     ASSERT_EQ("1. Test", *Iter++);
     ASSERT_EQ("2. Longer Test", *Iter++);