[ubsan] Add a nullability sanitizer

[vsk]

Summary

Teach UBSan how to detect violations of the _Nonnull annotation when
passing arguments to callees, in assignments, and in return stmts.

Because _Nonnull does not affect IRGen, the new checks are disabled by
default. The new driver flags are:

-fsanitize=nullability-arg (_Nonnull violation in call)
-fsanitize=nullability-assign (_Nonnull violation in assignment)
-fsanitize=nullability-return (_Nonnull violation in return stmt)
-fsanitize=nullability (all of the above)

This patch builds on top of UBSan's existing support for detecting
violations of the nonnull attributes ('nonnull' and 'returns_nonnull').
I am reusing the compiler-rt support for the existing checks for now.
I have FIXME's for this, and plan on handling them in a follow-up.

One point of note is that the nullability-return check is only allowed
to kick in if all arguments to the function satisfy their nullability
preconditions. This makes it necessary to emit some null checks in the
function body itself.

Testing: check-clang and check-ubsan. I also built some Apple ObjC
frameworks with an asserts-enabled compiler, and verified that we get
valid reports.

[aprantl]

docs/UndefinedBehaviorSanitizer.rst

98

s/a/an/

99

Can you rephrase this without using 2nd person?

lib/CodeGen/CGCall.cpp

2926

This function could use some comments.

3319

comment what this early return means?

lib/CodeGen/CGDecl.cpp

676

Comment what the check is doing?

1913

Is this a typo, or do you. really want Nullability here? Asking because everywhere else there is a check for ! Nullability && *Nullability == NullabilityKind::NonNull.
Ans should this condition be factored out?

lib/CodeGen/CodeGenFunction.h

1766

The \brief is redundant. We run doxygen with autobrief.

3474

ditto

[vsk]

• Improve the wording of the docs.
• Drop a weak test from 'ubsan-null-retval.m'.

[jroelofs]

lib/CodeGen/CGDecl.cpp

1913

