[PATCH] D123309: [SafeStack] Don't create SCEV min between pointer and integer (PR54784)

[Nikita Popov via Phabricator via llvm-commits]

nikic created this revision.
nikic added reviewers: vitalybuka, tstellar, efriedma, fhahn.
Herald added a subscriber: hiraditya.
Herald added a project: All.
nikic requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

Rather than rewriting the alloca pointer to zero, use removePointerBase() to drop the base pointer. This will simply bail if the base pointer is not the alloca. We could try doing something more fancy here (like dropping the sources not based on the alloca on the premise that they aren't SafeStack-relevant or something), but I don't think that's worthwhile.

Fixes https://github.com/llvm/llvm-project/issues/54784.


https://reviews.llvm.org/D123309

Files:
  llvm/lib/CodeGen/SafeStack.cpp
  llvm/test/Transforms/SafeStack/pr54784.ll


Index: llvm/test/Transforms/SafeStack/pr54784.ll
===================================================================
--- /dev/null
+++ llvm/test/Transforms/SafeStack/pr54784.ll
@@ -0,0 +1,26 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -safe-stack < %s | FileCheck %s
+
+target triple = "x86_64-unknown-unknown"
+
+define void @min_of_pointers(i32* %p) safestack {
+; CHECK-LABEL: @min_of_pointers(
+; CHECK-NEXT:    [[UNSAFE_STACK_PTR:%.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr, align 8
+; CHECK-NEXT:    [[UNSAFE_STACK_STATIC_TOP:%.*]] = getelementptr i8, i8* [[UNSAFE_STACK_PTR]], i32 -16
+; CHECK-NEXT:    store i8* [[UNSAFE_STACK_STATIC_TOP]], i8** @__safestack_unsafe_stack_ptr, align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i8, i8* [[UNSAFE_STACK_PTR]], i32 -4
+; CHECK-NEXT:    [[A_UNSAFE1:%.*]] = bitcast i8* [[TMP1]] to i32*
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32* [[A_UNSAFE1]], [[P:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr i8, i8* [[UNSAFE_STACK_PTR]], i32 -4
+; CHECK-NEXT:    [[A_UNSAFE:%.*]] = bitcast i8* [[TMP2]] to i32*
+; CHECK-NEXT:    [[S:%.*]] = select i1 [[CMP]], i32* [[A_UNSAFE]], i32* [[P]]
+; CHECK-NEXT:    store i32 0, i32* [[S]], align 4
+; CHECK-NEXT:    store i8* [[UNSAFE_STACK_PTR]], i8** @__safestack_unsafe_stack_ptr, align 8
+; CHECK-NEXT:    ret void
+;
+  %a = alloca i32
+  %cmp = icmp ult i32* %a, %p
+  %s = select i1 %cmp, i32* %a, i32* %p
+  store i32 0, i32* %s
+  ret void
+}
Index: llvm/lib/CodeGen/SafeStack.cpp
===================================================================
--- llvm/lib/CodeGen/SafeStack.cpp
+++ llvm/lib/CodeGen/SafeStack.cpp
@@ -101,24 +101,6 @@
 
 namespace {
 
-/// Rewrite an SCEV expression for a memory access address to an expression that
-/// represents offset from the given alloca.
-///
-/// The implementation simply replaces all mentions of the alloca with zero.
-class AllocaOffsetRewriter : public SCEVRewriteVisitor<AllocaOffsetRewriter> {
-  const Value *AllocaPtr;
-
-public:
-  AllocaOffsetRewriter(ScalarEvolution &SE, const Value *AllocaPtr)
-      : SCEVRewriteVisitor(SE), AllocaPtr(AllocaPtr) {}
-
-  const SCEV *visitUnknown(const SCEVUnknown *Expr) {
-    if (Expr->getValue() == AllocaPtr)
-      return SE.getZero(Expr->getType());
-    return Expr;
-  }
-};
-
 /// The SafeStack pass splits the stack of each function into the safe
 /// stack, which is only accessed through memory safe dereferences (as
 /// determined statically), and the unsafe stack, which contains all
@@ -233,9 +215,18 @@
 
 bool SafeStack::IsAccessSafe(Value *Addr, uint64_t AccessSize,
                              const Value *AllocaPtr, uint64_t AllocaSize) {
-  AllocaOffsetRewriter Rewriter(SE, AllocaPtr);
-  const SCEV *Expr = Rewriter.visit(SE.getSCEV(Addr));
+  const SCEV *AddrExpr = SE.getSCEV(Addr);
+  const auto *Base = dyn_cast<SCEVUnknown>(SE.getPointerBase(AddrExpr));
+  if (!Base || Base->getValue() != AllocaPtr) {
+    LLVM_DEBUG(
+        dbgs() << "[SafeStack] "
+               << (isa<AllocaInst>(AllocaPtr) ? "Alloca " : "ByValArgument ")
+               << *AllocaPtr << "\n"
+               << "SCEV " << *AddrExpr << " not directly based on alloca\n");
+    return false;
+  }
 
+  const SCEV *Expr = SE.removePointerBase(AddrExpr);
   uint64_t BitWidth = SE.getTypeSizeInBits(Expr->getType());
   ConstantRange AccessStartRange = SE.getUnsignedRange(Expr);
   ConstantRange SizeRange =


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D123309.421184.patch
Type: text/x-patch
Size: 3476 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20220407/9172cdee/attachment.bin>