fix(python): fix stub decoding routine

asm/amd/insn.go

@@ -0,0 +1,20 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package amd // import "go.opentelemetry.io/ebpf-profiler/asm/amd"
+import "bytes"
+
+// https://www.felixcloutier.com/x86/endbr64
+var opcodeEndBr64 = []byte{0xf3, 0x0f, 0x1e, 0xfa}
+
+// DecodeSkippable decodes an instruction that we don't care much about and are going to skip,
+// as golang.org/x/arch/x86/x86asm fails to decode it.
+// The second returned argument is the size of the decoded instruction to skip.
+func DecodeSkippable(code []byte) (ok bool, size int) {
+	switch {
+	case bytes.HasPrefix(code, opcodeEndBr64):
+		return true, len(opcodeEndBr64)
+	default:
+		return false, 0
+	}
+}

asm/amd/insn_test.go

@@ -0,0 +1,19 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package amd
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestEndBr64(t *testing.T) {
+	res, n := DecodeSkippable([]byte{0xF3, 0x0F, 0x1E, 0xFA})
+	assert.True(t, res)
+	assert.Equal(t, 4, n)
+
+	res, _ = DecodeSkippable([]byte{})
+	assert.False(t, res)
+}

asm/amd/regs_state.go

@@ -0,0 +1,66 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package amd // import "go.opentelemetry.io/ebpf-profiler/asm/amd"
+
+import "golang.org/x/arch/x86/x86asm"
+
+// regIndex returns index into RegsState.regs
+func regIndex(reg x86asm.Reg) int {
+	switch reg {
+	case x86asm.RAX, x86asm.EAX:
+		return 1
+	case x86asm.RBX, x86asm.EBX:
+		return 2
+	case x86asm.RCX, x86asm.ECX:
+		return 3
+	case x86asm.RDX, x86asm.EDX:
+		return 4
+	case x86asm.RDI, x86asm.EDI:
+		return 5
+	case x86asm.RSI, x86asm.ESI:
+		return 6
+	case x86asm.RBP, x86asm.EBP:
+		return 7
+	case x86asm.R8, x86asm.R8L:
+		return 8
+	case x86asm.R9, x86asm.R9L:
+		return 9
+	case x86asm.R10, x86asm.R10L:
+		return 10
+	case x86asm.R11, x86asm.R11L:
+		return 11
+	case x86asm.R12, x86asm.R12L:
+		return 12
+	case x86asm.R13, x86asm.R13L:
+		return 13
+	case x86asm.R14, x86asm.R14L:
+		return 14
+	case x86asm.R15, x86asm.R15L:
+		return 15
+	case x86asm.RSP, x86asm.ESP:
+		return 16
+	case x86asm.RIP:
+		return 17
+	default:
+		return 0
+	}
+}
+
+type RegsState struct {
+	regs [18]regState
+}
+
+func (r *RegsState) Set(reg x86asm.Reg, value, loadedFrom uint64) {
+	r.regs[regIndex(reg)].Value = value
+	r.regs[regIndex(reg)].LoadedFrom = loadedFrom
+}
+
+func (r *RegsState) Get(reg x86asm.Reg) (value, loadedFrom uint64) {
+	return r.regs[regIndex(reg)].Value, r.regs[regIndex(reg)].LoadedFrom
+}
+
+type regState struct {
+	LoadedFrom uint64
+	Value      uint64
+}

interpreter/python/decode.go

@@ -4,16 +4,22 @@
 package python // import "go.opentelemetry.io/ebpf-profiler/interpreter/python"
 
 import (
-	ah "go.opentelemetry.io/ebpf-profiler/armhelpers"
-	aa "golang.org/x/arch/arm64/arm64asm"
+	"errors"
+	"fmt"
+	"runtime"
 
+	ah "go.opentelemetry.io/ebpf-profiler/armhelpers"
+	"go.opentelemetry.io/ebpf-profiler/asm/amd"
 	"go.opentelemetry.io/ebpf-profiler/libpf"
+	aa "golang.org/x/arch/arm64/arm64asm"
+	"golang.org/x/arch/x86/x86asm"
 )
 