Could just stick if (auto Nullability = in there, and remove the Nullability &&.

lib/CodeGen/CodeGenFunction.cpp

829

Hasn't the if (auto Nullability = already checked the truthiness of the var, making the Nullability && part of the second check redundant?

test/CodeGenObjC/ubsan-null-retval.m

14 ↗

(On Diff #91100)

The:

alloca
store
load

part of this looks like it's just making extra work for mem2reg. Is the alloca actually necessary?

[vsk]

Thanks for your comments, and sorry for jumping the gun earlier with an updated diff. I'll attach a fixed-up diff shortly.

lib/CodeGen/CGDecl.cpp

1913

'Nullability' should not have been checked for truthiness twice. By factoring out, do you mean moving the check's logic out of EmitParmDecl and into a helper? If so, I'm not sure if that would make the code cleaner, but I'm open to it.

test/CodeGenObjC/ubsan-null-retval.m

14 ↗

(On Diff #91100)

Clang likes its allocas :). This is typical -O0 behavior: IIUC it simplifies IRGen. E.g for expressions which take the address of an argument, because there is guaranteed storage for the arg.

[vsk]

• Add a test for the mixed _Nonnull arg + attribute((nonnull)) arg case.
• Reword docs per Adrian's comments, fix up doxygen comments, add better code comments, drop redundant "Nullability" truthiness checks.

[zaks.anna]

docs/UndefinedBehaviorSanitizer.rst

101

Have you investigated if -fsanitize=nullability-arg is too noisy? I think we should somehow "recommend" return and assignment check to most users. This is not clear from the help and options provided. The main concern here is that someone would try -fsanitize=nullability, see a ton of warnings they don't care about and think that the whole check is too noisy.

I cannot come up with a great solution. Options that come into mind are:

• Drop "arg" check completely.
• Remove arg from the nullability group.
• Have 2 groups for nullability.

We can also wait for concrete data before making a decision here and address this in subsequent commits.

lib/CodeGen/CGDecl.cpp

1909

I would add a comment explaining why this is needed, similar to what you included in the commit message:
"One point of note is that the nullability-return check is only allowed
to kick in if all arguments to the function satisfy their nullability
preconditions. This makes it necessary to emit some null checks in the
function body itself."

Maybe even rename CanCheckRetValNullability into RetValNullabilityPrecondition. I like this more than "CanCheck" because it's just a precondition value. You check if it is satisfied (evaluates to true or false) later when it's used in a branch.

[filcab]

Please make the tests tighter using CHECK-NEXT when possible. Much easier if later anyone needs to debug differences in IR.

docs/UndefinedBehaviorSanitizer.rst

102

This is weird. Please just document each of the nullability-* in isolation. None of the other flags is documenting more than one flag per bullet point, and I think splitting them up here is easy. It also makes it easier to read.
You already added the nullability group below, which mentions it adds the three checks.

141

... and the ``nullability`` group.

test/CodeGenObjC/ubsan-nullability.m

114

CHECK-NEXT?

[vsk]

One question for Anna (inline). I will update the diff with the documentation/code comments/renaming fixes once I hear back. Thanks again for the comments!

docs/UndefinedBehaviorSanitizer.rst

101

The nullability-arg check was not noisy on a small sample (n=4) of Apple-internal projects. I only collected 11 diagnostics in total, all of which seemed actionable. I.e the issues were not isolated to system headers: the project code actually did violate arg nullability. Based on what I've seen so far, I'd like to start off with including the arg check in a "nullability-full" group. If arg check noisiness becomes an issue, then I'd like to create a new "nullability-partial" group with just the return/assignment checks.

TL;DR: let's rename -fsanitize=nullability to -fsanitize=nullability-full, keep the arg check, and introduce a nullability-partial group later if we need to. Wdyt?

[zaks.anna]

docs/UndefinedBehaviorSanitizer.rst

101

Keeping things as is and introducing a new group if this becomes a problem sounds good to me. You could use "nullability" and "nullability-pedantic". That way you do not need to change the name and get to use "nullability".

[vsk]

The plan is to start off with -fsanitize=nullability, and then create a nullability-pedantic group later if it's really necessary. I think I've addressed all of the inline comments, and will upload a new diff shortly.

docs/UndefinedBehaviorSanitizer.rst

101

Sgtm.

141

I define the group after the 'undefined' group, so I used a slightly different spelling.

lib/CodeGen/CGDecl.cpp

1909

I added the comment and did the rename. I like 'RetValNullabilityPrecondition'.

test/CodeGenObjC/ubsan-nullability.m

114

I've promoted all the existing 'CHECK' lines to 'CHECK-NEXT's where possible.

[vsk]

• Rework documentation, add better code comments, and tighten up some check lines.

docs/UndefinedBehaviorSanitizer.rst

-  ``-fsanitize=float-divide-by-zero``: Floating point division by
      zero.
   -  ``-fsanitize=function``: Indirect call of a function through a
      function pointer of the wrong type (Linux, C++ and x86/x86_64 only).
   -  ``-fsanitize=integer-divide-by-zero``: Integer division by zero.
   -  ``-fsanitize=nonnull-attribute``: Passing null pointer as a function
      parameter which is declared to never be null.
   -  ``-fsanitize=null``: Use of a null pointer or creation of a null
      reference.
+  -  ``-fsanitize=nullability``: Passing null to a ``_Nonnull`` function
+     parameter, assigning null to to an lvalue marked ``_Nonnull``, or
+     returning null from a function with a ``_Nonnull`` return type.
+     ``-fsanitize=nullability-return`` enables just the return check,
+     ``-fsanitize=nullability-assign`` enables just the assignment check, and
+     ``-fsanitize=nullability-arg`` enables just the argument check. While
+     violating nullability rules does not result in undefined behavior, it
+     is often unintentional, so UBSan offers to catch it.
   -  ``-fsanitize=object-size``: An attempt to potentially use bytes which
      the optimizer can determine are not part of the object being accessed.
      This will also detect some types of undefined behavior that may not
      directly access memory, but are provably incorrect given the size of
      the objects involved, such as invalid downcasts and calling methods on
      invalid pointers. These checks are made in terms of
      ``__builtin_object_size``, and consequently may be able to detect more
      problems at higher optimization levels.
Available checks are:
   -  ``-fsanitize=vptr``: Use of an object whose vptr indicates that
      it is of the wrong dynamic type, or that its lifetime has not
      begun or has ended. Incompatible with ``-fno-rtti``. Link must
      be performed by ``clang++``, not ``clang``, to make sure C++-specific
      parts of the runtime library and C++ standard libraries are present.
 
 You can also use the following check groups:
   -  ``-fsanitize=undefined``: All of the checks listed above other than
-     ``unsigned-integer-overflow``.
+     ``unsigned-integer-overflow`` and ``nullability``.
   -  ``-fsanitize=undefined-trap``: Deprecated alias of
      ``-fsanitize=undefined``.
   -  ``-fsanitize=integer``: Checks for undefined or suspicious integer
      behavior (e.g. unsigned integer overflow).
+  -  ``-fsanitize=nullability``: Enables ``-fsanitize=nullability-arg``,
+     ``-fsanitize=nullability-assign``, and ``-fsanitize=nullability-return``.
 
 Stack traces and report symbolization
 =====================================
 If you want UBSan to print symbolized stack trace for each error report, you
 will need to:
 
 #. Compile with ``-g`` and ``-fno-omit-frame-pointer`` to get proper debug
    information in your binary.

include/clang/Basic/Sanitizers.def

 SANITIZER("bool", Bool)
 SANITIZER("enum", Enum)
 SANITIZER("float-cast-overflow", FloatCastOverflow)
 SANITIZER("float-divide-by-zero", FloatDivideByZero)
 SANITIZER("function", Function)
 SANITIZER("integer-divide-by-zero", IntegerDivideByZero)
 SANITIZER("nonnull-attribute", NonnullAttribute)
 SANITIZER("null", Null)
+SANITIZER("nullability-arg", NullabilityArg)
+SANITIZER("nullability-assign", NullabilityAssign)
+SANITIZER("nullability-return", NullabilityReturn)
+SANITIZER_GROUP("nullability", Nullability,
+                NullabilityArg | NullabilityAssign | NullabilityReturn)
 SANITIZER("object-size", ObjectSize)
 SANITIZER("return", Return)
 SANITIZER("returns-nonnull-attribute", ReturnsNonnullAttribute)
 SANITIZER("shift-base", ShiftBase)
 SANITIZER("shift-exponent", ShiftExponent)
 SANITIZER_GROUP("shift", Shift, ShiftBase | ShiftExponent)
 SANITIZER("signed-integer-overflow", SignedIntegerOverflow)
 SANITIZER("unreachable", Unreachable)

lib/CodeGen/CGCall.cpp

#endif
   }
 
   case ABIArgInfo::Expand:
     llvm_unreachable("Invalid ABI kind for return argument");
   }
 
   llvm::Instruction *Ret;
   if (RV) {
-    if (CurCodeDecl && SanOpts.has(SanitizerKind::ReturnsNonnullAttribute)) {
-      if (auto RetNNAttr = CurCodeDecl->getAttr<ReturnsNonNullAttr>()) {
-        SanitizerScope SanScope(this);
-        llvm::Value *Cond = Builder.CreateICmpNE(
-            RV, llvm::Constant::getNullValue(RV->getType()));
-        llvm::Constant *StaticData[] = {
-            EmitCheckSourceLocation(EndLoc),
-            EmitCheckSourceLocation(RetNNAttr->getLocation()),
-        };
-        EmitCheck(std::make_pair(Cond, SanitizerKind::ReturnsNonnullAttribute),
-                  SanitizerHandler::NonnullReturn, StaticData, None);
-      }
-    }
+    EmitReturnValueCheck(RV, EndLoc);
     Ret = Builder.CreateRet(RV);
   } else {
     Ret = Builder.CreateRetVoid();
   }
 
   if (RetDbgLoc)
     Ret->setDebugLoc(std::move(RetDbgLoc));
 }
 
