[AddressSanitizer] Fix for wrong argument values appearing in backtraces

[vsk]

Summary

In some cases, ASan may insert instrumentation before function arguments
have been stored into their allocas. This causes two issues:

1. The argument value must be spilled until it can be stored into the reserved alloca, wasting a stack slot.

2. Until the store occurs in a later basic block, the debug location will point to the wrong frame offset, and backtraces will show an uninitialized value.

The proposed solution is to move instructions which initialize allocas
for arguments up into the entry block, before the position where ASan
starts inserting its instrumentation.

For the motivating test case, before the patch we see:

| 0033: movq %rdi, 0x68(%rbx)  |   | DW_TAG_formal_parameter     |
| ...                          |   |   DW_AT_name ("a")          |
| 00d1: movq 0x68(%rbx), %rsi  |   |   DW_AT_location (RBX+0x90) |
| 00d5: movq %rsi, 0x90(%rbx)  |   |        ^ not correct ...    |

and after the patch we see:

| 002f: movq %rdi, 0x70(%rbx)  |   | DW_TAG_formal_parameter     |
|                              |   |   DW_AT_name ("a")          |
|                              |   |   DW_AT_location (RBX+0x70) |

rdar://61122691

[aprantl]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

3003

It's counterintuitive that this isn't !isInterestingAlloca(). Perhaps commen why we skip the interesting ones?

[eugenis]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

3023

What happens with

a = alloca
store arg, a
load a
store arg2, a

will the second store be moved across the aliasing load?

[vsk]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

3023

Yes, thanks for catching this.

I plan to address this by stopping the loop when an unknown instruction is seen. I.e., make it bail out if "I" is not a StoreInst or a CastInst that matches the known argument init cases. Does that sound ok?

Another more expensive/general option is to use the isSafeToMoveBefore utility from CodeMoverUtils (this requires DependenceInfo, DomTree, and PostDomTree).

[eugenis]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

3023

Sounds fine.

[vsk]

Do not reorder past unknown instructions.

[vsk]

Hm, this will not work if there's more than one interesting alloca for ASan to instrument. Should have a fix shortly.

[vsk]

Handle case where there is more than one interesting alloca.

[vsk]

@eugenis do you have any further feedback on this one? Thanks.

[eugenis]

LGTM with 2 comments

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

2991

Mention "uninstrumented" (alloca) or something like that in the function name to make it clear that this function does not break regular alloca instrumentation by moving stores before stack poisoning.

This only really affects -O0 compilation, right?

3008

It looks like it would be simpler and cheaper to move this check down to the initializer of IsArgInitViaCast, and get rid of the vector update and lookup.

[vsk]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

2991

Yes, with optimization enabled most allocas for arguments get promoted away.

3008

Sounds good. The only gotcha is to make sure the cast appears after InsBefore. I left a comment about this.

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

 //===- AddressSanitizer.cpp - memory error detector -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
void FunctionStackPoisoner::processDynamicAllocas() {
 
   // Handle dynamic allocas.
   createDynamicAllocasInitStorage();
   for (auto &AI : DynamicAllocaVec)
     handleDynamicAllocaCall(AI);
   unpoisonDynamicAllocas();
 }
 
