diff --git a/include/llvm/ProfileData/InstrProf.h b/include/llvm/ProfileData/InstrProf.h index 778fe16c2e6..3a122666200 100644 --- a/include/llvm/ProfileData/InstrProf.h +++ b/include/llvm/ProfileData/InstrProf.h @@ -506,6 +506,9 @@ typedef struct ValueProfData { /// Return a pointer to \c ValueProfData instance ready to be streamed. static std::unique_ptr serializeFrom(const InstrProfRecord &Record); + /// Check the integrity of the record. Return the error code when + /// an error is detected, otherwise return instrprof_error::success. + instrprof_error checkIntegrity(); /// Return a pointer to \c ValueProfileData instance ready to be read. /// All data in the instance are properly byte swapped. The input /// data is assumed to be in little endian order. diff --git a/lib/ProfileData/InstrProf.cpp b/lib/ProfileData/InstrProf.cpp index 86cd76c9a20..54c37e30dba 100644 --- a/lib/ProfileData/InstrProf.cpp +++ b/lib/ProfileData/InstrProf.cpp @@ -442,6 +442,24 @@ static std::unique_ptr allocValueProfData(uint32_t TotalSize) { ValueProfData()); } +instrprof_error ValueProfData::checkIntegrity() { + if (NumValueKinds > IPVK_Last + 1) + return instrprof_error::malformed; + // Total size needs to be mulltiple of quadword size. + if (TotalSize % sizeof(uint64_t)) + return instrprof_error::malformed; + + ValueProfRecord *VR = getFirstValueProfRecord(this); + for (uint32_t K = 0; K < this->NumValueKinds; K++) { + if (VR->Kind > IPVK_Last) + return instrprof_error::malformed; + VR = getValueProfRecordNext(VR); + if ((char *)VR - (char *)this > (ptrdiff_t)TotalSize) + return instrprof_error::malformed; + } + return instrprof_error::success; +} + ErrorOr> ValueProfData::getValueProfData(const unsigned char *D, const unsigned char *const BufferEnd, @@ -452,31 +470,17 @@ ValueProfData::getValueProfData(const unsigned char *D, const unsigned char *Header = D; uint32_t TotalSize = swapToHostOrder(Header, Endianness); - uint32_t NumValueKinds = swapToHostOrder(Header, Endianness); - if (D + TotalSize > BufferEnd) return instrprof_error::too_large; - if (NumValueKinds > IPVK_Last + 1) - return instrprof_error::malformed; - // Total size needs to be mulltiple of quadword size. - if (TotalSize % sizeof(uint64_t)) - return instrprof_error::malformed; std::unique_ptr VPD = allocValueProfData(TotalSize); - memcpy(VPD.get(), D, TotalSize); // Byte swap. VPD->swapBytesToHost(Endianness); - // Data integrity check: - ValueProfRecord *VR = getFirstValueProfRecord(VPD.get()); - for (uint32_t K = 0; K < VPD->NumValueKinds; K++) { - if (VR->Kind > IPVK_Last) - return instrprof_error::malformed; - VR = getValueProfRecordNext(VR); - if ((char *)VR - (char *)VPD.get() > (ptrdiff_t)TotalSize) - return instrprof_error::malformed; - } + instrprof_error EC = VPD->checkIntegrity(); + if (EC != instrprof_error::success) + return EC; return std::move(VPD); }