From 818581cd0c19a196ba2db956a7a8b486b3e85d56 Mon Sep 17 00:00:00 2001
From: David Majnemer <david.majnemer@gmail.com>
Date: Wed, 1 Jun 2016 18:13:06 +0000
Subject: [PATCH] [CodeView] Make sure StreamRef::readBytes doesn't read too
 much

llvm-svn: 271418
---
 include/llvm/DebugInfo/CodeView/StreamRef.h   |  5 ++++-
 .../DebugInfo/PDB/MappedBlockStreamTest.cpp   | 22 ++++++++++++++-----
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/include/llvm/DebugInfo/CodeView/StreamRef.h b/include/llvm/DebugInfo/CodeView/StreamRef.h
index 7fdfb2ad8f7..2242b6783e3 100644
--- a/include/llvm/DebugInfo/CodeView/StreamRef.h
+++ b/include/llvm/DebugInfo/CodeView/StreamRef.h
@@ -10,6 +10,7 @@
 #ifndef LLVM_DEBUGINFO_CODEVIEW_STREAMREF_H
 #define LLVM_DEBUGINFO_CODEVIEW_STREAMREF_H
 
+#include "llvm/DebugInfo/CodeView/CodeViewError.h"
 #include "llvm/DebugInfo/CodeView/StreamInterface.h"
 
 namespace llvm {
@@ -28,6 +29,8 @@ public:
 
   Error readBytes(uint32_t Offset, uint32_t Size,
                   ArrayRef<uint8_t> &Buffer) const override {
+    if (Size > Length)
+      return make_error<CodeViewError>(cv_error_code::insufficient_buffer);
     return Stream->readBytes(ViewOffset + Offset, Size, Buffer);
   }
 
@@ -74,4 +77,4 @@ private:
 }
 }
 
-#endif // LLVM_DEBUGINFO_CODEVIEW_STREAMREF_H
\ No newline at end of file
+#endif // LLVM_DEBUGINFO_CODEVIEW_STREAMREF_H
diff --git a/unittests/DebugInfo/PDB/MappedBlockStreamTest.cpp b/unittests/DebugInfo/PDB/MappedBlockStreamTest.cpp
index 03f2befaed5..aa579a3700a 100644
--- a/unittests/DebugInfo/PDB/MappedBlockStreamTest.cpp
+++ b/unittests/DebugInfo/PDB/MappedBlockStreamTest.cpp
@@ -71,14 +71,14 @@ private:
 
 // Tests that a read which is entirely contained within a single block works
 // and does not allocate.
-TEST(MappedBlockStreamTest, ZeroCopyReadNoBreak) {
+TEST(MappedBlockStreamTest, ReadBeyondEndOfStreamRef) {
   DiscontiguousFile F;
   MappedBlockStream S(0, F);
   StreamReader R(S);
-  StringRef Str;
-  EXPECT_NO_ERROR(R.readFixedString(Str, 1));
-  EXPECT_EQ(Str, StringRef("A"));
-  EXPECT_EQ(0U, S.getNumBytesCopied());
+  StreamRef SR;
+  EXPECT_NO_ERROR(R.readStreamRef(SR, 0U));
+  ArrayRef<uint8_t> Buffer;
+  EXPECT_ERROR(SR.readBytes(0U, 1U, Buffer));
 }
 
 // Tests that a read which outputs into a full destination buffer works and
@@ -162,4 +162,16 @@ TEST(MappedBlockStreamTest, InvalidReadSizeNonContiguousBreak) {
   EXPECT_EQ(0U, S.getNumBytesCopied());
 }
 
+// Tests that a read which is entirely contained within a single block but
+// beyond the end of a StreamRef fails.
+TEST(MappedBlockStreamTest, ZeroCopyReadNoBreak) {
+  DiscontiguousFile F;
+  MappedBlockStream S(0, F);
+  StreamReader R(S);
+  StringRef Str;
+  EXPECT_NO_ERROR(R.readFixedString(Str, 1));
+  EXPECT_EQ(Str, StringRef("A"));
+  EXPECT_EQ(0U, S.getNumBytesCopied());
+}
+
 } // end anonymous namespace