[safestack] Sink unsafe address computation to each use.

[eugenis]

Summary

This is a fix for PR27844.
When replacing uses of unsafe allocas, emit the new location immediately after each use.
Without this, the pointer stays live from the function entry to the last use, while it's usually cheaper to recalculate.

[eugenis]

ping

[pcc]

lib/CodeGen/SafeStack.cpp

615

Why bother setting the name if we're about to remove the instruction?

638

This seems a little weird to me. I wouldn't expect anything to be holding a value handle here (and if it did, it probably wouldn't want the replacement for whichever instruction happened to be at the end of the use list) and the only important metadata is debug info, which is handled by replaceDbgDeclareForAlloca, I believe. Can we just remove this part?

[eugenis]

lib/CodeGen/SafeStack.cpp

615

Because otherwise CreateBitCast(..., Name) would de-duplicate the name and the new instruction would end up with something like Name.unsafe.12345.

638

The new debug metadata generated in replaceDbgDeclareForAlloca is still tied to the alloca that is being replaced. Without this line, it gets replaced with an empty metadata.

The condition (AI->hasValueHandle() || AI->isUsedByMetadata()) is taken from Value::replaceAllUsesWith, and it describes all cases where that function would do any work that is not yet done at this point (i.e. replacing the non-IR uses).

[pcc]

lib/CodeGen/SafeStack.cpp

615

Aren't you achieving that by appending the ".unsafe" suffix?

638

The new debug metadata generated in replaceDbgDeclareForAlloca is still tied to the alloca that is being replaced. Without this line, it gets replaced with an empty metadata.

That's surprising.
http://llvm-cs.pcc.me.uk/lib/Transforms/Utils/Local.cpp#1223
http://llvm-cs.pcc.me.uk/lib/Transforms/Utils/Local.cpp#1185
It certainly looks like the intent is to replace the address. Is there some case where we aren't replacing it correctly?

The condition (AI->hasValueHandle() || AI->isUsedByMetadata()) is taken from Value::replaceAllUsesWith, and it describes all cases where that function would do any work that is not yet done at this point (i.e. replacing the non-IR uses).

We shouldn't cargo cult that logic from Value::replaceAllUsesWith, as it only makes sense when the replacement is uniform.

[eugenis]

lib/CodeGen/SafeStack.cpp

615

Right. Will remove.

638

That case is llvm.dbg.value.

replaceDbgDeclareForAlloca does not touch dbg.value. They are updated by RAUW, but that does not give them the same "good" debug location as replaceDbgDeclareForAlloca does for dbg.declare.

I'll fix it in replaceDbgDeclareForAlloca

[eugenis]

Rebased on top of D20986

lib/CodeGen/SafeStack.cpp

615

Removed.

649

Fixed in D20986.

[eugenis]

ping

[pcc]

lib/CodeGen/SafeStack.cpp

613

Explain why in the comment.

631

I believe this can happen when a function has multiple identical control flow edges, e.g.

define void @main() {
entry:
  br i1 undef, label %x, label %x

x:
  %p = phi i32 [ 0, %entry ], [ 0, %entry ]
  ret void
}

Can you please add tests that cover phi nodes, including this case?

[pcc]

lib/CodeGen/SafeStack.cpp

613

Is that the right PR number?

[eugenis]

fix PR number in comments

[pcc]

LGTM

[eugenis]

r272969

lib/CodeGen/SafeStack.cpp

