[BOLT][HACK] Add memory unpoisoning

[maksfb]

Summary

Prototype for memsan instrumentation.

bolt/include/bolt/Core/MCPlusBuilder.h

 
   /// Creates inline memcpy instruction. If \p ReturnEnd is true, then return
   /// (dest + n) instead of dest.
   virtual InstructionListType createInlineMemcpy(bool ReturnEnd) const {
     llvm_unreachable("not implemented");
     return {};
   }
 
+  /// Create code for unpoisoning memory stored by \p StoreInst.
+  virtual std::vector<MCInst> 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<Relocation> createRelocation(const MCFixup &Fixup,
                                                 const MCAsmBackend &MAB) const {
     llvm_unreachable("not implemented");
     return Relocation();
   }
 

bolt/include/bolt/Passes/BinaryPasses.h

explicit RemoveNops(const cl::opt<bool> &PrintPass)
       : BinaryFunctionPass(PrintPass) {}
 
   const char *getName() const override { return "remove-nops"; }
 
   /// Pass entry point
   void runOnFunctions(BinaryContext &BC) override;
 };
 
+class SanitizeMemory : public BinaryFunctionPass {
+public:
+  explicit SanitizeMemory(const cl::opt<bool> &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.
   FOP_ALL   /// Perform FOP on all functions.
 };
 
 } // namespace bolt
 } // namespace llvm
 
 #endif

bolt/lib/Passes/BinaryPasses.cpp

ParallelUtilities::PredicateTy SkipFunc = [&](const BinaryFunction &BF) {
     return BF.shouldPreserveNops();
   };
 
   ParallelUtilities::runOnEachFunction(
       BC, ParallelUtilities::SchedulingPolicy::SP_INST_LINEAR, WorkFun,
       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<MCInst> 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

bolt/lib/Rewrite/BinaryPassManager.cpp

PrintUCE("print-uce",
              cl::Hidden, cl::cat(BoltOptCategory));
 
 static cl::opt<bool> RegReAssign(
     "reg-reassign",
     cl::desc(
         "reassign registers so as to avoid using REX prefixes in hot code"),
     cl::cat(BoltOptCategory));
 
+static cl::opt<bool> SanitizeMemory(
+    "sanitize-memory",
+    cl::desc("add instrumentation code for memory sanitizer (X86-only)"),
+    cl::init(false), cl::ZeroOrMore, cl::cat(BoltCategory));
+
 static cl::opt<bool> SimplifyConditionalTailCalls(
     "simplify-conditional-tail-calls",
     cl::desc("simplify conditional tail calls by removing unnecessary jumps"),
     cl::init(true), cl::cat(BoltOptCategory));
 
 static cl::opt<bool> SimplifyRODataLoads(
     "simplify-rodata-loads",
     cl::desc("simplify loads from read-only sections by replacing the memory "
void BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
 
   Manager.registerPass(std::make_unique<RemoveNops>(NeverPrint));
 
   Manager.registerPass(std::make_unique<NormalizeCFG>(PrintNormalized));
 
   Manager.registerPass(std::make_unique<StripRepRet>(NeverPrint),
                        opts::StripRepRet);
 
+  Manager.registerPass(std::make_unique<SanitizeMemory>(NeverPrint),
+                       opts::SanitizeMemory);
+
   Manager.registerPass(std::make_unique<IdenticalCodeFolding>(PrintICF),
                        opts::ICF);
 
   if (BC.isAArch64())
     Manager.registerPass(
         std::make_unique<VeneerElimination>(PrintVeneerElimination));
 
   Manager.registerPass(

bolt/lib/Target/X86/X86MCPlusBuilder.cpp

Code.emplace_back(MCInstBuilder(X86::MOV8mr)
                           .addReg(X86::NoRegister)
                           .addReg(X86::CL));
     Code.emplace_back(MCInstBuilder(X86::MOV64rr)
                           .addReg(X86::RAX)
                           .addReg(X86::RDI));
     return Code;
   }
 
+  std::vector<MCInst> 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<MCInst> 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));
     Code.emplace_back(MCInstBuilder(X86::JCC_1)
                           .addExpr(MCSymbolRefExpr::create(
                               Target, MCSymbolRefExpr::VK_None, *Ctx))
                           .addImm(X86::COND_E));
     return Code;