[DX] Improve parse error messages

This change refactors the parte parsing logic to operate on StringRefs
of the part data rather than starting from an offset and splicing down.
It also improves some of the error reporting around part layout.

Specifically, this code now reports a distinct error if there isn't
enough data in the buffer to store the part size and it reports an
error if the parts overlap.

Reviewed By: bob80905

Differential Revision: https://reviews.llvm.org/D139681

diff --git a/llvm/include/llvm/Object/DXContainer.h b/llvm/include/llvm/Object/DXContainer.h
index ae819e9..ffa2db4 100644 (file)
--- a/llvm/include/llvm/Object/DXContainer.h
+++ b/llvm/include/llvm/Object/DXContainer.h
@@ -39,9 +39,9 @@ private:
 
   Error parseHeader();
   Error parsePartOffsets();
-  Error parseDXILHeader(uint32_t Offset);
-  Error parseShaderFlags(uint32_t Offset);
-  Error parseHash(uint32_t Offset);
+  Error parseDXILHeader(StringRef Part);
+  Error parseShaderFlags(StringRef Part);
+  Error parseHash(StringRef Part);
   friend class PartIterator;
 
 public:

diff --git a/llvm/lib/Object/DXContainer.cpp b/llvm/lib/Object/DXContainer.cpp
index 35a58ed..4d8f261 100644 (file)
--- a/llvm/lib/Object/DXContainer.cpp
+++ b/llvm/lib/Object/DXContainer.cpp
@@ -9,6 +9,7 @@
 #include "llvm/Object/DXContainer.h"
 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/Object/Error.h"
+#include "llvm/Support/FormatVariadic.h"
 
 using namespace llvm;
 using namespace llvm::object;
@@ -31,12 +32,13 @@ static Error readStruct(StringRef Buffer, const char *Src, T &Struct) {
 }
 
 template <typename T>
