Skip to content

fix(python): fix stub decoding routine #412

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.

Sign up for GitHub

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

Merged
florianl merged 28 commits into from
Apr 22, 2025
Merged

fix(python): fix stub decoding routine #412

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
28 commits
Select commit Hold shift + click to select a range
b8f926a
python stub decoding improvements
korniltsev Mar 20, 2025
6a8e9fc
lint
korniltsev Mar 20, 2025
e4f1c2a
Merge branch 'main' into python-fix
korniltsev Mar 20, 2025
aa5ef11
more tests, debug dump code on failure
korniltsev Mar 20, 2025
b0d9b46
lint
korniltsev Mar 21, 2025
d1bf67c
rewrite in go
korniltsev Mar 21, 2025
3eff3a8
lint
korniltsev Apr 4, 2025
93e569e
legal
korniltsev Apr 4, 2025
3c066cc
Merge branch 'main' into python-fix
korniltsev Apr 5, 2025
b4e7b57
remove debug printings
korniltsev Apr 6, 2025
f42edc8
decodeStub: return error, include hexdump of the code into the error
korniltsev Apr 6, 2025
893a170
merge tests into decode_tet.go
korniltsev Apr 6, 2025
2bd6bdd
extract regs state into a new amd package
korniltsev Apr 6, 2025
d0c1b8c
extract endbr64 to amd package
korniltsev Apr 6, 2025
9406305
fmt
korniltsev Apr 6, 2025
507bfd5
get rid of platform specific decode files
korniltsev Apr 6, 2025
f568cda
lint
korniltsev Apr 6, 2025
576fdf4
handle index-scale insns
korniltsev Apr 6, 2025
06daab6
rm printf
korniltsev Apr 6, 2025
6def09d
add coredump tests
korniltsev Apr 6, 2025
36bb7a0
update modulestore to include meaningfull errors
korniltsev Apr 6, 2025
303cd1a
lint
korniltsev Apr 7, 2025
ab12dbd
Apply suggestions from code review
korniltsev Apr 7, 2025
796a617
revert modulestore changes
korniltsev Apr 7, 2025
4372fc3
add endbr64 comment
korniltsev Apr 7, 2025
6573aef
add extra test for lea edit (32bit)
korniltsev Apr 7, 2025
7a1593b
review fixes
korniltsev Apr 14, 2025
6c6f8ca
add r8l-r15l regs
korniltsev Apr 14, 2025
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
152 changes: 104 additions & 48 deletions interpreter/python/decode_amd64.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,40 @@
#include "../../zydis/Zydis.h"
#include "decode_amd64.h"

// #define DECODE_AMD_DEBUG

#if defined(DECODE_AMD_DEBUG)
#include <stdio.h>
#endif

static int reg_index(ZydisRegister reg)
{
switch (reg) {
case ZYDIS_REGISTER_RAX:
case ZYDIS_REGISTER_EAX: return 1;
case ZYDIS_REGISTER_RBX:
case ZYDIS_REGISTER_EBX: return 2;
case ZYDIS_REGISTER_RCX:
case ZYDIS_REGISTER_ECX: return 3;
case ZYDIS_REGISTER_RDX:
case ZYDIS_REGISTER_EDX: return 4;
case ZYDIS_REGISTER_RDI:
case ZYDIS_REGISTER_EDI: return 5;
case ZYDIS_REGISTER_RSI:
case ZYDIS_REGISTER_ESI: return 6;
case ZYDIS_REGISTER_RBP:
case ZYDIS_REGISTER_EBP: return 7;
case ZYDIS_REGISTER_RSP:
case ZYDIS_REGISTER_ESP: return 8;
case ZYDIS_REGISTER_RIP: return 9;
default: return 0;
}
}

struct reg_state {
ZyanU64 loaded_from;
ZyanU64 value;
};

