Lex: Return "" when HeaderMap::lookupFilename fails

Change getString() to return Optional<StringRef>, and change
lookupFilename() to return an empty string if either one of the prefix
and suffix can't be found.

This is a more robust follow-up to r261461, but it's still not entirely
satisfactory.  Ideally we'd report that the header map is corrupt;
perhaps something for a follow-up.

llvm-svn: 261596

clang/include/clang/Lex/HeaderMap.h

@@ -15,6 +15,7 @@
 #define LLVM_CLANG_LEX_HEADERMAP_H
 
 #include "clang/Basic/LLVM.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include <memory>
@@ -53,7 +54,10 @@ private:
   unsigned getEndianAdjustedWord(unsigned X) const;
   const HMapHeader &getHeader() const;
   HMapBucket getBucket(unsigned BucketNo) const;
-  StringRef getString(unsigned StrTabIdx) const;
+
+  /// Look up the specified string in the string table.  If the string index is
+  /// not valid, return None.
+  Optional<StringRef> getString(unsigned StrTabIdx) const;
 };
 
 /// This class represents an Apple concept known as a 'header map'.  To the

clang/lib/Lex/HeaderMap.cpp

@@ -16,6 +16,7 @@
 #include "clang/Basic/CharInfo.h"
 #include "clang/Basic/FileManager.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/Support/Compiler.h"
 #include "llvm/Support/DataTypes.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/MemoryBuffer.h"
@@ -144,15 +145,13 @@ HMapBucket HeaderMapImpl::getBucket(unsigned BucketNo) const {
   return Result;
 }
 
