[llvm] 3fb0621 - [Attributes] Assert correct attribute constructor is used (NFCI)

[Nikita Popov via llvm-commits]

Author: Nikita Popov
Date: 2021-07-12T21:11:59+02:00
New Revision: 3fb0621fafbb16ea157830b308c26a77ba0a04e3

URL: https://github.com/llvm/llvm-project/commit/3fb0621fafbb16ea157830b308c26a77ba0a04e3
DIFF: https://github.com/llvm/llvm-project/commit/3fb0621fafbb16ea157830b308c26a77ba0a04e3.diff

LOG: [Attributes] Assert correct attribute constructor is used (NFCI)

Assert that enum/int/type attributes go through the constructor
they are supposed to use.

To make sure this can't happen via invalid bitcode, explicitly
verify that the attribute kind if correct there.

Added: 
    

Modified: 
    llvm/include/llvm/IR/Attributes.h
    llvm/lib/Bitcode/Reader/BitcodeReader.cpp
    llvm/lib/IR/Attributes.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/IR/Attributes.h b/llvm/include/llvm/IR/Attributes.h
index 4c2a3665f7d46..2f04c02e4b760 100644
--- a/llvm/include/llvm/IR/Attributes.h
+++ b/llvm/include/llvm/IR/Attributes.h
@@ -80,6 +80,12 @@ class Attribute {
 
   static const unsigned NumTypeAttrKinds = LastTypeAttr - FirstTypeAttr + 1;
 
+  static bool isEnumAttrKind(AttrKind Kind) {
+    return Kind >= FirstEnumAttr && Kind <= LastEnumAttr;
+  }
+  static bool isIntAttrKind(AttrKind Kind) {
+    return Kind >= FirstIntAttr && Kind <= LastIntAttr;
+  }
   static bool isTypeAttrKind(AttrKind Kind) {
     return Kind >= FirstTypeAttr && Kind <= LastTypeAttr;
   }

diff  --git a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
index 585527fb7c007..c9c1cef3c8793 100644
--- a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
+++ b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
@@ -1599,12 +1599,16 @@ Error BitcodeReader::parseAttributeGroupBlock() {
             B.addStructRetAttr(nullptr);
           else if (Kind == Attribute::InAlloca)
             B.addInAllocaAttr(nullptr);
-
-          B.addAttribute(Kind);
+          else if (Attribute::isEnumAttrKind(Kind))
+            B.addAttribute(Kind);
+          else
+            return error("Not an enum attribute");
         } else if (Record[i] == 1) { // Integer attribute
           Attribute::AttrKind Kind;
           if (Error Err = parseAttrKind(Record[++i], &Kind))
             return Err;
+          if (!Attribute::isIntAttrKind(Kind))
+            return error("Not an int attribute");
           if (Kind == Attribute::Alignment)
             B.addAlignmentAttr(Record[++i]);
           else if (Kind == Attribute::StackAlignment)

diff  --git a/llvm/lib/IR/Attributes.cpp b/llvm/lib/IR/Attributes.cpp
index 45f7f183d4b02..3c3ffeabdbb48 100644
--- a/llvm/lib/IR/Attributes.cpp
+++ b/llvm/lib/IR/Attributes.cpp
@@ -91,6 +91,11 @@ static std::pair<unsigned, unsigned> unpackVScaleRangeArgs(uint64_t Value) {
 
 Attribute Attribute::get(LLVMContext &Context, Attribute::AttrKind Kind,
                          uint64_t Val) {
+  if (Val)
+    assert(Attribute::isIntAttrKind(Kind) && "Not an int attribute");
+  else
+    assert(Attribute::isEnumAttrKind(Kind) && "Not an enum attribute");
+
   LLVMContextImpl *pImpl = Context.pImpl;
   FoldingSetNodeID ID;
   ID.AddInteger(Kind);