[llvm-profdata] Make profile reader behavior consistent when encountering malformed profiles

[huangjd]

Summary

Changes to the profile reader behavior (previously the behavior is bugged when encountering such scenario);

1. For extbinary format, when reading a second name table, the previous one is cleared (previously causes array out of bound)

2. Check if profile mixes MD5 and nonMD5 functions

[davidxl]

llvm/lib/ProfileData/SampleProfReader.cpp

745–746

why changing FixedLengthMD5 into a local variable?

1314

Should the code be guarded with #ifdef NDEBUG?

llvm/test/tools/llvm-profdata/sample-nametable.test

2

Are the test cases (profile data) ill-formed data or valid data? If valid, is it possible to check in the text format and convert into binary format in the testing itself?

5

Should this succeed instead? I assume without the fix it will crash?

8

Are there any expected strings?

[huangjd]

llvm/lib/ProfileData/SampleProfReader.cpp

745–746

The profile reader does not need to keep this info. It only matters while reading the section. In the future we plan to support having multiple name tables possibly with different flags set, so the reader should minimize keeping any states as a result of reading a section.

[huangjd]

llvm/test/tools/llvm-profdata/sample-nametable.test

2

They are all intentionally constructed malformed data and cannot be represented in text form, since the principle is not to trust user input, so the compiler should safeguard against such cases

5

