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.

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@261449 91177308-0d34-0410-b5e6-96231b3b80d8

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);

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