// decode_stub_argument() will decode instructions from given code blob until an assignment
// for the given argument register is found. The value loaded is then determined from the
Expand All @@ -18,66 +51,89 @@
// 3) Loading via pointer + displacement. Happens when the main state is given as argument,
// and the value is loaded from it. In this case 'memory_base' should be the address of
// the global state variable.
uint64_t decode_stub_argument(const uint8_t* code, size_t codesz, uint8_t argument_no,
uint64_t rip_base, uint64_t memory_base) {
// todo update comment
// todo rewrite in go

@gnurizen gnurizen Mar 20, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For what its worth we've gotten pretty good mileage out of go's builtin disassembler: https://github.com/parca-dev/opentelemetry-ebpf-profiler/blob/main/interpreter/luajit/extractor_x86.go

Its shortcomings don't really manifest for simple cases.

@fabled fabled Mar 20, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The long history is that we needed Zydis originally for the stack delta extraction filtering, and the go disassembler did not recognize all opcodes. We wanted to keep only single disassembler so Zydis was used for all other purposes as well. I think the go built in disassembler has improved since, and the stub disassembly is more constrained, this could probably an option. Especially since we use the go builtin disassembler for all arm64 side.

@gnurizen gnurizen Mar 20, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The go x86 disassembler is definitely still broken for anything meaty but I found for simple tasks like this it works well. Personally I find avoiding cgo when possible is worth it and if I was gonna undertake any new decoder tasks I'd probably take a hard look at https://github.com/zyantific/zydis-go.

@korniltsev korniltsev Mar 21, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've drafted a go version.

It's a bit slower but I think it's worth having less unsafe code in the program running as privileged root parsing untrusted user data.

cpu: 13th Gen Intel(R) Core(TM) i7-1365U
               │   old.txt    │               new.txt               │
               │    sec/op    │   sec/op     vs base                │
DecodeAmd64-12   2.000µ ± 13%   2.732µ ± 8%  +36.63% (p=0.000 n=10)

I think rewriting in go should be a separate issue though. I will submit a separate change if we merge #412
Did a quick look into zydis-go and If I understand correctly - it embeds some precompiled native binaries and I'm not excited to use this for the same reasons described above to not use memory unsafe languages and lack of transparency.

@korniltsev korniltsev Mar 21, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Although I'm happy to apply it here in this PR as well. Both go and C versions looks like almost complete rewrite (touches all the lines )

@fabled fabled Apr 4, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@christos68k @florianl @athre0z Do you see any blockers for not using the built-in Go disassembler for amd64?

Currently we have Zydis used in interpreter/php, interpreter/python, and tpbase (for kernel/fsbase and libc/pthread).

The commit referenced would fix php. I'm available to rewrite the three others which are simpler. This would allow cross-running amd64 tests on arm64 builds, and remove the 12M .c drop from this repository.

I believe the Go code will be more maintainable and safer. The original reason for including Zydis was the stack delta extraction code inspecting call return locations to determine which stack deltas to keep and thus had requirements the Go disassembler was not able to handle. But this code is long gone, and I'd like to simplify our codebase now.

Does any anyone object this? If no, I would prefer to review the Go version for this fix.

@christos68k christos68k Apr 4, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SGTM

// todo add coredump tests
uint64_t decode_stub_argument(
const uint8_t *code,
size_t code_sz,
uint64_t code_address,
uint64_t memory_base)
{
ZydisDecoder decoder;
ZydisDecoderInit(&decoder, ZYDIS_MACHINE_MODE_LONG_64, ZYDIS_STACK_WIDTH_64);

// Argument number to x86_64 calling convention register mapping.
ZydisRegister target_register64, target_register32;
switch (argument_no) {
case 0:
target_register64 = ZYDIS_REGISTER_RDI;
target_register32 = ZYDIS_REGISTER_EDI;
break;
case 1:
target_register64 = ZYDIS_REGISTER_RSI;
target_register32 = ZYDIS_REGISTER_ESI;
break;
case 2:
target_register64 = ZYDIS_REGISTER_RDX;
target_register32 = ZYDIS_REGISTER_EDX;
break;
default:
return 0;
}
ZydisRegister target_register64 = ZYDIS_REGISTER_RDI;

// Iterate instructions
ZydisDecodedInstruction instr;
ZydisDecodedOperand operands[ZYDIS_MAX_OPERAND_COUNT];
ZyanUSize instruction_offset = 0;
while (ZYAN_SUCCESS(ZydisDecoderDecodeFull(&decoder, code + instruction_offset,
codesz - instruction_offset, &instr, operands))) {
instruction_offset += instr.length;
if (instr.mnemonic == ZYDIS_MNEMONIC_CALL ||
instr.mnemonic == ZYDIS_MNEMONIC_JMP) {
// Unexpected call/jmp indicating end of stub code
return 0;
}
if (!(instr.mnemonic == ZYDIS_MNEMONIC_LEA ||
instr.mnemonic == ZYDIS_MNEMONIC_MOV) ||
operands[0].type != ZYDIS_OPERAND_TYPE_REGISTER ||
(operands[0].reg.value != target_register64 &&
operands[0].reg.value != target_register32)) {
// Only "LEA/MOV target_reg, ..." meaningful
continue;
struct reg_state regs[32] = {};
while (ZYAN_SUCCESS(
ZydisDecoderDecodeFull(
&decoder, code + instruction_offset, code_sz - instruction_offset, &instr, operands))) {
#if defined(DECODE_AMD_DEBUG)
ZydisDisassembledInstruction dbgi = {};
if (ZYAN_SUCCESS(
ZydisDisassembleIntel(
ZYDIS_MACHINE_MODE_LONG_64,
code_address + instruction_offset,
code + instruction_offset,
code_sz - instruction_offset,
&dbgi))) {
printf("%-12p %s\n", (void *)(code_address + instruction_offset), dbgi.text);
fflush(stdout);
}
if (operands[1].type == ZYDIS_OPERAND_TYPE_IMMEDIATE) {
// MOV target_reg, immediate
return operands[1].imm.value.u;
#endif
instruction_offset += instr.length;
regs[reg_index(ZYDIS_REGISTER_RIP)].value = code_address + instruction_offset;
if (instr.mnemonic == ZYDIS_MNEMONIC_CALL || instr.mnemonic == ZYDIS_MNEMONIC_JMP) {
if (regs[reg_index(target_register64)].loaded_from) {
return regs[reg_index(target_register64)].loaded_from;
}
return regs[reg_index(target_register64)].value;
}
if (operands[1].type == ZYDIS_OPERAND_TYPE_MEMORY &&
operands[1].mem.disp.has_displacement) {
if (operands[1].mem.base == ZYDIS_REGISTER_RIP) {
// MOV/LEA target_reg, [RIP + XXXX]
return rip_base + instruction_offset + operands[1].mem.disp.value;
} else if (memory_base) {
// MOV/LEA target_reg, [REG + XXXX]
return memory_base + operands[1].mem.disp.value;
if (
(instr.mnemonic == ZYDIS_MNEMONIC_LEA || instr.mnemonic == ZYDIS_MNEMONIC_MOV) &&
operands[0].type == ZYDIS_OPERAND_TYPE_REGISTER) {

ZyanU64 v = 0;
ZyanU64 loaded_from = 0;
if (operands[1].type == ZYDIS_OPERAND_TYPE_IMMEDIATE) {
v = operands[1].imm.value.u;
}
if (operands[1].type == ZYDIS_OPERAND_TYPE_MEMORY && operands[1].mem.disp.has_displacement) {
ZyanU64 at = regs[reg_index(operands[1].mem.base)].value + operands[1].mem.disp.value;
if (instr.mnemonic == ZYDIS_MNEMONIC_MOV) {
// todo: do not assume that we're reading a ptr to memory_base
v = memory_base;
loaded_from = at;
}
if (instr.mnemonic == ZYDIS_MNEMONIC_LEA) {
v = at;
}
}
continue;
if (operands[1].type == ZYDIS_OPERAND_TYPE_REGISTER) {
v = regs[reg_index(operands[1].reg.value)].value;
}
#if defined(DECODE_AMD_DEBUG)
printf(" | regs[%d] = %lx\n", reg_index(operands[0].reg.value), v);
#endif
regs[reg_index(operands[0].reg.value)].value = v;
regs[reg_index(operands[0].reg.value)].loaded_from = loaded_from;
}
if (instr.mnemonic == ZYDIS_MNEMONIC_ADD && instr.operand_count == 3 && operands[0].type ==
ZYDIS_OPERAND_TYPE_REGISTER && operands[1].type == ZYDIS_OPERAND_TYPE_MEMORY) {
// todo: do not assume that we're reading a ptr to memory_base
ZyanU64 v = regs[reg_index(operands[0].reg.value)].value + memory_base;
regs[reg_index(operands[0].reg.value)].value = v;
regs[reg_index(operands[0].reg.value)].loaded_from = 0;
#if defined(DECODE_AMD_DEBUG)
printf(" | regs[%d] = %lx\n", reg_index(operands[0].reg.value), v);
#endif
}
}

return 0;
}
}
13 changes: 7 additions & 6 deletions interpreter/python/decode_amd64.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,15 @@ import (
// #include "../../support/ebpf/types.h"
import "C"

func decodeStubArgumentWrapperX64(code []byte, argNumber uint8, symbolValue,
func decodeStubArgumentWrapperX64(code []byte, symbolValue,
addrBase libpf.SymbolValue) libpf.SymbolValue {
if len(code) == 0 {
return 0
}
return libpf.SymbolValue(C.decode_stub_argument(
(*C.uint8_t)(unsafe.Pointer(&code[0])), C.size_t(len(code)),
C.uint8_t(argNumber), C.uint64_t(symbolValue), C.uint64_t(addrBase)))
(*C.uint8_t)(unsafe.Pointer(&code[0])), C.size_t(len(code)), C.uint64_t(symbolValue), C.uint64_t(addrBase)))
}

func decodeStubArgumentWrapper(code []byte, argNumber uint8, symbolValue,
addrBase libpf.SymbolValue) libpf.SymbolValue {
return decodeStubArgumentWrapperX64(code, argNumber, symbolValue, addrBase)
func decodeStubArgumentWrapper(code []byte, symbolValue, addrBase libpf.SymbolValue) libpf.SymbolValue {
return decodeStubArgumentWrapperX64(code, symbolValue, addrBase)
}
2 changes: 1 addition & 1 deletion interpreter/python/decode_amd64.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,6 @@

#include <stdint.h>

uint64_t decode_stub_argument(const uint8_t* code, size_t codesz, uint8_t argument_no, uint64_t rip_base, uint64_t memory_base);
uint64_t decode_stub_argument(const uint8_t* code, size_t code_sz, uint64_t code_address, uint64_t memory_base);

#endif
4 changes: 2 additions & 2 deletions interpreter/python/decode_arm64.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import (
"go.opentelemetry.io/ebpf-profiler/libpf"
)

func decodeStubArgumentWrapper(code []byte, argNumber uint8, symbolValue,
func decodeStubArgumentWrapper(code []byte, symbolValue,
addrBase libpf.SymbolValue) libpf.SymbolValue {
return decodeStubArgumentWrapperARM64(code, argNumber, symbolValue, addrBase)
return decodeStubArgumentWrapperARM64(code, 0, symbolValue, addrBase)
}
102 changes: 102 additions & 0 deletions interpreter/python/decode_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,3 +42,105 @@ func TestAnalyzeArm64Stubs(t *testing.T) {
0, 0, 0)
assert.Equal(t, libpf.SymbolValue(604), val, "PyGILState_GetThisThreadState test")
}

func TestAmd64(t *testing.T) {
testdata := []struct {
name string
code []byte
mem uint64
rip uint64
expected uint64
}{
{
name: "cpython 3.12 ",
code: []byte{

0xF3, 0x0F, 0x1E, 0xFA, // endbr64
0x53, // push rbx
0x48, 0x8D, 0x1D, 0x1C, 0x42, 0x37, // lea rbx, __TMC_END__.autoTSSkey
0x00, //
0x48, 0x89, 0xDF, // mov rdi, rbx
0xE8, 0x8C, 0x83, 0x01, 0x00, // call PyThread_tss_is_created
0x85, 0xC0, // test eax, eax
0x74, 0x10, // jz short loc_2F1928
0x48, 0x89, 0xDF, // mov rdi, rbx
0x5B, // pop rbx
},
mem: 0, // not unused
rip: 0x2F1900,
expected: 0x665B28,
},
{
name: "cpython 3.10 PyGILState_GetThisThreadState",
code: []byte{
0xF3, 0x0F, 0x1E, 0xFA, // endbr64
0x48, 0x8B, 0x05, 0xB5, 0x2D, // mov rax, cs:runtime << mem
0x18, 0x00, //
0x48, 0x83, 0xB8, 0x40, 0x02, // cmp qword ptr [rax+240h], 0
0x00, 0x00, 0x00, //
0x74, 0x13, // jz short loc_209F98
0x48, 0x8D, 0xB8, 0x48, 0x02, // lea rdi, [rax+248h] ; key
0x00, 0x00, //
0xE9, 0xFF, 0xF9, 0xE5, 0xFF, // jmp _PyThread_tss_get
},
mem: 0x3C1680,
rip: 0x209F70,
expected: 0x3C18C8,
},
{
name: "cpython 3.11.2 PyGILState_GetThisThreadState google/cloud-sdk:502.0.0-slim",
code: []byte{
0x48, 0x83, 0x3D, 0x00, 0x47, // cmp cs:qword_A5C968, 0
0x56, 0x00, 0x00, //
0x0F, 0x84, 0xEF, 0xC1, 0xF2, // jz loc_42445D
0xFF, //
0x8B, 0x3D, 0x00, 0x47, 0x56, // mov edi, cs:dword_A5C974
0x00, //
0xE9, 0x77, 0x84, 0xF2, 0xFF, // jmp _pthread_getspecific
},
mem: 0, // not used
rip: 0x4F8260,
expected: 0xA5C974,
},
{
name: "gcloud-sdk 515.0.0 3.12 bundled",
code: []byte{
0x53, // push rbx
0xBB, 0x08, 0x06, 0x00, 0x00, // mov ebx, 608h
0x48, 0x03, 0x1D, 0xBB, 0x10, 0x1A, 0x01, // add rbx, cs:_PyRuntime_ptr << mem
0x48, 0x89, 0xDF, // mov rdi, rbx
0xE8, 0x6B, 0x81, 0xE5, 0xFF, // call _PyThread_tss_is_created
0x85, 0xC0, // test eax, eax
0x74, 0x09, // jz short loc_3C89E2
0x48, 0x89, 0xDF, // mov rdi, rbx
0x5B, // pop rbx
0xE9, 0x3E, 0xA3, 0xE5, 0xFF, // jmp _PyThread_tss_get
},
mem: 0x16905C0,
rip: 0x3C89C0,
expected: 0x16905C0 + 0x608,
},
{
name: "gcloud-sdk 502 3.11 bundled",
code: []byte{
0x48, 0x8B, 0x05, 0x61, 0x9D, 0x12, 0x01, // mov rax, cs:_PyRuntime_ptr
0x48, 0x83, 0xB8, 0x48, 0x02, 0x00, 0x00, 0x00, // cmp qword ptr [rax+248h], 0
0x74, 0x11, // jz short loc_377D42
0xBF, 0x50, 0x02, 0x00, 0x00, // mov edi, 250h
0x48, 0x03, 0x3D, 0x4B, 0x9D, 0x12, 0x01, // add rdi, cs:_PyRuntime_ptr << mem
0xE9, 0x5E, 0x6E, 0xE9, 0xFF, // jmp _PyThread_tss_get
},
mem: 0x15D36F0,
rip: 0x77D20,
expected: 0x15D36F0 + 0x250,
},
}

for _, td := range testdata {
t.Run(td.name, func(t *testing.T) {
val := decodeStubArgumentWrapperX64(
td.code, libpf.SymbolValue(td.rip), libpf.SymbolValue(td.mem))
assert.Equal(t, libpf.SymbolValue(td.expected), val)
})
}
}
7 changes: 3 additions & 4 deletions interpreter/python/python.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -652,8 +652,7 @@ func (d *pythonData) readIntrospectionData(ef *pfelf.File, symbol libpf.SymbolNa

// decodeStub will resolve a given symbol, extract the code for it, and analyze
// the code to resolve specified argument parameter to the first jump/call.
func decodeStub(ef *pfelf.File, addrBase libpf.SymbolValue, symbolName libpf.SymbolName,
argNumber uint8) libpf.SymbolValue {
func decodeStub(ef *pfelf.File, addrBase libpf.SymbolValue, symbolName libpf.SymbolName) libpf.SymbolValue {
symbolValue, err := ef.LookupSymbolAddress(symbolName)
if err != nil {
return libpf.SymbolValueInvalid
Expand All @@ -664,7 +663,7 @@ func decodeStub(ef *pfelf.File, addrBase libpf.SymbolValue, symbolName libpf.Sym
return libpf.SymbolValueInvalid
}

value := decodeStubArgumentWrapper(code, argNumber, symbolValue, addrBase)
value := decodeStubArgumentWrapper(code, symbolValue, addrBase)

// Sanity check the value range and alignment
if value%4 != 0 {
Expand Down Expand Up @@ -733,7 +732,7 @@ func Loader(ebpf interpreter.EbpfHandler, info *interpreter.LoaderInfo) (interpr
}

// Calls first: PyThread_tss_get(autoTSSKey)
autoTLSKey = decodeStub(ef, pyruntimeAddr, "PyGILState_GetThisThreadState", 0)
autoTLSKey = decodeStub(ef, pyruntimeAddr, "PyGILState_GetThisThreadState")
if autoTLSKey == libpf.SymbolValueInvalid {
return nil, errors.New("unable to resolve autoTLSKey")
}
Expand Down