[PDB] Fix use after free.

[zturner]

Summary

Previously MappedBlockStream owned its own BumpPtrAllocator that
it would allocate from when a read crossed a block boundary. This
way it could still return the user a contiguous buffer of the
requested size. However, It's not uncommon to open a stream, read
some stuff, close it, and then save the information for later.
After all, since the entire file is mapped into memory, the data
should always be available as long as the file is open.

Of course, the exception to this is when the data isn't *in* the
file, but rather in some buffer that we temporarily allocated to
present this contiguous view. And this buffer would get destroyed
as soon as the strema was closed.

The fix here is to force the user to specify the allocator, this
way it can provide an allocator that has whatever lifetime it
chooses.

[zturner]

Also, no current code was exercising this bug which is why it never surfaced on ASAN bots. It's difficult to generate a PDB that is guaranteed to exhibit this problem, so instead a unit test is used. The unit test arranges the blocks of a fabricated MappedBlockStream so that it's guaranteed to allocate from the temporary allocator, and then it ensures that this data lives after the stream is destroyed.

[inglorion]

lgtm. Thanks!

llvm/include/llvm/DebugInfo/MSF/MappedBlockStream.h

 /// the MSF.  MappedBlockStream provides methods for reading from and writing
 /// to one of these streams transparently, as if it were a contiguous sequence
 /// of bytes.
 class MappedBlockStream : public BinaryStream {
   friend class WritableMappedBlockStream;
 public:
   static std::unique_ptr<MappedBlockStream>
   createStream(uint32_t BlockSize, const MSFStreamLayout &Layout,
-               BinaryStreamRef MsfData);
+               BinaryStreamRef MsfData, BumpPtrAllocator &Allocator);
 
   static std::unique_ptr<MappedBlockStream>
   createIndexedStream(const MSFLayout &Layout, BinaryStreamRef MsfData,
-                      uint32_t StreamIndex);
+                      uint32_t StreamIndex, BumpPtrAllocator &Allocator);
 
   static std::unique_ptr<MappedBlockStream>
-  createFpmStream(const MSFLayout &Layout, BinaryStreamRef MsfData);
+  createFpmStream(const MSFLayout &Layout, BinaryStreamRef MsfData,
+                  BumpPtrAllocator &Allocator);
 
   static std::unique_ptr<MappedBlockStream>
-  createDirectoryStream(const MSFLayout &Layout, BinaryStreamRef MsfData);
+  createDirectoryStream(const MSFLayout &Layout, BinaryStreamRef MsfData,
+                        BumpPtrAllocator &Allocator);
 
   llvm::support::endianness getEndian() const override {
     return llvm::support::little;
   }
 
   Error readBytes(uint32_t Offset, uint32_t Size,
                   ArrayRef<uint8_t> &Buffer) override;
   Error readLongestContiguousChunk(uint32_t Offset,
                                    ArrayRef<uint8_t> &Buffer) override;
 
   uint32_t getLength() override;
 
-  uint32_t getNumBytesCopied() const;
-
-  llvm::BumpPtrAllocator &getAllocator() { return Pool; }
+  llvm::BumpPtrAllocator &getAllocator() { return Allocator; }
 
   void invalidateCache();
 
   uint32_t getBlockSize() const { return BlockSize; }
   uint32_t getNumBlocks() const { return StreamLayout.Blocks.size(); }
   uint32_t getStreamLength() const { return StreamLayout.Length; }
 
 protected:
   MappedBlockStream(uint32_t BlockSize, const MSFStreamLayout &StreamLayout,
-                    BinaryStreamRef MsfData);
+                    BinaryStreamRef MsfData, BumpPtrAllocator &Allocator);
 
 private:
   const MSFStreamLayout &getStreamLayout() const { return StreamLayout; }
   void fixCacheAfterWrite(uint32_t Offset, ArrayRef<uint8_t> Data) const;
 
   Error readBytes(uint32_t Offset, MutableArrayRef<uint8_t> Buffer);
   bool tryReadContiguously(uint32_t Offset, uint32_t Size,
                            ArrayRef<uint8_t> &Buffer);
 
   const uint32_t BlockSize;
   const MSFStreamLayout StreamLayout;
   BinaryStreamRef MsfData;
 
   typedef MutableArrayRef<uint8_t> CacheEntry;
-  llvm::BumpPtrAllocator Pool;
+
+  // We just store the allocator by reference.  We use this to allocate
+  // contiguous memory for things like arrays or strings that cross a block
+  // boundary, and this memory is expected to outlive the stream.  For example,
+  // someone could create a stream, read some stuff, then close the stream, and
+  // we would like outstanding references to fields to remain valid since the
+  // entire file is mapped anyway.  Because of that, the user must supply the
+  // allocator to allocate broken records from.
+  BumpPtrAllocator &Allocator;
   DenseMap<uint32_t, std::vector<CacheEntry>> CacheMap;
 };
 
 class WritableMappedBlockStream : public WritableBinaryStream {
 public:
   static std::unique_ptr<WritableMappedBlockStream>
   createStream(uint32_t BlockSize, const MSFStreamLayout &Layout,
-               WritableBinaryStreamRef MsfData);
+               WritableBinaryStreamRef MsfData, BumpPtrAllocator &Allocator);
 
   static std::unique_ptr<WritableMappedBlockStream>
   createIndexedStream(const MSFLayout &Layout, WritableBinaryStreamRef MsfData,
-                      uint32_t StreamIndex);
+                      uint32_t StreamIndex, BumpPtrAllocator &Allocator);
 
   static std::unique_ptr<WritableMappedBlockStream>
   createDirectoryStream(const MSFLayout &Layout,
-                        WritableBinaryStreamRef MsfData);
+                        WritableBinaryStreamRef MsfData,
+                        BumpPtrAllocator &Allocator);
 
   static std::unique_ptr<WritableMappedBlockStream>
