[sancov] Introduce optional callback for stack-depth tracking

Normally -fsanitize-coverage=stack-depth inserts inline arithmetic to
update thread_local __sancov_lowest_stack. To support stack depth
tracking in the Linux kernel, which does not implement traditional
thread_local storage, provide the option to call a function instead.

This matches the existing "stackleak" implementation that is supported
in Linux via a GCC plugin. To make this coverage more performant, a
minimum estimated stack depth can be chosen to enable the callback mode,
skipping instrumentation of functions with smaller stacks.

With -fsanitize-coverage-stack-depth-callback-min set greater than 0,
the __sanitize_cov_stack_depth() callback will be injected when the
estimated stack depth is greater than or equal to the given minimum.

Author: kees

clang/include/clang/Basic/CodeGenOptions.def

@@ -305,6 +305,7 @@ CODEGENOPT(SanitizeCoveragePCTable, 1, 0) ///< Create a PC Table.
 CODEGENOPT(SanitizeCoverageControlFlow, 1, 0) ///< Collect control flow
 CODEGENOPT(SanitizeCoverageNoPrune, 1, 0) ///< Disable coverage pruning.
 CODEGENOPT(SanitizeCoverageStackDepth, 1, 0) ///< Enable max stack depth tracing
+VALUE_CODEGENOPT(SanitizeCoverageStackDepthCallbackMin , 32, 0) ///< Enable stack depth tracing callbacks.
 CODEGENOPT(SanitizeCoverageTraceLoads, 1, 0) ///< Enable tracing of loads.
 CODEGENOPT(SanitizeCoverageTraceStores, 1, 0) ///< Enable tracing of stores.
 CODEGENOPT(SanitizeBinaryMetadataCovered, 1, 0) ///< Emit PCs for covered functions.

clang/include/clang/Driver/Options.td

@@ -2361,6 +2361,13 @@ def fsanitize_coverage_ignorelist : Joined<["-"], "fsanitize-coverage-ignorelist
     HelpText<"Disable sanitizer coverage instrumentation for modules and functions "
              "that match the provided special case list, even the allowed ones">,
     MarshallingInfoStringVector<CodeGenOpts<"SanitizeCoverageIgnorelistFiles">>;
+def fsanitize_coverage_stack_depth_callback_min_EQ
+    : Joined<["-"], "fsanitize-coverage-stack-depth-callback-min=">,
+      Group<f_clang_Group>,
+      MetaVarName<"<M>">,
+      HelpText<"Use callback for max stack depth tracing with minimum stack "
+               "depth M">,
+      MarshallingInfoInt<CodeGenOpts<"SanitizeCoverageStackDepthCallbackMin">>;
 def fexperimental_sanitize_metadata_EQ : CommaJoined<["-"], "fexperimental-sanitize-metadata=">,
   Group<f_Group>,
   HelpText<"Specify the type of metadata to emit for binary analysis sanitizers">;

clang/include/clang/Driver/SanitizerArgs.h