+void CodeGenFunction::EmitReturnValueCheck(llvm::Value *RV,
+                                           SourceLocation EndLoc) {
+  // A current decl may not be available when emitting vtable thunks.
+  if (!CurCodeDecl)
+    return;
+
+  ReturnsNonNullAttr *RetNNAttr = nullptr;
+  if (SanOpts.has(SanitizerKind::ReturnsNonnullAttribute))
+    RetNNAttr = CurCodeDecl->getAttr<ReturnsNonNullAttr>();
+
+  if (!RetNNAttr && !requiresReturnValueNullabilityCheck())
+    return;
+
+  // Prefer the returns_nonnull attribute if it's present.
+  SourceLocation AttrLoc;
+  SanitizerMask CheckKind;
+  if (RetNNAttr) {
+    assert(!requiresReturnValueNullabilityCheck() &&
+           "Cannot check nullability and the nonnull attribute");
+    AttrLoc = RetNNAttr->getLocation();
+    CheckKind = SanitizerKind::ReturnsNonnullAttribute;
+  } else {
+    // FIXME: The runtime shouldn't refer to the 'returns_nonnull' attribute.
+    if (auto *DD = dyn_cast<DeclaratorDecl>(CurCodeDecl))
+      if (auto FTL =
+              DD->getTypeSourceInfo()->getTypeLoc().castAs<FunctionTypeLoc>())
+        AttrLoc = FTL.getReturnLoc().findNullabilityLoc();
+    CheckKind = SanitizerKind::NullabilityReturn;
+  }
+
+  SanitizerScope SanScope(this);
+
+  llvm::BasicBlock *Check = nullptr;
+  llvm::BasicBlock *NoCheck = nullptr;
+  if (requiresReturnValueNullabilityCheck()) {
+    // Before doing the nullability check, make sure that the preconditions for
+    // the check are met.
+    Check = createBasicBlock("nullcheck");
+    NoCheck = createBasicBlock("no.nullcheck");
+    Builder.CreateCondBr(CanCheckRetValNullability, Check, NoCheck);
+    EmitBlock(Check);
+  }
+
+  // Now do the null check. If the returns_nonnull attribute is present, this
+  // is done unconditionally.
+  llvm::Value *Cond = Builder.CreateIsNotNull(RV);
+  llvm::Constant *StaticData[] = {
+      EmitCheckSourceLocation(EndLoc), EmitCheckSourceLocation(AttrLoc),
+  };
+  EmitCheck(std::make_pair(Cond, CheckKind), SanitizerHandler::NonnullReturn,
+            StaticData, None);
+
+  if (requiresReturnValueNullabilityCheck())
+    EmitBlock(NoCheck);
+}
+
 static bool isInAllocaArgument(CGCXXABI &ABI, QualType type) {
   const CXXRecordDecl *RD = type->getAsCXXRecordDecl();
   return RD && ABI.getRecordArgABI(RD) == CGCXXABI::RAA_DirectInMemory;
 }
 
 static AggValueSlot createPlaceholderSlot(CodeGenFunction &CGF,
                                           QualType Ty) {
   // FIXME: Generate IR in one pass, rather than going back and fixing up these
if (StackBase) {
     CGF.Builder.CreateCall(F, StackBase);
   }
 }
 
 void CodeGenFunction::EmitNonNullArgCheck(RValue RV, QualType ArgType,
                                           SourceLocation ArgLoc,
                                           AbstractCallee AC,
                                           unsigned ParmNum) {
-  if (!SanOpts.has(SanitizerKind::NonnullAttribute) || !AC.getDecl())
+  if (!AC.getDecl() || !(SanOpts.has(SanitizerKind::NonnullAttribute) ||
+                         SanOpts.has(SanitizerKind::NullabilityArg)))
     return;
+
+  // The param decl may be missing in a variadic function.
   auto PVD = ParmNum < AC.getNumParams() ? AC.getParamDecl(ParmNum) : nullptr;
   unsigned ArgNo = PVD ? PVD->getFunctionScopeIndex() : ParmNum;
-  auto NNAttr = getNonNullAttr(AC.getDecl(), PVD, ArgType, ArgNo);
-  if (!NNAttr)
+
+  // Prefer the nonnull attribute if it's present. 
+  const NonNullAttr *NNAttr = nullptr;
+  if (SanOpts.has(SanitizerKind::NonnullAttribute))
+    NNAttr = getNonNullAttr(AC.getDecl(), PVD, ArgType, ArgNo);
+
+  bool CanCheckNullability = false;
+  if (SanOpts.has(SanitizerKind::NullabilityArg) && !NNAttr && PVD) {
+    auto Nullability = PVD->getType()->getNullability(getContext());
+    CanCheckNullability =
+        Nullability && *Nullability == NullabilityKind::NonNull;
+  }
+
+  if (!NNAttr && !CanCheckNullability)
     return;
+
+  SourceLocation AttrLoc;
+  SanitizerMask CheckKind;
+  if (NNAttr) {
+    AttrLoc = NNAttr->getLocation();
+    CheckKind = SanitizerKind::NonnullAttribute;
+  } else {
+    AttrLoc = PVD->getTypeSourceInfo()->getTypeLoc().findNullabilityLoc();
+    CheckKind = SanitizerKind::NullabilityArg;
+  }
+
   SanitizerScope SanScope(this);
   assert(RV.isScalar());
   llvm::Value *V = RV.getScalarVal();
   llvm::Value *Cond =
       Builder.CreateICmpNE(V, llvm::Constant::getNullValue(V->getType()));
   llvm::Constant *StaticData[] = {
-      EmitCheckSourceLocation(ArgLoc),
-      EmitCheckSourceLocation(NNAttr->getLocation()),
+      EmitCheckSourceLocation(ArgLoc), EmitCheckSourceLocation(AttrLoc),
       llvm::ConstantInt::get(Int32Ty, ArgNo + 1),
   };
-  EmitCheck(std::make_pair(Cond, SanitizerKind::NonnullAttribute),
-            SanitizerHandler::NonnullArg, StaticData, None);
+  EmitCheck(std::make_pair(Cond, CheckKind), SanitizerHandler::NonnullArg,
+            StaticData, None);
 }
 
 void CodeGenFunction::EmitCallArgs(
     CallArgList &Args, ArrayRef<QualType> ArgTypes,
     llvm::iterator_range<CallExpr::const_arg_iterator> ArgRange,
     AbstractCallee AC, unsigned ParamsToSkip, EvaluationOrder Order) {
   assert((int)ArgTypes.size() == (ArgRange.end() - ArgRange.begin()));
 

lib/CodeGen/CGDecl.cpp

 }
 
 static void drillIntoBlockVariable(CodeGenFunction &CGF,
                                    LValue &lvalue,
                                    const VarDecl *var) {
   lvalue.setAddress(CGF.emitBlockByrefAddress(lvalue.getAddress(), var));
 }
 
