lib/Target/X86/AsmParser/X86AsmInstrumentation.cpp

   1 //===-- X86AsmInstrumentation.cpp - Instrument X86 inline assembly C++ -*-===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9
  10 #include "MCTargetDesc/X86BaseInfo.h"
  11 #include "X86AsmInstrumentation.h"
  12 #include "X86Operand.h"
  13 #include "X86RegisterInfo.h"
  14 #include "llvm/ADT/StringExtras.h"
  15 #include "llvm/ADT/Triple.h"
  16 #include "llvm/CodeGen/MachineValueType.h"
  17 #include "llvm/IR/Function.h"
  18 #include "llvm/MC/MCAsmInfo.h"
  19 #include "llvm/MC/MCContext.h"
  20 #include "llvm/MC/MCInst.h"
  21 #include "llvm/MC/MCInstBuilder.h"
  22 #include "llvm/MC/MCInstrInfo.h"
  23 #include "llvm/MC/MCParser/MCParsedAsmOperand.h"
  24 #include "llvm/MC/MCStreamer.h"
  25 #include "llvm/MC/MCSubtargetInfo.h"
  26 #include "llvm/MC/MCTargetAsmParser.h"
  27 #include "llvm/MC/MCTargetOptions.h"
  28 #include "llvm/Support/CommandLine.h"
  29
  30 namespace llvm {
  31 namespace {
  32
  33 static cl::opt<bool> ClAsanInstrumentAssembly(
  34     "asan-instrument-assembly",
  35     cl::desc("instrument assembly with AddressSanitizer checks"), cl::Hidden,
  36     cl::init(false));
  37
  38 bool IsStackReg(unsigned Reg) {
  39   return Reg == X86::RSP || Reg == X86::ESP || Reg == X86::SP;
  40 }
  41
  42 bool IsSmallMemAccess(unsigned AccessSize) { return AccessSize < 8; }
  43
  44 std::string FuncName(unsigned AccessSize, bool IsWrite) {
  45   return std::string("__asan_report_") + (IsWrite ? "store" : "load") +
  46          utostr(AccessSize);
  47 }
  48
  49 class X86AddressSanitizer : public X86AsmInstrumentation {
  50 public:
  51   struct RegisterContext {
  52     RegisterContext(unsigned AddressReg, unsigned ShadowReg,
  53                     unsigned ScratchReg)
  54         : AddressReg(AddressReg), ShadowReg(ShadowReg), ScratchReg(ScratchReg) {
  55     }
  56
  57     unsigned addressReg(MVT::SimpleValueType VT) const {
  58       return getX86SubSuperRegister(AddressReg, VT);
  59     }
  60
  61     unsigned shadowReg(MVT::SimpleValueType VT) const {
  62       return getX86SubSuperRegister(ShadowReg, VT);
  63     }
  64
  65     unsigned scratchReg(MVT::SimpleValueType VT) const {
  66       return getX86SubSuperRegister(ScratchReg, VT);
  67     }
  68
  69     const unsigned AddressReg;
  70     const unsigned ShadowReg;
  71     const unsigned ScratchReg;
  72   };
  73
  74   X86AddressSanitizer(const MCSubtargetInfo &STI)
  75       : X86AsmInstrumentation(STI), RepPrefix(false) {}
  76   virtual ~X86AddressSanitizer() {}
  77
  78   // X86AsmInstrumentation implementation:
  79   virtual void InstrumentAndEmitInstruction(const MCInst &Inst,
  80                                             OperandVector &Operands,
  81                                             MCContext &Ctx,
  82                                             const MCInstrInfo &MII,
  83                                             MCStreamer &Out) override {
  84     InstrumentMOVS(Inst, Operands, Ctx, MII, Out);
  85     if (RepPrefix)
  86       EmitInstruction(Out, MCInstBuilder(X86::REP_PREFIX));
  87
  88     InstrumentMOV(Inst, Operands, Ctx, MII, Out);
  89
  90     RepPrefix = (Inst.getOpcode() == X86::REP_PREFIX);
  91     if (!RepPrefix)
  92       EmitInstruction(Out, Inst);
  93   }
  94
  95   // Should be implemented differently in x86_32 and x86_64 subclasses.
  96   virtual void StoreFlags(MCStreamer &Out) = 0;
  97
  98   virtual void RestoreFlags(MCStreamer &Out) = 0;
  99
 100   // Adjusts up stack and saves all registers used in instrumentation.
 101   virtual void InstrumentMemOperandPrologue(const RegisterContext &RegCtx,
 102                                             MCContext &Ctx,
 103                                             MCStreamer &Out) = 0;
 104
 105   // Restores all registers used in instrumentation and adjusts stack.
 106   virtual void InstrumentMemOperandEpilogue(const RegisterContext &RegCtx,
 107                                             MCContext &Ctx,
 108                                             MCStreamer &Out) = 0;
 109
 110   virtual void InstrumentMemOperandSmall(X86Operand &Op, unsigned AccessSize,
 111                                          bool IsWrite,
 112                                          const RegisterContext &RegCtx,
 113                                          MCContext &Ctx, MCStreamer &Out) = 0;
 114   virtual void InstrumentMemOperandLarge(X86Operand &Op, unsigned AccessSize,
 115                                          bool IsWrite,
 116                                          const RegisterContext &RegCtx,
 117                                          MCContext &Ctx, MCStreamer &Out) = 0;
 118
 119   virtual void InstrumentMOVSImpl(unsigned AccessSize, MCContext &Ctx,
 120                                   MCStreamer &Out) = 0;
 121
 122   void InstrumentMemOperand(X86Operand &Op, unsigned AccessSize, bool IsWrite,
 123                             const RegisterContext &RegCtx, MCContext &Ctx,
 124                             MCStreamer &Out);
 125   void InstrumentMOVSBase(unsigned DstReg, unsigned SrcReg, unsigned CntReg,
 126                           unsigned AccessSize, MCContext &Ctx, MCStreamer &Out);
 127
 128   void InstrumentMOVS(const MCInst &Inst, OperandVector &Operands,
 129                       MCContext &Ctx, const MCInstrInfo &MII, MCStreamer &Out);
 130   void InstrumentMOV(const MCInst &Inst, OperandVector &Operands,
 131                      MCContext &Ctx, const MCInstrInfo &MII, MCStreamer &Out);
 132
 133 protected:
 134   void EmitLabel(MCStreamer &Out, MCSymbol *Label) { Out.EmitLabel(Label); }
 135
 136   // True when previous instruction was actually REP prefix.
 137   bool RepPrefix;
 138 };
 139
 140 void X86AddressSanitizer::InstrumentMemOperand(
 141     X86Operand &Op, unsigned AccessSize, bool IsWrite,
 142     const RegisterContext &RegCtx, MCContext &Ctx, MCStreamer &Out) {
 143   assert(Op.isMem() && "Op should be a memory operand.");
 144   assert((AccessSize & (AccessSize - 1)) == 0 && AccessSize <= 16 &&
 145          "AccessSize should be a power of two, less or equal than 16.");
 146   // FIXME: take into account load/store alignment.
 147   if (IsSmallMemAccess(AccessSize))
 148     InstrumentMemOperandSmall(Op, AccessSize, IsWrite, RegCtx, Ctx, Out);
 149   else
 150     InstrumentMemOperandLarge(Op, AccessSize, IsWrite, RegCtx, Ctx, Out);
 151 }
 152
 153 void X86AddressSanitizer::InstrumentMOVSBase(unsigned DstReg, unsigned SrcReg,
 154                                              unsigned CntReg,
 155                                              unsigned AccessSize,
 156                                              MCContext &Ctx, MCStreamer &Out) {
 157   // FIXME: check whole ranges [DstReg .. DstReg + AccessSize * (CntReg - 1)]
 158   // and [SrcReg .. SrcReg + AccessSize * (CntReg - 1)].
 159   RegisterContext RegCtx(X86::RDX /* AddressReg */, X86::RAX /* ShadowReg */,
 160                          IsSmallMemAccess(AccessSize)
 161                              ? X86::RBX
 162                              : X86::NoRegister /* ScratchReg */);
 163
 164   InstrumentMemOperandPrologue(RegCtx, Ctx, Out);
 165
 166   // Test (%SrcReg)
 167   {
 168     const MCExpr *Disp = MCConstantExpr::Create(0, Ctx);
 169     std::unique_ptr<X86Operand> Op(X86Operand::CreateMem(
 170         0, Disp, SrcReg, 0, AccessSize, SMLoc(), SMLoc()));
 171     InstrumentMemOperand(*Op, AccessSize, false /* IsWrite */, RegCtx, Ctx,
 172                          Out);
 173   }
 174
 175   // Test -1(%SrcReg, %CntReg, AccessSize)
 176   {
 177     const MCExpr *Disp = MCConstantExpr::Create(-1, Ctx);
 178     std::unique_ptr<X86Operand> Op(X86Operand::CreateMem(
 179         0, Disp, SrcReg, CntReg, AccessSize, SMLoc(), SMLoc()));
 180     InstrumentMemOperand(*Op, AccessSize, false /* IsWrite */, RegCtx, Ctx,
 181                          Out);
 182   }
 183
 184   // Test (%DstReg)
 185   {
 186     const MCExpr *Disp = MCConstantExpr::Create(0, Ctx);
 187     std::unique_ptr<X86Operand> Op(X86Operand::CreateMem(
 188         0, Disp, DstReg, 0, AccessSize, SMLoc(), SMLoc()));
 189     InstrumentMemOperand(*Op, AccessSize, true /* IsWrite */, RegCtx, Ctx, Out);
 190   }
 191
 192   // Test -1(%DstReg, %CntReg, AccessSize)
 193   {
 194     const MCExpr *Disp = MCConstantExpr::Create(-1, Ctx);
 195     std::unique_ptr<X86Operand> Op(X86Operand::CreateMem(
 196         0, Disp, DstReg, CntReg, AccessSize, SMLoc(), SMLoc()));
 197     InstrumentMemOperand(*Op, AccessSize, true /* IsWrite */, RegCtx, Ctx, Out);
 198   }
 199
 200   InstrumentMemOperandEpilogue(RegCtx, Ctx, Out);
 201 }
 202
 203 void X86AddressSanitizer::InstrumentMOVS(const MCInst &Inst,
 204                                          OperandVector &Operands,
 205                                          MCContext &Ctx, const MCInstrInfo &MII,
 206                                          MCStreamer &Out) {
 207   // Access size in bytes.
 208   unsigned AccessSize = 0;
 209
 210   switch (Inst.getOpcode()) {
 211   case X86::MOVSB:
 212     AccessSize = 1;
 213     break;
 214   case X86::MOVSW:
 215     AccessSize = 2;
 216     break;
 217   case X86::MOVSL:
 218     AccessSize = 4;
 219     break;
 220   case X86::MOVSQ:
 221     AccessSize = 8;
 222     break;
 223   default:
 224     return;
 225   }
 226
 227   InstrumentMOVSImpl(AccessSize, Ctx, Out);
 228 }
 229
 230 void X86AddressSanitizer::InstrumentMOV(const MCInst &Inst,
 231                                         OperandVector &Operands, MCContext &Ctx,
 232                                         const MCInstrInfo &MII,
 233                                         MCStreamer &Out) {
 234   // Access size in bytes.
 235   unsigned AccessSize = 0;
 236
 237   switch (Inst.getOpcode()) {
 238   case X86::MOV8mi:
 239   case X86::MOV8mr:
 240   case X86::MOV8rm:
 241     AccessSize = 1;
 242     break;
 243   case X86::MOV16mi:
 244   case X86::MOV16mr:
 245   case X86::MOV16rm:
 246     AccessSize = 2;
 247     break;
 248   case X86::MOV32mi:
 249   case X86::MOV32mr:
 250   case X86::MOV32rm:
 251     AccessSize = 4;
 252     break;
 253   case X86::MOV64mi32:
 254   case X86::MOV64mr:
 255   case X86::MOV64rm:
 256     AccessSize = 8;
 257     break;
 258   case X86::MOVAPDmr:
 259   case X86::MOVAPSmr:
 260   case X86::MOVAPDrm:
 261   case X86::MOVAPSrm:
 262     AccessSize = 16;
 263     break;
 264   default:
 265     return;
 266   }
 267
 268   const bool IsWrite = MII.get(Inst.getOpcode()).mayStore();
 269   RegisterContext RegCtx(X86::RDI /* AddressReg */, X86::RAX /* ShadowReg */,
 270                          IsSmallMemAccess(AccessSize)
 271                              ? X86::RCX
 272                              : X86::NoRegister /* ScratchReg */);
 273
 274   for (unsigned Ix = 0; Ix < Operands.size(); ++Ix) {
 275     assert(Operands[Ix]);
 276     MCParsedAsmOperand &Op = *Operands[Ix];
 277     if (Op.isMem()) {
 278       X86Operand &MemOp = static_cast<X86Operand &>(Op);
 279       // FIXME: get rid of this limitation.
 280       if (IsStackReg(MemOp.getMemBaseReg()) ||
 281           IsStackReg(MemOp.getMemIndexReg())) {
 282         continue;
 283       }
 284
 285       InstrumentMemOperandPrologue(RegCtx, Ctx, Out);
 286       InstrumentMemOperand(MemOp, AccessSize, IsWrite, RegCtx, Ctx, Out);
 287       InstrumentMemOperandEpilogue(RegCtx, Ctx, Out);
 288     }
 289   }
 290 }
 291
 292 class X86AddressSanitizer32 : public X86AddressSanitizer {
 293 public:
 294   static const long kShadowOffset = 0x20000000;
 295
 296   X86AddressSanitizer32(const MCSubtargetInfo &STI)
 297       : X86AddressSanitizer(STI) {}
 298
 299   virtual ~X86AddressSanitizer32() {}
 300
 301   unsigned GetFrameReg(const MCContext &Ctx, MCStreamer &Out) {
 302     unsigned FrameReg = GetFrameRegGeneric(Ctx, Out);
 303     if (FrameReg == X86::NoRegister)
 304       return FrameReg;
 305     return getX86SubSuperRegister(FrameReg, MVT::i32);
 306   }
 307
 308   virtual void StoreFlags(MCStreamer &Out) override {
 309     EmitInstruction(Out, MCInstBuilder(X86::PUSHF32));
 310   }
 311
 312   virtual void RestoreFlags(MCStreamer &Out) override {
 313     EmitInstruction(Out, MCInstBuilder(X86::POPF32));
 314   }
 315
 316   virtual void InstrumentMemOperandPrologue(const RegisterContext &RegCtx,
 317                                             MCContext &Ctx,
 318                                             MCStreamer &Out) override {
 319     const MCRegisterInfo *MRI = Ctx.getRegisterInfo();
 320     unsigned FrameReg = GetFrameReg(Ctx, Out);
 321     if (MRI && FrameReg != X86::NoRegister) {
 322       EmitInstruction(
 323           Out, MCInstBuilder(X86::PUSH32r).addReg(X86::EBP));
 324       if (FrameReg == X86::ESP) {
 325         Out.EmitCFIAdjustCfaOffset(4 /* byte size of the FrameReg */);
 326         Out.EmitCFIRelOffset(
 327             MRI->getDwarfRegNum(X86::EBP, true /* IsEH */), 0);
 328       }
 329       EmitInstruction(
 330           Out, MCInstBuilder(X86::MOV32rr).addReg(X86::EBP).addReg(FrameReg));
 331       Out.EmitCFIRememberState();
 332       Out.EmitCFIDefCfaRegister(
 333           MRI->getDwarfRegNum(X86::EBP, true /* IsEH */));
 334     }
 335
 336     EmitInstruction(
 337         Out, MCInstBuilder(X86::PUSH32r).addReg(RegCtx.addressReg(MVT::i32)));
 338     EmitInstruction(
 339         Out, MCInstBuilder(X86::PUSH32r).addReg(RegCtx.shadowReg(MVT::i32)));
 340     if (RegCtx.ScratchReg != X86::NoRegister) {
 341       EmitInstruction(
 342           Out, MCInstBuilder(X86::PUSH32r).addReg(RegCtx.scratchReg(MVT::i32)));
 343     }
 344     StoreFlags(Out);
 345   }
 346
 347   virtual void InstrumentMemOperandEpilogue(const RegisterContext &RegCtx,
 348                                             MCContext &Ctx,
 349                                             MCStreamer &Out) override {
 350     RestoreFlags(Out);
 351     if (RegCtx.ScratchReg != X86::NoRegister) {
 352       EmitInstruction(
 353           Out, MCInstBuilder(X86::POP32r).addReg(RegCtx.scratchReg(MVT::i32)));
 354     }
 355     EmitInstruction(
 356         Out, MCInstBuilder(X86::POP32r).addReg(RegCtx.shadowReg(MVT::i32)));
 357     EmitInstruction(
 358         Out, MCInstBuilder(X86::POP32r).addReg(RegCtx.addressReg(MVT::i32)));
 359
 360     unsigned FrameReg = GetFrameReg(Ctx, Out);
 361     if (Ctx.getRegisterInfo() && FrameReg != X86::NoRegister) {
 362       EmitInstruction(
 363           Out, MCInstBuilder(X86::POP32r).addReg(X86::EBP));
 364       Out.EmitCFIRestoreState();
 365       if (FrameReg == X86::ESP)
 366         Out.EmitCFIAdjustCfaOffset(-4 /* byte size of the FrameReg */);
 367     }
 368   }
 369
 370   virtual void InstrumentMemOperandSmall(X86Operand &Op, unsigned AccessSize,
 371                                          bool IsWrite,
 372                                          const RegisterContext &RegCtx,
 373                                          MCContext &Ctx,
 374                                          MCStreamer &Out) override;
 375   virtual void InstrumentMemOperandLarge(X86Operand &Op, unsigned AccessSize,
 376                                          bool IsWrite,
 377                                          const RegisterContext &RegCtx,
 378                                          MCContext &Ctx,
 379                                          MCStreamer &Out) override;
 380   virtual void InstrumentMOVSImpl(unsigned AccessSize, MCContext &Ctx,
 381                                   MCStreamer &Out) override;
 382
 383 private:
 384   void EmitCallAsanReport(unsigned AccessSize, bool IsWrite, MCContext &Ctx,
 385                           MCStreamer &Out, const RegisterContext &RegCtx) {
 386     EmitInstruction(Out, MCInstBuilder(X86::CLD));
 387     EmitInstruction(Out, MCInstBuilder(X86::MMX_EMMS));
 388
 389     EmitInstruction(Out, MCInstBuilder(X86::AND64ri8)
 390                              .addReg(X86::ESP)
 391                              .addReg(X86::ESP)
 392                              .addImm(-16));
 393     EmitInstruction(
 394         Out, MCInstBuilder(X86::PUSH32r).addReg(RegCtx.addressReg(MVT::i32)));
 395
 396     const std::string &Fn = FuncName(AccessSize, IsWrite);
 397     MCSymbol *FnSym = Ctx.GetOrCreateSymbol(StringRef(Fn));
 398     const MCSymbolRefExpr *FnExpr =
 399         MCSymbolRefExpr::Create(FnSym, MCSymbolRefExpr::VK_PLT, Ctx);
 400     EmitInstruction(Out, MCInstBuilder(X86::CALLpcrel32).addExpr(FnExpr));
 401   }
 402 };
 403
 404 void X86AddressSanitizer32::InstrumentMemOperandSmall(
 405     X86Operand &Op, unsigned AccessSize, bool IsWrite,
 406     const RegisterContext &RegCtx, MCContext &Ctx, MCStreamer &Out) {
 407   unsigned AddressRegI32 = RegCtx.addressReg(MVT::i32);
 408   unsigned ShadowRegI32 = RegCtx.shadowReg(MVT::i32);
 409   unsigned ShadowRegI8 = RegCtx.shadowReg(MVT::i8);
 410
 411   assert(RegCtx.ScratchReg != X86::NoRegister);
 412   unsigned ScratchRegI32 = RegCtx.scratchReg(MVT::i32);
 413
 414   {
 415     MCInst Inst;
 416     Inst.setOpcode(X86::LEA32r);
 417     Inst.addOperand(MCOperand::CreateReg(AddressRegI32));
 418     Op.addMemOperands(Inst, 5);
 419     EmitInstruction(Out, Inst);
 420   }
 421
 422   EmitInstruction(Out, MCInstBuilder(X86::MOV32rr).addReg(ShadowRegI32).addReg(
 423                            AddressRegI32));
 424   EmitInstruction(Out, MCInstBuilder(X86::SHR32ri)
 425                            .addReg(ShadowRegI32)
 426                            .addReg(ShadowRegI32)
 427                            .addImm(3));
 428
 429   {
 430     MCInst Inst;
 431     Inst.setOpcode(X86::MOV8rm);
 432     Inst.addOperand(MCOperand::CreateReg(ShadowRegI8));
 433     const MCExpr *Disp = MCConstantExpr::Create(kShadowOffset, Ctx);
 434     std::unique_ptr<X86Operand> Op(
 435         X86Operand::CreateMem(0, Disp, ShadowRegI32, 0, 1, SMLoc(), SMLoc()));
 436     Op->addMemOperands(Inst, 5);
 437     EmitInstruction(Out, Inst);
 438   }
 439
 440   EmitInstruction(
 441       Out, MCInstBuilder(X86::TEST8rr).addReg(ShadowRegI8).addReg(ShadowRegI8));
 442   MCSymbol *DoneSym = Ctx.CreateTempSymbol();
 443   const MCExpr *DoneExpr = MCSymbolRefExpr::Create(DoneSym, Ctx);
 444   EmitInstruction(Out, MCInstBuilder(X86::JE_4).addExpr(DoneExpr));
 445
 446   EmitInstruction(Out, MCInstBuilder(X86::MOV32rr).addReg(ScratchRegI32).addReg(
 447                            AddressRegI32));
 448   EmitInstruction(Out, MCInstBuilder(X86::AND32ri)
 449                            .addReg(ScratchRegI32)
 450                            .addReg(ScratchRegI32)
 451                            .addImm(7));
 452
 453   switch (AccessSize) {
 454   case 1:
 455     break;
 456   case 2: {
 457     MCInst Inst;
 458     Inst.setOpcode(X86::LEA32r);
 459     Inst.addOperand(MCOperand::CreateReg(ScratchRegI32));
 460
 461     const MCExpr *Disp = MCConstantExpr::Create(1, Ctx);
 462     std::unique_ptr<X86Operand> Op(
 463         X86Operand::CreateMem(0, Disp, ScratchRegI32, 0, 1, SMLoc(), SMLoc()));
 464     Op->addMemOperands(Inst, 5);
 465     EmitInstruction(Out, Inst);
 466     break;
 467   }
 468   case 4:
 469     EmitInstruction(Out, MCInstBuilder(X86::ADD32ri8)
 470                              .addReg(ScratchRegI32)
 471                              .addReg(ScratchRegI32)
 472                              .addImm(3));
 473     break;
 474   default:
 475     assert(false && "Incorrect access size");
 476     break;
 477   }
 478
 479   EmitInstruction(
 480       Out,
 481       MCInstBuilder(X86::MOVSX32rr8).addReg(ShadowRegI32).addReg(ShadowRegI8));
 482   EmitInstruction(Out, MCInstBuilder(X86::CMP32rr).addReg(ScratchRegI32).addReg(
 483                            ShadowRegI32));
 484   EmitInstruction(Out, MCInstBuilder(X86::JL_4).addExpr(DoneExpr));
 485
 486   EmitCallAsanReport(AccessSize, IsWrite, Ctx, Out, RegCtx);
 487   EmitLabel(Out, DoneSym);
 488 }
 489
 490 void X86AddressSanitizer32::InstrumentMemOperandLarge(
 491     X86Operand &Op, unsigned AccessSize, bool IsWrite,
 492     const RegisterContext &RegCtx, MCContext &Ctx, MCStreamer &Out) {
 493   unsigned AddressRegI32 = RegCtx.addressReg(MVT::i32);
 494   unsigned ShadowRegI32 = RegCtx.shadowReg(MVT::i32);
 495
 496   {
 497     MCInst Inst;
 498     Inst.setOpcode(X86::LEA32r);
 499     Inst.addOperand(MCOperand::CreateReg(AddressRegI32));
 500     Op.addMemOperands(Inst, 5);
 501     EmitInstruction(Out, Inst);
 502   }
 503
 504   EmitInstruction(Out, MCInstBuilder(X86::MOV32rr).addReg(ShadowRegI32).addReg(
 505                            AddressRegI32));
 506   EmitInstruction(Out, MCInstBuilder(X86::SHR32ri)
 507                            .addReg(ShadowRegI32)
 508                            .addReg(ShadowRegI32)
 509                            .addImm(3));
 510   {
 511     MCInst Inst;
 512     switch (AccessSize) {
 513     case 8:
 514       Inst.setOpcode(X86::CMP8mi);
 515       break;
 516     case 16:
 517       Inst.setOpcode(X86::CMP16mi);
 518       break;
 519     default:
 520       assert(false && "Incorrect access size");
 521       break;
 522     }
 523     const MCExpr *Disp = MCConstantExpr::Create(kShadowOffset, Ctx);
 524     std::unique_ptr<X86Operand> Op(
 525         X86Operand::CreateMem(0, Disp, ShadowRegI32, 0, 1, SMLoc(), SMLoc()));
 526     Op->addMemOperands(Inst, 5);
 527     Inst.addOperand(MCOperand::CreateImm(0));
 528     EmitInstruction(Out, Inst);
 529   }
 530   MCSymbol *DoneSym = Ctx.CreateTempSymbol();
 531   const MCExpr *DoneExpr = MCSymbolRefExpr::Create(DoneSym, Ctx);
 532   EmitInstruction(Out, MCInstBuilder(X86::JE_4).addExpr(DoneExpr));
 533
 534   EmitCallAsanReport(AccessSize, IsWrite, Ctx, Out, RegCtx);
 535   EmitLabel(Out, DoneSym);
 536 }
 537
 538 void X86AddressSanitizer32::InstrumentMOVSImpl(unsigned AccessSize,
 539                                                MCContext &Ctx,
 540                                                MCStreamer &Out) {
 541   StoreFlags(Out);
 542
 543   // No need to test when ECX is equals to zero.
 544   MCSymbol *DoneSym = Ctx.CreateTempSymbol();
 545   const MCExpr *DoneExpr = MCSymbolRefExpr::Create(DoneSym, Ctx);
 546   EmitInstruction(
 547       Out, MCInstBuilder(X86::TEST32rr).addReg(X86::ECX).addReg(X86::ECX));
 548   EmitInstruction(Out, MCInstBuilder(X86::JE_4).addExpr(DoneExpr));
 549
 550   // Instrument first and last elements in src and dst range.
 551   InstrumentMOVSBase(X86::EDI /* DstReg */, X86::ESI /* SrcReg */,
 552                      X86::ECX /* CntReg */, AccessSize, Ctx, Out);
 553
 554   EmitLabel(Out, DoneSym);
 555   RestoreFlags(Out);
 556 }
 557
 558 class X86AddressSanitizer64 : public X86AddressSanitizer {
 559 public:
 560   static const long kShadowOffset = 0x7fff8000;
 561
 562   X86AddressSanitizer64(const MCSubtargetInfo &STI)
 563       : X86AddressSanitizer(STI) {}
 564
 565   virtual ~X86AddressSanitizer64() {}
 566
 567   unsigned GetFrameReg(const MCContext &Ctx, MCStreamer &Out) {
 568     unsigned FrameReg = GetFrameRegGeneric(Ctx, Out);
 569     if (FrameReg == X86::NoRegister)
 570       return FrameReg;
 571     return getX86SubSuperRegister(FrameReg, MVT::i64);
 572   }
 573
 574   virtual void StoreFlags(MCStreamer &Out) override {
 575     EmitInstruction(Out, MCInstBuilder(X86::PUSHF64));
 576   }
 577
 578   virtual void RestoreFlags(MCStreamer &Out) override {
 579     EmitInstruction(Out, MCInstBuilder(X86::POPF64));
 580   }
 581
 582   virtual void InstrumentMemOperandPrologue(const RegisterContext &RegCtx,
 583                                             MCContext &Ctx,
 584                                             MCStreamer &Out) override {
 585     const MCRegisterInfo *MRI = Ctx.getRegisterInfo();
 586     unsigned FrameReg = GetFrameReg(Ctx, Out);
 587     if (MRI && FrameReg != X86::NoRegister) {
 588       EmitInstruction(Out, MCInstBuilder(X86::PUSH64r).addReg(X86::RBP));
 589       if (FrameReg == X86::RSP) {
 590         Out.EmitCFIAdjustCfaOffset(8 /* byte size of the FrameReg */);
 591         Out.EmitCFIRelOffset(
 592             MRI->getDwarfRegNum(X86::RBP, true /* IsEH */), 0);
 593       }
 594       EmitInstruction(
 595           Out, MCInstBuilder(X86::MOV64rr).addReg(X86::RBP).addReg(FrameReg));
 596       Out.EmitCFIRememberState();
 597       Out.EmitCFIDefCfaRegister(
 598           MRI->getDwarfRegNum(X86::RBP, true /* IsEH */));
 599     }
 600
 601     EmitAdjustRSP(Ctx, Out, -128);
 602     EmitInstruction(
 603         Out, MCInstBuilder(X86::PUSH64r).addReg(RegCtx.shadowReg(MVT::i64)));
 604     EmitInstruction(
 605         Out, MCInstBuilder(X86::PUSH64r).addReg(RegCtx.addressReg(MVT::i64)));
 606     if (RegCtx.ScratchReg != X86::NoRegister) {
 607       EmitInstruction(
 608           Out, MCInstBuilder(X86::PUSH64r).addReg(RegCtx.scratchReg(MVT::i64)));
 609     }
 610     StoreFlags(Out);
 611   }
 612
 613   virtual void InstrumentMemOperandEpilogue(const RegisterContext &RegCtx,
 614                                             MCContext &Ctx,
 615                                             MCStreamer &Out) override {
 616     RestoreFlags(Out);
 617     if (RegCtx.ScratchReg != X86::NoRegister) {
 618       EmitInstruction(
 619           Out, MCInstBuilder(X86::POP64r).addReg(RegCtx.scratchReg(MVT::i64)));
 620     }
 621     EmitInstruction(
 622         Out, MCInstBuilder(X86::POP64r).addReg(RegCtx.addressReg(MVT::i64)));
 623     EmitInstruction(
 624         Out, MCInstBuilder(X86::POP64r).addReg(RegCtx.shadowReg(MVT::i64)));
 625     EmitAdjustRSP(Ctx, Out, 128);
 626
 627     unsigned FrameReg = GetFrameReg(Ctx, Out);
 628     if (Ctx.getRegisterInfo() && FrameReg != X86::NoRegister) {
 629       EmitInstruction(
 630           Out, MCInstBuilder(X86::POP64r).addReg(X86::RBP));
 631       Out.EmitCFIRestoreState();
 632       if (FrameReg == X86::RSP)
 633         Out.EmitCFIAdjustCfaOffset(-8 /* byte size of the FrameReg */);
 634     }
 635   }
 636
 637   virtual void InstrumentMemOperandSmall(X86Operand &Op, unsigned AccessSize,
 638                                          bool IsWrite,
 639                                          const RegisterContext &RegCtx,
 640                                          MCContext &Ctx,
 641                                          MCStreamer &Out) override;
 642   virtual void InstrumentMemOperandLarge(X86Operand &Op, unsigned AccessSize,
 643                                          bool IsWrite,
 644                                          const RegisterContext &RegCtx,
 645                                          MCContext &Ctx,
 646                                          MCStreamer &Out) override;
 647   virtual void InstrumentMOVSImpl(unsigned AccessSize, MCContext &Ctx,
 648                                   MCStreamer &Out) override;
 649
 650 private:
 651   void EmitAdjustRSP(MCContext &Ctx, MCStreamer &Out, long Offset) {
 652     MCInst Inst;
 653     Inst.setOpcode(X86::LEA64r);
 654     Inst.addOperand(MCOperand::CreateReg(X86::RSP));
 655
 656     const MCExpr *Disp = MCConstantExpr::Create(Offset, Ctx);
 657     std::unique_ptr<X86Operand> Op(
 658         X86Operand::CreateMem(0, Disp, X86::RSP, 0, 1, SMLoc(), SMLoc()));
 659     Op->addMemOperands(Inst, 5);
 660     EmitInstruction(Out, Inst);
 661   }
 662
 663   void EmitCallAsanReport(unsigned AccessSize, bool IsWrite, MCContext &Ctx,
 664                           MCStreamer &Out, const RegisterContext &RegCtx) {
 665     EmitInstruction(Out, MCInstBuilder(X86::CLD));
 666     EmitInstruction(Out, MCInstBuilder(X86::MMX_EMMS));
 667
 668     EmitInstruction(Out, MCInstBuilder(X86::AND64ri8)
 669                              .addReg(X86::RSP)
 670                              .addReg(X86::RSP)
 671                              .addImm(-16));
 672
 673     if (RegCtx.AddressReg != X86::RDI) {
 674       EmitInstruction(Out, MCInstBuilder(X86::MOV64rr).addReg(X86::RDI).addReg(
 675                                RegCtx.addressReg(MVT::i64)));
 676     }
 677     const std::string &Fn = FuncName(AccessSize, IsWrite);
 678     MCSymbol *FnSym = Ctx.GetOrCreateSymbol(StringRef(Fn));
 679     const MCSymbolRefExpr *FnExpr =
 680         MCSymbolRefExpr::Create(FnSym, MCSymbolRefExpr::VK_PLT, Ctx);
 681     EmitInstruction(Out, MCInstBuilder(X86::CALL64pcrel32).addExpr(FnExpr));
 682   }
 683 };
 684
 685 void X86AddressSanitizer64::InstrumentMemOperandSmall(
 686     X86Operand &Op, unsigned AccessSize, bool IsWrite,
 687     const RegisterContext &RegCtx, MCContext &Ctx, MCStreamer &Out) {
 688   unsigned AddressRegI64 = RegCtx.addressReg(MVT::i64);
 689   unsigned AddressRegI32 = RegCtx.addressReg(MVT::i32);
 690   unsigned ShadowRegI64 = RegCtx.shadowReg(MVT::i64);
 691   unsigned ShadowRegI32 = RegCtx.shadowReg(MVT::i32);
 692   unsigned ShadowRegI8 = RegCtx.shadowReg(MVT::i8);
 693
 694   assert(RegCtx.ScratchReg != X86::NoRegister);
 695   unsigned ScratchRegI32 = RegCtx.scratchReg(MVT::i32);
 696
 697   {
 698     MCInst Inst;
 699     Inst.setOpcode(X86::LEA64r);
 700     Inst.addOperand(MCOperand::CreateReg(AddressRegI64));
 701     Op.addMemOperands(Inst, 5);
 702     EmitInstruction(Out, Inst);
 703   }
 704   EmitInstruction(Out, MCInstBuilder(X86::MOV64rr).addReg(ShadowRegI64).addReg(
 705                            AddressRegI64));
 706   EmitInstruction(Out, MCInstBuilder(X86::SHR64ri)
 707                            .addReg(ShadowRegI64)
 708                            .addReg(ShadowRegI64)
 709                            .addImm(3));
 710   {
 711     MCInst Inst;
 712     Inst.setOpcode(X86::MOV8rm);
 713     Inst.addOperand(MCOperand::CreateReg(ShadowRegI8));
 714     const MCExpr *Disp = MCConstantExpr::Create(kShadowOffset, Ctx);
 715     std::unique_ptr<X86Operand> Op(
 716         X86Operand::CreateMem(0, Disp, ShadowRegI64, 0, 1, SMLoc(), SMLoc()));
 717     Op->addMemOperands(Inst, 5);
 718     EmitInstruction(Out, Inst);
 719   }
 720
 721   EmitInstruction(
 722       Out, MCInstBuilder(X86::TEST8rr).addReg(ShadowRegI8).addReg(ShadowRegI8));
 723   MCSymbol *DoneSym = Ctx.CreateTempSymbol();
 724   const MCExpr *DoneExpr = MCSymbolRefExpr::Create(DoneSym, Ctx);
 725   EmitInstruction(Out, MCInstBuilder(X86::JE_4).addExpr(DoneExpr));
 726
 727   EmitInstruction(Out, MCInstBuilder(X86::MOV32rr).addReg(ScratchRegI32).addReg(
 728                            AddressRegI32));
 729   EmitInstruction(Out, MCInstBuilder(X86::AND32ri)
 730                            .addReg(ScratchRegI32)
 731                            .addReg(ScratchRegI32)
 732                            .addImm(7));
 733
 734   switch (AccessSize) {
 735   case 1:
 736     break;
 737   case 2: {
 738     MCInst Inst;
 739     Inst.setOpcode(X86::LEA32r);
 740     Inst.addOperand(MCOperand::CreateReg(ScratchRegI32));
 741
 742     const MCExpr *Disp = MCConstantExpr::Create(1, Ctx);
 743     std::unique_ptr<X86Operand> Op(
 744         X86Operand::CreateMem(0, Disp, ScratchRegI32, 0, 1, SMLoc(), SMLoc()));
 745     Op->addMemOperands(Inst, 5);
 746     EmitInstruction(Out, Inst);
 747     break;
 748   }
 749   case 4:
 750     EmitInstruction(Out, MCInstBuilder(X86::ADD32ri8)
 751                              .addReg(ScratchRegI32)
 752                              .addReg(ScratchRegI32)
 753                              .addImm(3));
 754     break;
 755   default:
 756     assert(false && "Incorrect access size");
 757     break;
 758   }
 759
 760   EmitInstruction(
 761       Out,
 762       MCInstBuilder(X86::MOVSX32rr8).addReg(ShadowRegI32).addReg(ShadowRegI8));
 763   EmitInstruction(Out, MCInstBuilder(X86::CMP32rr).addReg(ScratchRegI32).addReg(
 764                            ShadowRegI32));
 765   EmitInstruction(Out, MCInstBuilder(X86::JL_4).addExpr(DoneExpr));
 766
 767   EmitCallAsanReport(AccessSize, IsWrite, Ctx, Out, RegCtx);
 768   EmitLabel(Out, DoneSym);
 769 }
 770
 771 void X86AddressSanitizer64::InstrumentMemOperandLarge(
 772     X86Operand &Op, unsigned AccessSize, bool IsWrite,
 773     const RegisterContext &RegCtx, MCContext &Ctx, MCStreamer &Out) {
 774   unsigned AddressRegI64 = RegCtx.addressReg(MVT::i64);
 775   unsigned ShadowRegI64 = RegCtx.shadowReg(MVT::i64);
 776
 777   {
 778     MCInst Inst;
 779     Inst.setOpcode(X86::LEA64r);
 780     Inst.addOperand(MCOperand::CreateReg(AddressRegI64));
 781     Op.addMemOperands(Inst, 5);
 782     EmitInstruction(Out, Inst);
 783   }
 784   EmitInstruction(Out, MCInstBuilder(X86::MOV64rr).addReg(ShadowRegI64).addReg(
 785                            AddressRegI64));
 786   EmitInstruction(Out, MCInstBuilder(X86::SHR64ri)
 787                            .addReg(ShadowRegI64)
 788                            .addReg(ShadowRegI64)
 789                            .addImm(3));
 790   {
 791     MCInst Inst;
 792     switch (AccessSize) {
 793     case 8:
 794       Inst.setOpcode(X86::CMP8mi);
 795       break;
 796     case 16:
 797       Inst.setOpcode(X86::CMP16mi);
 798       break;
 799     default:
 800       assert(false && "Incorrect access size");
 801       break;
 802     }
 803     const MCExpr *Disp = MCConstantExpr::Create(kShadowOffset, Ctx);
 804     std::unique_ptr<X86Operand> Op(
 805         X86Operand::CreateMem(0, Disp, ShadowRegI64, 0, 1, SMLoc(), SMLoc()));
 806     Op->addMemOperands(Inst, 5);
 807     Inst.addOperand(MCOperand::CreateImm(0));
 808     EmitInstruction(Out, Inst);
 809   }
 810
 811   MCSymbol *DoneSym = Ctx.CreateTempSymbol();
 812   const MCExpr *DoneExpr = MCSymbolRefExpr::Create(DoneSym, Ctx);
 813   EmitInstruction(Out, MCInstBuilder(X86::JE_4).addExpr(DoneExpr));
 814
 815   EmitCallAsanReport(AccessSize, IsWrite, Ctx, Out, RegCtx);
 816   EmitLabel(Out, DoneSym);
 817 }
 818
 819 void X86AddressSanitizer64::InstrumentMOVSImpl(unsigned AccessSize,
 820                                                MCContext &Ctx,
 821                                                MCStreamer &Out) {
 822   StoreFlags(Out);
 823
 824   // No need to test when RCX is equals to zero.
 825   MCSymbol *DoneSym = Ctx.CreateTempSymbol();
 826   const MCExpr *DoneExpr = MCSymbolRefExpr::Create(DoneSym, Ctx);
 827   EmitInstruction(
 828       Out, MCInstBuilder(X86::TEST64rr).addReg(X86::RCX).addReg(X86::RCX));
 829   EmitInstruction(Out, MCInstBuilder(X86::JE_4).addExpr(DoneExpr));
 830
 831   // Instrument first and last elements in src and dst range.
 832   InstrumentMOVSBase(X86::RDI /* DstReg */, X86::RSI /* SrcReg */,
 833                      X86::RCX /* CntReg */, AccessSize, Ctx, Out);
 834
 835   EmitLabel(Out, DoneSym);
 836   RestoreFlags(Out);
 837 }
 838
 839 } // End anonymous namespace
 840
 841 X86AsmInstrumentation::X86AsmInstrumentation(const MCSubtargetInfo &STI)
 842     : STI(STI), InitialFrameReg(0) {}
 843
 844 X86AsmInstrumentation::~X86AsmInstrumentation() {}
 845
 846 void X86AsmInstrumentation::InstrumentAndEmitInstruction(
 847     const MCInst &Inst, OperandVector &Operands, MCContext &Ctx,
 848     const MCInstrInfo &MII, MCStreamer &Out) {
 849   EmitInstruction(Out, Inst);
 850 }
 851
 852 void X86AsmInstrumentation::EmitInstruction(MCStreamer &Out,
 853                                             const MCInst &Inst) {
 854   Out.EmitInstruction(Inst, STI);
 855 }
 856
 857 unsigned X86AsmInstrumentation::GetFrameRegGeneric(const MCContext &Ctx,
 858                                                    MCStreamer &Out) {
 859   if (!Out.getNumFrameInfos()) // No active dwarf frame
 860     return X86::NoRegister;
 861   const MCDwarfFrameInfo &Frame = Out.getDwarfFrameInfos().back();
 862   if (Frame.End) // Active dwarf frame is closed
 863     return X86::NoRegister;
 864   const MCRegisterInfo *MRI = Ctx.getRegisterInfo();
 865   if (!MRI) // No register info
 866     return X86::NoRegister;
 867
 868   if (InitialFrameReg) {
 869     // FrameReg is set explicitly, we're instrumenting a MachineFunction.
 870     return InitialFrameReg;
 871   }
 872
 873   return MRI->getLLVMRegNum(Frame.CurrentCfaRegister, true /* IsEH */);
 874 }
 875
 876 X86AsmInstrumentation *
 877 CreateX86AsmInstrumentation(const MCTargetOptions &MCOptions,
 878                             const MCContext &Ctx, const MCSubtargetInfo &STI) {
 879   Triple T(STI.getTargetTriple());
 880   const bool hasCompilerRTSupport = T.isOSLinux();
 881   if (ClAsanInstrumentAssembly && hasCompilerRTSupport &&
 882       MCOptions.SanitizeAddress) {
 883     if ((STI.getFeatureBits() & X86::Mode32Bit) != 0)
 884       return new X86AddressSanitizer32(STI);
 885     if ((STI.getFeatureBits() & X86::Mode64Bit) != 0)
 886       return new X86AddressSanitizer64(STI);
 887   }
 888   return new X86AsmInstrumentation(STI);
 889 }
 890
 891 } // End llvm namespace