@@ -34,6 +34,7 @@ class SanitizerArgs {
   std::vector<std::string> CoverageIgnorelistFiles;
   std::vector<std::string> BinaryMetadataIgnorelistFiles;
   int CoverageFeatures = 0;
+  int CoverageStackDepthCallbackMin = 0;
   int BinaryMetadataFeatures = 0;
   int OverflowPatternExclusions = 0;
   int MsanTrackOrigins = 0;

clang/lib/CodeGen/BackendUtil.cpp

@@ -255,6 +255,7 @@ getSancovOptsFromCGOpts(const CodeGenOptions &CGOpts) {
   Opts.InlineBoolFlag = CGOpts.SanitizeCoverageInlineBoolFlag;
   Opts.PCTable = CGOpts.SanitizeCoveragePCTable;
   Opts.StackDepth = CGOpts.SanitizeCoverageStackDepth;
+  Opts.StackDepthCallbackMin = CGOpts.SanitizeCoverageStackDepthCallbackMin;
   Opts.TraceLoads = CGOpts.SanitizeCoverageTraceLoads;
   Opts.TraceStores = CGOpts.SanitizeCoverageTraceStores;
   Opts.CollectControlFlow = CGOpts.SanitizeCoverageControlFlow;

clang/lib/Driver/SanitizerArgs.cpp

@@ -751,6 +751,17 @@ SanitizerArgs::SanitizerArgs(const ToolChain &TC,
       options::OPT_fno_sanitize_ignorelist,
       clang::diag::err_drv_malformed_sanitizer_ignorelist, DiagnoseErrors);
 
+  // Verify that -fsanitize-coverage-stack-depth-callback-min is >= 0.
+  if (Arg *A = Args.getLastArg(
+          options::OPT_fsanitize_coverage_stack_depth_callback_min_EQ)) {
+    StringRef S = A->getValue();
+    if (S.getAsInteger(0, CoverageStackDepthCallbackMin) ||
+        CoverageStackDepthCallbackMin < 0) {
+      if (DiagnoseErrors)
+        D.Diag(clang::diag::err_drv_invalid_value) << A->getAsString(Args) << S;
+    }
+  }
+
   // Parse -f[no-]sanitize-memory-track-origins[=level] options.
   if (AllAddedKinds & SanitizerKind::Memory) {
     if (Arg *A =
@@ -1269,6 +1280,11 @@ void SanitizerArgs::addArgs(const ToolChain &TC, const llvm::opt::ArgList &Args,
   addSpecialCaseListOpt(Args, CmdArgs, "-fsanitize-coverage-ignorelist=",
                         CoverageIgnorelistFiles);
 
+  if (CoverageStackDepthCallbackMin)
+    CmdArgs.push_back(
+        Args.MakeArgString("-fsanitize-coverage-stack-depth-callback-min=" +
+                           Twine(CoverageStackDepthCallbackMin)));
+
   if (!GPUSanitize) {
     // Translate available BinaryMetadataFeatures to corresponding clang-cc1
     // flags. Does not depend on any other sanitizers. Unsupported on GPUs.

clang/test/Driver/fsanitize-coverage.c

@@ -93,6 +93,20 @@
 // CHECK-STACK-DEPTH-PC-GUARD: -fsanitize-coverage-trace-pc-guard
 // CHECK-STACK-DEPTH-PC-GUARD: -fsanitize-coverage-stack-depth
 
+// RUN: %clang --target=x86_64-linux-gnu \
+// RUN:     -fsanitize-coverage-stack-depth-callback-min=100 %s -### 2>&1 | \
+// RUN:     FileCheck %s --check-prefix=CHECK-STACK-DEPTH-CALLBACK
+// RUN: %clang --target=x86_64-linux-gnu \
+// RUN:     -fsanitize-coverage-stack-depth-callback-min=0 %s -### 2>&1 | \
+// RUN:     FileCheck %s --check-prefix=CHECK-STACK-DEPTH-CALLBACK-ZERO
+// RUN: not %clang --target=x86_64-linux-gnu \
+// RUN:     -fsanitize-coverage-stack-depth-callback-min=-10 %s -### 2>&1 | \
+// RUN:     FileCheck %s --check-prefix=CHECK-STACK-DEPTH-CALLBACK-NEGATIVE
+// CHECK-STACK-DEPTH-CALLBACK-NOT: error:
+// CHECK-STACK-DEPTH-CALLBACK: -fsanitize-coverage-stack-depth-callback-min=100
+// CHECK-STACK-DEPTH-CALLBACK-ZERO-NOT: -fsanitize-coverage-stack-depth-callback-min=0
+// CHECK-STACK-DEPTH-CALLBACK-NEGATIVE: error: invalid value '-10' in '-fsanitize-coverage-stack-depth-callback-min=-10'
+
 // RUN: %clang --target=x86_64-linux-gnu -fsanitize=address -fsanitize-coverage=trace-cmp,indirect-calls %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-NO-TYPE-NECESSARY
 // CHECK-NO-TYPE-NECESSARY-NOT: error:
 // CHECK-NO-TYPE-NECESSARY: -fsanitize-coverage-indirect-calls

llvm/include/llvm/Transforms/Utils/Instrumentation.h

@@ -158,6 +158,7 @@ struct SanitizerCoverageOptions {
   bool PCTable = false;
   bool NoPrune = false;
   bool StackDepth = false;
+  int StackDepthCallbackMin = 0;
   bool TraceLoads = false;
   bool TraceStores = false;
   bool CollectControlFlow = false;

llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp

@@ -86,6 +86,7 @@ const char SanCovPCsSectionName[] = "sancov_pcs";
 const char SanCovCFsSectionName[] = "sancov_cfs";
 const char SanCovCallbackGateSectionName[] = "sancov_gate";
 
+const char SanCovStackDepthCallbackName[] = "__sanitizer_cov_stack_depth";
 const char SanCovLowestStackName[] = "__sancov_lowest_stack";
 const char SanCovCallbackGateName[] = "__sancov_should_track";
 
@@ -152,6 +153,12 @@ static cl::opt<bool> ClStackDepth("sanitizer-coverage-stack-depth",
                                   cl::desc("max stack depth tracing"),
                                   cl::Hidden);
 
+static cl::opt<int> ClStackDepthCallbackMin(
+    "sanitizer-coverage-stack-depth-callback-min",
+    cl::desc("max stack depth tracing should use callback and only when "
+             "stack depth more than specified"),
+    cl::Hidden);
+
 static cl::opt<bool>
     ClCollectCF("sanitizer-coverage-control-flow",
                 cl::desc("collect control flow for each function"), cl::Hidden);
@@ -202,6 +209,8 @@ SanitizerCoverageOptions OverrideFromCL(SanitizerCoverageOptions Options) {
   Options.PCTable |= ClCreatePCTable;
   Options.NoPrune |= !ClPruneBlocks;
   Options.StackDepth |= ClStackDepth;
+  Options.StackDepthCallbackMin = std::max(Options.StackDepthCallbackMin,
+                                           ClStackDepthCallbackMin.getValue());
   Options.TraceLoads |= ClLoadTracing;
   Options.TraceStores |= ClStoreTracing;
   Options.GatedCallbacks |= ClGatedCallbacks;
@@ -271,6 +280,7 @@ class ModuleSanitizerCoverage {
   DomTreeCallback DTCallback;
   PostDomTreeCallback PDTCallback;
 
+  FunctionCallee SanCovStackDepthCallback;
   FunctionCallee SanCovTracePCIndir;
   FunctionCallee SanCovTracePC, SanCovTracePCGuard;
   std::array<FunctionCallee, 4> SanCovTraceCmpFunction;
@@ -514,6 +524,9 @@ bool ModuleSanitizerCoverage::instrumentModule() {
   SanCovTracePCGuard =
       M.getOrInsertFunction(SanCovTracePCGuardName, VoidTy, PtrTy);
 
+  SanCovStackDepthCallback =
+      M.getOrInsertFunction(SanCovStackDepthCallbackName, VoidTy);
+
   for (auto &F : M)
     instrumentFunction(F);
 
@@ -1029,6 +1042,8 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
   if (IsEntryBB) {
     if (auto SP = F.getSubprogram())
       EntryLoc = DILocation::get(SP->getContext(), SP->getScopeLine(), 0, SP);
+    // FIXME: stack-depth does not correctly instrument dynamic allocas.
+    //
     // Keep static allocas and llvm.localescape calls in the entry block.  Even
     // if we aren't splitting the block, it's nice for allocas to be before
     // calls.
@@ -1078,22 +1093,50 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
     Store->setNoSanitizeMetadata();
   }
   if (Options.StackDepth && IsEntryBB && !IsLeafFunc) {
-    // Check stack depth.  If it's the deepest so far, record it.
     Module *M = F.getParent();
-    auto FrameAddrPtr = IRB.CreateIntrinsic(
-        Intrinsic::frameaddress,
-        IRB.getPtrTy(M->getDataLayout().getAllocaAddrSpace()),
-        {Constant::getNullValue(Int32Ty)});
-    auto FrameAddrInt = IRB.CreatePtrToInt(FrameAddrPtr, IntptrTy);
-    auto LowestStack = IRB.CreateLoad(IntptrTy, SanCovLowestStack);
-    auto IsStackLower = IRB.CreateICmpULT(FrameAddrInt, LowestStack);
-    auto ThenTerm = SplitBlockAndInsertIfThen(
-        IsStackLower, &*IP, false,
-        MDBuilder(IRB.getContext()).createUnlikelyBranchWeights());
-    IRBuilder<> ThenIRB(ThenTerm);
-    auto Store = ThenIRB.CreateStore(FrameAddrInt, SanCovLowestStack);
-    LowestStack->setNoSanitizeMetadata();
-    Store->setNoSanitizeMetadata();
+    if (Options.StackDepthCallbackMin) {
+      // In callback mode, only add call when stack depth reaches minimum.
+      const DataLayout &DL = M->getDataLayout();
+      uint32_t EstimatedStackSize = 0;
+
+      // Make an estimate on the stack usage.
+      for (auto &I : BB) {
+        if (auto *AI = dyn_cast<AllocaInst>(&I)) {
+          if (AI->isStaticAlloca()) {
+            uint32_t Bytes = DL.getTypeAllocSize(AI->getAllocatedType());
+            if (AI->isArrayAllocation()) {
+              if (const ConstantInt *arraySize =
+                      dyn_cast<ConstantInt>(AI->getArraySize()))
+                Bytes *= arraySize->getZExtValue();
+            }
+            EstimatedStackSize += Bytes;
+          } else {
+            // Dynamic alloca: require we always perform callback.
+            EstimatedStackSize = Options.StackDepthCallbackMin;
+            break;
+          }
+        }
+      }
+
+      if (EstimatedStackSize >= Options.StackDepthCallbackMin)
+        IRB.CreateCall(SanCovStackDepthCallback)->setCannotMerge();
+    } else {
+      // Check stack depth.  If it's the deepest so far, record it.
+      auto FrameAddrPtr = IRB.CreateIntrinsic(
+          Intrinsic::frameaddress,
+          IRB.getPtrTy(M->getDataLayout().getAllocaAddrSpace()),
+          {Constant::getNullValue(Int32Ty)});
+      auto FrameAddrInt = IRB.CreatePtrToInt(FrameAddrPtr, IntptrTy);
+      auto LowestStack = IRB.CreateLoad(IntptrTy, SanCovLowestStack);
+      auto IsStackLower = IRB.CreateICmpULT(FrameAddrInt, LowestStack);
+      auto ThenTerm = SplitBlockAndInsertIfThen(
+          IsStackLower, &*IP, false,
+          MDBuilder(IRB.getContext()).createUnlikelyBranchWeights());
+      IRBuilder<> ThenIRB(ThenTerm);
+      auto Store = ThenIRB.CreateStore(FrameAddrInt, SanCovLowestStack);
+      LowestStack->setNoSanitizeMetadata();
+      Store->setNoSanitizeMetadata();
+    }
   }
 }