Add the ability to edit stream contents with llvm-pdbutil

[zturner]

Summary

For experimentation and diagnostic purposes, it's often useful to be able to replace the contents of one stream with a binary blob. For example, if you have one PDB that works and one that doesn't, this can be useful for helping understand why the broken one is broken. You can replace stream contents 1 by 1 until something starts working. This is trickier than it sounds though, because if a stream's size changes it changes the stream directory as well, and in some cases could even cause the file to grow (or shrink).

To do this we load the old MSF layout and create an MSFBuilder out of it with the exact same properties as the existing file. Then we load the new stream contents and update the stream size, which should automatically decide how big the new file should be and where the new stream's blocks will be. Finally, we re-open the original file in a special mode that allows partial writes, keeping any untouched data intact, then write out only the modified MSF layout and new stream contents. In some cases -- particularly if a stream shrinks -- there will be "stale" data in the PDB file, but this shouldn't matter because the stream directory keeps track of a stream's *actual* size, so any tool reading the PDB would just ignore the stale data.

[inglorion]

This is really useful. Thanks for doing this. I got a couple of questions and comments in-line.

llvm/include/llvm/DebugInfo/MSF/MSFBuilder.h

125

It's not clear to me from the names and the comments what exactly this method does and what the parameters mean. I think something like commitUpdate(StringRef Path, MSFLayout NewLayout) would make it clearer that we're committing an update to an existing MSF file, and the parameter is supposed to be the new layout and not the old one.

llvm/lib/DebugInfo/MSF/MSFBuilder.cpp

51

Why +2?

llvm/tools/llvm-pdbutil/llvm-pdbutil.cpp

688

nit: in llvm/Option/ArgList.h, we have a similar method, but it's called hasArg. I would use the same name here, in part because it's shorter. :-)

1250

Shouldn't this assert that StreamFile was _not_ passed?

1263

Is there a way to preserve the original error information here, so that one can distinguish between for example a non-existent file and a file one does not have permission to read?

1277

I think this should be ExitOnErr rather than cantFail, similar to other uses of FileBufferByteStream::commit().

llvm/include/llvm/DebugInfo/MSF/MSFBuilder.h

public:
   /// failed.  Currently the only way this can fail is if an invalid block size
   /// is specified, or `MinBlockCount` does not leave enough room for the
   /// mandatory reserved blocks required by an MSF file.
   static Expected<MSFBuilder> create(BumpPtrAllocator &Allocator,
                                      uint32_t BlockSize,
                                      uint32_t MinBlockCount = 0,
                                      bool CanGrow = true);
 
+  /// Create a new `MSFBuilder` from an existing layout.
+  static MSFBuilder create(BumpPtrAllocator &Allocator,
+                           const MSFLayout &Layout);
+
   /// Request the block map to be at a specific block address.  This is useful
   /// when editing a MSF and you want the layout to be as stable as possible.
   Error setBlockMapAddr(uint32_t Addr);
   Error setDirectoryBlocksHint(ArrayRef<uint32_t> DirBlocks);
   void setFreePageMap(uint32_t Fpm);
   void setUnknown1(uint32_t Unk1);
 
   /// Add a stream to the MSF file with the given size, occupying the given
public:
 
   /// Finalize the layout and build the headers and structures that describe the
   /// MSF layout and can be written directly to the MSF file.
   Expected<MSFLayout> generateLayout();
 
   /// Write the MSF layout to the underlying file.
   Expected<FileBufferByteStream> commit(StringRef Path, MSFLayout &Layout);
 
+  /// Write the MSF layout to the underlying file, leaving any contents of the
+  /// existing file untouched.
+  Expected<FileBufferByteStream> commitPartial(StringRef Path,
+                                               MSFLayout &Layout);
+
   BumpPtrAllocator &getAllocator() { return Allocator; }
 
 private:
+  MSFBuilder(BumpPtrAllocator &Allocator);
   MSFBuilder(uint32_t BlockSize, uint32_t MinBlockCount, bool CanGrow,
              BumpPtrAllocator &Allocator);
 
+  Expected<FileBufferByteStream>
+  commitInternal(StringRef Path, MSFLayout &Layout, unsigned Flags);
+
   Error allocateBlocks(uint32_t NumBlocks, MutableArrayRef<uint32_t> Blocks);
   uint32_t computeDirectoryByteSize() const;
 
   using BlockList = std::vector<uint32_t>;
 
   BumpPtrAllocator &Allocator;
 
   bool IsGrowable;

