1 files changed, 24 insertions, 14 deletions

diff --git a/include/llvm/DebugInfo/CodeView/CVRecord.h b/include/llvm/DebugInfo/CodeView/CVRecord.h
index 11ca9ff108de..784c47e3bf5d 100644
--- a/include/llvm/DebugInfo/CodeView/CVRecord.h
+++ b/include/llvm/DebugInfo/CodeView/CVRecord.h
@@ -1,9 +1,8 @@
-//===- RecordIterator.h -----------------------------------------*- C++ -*-===//
+//===- CVRecord.h -----------------------------------------------*- C++ -*-===//
 //
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
@@ -25,17 +24,31 @@ namespace llvm {
 
 namespace codeview {
 
+/// CVRecord is a fat pointer (base + size pair) to a symbol or type record.
+/// Carrying the size separately instead of trusting the size stored in the
+/// record prefix provides some extra safety and flexibility.
 template <typename Kind> class CVRecord {
 public:
-  CVRecord() : Type(static_cast<Kind>(0)) {}
+  CVRecord() = default;
+
+  CVRecord(ArrayRef<uint8_t> Data) : RecordData(Data) {}
 
-  CVRecord(Kind K, ArrayRef<uint8_t> Data) : Type(K), RecordData(Data) {}
+  CVRecord(const RecordPrefix *P, size_t Size)
+      : RecordData(reinterpret_cast<const uint8_t *>(P), Size) {}
 
-  bool valid() const { return Type != static_cast<Kind>(0); }
+  bool valid() const { return kind() != Kind(0); }
 
   uint32_t length() const { return RecordData.size(); }
-  Kind kind() const { return Type; }
+
+  Kind kind() const {
+    if (RecordData.size() < sizeof(RecordPrefix))
+      return Kind(0);
+    return static_cast<Kind>(static_cast<uint16_t>(
+        reinterpret_cast<const RecordPrefix *>(RecordData.data())->RecordKind));
+  }
+
   ArrayRef<uint8_t> data() const { return RecordData; }
+
   StringRef str_data() const {
     return StringRef(reinterpret_cast<const char *>(RecordData.data()),
                      RecordData.size());
@@ -45,7 +58,6 @@ public:
     return RecordData.drop_front(sizeof(RecordPrefix));
   }
 
-  Kind Type;
   ArrayRef<uint8_t> RecordData;
 };
 
@@ -72,8 +84,7 @@ Error forEachCodeViewRecord(ArrayRef<uint8_t> StreamBuffer, Func F) {
     ArrayRef<uint8_t> Data = StreamBuffer.take_front(RealLen);
     StreamBuffer = StreamBuffer.drop_front(RealLen);
 
-    Record R(static_cast<decltype(Record::Type)>((uint16_t)Prefix->RecordKind),
-             Data);
+    Record R(Data);
     if (auto EC = F(R))
       return EC;
   }
@@ -92,13 +103,12 @@ inline Expected<CVRecord<Kind>> readCVRecordFromStream(BinaryStreamRef Stream,
     return std::move(EC);
   if (Prefix->RecordLen < 2)
     return make_error<CodeViewError>(cv_error_code::corrupt_record);
-  Kind K = static_cast<Kind>(uint16_t(Prefix->RecordKind));
 
   Reader.setOffset(Offset);
   ArrayRef<uint8_t> RawData;
   if (auto EC = Reader.readBytes(RawData, Prefix->RecordLen + sizeof(uint16_t)))
     return std::move(EC);
-  return codeview::CVRecord<Kind>(K, RawData);
+  return codeview::CVRecord<Kind>(RawData);
 }
 
 } // end namespace codeview