Skip to content

"[bolt][aarch64] Fixed indirect call instrumentation snippet" #141918

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions bolt/include/bolt/Core/MCPlusBuilder.h
Original file line number Diff line number Diff line change
Expand Up @@ -511,6 +511,11 @@ class MCPlusBuilder {
llvm_unreachable("not implemented");
}

virtual void createDirectBranch(MCInst &Inst, const MCSymbol *Target,
MCContext *Ctx) {
llvm_unreachable("not implemented");
}

virtual MCPhysReg getX86R11() const { llvm_unreachable("not implemented"); }

virtual unsigned getShortBranchOpcode(unsigned Opcode) const {
Expand Down
18 changes: 12 additions & 6 deletions bolt/lib/Passes/Instrumentation.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -293,9 +293,12 @@ void Instrumentation::instrumentIndirectTarget(BinaryBasicBlock &BB,
BinaryBasicBlock::iterator &Iter,
BinaryFunction &FromFunction,
uint32_t From) {
auto L = FromFunction.getBinaryContext().scopeLock();
const size_t IndCallSiteID = Summary->IndCallDescriptions.size();
createIndCallDescription(FromFunction, From);
size_t IndCallSiteID;
{
auto L = FromFunction.getBinaryContext().scopeLock();
IndCallSiteID = Summary->IndCallDescriptions.size();
createIndCallDescription(FromFunction, From);
}

BinaryContext &BC = FromFunction.getBinaryContext();
bool IsTailCall = BC.MIB->isTailCall(*Iter);
Expand All @@ -305,9 +308,12 @@ void Instrumentation::instrumentIndirectTarget(BinaryBasicBlock &BB,
: IndCallHandlerExitBBFunction->getSymbol(),
IndCallSiteID, &*BC.Ctx);

Iter = BB.eraseInstruction(Iter);
Iter = insertInstructions(CounterInstrs, BB, Iter);
--Iter;
if (!BC.isAArch64()) {
Iter = BB.eraseInstruction(Iter);
Iter = insertInstructions(CounterInstrs, BB, Iter);
--Iter;
} else
Iter = insertInstructions(CounterInstrs, BB, Iter);
}

bool Instrumentation::instrumentOneTarget(
Expand Down
126 changes: 92 additions & 34 deletions bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1966,6 +1966,15 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
convertJmpToTailCall(Inst);
}

void createDirectBranch(MCInst &Inst, const MCSymbol *Target,
MCContext *Ctx) override {
Inst.setOpcode(AArch64::B);
Inst.clear();
Inst.addOperand(MCOperand::createExpr(getTargetExprFor(
Inst, MCSymbolRefExpr::create(Target, MCSymbolRefExpr::VK_None, *Ctx),
*Ctx, 0)));
}

bool analyzeBranch(InstructionIterator Begin, InstructionIterator End,
const MCSymbol *&TBB, const MCSymbol *&FBB,
MCInst *&CondBranch,
Expand Down Expand Up @@ -2328,21 +2337,26 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
}

InstructionListType createInstrumentedIndCallHandlerExitBB() const override {
InstructionListType Insts(5);
// Code sequence for instrumented indirect call handler:
// ldr x1, [sp, #16]
// msr nzcv, x1
// ldp x0, x1, [sp], #16
// ldr x16, [sp], #16
// ldp x0, x1, [sp], #16
// br x16
setSystemFlag(Insts[0], AArch64::X1);
createPopRegisters(Insts[1], AArch64::X0, AArch64::X1);
// Here we load address of the next function which should be called in the
// original binary to X16 register. Writing to X16 is permitted without
// needing to restore.
loadReg(Insts[2], AArch64::X16, AArch64::SP);
createPopRegisters(Insts[3], AArch64::X0, AArch64::X1);
createIndirectBranch(Insts[4], AArch64::X16, 0);
// ret

InstructionListType Insts;

Insts.emplace_back();
loadReg(Insts.back(), AArch64::X1, AArch64::SP);

Insts.emplace_back();
setSystemFlag(Insts.back(), AArch64::X1);

Insts.emplace_back();
createPopRegisters(Insts.back(), AArch64::X0, AArch64::X1);

Insts.emplace_back();
createReturn(Insts.back());

return Insts;
}

Expand Down Expand Up @@ -2418,39 +2432,69 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
MCSymbol *HandlerFuncAddr,
int CallSiteID,
MCContext *Ctx) override {
InstructionListType Insts;
// Code sequence used to enter indirect call instrumentation helper:
// stp x0, x1, [sp, #-16]! createPushRegisters
// stp x0, x1, [sp, #-16]! createPushRegisters (1)
// mov target x0 convertIndirectCallToLoad -> orr x0 target xzr
// mov x1 CallSiteID createLoadImmediate ->
// movk x1, #0x0, lsl #48
// movk x1, #0x0, lsl #32
// movk x1, #0x0, lsl #16
// movk x1, #0x0
// stp x0, x1, [sp, #-16]!
// bl *HandlerFuncAddr createIndirectCall ->
// stp x0, x1, [sp, #-16]! (2)
// adr x0 *HandlerFuncAddr -> adrp + add
// blr x0
// str x30, [sp, #-16]! (3)
// blr x0 (__bolt_instr_ind_call_handler_func)
// ldr x30, sp, #16 (3)
// ldp x0, x1, [sp], #16 (2)
// mov x0, x0 ; move target address to used register
// ldp x0, x1, [sp], #16 (1)

InstructionListType Insts;
Insts.emplace_back();
createPushRegisters(Insts.back(), AArch64::X0, AArch64::X1);
createPushRegisters(Insts.back(), getIntArgRegister(0),
getIntArgRegister(1));
Insts.emplace_back(CallInst);
convertIndirectCallToLoad(Insts.back(), AArch64::X0);
convertIndirectCallToLoad(Insts.back(), getIntArgRegister(0));
InstructionListType LoadImm =
createLoadImmediate(getIntArgRegister(1), CallSiteID);
Insts.insert(Insts.end(), LoadImm.begin(), LoadImm.end());
Insts.emplace_back();
createPushRegisters(Insts.back(), AArch64::X0, AArch64::X1);
createPushRegisters(Insts.back(), getIntArgRegister(0),
getIntArgRegister(1));
Insts.resize(Insts.size() + 2);
InstructionListType Addr =
materializeAddress(HandlerFuncAddr, Ctx, AArch64::X0);
InstructionListType Addr = materializeAddress(
HandlerFuncAddr, Ctx, CallInst.getOperand(0).getReg());
assert(Addr.size() == 2 && "Invalid Addr size");
std::copy(Addr.begin(), Addr.end(), Insts.end() - Addr.size());

Insts.emplace_back();
storeReg(Insts.back(), AArch64::LR, getSpRegister(/*Size*/ 8));

Insts.emplace_back();
createIndirectCallInst(Insts.back(), false,
CallInst.getOperand(0).getReg());

Insts.emplace_back();
createIndirectCallInst(Insts.back(), isTailCall(CallInst), AArch64::X0);
loadReg(Insts.back(), AArch64::LR, getSpRegister(/*Size*/ 8));

// Carry over metadata including tail call marker if present.
stripAnnotations(Insts.back());
moveAnnotations(std::move(CallInst), Insts.back());
Insts.emplace_back();
createPopRegisters(Insts.back(), getIntArgRegister(0),
getIntArgRegister(1));

// move x0 to indirect call register
Insts.emplace_back();
Insts.back().setOpcode(AArch64::ORRXrs);
Insts.back().insert(Insts.back().begin(),
MCOperand::createReg(CallInst.getOperand(0).getReg()));
Insts.back().insert(Insts.back().begin() + 1,
MCOperand::createReg(AArch64::XZR));
Insts.back().insert(Insts.back().begin() + 2,
MCOperand::createReg(getIntArgRegister(0)));
Insts.back().insert(Insts.back().begin() + 3, MCOperand::createImm(0));

Insts.emplace_back();
createPopRegisters(Insts.back(), getIntArgRegister(0),
getIntArgRegister(1));

return Insts;
}
Expand All @@ -2472,30 +2516,44 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
// ldr x30, [sp], #16
// b IndCallHandler
InstructionListType Insts;

