[DX] Improve parse error messages

[beanz]

Summary

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.

[bob80905]

llvm/lib/Object/DXContainer.cpp

108–109

This comment blocks seems pretty irrelevant now, so maybe it should be removed / updated.
I am also concerned this might change some error text behavior, since I don't see the PartOffset variable ever getting the header size value subtracted from it.
Although it's possible to exit early when PartOffset < LastOffset, what happens if, say, PartOffset = Data.getBufferSize() - sizeof(dxbc::PartHeader) + 1?
In the old code, it returns, but in this new code, it might not enter the if block, right?

[beanz]

llvm/lib/Object/DXContainer.cpp

108–109

Yep! I need to update that.

The PartOffset field is read from the buffer, so it shouldn't need to be modified. The offset values in the file _should_ point to valid offsets in the file from which we can read at least a part header (4 bytes for the part name and 4 bytes for the part size).

The ParsePartNoSize test below should cover the case where a part offset is too close to the end of the file to have a size field, but I think you're right that there is a missing case where the part size can't handle the part name field.

I'll update the patch.

Thanks!

[beanz]

Updating based on PR feedback from bob80905. Thank you!

[bob80905]

Looks good to me.
Only minor nit at llvm/lib/Object/DXContainer.cpp:114.
Changing the condition from if (PartOffset >= Data.getBufferSize() - sizeof(dxbc::PartHeader::Name)) to if (PartOffset >= Data.getBufferSize() - sizeof(dxbc::PartHeader::Name) && PartOffset < Data.getBufferSize()))
would allow the parsed error message to be more precise, by explicitly stating the error is due to the partOffset being unable to read the part header name.
Then, another condition, PartOffset >= Data.getBufferSize())), would be able to definitively state that the part offset is beyond the file boundary.

llvm/include/llvm/Object/DXContainer.h

private:
   dxbc::Header Header;
   SmallVector<uint32_t, 4> PartOffsets;
   std::optional<DXILData> DXIL;
   std::optional<uint64_t> ShaderFlags;
   std::optional<dxbc::ShaderHash> Hash;
 
   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:
   // The PartIterator is a wrapper around the iterator for the PartOffsets
   // member of the DXContainer. It contains a refernce to the container, and the
   // current iterator value, as well as storage for a parsed part header.
   class PartIterator {
     const DXContainer &Container;

llvm/lib/Object/DXContainer.cpp

 //===- DXContainer.cpp - DXContainer object file implementation -----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #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;
 
 static Error parseFailed(const Twine &Msg) {
   return make_error<GenericBinaryError>(Msg.str(), object_error::parse_failed);
 }
 
 template <typename T>
 static Error readStruct(StringRef Buffer, const char *Src, T &Struct) {
   // 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");
 
   memcpy(&Struct, Src, sizeof(T));
   // DXContainer is always little endian
   if (sys::IsBigEndianHost)
     Struct.swapBytes();
   return Error::success();
 }
 
 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
   // number of parts and part data itself is not required to be padded.
   if (reinterpret_cast<uintptr_t>(Src) % alignof(T) != 0)
     memcpy(reinterpret_cast<char *>(&Val), Src, sizeof(T));
   else
     Val = *reinterpret_cast<const T *>(Src);
   // DXContainer is always little endian
   if (sys::IsBigEndianHost)
     sys::swapByteOrder(Val);
   return Error::success();
 }
 
 DXContainer::DXContainer(MemoryBufferRef O) : Data(O) {}
 
 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:
       break;
     }
   }
   return Error::success();
 }

llvm/test/tools/obj2yaml/DXContainer/PartTooSmall.yaml

+# 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
+...

llvm/unittests/Object/DXContainerTest.cpp

uint8_t Buffer[] = {
       0x00, 0x00, 0x70, 0x0D, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
   };
   EXPECT_THAT_EXPECTED(
       DXContainer::create(getMemoryBuffer<32>(Buffer)),
       FailedWithMessage("Reading structure out of file bounds"));
 }
 
 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,
       0x70, 0x0D, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF,
   };
   EXPECT_THAT_EXPECTED(
       DXContainer::create(getMemoryBuffer<36>(Buffer)),
       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,
       0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
       0x70, 0x0D, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00, 0x3C, 0x00, 0x00, 0x00,
       0x44, 0x00, 0x00, 0x00, 0x4C, 0x00, 0x00, 0x00, 0x54, 0x00, 0x00, 0x00,
       0x5C, 0x00, 0x00, 0x00, 0x64, 0x00, 0x00, 0x00, 0x6C, 0x00, 0x00, 0x00,
       0x46, 0x4B, 0x45, 0x30, 0x00, 0x00, 0x00, 0x00, 0x46, 0x4B, 0x45, 0x31,