[asan-assembly-instrumentation] Added CFI directives to the generated instrumentation code.

Summary: [asan-assembly-instrumentation] Added CFI directives to the generated instrumentation code.

Reviewers: eugenis

Subscribers: llvm-commits

Differential Revision: http://reviews.llvm.org/D5189

llvm-svn: 217482

include/llvm/MC/MCTargetAsmParser.h

@@ -123,6 +123,9 @@ public:
   virtual bool ParseRegister(unsigned &RegNo, SMLoc &StartLoc,
                              SMLoc &EndLoc) = 0;
 
+  /// Sets frame register corresponding to the current MachineFunction.
+  virtual void SetFrameRegister(unsigned RegNo) {}
+
   /// ParseInstruction - Parse one assembly instruction.
   ///
   /// The parser is positioned following the instruction name. The target

lib/CodeGen/AsmPrinter/AsmPrinterInlineAsm.cpp

@@ -33,6 +33,7 @@
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
+#include "llvm/Target/TargetRegisterInfo.h"
 #include "llvm/Target/TargetSubtargetInfo.h"
 using namespace llvm;
 
@@ -146,6 +147,10 @@ void AsmPrinter::EmitInlineAsm(StringRef Str, const MDNode *LocMDNode,
                        " we don't have an asm parser for this target\n");
   Parser->setAssemblerDialect(Dialect);
   Parser->setTargetParser(*TAP.get());
