1 files changed, 97 insertions, 20 deletions

diff --git a/llvm/lib/XRay/FDRRecordProducer.cpp b/llvm/lib/XRay/FDRRecordProducer.cpp
index 122578010c4..25b3ee8af21 100644
--- a/llvm/lib/XRay/FDRRecordProducer.cpp
+++ b/llvm/lib/XRay/FDRRecordProducer.cpp
@@ -9,29 +9,32 @@
 #include "llvm/XRay/FDRRecordProducer.h"
 #include "llvm/Support/DataExtractor.h"
 
+#include <cstdint>
+
 namespace llvm {
 namespace xray {
 
 namespace {
 
+// Keep this in sync with the values written in the XRay FDR mode runtime in
+// compiler-rt.
+enum MetadataRecordKinds : uint8_t {
+  NewBufferKind,
+  EndOfBufferKind,
+  NewCPUIdKind,
+  TSCWrapKind,
+  WalltimeMarkerKind,
+  CustomEventMarkerKind,
+  CallArgumentKind,
+  BufferExtentsKind,
+  TypedEventMarkerKind,
+  PidKind,
+  // This is an end marker, used to identify the upper bound for this enum.
+  EnumEndMarker,
+};
+
 Expected<std::unique_ptr<Record>>
 metadataRecordType(const XRayFileHeader &Header, uint8_t T) {
-  // Keep this in sync with the values written in the XRay FDR mode runtime in
-  // compiler-rt.
-  enum MetadataRecordKinds : uint8_t {
-    NewBufferKind,
-    EndOfBufferKind,
-    NewCPUIdKind,
-    TSCWrapKind,
-    WalltimeMarkerKind,
-    CustomEventMarkerKind,
-    CallArgumentKind,
-    BufferExtentsKind,
-    TypedEventMarkerKind,
-    PidKind,
-    // This is an end marker, used to identify the upper bound for this enum.
-    EnumEndMarker,
-  };
 
   if (T >= static_cast<uint8_t>(MetadataRecordKinds::EnumEndMarker))
     return createStringError(std::make_error_code(std::errc::invalid_argument),
@@ -70,9 +73,70 @@ metadataRecordType(const XRayFileHeader &Header, uint8_t T) {
   llvm_unreachable("Unhandled MetadataRecordKinds enum value");
 }
 
+constexpr bool isMetadataIntroducer(uint8_t FirstByte) {
+  return FirstByte & 0x01u;
+}
+
 } // namespace
 
+Expected<std::unique_ptr<Record>>
+FileBasedRecordProducer::findNextBufferExtent() {
+  // We seek one byte at a time until we find a suitable buffer extents metadata
+  // record introducer.
+  std::unique_ptr<Record> R;
+  while (!R) {
+    auto PreReadOffset = OffsetPtr;
+    uint8_t FirstByte = E.getU8(&OffsetPtr);
+    if (OffsetPtr == PreReadOffset)
+      return createStringError(
+          std::make_error_code(std::errc::executable_format_error),
+          "Failed reading one byte from offset %d.", OffsetPtr);
+
+    if (isMetadataIntroducer(FirstByte)) {
+      auto LoadedType = FirstByte >> 1;
+      if (LoadedType == MetadataRecordKinds::BufferExtentsKind) {
+        auto MetadataRecordOrErr = metadataRecordType(Header, LoadedType);
+        if (!MetadataRecordOrErr)
+          return MetadataRecordOrErr.takeError();
+
+        R = std::move(MetadataRecordOrErr.get());
+        RecordInitializer RI(E, OffsetPtr);
+        if (auto Err = R->apply(RI))
+          return std::move(Err);
+        return std::move(R);
+      }
+    }
+  }
+  llvm_unreachable("Must always terminate with either an error or a record.");
+}
+
 Expected<std::unique_ptr<Record>> FileBasedRecordProducer::produce() {
+  // First, we set up our result record.
+  std::unique_ptr<Record> R;
+
+  // Before we do any further reading, we should check whether we're at the end
+  // of the current buffer we're been consuming. In FDR logs version >= 3, we
+  // rely on the buffer extents record to determine how many bytes we should be
+  // considering as valid records.
+  if (Header.Version >= 3 && CurrentBufferBytes == 0) {
+    // Find the next buffer extents record.
+    auto BufferExtentsOrError = findNextBufferExtent();
+    if (!BufferExtentsOrError)
+      return joinErrors(
+          BufferExtentsOrError.takeError(),
+          createStringError(
+              std::make_error_code(std::errc::executable_format_error),
+              "Failed to find the next BufferExtents record."));
+
+    R = std::move(BufferExtentsOrError.get());
+    assert(R != nullptr);
+    assert(isa<BufferExtents>(R.get()));
+    auto BE = dyn_cast<BufferExtents>(R.get());
+    CurrentBufferBytes = BE->size();
+    return std::move(R);
+  }
+
+  //
   // At the top level, we read one byte to determine the type of the record to
   // create. This byte will comprise of the following bits:
   //
@@ -90,11 +154,8 @@ Expected<std::unique_ptr<Record>> FileBasedRecordProducer::produce() {
         std::make_error_code(std::errc::executable_format_error),
         "Failed reading one byte from offset %d.", OffsetPtr);
 
-  // Set up our result record.
-  std::unique_ptr<Record> R;
-
   // For metadata records, handle especially here.
-  if (FirstByte & 0x01) {
+  if (isMetadataIntroducer(FirstByte)) {
     auto LoadedType = FirstByte >> 1;
     auto MetadataRecordOrErr = metadataRecordType(Header, LoadedType);
     if (!MetadataRecordOrErr)
@@ -113,6 +174,22 @@ Expected<std::unique_ptr<Record>> FileBasedRecordProducer::produce() {
   if (auto Err = R->apply(RI))
     return std::move(Err);
 
+  // If we encountered a BufferExtents record, we should record the remaining
+  // bytes for the current buffer, to determine when we should start ignoring
+  // potentially malformed data and looking for buffer extents records.
+  if (auto BE = dyn_cast<BufferExtents>(R.get())) {
+    CurrentBufferBytes = BE->size();
+  } else if (Header.Version >= 3) {
+    if (OffsetPtr - PreReadOffset > CurrentBufferBytes)
+      return createStringError(
+          std::make_error_code(std::errc::executable_format_error),
+          "Buffer over-read at offset %d (over-read by %d bytes); Record Type "
+          "= %s.",
+          OffsetPtr, (OffsetPtr - PreReadOffset) - CurrentBufferBytes,
+          Record::kindToString(R->getRecordType()).data());
+
+    CurrentBufferBytes -= OffsetPtr - PreReadOffset;
+  }
   assert(R != nullptr);
   return std::move(R);
 }