+void CodeGenFunction::EmitNullabilityCheck(LValue LHS, llvm::Value *RHS,
+                                           SourceLocation Loc) {
+  if (!SanOpts.has(SanitizerKind::NullabilityAssign))
+    return;
+
+  auto Nullability = LHS.getType()->getNullability(getContext());
+  if (!Nullability || *Nullability != NullabilityKind::NonNull)
+    return;
+
+  // Check if the right hand side of the assignment is nonnull, if the left
+  // hand side must be nonnull.
+  SanitizerScope SanScope(this);
+  llvm::Value *IsNotNull = Builder.CreateIsNotNull(RHS);
+  // FIXME: The runtime shouldn't refer to a 'reference'.
+  llvm::Constant *StaticData[] = {
+      EmitCheckSourceLocation(Loc), EmitCheckTypeDescriptor(LHS.getType()),
+      llvm::ConstantInt::get(Int8Ty, 1),
+      llvm::ConstantInt::get(Int8Ty, TCK_ReferenceBinding)};
+  EmitCheck({{IsNotNull, SanitizerKind::NullabilityAssign}},
+            SanitizerHandler::TypeMismatch, StaticData, RHS);
+}
+
 void CodeGenFunction::EmitScalarInit(const Expr *init, const ValueDecl *D,
                                      LValue lvalue, bool capturedByInit) {
   Qualifiers::ObjCLifetime lifetime = lvalue.getObjCLifetime();
   if (!lifetime) {
     llvm::Value *value = EmitScalarExpr(init);
     if (capturedByInit)
       drillIntoBlockVariable(*this, lvalue, cast<VarDecl>(D));
+    EmitNullabilityCheck(lvalue, value, init->getExprLoc());
     EmitStoreThroughLValue(RValue::get(value), lvalue, true);
     return;
   }
 
   if (const CXXDefaultInitExpr *DIE = dyn_cast<CXXDefaultInitExpr>(init))
     init = DIE->getExpr();
 
   // If we're emitting a value with lifetime, we have to do the
void CodeGenFunction::EmitScalarInit(const Expr *init, const ValueDecl *D,
 
   case Qualifiers::OCL_Autoreleasing:
     value = EmitARCRetainAutoreleaseScalarExpr(init);
     break;
   }
 
   if (capturedByInit) drillIntoBlockVariable(*this, lvalue, cast<VarDecl>(D));
 
+  EmitNullabilityCheck(lvalue, value, init->getExprLoc());
+
   // If the variable might have been accessed by its initializer, we
   // might have to initialize with a barrier.  We have to do this for
   // both __weak and __strong, but __weak got filtered out above.
   if (accessedByInit && lifetime == Qualifiers::OCL_Strong) {
     llvm::Value *oldValue = EmitLoadOfScalar(lvalue, init->getExprLoc());
     EmitStoreOfScalar(value, lvalue, /* isInitialization */ true);
     EmitARCRelease(oldValue, ARCImpreciseLifetime);
     return;
if (CGDebugInfo *DI = getDebugInfo()) {
     if (CGM.getCodeGenOpts().getDebugInfo() >=
         codegenoptions::LimitedDebugInfo) {
       DI->EmitDeclareOfArgVariable(&D, DeclPtr.getPointer(), ArgNo, Builder);
     }
   }
 
   if (D.hasAttr<AnnotateAttr>())
     EmitVarAnnotations(&D, DeclPtr.getPointer());
+
+  // Skip the return value nullability check if the nullability preconditions
+  // are broken.
+  if (requiresReturnValueNullabilityCheck()) {
+    auto Nullability = Ty->getNullability(getContext());
+    if (Nullability && *Nullability == NullabilityKind::NonNull) {
+      SanitizerScope SanScope(this);
+      CanCheckRetValNullability =
+          Builder.CreateAnd(CanCheckRetValNullability,
+                            Builder.CreateIsNotNull(Arg.getAnyValue()));
+    }
+  }
 }
 
 void CodeGenModule::EmitOMPDeclareReduction(const OMPDeclareReductionDecl *D,
                                             CodeGenFunction *CGF) {
   if (!LangOpts.OpenMP || (!LangOpts.EmitAllDecls && !D->isUsed()))
     return;
   getOpenMPRuntime().emitUserDefinedReduction(CGF, D);
 }

lib/CodeGen/CGExprScalar.cpp

case Qualifiers::OCL_None:
     // this should improve codegen just a little.
     RHS = Visit(E->getRHS());
     LHS = EmitCheckedLValue(E->getLHS(), CodeGenFunction::TCK_Store);
 
     // Store the value into the LHS.  Bit-fields are handled specially
     // because the result is altered by the store, i.e., [C99 6.5.16p1]
     // 'An assignment expression has the value of the left operand after
     // the assignment...'.
-    if (LHS.isBitField())
+    if (LHS.isBitField()) {
       CGF.EmitStoreThroughBitfieldLValue(RValue::get(RHS), LHS, &RHS);
-    else
+    } else {
+      CGF.EmitNullabilityCheck(LHS, RHS, E->getExprLoc());
       CGF.EmitStoreThroughLValue(RValue::get(RHS), LHS);
     }
+  }
 
   // If the result is clearly ignored, return now.
   if (Ignore)
     return nullptr;
 
   // The result of an assignment in C is the assigned r-value.
   if (!CGF.getLangOpts().CPlusPlus)
     return RHS;