-/// getString - Look up the specified string in the string table.  If the string
-/// index is not valid, it returns an empty string.
-StringRef HeaderMapImpl::getString(unsigned StrTabIdx) const {
+Optional<StringRef> HeaderMapImpl::getString(unsigned StrTabIdx) const {
   // Add the start of the string table to the idx.
   StrTabIdx += getEndianAdjustedWord(getHeader().StringsOffset);
 
   // Check for invalid index.
   if (StrTabIdx >= FileBuffer->getBufferSize())
-    return "";
+    return None;
 
   const char *Data = FileBuffer->getBufferStart() + StrTabIdx;
   unsigned MaxLen = FileBuffer->getBufferSize() - StrTabIdx;
@@ -160,7 +159,7 @@ StringRef HeaderMapImpl::getString(unsigned StrTabIdx) const {
 
   // Check whether the buffer is null-terminated.
   if (Len == MaxLen && Data[Len - 1])
-    return "";
+    return None;
 
   return StringRef(Data, Len);
 }
@@ -177,13 +176,19 @@ LLVM_DUMP_METHOD void HeaderMapImpl::dump() const {
   llvm::dbgs() << "Header Map " << getFileName() << ":\n  " << NumBuckets
                << ", " << getEndianAdjustedWord(Hdr.NumEntries) << "\n";
 
+  auto getStringOrInvalid = [this](unsigned Id) -> StringRef {
+    if (Optional<StringRef> S = getString(Id))
+      return *S;
+    return "<invalid>";
+  };
+
   for (unsigned i = 0; i != NumBuckets; ++i) {
     HMapBucket B = getBucket(i);
     if (B.Key == HMAP_EmptyBucketKey) continue;
 
-    StringRef Key    = getString(B.Key);
-    StringRef Prefix = getString(B.Prefix);
-    StringRef Suffix = getString(B.Suffix);
+    StringRef Key = getStringOrInvalid(B.Key);
+    StringRef Prefix = getStringOrInvalid(B.Prefix);
+    StringRef Suffix = getStringOrInvalid(B.Suffix);
     llvm::dbgs() << "  " << i << ". " << Key << " -> '" << Prefix << "' '"
                  << Suffix << "'\n";
   }
@@ -216,16 +221,22 @@ StringRef HeaderMapImpl::lookupFilename(StringRef Filename,
     if (B.Key == HMAP_EmptyBucketKey) return StringRef(); // Hash miss.
 
     // See if the key matches.  If not, probe on.
-    if (!Filename.equals_lower(getString(B.Key)))
+    Optional<StringRef> Key = getString(B.Key);
+    if (LLVM_UNLIKELY(!Key))
+      continue;
+    if (!Filename.equals_lower(*Key))
       continue;
 
     // If so, we have a match in the hash table.  Construct the destination
     // path.
-    StringRef Prefix = getString(B.Prefix);
-    StringRef Suffix = getString(B.Suffix);
+    Optional<StringRef> Prefix = getString(B.Prefix);
+    Optional<StringRef> Suffix = getString(B.Suffix);
+
     DestPath.clear();
-    DestPath.append(Prefix.begin(), Prefix.end());
-    DestPath.append(Suffix.begin(), Suffix.end());
+    if (LLVM_LIKELY(Prefix && Suffix)) {
+      DestPath.append(Prefix->begin(), Prefix->end());
+      DestPath.append(Suffix->begin(), Suffix->end());
+    }
     return StringRef(DestPath.begin(), DestPath.size());
   }
 }

clang/unittests/Lex/HeaderMapTest.cpp

@@ -185,7 +185,7 @@ template <class FileTy, class PaddingTy> struct PaddedFile {
   PaddingTy Padding;
 };
 
-TEST(HeaderMapTest, lookupFilenameTruncated) {
+TEST(HeaderMapTest, lookupFilenameTruncatedSuffix) {
   typedef MapFile<2, 64 - sizeof(HMapHeader) - 2 * sizeof(HMapBucket)> FileTy;
   static_assert(std::is_standard_layout<FileTy>::value,
                 "Expected standard layout");
@@ -215,10 +215,44 @@ TEST(HeaderMapTest, lookupFilenameTruncated) {
   HeaderMapImpl Map(File.getBuffer(), NeedsSwap);
 
   // The string for "c" runs to the end of File.  Check that the suffix
-  // ("cxxxx...") is ignored.  Another option would be to return an empty
-  // filename altogether.
+  // ("cxxxx...") is detected as truncated, and an empty string is returned.
   SmallString<24> DestPath;
-  ASSERT_EQ("b", Map.lookupFilename("a", DestPath));
+  ASSERT_EQ("", Map.lookupFilename("a", DestPath));
+}
+
+TEST(HeaderMapTest, lookupFilenameTruncatedPrefix) {
+  typedef MapFile<2, 64 - sizeof(HMapHeader) - 2 * sizeof(HMapBucket)> FileTy;
+  static_assert(std::is_standard_layout<FileTy>::value,
+                "Expected standard layout");
+  static_assert(sizeof(FileTy) == 64, "check the math");
+  PaddedFile<FileTy, uint64_t> P;
+  auto &File = P.File;
+  auto &Padding = P.Padding;
+  File.init();
+
+  FileMaker<FileTy> Maker(File);
+  auto a = Maker.addString("a");
+  auto c = Maker.addString("c");
+  auto b = Maker.addString("b"); // Store the prefix last.
+  Maker.addBucket(getHash("a"), a, b, c);
+
+  // Add 'x' characters to cause an overflow into Padding.
+  ASSERT_EQ('b', File.Bytes[5]);
+  for (unsigned I = 6; I < sizeof(File.Bytes); ++I) {
+    ASSERT_EQ(0, File.Bytes[I]);
+    File.Bytes[I] = 'x';
+  }
+  Padding = 0xffffffff; // Padding won't stop it either.
+
+  bool NeedsSwap;
+  ASSERT_TRUE(HeaderMapImpl::checkHeader(*File.getBuffer(), NeedsSwap));
+  ASSERT_FALSE(NeedsSwap);
+  HeaderMapImpl Map(File.getBuffer(), NeedsSwap);
+
+  // The string for "b" runs to the end of File.  Check that the prefix
+  // ("bxxxx...") is detected as truncated, and an empty string is returned.
+  SmallString<24> DestPath;
+  ASSERT_EQ("", Map.lookupFilename("a", DestPath));
 }
 
 } // end namespace