Insts.emplace_back();
createPushRegisters(Insts.back(), AArch64::X0, AArch64::X1);
createPushRegisters(Insts.back(), getIntArgRegister(0),
getIntArgRegister(1));

Insts.emplace_back();
getSystemFlag(Insts.back(), getIntArgRegister(1));

Insts.emplace_back();
storeReg(Insts.back(), getIntArgRegister(1), getSpRegister(/*Size*/ 8));

Insts.emplace_back();
Insts.emplace_back();
InstructionListType Addr =
materializeAddress(InstrTrampoline, Ctx, AArch64::X0);
materializeAddress(InstrTrampoline, Ctx, getIntArgRegister(0));
std::copy(Addr.begin(), Addr.end(), Insts.end() - Addr.size());
assert(Addr.size() == 2 && "Invalid Addr size");

Insts.emplace_back();
loadReg(Insts.back(), AArch64::X0, AArch64::X0);
loadReg(Insts.back(), getIntArgRegister(0), getIntArgRegister(0));

InstructionListType cmpJmp =
createCmpJE(AArch64::X0, 0, IndCallHandler, Ctx);
createCmpJE(getIntArgRegister(0), 0, IndCallHandler, Ctx);
Insts.insert(Insts.end(), cmpJmp.begin(), cmpJmp.end());

Insts.emplace_back();
storeReg(Insts.back(), AArch64::LR, AArch64::SP);
storeReg(Insts.back(), AArch64::LR, getSpRegister(/*Size*/ 8));

Insts.emplace_back();
Insts.back().setOpcode(AArch64::BLR);
Insts.back().addOperand(MCOperand::createReg(AArch64::X0));
Insts.back().addOperand(MCOperand::createReg(getIntArgRegister(0)));

Insts.emplace_back();
loadReg(Insts.back(), AArch64::LR, AArch64::SP);
loadReg(Insts.back(), AArch64::LR, getSpRegister(/*Size*/ 8));

Insts.emplace_back();
createDirectCall(Insts.back(), IndCallHandler, Ctx, /*IsTailCall*/ true);
createDirectBranch(Insts.back(), IndCallHandler, Ctx);

return Insts;
}

Expand Down
4 changes: 2 additions & 2 deletions bolt/runtime/instr.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1668,7 +1668,7 @@ extern "C" __attribute((naked)) void __bolt_instr_indirect_call()
#if defined(__aarch64__)
// clang-format off
__asm__ __volatile__(SAVE_ALL
"ldp x0, x1, [sp, #288]\n"
"ldp x0, x1, [sp, #320]\n"
"bl instrumentIndirectCall\n"
RESTORE_ALL
"ret\n"
Expand Down Expand Up @@ -1705,7 +1705,7 @@ extern "C" __attribute((naked)) void __bolt_instr_indirect_tailcall()
#if defined(__aarch64__)
// clang-format off
__asm__ __volatile__(SAVE_ALL
"ldp x0, x1, [sp, #288]\n"
"ldp x0, x1, [sp, #320]\n"
"bl instrumentIndirectCall\n"
RESTORE_ALL
"ret\n"
Expand Down
Loading