diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h --- a/bolt/include/bolt/Core/MCPlusBuilder.h +++ b/bolt/include/bolt/Core/MCPlusBuilder.h @@ -1571,6 +1571,13 @@ return {}; } + /// Create code for unpoisoning memory stored by \p StoreInst. + virtual std::vector createUnpoisonMemory(const MCInst &StoreInst, + MCContext *Ctx) { + llvm_unreachable("not implemented"); + return {}; + } + /// Create a target-specific relocation out of the \p Fixup. /// Note that not every fixup could be converted into a relocation. virtual Optional createRelocation(const MCFixup &Fixup, diff --git a/bolt/include/bolt/Passes/BinaryPasses.h b/bolt/include/bolt/Passes/BinaryPasses.h --- a/bolt/include/bolt/Passes/BinaryPasses.h +++ b/bolt/include/bolt/Passes/BinaryPasses.h @@ -466,6 +466,16 @@ void runOnFunctions(BinaryContext &BC) override; }; +class SanitizeMemory : public BinaryFunctionPass { +public: + explicit SanitizeMemory(const cl::opt &PrintPass) + : BinaryFunctionPass(PrintPass) {} + + const char *getName() const override { return "sanitize-memory"; } + + void runOnFunctions(BinaryContext &BC) override; +}; + enum FrameOptimizationType : char { FOP_NONE, /// Don't perform FOP. FOP_HOT, /// Perform FOP on hot functions. diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp --- a/bolt/lib/Passes/BinaryPasses.cpp +++ b/bolt/lib/Passes/BinaryPasses.cpp @@ -1824,5 +1824,53 @@ SkipFunc, "RemoveNops"); } +void SanitizeMemory::runOnFunctions(BinaryContext &BC) { + if (!BC.isX86()) + return; + + if (!BC.HasRelocations) { + errs() << "BOLT-ERROR: relocations required for instrumentation\n"; + exit(1); + } + + // TODO: check that the function is not already instrumented + uint64_t NumInstrumented = 0; + for (auto &BFI : BC.getBinaryFunctions()) { + BinaryFunction &BF = BFI.second; + for (BinaryBasicBlock &BB : BF) { + for (auto II = BB.begin(); II != BB.end(); ++II) { + MCInst &Inst = *II; + + // Skip instructions that do not store to memory. + const MCInstrDesc &Desc = BC.MII->get(Inst.getOpcode()); + if (!Desc.mayStore()) + continue; + + // Memory operand should be the first one. + int MemOp = BC.MIB->getMemoryOperandNo(Inst); + if (MemOp != 0) + continue; + + LLVM_DEBUG({ + dbgs() << "will instrument store: "; + BC.printInstruction(dbgs(), Inst, 0, &BF); + }); + + const std::vector InstrumentationCode = + BC.MIB->createUnpoisonMemory(Inst, BC.Ctx.get()); + II = BB.insertInstruction(std::next(II), MCInst()); + II = BB.replaceInstruction(II, InstrumentationCode); + std::advance(II, InstrumentationCode.size() - 1); + + ++NumInstrumented; + } + } + } + + if (NumInstrumented) + outs() << "BOLT-INFO: instrumented " << NumInstrumented + << " instructions for MemSan\n"; +} + } // namespace bolt } // namespace llvm diff --git a/bolt/lib/Rewrite/BinaryPassManager.cpp b/bolt/lib/Rewrite/BinaryPassManager.cpp --- a/bolt/lib/Rewrite/BinaryPassManager.cpp +++ b/bolt/lib/Rewrite/BinaryPassManager.cpp @@ -193,6 +193,11 @@ "reassign registers so as to avoid using REX prefixes in hot code"), cl::cat(BoltOptCategory)); +static cl::opt SanitizeMemory( + "sanitize-memory", + cl::desc("add instrumentation code for memory sanitizer (X86-only)"), + cl::init(false), cl::ZeroOrMore, cl::cat(BoltCategory)); + static cl::opt SimplifyConditionalTailCalls( "simplify-conditional-tail-calls", cl::desc("simplify conditional tail calls by removing unnecessary jumps"), @@ -336,6 +341,9 @@ Manager.registerPass(std::make_unique(NeverPrint), opts::StripRepRet); + Manager.registerPass(std::make_unique(NeverPrint), + opts::SanitizeMemory); + Manager.registerPass(std::make_unique(PrintICF), opts::ICF); diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp --- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp +++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp @@ -2568,10 +2568,44 @@ return Code; } + std::vector createUnpoisonMemory(const MCInst &StoreInst, + MCContext *Ctx) override { + assert(getMemoryOperandNo(StoreInst) == 0 && "memory expected in first op"); + + // push %rdi + // push %rsi + // lea (memory), %rdi + // movq $8, %rsi + // call __msan_unpoison + // pop %rsi + // pop %rdi + std::vector Code; + Code.resize(7); + unsigned I = 0; + + createPushRegister(Code[I++], X86::RDI, 8); + createPushRegister(Code[I++], X86::RSI, 8); + MCInst &LEA = Code[I++]; + LEA.setOpcode(X86::LEA64r); + LEA.addOperand(MCOperand::createReg(X86::RDI)); + for (unsigned Op = 0; Op < X86::AddrNumOperands; ++Op) + LEA.addOperand(StoreInst.getOperand(Op)); + // TODO: get the proper memory size + Code[I++] = MCInstBuilder(X86::MOV64ri32).addReg(X86::RSI).addImm(8); + MCSymbol *UnpoisonSym = Ctx->getOrCreateSymbol("__msan_unpoison"); + // TODO: check that __msan_unpoison doesn't clobber registers + createCall(Code[I++], UnpoisonSym, Ctx); + createPopRegister(Code[I++], X86::RSI, 8); + createPopRegister(Code[I++], X86::RDI, 8); + + return Code; + } + InstructionListType createCmpJE(MCPhysReg RegNo, int64_t Imm, const MCSymbol *Target, MCContext *Ctx) const override { InstructionListType Code; + Code.emplace_back(MCInstBuilder(X86::CMP64ri8) .addReg(RegNo) .addImm(Imm));