[ubsan] Extend the nonnull arg check to ObjC

UBSan's nonnull argument check applies when a parameter has the
"nonnull" attribute. The check currently works for FunctionDecls, but
not for ObjCMethodDecls. This patch extends the check to work for ObjC.

Differential Revision: https://reviews.llvm.org/D30599

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@296996 91177308-0d34-0410-b5e6-96231b3b80d8

lib/CodeGen/CGCall.cpp

@@ -3242,13 +3242,13 @@ void CallArgList::freeArgumentMemory(CodeGenFunction &CGF) const {
 
 void CodeGenFunction::EmitNonNullArgCheck(RValue RV, QualType ArgType,
                                           SourceLocation ArgLoc,
-                                          const FunctionDecl *FD,
+                                          AbstractCallee AC,
                                           unsigned ParmNum) {
-  if (!SanOpts.has(SanitizerKind::NonnullAttribute) || !FD)
+  if (!SanOpts.has(SanitizerKind::NonnullAttribute) || !AC.getDecl())
     return;
-  auto PVD = ParmNum < FD->getNumParams() ? FD->getParamDecl(ParmNum) : nullptr;
+  auto PVD = ParmNum < AC.getNumParams() ? AC.getParamDecl(ParmNum) : nullptr;
   unsigned ArgNo = PVD ? PVD->getFunctionScopeIndex() : ParmNum;
-  auto NNAttr = getNonNullAttr(FD, PVD, ArgType, ArgNo);
+  auto NNAttr = getNonNullAttr(AC.getDecl(), PVD, ArgType, ArgNo);
   if (!NNAttr)
     return;
   SanitizerScope SanScope(this);
@@ -3268,8 +3268,7 @@ void CodeGenFunction::EmitNonNullArgCheck(RValue RV, QualType ArgType,
 void CodeGenFunction::EmitCallArgs(
     CallArgList &Args, ArrayRef<QualType> ArgTypes,
     llvm::iterator_range<CallExpr::const_arg_iterator> ArgRange,
-    const FunctionDecl *CalleeDecl, unsigned ParamsToSkip,
-    EvaluationOrder Order) {
+    AbstractCallee AC, unsigned ParamsToSkip, EvaluationOrder Order) {
   assert((int)ArgTypes.size() == (ArgRange.end() - ArgRange.begin()));
 
   // We *have* to evaluate arguments from right to left in the MS C++ ABI,
@@ -3284,9 +3283,9 @@ void CodeGenFunction::EmitCallArgs(
 
   auto MaybeEmitImplicitObjectSize = [&](unsigned I, const Expr *Arg,
                                          RValue EmittedArg) {
-    if (CalleeDecl == nullptr || I >= CalleeDecl->getNumParams())
+    if (!AC.hasFunctionDecl() || I >= AC.getNumParams())
       return;
-    auto *PS = CalleeDecl->getParamDecl(I)->getAttr<PassObjectSizeAttr>();
+    auto *PS = AC.getParamDecl(I)->getAttr<PassObjectSizeAttr>();
     if (PS == nullptr)
       return;
 
@@ -3328,7 +3327,7 @@ void CodeGenFunction::EmitCallArgs(
            "The code below depends on only adding one arg per EmitCallArg");
     (void)InitialArgSize;
     RValue RVArg = Args.back().RV;
-    EmitNonNullArgCheck(RVArg, ArgTypes[Idx], (*Arg)->getExprLoc(), CalleeDecl,
+    EmitNonNullArgCheck(RVArg, ArgTypes[Idx], (*Arg)->getExprLoc(), AC,
                         ParamsToSkip + Idx);
     // @llvm.objectsize should never have side-effects and shouldn't need
     // destruction/cleanups, so we can safely "emit" it after its arg,

lib/CodeGen/CGExprCXX.cpp

@@ -1579,7 +1579,7 @@ llvm::Value *CodeGenFunction::EmitCXXNewExpr(const CXXNewExpr *E) {
 
     // FIXME: Why do we not pass a CalleeDecl here?
     EmitCallArgs(allocatorArgs, allocatorType, E->placement_arguments(),
-                 /*CalleeDecl*/nullptr, /*ParamsToSkip*/ParamsToSkip);
+                 /*AC*/AbstractCallee(), /*ParamsToSkip*/ParamsToSkip);
 
     RValue RV =
       EmitNewDeleteCall(*this, allocator, allocatorType, allocatorArgs);

lib/CodeGen/CGObjC.cpp

@@ -427,7 +427,7 @@ RValue CodeGenFunction::EmitObjCMessageExpr(const ObjCMessageExpr *E,
   QualType ResultType = method ? method->getReturnType() : E->getType();
 
   CallArgList Args;
-  EmitCallArgs(Args, method, E->arguments());
+  EmitCallArgs(Args, method, E->arguments(), /*AC*/AbstractCallee(method));
 
   // For delegate init calls in ARC, do an unsafe store of null into
   // self.  This represents the call taking direct ownership of that

lib/CodeGen/CodeGenFunction.h

@@ -305,6 +305,31 @@ public:
     ~CGCapturedStmtRAII() { CGF.CapturedStmtInfo = PrevCapturedStmtInfo; }
   };
 
+  /// An abstract representation of regular/ObjC call/message targets.
+  class AbstractCallee {
+    /// The function declaration of the callee.
+    const Decl *CalleeDecl;
+
+  public:
+    AbstractCallee() : CalleeDecl(nullptr) {}
+    AbstractCallee(const FunctionDecl *FD) : CalleeDecl(FD) {}
+    AbstractCallee(const ObjCMethodDecl *OMD) : CalleeDecl(OMD) {}
+    bool hasFunctionDecl() const {
+      return dyn_cast_or_null<FunctionDecl>(CalleeDecl);
+    }
+    const Decl *getDecl() const { return CalleeDecl; }
+    unsigned getNumParams() const {
+      if (const auto *FD = dyn_cast<FunctionDecl>(CalleeDecl))
+        return FD->getNumParams();
+      return cast<ObjCMethodDecl>(CalleeDecl)->param_size();
+    }
+    const ParmVarDecl *getParamDecl(unsigned I) const {
+      if (const auto *FD = dyn_cast<FunctionDecl>(CalleeDecl))
+        return FD->getParamDecl(I);
+      return *(cast<ObjCMethodDecl>(CalleeDecl)->param_begin() + I);
+    }
+  };
+
   /// \brief Sanitizers enabled for this function.
   SanitizerSet SanOpts;
 
@@ -3467,7 +3492,7 @@ public:
   /// \brief Create a check for a function parameter that may potentially be
   /// declared as non-null.
   void EmitNonNullArgCheck(RValue RV, QualType ArgType, SourceLocation ArgLoc,
-                           const FunctionDecl *FD, unsigned ParmNum);
+                           AbstractCallee AC, unsigned ParmNum);
 
   /// EmitCallArg - Emit a single call argument.
   void EmitCallArg(CallArgList &args, const Expr *E, QualType ArgType);
@@ -3569,7 +3594,7 @@ public:
   template <typename T>
   void EmitCallArgs(CallArgList &Args, const T *CallArgTypeInfo,
                     llvm::iterator_range<CallExpr::const_arg_iterator> ArgRange,
-                    const FunctionDecl *CalleeDecl = nullptr,
+                    AbstractCallee AC = AbstractCallee(),
                     unsigned ParamsToSkip = 0,
                     EvaluationOrder Order = EvaluationOrder::Default) {
     SmallVector<QualType, 16> ArgTypes;
@@ -3611,12 +3636,12 @@ public:
     for (auto *A : llvm::make_range(Arg, ArgRange.end()))
       ArgTypes.push_back(CallArgTypeInfo ? getVarArgType(A) : A->getType());
 
-    EmitCallArgs(Args, ArgTypes, ArgRange, CalleeDecl, ParamsToSkip, Order);
+    EmitCallArgs(Args, ArgTypes, ArgRange, AC, ParamsToSkip, Order);
   }
 
   void EmitCallArgs(CallArgList &Args, ArrayRef<QualType> ArgTypes,
                     llvm::iterator_range<CallExpr::const_arg_iterator> ArgRange,
-                    const FunctionDecl *CalleeDecl = nullptr,
+                    AbstractCallee AC = AbstractCallee(),
                     unsigned ParamsToSkip = 0,
                     EvaluationOrder Order = EvaluationOrder::Default);
 

test/CodeGenObjC/ubsan-nonnull.m

@@ -0,0 +1,48 @@
+// RUN: %clang_cc1 -x objective-c -emit-llvm -triple x86_64-apple-macosx10.10.0 -fsanitize=nonnull-attribute %s -o - -w | FileCheck %s
+
+@interface A
+
+-(void) one_arg: (__attribute__((nonnull)) int *) arg1;
+
+-(void) varargs: (__attribute__((nonnull)) int *) arg1, ...;
+
++(void) clsmethod: (__attribute__((nonnull)) int *) arg1;
+
+@end
+
+@implementation A
+
+// CHECK-LABEL: define internal void @"\01-[A one_arg:]"
+// CHECK-SAME: i32* nonnull
+-(void) one_arg: (__attribute__((nonnull)) int *) arg1 {}
+
+// CHECK-LABEL: define internal void @"\01-[A varargs:]"
+// CHECK-SAME: i32* nonnull
+-(void) varargs: (__attribute__((nonnull)) int *) arg1, ... {}
+
+// CHECK-LABEL: define internal void @"\01+[A clsmethod:]"
+// CHECK-SAME: i32* nonnull
++(void) clsmethod: (__attribute__((nonnull)) int *) arg1 {}
+
+@end
+
+// CHECK-LABEL: define void @call_A
+void call_A(A *a, int *p) {
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* [[P1:%.*]], null, !nosanitize
+  // CHECK: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_nonnull_arg{{.*}} !nosanitize
+  // CHECK: call void {{.*}} @objc_msgSend {{.*}} ({{.*}}, i32* [[P1]])
+  [a one_arg: p];
+
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* [[P2:%.*]], null, !nosanitize
+  // CHECK: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_nonnull_arg{{.*}} !nosanitize
+  // CHECK: call void {{.*}} @objc_msgSend {{.*}} ({{.*}}, i32* [[P2]], {{.*}})
+  [a varargs: p, p];
+
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* [[P3:%.*]], null, !nosanitize
+  // CHECK: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_nonnull_arg{{.*}} !nosanitize
+  // CHECK: call void {{.*}} @objc_msgSend {{.*}} ({{.*}}, i32* [[P3]])
+  [A clsmethod: p];
+}