[BOLT] Fix intermittent crash with instrumentation

[maksfb]

Summary

When createInstrumentedIndirectCall() was invoked for tail calls, we
attached annotation instruction twice to the new call instruction.
First in createDirectCall(), and then again while copying over the
metadata operands.

As a result, the annotations were not properly stripped for such calls
before the call to freeAnnotations() in LowerAnnotations pass. That lead
to use-after-free while restoring the offsets with setOffset() call.

[yota9]

bolt/include/bolt/Core/MCPlusBuilder.h

181

Maybe add assert(!getAnnotationInst(DstInst) just for extra safety

182

How about to make separate function smth like "getAnnotationOpIndex" and use it also in getAnnotationInst and stripAnnotations? Just to ensure everything is in sync

[maksfb]

bolt/include/bolt/Core/MCPlusBuilder.h

181

We have it at line 168.

182

Operand handling should be happening in MCPlusBuilder annotation interface (after this change). I can add removeAnnotationInst(), but it still will be a part of the internal interface. The direct handling of operands in createInstrumentedIndirectCall() was really unfortunate.

[yota9]

LGTM

[maksfb]

Add MCPlusBuilder::removeAnnotationInst()

bolt/include/bolt/Core/MCPlusBuilder.h

if (!LastOp.isInst())
       return nullptr;
 
     MCInst *AnnotationInst = const_cast<MCInst *>(LastOp.getInst());
     assert(AnnotationInst->getOpcode() == TargetOpcode::ANNOTATION_LABEL);
 
     return AnnotationInst;
   }
 