llvm/lib/DebugInfo/MSF/MSFBuilder.cpp

: Allocator(Allocator), IsGrowable(CanGrow),
       FreePageMap(kDefaultFreePageMap), BlockSize(BlockSize),
       BlockMapAddr(kDefaultBlockMapAddr), FreeBlocks(MinBlockCount, true) {
   FreeBlocks[kSuperBlockBlock] = false;
   FreeBlocks[kFreePageMap0Block] = false;
   FreeBlocks[kFreePageMap1Block] = false;
   FreeBlocks[BlockMapAddr] = false;
 }
 
+MSFBuilder::MSFBuilder(BumpPtrAllocator &Allocator)
+    : MSFBuilder(4096, msf::getMinimumBlockCount() + 2, true, Allocator) {}
+
 Expected<MSFBuilder> MSFBuilder::create(BumpPtrAllocator &Allocator,
                                         uint32_t BlockSize,
                                         uint32_t MinBlockCount, bool CanGrow) {
   if (!isValidBlockSize(BlockSize))
     return make_error<MSFError>(msf_error_code::invalid_format,
                                 "The requested block size is unsupported");
 
   return MSFBuilder(BlockSize,
                     std::max(MinBlockCount, msf::getMinimumBlockCount()),
                     CanGrow, Allocator);
 }
 
+MSFBuilder MSFBuilder::create(BumpPtrAllocator &Allocator,
+                              const MSFLayout &Layout) {
+  MSFBuilder Result(Allocator);
+  Result.IsGrowable = true;
+  Result.setFreePageMap(Layout.SB->FreeBlockMapBlock);
+  Result.setUnknown1(Layout.SB->Unknown1);
+  Result.BlockSize = Layout.SB->BlockSize;
+  Result.BlockMapAddr = Layout.SB->BlockMapAddr;
+  Result.FreeBlocks = Layout.FreePageMap;
+  Result.DirectoryBlocks.assign(Layout.DirectoryBlocks.begin(),
+                                Layout.DirectoryBlocks.end());
+
+  assert(Layout.StreamMap.size() == Layout.StreamSizes.size());
+  Result.StreamData.reserve(Layout.StreamMap.size());
+  for (size_t I = 0; I < Layout.StreamMap.size(); ++I) {
+    uint32_t Size = Layout.StreamSizes[I];
+    ArrayRef<support::ulittle32_t> Blocks = Layout.StreamMap[I];
+    std::vector<uint32_t> BlockVector{Blocks.begin(), Blocks.end()};
+    Result.StreamData.emplace_back(Size, std::move(BlockVector));
+  }
+  return std::move(Result);
+}
+
 Error MSFBuilder::setBlockMapAddr(uint32_t Addr) {
   if (Addr == BlockMapAddr)
     return Error::success();
 
   if (Addr >= FreeBlocks.size()) {
     if (!IsGrowable)
       return make_error<MSFError>(msf_error_code::insufficient_buffer,
                                   "Cannot grow the number of blocks");
for (uint32_t I = 0; I < 8; ++I) {
       ThisByte |= Mask;
       ++BI;
     }
     cantFail(FpmWriter.writeObject(ThisByte));
   }
   assert(FpmWriter.bytesRemaining() == 0);
 }
 