lib/CodeGen/CodeGenFunction.h

private:
   /// The current source location that should be used for exception
   /// handling code.
   SourceLocation CurEHLocation;
 
   /// BlockByrefInfos - For each __block variable, contains
   /// information about the layout of the variable.
   llvm::DenseMap<const ValueDecl *, BlockByrefInfo> BlockByrefInfos;
 
+  /// Used by -fsanitize=nullability-return to determine whether the return
+  /// value can be checked.
+  llvm::Value *CanCheckRetValNullability = nullptr;
+
+  /// Check if -fsanitize=nullability-return instrumentation is required for
+  /// this function.
+  bool requiresReturnValueNullabilityCheck() const {
+    return CanCheckRetValNullability;
+  }
+
   llvm::BasicBlock *TerminateLandingPad;
   llvm::BasicBlock *TerminateHandler;
   llvm::BasicBlock *TrapBB;
 
   /// True if we need emit the life-time markers.
   const bool ShouldEmitLifetimeMarkers;
 
   /// Add a kernel metadata node to the named metadata node 'opencl.kernels'.
void EmitFunctionProlog(const CGFunctionInfo &FI,
                           llvm::Function *Fn,
                           const FunctionArgList &Args);
 
   /// EmitFunctionEpilog - Emit the target specific LLVM code to return the
   /// given temporary.
   void EmitFunctionEpilog(const CGFunctionInfo &FI, bool EmitRetDbgLoc,
                           SourceLocation EndLoc);
 
+  /// Emit a test that checks if the return value \p RV is nonnull.
+  void EmitReturnValueCheck(llvm::Value *RV, SourceLocation EndLoc);
+
   /// EmitStartEHSpec - Emit the start of the exception spec.
   void EmitStartEHSpec(const Decl *D);
 
   /// EmitEndEHSpec - Emit the end of the exception spec.
   void EmitEndEHSpec(const Decl *D);
 
   /// getTerminateLandingPad - Return a landing pad that just calls terminate.
   llvm::BasicBlock *getTerminateLandingPad();
public:
   /// EmitBranchOnBoolExpr - Emit a branch on a boolean condition (e.g. for an
   /// if statement) to the specified blocks.  Based on the condition, this might
   /// try to simplify the codegen of the conditional based on the branch.
   /// TrueCount should be the number of times we expect the condition to
   /// evaluate to true based on PGO data.
   void EmitBranchOnBoolExpr(const Expr *Cond, llvm::BasicBlock *TrueBlock,
                             llvm::BasicBlock *FalseBlock, uint64_t TrueCount);
 