for (AllocaInst *AI : StaticAllocas) {
     unsigned Align =
         std::max((unsigned)DL->getPrefTypeAlignment(Ty), AI->getAlignment());
 
     // Add alignment.
     // NOTE: we ensure that BasePointer itself is aligned to >= Align.
     StaticOffset += Size;
     StaticOffset = alignTo(StaticOffset, Align);
 
-    Value *Off = IRB.CreateGEP(BasePointer, // BasePointer is i8*
-                               ConstantInt::get(Int32Ty, -StaticOffset));
-    Value *NewAI = IRB.CreateBitCast(Off, AI->getType(), AI->getName());
-    if (AI->hasName() && isa<Instruction>(NewAI))
-      cast<Instruction>(NewAI)->takeName(AI);
-
-    // Replace alloc with the new location.
     replaceDbgDeclareForAlloca(AI, BasePointer, DIB, /*Deref=*/true, -StaticOffset);
     replaceDbgValueForAlloca(AI, BasePointer, DIB, -StaticOffset);
-    AI->replaceAllUsesWith(NewAI);
+
+    // Replace uses of the alloca with the new location.
+    // Insert address calculation close to each use.
+    std::string Name = std::string(AI->getName()) + ".unsafe";
+    while (!AI->use_empty()) {
+      Use &U = *AI->use_begin();
+      Instruction *User = cast<Instruction>(U.getUser());
+
+      Instruction *InsertBefore;
+      if (auto *PHI = dyn_cast<PHINode>(User))
+        InsertBefore = PHI->getIncomingBlock(U)->getTerminator();
+      else
+        InsertBefore = User;
+
+      IRBuilder<> IRBUser(InsertBefore);
+      Value *Off = IRBUser.CreateGEP(BasePointer, // BasePointer is i8*
+                                     ConstantInt::get(Int32Ty, -StaticOffset));
+      Value *Replacement = IRBUser.CreateBitCast(Off, AI->getType(), Name);
+
+      if (auto *PHI = dyn_cast<PHINode>(User)) {
+        // PHI nodes may have multiple incoming edges from the same BB (why??),
+        // all must be updated at once with the same incoming value.
+        auto *BB = PHI->getIncomingBlock(U);
+        for (unsigned I = 0; I < PHI->getNumIncomingValues(); ++I)
+          if (PHI->getIncomingBlock(I) == BB)
+            PHI->setIncomingValue(I, Replacement);
+      } else {
+        U.set(Replacement);
+      }
+    }
+
     AI->eraseFromParent();
   }
 
   // Re-align BasePointer so that our callees would see it aligned as
   // expected.
   // FIXME: no need to update BasePointer in leaf functions.
   StaticOffset = alignTo(StaticOffset, StackAlignment);
 
   // Update shadow stack pointer in the function epilogue.
   IRB.SetInsertPoint(BasePointer->getNextNode());
 
   Value *StaticTop =
       IRB.CreateGEP(BasePointer, ConstantInt::get(Int32Ty, -StaticOffset),
                     "unsafe_stack_static_top");
   IRB.CreateStore(StaticTop, UnsafeStackPtr);
   return StaticTop;

test/Transforms/SafeStack/array-aligned.ll

 ; RUN: opt -safe-stack -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck %s
 ; RUN: opt -safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s
 
 @.str = private unnamed_addr constant [4 x i8] c"%s\0A\00", align 1
 
 ; array of [16 x i8]
 
 define void @foo(i8* %a) nounwind uwtable safestack {
 entry:
   ; CHECK: %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
 
   ; CHECK: %[[USST:.*]] = getelementptr i8, i8* %[[USP]], i32 -16
 
   ; CHECK: store i8* %[[USST]], i8** @__safestack_unsafe_stack_ptr
 
-  ; CHECK: %[[AADDR:.*]] = alloca i8*, align 8
   %a.addr = alloca i8*, align 8
-
-  ; CHECK: %[[BUFPTR:.*]] = getelementptr i8, i8* %[[USP]], i32 -16
-  ; CHECK: %[[BUFPTR2:.*]] = bitcast i8* %[[BUFPTR]] to [16 x i8]*
   %buf = alloca [16 x i8], align 16
 
+  ; CHECK: %[[AADDR:.*]] = alloca i8*, align 8
   ; CHECK: store i8* {{.*}}, i8** %[[AADDR]], align 8
   store i8* %a, i8** %a.addr, align 8
 
+  ; CHECK: %[[BUFPTR:.*]] = getelementptr i8, i8* %[[USP]], i32 -16
+  ; CHECK: %[[BUFPTR2:.*]] = bitcast i8* %[[BUFPTR]] to [16 x i8]*
   ; CHECK: %[[GEP:.*]] = getelementptr inbounds [16 x i8], [16 x i8]* %[[BUFPTR2]], i32 0, i32 0
   %gep = getelementptr inbounds [16 x i8], [16 x i8]* %buf, i32 0, i32 0
 
   ; CHECK: %[[A2:.*]] = load i8*, i8** %[[AADDR]], align 8
   %a2 = load i8*, i8** %a.addr, align 8
 
   ; CHECK: call i8* @strcpy(i8* %[[GEP]], i8* %[[A2]])
   %call = call i8* @strcpy(i8* %gep, i8* %a2)
 
   ; CHECK: store i8* %[[USP]], i8** @__safestack_unsafe_stack_ptr
   ret void
 }
 
 declare i8* @strcpy(i8*, i8*)