-Expected<FileBufferByteStream> MSFBuilder::commit(StringRef Path,
-                                                  MSFLayout &Layout) {
+Expected<FileBufferByteStream>
+MSFBuilder::commitInternal(StringRef Path, MSFLayout &Layout, unsigned Flags) {
   Expected<MSFLayout> L = generateLayout();
   if (!L)
     return L.takeError();
 
   Layout = std::move(*L);
 
   uint64_t FileSize = Layout.SB->BlockSize * Layout.SB->NumBlocks;
-  auto OutFileOrError = FileOutputBuffer::create(Path, FileSize);
+  auto OutFileOrError = FileOutputBuffer::create(Path, FileSize, Flags);
   if (auto EC = OutFileOrError.takeError())
     return std::move(EC);
 
   FileBufferByteStream Buffer(std::move(*OutFileOrError),
                               llvm::support::little);
   BinaryStreamWriter Writer(Buffer);
 
   if (auto EC = Writer.writeObject(*Layout.SB))
MSFBuilder::commitInternal(StringRef Path, MSFLayout &Layout, unsigned Flags) {
 
   for (const auto &Blocks : Layout.StreamMap) {
     if (auto EC = DW.writeArray(Blocks))
       return std::move(EC);
   }
 
   return std::move(Buffer);
 }
+
+/// Write the MSF layout to the underlying file, leaving any contents of the
+/// existing file untouched.
+Expected<FileBufferByteStream> MSFBuilder::commitPartial(StringRef Path,
+                                                         MSFLayout &Layout) {
+  return commitInternal(Path, Layout, FileOutputBuffer::F_modify);
+}
+
+Expected<FileBufferByteStream> MSFBuilder::commit(StringRef Path,
+                                                  MSFLayout &Layout) {
+  return commitInternal(Path, Layout, 0);
+}

llvm/tools/llvm-pdbutil/llvm-pdbutil.cpp

cl::SubCommand MergeSubcommand("merge",
                                "Merge multiple PDBs into a single PDB");
 
 cl::SubCommand ExplainSubcommand("explain",
                                  "Explain the meaning of a file offset");
 
 cl::SubCommand ExportSubcommand("export",
                                 "Write binary data from a stream to a file");
 
+cl::SubCommand EditStreamSubcommand("edit-stream",
+                                    "Edit the contents of a PDB stream");
+
 cl::OptionCategory TypeCategory("Symbol Type Options");
 cl::OptionCategory FilterCategory("Filtering and Sorting Options");
 cl::OptionCategory OtherOptions("Other Options");
 
 cl::ValuesClass ChunkValues = cl::values(
     clEnumValN(ModuleSubsection::CrossScopeExports, "cme",
                "Cross module exports (DEBUG_S_CROSSSCOPEEXPORTS subsection)"),
     clEnumValN(ModuleSubsection::CrossScopeImports, "cmi",
Stream("stream", cl::Required,
            cl::desc("The index or name of the stream whose contents to export"),
            cl::sub(ExportSubcommand));
 cl::opt<bool> ForceName("name",
                         cl::desc("Force the interpretation of -stream as a "
                                  "string, even if it is a valid integer"),
                         cl::sub(ExportSubcommand), cl::Optional,
                         cl::init(false));
 } // namespace exportstream
+
+namespace writestream {
+cl::list<std::string> OutputPdbFilename(cl::Positional,
+                                        cl::desc("<output PDB file>"),
+                                        cl::Required,
+                                        cl::sub(EditStreamSubcommand));
+cl::opt<std::string>
+    StreamFile("stream-file",
+               cl::desc("A file from which to set the stream contents"),
+               cl::Optional, cl::sub(EditStreamSubcommand));
+cl::opt<uint32_t>
+    StreamSize("set-size",
+               cl::desc("Set the size of the stream to the specified value"),
+               cl::Optional, cl::sub(EditStreamSubcommand));
+cl::opt<std::string>
+    Stream("stream", cl::Required,
+           cl::desc("The index or name of the stream whose contents to import"),
+           cl::sub(EditStreamSubcommand));
+cl::opt<bool> ForceName("name",
+                        cl::desc("Force the interpretation of -stream as a "
+                                 "string, even if it is a valid integer"),
+                        cl::sub(EditStreamSubcommand), cl::Optional,
+                        cl::init(false));
+} // namespace writestream
 }
 
 static ExitOnError ExitOnErr;
 
+template <typename OptType> static bool optWasPassed(const OptType &Opt) {
+  return Opt.getNumOccurrences() > 0;
+}
+
 static void yamlToPdb(StringRef Path) {
   BumpPtrAllocator Allocator;
   ErrorOr<std::unique_ptr<MemoryBuffer>> ErrorOrBuffer =
       MemoryBuffer::getFileOrSTDIN(Path, /*FileSize=*/-1,
                                    /*RequiresNullTerminator=*/false);
 
   if (ErrorOrBuffer.getError()) {
     ExitOnErr(make_error<GenericError>(generic_error_code::invalid_path, Path));
static void explain() {
 
   for (uint64_t Off : opts::explain::Offsets) {
     auto O = llvm::make_unique<ExplainOutputStyle>(IF, Off);
 
     ExitOnErr(O->dump());
   }
 }
 
+static Expected<uint32_t> getStreamIndex(PDBFile &File, StringRef Stream,
+                                         bool ForceName) {
+  uint32_t Index = 0;
+  // Make sure that the stream the user is requesting actually exists.
+  if (!ForceName) {
+    // First try to parse it as an integer, if it fails fall back to treating it
+    // as a named stream.
+    if (to_integer(Stream, Index)) {
+      if (Index >= File.getNumStreams()) {
+        errs() << "Error: " << Index << " is not a valid stream index.\n";
+        exit(1);
+      }
+      return Index;
+    }
+  }
+
+  InfoStream &IS = cantFail(File.getPDBInfoStream());
+  Index = ExitOnErr(IS.getNamedStreamIndex(opts::exportstream::Stream));
+  return Index;
+}
+
 static void exportStream() {
   std::unique_ptr<IPDBSession> Session;
   PDBFile &File = loadPDB(opts::exportstream::InputFilename.front(), Session);
 
   std::unique_ptr<MappedBlockStream> SourceStream;
   uint32_t Index = 0;
   bool Success = false;
   std::string OutFileName = opts::exportstream::OutputFile;
static void exportStream() {
   auto OutFile = ExitOnErr(
       FileOutputBuffer::create(OutFileName, SourceStream->getLength()));
   FileBufferByteStream DestStream(std::move(OutFile), llvm::support::little);
   BinaryStreamWriter Writer(DestStream);
   ExitOnErr(Writer.writeStreamRef(*SourceStream));
   ExitOnErr(DestStream.commit());
 }
 
+static void editStream() {
+  uint32_t StreamIndex = 0;
+  std::error_code EC;
+  std::string FileName = opts::writestream::OutputPdbFilename.front();
+
+  BumpPtrAllocator Allocator;
+  Optional<MSFBuilder> Builder;
+  // Load the file first to get the current MSF layout.
+  {
+    std::unique_ptr<IPDBSession> Session;
+    PDBFile &File = loadPDB(FileName, Session);
+    StreamIndex = ExitOnErr(getStreamIndex(File, opts::writestream::Stream,
+                                           opts::writestream::ForceName));
+    // We need to do this part before the session goes out of scope, since
+    // MSFLayout uses memory owned by its session's allocator, and we need
+    // that to live long enough to "clone" the layout in a new MSFBuilder.
+    Builder.emplace(MSFBuilder::create(Allocator, File.getMsfLayout()));
+  }
+
+  MSFLayout NewLayout;
+  if (optWasPassed(opts::writestream::StreamSize)) {
+    assert(optWasPassed(opts::writestream::StreamFile));
+    cantFail(
+        Builder->setStreamSize(StreamIndex, opts::writestream::StreamSize));
+    FileBufferByteStream OutputBuffer =
+        ExitOnErr(Builder->commitPartial(FileName, NewLayout));
+
+    cantFail(OutputBuffer.commit());
+    return;
+  }
+
+  auto ErrorOrFile = MemoryBuffer::getFile(opts::writestream::StreamFile);
+  if (ErrorOrFile.getError())
+    ExitOnErr(make_error<GenericError>(generic_error_code::invalid_path,
+                                       opts::writestream::StreamFile));
+
+  auto StreamContents = std::move(*ErrorOrFile);
+  cantFail(
+      Builder->setStreamSize(StreamIndex, StreamContents->getBufferSize()));
+
+  FileBufferByteStream OutputBuffer =
+      ExitOnErr(Builder->commitPartial(FileName, NewLayout));
+
+  auto Stream = WritableMappedBlockStream::createIndexedStream(
+      NewLayout, OutputBuffer, StreamIndex, Allocator);
+  BinaryStreamWriter Writer(*Stream);
+  cantFail(
+      Writer.writeBytes(arrayRefFromStringRef(StreamContents->getBuffer())));
+  cantFail(OutputBuffer.commit());
+}
+
 static bool parseRange(StringRef Str,
                        Optional<opts::bytes::NumberRange> &Parsed) {
   if (Str.empty())
     return true;
 
   llvm::Regex R("^([^-]+)(-([^-]+))?$");
   llvm::SmallVector<llvm::StringRef, 2> Matches;
   if (!R.match(Str, &Matches))
if (opts::merge::InputFilenames.size() < 2) {
       errs() << "merge subcommand requires at least 2 input files.\n";
       exit(1);
     }
     mergePdbs();
   } else if (opts::ExplainSubcommand) {
     explain();
   } else if (opts::ExportSubcommand) {
     exportStream();
+  } else if (opts::EditStreamSubcommand) {
+    if (optWasPassed(opts::writestream::StreamSize) ==
+        optWasPassed(opts::writestream::StreamFile)) {
+      errs() << "Error: Exactly one of `-stream-size` or `-stream-file` must "
+                "be specified.\n";
+      exit(1);
+    }
+
+    editStream();
   }
 
   outs().flush();
   return 0;
 }