[pdb] Try to fix use after free.

llvm-svn: 272078

5 files changed, 21 insertions, 2 deletions

diff --git a/llvm/lib/DebugInfo/PDB/Raw/DbiStream.cpp b/llvm/lib/DebugInfo/PDB/Raw/DbiStream.cpp
index 80441ccd825..90bc2a2e556 100644
--- a/llvm/lib/DebugInfo/PDB/Raw/DbiStream.cpp
+++ b/llvm/lib/DebugInfo/PDB/Raw/DbiStream.cpp
@@ -293,6 +293,9 @@ Error DbiStream::initializeSectionContributionData() {
 // Initializes this->SectionHeaders.
 Error DbiStream::initializeSectionHeadersData() {
   uint32_t StreamNum = getDebugStreamIndex(DbgHeaderType::SectionHdr);
+  if (StreamNum >= Pdb.getNumStreams())
+    return make_error<RawError>(raw_error_code::no_stream);
+
   SectionHeaderStream.reset(new MappedBlockStream(
       llvm::make_unique<IndexedStreamData>(StreamNum, Pdb), Pdb));
 
@@ -312,6 +315,9 @@ Error DbiStream::initializeSectionHeadersData() {
 // Initializes this->Fpos.
 Error DbiStream::initializeFpoRecords() {
   uint32_t StreamNum = getDebugStreamIndex(DbgHeaderType::NewFPO);
+  if (StreamNum >= Pdb.getNumStreams())
+    return make_error<RawError>(raw_error_code::no_stream);
+
   FpoStream.reset(new MappedBlockStream(
       llvm::make_unique<IndexedStreamData>(StreamNum, Pdb), Pdb));
 
diff --git a/llvm/lib/DebugInfo/PDB/Raw/PDBFile.cpp b/llvm/lib/DebugInfo/PDB/Raw/PDBFile.cpp
index 22094773624..2796abf3ea5 100644
--- a/llvm/lib/DebugInfo/PDB/Raw/PDBFile.cpp
+++ b/llvm/lib/DebugInfo/PDB/Raw/PDBFile.cpp
@@ -325,6 +325,9 @@ Expected<NameHashTable &> PDBFile::getStringTable() {
 
     if (NameStreamIndex == 0)
       return make_error<RawError>(raw_error_code::no_stream);
+    if (NameStreamIndex >= getNumStreams())
+      return make_error<RawError>(raw_error_code::no_stream);
+
     auto SD = llvm::make_unique<IndexedStreamData>(NameStreamIndex, *this);
     auto S = llvm::make_unique<MappedBlockStream>(std::move(SD), *this);
     codeview::StreamReader Reader(*S);
diff --git a/llvm/lib/DebugInfo/PDB/Raw/TpiStream.cpp b/llvm/lib/DebugInfo/PDB/Raw/TpiStream.cpp
index 386f8ac80a5..f34a513f499 100644
--- a/llvm/lib/DebugInfo/PDB/Raw/TpiStream.cpp
+++ b/llvm/lib/DebugInfo/PDB/Raw/TpiStream.cpp
@@ -104,6 +104,10 @@ Error TpiStream::reload() {
     return EC;
 
   // Hash indices, hash values, etc come from the hash stream.
+  if (Header->HashStreamIndex >= Pdb.getNumStreams())
+    return make_error<RawError>(raw_error_code::corrupt_file,
+                                "Invalid TPI hash stream index.");
+
   HashStream.reset(new MappedBlockStream(
       llvm::make_unique<IndexedStreamData>(Header->HashStreamIndex, Pdb), Pdb));
   codeview::StreamReader HSR(*HashStream);
diff --git a/llvm/tools/llvm-pdbdump/LLVMOutputStyle.cpp b/llvm/tools/llvm-pdbdump/LLVMOutputStyle.cpp
index b287243d0c4..20d08f5c8ff 100644
--- a/llvm/tools/llvm-pdbdump/LLVMOutputStyle.cpp
+++ b/llvm/tools/llvm-pdbdump/LLVMOutputStyle.cpp
@@ -191,10 +191,12 @@ Error LLVMOutputStyle::dumpStreamData() {
   uint32_t StreamCount = File.getNumStreams();
   StringRef DumpStreamStr = opts::DumpStreamDataIdx;
   uint32_t DumpStreamNum;
-  if (DumpStreamStr.getAsInteger(/*Radix=*/0U, DumpStreamNum) ||
-      DumpStreamNum >= StreamCount)
+  if (DumpStreamStr.getAsInteger(/*Radix=*/0U, DumpStreamNum))
     return Error::success();
 
+  if (DumpStreamNum >= StreamCount)
+    return make_error<RawError>(raw_error_code::no_stream);
+
   MappedBlockStream S(llvm::make_unique<IndexedStreamData>(DumpStreamNum, File),
                       File);
   codeview::StreamReader R(S);
@@ -238,6 +240,8 @@ Error LLVMOutputStyle::dumpNamedStream() {
   InfoStream &IS = InfoS.get();
 
   uint32_t NameStreamIndex = IS.getNamedStreamIndex(opts::DumpStreamDataName);
+  if (NameStreamIndex == 0 || NameStreamIndex >= File.getNumStreams())
+    return make_error<RawError>(raw_error_code::no_stream);
 
   if (NameStreamIndex != 0) {
     std::string Name("Stream '");
diff --git a/llvm/tools/llvm-pdbdump/llvm-pdbdump.cpp b/llvm/tools/llvm-pdbdump/llvm-pdbdump.cpp
index 7d4ba6471f1..2445664db21 100644
--- a/llvm/tools/llvm-pdbdump/llvm-pdbdump.cpp
+++ b/llvm/tools/llvm-pdbdump/llvm-pdbdump.cpp
@@ -294,6 +294,8 @@ bool isRawDumpEnabled() {
     return true;
   if (opts::DumpIpiRecordBytes)
     return true;
+  if (opts::DumpSectionHeaders)
+    return true;
   if (opts::DumpSectionContribs)
     return true;
   if (opts::DumpSectionMap)