-static Error readInteger(StringRef Buffer, const char *Src, T &Val) {
+static Error readInteger(StringRef Buffer, const char *Src, T &Val,
+                         Twine Str = "structure") {
   static_assert(std::is_integral_v<T>,
                 "Cannot call readInteger on non-integral type.");
   // Don't read before the beginning or past the end of the file
   if (Src < Buffer.begin() || Src + sizeof(T) > Buffer.end())
-    return parseFailed("Reading structure out of file bounds");
+    return parseFailed(Twine("Reading ") + Str + " out of file bounds");
 
   // The DXContainer offset table is comprised of uint32_t values but not padded
   // to a 64-bit boundary. So Parts may start unaligned if there is an odd
@@ -57,69 +59,85 @@ Error DXContainer::parseHeader() {
   return readStruct(Data.getBuffer(), Data.getBuffer().data(), Header);
 }
 
-Error DXContainer::parseDXILHeader(uint32_t Offset) {
+Error DXContainer::parseDXILHeader(StringRef Part) {
   if (DXIL)
     return parseFailed("More than one DXIL part is present in the file");
-  const char *Current = Data.getBuffer().data() + Offset;
+  const char *Current = Part.begin();
   dxbc::ProgramHeader Header;
-  if (Error Err = readStruct(Data.getBuffer(), Current, Header))
+  if (Error Err = readStruct(Part, Current, Header))
     return Err;
   Current += offsetof(dxbc::ProgramHeader, Bitcode) + Header.Bitcode.Offset;
   DXIL.emplace(std::make_pair(Header, Current));
   return Error::success();
 }
 
-Error DXContainer::parseShaderFlags(uint32_t Offset) {
+Error DXContainer::parseShaderFlags(StringRef Part) {
   if (ShaderFlags)
     return parseFailed("More than one SFI0 part is present in the file");
-  const char *Current = Data.getBuffer().data() + Offset;
   uint64_t FlagValue = 0;
-  if (Error Err = readInteger(Data.getBuffer(), Current, FlagValue))
+  if (Error Err = readInteger(Part, Part.begin(), FlagValue))
     return Err;
   ShaderFlags = FlagValue;
   return Error::success();
 }
 
-Error DXContainer::parseHash(uint32_t Offset) {
+Error DXContainer::parseHash(StringRef Part) {
   if (Hash)
     return parseFailed("More than one HASH part is present in the file");
-  const char *Current = Data.getBuffer().data() + Offset;
   dxbc::ShaderHash ReadHash;
-  if (Error Err = readStruct(Data.getBuffer(), Current, ReadHash))
+  if (Error Err = readStruct(Part, Part.begin(), ReadHash))
     return Err;
   Hash = ReadHash;
   return Error::success();
 }
 
 Error DXContainer::parsePartOffsets() {
+  uint32_t LastOffset =
+      sizeof(dxbc::Header) + (Header.PartCount * sizeof(uint32_t));
   const char *Current = Data.getBuffer().data() + sizeof(dxbc::Header);
   for (uint32_t Part = 0; Part < Header.PartCount; ++Part) {
     uint32_t PartOffset;
     if (Error Err = readInteger(Data.getBuffer(), Current, PartOffset))
       return Err;
+    if (PartOffset < LastOffset)
+      return parseFailed(
+          formatv(
+              "Part offset for part {0} begins before the previous part ends",
+              Part)
+              .str());
     Current += sizeof(uint32_t);
-    // We need to ensure that each part offset leaves enough space for a part
-    // header. To prevent overflow, we subtract the part header size from the
-    // buffer size, rather than adding to the offset. Since the file header is
-    // larger than the part header we can't reach this code unless the buffer
-    // is larger than the part header, so this can't underflow.
-    if (PartOffset > Data.getBufferSize() - sizeof(dxbc::PartHeader))
+    if (PartOffset >= Data.getBufferSize())
       return parseFailed("Part offset points beyond boundary of the file");
+    // To prevent overflow when reading the part name, we subtract the part name
+    // size from the buffer size, rather than adding to the offset. Since the
+    // file header is larger than the part header we can't reach this code
+    // unless the buffer is at least as large as a part header, so this
+    // subtraction can't underflow.
+    if (PartOffset >= Data.getBufferSize() - sizeof(dxbc::PartHeader::Name))
+      return parseFailed("File not large enough to read part name");
     PartOffsets.push_back(PartOffset);
 
     dxbc::PartType PT =
         dxbc::parsePartType(Data.getBuffer().substr(PartOffset, 4));
+    uint32_t PartDataStart = PartOffset + sizeof(dxbc::PartHeader);
+    uint32_t PartSize;
+    if (Error Err = readInteger(Data.getBuffer(),
+                                Data.getBufferStart() + PartOffset + 4,
+                                PartSize, "part size"))
+      return Err;
+    StringRef PartData = Data.getBuffer().substr(PartDataStart, PartSize);
+    LastOffset = PartOffset + PartSize;
     switch (PT) {
     case dxbc::PartType::DXIL:
-      if (Error Err = parseDXILHeader(PartOffset + sizeof(dxbc::PartHeader)))
+      if (Error Err = parseDXILHeader(PartData))
         return Err;
       break;
     case dxbc::PartType::SFI0:
-      if (Error Err = parseShaderFlags(PartOffset + sizeof(dxbc::PartHeader)))
+      if (Error Err = parseShaderFlags(PartData))
         return Err;
       break;
     case dxbc::PartType::HASH:
-      if (Error Err = parseHash(PartOffset + sizeof(dxbc::PartHeader)))
+      if (Error Err = parseHash(PartData))
         return Err;
       break;
     case dxbc::PartType::Unknown:

diff --git a/llvm/test/tools/obj2yaml/DXContainer/PartTooSmall.yaml b/llvm/test/tools/obj2yaml/DXContainer/PartTooSmall.yaml
new file mode 100644 (file)
index 0000000..9f54546
--- /dev/null
+++ b/llvm/test/tools/obj2yaml/DXContainer/PartTooSmall.yaml
@@ -0,0 +1,17 @@
+# RUN: yaml2obj %s | not obj2yaml 2>&1 | FileCheck %s 
+
+# In this test the hash part is too small to contain the hash data.
+
+# CHECK: Error reading file: <stdin>: Reading structure out of file bounds
+--- !dxcontainer
+Header:
+  Hash:            [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 
+                     0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
+  Version:
+    Major:           1
+    Minor:           0
+  PartCount:       1
+Parts:
+  - Name:            HASH
+    Size:            0
+...

diff --git a/llvm/unittests/Object/DXContainerTest.cpp b/llvm/unittests/Object/DXContainerTest.cpp
index 4145ed6..ba04152 100644 (file)
--- a/llvm/unittests/Object/DXContainerTest.cpp
+++ b/llvm/unittests/Object/DXContainerTest.cpp
@@ -71,6 +71,7 @@ TEST(DXCFile, ParsePartMissingOffsets) {
 }
 
 TEST(DXCFile, ParsePartInvalidOffsets) {
+  // This test covers a case where the part offset is beyond the buffer size.
   uint8_t Buffer[] = {
       0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
       0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
@@ -81,6 +82,48 @@ TEST(DXCFile, ParsePartInvalidOffsets) {
       FailedWithMessage("Part offset points beyond boundary of the file"));
 }
 
+TEST(DXCFile, ParsePartTooSmallBuffer) {
+  // This test covers a case where there is insufficent space to read a full
+  // part name, but the offset for the part is inside the buffer.
+  uint8_t Buffer[] = {0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00,
+                      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+                      0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+                      0x26, 0x0D, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+                      0x24, 0x00, 0x00, 0x00, 0x46, 0x4B};
+  EXPECT_THAT_EXPECTED(
+      DXContainer::create(getMemoryBuffer<38>(Buffer)),
+      FailedWithMessage("File not large enough to read part name"));
+}
+
+TEST(DXCFile, ParsePartNoSize) {
+  // This test covers a case where the part's header is readable, but the size
+  // the part extends beyond the boundaries of the file.
+  uint8_t Buffer[] = {0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00,
+                      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+                      0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x28, 0x0D, 0x00,
+                      0x00, 0x01, 0x00, 0x00, 0x00, 0x24, 0x00, 0x00, 0x00,
+                      0x46, 0x4B, 0x45, 0x30, 0x00, 0x00};
+  EXPECT_THAT_EXPECTED(
+      DXContainer::create(getMemoryBuffer<42>(Buffer)),
+      FailedWithMessage("Reading part size out of file bounds"));
+}
+
+TEST(DXCFile, ParseOverlappingParts) {
+  // This test covers a case where a part's offset is inside the size range
+  // covered by the previous part.
+  uint8_t Buffer[] = {
+      0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+      0x40, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x28, 0x00, 0x00, 0x00,
+      0x2C, 0x00, 0x00, 0x00, 0x46, 0x4B, 0x45, 0x30, 0x08, 0x00, 0x00, 0x00,
+      0x46, 0x4B, 0x45, 0x31, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+  };
+  EXPECT_THAT_EXPECTED(
+      DXContainer::create(getMemoryBuffer<60>(Buffer)),
+      FailedWithMessage(
+          "Part offset for part 1 begins before the previous part ends"));
+}
+
 TEST(DXCFile, ParseEmptyParts) {
   uint8_t Buffer[] = {
       0x44, 0x58, 0x42, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,