Lex: Check buckets on header map construction

If the number of buckets is not a power of two, immediately recognize
the header map as corrupt, rather than waiting for the first lookup.  I
converted the later check to an assert.

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

diff --git a/lib/Lex/HeaderMap.cpp b/lib/Lex/HeaderMap.cpp
index 26a179cf4e6d8083e1500fbc2545eaedad8a4aca..2b31ab9d04415a3b1bcb02afc6962dfd3ac83e66 100644 (file)
--- a/lib/Lex/HeaderMap.cpp
+++ b/lib/Lex/HeaderMap.cpp
@@ -19,6 +19,7 @@
 #include "llvm/Support/DataTypes.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/SwapByteOrder.h"
 #include <cstdio>
 #include <memory>
 using namespace clang;
@@ -82,6 +83,15 @@ bool HeaderMapImpl::checkHeader(const llvm::MemoryBuffer &File,
   if (Header->Reserved != 0)
     return false;
 
+  // Check the number of buckets.
+  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;
+
   // Okay, everything looks good.
   return true;
 }
@@ -191,10 +201,8 @@ StringRef HeaderMapImpl::lookupFilename(StringRef Filename,
   const HMapHeader &Hdr = getHeader();
   unsigned NumBuckets = getEndianAdjustedWord(Hdr.NumBuckets);
 
-  // If the number of buckets is not a power of two, the headermap is corrupt.
-  // Don't probe infinitely.
-  if (NumBuckets & (NumBuckets-1))
-    return StringRef();
+  // Don't probe infinitely.  This should be checked before constructing.
+  assert(!(NumBuckets & (NumBuckets - 1)) && "Expected power of 2");
 
   // Linearly probe the hash table.
   for (unsigned Bucket = HashHMapKey(Filename);; ++Bucket) {

diff --git a/unittests/Lex/HeaderMapTest.cpp b/unittests/Lex/HeaderMapTest.cpp
index 726e89c2751452dbec92b6227e610f65a500be4b..b8e1d6c46a05a89bc42818b48c550fa64d12e85f 100644 (file)
--- a/unittests/Lex/HeaderMapTest.cpp
+++ b/unittests/Lex/HeaderMapTest.cpp
@@ -91,4 +91,13 @@ TEST(HeaderMapTest, checkHeaderValidButEmpty) {
   ASSERT_TRUE(NeedsSwap);
 }
 
+TEST(HeaderMapTest, checkHeader3Buckets) {
+  MapFile<3, 1> File;
+  ASSERT_EQ(3 * sizeof(HMapBucket), sizeof(File.Buckets));
+
+  File.init();
+  bool NeedsSwap;
+  ASSERT_FALSE(HeaderMapImpl::checkHeader(*File.getBuffer(), NeedsSwap));
+}
+
 } // end namespace