adding few more tests

Author: joaosaffran

llvm/include/llvm/BinaryFormat/DXContainer.h

@@ -14,9 +14,12 @@
 #define LLVM_BINARYFORMAT_DXCONTAINER_H
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Object/Error.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/SwapByteOrder.h"
 #include "llvm/TargetParser/Triple.h"
 
+#include <cstdint>
 #include <stdint.h>
 
 namespace llvm {
@@ -63,15 +66,6 @@ struct ShaderHash {
   void swapBytes() { sys::swapByteOrder(Flags); }
 };
 
-struct RootSignatureDesc {
-  uint32_t Size;
-  uint32_t Flags;
-
-  void swapBytes() {
-    sys::swapByteOrder(Size);
-    sys::swapByteOrder(Flags);
-  }
-};
 
 struct ContainerVersion {
   uint16_t Major;
@@ -556,6 +550,21 @@ struct ProgramSignatureElement {
 static_assert(sizeof(ProgramSignatureElement) == 32,
               "ProgramSignatureElement is misaligned");
 
+struct RootSignatureValidations {
+
+    static Expected<uint32_t> validateRootFlag(uint32_t Flags) {
+      if ((Flags & ~0x80000fff) != 0)
+        return llvm::make_error<object::GenericBinaryError>("Invalid flag");
+      return Flags;
+    }
+
+    static Expected<uint32_t> validateVersion(uint32_t Version) {
+      if (Version < 1 || Version > 2)
+        return llvm::make_error<object::GenericBinaryError>("Invalid Version");
+      return Version;
+    }
+};
+
 } // namespace dxbc
 } // namespace llvm
 

llvm/include/llvm/Object/DXContainer.h

@@ -18,6 +18,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/BinaryFormat/DXContainer.h"
+#include "llvm/Object/Error.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/TargetParser/Triple.h"

llvm/lib/Object/DXContainer.cpp

@@ -12,7 +12,9 @@
 #include "llvm/Object/Error.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/Endian.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/FormatVariadic.h"
+#include <cstdint>
 
 using namespace llvm;
 using namespace llvm::object;
@@ -246,14 +248,20 @@ void DXContainer::PartIterator::updateIteratorImpl(const uint32_t Offset) {
 Error DirectX::RootSignature::parse(StringRef Data) {
   const char *Current = Data.begin();
 
+
   // Root Signature headers expects 6 integers to be present.
   if (Data.size() < 6 * sizeof(uint32_t)) {
     return parseFailed("Invalid data. Too small.");
   }
 
-  Version = support::endian::read<uint32_t, llvm::endianness::little>(Current);
+  uint32_t VValue = support::endian::read<uint32_t, llvm::endianness::little>(Current);
   Current += sizeof(uint32_t);
 
+  Expected<uint32_t> MaybeVersion = dxbc::RootSignatureValidations::validateVersion(VValue);
+  if(Error E = MaybeVersion.takeError())
+    return E;
+  Version = MaybeVersion.get();
+
   NumParameters =
       support::endian::read<uint32_t, llvm::endianness::little>(Current);
   Current += sizeof(uint32_t);
@@ -270,9 +278,14 @@ Error DirectX::RootSignature::parse(StringRef Data) {
       support::endian::read<uint32_t, llvm::endianness::little>(Current);
   Current += sizeof(uint32_t);
 
-  Flags = support::endian::read<uint32_t, llvm::endianness::little>(Current);
+  uint32_t FValue = support::endian::read<uint32_t, llvm::endianness::little>(Current);
   Current += sizeof(uint32_t);
 
+  Expected<uint32_t> MaybeFlag = dxbc::RootSignatureValidations::validateRootFlag(FValue);
+  if(Error E = MaybeFlag.takeError())
+    return E;
+  Flags = MaybeFlag.get();
+
   return Error::success();
 }
 

llvm/unittests/Object/DXContainerTest.cpp

@@ -833,7 +833,7 @@ TEST(RootSignature, ParseRootFlags) {
         0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
     };
     DXContainer C =
-        llvm::cantFail(DXContainer::create(getMemoryBuffer<180>(Buffer)));
+        llvm::cantFail(DXContainer::create(getMemoryBuffer<68>(Buffer)));
 
     const auto &RS = C.getRootSignature();
     ASSERT_TRUE(RS.has_value());
@@ -855,7 +855,33 @@ TEST(RootSignature, ParseRootFlags) {
         0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
         0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
     };
-    EXPECT_THAT_EXPECTED(DXContainer::create(getMemoryBuffer<44>(Buffer)),
+    EXPECT_THAT_EXPECTED(DXContainer::create(getMemoryBuffer<64>(Buffer)),
                          FailedWithMessage("Invalid data. Too small."));
   }
+  {
+    // Version has been changed to an invalid number.
+    uint8_t Buffer[] = {
+        0x44, 0x58, 0x42, 0x43, 0x32, 0x9A, 0x53, 0xD8, 0xEC, 0xBE, 0x35, 0x6F,
+        0x05, 0x39, 0xE1, 0xFE, 0x31, 0x20, 0xF0, 0xC1, 0x01, 0x00, 0x00, 0x00,
+        0x44, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x24, 0x00, 0x00, 0x00,
+        0x52, 0x54, 0x53, 0x30, 0x18, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
+    };
+        EXPECT_THAT_EXPECTED(DXContainer::create(getMemoryBuffer<68>(Buffer)),
+                         FailedWithMessage("Invalid Version"));
+  }
+  {
+    // Flag has been set to an invalid value
+    uint8_t Buffer[] = {
+        0x44, 0x58, 0x42, 0x43, 0x32, 0x9A, 0x53, 0xD8, 0xEC, 0xBE, 0x35, 0x6F,
+        0x05, 0x39, 0xE1, 0xFE, 0x31, 0x20, 0xF0, 0xC1, 0x01, 0x00, 0x00, 0x00,
+        0x44, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x24, 0x00, 0x00, 0x00,
+        0x52, 0x54, 0x53, 0x30, 0x18, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0xFF,
+    };
+        EXPECT_THAT_EXPECTED(DXContainer::create(getMemoryBuffer<68>(Buffer)),
+                         FailedWithMessage("Invalid flag"));
+  }
 }