test/Transforms/SafeStack/array.ll

 define void @foo(i8* %a) nounwind uwtable safestack {
 entry:
   ; CHECK: %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
 
   ; CHECK: %[[USST:.*]] = getelementptr i8, i8* %[[USP]], i32 -16
 
   ; CHECK: store i8* %[[USST]], i8** @__safestack_unsafe_stack_ptr
 
-  ; CHECK: %[[AADDR:.*]] = alloca i8*, align 8
   %a.addr = alloca i8*, align 8
-
-  ; CHECK: %[[BUFPTR:.*]] = getelementptr i8, i8* %[[USP]], i32 -4
-  ; CHECK: %[[BUFPTR2:.*]] = bitcast i8* %[[BUFPTR]] to [4 x i8]*
   %buf = alloca [4 x i8], align 1
 
+  ; CHECK: %[[AADDR:.*]] = alloca i8*, align 8
   ; CHECK: store i8* {{.*}}, i8** %[[AADDR]], align 8
   store i8* %a, i8** %a.addr, align 8
 
+  ; CHECK: %[[BUFPTR:.*]] = getelementptr i8, i8* %[[USP]], i32 -4
+  ; CHECK: %[[BUFPTR2:.*]] = bitcast i8* %[[BUFPTR]] to [4 x i8]*
   ; CHECK: %[[GEP:.*]] = getelementptr inbounds [4 x i8], [4 x i8]* %[[BUFPTR2]], i32 0, i32 0
   %gep = getelementptr inbounds [4 x i8], [4 x i8]* %buf, i32 0, i32 0
 
   ; CHECK: %[[A2:.*]] = load i8*, i8** %[[AADDR]], align 8
   %a2 = load i8*, i8** %a.addr, align 8
 
   ; CHECK: call i8* @strcpy(i8* %[[GEP]], i8* %[[A2]])
   %call = call i8* @strcpy(i8* %gep, i8* %a2)

test/Transforms/SafeStack/sink-to-use.ll

+; Test that unsafe alloca address calculation is done immediately before each use.
+; RUN: opt -safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s
+; RUN: opt -safe-stack -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck %s
+
+define void @f() safestack {
+entry:
+  %x0 = alloca i32, align 4
+  %x1 = alloca i32, align 4
+
+; CHECK: %[[A:.*]] = getelementptr i8, i8* %{{.*}}, i32 -4
+; CHECK: %[[X0:.*]] = bitcast i8* %[[A]] to i32*
+; CHECK: call void @use(i32* %[[X0]])
+  call void @use(i32* %x0)
+
+; CHECK: %[[B:.*]] = getelementptr i8, i8* %{{.*}}, i32 -8
+; CHECK: %[[X1:.*]] = bitcast i8* %[[B]] to i32*
+; CHECK: call void @use(i32* %[[X1]])
+  call void @use(i32* %x1)
+  ret void
+}
+
+declare void @use(i32*)

test/Transforms/SafeStack/struct.ll

 define void @foo(i8* %a) nounwind uwtable safestack {
 entry:
   ; CHECK: %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
 
   ; CHECK: %[[USST:.*]] = getelementptr i8, i8* %[[USP]], i32 -16
 
   ; CHECK: store i8* %[[USST]], i8** @__safestack_unsafe_stack_ptr
 
-  ; CHECK: %[[AADDR:.*]] = alloca i8*, align 8
   %a.addr = alloca i8*, align 8
-
-  ; CHECK: %[[BUFPTR:.*]] = getelementptr i8, i8* %[[USP]], i32 -16
-  ; CHECK: %[[BUFPTR2:.*]] = bitcast i8* %[[BUFPTR]] to %struct.foo*
   %buf = alloca %struct.foo, align 1
 
+  ; CHECK: %[[AADDR:.*]] = alloca i8*, align 8
   ; CHECK: store i8* {{.*}}, i8** %[[AADDR]], align 8
   store i8* %a, i8** %a.addr, align 8
 
+  ; CHECK: %[[BUFPTR:.*]] = getelementptr i8, i8* %[[USP]], i32 -16
+  ; CHECK: %[[BUFPTR2:.*]] = bitcast i8* %[[BUFPTR]] to %struct.foo*
   ; CHECK: %[[GEP:.*]] = getelementptr inbounds %struct.foo, %struct.foo* %[[BUFPTR2]], i32 0, i32 0, i32 0
   %gep = getelementptr inbounds %struct.foo, %struct.foo* %buf, i32 0, i32 0, i32 0
 
   ; CHECK: %[[A:.*]] = load i8*, i8** %[[AADDR]], align 8
   %a2 = load i8*, i8** %a.addr, align 8
 
   ; CHECK: call i8* @strcpy(i8* %[[GEP]], i8* %[[A]])
   %call = call i8* @strcpy(i8* %gep, i8* %a2)
 
   ; CHECK: store i8* %[[USP]], i8** @__safestack_unsafe_stack_ptr
   ret void
 }
 
 declare i8* @strcpy(i8*, i8*)