-  createFpmStream(const MSFLayout &Layout, WritableBinaryStreamRef MsfData);
+  createFpmStream(const MSFLayout &Layout, WritableBinaryStreamRef MsfData,
+                  BumpPtrAllocator &Allocator);
 
   llvm::support::endianness getEndian() const override {
     return llvm::support::little;
   }
 
   Error readBytes(uint32_t Offset, uint32_t Size,
                   ArrayRef<uint8_t> &Buffer) override;
   Error readLongestContiguousChunk(uint32_t Offset,
public:
   }
   uint32_t getBlockSize() const { return ReadInterface.getBlockSize(); }
   uint32_t getNumBlocks() const { return ReadInterface.getNumBlocks(); }
   uint32_t getStreamLength() const { return ReadInterface.getStreamLength(); }
 
 protected:
   WritableMappedBlockStream(uint32_t BlockSize,
                             const MSFStreamLayout &StreamLayout,
-                            WritableBinaryStreamRef MsfData);
+                            WritableBinaryStreamRef MsfData,
+                            BumpPtrAllocator &Allocator);
 
 private:
   MappedBlockStream ReadInterface;
 
   WritableBinaryStreamRef WriteInterface;
 };
 
 } // end namespace pdb
 } // end namespace llvm
 
 #endif // LLVM_DEBUGINFO_MSF_MAPPEDBLOCKSTREAM_H

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

 }
 namespace pdb {
 class PDBFile;
 
 class TpiStream {
   friend class TpiStreamBuilder;
 
 public:
-  TpiStream(const PDBFile &File,
-            std::unique_ptr<msf::MappedBlockStream> Stream);
+  TpiStream(PDBFile &File, std::unique_ptr<msf::MappedBlockStream> Stream);
   ~TpiStream();
   Error reload();
 
   PdbRaw_TpiVer getTpiVersion() const;
 
   uint32_t TypeIndexBegin() const;
   uint32_t TypeIndexEnd() const;
   uint32_t getNumTypeRecords() const;
public:
   codeview::CVTypeRange types(bool *HadError) const;
   const codeview::CVTypeArray &typeArray() const { return TypeRecords; }
 
   codeview::LazyRandomTypeCollection &typeCollection() { return *Types; }
 
   Error commit();
 
 private:
-  const PDBFile &Pdb;
+  PDBFile &Pdb;
   std::unique_ptr<msf::MappedBlockStream> Stream;
 
   std::unique_ptr<codeview::LazyRandomTypeCollection> Types;
 
   codeview::CVTypeArray TypeRecords;
 
   std::unique_ptr<BinaryStream> HashStream;
   FixedStreamArray<support::ulittle32_t> HashValues;

llvm/lib/DebugInfo/MSF/MappedBlockStream.cpp

 typedef std::pair<uint32_t, uint32_t> Interval;
 static Interval intersect(const Interval &I1, const Interval &I2) {
   return std::make_pair(std::max(I1.first, I2.first),
                         std::min(I1.second, I2.second));
 }
 
 MappedBlockStream::MappedBlockStream(uint32_t BlockSize,
                                      const MSFStreamLayout &Layout,
-                                     BinaryStreamRef MsfData)
-    : BlockSize(BlockSize), StreamLayout(Layout), MsfData(MsfData) {}
-
-std::unique_ptr<MappedBlockStream>
-MappedBlockStream::createStream(uint32_t BlockSize,
-                                const MSFStreamLayout &Layout,
-                                BinaryStreamRef MsfData) {
+                                     BinaryStreamRef MsfData,
+                                     BumpPtrAllocator &Allocator)
+    : BlockSize(BlockSize), StreamLayout(Layout), MsfData(MsfData),
+      Allocator(Allocator) {}
+
+std::unique_ptr<MappedBlockStream> MappedBlockStream::createStream(
+    uint32_t BlockSize, const MSFStreamLayout &Layout, BinaryStreamRef MsfData,
+    BumpPtrAllocator &Allocator) {
   return llvm::make_unique<MappedBlockStreamImpl<MappedBlockStream>>(
-      BlockSize, Layout, MsfData);
+      BlockSize, Layout, MsfData, Allocator);
 }
 
 std::unique_ptr<MappedBlockStream> MappedBlockStream::createIndexedStream(
-    const MSFLayout &Layout, BinaryStreamRef MsfData, uint32_t StreamIndex) {
+    const MSFLayout &Layout, BinaryStreamRef MsfData, uint32_t StreamIndex,
+    BumpPtrAllocator &Allocator) {
   assert(StreamIndex < Layout.StreamMap.size() && "Invalid stream index");
   MSFStreamLayout SL;
   SL.Blocks = Layout.StreamMap[StreamIndex];
   SL.Length = Layout.StreamSizes[StreamIndex];
   return llvm::make_unique<MappedBlockStreamImpl<MappedBlockStream>>(
-      Layout.SB->BlockSize, SL, MsfData);
+      Layout.SB->BlockSize, SL, MsfData, Allocator);
 }
 
 std::unique_ptr<MappedBlockStream>
 MappedBlockStream::createDirectoryStream(const MSFLayout &Layout,
-                                         BinaryStreamRef MsfData) {
+                                         BinaryStreamRef MsfData,
+                                         BumpPtrAllocator &Allocator) {
   MSFStreamLayout SL;
   SL.Blocks = Layout.DirectoryBlocks;
   SL.Length = Layout.SB->NumDirectoryBytes;
-  return createStream(Layout.SB->BlockSize, SL, MsfData);
+  return createStream(Layout.SB->BlockSize, SL, MsfData, Allocator);
 }
 
 std::unique_ptr<MappedBlockStream>
 MappedBlockStream::createFpmStream(const MSFLayout &Layout,
-                                   BinaryStreamRef MsfData) {
+                                   BinaryStreamRef MsfData,
+                                   BumpPtrAllocator &Allocator) {
   MSFStreamLayout SL;
   initializeFpmStreamLayout(Layout, SL);
-  return createStream(Layout.SB->BlockSize, SL, MsfData);
+  return createStream(Layout.SB->BlockSize, SL, MsfData, Allocator);
 }
 
 Error MappedBlockStream::readBytes(uint32_t Offset, uint32_t Size,
                                    ArrayRef<uint8_t> &Buffer) {
   // Make sure we aren't trying to read beyond the end of the stream.
   if (auto EC = checkOffset(Offset, Size))
     return EC;
 
for (auto &CacheItem : CacheMap) {
     Buffer = CachedAlloc.slice(CacheRangeOffset, Size);
     return Error::success();
   }
 
   // Otherwise allocate a large enough buffer in the pool, memcpy the data
   // into it, and return an ArrayRef to that.  Do not touch existing pool
   // allocations, as existing clients may be holding a pointer which must
   // not be invalidated.
-  uint8_t *WriteBuffer = static_cast<uint8_t *>(Pool.Allocate(Size, 8));
+  uint8_t *WriteBuffer = static_cast<uint8_t *>(Allocator.Allocate(Size, 8));
   if (auto EC = readBytes(Offset, MutableArrayRef<uint8_t>(WriteBuffer, Size)))
     return EC;
 
   if (CacheIter != CacheMap.end()) {
     CacheIter->second.emplace_back(WriteBuffer, Size);
   } else {
     std::vector<CacheEntry> List;
     List.emplace_back(WriteBuffer, Size);
while (BytesLeft > 0) {
     BytesLeft -= BytesInChunk;
     ++BlockNum;
     OffsetInBlock = 0;
   }
 
   return Error::success();
 }
 
-uint32_t MappedBlockStream::getNumBytesCopied() const {
-  return static_cast<uint32_t>(Pool.getBytesAllocated());
-}
-
 void MappedBlockStream::invalidateCache() { CacheMap.shrink_and_clear(); }
 
 void MappedBlockStream::fixCacheAfterWrite(uint32_t Offset,
                                            ArrayRef<uint8_t> Data) const {
   // If this write overlapped a read which previously came from the pool,
   // someone may still be holding a pointer to that alloc which is now invalid.
   // Compute the overlapping range and update the cache entry, so any
   // outstanding buffers are automatically updated.
for (const auto &Alloc : MapEntry.second) {
           AbsoluteDifference(CachedInterval.first, Intersection.first);
       ::memcpy(Alloc.data() + DestOffset, Data.data() + SrcOffset, Length);
     }
   }
 }
 
 WritableMappedBlockStream::WritableMappedBlockStream(
     uint32_t BlockSize, const MSFStreamLayout &Layout,
-    WritableBinaryStreamRef MsfData)
-    : ReadInterface(BlockSize, Layout, MsfData), WriteInterface(MsfData) {}
+    WritableBinaryStreamRef MsfData, BumpPtrAllocator &Allocator)
+    : ReadInterface(BlockSize, Layout, MsfData, Allocator),
+      WriteInterface(MsfData) {}
 
 std::unique_ptr<WritableMappedBlockStream>
 WritableMappedBlockStream::createStream(uint32_t BlockSize,
                                         const MSFStreamLayout &Layout,
-                                        WritableBinaryStreamRef MsfData) {
+                                        WritableBinaryStreamRef MsfData,
+                                        BumpPtrAllocator &Allocator) {
   return llvm::make_unique<MappedBlockStreamImpl<WritableMappedBlockStream>>(
-      BlockSize, Layout, MsfData);
+      BlockSize, Layout, MsfData, Allocator);
 }
 
 std::unique_ptr<WritableMappedBlockStream>
 WritableMappedBlockStream::createIndexedStream(const MSFLayout &Layout,
                                                WritableBinaryStreamRef MsfData,
-                                               uint32_t StreamIndex) {
+                                               uint32_t StreamIndex,
+                                               BumpPtrAllocator &Allocator) {
   assert(StreamIndex < Layout.StreamMap.size() && "Invalid stream index");
   MSFStreamLayout SL;
   SL.Blocks = Layout.StreamMap[StreamIndex];
   SL.Length = Layout.StreamSizes[StreamIndex];
-  return createStream(Layout.SB->BlockSize, SL, MsfData);
+  return createStream(Layout.SB->BlockSize, SL, MsfData, Allocator);
 }
 
 std::unique_ptr<WritableMappedBlockStream>
 WritableMappedBlockStream::createDirectoryStream(
-    const MSFLayout &Layout, WritableBinaryStreamRef MsfData) {
+    const MSFLayout &Layout, WritableBinaryStreamRef MsfData,
+    BumpPtrAllocator &Allocator) {
   MSFStreamLayout SL;
   SL.Blocks = Layout.DirectoryBlocks;
   SL.Length = Layout.SB->NumDirectoryBytes;
-  return createStream(Layout.SB->BlockSize, SL, MsfData);
+  return createStream(Layout.SB->BlockSize, SL, MsfData, Allocator);
 }
 
 std::unique_ptr<WritableMappedBlockStream>
 WritableMappedBlockStream::createFpmStream(const MSFLayout &Layout,
-                                           WritableBinaryStreamRef MsfData) {
+                                           WritableBinaryStreamRef MsfData,
+                                           BumpPtrAllocator &Allocator) {
   MSFStreamLayout SL;
   initializeFpmStreamLayout(Layout, SL);
-  return createStream(Layout.SB->BlockSize, SL, MsfData);
+  return createStream(Layout.SB->BlockSize, SL, MsfData, Allocator);
 }
 
 Error WritableMappedBlockStream::readBytes(uint32_t Offset, uint32_t Size,
                                            ArrayRef<uint8_t> &Buffer) {
   return ReadInterface.readBytes(Offset, Size, Buffer);
 }
 
 Error WritableMappedBlockStream::readLongestContiguousChunk(

llvm/lib/DebugInfo/PDB/Native/DbiModuleDescriptorBuilder.cpp

if (auto EC = ModiWriter.writeCString(ModuleName))
     return EC;
   if (auto EC = ModiWriter.writeCString(ObjFileName))
     return EC;
   if (auto EC = ModiWriter.padToAlignment(sizeof(uint32_t)))
     return EC;
 
   if (Layout.ModDiStream != kInvalidStreamIndex) {
     auto NS = WritableMappedBlockStream::createIndexedStream(
-        MsfLayout, MsfBuffer, Layout.ModDiStream);
+        MsfLayout, MsfBuffer, Layout.ModDiStream, MSF.getAllocator());
     WritableBinaryStreamRef Ref(*NS);
     BinaryStreamWriter SymbolWriter(Ref);
     // Write the symbols.
     if (auto EC =
             SymbolWriter.writeInteger<uint32_t>(COFF::DEBUG_SECTION_MAGIC))
       return EC;
     BinaryItemStream<CVSymbol> Records(llvm::support::endianness::little);
     Records.setItems(Symbols);

llvm/lib/DebugInfo/PDB/Native/DbiStream.cpp

Error DbiStream::initializeSectionHeadersData() {
   if (DbgStreams.size() == 0)
     return Error::success();
 
   uint32_t StreamNum = getDebugStreamIndex(DbgHeaderType::SectionHdr);
   if (StreamNum >= Pdb.getNumStreams())
     return make_error<RawError>(raw_error_code::no_stream);
 
   auto SHS = MappedBlockStream::createIndexedStream(
-      Pdb.getMsfLayout(), Pdb.getMsfBuffer(), StreamNum);
+      Pdb.getMsfLayout(), Pdb.getMsfBuffer(), StreamNum, Pdb.getAllocator());
 
   size_t StreamLen = SHS->getLength();
   if (StreamLen % sizeof(object::coff_section))
     return make_error<RawError>(raw_error_code::corrupt_file,
                                 "Corrupted section header stream.");
 
   size_t NumSections = StreamLen / sizeof(object::coff_section);
   BinaryStreamReader Reader(*SHS);
Error DbiStream::initializeFpoRecords() {
   // This means there is no FPO data.
   if (StreamNum == kInvalidStreamIndex)
     return Error::success();
 
   if (StreamNum >= Pdb.getNumStreams())
     return make_error<RawError>(raw_error_code::no_stream);
 
   auto FS = MappedBlockStream::createIndexedStream(
-      Pdb.getMsfLayout(), Pdb.getMsfBuffer(), StreamNum);
+      Pdb.getMsfLayout(), Pdb.getMsfBuffer(), StreamNum, Pdb.getAllocator());
 
   size_t StreamLen = FS->getLength();
   if (StreamLen % sizeof(object::FpoData))
     return make_error<RawError>(raw_error_code::corrupt_file,
                                 "Corrupted New FPO stream.");
 
   size_t NumRecords = StreamLen / sizeof(object::FpoData);
   BinaryStreamReader Reader(*FS);

llvm/lib/DebugInfo/PDB/Native/DbiStreamBuilder.cpp

std::vector<SecMapEntry> DbiStreamBuilder::createSectionMap(
   return Ret;
 }
 
 Error DbiStreamBuilder::commit(const msf::MSFLayout &Layout,
                                WritableBinaryStreamRef MsfBuffer) {
   if (auto EC = finalize())
     return EC;
 
-  auto DbiS = WritableMappedBlockStream::createIndexedStream(Layout, MsfBuffer,
-                                                             StreamDBI);
+  auto DbiS = WritableMappedBlockStream::createIndexedStream(
+      Layout, MsfBuffer, StreamDBI, Allocator);
 
   BinaryStreamWriter Writer(*DbiS);
   if (auto EC = Writer.writeObject(*Header))
     return EC;
 
   for (auto &M : ModiList) {
     if (auto EC = M->commit(Writer, Layout, MsfBuffer))
       return EC;
Error DbiStreamBuilder::commit(const msf::MSFLayout &Layout,
   for (auto &Stream : DbgStreams)
     if (auto EC = Writer.writeInteger(Stream.StreamNumber))
       return EC;
 
   for (auto &Stream : DbgStreams) {
     if (Stream.StreamNumber == kInvalidStreamIndex)
       continue;
     auto WritableStream = WritableMappedBlockStream::createIndexedStream(
-        Layout, MsfBuffer, Stream.StreamNumber);
+        Layout, MsfBuffer, Stream.StreamNumber, Allocator);
     BinaryStreamWriter DbgStreamWriter(*WritableStream);
     if (auto EC = DbgStreamWriter.writeArray(Stream.Data))
       return EC;
   }
 
   if (Writer.bytesRemaining() > 0)
     return make_error<RawError>(raw_error_code::invalid_format,
                                 "Unexpected bytes found in DBI Stream");
   return Error::success();
 }

llvm/lib/DebugInfo/PDB/Native/InfoStreamBuilder.cpp

uint32_t Length = sizeof(InfoStreamHeader) + NamedStreams.finalize() +
                     (Features.size() + 1) * sizeof(uint32_t);
   if (auto EC = Msf.setStreamSize(StreamPDB, Length))
     return EC;
   return Error::success();
 }
 
 Error InfoStreamBuilder::commit(const msf::MSFLayout &Layout,
                                 WritableBinaryStreamRef Buffer) const {
-  auto InfoS =
-      WritableMappedBlockStream::createIndexedStream(Layout, Buffer, StreamPDB);
+  auto InfoS = WritableMappedBlockStream::createIndexedStream(
+      Layout, Buffer, StreamPDB, Msf.getAllocator());
   BinaryStreamWriter Writer(*InfoS);
 
   InfoStreamHeader H;
   H.Age = Age;
   H.Signature = Sig;
   H.Version = Ver;
   H.Guid = Guid;
   if (auto EC = Writer.writeObject(H))

llvm/lib/DebugInfo/PDB/Native/PDBFile.cpp

Error PDBFile::parseFileHeaders() {
   // the Fpm is split across the file at `getBlockSize()` intervals.  As a
   // result, every block whose index is of the form |{1,2} + getBlockSize() * k|
   // for any non-negative integer k is an Fpm block.  In theory, we only really
   // need to reserve blocks of the form |{1,2} + getBlockSize() * 8 * k|, but
   // current versions of the MSF format already expect the Fpm to be arranged
   // at getBlockSize() intervals, so we have to be compatible.
   // See the function fpmPn() for more information:
   // https://github.com/Microsoft/microsoft-pdb/blob/master/PDB/msf/msf.cpp#L489
-  auto FpmStream = MappedBlockStream::createFpmStream(ContainerLayout, *Buffer);
+  auto FpmStream =
+      MappedBlockStream::createFpmStream(ContainerLayout, *Buffer, Allocator);
   BinaryStreamReader FpmReader(*FpmStream);
   ArrayRef<uint8_t> FpmBytes;
   if (auto EC = FpmReader.readBytes(FpmBytes,
                                     msf::getFullFpmByteSize(ContainerLayout)))
     return EC;
   uint32_t BlocksRemaining = getBlockCount();
   uint32_t BI = 0;
   for (auto Byte : FpmBytes) {
Error PDBFile::parseStreamData() {
 
   uint32_t NumStreams = 0;
 
   // Normally you can't use a MappedBlockStream without having fully parsed the
   // PDB file, because it accesses the directory and various other things, which
   // is exactly what we are attempting to parse.  By specifying a custom
   // subclass of IPDBStreamData which only accesses the fields that have already
   // been parsed, we can avoid this and reuse MappedBlockStream.
-  auto DS = MappedBlockStream::createDirectoryStream(ContainerLayout, *Buffer);
+  auto DS = MappedBlockStream::createDirectoryStream(ContainerLayout, *Buffer,
+                                                     Allocator);
   BinaryStreamReader Reader(*DS);
   if (auto EC = Reader.readInteger(NumStreams))
     return EC;
 
   if (auto EC = Reader.readArray(ContainerLayout.StreamSizes, NumStreams))
     return EC;
   for (uint32_t I = 0; I < NumStreams; ++I) {
     uint32_t StreamSize = getStreamByteSize(I);

 /// will have an MSFError with code msf_error_code::no_stream.  Else, the return
 /// value will contain the stream returned by createIndexedStream().
 Expected<std::unique_ptr<MappedBlockStream>>
 PDBFile::safelyCreateIndexedStream(const MSFLayout &Layout,
                                    BinaryStreamRef MsfData,
                                    uint32_t StreamIndex) const {
   if (StreamIndex >= getNumStreams())
     return make_error<RawError>(raw_error_code::no_stream);
-  return MappedBlockStream::createIndexedStream(Layout, MsfData, StreamIndex);
+  return MappedBlockStream::createIndexedStream(Layout, MsfData, StreamIndex,
+                                                Allocator);
 }

llvm/lib/DebugInfo/PDB/Native/PDBFileBuilder.cpp

Error PDBFileBuilder::commit(StringRef Filename) {
   if (auto EC = Writer.writeObject(*Layout.SB))
     return EC;
   uint32_t BlockMapOffset =
       msf::blockToOffset(Layout.SB->BlockMapAddr, Layout.SB->BlockSize);
   Writer.setOffset(BlockMapOffset);
   if (auto EC = Writer.writeArray(Layout.DirectoryBlocks))
     return EC;
 
-  auto DirStream =
-      WritableMappedBlockStream::createDirectoryStream(Layout, Buffer);
+  auto DirStream = WritableMappedBlockStream::createDirectoryStream(
+      Layout, Buffer, Allocator);
   BinaryStreamWriter DW(*DirStream);
   if (auto EC = DW.writeInteger<uint32_t>(Layout.StreamSizes.size()))
     return EC;
 
   if (auto EC = DW.writeArray(Layout.StreamSizes))
     return EC;
 
   for (const auto &Blocks : Layout.StreamMap) {
     if (auto EC = DW.writeArray(Blocks))
       return EC;
   }
 
   auto ExpectedSN = getNamedStreamIndex("/names");
   if (!ExpectedSN)
     return ExpectedSN.takeError();
 
-  auto NS = WritableMappedBlockStream::createIndexedStream(Layout, Buffer,
-                                                           *ExpectedSN);
+  auto NS = WritableMappedBlockStream::createIndexedStream(
+      Layout, Buffer, *ExpectedSN, Allocator);
   BinaryStreamWriter NSWriter(*NS);
   if (auto EC = Strings.commit(NSWriter))
     return EC;
 
   if (Info) {
     if (auto EC = Info->commit(Layout, Buffer))
       return EC;
   }

llvm/lib/DebugInfo/PDB/Native/TpiStream.cpp

 #include <vector>
 
 using namespace llvm;
 using namespace llvm::codeview;
 using namespace llvm::support;
 using namespace llvm::msf;
 using namespace llvm::pdb;
 
-TpiStream::TpiStream(const PDBFile &File,
-                     std::unique_ptr<MappedBlockStream> Stream)
+TpiStream::TpiStream(PDBFile &File, std::unique_ptr<MappedBlockStream> Stream)
     : Pdb(File), Stream(std::move(Stream)) {}
 
 TpiStream::~TpiStream() = default;
 
 Error TpiStream::reload() {
   BinaryStreamReader Reader(*Stream);
 
   if (Reader.bytesRemaining() < sizeof(TpiStreamHeader))
Error TpiStream::reload() {
 
   // Hash indices, hash values, etc come from the hash stream.
   if (Header->HashStreamIndex != kInvalidStreamIndex) {
     if (Header->HashStreamIndex >= Pdb.getNumStreams())
       return make_error<RawError>(raw_error_code::corrupt_file,
                                   "Invalid TPI hash stream index.");
 
     auto HS = MappedBlockStream::createIndexedStream(
-        Pdb.getMsfLayout(), Pdb.getMsfBuffer(), Header->HashStreamIndex);
+        Pdb.getMsfLayout(), Pdb.getMsfBuffer(), Header->HashStreamIndex,
+        Pdb.getAllocator());
     BinaryStreamReader HSR(*HS);
 
     // There should be a hash value for every type record, or no hashes at all.
     uint32_t NumHashValues =
         Header->HashValueBuffer.Length / sizeof(ulittle32_t);
     if (NumHashValues != getNumTypeRecords() && NumHashValues != 0)
       return make_error<RawError>(
           raw_error_code::corrupt_file,

llvm/lib/DebugInfo/PDB/Native/TpiStreamBuilder.cpp

Error TpiStreamBuilder::finalizeMsfLayout() {
   return Error::success();
 }
 
 Error TpiStreamBuilder::commit(const msf::MSFLayout &Layout,
                                WritableBinaryStreamRef Buffer) {
   if (auto EC = finalize())
     return EC;
 
-  auto InfoS =
-      WritableMappedBlockStream::createIndexedStream(Layout, Buffer, Idx);
+  auto InfoS = WritableMappedBlockStream::createIndexedStream(Layout, Buffer,
+                                                              Idx, Allocator);
 
   BinaryStreamWriter Writer(*InfoS);
   if (auto EC = Writer.writeObject(*Header))
     return EC;
 
   for (auto Rec : TypeRecords)
     if (auto EC = Writer.writeBytes(Rec))
       return EC;
 
   if (HashStreamIndex != kInvalidStreamIndex) {
-    auto HVS = WritableMappedBlockStream::createIndexedStream(Layout, Buffer,
-                                                              HashStreamIndex);
+    auto HVS = WritableMappedBlockStream::createIndexedStream(
+        Layout, Buffer, HashStreamIndex, Allocator);
     BinaryStreamWriter HW(*HVS);
     if (HashValueStream) {
       if (auto EC = HW.writeStreamRef(*HashValueStream))
         return EC;
     }
 
     for (auto &IndexOffset : TypeIndexOffsets) {
       if (auto EC = HW.writeObject(IndexOffset))
         return EC;
     }
   }
 
   return Error::success();
 }

llvm/tools/llvm-pdbdump/LLVMOutputStyle.cpp

for (auto &Str : opts::raw::DumpStreamData) {
     uint32_t End = 0;
 
     if (auto EC = parseStreamSpec(Str, SI, Begin, Size))
       return EC;
 
     if (SI >= File.getNumStreams())
       return make_error<RawError>(raw_error_code::no_stream);
 
-    auto S = MappedBlockStream::createIndexedStream(File.getMsfLayout(),
-                                                    File.getMsfBuffer(), SI);
+    auto S = MappedBlockStream::createIndexedStream(
+        File.getMsfLayout(), File.getMsfBuffer(), SI, File.getAllocator());
     if (!S)
       continue;
     DictScope DD(P, "Stream");
     if (Size == 0)
       End = S->getLength();
     else {
       End = Begin + Size;
       if (End >= S->getLength())
for (uint32_t I = 0; I < Modules.getModuleCount(); ++I) {
           P.printString(File);
       }
       bool HasModuleDI = (Modi.getModuleStreamIndex() < File.getNumStreams());
       bool ShouldDumpSymbols =
           (opts::raw::DumpModuleSyms || opts::raw::DumpSymRecordBytes);
       if (HasModuleDI && (ShouldDumpSymbols || opts::raw::DumpLineInfo)) {
         auto ModStreamData = MappedBlockStream::createIndexedStream(
             File.getMsfLayout(), File.getMsfBuffer(),
-            Modi.getModuleStreamIndex());
+            Modi.getModuleStreamIndex(), File.getAllocator());
 
         ModuleDebugStreamRef ModS(Modi, std::move(ModStreamData));
         if (auto EC = ModS.reload())
           return EC;
 
         if (ShouldDumpSymbols) {
           auto ExpectedTypes = initializeTypeDatabase(StreamTPI);
           if (!ExpectedTypes)

llvm/tools/llvm-pdbdump/YAMLOutputStyle.cpp

for (uint32_t I = 0; I < Modules.getModuleCount(); ++I) {
         DMI.SourceFiles.assign(Files.begin(), Files.end());
       }
 
       uint16_t ModiStream = MI.getModuleStreamIndex();
       if (ModiStream == kInvalidStreamIndex)
         continue;
 
       auto ModStreamData = msf::MappedBlockStream::createIndexedStream(
-          File.getMsfLayout(), File.getMsfBuffer(), ModiStream);
+          File.getMsfLayout(), File.getMsfBuffer(), ModiStream,
+          File.getAllocator());
 
       pdb::ModuleDebugStreamRef ModS(MI, std::move(ModStreamData));
       if (auto EC = ModS.reload())
         return EC;
 
       auto ExpectedST = File.getStringTable();
       if (!ExpectedST)
         return ExpectedST.takeError();

llvm/tools/llvm-pdbdump/fuzzer/llvm-pdbdump-fuzzer.cpp

extern "C" int LLVMFuzzerTestOneInput(uint8_t *data, size_t size) {
   if (auto E = InfoS.takeError()) {
     consumeError(std::move(E));
     return 0;
   }
   pdb::DbiStream &DS = DbiS.get();
 
   for (auto &Modi : DS.modules()) {
     auto ModStreamData = pdb::MappedBlockStream::createIndexedStream(
-      Modi.Info.getModuleStreamIndex(), *File);
+        Modi.Info.getModuleStreamIndex(), *File, File->getAllocator());
     if (!ModStreamData) {
       consumeError(ModStreamData.takeError());
       return 0;
     }
     pdb::ModuleDebugStreamRef ModS(Modi.Info, std::move(*ModStreamData));
     if (auto E = ModS.reload()) {
       consumeError(std::move(E));
       return 0;

llvm/unittests/DebugInfo/PDB/MappedBlockStreamTest.cpp

Error writeBytes(uint32_t Offset, ArrayRef<uint8_t> SrcData) override {
     return Error::success();
   }
   Error commit() override { return Error::success(); }
 
   MSFStreamLayout layout() const {
     return MSFStreamLayout{static_cast<uint32_t>(Data.size()), Blocks};
   }
 
+  BumpPtrAllocator Allocator;
+
 private:
   std::vector<support::ulittle32_t> Blocks;
   MutableArrayRef<uint8_t> Data;
 };
 
 TEST(MappedBlockStreamTest, NumBlocks) {
   DiscontiguousStream F(BlocksAry, DataAry);
-  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F);
+  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F,
+                                           F.Allocator);
   EXPECT_EQ(F.block_size(), S->getBlockSize());
   EXPECT_EQ(F.layout().Blocks.size(), S->getNumBlocks());
 
 }
 
 // Tests that a read which is entirely contained within a single block works
 // and does not allocate.
 TEST(MappedBlockStreamTest, ReadBeyondEndOfStreamRef) {
   DiscontiguousStream F(BlocksAry, DataAry);
-  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F);
+  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F,
+                                           F.Allocator);
 
   BinaryStreamReader R(*S);
   BinaryStreamRef SR;
   EXPECT_NO_ERROR(R.readStreamRef(SR, 0U));
   ArrayRef<uint8_t> Buffer;
   EXPECT_ERROR(SR.readBytes(0U, 1U, Buffer));
   EXPECT_NO_ERROR(R.readStreamRef(SR, 1U));
   EXPECT_ERROR(SR.readBytes(1U, 1U, Buffer));
 }
 
 // Tests that a read which outputs into a full destination buffer works and
 // does not fail due to the length of the output buffer.
 TEST(MappedBlockStreamTest, ReadOntoNonEmptyBuffer) {
   DiscontiguousStream F(BlocksAry, DataAry);
-  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F);
+  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F,
+                                           F.Allocator);
 
   BinaryStreamReader R(*S);
   StringRef Str = "ZYXWVUTSRQPONMLKJIHGFEDCBA";
   EXPECT_NO_ERROR(R.readFixedString(Str, 1));
   EXPECT_EQ(Str, StringRef("A"));
-  EXPECT_EQ(0U, S->getNumBytesCopied());
+  EXPECT_EQ(0U, F.Allocator.getBytesAllocated());
 }
 
 // Tests that a read which crosses a block boundary, but where the subsequent
 // blocks are still contiguous in memory to the previous block works and does
 // not allocate memory.
 TEST(MappedBlockStreamTest, ZeroCopyReadContiguousBreak) {
   DiscontiguousStream F(BlocksAry, DataAry);
-  auto S = MappedBlockStream::createStream(F.block_size(),
-                                           F.layout(), F);
+  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F,
+                                           F.Allocator);
   BinaryStreamReader R(*S);
   StringRef Str;
   EXPECT_NO_ERROR(R.readFixedString(Str, 2));
   EXPECT_EQ(Str, StringRef("AB"));
-  EXPECT_EQ(0U, S->getNumBytesCopied());
+  EXPECT_EQ(0U, F.Allocator.getBytesAllocated());
 
   R.setOffset(6);
   EXPECT_NO_ERROR(R.readFixedString(Str, 4));
   EXPECT_EQ(Str, StringRef("GHIJ"));
-  EXPECT_EQ(0U, S->getNumBytesCopied());
+  EXPECT_EQ(0U, F.Allocator.getBytesAllocated());
 }
 
 // Tests that a read which crosses a block boundary and cannot be referenced
 // contiguously works and allocates only the precise amount of bytes
 // requested.
 TEST(MappedBlockStreamTest, CopyReadNonContiguousBreak) {
   DiscontiguousStream F(BlocksAry, DataAry);
-  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F);
+  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F,
+                                           F.Allocator);
   BinaryStreamReader R(*S);
   StringRef Str;
   EXPECT_NO_ERROR(R.readFixedString(Str, 10));
   EXPECT_EQ(Str, StringRef("ABCDEFGHIJ"));
-  EXPECT_EQ(10U, S->getNumBytesCopied());
+  EXPECT_EQ(10U, F.Allocator.getBytesAllocated());
 }
 
 // Test that an out of bounds read which doesn't cross a block boundary
 // fails and allocates no memory.
 TEST(MappedBlockStreamTest, InvalidReadSizeNoBreak) {
   DiscontiguousStream F(BlocksAry, DataAry);
-  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F);
+  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F,
+                                           F.Allocator);
   BinaryStreamReader R(*S);
   StringRef Str;
 
   R.setOffset(10);
   EXPECT_ERROR(R.readFixedString(Str, 1));
-  EXPECT_EQ(0U, S->getNumBytesCopied());
+  EXPECT_EQ(0U, F.Allocator.getBytesAllocated());
 }
 
 // Test that an out of bounds read which crosses a contiguous block boundary
 // fails and allocates no memory.
 TEST(MappedBlockStreamTest, InvalidReadSizeContiguousBreak) {
   DiscontiguousStream F(BlocksAry, DataAry);
-  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F);
+  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F,
+                                           F.Allocator);
   BinaryStreamReader R(*S);
   StringRef Str;
 
   R.setOffset(6);
   EXPECT_ERROR(R.readFixedString(Str, 5));
-  EXPECT_EQ(0U, S->getNumBytesCopied());
+  EXPECT_EQ(0U, F.Allocator.getBytesAllocated());
 }
 
 // Test that an out of bounds read which crosses a discontiguous block
 // boundary fails and allocates no memory.
 TEST(MappedBlockStreamTest, InvalidReadSizeNonContiguousBreak) {
   DiscontiguousStream F(BlocksAry, DataAry);
-  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F);
+  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F,
+                                           F.Allocator);
   BinaryStreamReader R(*S);
   StringRef Str;
 
   EXPECT_ERROR(R.readFixedString(Str, 11));
