diff --git a/lib/Lex/HeaderMap.cpp b/lib/Lex/HeaderMap.cpp index 2b31ab9d04415a3b1bcb02afc6962dfd3ac83e66..220e70d8d589e95d571c3ee986fa5e604b8c1a82 100644 --- a/lib/Lex/HeaderMap.cpp +++ b/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/unittests/Lex/HeaderMapTest.cpp b/unittests/Lex/HeaderMapTest.cpp index b8e1d6c46a05a89bc42818b48c550fa64d12e85f..33befd8112f49db25108a8624e9f5e5fe9dbc584 100644 --- a/unittests/Lex/HeaderMapTest.cpp +++ b/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