r261448 - Lex: Check buckets on header map construction

[Duncan P. N. Exon Smith via cfe-commits]

Author: dexonsmith
Date: Sat Feb 20 15:00:58 2016
New Revision: 261448

URL: http://llvm.org/viewvc/llvm-project?rev=261448&view=rev
Log:
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.

Modified:
    cfe/trunk/lib/Lex/HeaderMap.cpp
    cfe/trunk/unittests/Lex/HeaderMapTest.cpp

Modified: cfe/trunk/lib/Lex/HeaderMap.cpp
URL: http://llvm.org/viewvc/llvm-project/cfe/trunk/lib/Lex/HeaderMap.cpp?rev=261448&r1=261447&r2=261448&view=diff
==============================================================================
--- cfe/trunk/lib/Lex/HeaderMap.cpp (original)
+++ cfe/trunk/lib/Lex/HeaderMap.cpp Sat Feb 20 15:00:58 2016
@@ -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 ll
   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(
   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) {

Modified: cfe/trunk/unittests/Lex/HeaderMapTest.cpp
URL: http://llvm.org/viewvc/llvm-project/cfe/trunk/unittests/Lex/HeaderMapTest.cpp?rev=261448&r1=261447&r2=261448&view=diff
==============================================================================
--- cfe/trunk/unittests/Lex/HeaderMapTest.cpp (original)
+++ cfe/trunk/unittests/Lex/HeaderMapTest.cpp Sat Feb 20 15:00:58 2016
@@ -91,4 +91,13 @@ TEST(HeaderMapTest, checkHeaderValidButE
   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