-  EXPECT_EQ(0U, S->getNumBytesCopied());
+  EXPECT_EQ(0U, F.Allocator.getBytesAllocated());
 }
 
 // Tests that a read which is entirely contained within a single block but
 // beyond the end of a StreamRef fails.
 TEST(MappedBlockStreamTest, ZeroCopyReadNoBreak) {
   DiscontiguousStream F(BlocksAry, DataAry);
-  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F);
+  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F,
+                                           F.Allocator);
   BinaryStreamReader R(*S);
   StringRef Str;
   EXPECT_NO_ERROR(R.readFixedString(Str, 1));
   EXPECT_EQ(Str, StringRef("A"));
-  EXPECT_EQ(0U, S->getNumBytesCopied());
+  EXPECT_EQ(0U, F.Allocator.getBytesAllocated());
 }
 
 // Tests that a read which is not aligned on the same boundary as a previous
 // cached request, but which is known to overlap that request, shares the
 // previous allocation.
 TEST(MappedBlockStreamTest, UnalignedOverlappingRead) {
   DiscontiguousStream F(BlocksAry, DataAry);
-  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F);
+  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F,
+                                           F.Allocator);
   BinaryStreamReader R(*S);
   StringRef Str1;
   StringRef Str2;
   EXPECT_NO_ERROR(R.readFixedString(Str1, 7));
   EXPECT_EQ(Str1, StringRef("ABCDEFG"));