+  /// Given an assignment `*LHS = RHS`, emit a test that checks if \p RHS is
+  /// nonnull, if \p LHS is marked _Nonnull.
+  void EmitNullabilityCheck(LValue LHS, llvm::Value *RHS, SourceLocation Loc);
+
   /// \brief Emit a description of a type in a format suitable for passing to
   /// a runtime sanitizer handler.
   llvm::Constant *EmitCheckTypeDescriptor(QualType T);
 
   /// \brief Convert a value into a format suitable for passing to a runtime
   /// sanitizer handler.
   llvm::Value *EmitCheckValue(llvm::Value *V);
 

lib/CodeGen/CodeGenFunction.cpp

if (const FunctionDecl *FD = dyn_cast_or_null<FunctionDecl>(D)) {
         llvm::Constant *PrologueStructElems[] = { PrologueSig, FTRTTIConst };
         llvm::Constant *PrologueStructConst =
             llvm::ConstantStruct::getAnon(PrologueStructElems, /*Packed=*/true);
         Fn->setPrologueData(PrologueStructConst);
       }
     }
   }
 
+  // If we're checking nullability, we need to know whether we can check the
+  // return value. Initialize the flag to 'true' and refine it in EmitParmDecl.
+  if (SanOpts.has(SanitizerKind::NullabilityReturn)) {
+    auto Nullability = FnRetTy->getNullability(getContext());
+    if (Nullability && *Nullability == NullabilityKind::NonNull) {
+      if (!(SanOpts.has(SanitizerKind::ReturnsNonnullAttribute) &&
+            CurCodeDecl && CurCodeDecl->getAttr<ReturnsNonNullAttr>()))
+        CanCheckRetValNullability =
+            llvm::ConstantInt::getTrue(getLLVMContext());
+    }
+  }
+
   // If we're in C++ mode and the function name is "main", it is guaranteed
   // to be norecurse by the standard (3.6.1.3 "The function main shall not be
   // used within a program").
   if (getLangOpts().CPlusPlus)
     if (const FunctionDecl *FD = dyn_cast_or_null<FunctionDecl>(D))
       if (FD->isMain())
         Fn->addFnAttr(llvm::Attribute::NoRecurse);
 

