[clang] [llvm] [sancov] Introduce optional callback for stack-depth tracking (PR #138323)

[Kees Cook via llvm-commits]

================
@@ -1078,22 +1091,61 @@ 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();
+
+    // Find an insertion point after last "alloca".
+    llvm::Instruction *InsertBefore = NULL;
+    for (auto &I : BB) {
+      if (llvm::isa<llvm::AllocaInst>(I))
+        InsertBefore = I.getNextNode(); // Move past the "alloca".
+    }
+    // But only use it if we actually found an "alloca".
+    if (InsertBefore)
+      IRB.SetInsertPoint(InsertBefore);
+
+    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;
----------------
kees wrote:

I've added a more detailed comment to the accounting loop (and moved the insertion calculation there since it's only needed for the callback case). And I've tweaked the documentation a bit more with a short example.

https://github.com/llvm/llvm-project/pull/138323