-  EXPECT_EQ(7U, S->getNumBytesCopied());
+  EXPECT_EQ(7U, F.Allocator.getBytesAllocated());
 
   R.setOffset(2);
   EXPECT_NO_ERROR(R.readFixedString(Str2, 3));
   EXPECT_EQ(Str2, StringRef("CDE"));
   EXPECT_EQ(Str1.data() + 2, Str2.data());
-  EXPECT_EQ(7U, S->getNumBytesCopied());
+  EXPECT_EQ(7U, F.Allocator.getBytesAllocated());
 }
 
 // Tests that a read which is not aligned on the same boundary as a previous
 // cached request, but which only partially overlaps a previous cached request,
 // still works correctly and allocates again from the shared pool.
 TEST(MappedBlockStreamTest, UnalignedOverlappingReadFail) {
   DiscontiguousStream F(BlocksAry, DataAry);
-  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F);
+  auto S = MappedBlockStream::createStream(F.block_size(), F.layout(), F,
+                                           F.Allocator);
   BinaryStreamReader R(*S);
   StringRef Str1;
   StringRef Str2;
   EXPECT_NO_ERROR(R.readFixedString(Str1, 6));
   EXPECT_EQ(Str1, StringRef("ABCDEF"));
-  EXPECT_EQ(6U, S->getNumBytesCopied());
+  EXPECT_EQ(6U, F.Allocator.getBytesAllocated());
 
   R.setOffset(4);
   EXPECT_NO_ERROR(R.readFixedString(Str2, 4));
   EXPECT_EQ(Str2, StringRef("EFGH"));
