From 8d6a31c023a3eaf0817cf83526f87bd53997ffc0 Mon Sep 17 00:00:00 2001
From: "Duncan P. N. Exon Smith" <dexonsmith@apple.com>
Date: Sat, 20 Feb 2016 21:24:31 +0000
Subject: [PATCH] Lex: Check whether the header map buffer has space for the
 buckets

Check up front whether the header map buffer has space for all of its
declared buckets.

There was already a check in `getBucket()`, but it had UB (comparing
pointers that were outside of objects in the error path) and was
insufficient (only checking for a single byte of the relevant bucket).
I fixed the check, moved it to `checkHeader()`, and left a fixed version
behind as an assertion.

llvm-svn: 261449
---
 clang/lib/Lex/HeaderMap.cpp           | 20 ++++++++++----------
 clang/unittests/Lex/HeaderMapTest.cpp |  8 ++++++++
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/clang/lib/Lex/HeaderMap.cpp b/clang/lib/Lex/HeaderMap.cpp
index 2b31ab9..220e70d 100644
--- a/clang/lib/Lex/HeaderMap.cpp
+++ b/clang/lib/Lex/HeaderMap.cpp
@@ -83,14 +83,16 @@ bool HeaderMapImpl::checkHeader(const llvm::MemoryBuffer &File,
   if (Header->Reserved != 0)
     return false;
 
-  // Check the number of buckets.
+  // Check the number of buckets.  It should be a power of two, and there
+  // should be enough space in the file for all of them.
   auto NumBuckets = NeedsByteSwap
                         ? llvm::sys::getSwappedBytes(Header->NumBuckets)
                         : Header->NumBuckets;
-
-  // If the number of buckets is not a power of two, the headermap is corrupt.
   if (NumBuckets & (NumBuckets - 1))
     return false;
+  if (File.getBufferSize() <
+      sizeof(HMapHeader) + sizeof(HMapBucket) * NumBuckets)
+    return false;
 
   // Okay, everything looks good.
   return true;
@@ -122,21 +124,19 @@ const HMapHeader &HeaderMapImpl::getHeader() const {
 /// bswap'ing its fields as appropriate.  If the bucket number is not valid,
 /// this return a bucket with an empty key (0).
 HMapBucket HeaderMapImpl::getBucket(unsigned BucketNo) const {
+  assert(FileBuffer->getBufferSize() >=
+             sizeof(HMapHeader) + sizeof(HMapBucket) * BucketNo &&
+         "Expected bucket to be in range");
+
   HMapBucket Result;
   Result.Key = HMAP_EmptyBucketKey;
 
   const HMapBucket *BucketArray =
     reinterpret_cast<const HMapBucket*>(FileBuffer->getBufferStart() +
                                         sizeof(HMapHeader));
-
   const HMapBucket *BucketPtr = BucketArray+BucketNo;
-  if ((const char*)(BucketPtr+1) > FileBuffer->getBufferEnd()) {
-    Result.Prefix = 0;
-    Result.Suffix = 0;
-    return Result;  // Invalid buffer, corrupt hmap.
-  }
 
-  // Otherwise, the bucket is valid.  Load the values, bswapping as needed.
+  // Load the values, bswapping as needed.
   Result.Key    = getEndianAdjustedWord(BucketPtr->Key);
   Result.Prefix = getEndianAdjustedWord(BucketPtr->Prefix);
   Result.Suffix = getEndianAdjustedWord(BucketPtr->Suffix);
diff --git a/clang/unittests/Lex/HeaderMapTest.cpp b/clang/unittests/Lex/HeaderMapTest.cpp
index b8e1d6c..33befd8 100644
--- a/clang/unittests/Lex/HeaderMapTest.cpp
+++ b/clang/unittests/Lex/HeaderMapTest.cpp
@@ -100,4 +100,12 @@ TEST(HeaderMapTest, checkHeader3Buckets) {
   ASSERT_FALSE(HeaderMapImpl::checkHeader(*File.getBuffer(), NeedsSwap));
 }
 
+TEST(HeaderMapTest, checkHeaderNotEnoughBuckets) {
+  MapFile<1, 1> File;
+  File.init();
+  File.Header.NumBuckets = 8;
+  bool NeedsSwap;
+  ASSERT_FALSE(HeaderMapImpl::checkHeader(*File.getBuffer(), NeedsSwap));
+}
+
 } // end namespace
-- 
2.7.4