The design decision (which I have previously asked for clarification and got no response https://discourse.llvm.org/t/llvm-profdata-what-behavior-was-intended-if-multiple-name-tables-are-present-in-an-extensible-sample-profile/69145) here is to make the reader clear a previously read name table when reading a new one. In this case the second name table only have one entry, so the function sample, being read after reading the second name table, has a string index of 1, which is out of bounds, so it should be rejected. Before the patch, it appends the second name table on the first one, while changing the fixed lenght MD5 base address to the second table, and index (out of bounds) into the sample's section, reading garbage data instead of crashing

8

This one is similar to the previous test, also out of bounds when indexing into the second table. In the original code it crashes because the fixed length MD5 base address is not set correctly

[huangjd]

Add ifndef guard to debug block

[wenlei]

llvm/lib/ProfileData/SampleProfReader.cpp

743–744

Why is assertion removed? Is it no longer the case?

[wenlei]

llvm/lib/ProfileData/SampleProfReader.cpp

557–558

What is the bogus End used for? Defend against call site that didn't check error and tries to continue reading?

[huangjd]

llvm/lib/ProfileData/SampleProfReader.cpp

557–558

Even though the section header of the name table should appear before the section header of the function offset table or profiles, the actual section data can be placed in arbitrary order. If the name table section is placed after the profile data section, when a string is read, End still points to the end of the profile data section which is before Data here, and readUnencodedNumber fails even if the name table index is valid.

Since we already checked a fixed length MD5 name table contains actually enough entries when reading the name table, we don't need to check for bounds here again, so setting End to max int so that readUnencodedNumber never fails bounds check

743–744

In release mode assert is ignored, and test case 2 causes a bug with a flag of FixedLengthMD5 & ~UseMD5 . I fixed the logic so that even in this case the logic is correct, so the check is no longer needed.

[wenlei]

llvm/lib/ProfileData/SampleProfReader.cpp

557–558

I think it's better to keep the invariant that End >= Data everywhere instead trying to get around it. To me, the problem here is the caller of readUnencodedNumber only set Data, but not End which leave the two in an inconsistent state. Can we also set/save/restore End there, and avoid the trick here to bypass the check?

const uint8_t *SavedData = Data;
Data = MD5NameMemStart + ((*Idx) * sizeof(uint64_t));
auto FID = readUnencodedNumber<uint64_t>();

743–744

Perhaps we should fix the logic/testcase that produced the situation FixedLengthMD5 is true while UseMD5 is false, instead of removing the assert and tolerate such cases.

In this function below, it also assumed IsMD5 needs to be true if FixedLengthMD5. I think it's reasonable to keep the assertion.

SampleProfileReaderExtBinaryBase::readNameTableSec(bool IsMD5,
                                                   bool FixedLengthMD5)
  if (!IsMD5)
    return SampleProfileReaderBinary::readNameTable();
...
  if (FixedLengthMD5)

llvm/include/llvm/ProfileData/SampleProfReader.h

protected:
   std::error_code readSecHdrTableEntry(uint64_t Idx);
   std::error_code readSecHdrTable();
 
   std::error_code readFuncMetadata(bool ProfileHasAttribute);
   std::error_code readFuncMetadata(bool ProfileHasAttribute,
                                    FunctionSamples *FProfile);
   std::error_code readFuncOffsetTable();
   std::error_code readFuncProfiles();
-  std::error_code readMD5NameTable();
-  std::error_code readNameTableSec(bool IsMD5);
+  std::error_code readNameTableSec(bool IsMD5, bool FixedLengthMD5);
   std::error_code readCSNameTableSec();
   std::error_code readProfileSymbolList();
 
   std::error_code readHeader() override;
   std::error_code verifySPMagic(uint64_t Magic) override = 0;
   virtual std::error_code readOneSection(const uint8_t *Start, uint64_t Size,
                                          const SecHdrTableEntry &Entry);
   // placeholder for subclasses to dispatch their own section readers.
protected:
 
   /// Function offset mapping ordered by contexts.
   std::unique_ptr<std::vector<std::pair<SampleContext, uint64_t>>>
       OrderedFuncOffsets;
 
   /// The set containing the functions to use when compiling a module.
   DenseSet<StringRef> FuncsToUse;
 
-  /// Use fixed length MD5 instead of ULEB128 encoding so NameTable doesn't
-  /// need to be read in up front and can be directly accessed using index.
-  bool FixedLengthMD5 = false;
   /// The starting address of NameTable containing fixed length MD5.
   const uint8_t *MD5NameMemStart = nullptr;
 
   /// If MD5 is used in NameTable section, the section saves uint64_t data.
   /// The uint64_t data has to be converted to a string and then the string
   /// will be used to initialize StringRef in NameTable.
   /// Note NameTable contains StringRef so it needs another buffer to own
   /// the string data. MD5StringBuf serves as the string buffer that is

llvm/lib/ProfileData/SampleProfReader.cpp

 ErrorOr<SampleContext> SampleProfileReaderBinary::readSampleContextFromTable() {
   auto FName(readStringFromTable());
   if (std::error_code EC = FName.getError())
     return EC;
   return SampleContext(*FName);
 }
 
 ErrorOr<StringRef> SampleProfileReaderExtBinaryBase::readStringFromTable() {
-  if (!FixedLengthMD5)
-    return SampleProfileReaderBinary::readStringFromTable();
-
   // read NameTable index.
   auto Idx = readStringIndex(NameTable);
   if (std::error_code EC = Idx.getError())
     return EC;
 
   // Check whether the name to be accessed has been accessed before,
   // if not, read it from memory directly.
   StringRef &SR = NameTable[*Idx];
-  if (SR.empty()) {
+  if (!SR.data()) {
+    assert(MD5NameMemStart);
     const uint8_t *SavedData = Data;
+    const uint8_t *SavedEnd = End;
     Data = MD5NameMemStart + ((*Idx) * sizeof(uint64_t));
+    End = reinterpret_cast<const uint8_t *>(
+        std::numeric_limits<uintptr_t>::max());
     auto FID = readUnencodedNumber<uint64_t>();
     if (std::error_code EC = FID.getError())
       return EC;
     // Save the string converted from uint64_t in MD5StringBuf. All the
     // references to the name are all StringRefs refering to the string
     // in MD5StringBuf.
     MD5StringBuf->push_back(std::to_string(*FID));
     SR = MD5StringBuf->back();
     Data = SavedData;
+    End = SavedEnd;
   }
   return SR;
 }
 
 ErrorOr<StringRef> SampleProfileReaderCompactBinary::readStringFromTable() {
   auto Idx = readStringIndex(NameTable);
   if (std::error_code EC = Idx.getError())
     return EC;
case SecProfSummary:
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagFullContext))
       FunctionSamples::ProfileIsCS = ProfileIsCS = true;
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagIsPreInlined))
       FunctionSamples::ProfileIsPreInlined = ProfileIsPreInlined = true;
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagFSDiscriminator))
       FunctionSamples::ProfileIsFS = ProfileIsFS = true;
     break;
   case SecNameTable: {
-    FixedLengthMD5 =
+    /// Use fixed length MD5 instead of ULEB128 encoding so NameTable doesn't
+    /// need to be read in up front and can be directly accessed using index.
+    bool FixedLengthMD5 =
         hasSecFlag(Entry, SecNameTableFlags::SecFlagFixedLengthMD5);
-    bool UseMD5 = hasSecFlag(Entry, SecNameTableFlags::SecFlagMD5Name);
-    assert((!FixedLengthMD5 || UseMD5) &&
-           "If FixedLengthMD5 is true, UseMD5 has to be true");
+    bool UseMD5 =
+        FixedLengthMD5 || hasSecFlag(Entry, SecNameTableFlags::SecFlagMD5Name);
     FunctionSamples::HasUniqSuffix =
         hasSecFlag(Entry, SecNameTableFlags::SecFlagUniqSuffix);
-    if (std::error_code EC = readNameTableSec(UseMD5))
+    if (std::error_code EC = readNameTableSec(UseMD5, FixedLengthMD5))
       return EC;
     break;
   }
   case SecCSNameTable: {
     if (std::error_code EC = readCSNameTableSec())
       return EC;
     break;
   }