-  EXPECT_EQ(10U, S->getNumBytesCopied());
+  EXPECT_EQ(10U, F.Allocator.getBytesAllocated());
 }
 
 TEST(MappedBlockStreamTest, WriteBeyondEndOfStream) {
   static uint8_t Data[] = {'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J'};
   static uint8_t LargeBuffer[] = {'0', '1', '2', '3', '4', '5',
                                   '6', '7', '8', '9', 'A'};
   static uint8_t SmallBuffer[] = {'0', '1', '2'};
   static_assert(sizeof(LargeBuffer) > sizeof(Data),
                 "LargeBuffer is not big enough");
 
   DiscontiguousStream F(BlocksAry, Data);
-  auto S = WritableMappedBlockStream::createStream(
-      F.block_size(), F.layout(), F);
+  auto S = WritableMappedBlockStream::createStream(F.block_size(), F.layout(),
+                                                   F, F.Allocator);
   ArrayRef<uint8_t> Buffer;
 
   EXPECT_ERROR(S->writeBytes(0, ArrayRef<uint8_t>(LargeBuffer)));
   EXPECT_NO_ERROR(S->writeBytes(0, ArrayRef<uint8_t>(SmallBuffer)));
   EXPECT_NO_ERROR(S->writeBytes(7, ArrayRef<uint8_t>(SmallBuffer)));
   EXPECT_ERROR(S->writeBytes(8, ArrayRef<uint8_t>(SmallBuffer)));
 }
 
 TEST(MappedBlockStreamTest, TestWriteBytesNoBreakBoundary) {
   static uint8_t Data[] = {'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J'};
   DiscontiguousStream F(BlocksAry, Data);
-  auto S = WritableMappedBlockStream::createStream(
-      F.block_size(), F.layout(), F);
+  auto S = WritableMappedBlockStream::createStream(F.block_size(), F.layout(),
+                                                   F, F.Allocator);
   ArrayRef<uint8_t> Buffer;
 
   EXPECT_NO_ERROR(S->readBytes(0, 1, Buffer));
   EXPECT_EQ(Buffer, ArrayRef<uint8_t>('A'));
   EXPECT_NO_ERROR(S->readBytes(9, 1, Buffer));
   EXPECT_EQ(Buffer, ArrayRef<uint8_t>('J'));
 
   EXPECT_NO_ERROR(S->writeBytes(0, ArrayRef<uint8_t>('J')));

 
 TEST(MappedBlockStreamTest, TestWriteBytesBreakBoundary) {
   static uint8_t Data[] = {'0', '0', '0', '0', '0', '0', '0', '0', '0', '0'};
   static uint8_t TestData[] = {'T', 'E', 'S', 'T', 'I', 'N', 'G', '.'};
   static uint8_t Expected[] = {'T', 'E', 'S', 'N', 'I',
                                'T', 'G', '.', '0', '0'};
 
   DiscontiguousStream F(BlocksAry, Data);
-  auto S = WritableMappedBlockStream::createStream(
-      F.block_size(), F.layout(), F);
+  auto S = WritableMappedBlockStream::createStream(F.block_size(), F.layout(),
+                                                   F, F.Allocator);
   ArrayRef<uint8_t> Buffer;
 
   EXPECT_NO_ERROR(S->writeBytes(0, TestData));
   // First just compare the memory, then compare the result of reading the
   // string out.
   EXPECT_EQ(ArrayRef<uint8_t>(Data), ArrayRef<uint8_t>(Expected));
 
   EXPECT_NO_ERROR(S->readBytes(0, 8, Buffer));
   EXPECT_EQ(Buffer, ArrayRef<uint8_t>(TestData));
 }
 
 TEST(MappedBlockStreamTest, TestWriteThenRead) {
   std::vector<uint8_t> DataBytes(10);
   MutableArrayRef<uint8_t> Data(DataBytes);
   const uint32_t Blocks[] = {2, 1, 0, 6, 3, 4, 5, 7, 9, 8};
 
   DiscontiguousStream F(Blocks, Data);
-  auto S = WritableMappedBlockStream::createStream(
-      F.block_size(), F.layout(), F);
+  auto S = WritableMappedBlockStream::createStream(F.block_size(), F.layout(),
+                                                   F, F.Allocator);
 
   enum class MyEnum : uint32_t { Val1 = 2908234, Val2 = 120891234 };
   using support::ulittle32_t;
 
   uint16_t u16[] = {31468, 0};
   uint32_t u32[] = {890723408, 0};
   MyEnum Enum[] = {MyEnum::Val1, MyEnum::Val2};
   StringRef ZStr[] = {"Zero Str", ""};
TEST(MappedBlockStreamTest, TestWriteContiguousStreamRef) {
   MutableArrayRef<uint8_t> DestData(DestDataBytes);
   const uint32_t DestBlocks[] = {2, 1, 0, 6, 3, 4, 5, 7, 9, 8};
 
   std::vector<uint8_t> SrcDataBytes(10);
   MutableArrayRef<uint8_t> SrcData(SrcDataBytes);
 
   DiscontiguousStream F(DestBlocks, DestData);
   auto DestStream = WritableMappedBlockStream::createStream(
-      F.block_size(), F.layout(), F);
+      F.block_size(), F.layout(), F, F.Allocator);
 
   // First write "Test Str" into the source stream.
   MutableBinaryByteStream SourceStream(SrcData, little);
   BinaryStreamWriter SourceWriter(SourceStream);
   EXPECT_NO_ERROR(SourceWriter.writeCString("Test Str"));
   EXPECT_EQ(SrcDataBytes, std::vector<uint8_t>(
                               {'T', 'e', 's', 't', ' ', 'S', 't', 'r', 0, 0}));
 
TEST(MappedBlockStreamTest, TestWriteDiscontiguousStreamRef) {
   std::vector<uint8_t> SrcDataBytes(10);
   MutableArrayRef<uint8_t> SrcData(SrcDataBytes);
   const uint32_t SrcBlocks[] = {1, 0, 6, 3, 4, 5, 2, 7, 8, 9};
 
   DiscontiguousStream DestF(DestBlocks, DestData);
   DiscontiguousStream SrcF(SrcBlocks, SrcData);
 
   auto Dest = WritableMappedBlockStream::createStream(
-      DestF.block_size(), DestF.layout(), DestF);
+      DestF.block_size(), DestF.layout(), DestF, DestF.Allocator);
   auto Src = WritableMappedBlockStream::createStream(
-      SrcF.block_size(), SrcF.layout(), SrcF);
+      SrcF.block_size(), SrcF.layout(), SrcF, SrcF.Allocator);
 
   // First write "Test Str" into the source stream.
   BinaryStreamWriter SourceWriter(*Src);
   EXPECT_NO_ERROR(SourceWriter.writeCString("Test Str"));
   EXPECT_EQ(SrcDataBytes, std::vector<uint8_t>(
                               {'e', 'T', 't', 't', ' ', 'S', 's', 'r', 0, 0}));
 
   // Then write the source stream into the dest stream.
   BinaryStreamWriter DestWriter(*Dest);
   EXPECT_NO_ERROR(DestWriter.writeStreamRef(*Src));
   EXPECT_EQ(DestDataBytes, std::vector<uint8_t>(
                                {'s', 'e', 'T', ' ', 'S', 't', 't', 'r', 0, 0}));
 
   // Then read the string back out of the dest stream.
   StringRef Result;
   BinaryStreamReader DestReader(*Dest);
   EXPECT_NO_ERROR(DestReader.readCString(Result));
   EXPECT_EQ(Result, "Test Str");
 }
 
+TEST(MappedBlockStreamTest, DataLivesAfterStreamDestruction) {
+  std::vector<uint8_t> DataBytes(10);
+  MutableArrayRef<uint8_t> Data(DataBytes);
+  const uint32_t Blocks[] = {2, 1, 0, 6, 3, 4, 5, 7, 9, 8};
+
+  StringRef Str[] = {"Zero Str", ""};
+
+  DiscontiguousStream F(Blocks, Data);
+  {
+    auto S = WritableMappedBlockStream::createStream(F.block_size(), F.layout(),
+                                                     F, F.Allocator);
+
+    BinaryStreamReader Reader(*S);
+    BinaryStreamWriter Writer(*S);
+    ::memset(DataBytes.data(), 0, 10);
+    EXPECT_NO_ERROR(Writer.writeCString(Str[0]));
+    EXPECT_NO_ERROR(Reader.readCString(Str[1]));
+    EXPECT_EQ(Str[0], Str[1]);
+  }
+
+  EXPECT_EQ(Str[0], Str[1]);
+}
+
 } // end anonymous namespace