lib/Driver/SanitizerArgs.cpp

 #include <memory>
 
 using namespace clang;
 using namespace clang::SanitizerKind;
 using namespace clang::driver;
 using namespace llvm::opt;
 
 enum : SanitizerMask {
-  NeedsUbsanRt = Undefined | Integer | CFI,
+  NeedsUbsanRt = Undefined | Integer | Nullability | CFI,
   NeedsUbsanCxxRt = Vptr | CFI,
   NotAllowedWithTrap = Vptr,
   RequiresPIE = DataFlow,
   NeedsUnwindTables = Address | Thread | Memory | DataFlow,
-  SupportsCoverage = Address | Memory | Leak | Undefined | Integer | DataFlow,
-  RecoverableByDefault = Undefined | Integer,
+  SupportsCoverage =
+      Address | Memory | Leak | Undefined | Integer | Nullability | DataFlow,
+  RecoverableByDefault = Undefined | Integer | Nullability,
   Unrecoverable = Unreachable | Return,
   LegacyFsanitizeRecoverMask = Undefined | Integer,
   NeedsLTO = CFI,
-  TrappingSupported =
-      (Undefined & ~Vptr) | UnsignedIntegerOverflow | LocalBounds | CFI,
+  TrappingSupported = (Undefined & ~Vptr) | UnsignedIntegerOverflow |
+                      Nullability | LocalBounds | CFI,
   TrappingDefault = CFI,
   CFIClasses = CFIVCall | CFINVCall | CFIDerivedCast | CFIUnrelatedCast,
 };
 
 enum CoverageFeature {
   CoverageFunc = 1 << 0,
   CoverageBB = 1 << 1,
   CoverageEdge = 1 << 2,

lib/Driver/ToolChain.cpp

bool ToolChain::AddFastMathRuntimeIfAvailable(const ArgList &Args,
   return true;
 }
 
 SanitizerMask ToolChain::getSupportedSanitizers() const {
   // Return sanitizers which don't require runtime support and are not
   // platform dependent.
   using namespace SanitizerKind;
   SanitizerMask Res = (Undefined & ~Vptr & ~Function) | (CFI & ~CFIICall) |
-                      CFICastStrict | UnsignedIntegerOverflow | LocalBounds;
+                      CFICastStrict | UnsignedIntegerOverflow | Nullability |
+                      LocalBounds;
   if (getTriple().getArch() == llvm::Triple::x86 ||
       getTriple().getArch() == llvm::Triple::x86_64 ||
       getTriple().getArch() == llvm::Triple::arm ||
       getTriple().getArch() == llvm::Triple::aarch64 ||
       getTriple().getArch() == llvm::Triple::wasm32 ||
       getTriple().getArch() == llvm::Triple::wasm64)
     Res |= CFIICall;
   return Res;

test/CodeGenObjC/ubsan-nonnull-and-nullability.m

+// REQUIRES: asserts
+// RUN: %clang_cc1 -x objective-c -emit-llvm -triple x86_64-apple-macosx10.10.0 -fsanitize=nullability-return,returns-nonnull-attribute,nullability-arg,nonnull-attribute %s -o - -w | FileCheck %s
+
+// If both the annotation and the attribute are present, prefer the attribute,
+// since it actually affects IRGen.
+
+// CHECK-LABEL: define nonnull i32* @f1
+__attribute__((returns_nonnull)) int *_Nonnull f1(int *_Nonnull p) {
+  // CHECK: entry:
+  // CHECK-NEXT: [[ADDR:%.*]] = alloca i32*
+  // CHECK-NEXT: store i32* [[P:%.*]], i32** [[ADDR]]
+  // CHECK-NEXT: [[ARG:%.*]] = load i32*, i32** [[ADDR]]
+  // CHECK-NEXT: [[ICMP:%.*]] = icmp ne i32* [[ARG]], null, !nosanitize
+  // CHECK-NEXT: br i1 [[ICMP]], label %[[CONT:.+]], label %[[HANDLE:[^,]+]]
+  // CHECK: [[HANDLE]]:
+  // CHECK-NEXT:   call void @__ubsan_handle_nonnull_return_abort
+  // CHECK-NEXT:   unreachable, !nosanitize
+  // CHECK: [[CONT]]:
+  // CHECK-NEXT:   ret i32*
+  return p;
+}
+
+// CHECK-LABEL: define void @f2
+void f2(int *_Nonnull __attribute__((nonnull)) p) {}
+
+// CHECK-LABEL: define void @call_f2
+void call_f2() {
+  // CHECK: call void @__ubsan_handle_nonnull_arg_abort
+  // CHECK-NOT: call void @__ubsan_handle_nonnull_arg_abort
+  f2((void *)0);
+}

test/CodeGenObjC/ubsan-nullability.m

+// REQUIRES: asserts
+// RUN: %clang_cc1 -x objective-c -emit-llvm -triple x86_64-apple-macosx10.10.0 -fsanitize=nullability-arg,nullability-assign,nullability-return -w %s -o - | FileCheck %s
+
+// CHECK: [[NONNULL_RV_LOC1:@.*]] = private unnamed_addr global {{.*}} i32 109, i32 1 {{.*}} i32 100, i32 6
+// CHECK: [[NONNULL_ARG_LOC:@.*]] = private unnamed_addr global {{.*}} i32 204, i32 15 {{.*}} i32 190, i32 23
+// CHECK: [[NONNULL_ASSIGN1_LOC:@.*]] = private unnamed_addr global {{.*}} i32 305, i32 9
+// CHECK: [[NONNULL_ASSIGN2_LOC:@.*]] = private unnamed_addr global {{.*}} i32 405, i32 10
+// CHECK: [[NONNULL_ASSIGN3_LOC:@.*]] = private unnamed_addr global {{.*}} i32 505, i32 10
+// CHECK: [[NONNULL_INIT1_LOC:@.*]] = private unnamed_addr global {{.*}} i32 604, i32 25
+// CHECK: [[NONNULL_INIT2_LOC1:@.*]] = private unnamed_addr global {{.*}} i32 707, i32 26
+// CHECK: [[NONNULL_INIT2_LOC2:@.*]] = private unnamed_addr global {{.*}} i32 707, i32 29
+// CHECK: [[NONNULL_RV_LOC2:@.*]] = private unnamed_addr global {{.*}} i32 817, i32 1 {{.*}} i32 800, i32 6
+
+#define NULL ((void *)0)
+
+// CHECK-LABEL: define i32* @nonnull_retval1
+#line 100
+int *_Nonnull nonnull_retval1(int *p) {
+  // CHECK: br i1 true, label %[[NULL:.*]], label %[[NONULL:.*]], !nosanitize
+  // CHECK: [[NULL]]:
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* {{.*}}, null, !nosanitize
+  // CHECK-NEXT: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_nonnull_return{{.*}}[[NONNULL_RV_LOC1]]
+  return p;
+  // CHECK: [[NONULL]]:
+  // CHECK: ret i32*
+}
+
+#line 190
+void nonnull_arg(int *_Nonnull p) {}
+
+// CHECK-LABEL: define void @call_func_with_nonnull_arg
+#line 200
+void call_func_with_nonnull_arg(int *_Nonnull p) {
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* {{.*}}, null, !nosanitize
+  // CHECK-NEXT: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_nonnull_arg{{.*}}[[NONNULL_ARG_LOC]]
+  nonnull_arg(p);
+}
+
+// CHECK-LABEL: define void @nonnull_assign1
+#line 300
+void nonnull_assign1(int *p) {
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* {{.*}}, null, !nosanitize
+  // CHECK-NEXT: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_type_mismatch{{.*}}[[NONNULL_ASSIGN1_LOC]]
+  int *_Nonnull local;
+  local = p;
+}
+
+// CHECK-LABEL: define void @nonnull_assign2
+#line 400
+void nonnull_assign2(int *p) {
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* %{{.*}}, null, !nosanitize
+  // CHECK-NEXT: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_type_mismatch{{.*}}[[NONNULL_ASSIGN2_LOC]]
+  int *_Nonnull arr[1];
+  arr[0] = p;
+}
+
+struct S1 {
+  int *_Nonnull mptr;
+};
+
+// CHECK-LABEL: define void @nonnull_assign3
+#line 500
+void nonnull_assign3(int *p) {
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* %{{.*}}, null, !nosanitize
+  // CHECK-NEXT: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_type_mismatch{{.*}}[[NONNULL_ASSIGN3_LOC]]
+  struct S1 s;
+  s.mptr = p;
+}
+
+// CHECK-LABEL: define void @nonnull_init1
+#line 600
+void nonnull_init1(int *p) {
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* %{{.*}}, null, !nosanitize
+  // CHECK-NEXT: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_type_mismatch{{.*}}[[NONNULL_INIT1_LOC]]
+  int *_Nonnull local = p;
+}
+
+// CHECK-LABEL: define void @nonnull_init2
+#line 700
+void nonnull_init2(int *p) {
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* %{{.*}}, null, !nosanitize
+  // CHECK-NEXT: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_type_mismatch{{.*}}[[NONNULL_INIT2_LOC1]]
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* %{{.*}}, null, !nosanitize
+  // CHECK-NEXT: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_type_mismatch{{.*}}[[NONNULL_INIT2_LOC2]]
+  int *_Nonnull arr[] = {p, p};
+}
+
+// CHECK-LABEL: define i32* @nonnull_retval2
+#line 800
+int *_Nonnull nonnull_retval2(int *_Nonnull arg1,  //< Test this.
+                              int *_Nonnull arg2,  //< Test this.
+                              int *_Nullable arg3, //< Don't test the rest.
+                              int *arg4,
+                              int arg5, ...) {
+  // CHECK: [[ARG1CMP:%.*]] = icmp ne i32* %arg1, null, !nosanitize
+  // CHECK: [[DO_RV_CHECK_1:%.*]] = and i1 true, [[ARG1CMP]], !nosanitize
+  // CHECK: [[ARG2CMP:%.*]] = icmp ne i32* %arg2, null, !nosanitize
+  // CHECK: [[DO_RV_CHECK_2:%.*]] = and i1 [[DO_RV_CHECK_1]], [[ARG2CMP]]
+  // CHECK: br i1 [[DO_RV_CHECK_2]], label %[[NULL:.*]], label %[[NONULL:.*]], !nosanitize
+  // CHECK: [[NULL]]:
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* {{.*}}, null, !nosanitize
+  // CHECK-NEXT: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_nonnull_return{{.*}}[[NONNULL_RV_LOC2]]
+  return arg1;
+  // CHECK: [[NONULL]]:
+  // CHECK: ret i32*
+}
+
+@interface A
++(int *_Nonnull) objc_clsmethod: (int *_Nonnull) arg1;
+-(int *_Nonnull) objc_method: (int *_Nonnull) arg1;
+@end
+
+@implementation A
+
+// CHECK-LABEL: define internal i32* @"\01+[A objc_clsmethod:]"
++(int *_Nonnull) objc_clsmethod: (int *_Nonnull) arg1 {
+  // CHECK: [[ARG1CMP:%.*]] = icmp ne i32* %arg1, null, !nosanitize
+  // CHECK: [[DO_RV_CHECK:%.*]] = and i1 true, [[ARG1CMP]]
+  // CHECK: br i1 [[DO_RV_CHECK]], label %[[NULL:.*]], label %[[NONULL:.*]], !nosanitize
+  // CHECK: [[NULL]]:
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* {{.*}}, null, !nosanitize
+  // CHECK-NEXT: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_nonnull_return{{.*}}
+  return arg1;
+  // CHECK: [[NONULL]]:
+  // CHECK: ret i32*
+}
+
+// CHECK-LABEL: define internal i32* @"\01-[A objc_method:]"
+-(int *_Nonnull) objc_method: (int *_Nonnull) arg1 {
+  // CHECK: [[ARG1CMP:%.*]] = icmp ne i32* %arg1, null, !nosanitize
+  // CHECK: [[DO_RV_CHECK:%.*]] = and i1 true, [[ARG1CMP]]
+  // CHECK: br i1 [[DO_RV_CHECK]], label %[[NULL:.*]], label %[[NONULL:.*]], !nosanitize
+  // CHECK: [[NULL]]:
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* {{.*}}, null, !nosanitize
+  // CHECK-NEXT: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_nonnull_return{{.*}}
+  return arg1;
+  // CHECK: [[NONULL]]:
+  // CHECK: ret i32*
+}
+@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 i32* {{.*}} @objc_msgSend to i32* {{.*}}({{.*}}, i32* [[P1]])
+  [a objc_method: p];
+
+  // CHECK: [[ICMP:%.*]] = icmp ne i32* [[P2:%.*]], null, !nosanitize
+  // CHECK: br i1 [[ICMP]], {{.*}}, !nosanitize
+  // CHECK: call void @__ubsan_handle_nonnull_arg{{.*}} !nosanitize
+  // CHECK: call i32* {{.*}} @objc_msgSend to i32* {{.*}}({{.*}}, i32* [[P2]])
+  [A objc_clsmethod: p];
+}
+
+void dont_crash(int *_Nonnull p, ...) {}
+
+int main() {
+  nonnull_retval1(NULL);
+  nonnull_retval2(NULL, NULL, NULL, NULL, 0, 0, 0, 0);
+  call_func_with_nonnull_arg(NULL);
+  nonnull_assign1(NULL);
+  nonnull_assign2(NULL);
+  nonnull_assign3(NULL);
+  nonnull_init1(NULL);
+  nonnull_init2(NULL);
+  call_A(NULL, NULL);
+  dont_crash(NULL, NULL);
+  return 0;
+}