-// decodeStubArgumentWrapperARM64 disassembles arm64 code and decodes the assumed value
+// decodeStubArgumentARM64 disassembles arm64 code and decodes the assumed value
 // of requested argument.
-func decodeStubArgumentWrapperARM64(code []byte, argNumber uint8, _,
+func decodeStubArgumentARM64(code []byte,
 	addrBase libpf.SymbolValue) libpf.SymbolValue {
+	const argNumber uint8 = 0
 	// The concept is to track the latest load offset for all X0..X30 registers.
 	// These registers are used as the function arguments. Once the first branch
 	// instruction (function call/tail jump) is found, the state of the requested
@@ -100,3 +106,96 @@ func decodeStubArgumentWrapperARM64(code []byte, argNumber uint8, _,
 
 	return libpf.SymbolValueInvalid
 }
+
+func decodeStubArgumentAMD64(code []byte, codeAddress, memoryBase uint64) (
+	libpf.SymbolValue, error) {
+	targetRegister := x86asm.RDI
+
+	instructionOffset := 0
+	regs := amd.RegsState{}
+
+	for instructionOffset < len(code) {
+		rem := code[instructionOffset:]
+		if ok, insnLen := amd.DecodeSkippable(rem); ok {
+			instructionOffset += insnLen
+			continue
+		}
+
+		inst, err := x86asm.Decode(rem, 64)
+		if err != nil {
+			return 0, fmt.Errorf("failed to decode instruction at 0x%x : %w",
+				instructionOffset, err)
+		}
+		instructionOffset += inst.Len
+		regs.Set(x86asm.RIP, codeAddress+uint64(instructionOffset), 0)
+
+		if inst.Op == x86asm.CALL || inst.Op == x86asm.JMP {
+			value, loadedFrom := regs.Get(targetRegister)
+			if loadedFrom != 0 {
+				return libpf.SymbolValue(loadedFrom), nil
+			}
+			return libpf.SymbolValue(value), nil
+		}
+
+		if (inst.Op == x86asm.LEA || inst.Op == x86asm.MOV) && inst.Args[0] != nil {
+			if reg, ok := inst.Args[0].(x86asm.Reg); ok {
+				var value uint64
+				var loadedFrom uint64
+
+				switch src := inst.Args[1].(type) {
+				case x86asm.Imm:
+					value = uint64(src)
+				case x86asm.Mem:
+					baseAddr, _ := regs.Get(src.Base)
+					displacement := uint64(src.Disp)
+
+					if inst.Op == x86asm.MOV {
+						value = memoryBase
+						loadedFrom = baseAddr + displacement
+						if src.Index != 0 {
+							indexValue, _ := regs.Get(src.Index)
+							loadedFrom += indexValue * uint64(src.Scale)
+						}
+					} else if inst.Op == x86asm.LEA {
+						value = baseAddr + displacement
+						if src.Index != 0 {
+							indexValue, _ := regs.Get(src.Index)
+							value += indexValue * uint64(src.Scale)
+						}
+					}
+
+				case x86asm.Reg:
+					value, _ = regs.Get(src)
+				}
+
+				regs.Set(reg, value, loadedFrom)
+			}
+		}
+
+		if inst.Op == x86asm.ADD && inst.Args[0] != nil && inst.Args[1] != nil {
+			if reg, ok0 := inst.Args[0].(x86asm.Reg); ok0 {
+				if _, ok1 := inst.Args[1].(x86asm.Mem); ok1 {
+					oldValue, _ := regs.Get(reg)
+					value := oldValue + memoryBase
+					regs.Set(reg, value, 0)
+				}
+			}
+		}
+	}
+	return 0, errors.New("no call/jump instructions found")
+}
+
+func decodeStubArgumentWrapper(
+	code []byte,
+	codeAddress libpf.SymbolValue,
+	memoryBase libpf.SymbolValue,
+) (libpf.SymbolValue, error) {
+	switch runtime.GOARCH {
+	case "arm64":
+		return decodeStubArgumentARM64(code, memoryBase), nil
+	case "amd64":
+		return decodeStubArgumentAMD64(code, uint64(codeAddress), uint64(memoryBase))
+	default:
+		return libpf.SymbolValueInvalid, fmt.Errorf("unsupported arch %s", runtime.GOARCH)
+	}
+}