Add additional guardrails against unsigned underflows caused by wrong WordCount (#3787)

Following up on #3671, this adds additional guardrails for cases that
were still unhandled by the previous patch. This PR also includes
multiple new tests for invalid wordcounts, and joins them all --the new,
and the ones added by #3671-- in their own directory
(test/negative/invalid-wordcount).

Author: maarquitos14

lib/SPIRV/libSPIRV/SPIRVDecorate.h

@@ -125,6 +125,7 @@ class SPIRVDecorate : public SPIRVDecorateGeneric {
 public:
   static const Op OC = OpDecorate;
   static const SPIRVWord FixedWC = 3;
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
   // Complete constructor for decorations without literals
   SPIRVDecorate(Decoration TheDec, SPIRVEntry *TheTarget)
       : SPIRVDecorateGeneric(OC, 3, TheDec, TheTarget) {}
@@ -240,6 +241,7 @@ class SPIRVDecorateId : public SPIRVDecorateGeneric {
 public:
   static const Op OC = OpDecorateId;
   static const SPIRVWord FixedWC = 3;
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
   // Complete constructor for decorations with one id operand
   SPIRVDecorateId(Decoration TheDec, SPIRVEntry *TheTarget, SPIRVId V)
       : SPIRVDecorateGeneric(OC, 4, TheDec, TheTarget, V) {}
@@ -330,6 +332,7 @@ class SPIRVMemberDecorate : public SPIRVDecorateGeneric {
 public:
   static const Op OC = OpMemberDecorate;
   static const SPIRVWord FixedWC = 4;
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
   // Complete constructor for decorations without literals
   SPIRVMemberDecorate(Decoration TheDec, SPIRVWord Member,
                       SPIRVEntry *TheTarget)
@@ -447,6 +450,8 @@ class SPIRVGroupDecorateGeneric : public SPIRVEntryNoIdGeneric {
     SPIRVCK(WordCount >= FixedWC, InvalidWordCount, "");
     Targets.resize(WordCount - FixedWC);
   }
+
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
   virtual void decorateTargets() = 0;
   _SPIRV_DCL_ENCDEC
 protected:

lib/SPIRV/libSPIRV/SPIRVEntry.h

@@ -380,6 +380,10 @@ class SPIRVEntry {
   /// its variable sized member before decoding the remaining words.
   virtual void setWordCount(SPIRVWord TheWordCount);
 
+  /// Return the minimum valid WordCount (1, the OpCode). Other entries will
+  /// override to return their respective fixed word counts.
+  virtual SPIRVWord getFixedWordCount() const { return 1; }
+
   /// Create an empty SPIRV object by op code, e.g. OpTypeInt creates
   /// SPIRVTypeInt.
   static SPIRVEntry *create(Op);
@@ -541,6 +545,7 @@ class SPIRVAnnotation : public SPIRVAnnotationGeneric {
 class SPIRVEntryPoint : public SPIRVAnnotation {
 public:
   static const SPIRVWord FixedWC = 4;
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
   SPIRVEntryPoint(SPIRVModule *TheModule, SPIRVExecutionModelKind,
                   SPIRVId TheId, const std::string &TheName,
                   std::vector<SPIRVId> Variables);

lib/SPIRV/libSPIRV/SPIRVFnVar.h

@@ -51,6 +51,7 @@ bool specializeFnVariants(SPIRVModule *BM, std::string &ErrMsg);
 class SPIRVConditionalEntryPointINTEL : public SPIRVAnnotation {
 public:
   static const SPIRVWord FixedWC = 5;
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
   SPIRVConditionalEntryPointINTEL(SPIRVModule *TheModule, SPIRVId Condition,
                                   SPIRVExecutionModelKind TheExecModel,
                                   SPIRVId TheId, const std::string &TheName,
@@ -249,6 +250,7 @@ class SPIRVSpecConstantTargetINTEL : public SPIRVValue {
 
     return Res;
   }
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
 
 protected:
   _SPIRV_DEF_ENCDEC4(Type, Id, Target, Features);
@@ -414,6 +416,7 @@ class SPIRVSpecConstantCapabilitiesINTEL : public SPIRVValue {
 
     return Res;
   }
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
 
 protected:
   _SPIRV_DEF_ENCDEC3(Type, Id, Capabilities);

lib/SPIRV/libSPIRV/SPIRVInstruction.h

@@ -280,8 +280,14 @@ class SPIRVInstTemplateBase : public SPIRVInstruction {
     // instructions.
     updateModuleVersion();
   }
+  SPIRVWord getFixedWordCount() const override {
+    // OpCode word plus the optional result-type and
+    // result-id words. Operands beyond this are variable length.
+    return 1 + (hasType() ? 1 : 0) + (hasId() ? 1 : 0);
+  }
   void setWordCount(SPIRVWord TheWordCount) override {
     SPIRVEntry::setWordCount(TheWordCount);
+    SPIRVCK(WordCount >= getFixedWordCount(), InvalidWordCount, "");
     auto NumOps = WordCount - 1;
     if (hasId())
       --NumOps;
@@ -512,6 +518,7 @@ class SPIRVVariableBase : public SPIRVInstruction {
       return std::vector<SPIRVEntry *>(1, V);
     return std::vector<SPIRVEntry *>();
   }
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
 
 protected:
   void validate() const override {

lib/SPIRV/libSPIRV/SPIRVModule.cpp

@@ -2484,6 +2484,15 @@ static SPIRVEntry *parseAndCreateSPIRVEntry(SPIRVWord &WordCount, Op &OpCode,
   if (Scope && !isModuleScopeAllowedOpCode(OpCode)) {
     Entry->setScope(Scope);
   }
+  // Reject a WordCount below the instruction's minimum before setWordCount().
+  if (!M.getErrorLog().checkError(WordCount >= Entry->getFixedWordCount(),
+                                  SPIRVEC_InvalidWordCount,
+                                  "WordCount is smaller than the instruction's "
+                                  "minimum word count")) {
+    M.setInvalid();
+    delete Entry;
+    return nullptr;
+  }
   Entry->setWordCount(WordCount);
   if (OpCode != OpLine)
     Entry->setLine(M.getCurrentLine());

lib/SPIRV/libSPIRV/SPIRVType.h

@@ -660,6 +660,7 @@ class SPIRVTypeImage : public SPIRVType {
   std::vector<SPIRVEntry *> getNonLiteralOperands() const override {
     return std::vector<SPIRVEntry *>(1, get<SPIRVType>(SampledType));
   }
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
 
 protected:
   _SPIRV_DEF_ENCDEC9(Id, SampledType, Desc.Dim, Desc.Depth, Desc.Arrayed,
@@ -805,6 +806,8 @@ class SPIRVTypeStruct : public SPIRVType {
     MemberTypeIdVec.resize(WordCount - FixedWC);
   }
 
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
+
   // TODO: Should we attach operands of continued instructions as well?
   std::vector<SPIRVEntry *> getNonLiteralOperands() const override {
     std::vector<SPIRVEntry *> Operands(MemberTypeIdVec.size());
@@ -877,6 +880,7 @@ class SPIRVTypeFunction : public SPIRVType {
       Operands.push_back(getEntry(I));
     return Operands;
   }
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
 
 protected:
   _SPIRV_DEF_ENCDEC3(Id, ReturnType, ParamTypeIdVec)

lib/SPIRV/libSPIRV/SPIRVValue.h

@@ -174,6 +174,7 @@ template <spv::Op OC> class SPIRVConstantBase : public SPIRVValue {
   double getDoubleValue() const { return getValue<double>(); }
   unsigned getNumWords() const { return NumWords; }
   const std::vector<SPIRVWord> &getSPIRVWords() { return Words; }
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
 
 protected:
   constexpr static SPIRVWord FixedWC = 3;
@@ -210,6 +211,7 @@ template <spv::Op OC> class SPIRVConstantBase : public SPIRVValue {
   }
   void setWordCount(SPIRVWord WordCount) override {
     SPIRVValue::setWordCount(WordCount);
+    SPIRVCK(WordCount >= FixedWC, InvalidWordCount, "");
     NumWords = WordCount - FixedWC;
   }
   void decode(std::istream &I) override {
@@ -358,6 +360,7 @@ template <spv::Op OC> class SPIRVConstantCompositeBase : public SPIRVValue {
     for (auto &I : ContinuedInstructions)
       O << *I;
   }
+  SPIRVWord getFixedWordCount() const override { return FixedWC; }
 
 protected:
   void validate() const override {

test/negative/invalid-wordcount/invalid-wordcount-aliasscopedecl.spt

@@ -0,0 +1,12 @@
+; OpAliasScopeDeclINTEL minimum = 2 (header + result-id); the scope arguments
+; are variable length. WordCount=1 is malformed.
+; Test that malformed input is rejected early rather than silently causing
+; integer underflow and potential process hang.
+
+; RUN: not llvm-spirv %s -to-binary -o %t.spv 2>&1 | FileCheck %s
+; CHECK: InvalidWordCount:
+
+119734787 65536 0 10 0
+2 Capability Kernel
+3 MemoryModel 1 2
+1 AliasScopeDeclINTEL 1

test/negative/invalid-wordcount/invalid-wordcount-asmcallintel.spt

@@ -0,0 +1,12 @@
+; OpAsmCallINTEL minimum = 4 (header + result-type + result-id + asm-target);
+; the arguments are variable length. WordCount=3 is malformed.
+; Test that malformed input is rejected early rather than silently causing
+; integer underflow and potential process hang.
+
+; RUN: not llvm-spirv %s -to-binary -o %t.spv 2>&1 | FileCheck %s
+; CHECK: InvalidWordCount:
+
+119734787 65536 0 10 0
+2 Capability Kernel
+3 MemoryModel 1 2
+3 AsmCallINTEL 1 2 3