+  void removeAnnotationInst(MCInst &Inst) const {
+    assert(getAnnotationInst(Inst) && "Expected annotation instruction.");
+    Inst.erase(std::prev(Inst.end()));
+    assert(!getAnnotationInst(Inst) &&
+           "More than one annotation instruction detected.");
+  }
+
   void setAnnotationOpValue(MCInst &Inst, unsigned Index, int64_t Value,
                             AllocatorIdTy AllocatorId = 0) {
     MCInst *AnnotationInst = getAnnotationInst(Inst);
     if (!AnnotationInst) {
       AnnotationAllocator &Allocator = getAnnotationAllocator(AllocatorId);
       AnnotationInst = new (Allocator.MCInstAllocator.Allocate()) MCInst();
       AnnotationInst->setOpcode(TargetOpcode::ANNOTATION_LABEL);
       Inst.addOperand(MCOperand::createInst(AnnotationInst));
protected:
 
   /// Names of non-standard annotations.
   SmallVector<std::string, 8> AnnotationNames;
 
   /// Allocate the TailCall annotation value. Clients of the target-specific
   /// MCPlusBuilder classes must use convert/lower/create* interfaces instead.
   void setTailCall(MCInst &Inst);
 
+  /// Transfer annotations from \p SrcInst to \p DstInst.
+  void moveAnnotations(MCInst &&SrcInst, MCInst &DstInst) const {
+    assert(!getAnnotationInst(DstInst) &&
+           "Destination instruction should not have annotations.");
+    const MCInst *AnnotationInst = getAnnotationInst(SrcInst);
+    if (!AnnotationInst)
+      return;
+
+    DstInst.addOperand(MCOperand::createInst(AnnotationInst));
+    removeAnnotationInst(SrcInst);
+  }
+
 public:
   class InstructionIterator {
   public:
     using iterator_category = std::bidirectional_iterator_tag;
     using value_type = MCInst;
     using difference_type = std::ptrdiff_t;
     using pointer = value_type *;
     using reference = value_type &;
if (!Index)
       return false;
     return removeAnnotation(Inst, *Index);
   }
 
   /// Remove meta-data, but don't destroy it.
   void stripAnnotations(MCInst &Inst, bool KeepTC = false);
 
   virtual InstructionListType
-  createInstrumentedIndirectCall(const MCInst &CallInst, bool TailCall,
-                                 MCSymbol *HandlerFuncAddr, int CallSiteID,
-                                 MCContext *Ctx) {
+  createInstrumentedIndirectCall(MCInst &&CallInst, MCSymbol *HandlerFuncAddr,
+                                 int CallSiteID, MCContext *Ctx) {
     llvm_unreachable("not implemented");
     return InstructionListType();
   }
 
   virtual InstructionListType createInstrumentedIndCallHandlerExitBB() const {
     llvm_unreachable("not implemented");
     return InstructionListType();
   }

bolt/lib/Core/MCPlusBuilder.cpp

 
 void MCPlusBuilder::stripAnnotations(MCInst &Inst, bool KeepTC) {
   MCInst *AnnotationInst = getAnnotationInst(Inst);
   if (!AnnotationInst)
     return;
   // Preserve TailCall annotation.
   auto IsTC = hasAnnotation(Inst, MCAnnotation::kTailCall);
 
-  Inst.erase(std::prev(Inst.end()));
+  removeAnnotationInst(Inst);
+
   if (KeepTC && IsTC)
     setTailCall(Inst);
 }
 
 void MCPlusBuilder::printAnnotations(const MCInst &Inst,
                                      raw_ostream &OS) const {
   const MCInst *AnnotationInst = getAnnotationInst(Inst);
   if (!AnnotationInst)

bolt/lib/Passes/Instrumentation.cpp

void Instrumentation::instrumentIndirectTarget(BinaryBasicBlock &BB,
                                                uint32_t From) {
   auto L = FromFunction.getBinaryContext().scopeLock();
   const size_t IndCallSiteID = Summary->IndCallDescriptions.size();
   createIndCallDescription(FromFunction, From);
 
   BinaryContext &BC = FromFunction.getBinaryContext();
   bool IsTailCall = BC.MIB->isTailCall(*Iter);
   InstructionListType CounterInstrs = BC.MIB->createInstrumentedIndirectCall(
-      *Iter, IsTailCall,
+      std::move(*Iter),
       IsTailCall ? IndTailCallHandlerExitBBFunction->getSymbol()
                  : IndCallHandlerExitBBFunction->getSymbol(),
       IndCallSiteID, &*BC.Ctx);
 
   Iter = BB.eraseInstruction(Iter);
   Iter = insertInstructions(CounterInstrs, BB, Iter);
   --Iter;
 }

bolt/lib/Target/X86/X86MCPlusBuilder.cpp

void createIndirectBranch(MCInst &Inst, MCPhysReg MemBaseReg,
     Inst.setOpcode(X86::JMP64m);
     Inst.addOperand(MCOperand::createReg(MemBaseReg));      // BaseReg
     Inst.addOperand(MCOperand::createImm(1));               // ScaleAmt
     Inst.addOperand(MCOperand::createReg(X86::NoRegister)); // IndexReg
     Inst.addOperand(MCOperand::createImm(Disp));            // Displacement
     Inst.addOperand(MCOperand::createReg(X86::NoRegister)); // AddrSegmentReg
   }
 
-  InstructionListType createInstrumentedIndirectCall(const MCInst &CallInst,
-                                                     bool TailCall,
+  InstructionListType createInstrumentedIndirectCall(MCInst &&CallInst,
                                                      MCSymbol *HandlerFuncAddr,
                                                      int CallSiteID,
                                                      MCContext *Ctx) override {
     // Check if the target address expression used in the original indirect call
     // uses the stack pointer, which we are going to clobber.
     static BitVector SPAliases(getAliases(X86::RSP));
     bool UsesSP = false;
     // Skip defs.
if (UsesSP) {
       createStackPointerIncrement(Insts.back(), 8, /*NoFlagsClobber=*/false);
     }
     Insts.emplace_back();
     createPushRegister(Insts.back(), TempReg, 8);
     Insts.emplace_back();
     createLoadImmediate(Insts.back(), TempReg, CallSiteID);
     Insts.emplace_back();
     createPushRegister(Insts.back(), TempReg, 8);
-    Insts.emplace_back();
-    createDirectCall(Insts.back(), HandlerFuncAddr, Ctx,
-                     /*TailCall=*/TailCall);
-    // Carry over metadata
-    for (int I = MCPlus::getNumPrimeOperands(CallInst),
-             E = CallInst.getNumOperands();
-         I != E; ++I)
-      Insts.back().addOperand(CallInst.getOperand(I));
+
+    MCInst &NewCallInst = Insts.emplace_back();
+    createDirectCall(NewCallInst, HandlerFuncAddr, Ctx, isTailCall(CallInst));
+
+    // Carry over metadata including tail call marker if present.
+    stripAnnotations(NewCallInst);
+    moveAnnotations(std::move(CallInst), NewCallInst);
 
     return Insts;
   }
 
   InstructionListType createInstrumentedIndCallHandlerExitBB() const override {
     const MCPhysReg TempReg = getIntArgRegister(0);
     // We just need to undo the sequence created for every ind call in
     // instrumentIndirectTarget(), which can be accomplished minimally with: