Add additional guardrails against unsigned underflows caused by wrong WordCount

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

@@ -655,6 +655,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,
@@ -800,6 +801,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());
@@ -872,6 +875,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

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

@@ -0,0 +1,12 @@
+; OpCompositeConstruct minimum = 3 (header + result-type + result-id); the
+; constituents are variable length. WordCount=2 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
+2 CompositeConstruct 1 2

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

@@ -0,0 +1,12 @@
+; OpConditionalCopyObjectINTEL minimum = 3 (header + result-type + result-id);
+; the condition/operand pairs are variable length. WordCount=2 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
+2 ConditionalCopyObjectINTEL 1 2

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

@@ -0,0 +1,13 @@
+; OpConditionalEntryPointINTEL minimum = 5 (header + condition + execution-model
+; + entry-point + name); the interface ids are variable length. WordCount=4 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
+4 ConditionalEntryPointINTEL 1 6 2 3

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

@@ -0,0 +1,13 @@
+; OpConstant minimum = 3 (header + result-type + result-id); the value words
+; are variable length. WordCount=2 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 6 0
+2 Capability Kernel
+3 MemoryModel 1 2
+4 TypeInt 3 32 0
+2 Constant 3 5

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

@@ -0,0 +1,12 @@
+; OpCopyMemory minimum = 3 (header + target + source); the optional
+; memory-access operands are variable length. WordCount=2 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
+2 CopyMemory 1 2

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

@@ -0,0 +1,12 @@
+; OpCopyMemorySized minimum = 4 (header + target + source + size); the optional
+; memory-access operands 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 CopyMemorySized 1 2 3

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

@@ -0,0 +1,12 @@
+; OpDecorateId minimum = 3 (header + target + decoration); the decoration
+; literals are variable length. WordCount=2 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
+2 DecorateId 1 2

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

@@ -0,0 +1,12 @@
+; OpGroupMemberDecorate minimum = 2 (header + decoration-group); the
+; target/member pairs 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 GroupMemberDecorate 1

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

@@ -0,0 +1,14 @@
+; Generic SPIRVInstTemplate decode path: minimum = 3 for an instruction with
+; both a result type and a result id (header + result-type + result-id);
+; operands are variable length. WordCount=2 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 6 0
+2 Capability Kernel
+3 MemoryModel 1 2
+4 TypeInt 3 32 0
+2 IAdd 3 5

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

@@ -0,0 +1,12 @@
+; OpLoad minimum = 4 (header + result-type + result-id + pointer); the optional
+; memory-access operands 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 Load 1 2 3

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

@@ -0,0 +1,12 @@
+; OpLoopControlINTEL minimum = 2 (header + loop-control); the loop-control
+; parameters 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 LoopControlINTEL 0

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

@@ -0,0 +1,13 @@
+; OpLoopMerge minimum = 4 (header + merge-block + continue-target +
+; loop-control); the loop-control parameters 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 LoopMerge 1 2 0

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

@@ -0,0 +1,12 @@
+; OpMemberDecorate minimum = 4 (header + struct-type + member + decoration);
+; the decoration literals 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 MemberDecorate 1 0 2

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

@@ -0,0 +1,12 @@
+; OpPhi minimum = 3 (header + result-type + result-id); the variable/parent
+; pairs are variable length. WordCount=2 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
+2 Phi 1 2

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

@@ -0,0 +1,13 @@
+; OpSpecConstantCapabilitiesINTEL minimum = 3 (header + result-type +
+; result-id); the capability operands are variable length. WordCount=2 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
+2 SpecConstantCapabilitiesINTEL 1 2