+  if (MF) {
+    const TargetRegisterInfo *TRI = MF->getSubtarget().getRegisterInfo();
+    TAP->SetFrameRegister(TRI->getFrameRegister(*MF));
+  }
 
   // Don't implicitly switch to the text section before the asm.
   int Res = Parser->Run(/*NoInitialTextSection*/ true,

lib/Target/X86/AsmParser/X86AsmInstrumentation.cpp

@@ -294,6 +294,7 @@ public:
 
   X86AddressSanitizer32(const MCSubtargetInfo &STI)
       : X86AddressSanitizer(STI) {}
+
   virtual ~X86AddressSanitizer32() {}
 
   virtual void StoreFlags(MCStreamer &Out) override {
@@ -307,6 +308,22 @@ public:
   virtual void InstrumentMemOperandPrologue(const RegisterContext &RegCtx,
                                             MCContext &Ctx,
                                             MCStreamer &Out) override {
+    const MCRegisterInfo* MRI = Ctx.getRegisterInfo();
+    if (MRI && FrameReg != X86::NoRegister) {
+      EmitInstruction(
+          Out, MCInstBuilder(X86::PUSH32r).addReg(X86::EBP));
+      if (FrameReg == X86::ESP) {
+        Out.EmitCFIAdjustCfaOffset(4 /* byte size of the FrameReg */);
+        Out.EmitCFIRelOffset(
+            MRI->getDwarfRegNum(X86::EBP, true /* IsEH */), 0);
+      }
+      EmitInstruction(
+          Out, MCInstBuilder(X86::MOV32rr).addReg(X86::EBP).addReg(FrameReg));
+      Out.EmitCFIRememberState();
+      Out.EmitCFIDefCfaRegister(
+          MRI->getDwarfRegNum(X86::EBP, true /* IsEH */));
+    }
+
     EmitInstruction(
         Out, MCInstBuilder(X86::PUSH32r).addReg(RegCtx.addressReg(MVT::i32)));
     EmitInstruction(
@@ -330,6 +347,14 @@ public:
         Out, MCInstBuilder(X86::POP32r).addReg(RegCtx.shadowReg(MVT::i32)));
     EmitInstruction(
         Out, MCInstBuilder(X86::POP32r).addReg(RegCtx.addressReg(MVT::i32)));
+
+    if (Ctx.getRegisterInfo() && FrameReg != X86::NoRegister) {
+      EmitInstruction(
+          Out, MCInstBuilder(X86::POP32r).addReg(X86::EBP));
+      Out.EmitCFIRestoreState();
+      if (FrameReg == X86::ESP)
+        Out.EmitCFIAdjustCfaOffset(-4 /* byte size of the FrameReg */);
+    }
   }
 
   virtual void InstrumentMemOperandSmall(X86Operand &Op, unsigned AccessSize,
@@ -526,6 +551,7 @@ public:
 
   X86AddressSanitizer64(const MCSubtargetInfo &STI)
       : X86AddressSanitizer(STI) {}
+
   virtual ~X86AddressSanitizer64() {}
 
   virtual void StoreFlags(MCStreamer &Out) override {
@@ -539,6 +565,21 @@ public:
   virtual void InstrumentMemOperandPrologue(const RegisterContext &RegCtx,
                                             MCContext &Ctx,
                                             MCStreamer &Out) override {
+    const MCRegisterInfo *RegisterInfo = Ctx.getRegisterInfo();
+    if (RegisterInfo && FrameReg != X86::NoRegister) {
+      EmitInstruction(Out, MCInstBuilder(X86::PUSH64r).addReg(X86::RBP));
+      if (FrameReg == X86::RSP) {
+        Out.EmitCFIAdjustCfaOffset(8 /* byte size of the FrameReg */);
+        Out.EmitCFIRelOffset(
+            RegisterInfo->getDwarfRegNum(X86::RBP, true /* IsEH */), 0);
+      }
+      EmitInstruction(
+          Out, MCInstBuilder(X86::MOV64rr).addReg(X86::RBP).addReg(FrameReg));
+      Out.EmitCFIRememberState();
+      Out.EmitCFIDefCfaRegister(
+          RegisterInfo->getDwarfRegNum(X86::RBP, true /* IsEH */));
+    }
+
     EmitAdjustRSP(Ctx, Out, -128);
     EmitInstruction(
         Out, MCInstBuilder(X86::PUSH64r).addReg(RegCtx.shadowReg(MVT::i64)));
@@ -564,6 +605,14 @@ public:
     EmitInstruction(
         Out, MCInstBuilder(X86::POP64r).addReg(RegCtx.shadowReg(MVT::i64)));
     EmitAdjustRSP(Ctx, Out, 128);
+
+    if (Ctx.getRegisterInfo() && FrameReg != X86::NoRegister) {
+      EmitInstruction(
+          Out, MCInstBuilder(X86::POP64r).addReg(X86::RBP));
+      Out.EmitCFIRestoreState();
+      if (FrameReg == X86::RSP)
+        Out.EmitCFIAdjustCfaOffset(-8 /* byte size of the FrameReg */);
+    }
   }
 
   virtual void InstrumentMemOperandSmall(X86Operand &Op, unsigned AccessSize,
@@ -771,7 +820,7 @@ void X86AddressSanitizer64::InstrumentMOVSImpl(unsigned AccessSize,
 } // End anonymous namespace
 
 X86AsmInstrumentation::X86AsmInstrumentation(const MCSubtargetInfo &STI)
-    : STI(STI) {}
+    : STI(STI), FrameReg(X86::NoRegister) {}
 
 X86AsmInstrumentation::~X86AsmInstrumentation() {}
 

lib/Target/X86/AsmParser/X86AsmInstrumentation.h

@@ -34,6 +34,10 @@ class X86AsmInstrumentation {
 public:
   virtual ~X86AsmInstrumentation();
 
+  void SetFrameRegister(unsigned RegNo) {
+    FrameReg = RegNo;
+  }
+
   // Tries to instrument and emit instruction.
   virtual void InstrumentAndEmitInstruction(
       const MCInst &Inst,
@@ -50,6 +54,8 @@ protected:
   void EmitInstruction(MCStreamer &Out, const MCInst &Inst);
 
   const MCSubtargetInfo &STI;
+
+  unsigned FrameReg;
 };
 
 } // End llvm namespace

lib/Target/X86/AsmParser/X86AsmParser.cpp

@@ -788,6 +788,8 @@ public:
 
   bool ParseRegister(unsigned &RegNo, SMLoc &StartLoc, SMLoc &EndLoc) override;
 
+  void SetFrameRegister(unsigned RegNo) override;
+
   bool ParseInstruction(ParseInstructionInfo &Info, StringRef Name,
                         SMLoc NameLoc, OperandVector &Operands) override;
 
@@ -970,6 +972,10 @@ bool X86AsmParser::ParseRegister(unsigned &RegNo,
   return false;
 }
 
+void X86AsmParser::SetFrameRegister(unsigned RegNo) {
+  Instrumentation->SetFrameRegister(RegNo);
+}
+
 std::unique_ptr<X86Operand> X86AsmParser::DefaultMemSIOperand(SMLoc Loc) {
   unsigned basereg =
     is64BitMode() ? X86::RSI : (is32BitMode() ? X86::ESI : X86::SI);

test/Instrumentation/AddressSanitizer/X86/asm_cfi.ll

@@ -0,0 +1,45 @@
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu -mcpu=corei7 -mattr=+sse2 -asm-instrumentation=address -asan-instrument-assembly | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; CHECK-LABEL: mov8b_rbp
+; CHECK: pushq %rbp
+; CHECK-NOT: .cfi_adjust_cfa_offset 8
+; CHECK: movq %rbp, %rbp
+; CHECK: .cfi_remember_state
+; CHECK: .cfi_def_cfa_register %rbp
+; CHECK: leaq -128(%rsp)
+; CHECK: callq __asan_report_load8@PLT
+; CHECK: leaq 128(%rsp)
+; CHECK: popq %rbp
+; CHECK: .cfi_restore_state
+; CHECK-NOT: .cfi_adjust_cfa_offset -8
+; CHECK: retq
+define void @mov8b_rbp(i64* %dst, i64* %src) #0 {
+entry:
+  tail call void asm sideeffect "movq ($0), %rax \0A\09movq %rax, ($1) \0A\09", "r,r,~{rax},~{memory},~{dirflag},~{fpsr},~{flags}"(i64* %src, i64* %dst)
+  ret void
+}
+
+; CHECK-LABEL: mov8b_rsp
+; CHECK: pushq %rbp
+; CHECK: .cfi_adjust_cfa_offset 8
+; CHECK: movq %rsp, %rbp
+; CHECK: .cfi_remember_state
+; CHECK: .cfi_def_cfa_register %rbp
+; CHECK: leaq -128(%rsp)
+; CHECK: callq __asan_report_load8@PLT
+; CHECK: leaq 128(%rsp)
+; CHECK: popq %rbp
+; CHECK: .cfi_restore_state
+; CHECK: .cfi_adjust_cfa_offset -8
+; CHECK: retq
+define void @mov8b_rsp(i64* %dst, i64* %src) #1 {
+entry:
+  tail call void asm sideeffect "movq ($0), %rax \0A\09movq %rax, ($1) \0A\09", "r,r,~{rax},~{memory},~{dirflag},~{fpsr},~{flags}"(i64* %src, i64* %dst)
+  ret void
+}
+
+attributes #0 = { nounwind sanitize_address uwtable "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" }
+attributes #1 = { nounwind sanitize_address uwtable "no-frame-pointer-elim"="false" }