Skip to content

Commit cbd7c5a

Browse files
kylewanginchinakylewanginchina
committed
refactor: Extend autoTLSkey disassembly for python unwind
1 parent ee17409 commit cbd7c5a

1 file changed

Lines changed: 278 additions & 34 deletions

File tree

‎agent/crates/trace-utils/src/unwind/python.rs‎

Lines changed: 278 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ use std::{
1818
cell::OnceCell, collections::HashMap, ffi::CStr, fs, io::Write, mem, path::PathBuf, slice,
1919
};
2020

21+
use ahash::AHashMap;
2122
use libc::c_void;
2223
use log::{debug, trace, warn};
2324
use object::{Object, ObjectSection, ObjectSymbol};
@@ -41,6 +42,62 @@ use super::elf_utils::MappedFile;
4142
#[cfg(target_arch = "x86_64")]
4243
use iced_x86::{Decoder, DecoderOptions, Mnemonic, OpKind, Register};
4344

45+
#[cfg(target_arch = "x86_64")]
46+
#[derive(Clone, Copy, Debug)]
47+
enum Value {
48+
Known(u64),
49+
/// Address expressed as _PyRuntime + offset (can be negative)
50+
RuntimeBase(i64),
51+
Unknown,
52+
}
53+
54+
#[cfg(target_arch = "x86_64")]
55+
impl Value {
56+
fn add(self, rhs: i64) -> Self {
57+
match self {
58+
Value::Known(v) => Value::Known(v.wrapping_add(rhs as u64)),
59+
Value::RuntimeBase(off) => Value::RuntimeBase(off.saturating_add(rhs)),
60+
Value::Unknown => Value::Unknown,
61+
}
62+
}
63+
64+
fn sub(self, rhs: i64) -> Self {
65+
match self {
66+
Value::Known(v) => Value::Known(v.wrapping_sub(rhs as u64)),
67+
Value::RuntimeBase(off) => Value::RuntimeBase(off.saturating_sub(rhs)),
68+
Value::Unknown => Value::Unknown,
69+
}
70+
}
71+
72+
fn shl(self, shift: u8) -> Self {
73+
if shift >= 63 {
74+
return Value::Unknown;
75+
}
76+
match self {
77+
Value::Known(v) => Value::Known(v << shift),
78+
Value::RuntimeBase(off) => {
79+
Value::RuntimeBase(off.checked_shl(shift as u32).unwrap_or(0))
80+
}
81+
Value::Unknown => Value::Unknown,
82+
}
83+
}
84+
85+
fn or(self, rhs: u64) -> Self {
86+
match self {
87+
Value::Known(v) => Value::Known(v | rhs),
88+
_ => Value::Unknown,
89+
}
90+
}
91+
92+
fn to_runtime_addr(self, runtime_addr: u64) -> Option<u64> {
93+
match self {
94+
Value::Known(v) => Some(v),
95+
Value::RuntimeBase(off) => Some(runtime_addr.wrapping_add(off as u64)),
96+
Value::Unknown => None,
97+
}
98+
}
99+
}
100+
44101
fn error_not_python(pid: u32) -> Error {
45102
Error::BadInterpreterType(pid, "python")
46103
}
@@ -275,48 +332,65 @@ impl Interpreter {
275332
runtime_addr: u64,
276333
) -> Result<u64> {
277334
let decoder = Decoder::with_ip(64, code, code_addr, DecoderOptions::NONE);
335+
let mut regs = AHashMap::new();
336+
regs.insert(Register::RAX, Value::RuntimeBase(0));
278337
for instr in decoder {
279-
if instr.mnemonic() == Mnemonic::Lea && instr.op0_register() == Register::RDI {
280-
if let Some(target_addr) = Self::decode_mem_target(&instr) {
338+
Self::propagate_known(&mut regs, &instr);
339+
340+
let op0 = Self::canon_reg(instr.op0_register());
341+
if op0 == Register::RDI {
342+
// Fast path: mov/lea from [_PyRuntime + disp] into EDI/RDI
343+
if instr.op1_kind() == OpKind::Memory && instr.memory_base() == Register::RAX {
344+
let disp = instr.memory_displacement64() as u64;
345+
let target_addr = runtime_addr.wrapping_add(disp);
281346
if let Some(addr) = self.validate_auto_tls_candidate(target_addr, runtime_addr)
282347
{
283-
debug!("Found autoTLSkey address {:#x} (from LEA)", addr);
348+
debug!(
349+
"Found autoTLSkey address {:#x} (mov/lea from RAX base)",
350+
addr
351+
);
284352
return Ok(addr);
285353
}
286354
}
287-
}
288355

289-
if instr.mnemonic() == Mnemonic::Mov && instr.op0_register() == Register::RDI {
290-
match instr.op1_kind() {
291-
OpKind::Immediate64 => {
292-
let target_addr = instr.immediate64();
293-
if let Some(addr) =
294-
self.validate_auto_tls_candidate(target_addr, runtime_addr)
295-
{
296-
debug!("Found autoTLSkey address {:#x} (imm64)", addr);
297-
return Ok(addr);
298-
}
299-
}
300-
OpKind::Immediate32 => {
301-
let target_addr = instr.immediate32() as u64;
302-
if let Some(addr) =
303-
self.validate_auto_tls_candidate(target_addr, runtime_addr)
304-
{
305-
debug!("Found autoTLSkey address {:#x} (imm32)", addr);
306-
return Ok(addr);
356+
if instr.mnemonic() == Mnemonic::Lea && instr.op1_kind() == OpKind::Memory {
357+
if let Some(val) = Self::compute_mem_addr(&instr, &regs).or_else(|| {
358+
Self::assume_runtime_base(&instr, runtime_addr).map(Value::Known)
359+
}) {
360+
if let Some(addr) = val.to_runtime_addr(runtime_addr) {
361+
if let Some(valid) =
362+
self.validate_auto_tls_candidate(addr, runtime_addr)
363+
{
364+
debug!("Found autoTLSkey address {:#x} (LEA)", valid);
365+
return Ok(valid);
366+
}
307367
}
308368
}
309-
OpKind::Memory => {
310-
if let Some(target_addr) = Self::decode_mem_target(&instr) {
311-
if let Some(addr) =
312-
self.validate_auto_tls_candidate(target_addr, runtime_addr)
369+
}
370+
371+
if instr.mnemonic() == Mnemonic::Mov && instr.op1_kind() == OpKind::Memory {
372+
if let Some(val) = Self::compute_mem_addr(&instr, &regs).or_else(|| {
373+
Self::assume_runtime_base(&instr, runtime_addr).map(Value::Known)
374+
}) {
375+
if let Some(addr) = val.to_runtime_addr(runtime_addr) {
376+
if let Some(valid) =
377+
self.validate_auto_tls_candidate(addr, runtime_addr)
313378
{
314-
debug!("Found autoTLSkey address {:#x} (from MOV)", addr);
315-
return Ok(addr);
379+
debug!("Found autoTLSkey address {:#x} (MOV)", valid);
380+
return Ok(valid);
316381
}
317382
}
318383
}
319-
_ => {}
384+
}
385+
386+
if let Some(val) = regs
387+
.get(&Register::RDI)
388+
.and_then(|v| v.to_runtime_addr(runtime_addr))
389+
{
390+
if let Some(addr) = self.validate_auto_tls_candidate(val, runtime_addr) {
391+
debug!("Found autoTLSkey address {:#x}", addr);
392+
return Ok(addr);
393+
}
320394
}
321395
}
322396
}
@@ -353,13 +427,183 @@ impl Interpreter {
353427
}
354428

355429
#[cfg(target_arch = "x86_64")]
356-
fn decode_mem_target(instr: &iced_x86::Instruction) -> Option<u64> {
430+
fn canon_reg(reg: Register) -> Register {
431+
match reg {
432+
Register::RAX | Register::EAX => Register::RAX,
433+
Register::RBX | Register::EBX => Register::RBX,
434+
Register::RCX | Register::ECX => Register::RCX,
435+
Register::RDX | Register::EDX => Register::RDX,
436+
Register::RSI | Register::ESI => Register::RSI,
437+
Register::RDI | Register::EDI => Register::RDI,
438+
Register::R8 | Register::R8D => Register::R8,
439+
Register::R9 | Register::R9D => Register::R9,
440+
Register::R10 | Register::R10D => Register::R10,
441+
Register::R11 | Register::R11D => Register::R11,
442+
Register::R12 | Register::R12D => Register::R12,
443+
Register::R13 | Register::R13D => Register::R13,
444+
Register::R14 | Register::R14D => Register::R14,
445+
Register::R15 | Register::R15D => Register::R15,
446+
_ => reg,
447+
}
448+
}
449+
450+
#[cfg(target_arch = "x86_64")]
451+
fn propagate_known(map: &mut AHashMap<Register, Value>, instr: &iced_x86::Instruction) {
452+
match instr.mnemonic() {
453+
Mnemonic::Mov => {
454+
if instr.op0_kind() != OpKind::Register {
455+
return;
456+
}
457+
let dst = Self::canon_reg(instr.op0_register());
458+
let val = match instr.op1_kind() {
459+
OpKind::Immediate32 => Some(Value::Known(instr.immediate32() as u64)),
460+
OpKind::Immediate64 => Some(Value::Known(instr.immediate64())),
461+
OpKind::Register => map.get(&Self::canon_reg(instr.op1_register())).copied(),
462+
OpKind::Memory => Self::compute_mem_addr(instr, map),
463+
_ => None,
464+
};
465+
if let Some(v) = val {
466+
map.insert(dst, v);
467+
}
468+
}
469+
Mnemonic::Lea => {
470+
if instr.op0_kind() != OpKind::Register || instr.op1_kind() != OpKind::Memory {
471+
return;
472+
}
473+
if let Some(addr) = Self::compute_mem_addr(instr, map) {
474+
map.insert(Self::canon_reg(instr.op0_register()), addr);
475+
}
476+
}
477+
Mnemonic::Add => {
478+
if instr.op0_kind() != OpKind::Register {
479+
return;
480+
}
481+
let dst = Self::canon_reg(instr.op0_register());
482+
if let Some(base) = map.get(&dst).copied() {
483+
let delta = match instr.op1_kind() {
484+
OpKind::Immediate32 => instr.immediate32() as i64,
485+
OpKind::Immediate8 => instr.immediate8() as i64,
486+
OpKind::Register => {
487+
match map.get(&Self::canon_reg(instr.op1_register())).copied() {
488+
Some(Value::Known(v)) => v as i64,
489+
Some(Value::RuntimeBase(off)) => off,
490+
_ => 0,
491+
}
492+
}
493+
_ => 0,
494+
};
495+
map.insert(dst, base.add(delta));
496+
}
497+
}
498+
Mnemonic::Sub => {
499+
if instr.op0_kind() != OpKind::Register {
500+
return;
501+
}
502+
let dst = Self::canon_reg(instr.op0_register());
503+
if let Some(base) = map.get(&dst).copied() {
504+
let delta = match instr.op1_kind() {
505+
OpKind::Immediate32 => instr.immediate32() as i64,
506+
OpKind::Immediate8 => instr.immediate8() as i64,
507+
OpKind::Register => {
508+
match map.get(&Self::canon_reg(instr.op1_register())).copied() {
509+
Some(Value::Known(v)) => v as i64,
510+
Some(Value::RuntimeBase(off)) => off,
511+
_ => 0,
512+
}
513+
}
514+
_ => 0,
515+
};
516+
map.insert(dst, base.sub(delta));
517+
}
518+
}
519+
Mnemonic::Shl => {
520+
if instr.op0_kind() != OpKind::Register || instr.op1_kind() != OpKind::Immediate8 {
521+
return;
522+
}
523+
let dst = Self::canon_reg(instr.op0_register());
524+
if let Some(base) = map.get(&dst).copied() {
525+
let shift = instr.immediate8();
526+
map.insert(dst, base.shl(shift));
527+
}
528+
}
529+
Mnemonic::And => {
530+
if instr.op0_kind() != OpKind::Register {
531+
return;
532+
}
533+
let dst = Self::canon_reg(instr.op0_register());
534+
if let Some(base) = map.get(&dst).copied() {
535+
let mask = match instr.op1_kind() {
536+
OpKind::Immediate32 => instr.immediate32() as u64,
537+
OpKind::Immediate8 => instr.immediate8() as u64,
538+
_ => return,
539+
};
540+
map.insert(dst, base.or(mask));
541+
}
542+
}
543+
Mnemonic::Or => {
544+
if instr.op0_kind() != OpKind::Register {
545+
return;
546+
}
547+
let dst = Self::canon_reg(instr.op0_register());
548+
if let Some(base) = map.get(&dst).copied() {
549+
let val = match instr.op1_kind() {
550+
OpKind::Immediate32 => instr.immediate32() as u64,
551+
OpKind::Immediate8 => instr.immediate8() as u64,
552+
_ => return,
553+
};
554+
map.insert(dst, base.or(val));
555+
}
556+
}
557+
_ => {}
558+
}
559+
}
560+
561+
#[cfg(target_arch = "x86_64")]
562+
fn compute_mem_addr(
563+
instr: &iced_x86::Instruction,
564+
map: &AHashMap<Register, Value>,
565+
) -> Option<Value> {
357566
let disp = instr.memory_displacement64() as i64;
358-
match instr.memory_base() {
359-
Register::RIP => Some(instr.next_ip().wrapping_add(disp as u64)),
360-
Register::None => Some(disp as u64),
361-
_ => None,
567+
let base_val = match instr.memory_base() {
568+
Register::RIP => Value::Known(instr.next_ip()),
569+
Register::None => Value::Known(0),
570+
base => map
571+
.get(&Self::canon_reg(base))
572+
.copied()
573+
.unwrap_or(Value::Unknown),
574+
};
575+
576+
let with_disp = match base_val {
577+
Value::Known(v) => Value::Known(v.wrapping_add(disp as u64)),
578+
Value::RuntimeBase(off) => Value::RuntimeBase(off.saturating_add(disp)),
579+
Value::Unknown => Value::Unknown,
580+
};
581+
582+
let result = if instr.memory_index() != Register::None {
583+
let idx_val = map.get(&Self::canon_reg(instr.memory_index())).copied()?;
584+
let scale = instr.memory_index_scale() as i64;
585+
match (with_disp, idx_val) {
586+
(Value::Known(a), Value::Known(b)) => {
587+
Value::Known(a.wrapping_add((b as i64 * scale) as u64))
588+
}
589+
(Value::RuntimeBase(off), Value::Known(b)) => {
590+
Value::RuntimeBase(off.saturating_add((b as i64 * scale) as i64))
591+
}
592+
_ => Value::Unknown,
593+
}
594+
} else {
595+
with_disp
596+
};
597+
598+
Some(result)
599+
}
600+
601+
#[cfg(target_arch = "x86_64")]
602+
fn assume_runtime_base(instr: &iced_x86::Instruction, runtime_addr: u64) -> Option<u64> {
603+
if instr.memory_base() == Register::RAX {
604+
return Some(runtime_addr.wrapping_add(instr.memory_displacement64() as u64));
362605
}
606+
None
363607
}
364608

365609
/// Decode autoTLSkey address from ARM64 assembly

0 commit comments

Comments
 (0)