if (Magic == SPMagic(SPF_Compact_Binary))
     return sampleprof_error::success;
   return sampleprof_error::bad_magic;
 }
 
 std::error_code SampleProfileReaderBinary::readNameTable() {
   auto Size = readNumber<size_t>();
   if (std::error_code EC = Size.getError())
     return EC;
-  NameTable.reserve(*Size + NameTable.size());
+  NameTable.clear();
+  NameTable.reserve(*Size);
   for (size_t I = 0; I < *Size; ++I) {
     auto Name(readString());
     if (std::error_code EC = Name.getError())
       return EC;
     NameTable.push_back(*Name);
   }
 
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderExtBinaryBase::readMD5NameTable() {
+std::error_code
+SampleProfileReaderExtBinaryBase::readNameTableSec(bool IsMD5,
+                                                   bool FixedLengthMD5) {
+  if (!IsMD5)
+    return SampleProfileReaderBinary::readNameTable();
+
   auto Size = readNumber<uint64_t>();
   if (std::error_code EC = Size.getError())
     return EC;
   MD5StringBuf = std::make_unique<std::vector<std::string>>();
   MD5StringBuf->reserve(*Size);
   if (FixedLengthMD5) {
+    assert(Data + (*Size) * sizeof(uint64_t) == End &&
+           "Fixed length MD5 name table does not contain specified number of "
+           "entries");
+    if (Data + (*Size) * sizeof(uint64_t) > End)
+      return sampleprof_error::truncated;
+
     // Preallocate and initialize NameTable so we can check whether a name
     // index has been read before by checking whether the element in the
     // NameTable is empty, meanwhile readStringIndex can do the boundary
     // check using the size of NameTable.
-    NameTable.resize(*Size + NameTable.size());
+    NameTable.resize(*Size);
 
     MD5NameMemStart = Data;
     Data = Data + (*Size) * sizeof(uint64_t);
     return sampleprof_error::success;
   }
+  NameTable.clear();
   NameTable.reserve(*Size);
   for (uint64_t I = 0; I < *Size; ++I) {
     auto FID = readNumber<uint64_t>();
     if (std::error_code EC = FID.getError())
       return EC;
     MD5StringBuf->push_back(std::to_string(*FID));
     // NameTable is a vector of StringRef. Here it is pushing back a
     // StringRef initialized with the last string in MD5stringBuf.
     NameTable.push_back(MD5StringBuf->back());
   }
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderExtBinaryBase::readNameTableSec(bool IsMD5) {
-  if (IsMD5)
-    return readMD5NameTable();
-  return SampleProfileReaderBinary::readNameTable();
-}
-
 // Read in the CS name table section, which basically contains a list of context
 // vectors. Each element of a context vector, aka a frame, refers to the
 // underlying raw function names that are stored in the name table, as well as
 // a callsite identifier that only makes sense for non-leaf frames.
 std::error_code SampleProfileReaderExtBinaryBase::readCSNameTableSec() {
   auto Size = readNumber<size_t>();
   if (std::error_code EC = Size.getError())
     return EC;
SampleProfileReaderExtBinaryBase::readFuncMetadata(bool ProfileHasAttribute) {
   assert(Data == End && "More data is read than expected");
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderCompactBinary::readNameTable() {
   auto Size = readNumber<uint64_t>();
   if (std::error_code EC = Size.getError())
     return EC;
+  NameTable.clear();
   NameTable.reserve(*Size);
   for (uint64_t I = 0; I < *Size; ++I) {
     auto FID = readNumber<uint64_t>();
     if (std::error_code EC = FID.getError())
       return EC;
     NameTable.push_back(std::to_string(*FID));
   }
   return sampleprof_error::success;
std::error_code SampleProfileReaderExtBinaryBase::readHeader() {
   End = BufStart + Buffer->getBufferSize();
 
   if (std::error_code EC = readMagicIdent())
     return EC;
 
   if (std::error_code EC = readSecHdrTable())
     return EC;
 
+#ifndef NDEBUG
+  bool HasMD5 = false;
+  bool HasNonMD5 = false;
+  for (auto &Entry : SecHdrTable) {
+    if (Entry.Size && Entry.Type == SecNameTable) {
+      bool IsMD5 = hasSecFlag(Entry, SecNameTableFlags::SecFlagMD5Name) ||
+                   hasSecFlag(Entry, SecNameTableFlags::SecFlagFixedLengthMD5);
+      if (IsMD5) {
+        HasMD5 = true;
+      } else
+        HasNonMD5 = true;
+    }
+  }
+  if (HasMD5 != HasNonMD5) {
+    assert(!(HasMD5 && HasNonMD5) &&
+           "Profile contains both MD5 and non-MD5 function names, non-MD5 "
+           "function names will be dropped");
+    assert((HasMD5 || HasNonMD5) && "Profile contains no name table");
+  }
+#endif
+
   return sampleprof_error::success;
 }
 
 uint64_t SampleProfileReaderExtBinaryBase::getSectionSize(SecType Type) {
   uint64_t Size = 0;
   for (auto &Entry : SecHdrTable) {
     if (Entry.Type == Type)
       Size += Entry.Size;

llvm/test/tools/llvm-profdata/sample-nametable.test

+Test several edge cases with unusual name table data in ExtBinary format.
+
+1- Multiple fixed-length MD5 name tables. Reading a new table should clear the content from old table, and a valid name index for the old name table should become invalid if the new name table has fewer entries.
+RUN: not llvm-profdata show --sample %p/Inputs/sample-multiple-nametables.profdata
+
+2- Multiple name tables, the first one has an empty string, the second one tricks the reader into expecting fixed-length MD5 values. Reader should not attempt "lazy loading" of the MD5 string in this case.
+RUN: not llvm-profdata show --sample %p/Inputs/sample-nametable-empty-string.profdata
+
+3- The data of the name table is placed after the data of the profiles. The reader should handle it correctly.
+RUN: llvm-profdata merge --sample --text %p/Inputs/sample-nametable-after-samples.profdata | FileCheck %s
+CHECK: 18446744073709551615:2:9
+