[clang] Function type attribute to prevent CFI instrumentation

This introduces the attribute discussed in
https://discourse.llvm.org/t/rfc-function-type-attribute-to-prevent-cfi-instrumentation/85458.

The proposed name has been changed from `no_cfi` to
`cfi_unchecked_callee` to help differentiate from `no_sanitize("cfi")`
more easily. The proposed attribute has the following semantics:

1. Indirect calls to a function type with this attribute will not be
   instrumented with CFI. That is, the indirect call will not be
   checked. Note that this only changes the behavior for indirect calls
   on pointers to function types having this attribute. It does not
   prevent all indirect function calls for a given type from being checked.
2. All direct references to a function whose type has this attribute will
   always reference the true function definition rather than an entry
   in the CFI jump table.
3. When a pointer to a function with this attribute is implicitly cast
   to a pointer to a function without this attribute, the compiler
   will give a warning saying this attribute is discarded. This warning
   can be silenced with an explicit C-style cast or C++ static_cast.

Author: PiJoules

clang/include/clang/AST/ASTContext.h

@@ -262,6 +262,8 @@ class ASTContext : public RefCountedBase<ASTContext> {
 
   mutable llvm::FoldingSet<CountAttributedType> CountAttributedTypes;
 
+  mutable llvm::FoldingSet<CFIUncheckedCalleeType> CFIUncheckedCalleeTypes;
+
   mutable llvm::FoldingSet<QualifiedTemplateName> QualifiedTemplateNames;
   mutable llvm::FoldingSet<DependentTemplateName> DependentTemplateNames;
   mutable llvm::FoldingSet<SubstTemplateTemplateParmStorage>
@@ -1464,6 +1466,9 @@ class ASTContext : public RefCountedBase<ASTContext> {
                          bool OrNull,
                          ArrayRef<TypeCoupledDeclRefInfo> DependentDecls) const;
 
+  /// Return a type wrapped with the `cfi_unchecked_callee` attribute.
+  QualType getCFIUncheckedCalleeType(QualType Wrapped) const;
+
   /// Return the uniqued reference to a type adjusted from the original
   /// type to a new type.
   QualType getAdjustedType(QualType Orig, QualType New) const;

clang/include/clang/AST/ASTNodeTraverser.h

@@ -445,6 +445,9 @@ class ASTNodeTraverser
   void VisitBTFTagAttributedType(const BTFTagAttributedType *T) {
     Visit(T->getWrappedType());
   }
+  void VisitCFIUncheckedCalleeType(const CFIUncheckedCalleeType *T) {
+    Visit(T->getWrappedType());
+  }
   void VisitHLSLAttributedResourceType(const HLSLAttributedResourceType *T) {
     QualType Contained = T->getContainedType();
     if (!Contained.isNull())

clang/include/clang/AST/RecursiveASTVisitor.h

@@ -1142,6 +1142,9 @@ DEF_TRAVERSE_TYPE(InjectedClassNameType, {})
 DEF_TRAVERSE_TYPE(AttributedType,
                   { TRY_TO(TraverseType(T->getModifiedType())); })
 
+DEF_TRAVERSE_TYPE(CFIUncheckedCalleeType,
+                  { TRY_TO(TraverseType(T->getWrappedType())); })
+
 DEF_TRAVERSE_TYPE(CountAttributedType, {
   if (T->getCountExpr())
     TRY_TO(TraverseStmt(T->getCountExpr()));
@@ -1448,6 +1451,9 @@ DEF_TRAVERSE_TYPELOC(MacroQualifiedType,
 DEF_TRAVERSE_TYPELOC(AttributedType,
                      { TRY_TO(TraverseTypeLoc(TL.getModifiedLoc())); })
 
+DEF_TRAVERSE_TYPELOC(CFIUncheckedCalleeType,
+                     { TRY_TO(TraverseTypeLoc(TL.getWrappedLoc())); })
+
 DEF_TRAVERSE_TYPELOC(CountAttributedType,
                      { TRY_TO(TraverseTypeLoc(TL.getInnerLoc())); })
 

clang/include/clang/AST/Type.h

@@ -2566,6 +2566,8 @@ class alignas(TypeAlignment) Type : public ExtQualsTypeCommonBase {
   bool isSignableType() const;
   bool isAnyPointerType() const;   // Any C pointer or ObjC object pointer
   bool isCountAttributedType() const;
+  bool isCFIUncheckedCalleeType() const;
+  bool isPointerToCFIUncheckedCalleeType() const;
   bool isBlockPointerType() const;
   bool isVoidPointerType() const;
   bool isReferenceType() const;
@@ -3060,6 +3062,10 @@ template <> const TemplateSpecializationType *Type::getAs() const;
 /// until it reaches an AttributedType or a non-sugared type.
 template <> const AttributedType *Type::getAs() const;
 
+/// This will check for an CFIUncheckedCalleeType by removing any existing sugar
+/// until it reaches an CFIUncheckedCalleeType or a non-sugared type.
+template <> const CFIUncheckedCalleeType *Type::getAs() const;
+
 /// This will check for a BoundsAttributedType by removing any existing
 /// sugar until it reaches an BoundsAttributedType or a non-sugared type.
 template <> const BoundsAttributedType *Type::getAs() const;
@@ -3349,6 +3355,30 @@ class BoundsAttributedType : public Type, public llvm::FoldingSetNode {
   }
 };
 
+class CFIUncheckedCalleeType : public Type, public llvm::FoldingSetNode {
+public:
+  bool isSugared() const { return true; }
+  QualType desugar() const { return WrappedTy; }
+
+  QualType getWrappedType() const { return WrappedTy; }
+
+  static bool classof(const Type *T) {
+    return T->getTypeClass() == CFIUncheckedCallee;
+  }
+
+  static void Profile(llvm::FoldingSetNodeID &ID, QualType WrappedTy);
+
+  void Profile(llvm::FoldingSetNodeID &ID) { Profile(ID, WrappedTy); }
+
+protected:
+  friend class ASTContext;
+
+  CFIUncheckedCalleeType(QualType Wrapped, QualType Canon);
+
+private:
+  QualType WrappedTy;
+};
+
 /// Represents a sugar type with `__counted_by` or `__sized_by` annotations,
 /// including their `_or_null` variants.
 class CountAttributedType final
@@ -8244,6 +8274,20 @@ inline bool Type::isObjectPointerType() const {
     return false;
 }
 
+inline bool Type::isCFIUncheckedCalleeType() const {
+  return getAs<CFIUncheckedCalleeType>();
+}
+
+inline bool Type::isPointerToCFIUncheckedCalleeType() const {
+  if (const auto *MFT = getAs<MemberPointerType>()) {
+    if (MFT->isMemberFunctionPointer())
+      return MFT->getPointeeType()->getAs<CFIUncheckedCalleeType>();
+  }
+  if (const auto *T = getAs<PointerType>())
+    return T->getPointeeType()->getAs<CFIUncheckedCalleeType>();
+  return false;
+}
+
 inline bool Type::isFunctionPointerType() const {
   if (const auto *T = getAs<PointerType>())
     return T->getPointeeType()->isFunctionType();
@@ -8799,6 +8843,8 @@ template <typename T> const T *Type::getAsAdjusted() const {
       Ty = A->desugar().getTypePtr();
     else if (const auto *M = dyn_cast<MacroQualifiedType>(Ty))
       Ty = M->desugar().getTypePtr();
+    else if (const auto *M = dyn_cast<CFIUncheckedCalleeType>(Ty))
+      Ty = M->desugar().getTypePtr();
     else
       break;
   }

clang/include/clang/AST/TypeLoc.h

@@ -973,6 +973,31 @@ class HLSLAttributedResourceTypeLoc
   }
 };
 
+struct CFIUncheckedCalleeLocInfo {
+  SourceLocation AttrLoc;
+};
+
+/// Type source information for the CFIUncheckedCallee Type
+class CFIUncheckedCalleeTypeLoc
+    : public ConcreteTypeLoc<UnqualTypeLoc, CFIUncheckedCalleeTypeLoc,
+                             CFIUncheckedCalleeType,
+                             CFIUncheckedCalleeLocInfo> {
+public:
+  TypeLoc getWrappedLoc() const { return getInnerTypeLoc(); }
+
+  void setAttrLoc(const SourceLocation &Loc) { getLocalData()->AttrLoc = Loc; }
+  SourceLocation getAttrLoc() const { return getLocalData()->AttrLoc; }
+  void initializeLocal(ASTContext &Context, SourceLocation Loc) {
+    setAttrLoc(Loc);
+  }
+  QualType getWrappedType() const { return getTypePtr()->getWrappedType(); }
+  QualType getInnerType() const { return getTypePtr()->getWrappedType(); }
+
+  SourceRange getLocalSourceRange() const {
+    return getInnerTypeLoc().getLocalSourceRange();
+  }
+};
+
 struct ObjCObjectTypeLocInfo {
   SourceLocation TypeArgsLAngleLoc;
   SourceLocation TypeArgsRAngleLoc;
@@ -2736,6 +2761,8 @@ inline T TypeLoc::getAsAdjusted() const {
       Cur = ATL.getOriginalLoc();
     else if (auto MQL = Cur.getAs<MacroQualifiedTypeLoc>())
       Cur = MQL.getInnerLoc();
+    else if (auto CFIUCL = Cur.getAs<CFIUncheckedCalleeTypeLoc>())
+      Cur = CFIUCL.getWrappedLoc();
     else
       break;
   }

clang/include/clang/AST/TypeProperties.td

@@ -633,6 +633,15 @@ let Class = ParenType in {
   }]>;
 }
 
+let Class = CFIUncheckedCalleeType in {
+  def : Property<"WrappedTy", QualType> {
+    let Read = [{ node->getWrappedType() }];
+  }
+  def : Creator<[{
+    return ctx.getCFIUncheckedCalleeType(WrappedTy);
+  }]>;
+}
+
 let Class = MacroQualifiedType in {
   def : Property<"underlyingType", QualType> {
     let Read = [{ node->getUnderlyingType() }];

clang/include/clang/Basic/Attr.td

@@ -3082,6 +3082,11 @@ def NoDeref : TypeAttr {
   let Documentation = [NoDerefDocs];
 }
 
+def CFIUncheckedCallee : TypeAttr {
+  let Spellings = [Clang<"cfi_unchecked_callee">];
+  let Documentation = [CFIUncheckedCalleeDocs];
+}
+
 def ReqdWorkGroupSize : InheritableAttr {
   // Does not have a [[]] spelling because it is an OpenCL-related attribute.
   let Spellings = [GNU<"reqd_work_group_size">];

clang/include/clang/Basic/AttrDocs.td

@@ -6869,6 +6869,54 @@ for references or Objective-C object pointers.
 }];
 }
 
+def CFIUncheckedCalleeDocs : Documentation {
+  let Category = DocCatType;
+  let Content = [{
+``cfi_unchecked_callee`` is a function type attribute which prevents the compiler from instrumenting
+`Control Flow Integrity <https://clang.llvm.org/docs/ControlFlowIntegrity.html>`_ checks on indirect
+function calls. Specifically, the attribute has the following semantics:
+
+1. Indirect calls to a function type with this attribute will not be instrumented with CFI. That is,
+   the indirect call will not be checked. Note that this only changes the behavior for indirect calls
+   on pointers to function types having this attribute. It does not prevent all indirect function calls
+   for a given type from being checked.
+2. All direct references to a function whose type has this attribute will always reference the
+   function definition rather than an entry in the CFI jump table.
+3. When a pointer to a function with this attribute is implicitly cast to a pointer to a function
+   without this attribute, the compiler will give a warning saying this attribute is discarded. This
+   warning can be silenced with an explicit cast. Note an explicit cast just disables the warning, so
+   direct references to a function with a ``cfi_unchecked_callee`` attribute will still reference the
+   function definition rather than the CFI jump table.
+
+.. code-block:: c
+
+  #define CFI_UNCHECKED_CALLEE __attribute__((cfi_unchecked_callee))
+
+  void no_cfi() CFI_UNCHECKED_CALLEE {}
+
+  void (*with_cfi)() = no_cfi;  // warning: implicit conversion discards `cfi_unchecked_callee` attribute.
+                                // `with_cfi` also points to the actual definition of `no_cfi` rather than
+                                // its jump table entry.
+
+  void invoke(void (CFI_UNCHECKED_CALLEE *func)()) {
+    func();  // CFI will not instrument this indirect call.
+
+    void (*func2)() = func;  // warning: implicit conversion discards `cfi_unchecked_callee` attribute.
+
+    func2();  // CFI will instrument this indirect call. Users should be careful however because if this
+              // references a function with type `cfi_unchecked_callee`, then the CFI check may incorrectly
+              // fail because the reference will be to the function definition rather than the CFI jump
+              // table entry.
+  }
+
+This attribute can only be applied on functions or member functions. This attribute can be a good
+alternative to ``no_sanitize("cfi")`` if you only want to disable innstrumentation for specific indirect
+calls rather than applying ``no_sanitize("cfi")`` on the whole function containing indirect call. Note
+that ``cfi_unchecked_attribute`` is a type attribute doesn't disable CFI instrumentation on a function
+body.
+}];
+}
+
 def ReinitializesDocs : Documentation {
   let Category = DocCatFunction;
   let Content = [{

clang/include/clang/Basic/DiagnosticGroups.td

@@ -1576,6 +1576,8 @@ def FunctionMultiVersioning
 
 def NoDeref : DiagGroup<"noderef">;
 
+def CFIUncheckedCallee : DiagGroup<"cfi-unchecked-callee">;
+
 // -fbounds-safety and bounds annotation related warnings
 def BoundsSafetyCountedByEltTyUnknownSize :
   DiagGroup<"bounds-safety-counted-by-elt-type-unknown-size">;

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -12628,6 +12628,15 @@ def warn_noderef_on_non_pointer_or_array : Warning<
 def warn_noderef_to_dereferenceable_pointer : Warning<
   "casting to dereferenceable pointer removes 'noderef' attribute">, InGroup<NoDeref>;
 
+def warn_cfi_unchecked_callee_on_non_function
+    : Warning<"use of `cfi_unchecked_callee` on %0; can only be used on "
+              "function types">,
+      InGroup<CFIUncheckedCallee>;
+def warn_cast_discards_cfi_unchecked_callee
+    : Warning<"implicit conversion from %0 to %1 discards "
+              "`cfi_unchecked_callee` attribute">,
+      InGroup<CFIUncheckedCallee>;
+
 def err_builtin_launder_invalid_arg : Error<
   "%select{non-pointer|function pointer|void pointer}0 argument to "
   "'__builtin_launder' is not allowed">;