1 files changed, 63 insertions, 30 deletions

diff --git a/llvm/lib/ProfileData/Coverage/CoverageMappingReader.cpp b/llvm/lib/ProfileData/Coverage/CoverageMappingReader.cpp
index 0573732..bee9c8d 100644
--- a/llvm/lib/ProfileData/Coverage/CoverageMappingReader.cpp
+++ b/llvm/lib/ProfileData/Coverage/CoverageMappingReader.cpp
@@ -822,8 +822,6 @@ static Error readCoverageMappingData(
   return Error::success();
 }
 
-static const char *TestingFormatMagic = "llvmcovmtestdata";
-
 Expected<std::unique_ptr<BinaryCoverageReader>>
 BinaryCoverageReader::createCoverageReaderFromBuffer(
     StringRef Coverage, FuncRecordsStorage &&FuncRecords,
@@ -865,7 +863,16 @@ loadTestingFormat(StringRef Data, StringRef CompilationDir) {
   uint8_t BytesInAddress = 8;
   support::endianness Endian = support::endianness::little;
 
-  Data = Data.substr(StringRef(TestingFormatMagic).size());
+  // Read the magic and version.
+  Data = Data.substr(sizeof(TestingFormatMagic));
+  if (Data.size() < sizeof(uint64_t))
+    return make_error<CoverageMapError>(coveragemap_error::malformed);
+  auto TestingVersion =
+      support::endian::byte_swap<uint64_t, support::endianness::little>(
+          *reinterpret_cast<const uint64_t *>(Data.data()));
+  Data = Data.substr(sizeof(uint64_t));
+
+  // Read the ProfileNames data.
   if (Data.empty())
     return make_error<CoverageMapError>(coveragemap_error::truncated);
   unsigned N = 0;
@@ -886,8 +893,23 @@ loadTestingFormat(StringRef Data, StringRef CompilationDir) {
   if (Error E = ProfileNames.create(Data.substr(0, ProfileNamesSize), Address))
     return std::move(E);
   Data = Data.substr(ProfileNamesSize);
+
+  // In Version2, the size of CoverageMapping is stored directly.
+  uint64_t CoverageMappingSize;
+  if (TestingVersion == uint64_t(TestingFormatVersion::Version2)) {
+    N = 0;
+    CoverageMappingSize = decodeULEB128(Data.bytes_begin(), &N);
+    if (N > Data.size())
+      return make_error<CoverageMapError>(coveragemap_error::malformed);
+    Data = Data.substr(N);
+    if (CoverageMappingSize < sizeof(CovMapHeader))
+      return make_error<CoverageMapError>(coveragemap_error::malformed);
+  } else if (TestingVersion != uint64_t(TestingFormatVersion::Version1)) {
+    return make_error<CoverageMapError>(coveragemap_error::unsupported_version);
+  }
+
   // Skip the padding bytes because coverage map data has an alignment of 8.
-  size_t Pad = offsetToAlignedAddr(Data.data(), Align(8));
+  auto Pad = offsetToAlignedAddr(Data.data(), Align(8));
   if (Data.size() < Pad)
     return make_error<CoverageMapError>(coveragemap_error::malformed);
   Data = Data.substr(Pad);
@@ -895,32 +917,38 @@ loadTestingFormat(StringRef Data, StringRef CompilationDir) {
     return make_error<CoverageMapError>(coveragemap_error::malformed);
   auto const *CovHeader = reinterpret_cast<const CovMapHeader *>(
       Data.substr(0, sizeof(CovMapHeader)).data());
-  CovMapVersion Version =
-      (CovMapVersion)CovHeader->getVersion<support::endianness::little>();
-  StringRef CoverageMapping;
-  BinaryCoverageReader::FuncRecordsStorage CoverageRecords;
+  auto Version =
+      CovMapVersion(CovHeader->getVersion<support::endianness::little>());
+
+  // In Version1, the size of CoverageMapping is calculated.
+  if (TestingVersion == uint64_t(TestingFormatVersion::Version1)) {
+    if (Version < CovMapVersion::Version4) {
+      CoverageMappingSize = Data.size();
+    } else {
+      auto FilenamesSize =
+          CovHeader->getFilenamesSize<support::endianness::little>();
+      CoverageMappingSize = sizeof(CovMapHeader) + FilenamesSize;
+    }
+  }
+
+  auto CoverageMapping = Data.substr(0, CoverageMappingSize);
+  Data = Data.substr(CoverageMappingSize);
+
+  // Read the CoverageRecords data.
   if (Version < CovMapVersion::Version4) {
-    CoverageMapping = Data;
-    if (CoverageMapping.empty())
-      return make_error<CoverageMapError>(coveragemap_error::truncated);
-    CoverageRecords = MemoryBuffer::getMemBuffer("");
+    if (!Data.empty())
+      return make_error<CoverageMapError>(coveragemap_error::malformed);
   } else {
-    uint32_t FilenamesSize =
-        CovHeader->getFilenamesSize<support::endianness::little>();
-    uint32_t CoverageMappingSize = sizeof(CovMapHeader) + FilenamesSize;
-    CoverageMapping = Data.substr(0, CoverageMappingSize);
-    if (CoverageMapping.empty())
-      return make_error<CoverageMapError>(coveragemap_error::truncated);
-    Data = Data.substr(CoverageMappingSize);
     // Skip the padding bytes because coverage records data has an alignment
     // of 8.
     Pad = offsetToAlignedAddr(Data.data(), Align(8));
     if (Data.size() < Pad)
       return make_error<CoverageMapError>(coveragemap_error::malformed);
-    CoverageRecords = MemoryBuffer::getMemBuffer(Data.substr(Pad));
-    if (CoverageRecords->getBufferSize() == 0)
-      return make_error<CoverageMapError>(coveragemap_error::truncated);
+    Data = Data.substr(Pad);
   }
+  BinaryCoverageReader::FuncRecordsStorage CoverageRecords =
+      MemoryBuffer::getMemBuffer(Data);
+
   return BinaryCoverageReader::createCoverageReaderFromBuffer(
       CoverageMapping, std::move(CoverageRecords), std::move(ProfileNames),
       BytesInAddress, Endian, CompilationDir);
@@ -1081,14 +1109,19 @@ BinaryCoverageReader::create(
     StringRef CompilationDir, SmallVectorImpl<object::BuildIDRef> *BinaryIDs) {
   std::vector<std::unique_ptr<BinaryCoverageReader>> Readers;
 
-  if (ObjectBuffer.getBuffer().startswith(TestingFormatMagic)) {
-    // This is a special format used for testing.
-    auto ReaderOrErr =
-        loadTestingFormat(ObjectBuffer.getBuffer(), CompilationDir);
-    if (!ReaderOrErr)
-      return ReaderOrErr.takeError();
-    Readers.push_back(std::move(ReaderOrErr.get()));
-    return std::move(Readers);
+  if (ObjectBuffer.getBuffer().size() > sizeof(TestingFormatMagic)) {
+    uint64_t Magic =
+        support::endian::byte_swap<uint64_t, support::endianness::little>(
+            *reinterpret_cast<const uint64_t *>(ObjectBuffer.getBufferStart()));
+    if (Magic == TestingFormatMagic) {
+      // This is a special format used for testing.
+      auto ReaderOrErr =
+          loadTestingFormat(ObjectBuffer.getBuffer(), CompilationDir);
+      if (!ReaderOrErr)
+        return ReaderOrErr.takeError();
+      Readers.push_back(std::move(ReaderOrErr.get()));
+      return std::move(Readers);
+    }
   }
 
   auto BinOrErr = createBinary(ObjectBuffer);