Add Check for WordCount Against SPIRV File Size (#3266)

A check is inserted in `parseAndCreateSPIRVEntry` function to ensure the
`WordCount` for the current SPIR-V instruction is not smaller than the
remaining input stream size. If the stream does not contain enough data
for the specified `WordCount`, an SPIRV error of `InvalidWordCount` is
raised.

Author: YixingZhang007

lib/SPIRV/SPIRVWriter.cpp

@@ -2088,7 +2088,10 @@ LLVMToSPIRVBase::transValueWithoutDecoration(Value *V, SPIRVBasicBlock *BB,
                << "have (" << MaxNumElements << "). Should the array be "
                << "split?\n Original LLVM value:\n"
                << toString(GV);
-            getErrorLog().checkError(false, SPIRVEC_InvalidWordCount, SS.str());
+            getErrorLog().checkError(false, SPIRVEC_InvalidWordCount,
+                                     "Can't encode instruction with word count "
+                                     "greater than 65535:\n" +
+                                         SS.str());
           }
         }
       }

lib/SPIRV/libSPIRV/SPIRVEntry.h

@@ -412,7 +412,10 @@ class SPIRVEntry {
       std::stringstream SS;
       SS << "Id: " << Id << ", OpCode: " << OpCodeNameMap::map(OpCode)
          << ", Name: \"" << Name << "\"\n";
-      getErrorLog().checkError(false, SPIRVEC_InvalidWordCount, SS.str());
+      getErrorLog().checkError(
+          false, SPIRVEC_InvalidWordCount,
+          "Can't encode instruction with word count greater than 65535:\n" +
+              SS.str());
     }
   }
   void validateFunctionControlMask(SPIRVWord FCtlMask) const;

lib/SPIRV/libSPIRV/SPIRVErrorEnum.h

@@ -16,8 +16,7 @@ _SPIRV_OP(InvalidLlvmModule, "Invalid LLVM module:")
 _SPIRV_OP(UnimplementedOpCode, "Unimplemented opcode")
 _SPIRV_OP(FunctionPointers, "Can't translate function pointer:\n")
 _SPIRV_OP(InvalidInstruction, "Can't translate llvm instruction:\n")
-_SPIRV_OP(InvalidWordCount,
-          "Can't encode instruction with word count greater than 65535:\n")
+_SPIRV_OP(InvalidWordCount, "")
 _SPIRV_OP(Requires1_1, "Feature requires SPIR-V 1.1 or greater:")
 _SPIRV_OP(RequiresVersion, "Cannot fulfill SPIR-V version restriction:\n")
 _SPIRV_OP(RequiresExtension,

lib/SPIRV/libSPIRV/SPIRVModule.cpp

@@ -2385,6 +2385,23 @@ SPIRVEntry *parseAndCreateSPIRVEntry(SPIRVWord &WordCount, Op &OpCode,
   if (WordCount == 0 || OpCode == OpNop) {
     return nullptr;
   }
+  if (!SPIRVUseTextFormat) {
+    std::streampos CurrentPos = IS.tellg();
+    IS.seekg(0, std::ios::end);
+    std::streamoff RemainingBytes = IS.tellg() - CurrentPos;
+    IS.clear();
+    IS.seekg(CurrentPos);
+    std::streamoff ExpectedBytes =
+        static_cast<std::streamoff>((WordCount - 1) * sizeof(SPIRVWord));
+    if (RemainingBytes < ExpectedBytes) {
+      M.getErrorLog().checkError(
+          false, SPIRVEC_InvalidWordCount,
+          "WordCount exceeds remaining input stream size: expected size = " +
+              std::to_string(ExpectedBytes) + " bytes, remaining size = " +
+              std::to_string(RemainingBytes) + " bytes");
+      M.setInvalid();
+    }
+  }
   SPIRVEntry *Entry = SPIRVEntry::create(OpCode);
   assert(Entry);
   Entry->setModule(&M);