1 files changed, 29 insertions, 4 deletions

diff --git a/lib/ProfileData/InstrProfReader.cpp b/lib/ProfileData/InstrProfReader.cpp
index 2ab0eb9449..7014f5e5cc 100644
--- a/lib/ProfileData/InstrProfReader.cpp
+++ b/lib/ProfileData/InstrProfReader.cpp
@@ -172,6 +172,29 @@ error_code RawInstrProfReader<IntPtrT>::readHeader() {
   return readHeader(*Header);
 }
 
+template <class IntPtrT>
+error_code RawInstrProfReader<IntPtrT>::readNextHeader(const char *CurrentPos) {
+  const char *End = DataBuffer->getBufferEnd();
+  // Skip zero padding between profiles.
+  while (CurrentPos != End && *CurrentPos == 0)
+    ++CurrentPos;
+  // If there's nothing left, we're done.
+  if (CurrentPos == End)
+    return instrprof_error::eof;
+  // If there isn't enough space for another header, this is probably just
+  // garbage at the end of the file.
+  if (CurrentPos + sizeof(RawHeader) > End)
+    return instrprof_error::malformed;
+  // The magic should have the same byte order as in the previous header.
+  uint64_t Magic = *reinterpret_cast<const uint64_t *>(CurrentPos);
+  if (Magic != swap(getRawMagic<IntPtrT>()))
+    return instrprof_error::bad_magic;
+
+  // There's another profile to read, so we need to process the header.
+  auto *Header = reinterpret_cast<const RawHeader *>(CurrentPos);
+  return readHeader(*Header);
+}
+
 static uint64_t getRawVersion() {
   return 1;
 }
@@ -190,16 +213,17 @@ error_code RawInstrProfReader<IntPtrT>::readHeader(const RawHeader &Header) {
   ptrdiff_t DataOffset = sizeof(RawHeader);
   ptrdiff_t CountersOffset = DataOffset + sizeof(ProfileData) * DataSize;
   ptrdiff_t NamesOffset = CountersOffset + sizeof(uint64_t) * CountersSize;
-  size_t FileSize = NamesOffset + sizeof(char) * NamesSize;
+  size_t ProfileSize = NamesOffset + sizeof(char) * NamesSize;
 
-  if (FileSize != DataBuffer->getBufferSize())
+  auto *Start = reinterpret_cast<const char *>(&Header);
+  if (Start + ProfileSize > DataBuffer->getBufferEnd())
     return error(instrprof_error::bad_header);
 
-  const char *Start = DataBuffer->getBufferStart();
   Data = reinterpret_cast<const ProfileData *>(Start + DataOffset);
   DataEnd = Data + DataSize;
   CountersStart = reinterpret_cast<const uint64_t *>(Start + CountersOffset);
   NamesStart = Start + NamesOffset;
+  ProfileEnd = Start + ProfileSize;
 
   return success();
 }
@@ -208,7 +232,8 @@ template <class IntPtrT>
 error_code
 RawInstrProfReader<IntPtrT>::readNextRecord(InstrProfRecord &Record) {
   if (Data == DataEnd)
-    return error(instrprof_error::eof);
+    if (error_code EC = readNextHeader(ProfileEnd))
+      return EC;
 
   // Get the raw data.
   StringRef RawName(getName(Data->NamePtr), swap(Data->NameSize));