+/// Collect instructions in the entry block after \p InsBefore which initialize
+/// permanent storage for a function argument. These instructions must remain in
+/// the entry block so that uninitialized values do not appear in backtraces. An
+/// added benefit is that this conserves spill slots.
+static void
+collectArgumentAllocaInitInsts(Instruction &InsBefore,
+                               SmallVectorImpl<Instruction *> &InitInsts) {
+  SmallVector<Instruction *, 4> ArgCasts;
+  Instruction *Start = InsBefore.getNextNonDebugInstruction();
+  for (Instruction *It = Start; It; It = It->getNextNonDebugInstruction()) {
+    // Argument initialization looks like:
+    // 1) store <Argument>, <Alloca> OR
+    // 2) <CastArgument> = cast <Argument> to ...
+    //    store <CastArgument> to <Alloca>
+    // Do not consider any other kind of instruction.
+    //
+    // Note: This covers all known cases, but may not be exhaustive. An
+    // alternative to pattern-matching stores is to DFS over all Argument uses:
+    // this might be more general, but is probably much more complicated.
+    if (auto *Cast = dyn_cast<CastInst>(It)) {
+      if (isa<Argument>(Cast->getOperand(0)))
+        ArgCasts.push_back(Cast);
+      continue;
+    }
+    if (auto *Store = dyn_cast<StoreInst>(It)) {
+      // The store destination must be an alloca. It is defined before
+      // InsBefore, otherwise we would have returned early.
+      auto *Alloca = dyn_cast<AllocaInst>(Store->getPointerOperand());
+      if (!Alloca)
+        continue;
+
+      Value *Val = Store->getValueOperand();
+      bool IsDirectArgInit = isa<Argument>(Val);
+      bool IsArgInitViaCast = find(ArgCasts, Val) != ArgCasts.end();
+      bool IsArgInit = IsDirectArgInit || IsArgInitViaCast;
+      if (!IsArgInit)
+        continue;
+
+      if (IsArgInitViaCast)
+        InitInsts.push_back(cast<Instruction>(Val));
+      InitInsts.push_back(Store);
+      continue;
+    }
+
+    // Do not reorder past unknown instructions: argument initialization should
+    // only involve casts and stores.
+    return;
+  }
+}
+
 void FunctionStackPoisoner::processStaticAllocas() {
   if (AllocaVec.empty()) {
     assert(StaticAllocaPoisonCallVec.empty());
     return;
   }
 
   int StackMallocIdx = -1;
   DebugLoc EntryDebugLocation;
   if (auto SP = F.getSubprogram())
     EntryDebugLocation = DebugLoc::get(SP->getScopeLine(), 0, SP);
 
   Instruction *InsBefore = AllocaVec[0];
   IRBuilder<> IRB(InsBefore);
 
   // Make sure non-instrumented allocas stay in the entry block. Otherwise,
   // debug info is broken, because only entry-block allocas are treated as
   // regular stack slots.
   auto InsBeforeB = InsBefore->getParent();
   assert(InsBeforeB == &F.getEntryBlock());
   for (auto *AI : StaticAllocasToMoveUp)
     if (AI->getParent() == InsBeforeB)
       AI->moveBefore(InsBefore);
 
+  // Move stores of arguments into entry-block allocas as well. This prevents
+  // extra stack slots from being generated (to house the argument values until
+  // they can be stored into the allocas). This also prevents uninitialized
+  // values from being shown in backtraces.
+  SmallVector<Instruction *, 8> ArgInitInsts;
+  collectArgumentAllocaInitInsts(*InsBefore, ArgInitInsts);
+  for (Instruction *ArgInitInst : ArgInitInsts)
+    ArgInitInst->moveBefore(InsBefore);
+
   // If we have a call to llvm.localescape, keep it in the entry block.
   if (LocalEscapeCall) LocalEscapeCall->moveBefore(InsBefore);
 
   SmallVector<ASanStackVariableDescription, 16> SVD;
   SVD.reserve(AllocaVec.size());
   for (AllocaInst *AI : AllocaVec) {
     ASanStackVariableDescription D = {AI->getName().data(),
                                       ASan.getAllocaSizeInBytes(*AI),

llvm/test/Instrumentation/AddressSanitizer/hoist-argument-init-insts.ll

+; RUN: opt < %s -asan -asan-module -asan-use-after-return -S | FileCheck %s
+
+; Source (-O0 -fsanitize=address -fsanitize-address-use-after-scope):
+;; struct S { int x, y; };
+;; void swap(S *a, S *b, bool doit) {
+;;   if (!doit)
+;;     return;
+;;   auto tmp = *a;
+;;   *a = *b;
+;;   *b = tmp;
+;; }
+
+target datalayout = "e-m:o-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.14.0"
+
+%struct.S = type { i32, i32 }
+
+; CHECK-LABEL: define {{.*}} @_Z4swapP1SS0_b(
+
+; First come the argument allocas.
+; CHECK:      [[argA:%.*]] = alloca %struct.S*,
+; CHECK-NEXT: [[argB:%.*]] = alloca %struct.S*,
+; CHECK-NEXT: [[argDoit:%.*]] = alloca i8,
+
+; Next, the stores into the argument allocas.
+; CHECK-NEXT: store %struct.S* {{.*}}, %struct.S** [[argA]]
+; CHECK-NEXT: store %struct.S* {{.*}}, %struct.S** [[argB]]
+; CHECK-NEXT: [[frombool:%.*]] = zext i1 {{.*}} to i8
+; CHECK-NEXT: store i8 [[frombool]], i8* [[argDoit]]
+
+define void @_Z4swapP1SS0_b(%struct.S* %a, %struct.S* %b, i1 zeroext %doit) sanitize_address {
+entry:
+  %a.addr = alloca %struct.S*, align 8
+  %b.addr = alloca %struct.S*, align 8
+  %doit.addr = alloca i8, align 1
+  %tmp = alloca %struct.S, align 4
+  store %struct.S* %a, %struct.S** %a.addr, align 8
+  store %struct.S* %b, %struct.S** %b.addr, align 8
+  %frombool = zext i1 %doit to i8
+  store i8 %frombool, i8* %doit.addr, align 1
+  %0 = load i8, i8* %doit.addr, align 1
+  %tobool = trunc i8 %0 to i1
+  br i1 %tobool, label %if.end, label %if.then
+
+if.then:                                          ; preds = %entry
+  br label %return
+
+if.end:                                           ; preds = %entry
+  %1 = load %struct.S*, %struct.S** %a.addr, align 8
+  %2 = bitcast %struct.S* %tmp to i8*
+  %3 = bitcast %struct.S* %1 to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 4 %2, i8* align 4 %3, i64 8, i1 false)
+  %4 = load %struct.S*, %struct.S** %b.addr, align 8
+  %5 = load %struct.S*, %struct.S** %a.addr, align 8
+  %6 = bitcast %struct.S* %5 to i8*
+  %7 = bitcast %struct.S* %4 to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 4 %6, i8* align 4 %7, i64 8, i1 false)
+  %8 = load %struct.S*, %struct.S** %b.addr, align 8
+  %9 = bitcast %struct.S* %8 to i8*
+  %10 = bitcast %struct.S* %tmp to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 4 %9, i8* align 4 %10, i64 8, i1 false)
+  br label %return
+
+return:                                           ; preds = %if.end, %if.then
+  ret void
+}
+
+declare void @llvm.memcpy.p0i8.p0i8.i64(i8* noalias nocapture writeonly, i8* noalias nocapture readonly, i64, i1 immarg)
+
+; CHECK-LABEL: define {{.*}} @func_with_load_in_arginit_sequence
+; CHECK:      [[argA:%.*]] = alloca %struct.S*,
+; CHECK-NEXT: [[argB:%.*]] = alloca %struct.S*,
+; CHECK-NEXT: [[argDoit:%.*]] = alloca i8,
+; CHECK-NEXT: store %struct.S* {{.*}}, %struct.S** [[argA]]
+; CHECK-NEXT: store %struct.S* {{.*}}, %struct.S** [[argB]]
+; CHECK-NEXT: [[stack_base:%.*]] = alloca i64
+
+; Synthetic test case, meant to check that we do not reorder instructions past
+; a load when attempting to hoist argument init insts.
+define void @func_with_load_in_arginit_sequence(%struct.S* %a, %struct.S* %b, i1 zeroext %doit) sanitize_address {
+entry:
+  %a.addr = alloca %struct.S*, align 8
+  %b.addr = alloca %struct.S*, align 8
+  %doit.addr = alloca i8, align 1
+  %tmp = alloca %struct.S, align 4
+  store %struct.S* %a, %struct.S** %a.addr, align 8
+  store %struct.S* %b, %struct.S** %b.addr, align 8
+
+  ; This load prevents the next argument init sequence from being moved.
+  %0 = load i8, i8* %doit.addr, align 1 
+
+  %frombool = zext i1 %doit to i8
+  store i8 %frombool, i8* %doit.addr, align 1
+  %tobool = trunc i8 %0 to i1
+  br i1 %tobool, label %if.end, label %if.then
+
+if.then:                                          ; preds = %entry
+  br label %return
+
+if.end:                                           ; preds = %entry
+  %1 = load %struct.S*, %struct.S** %a.addr, align 8
+  %2 = bitcast %struct.S* %tmp to i8*
+  %3 = bitcast %struct.S* %1 to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 4 %2, i8* align 4 %3, i64 8, i1 false)
+  %4 = load %struct.S*, %struct.S** %b.addr, align 8
+  %5 = load %struct.S*, %struct.S** %a.addr, align 8
+  %6 = bitcast %struct.S* %5 to i8*
+  %7 = bitcast %struct.S* %4 to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 4 %6, i8* align 4 %7, i64 8, i1 false)
+  %8 = load %struct.S*, %struct.S** %b.addr, align 8
+  %9 = bitcast %struct.S* %8 to i8*
+  %10 = bitcast %struct.S* %tmp to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 4 %9, i8* align 4 %10, i64 8, i1 false)
+  br label %return
+
+return:                                           ; preds = %if.end, %if.then
+  ret void
+}
+
+