gasm-sdk/debug/ptrace_linux_amd64.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause

//go:build linux && amd64

// Package debug implements the interactive debugger for gasm (Phase 4):
// single-stepping, breakpoints, register and memory inspection for
// JIT-assembled Plan 9 amd64 functions, controlled via ptrace.
package debug

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"
	"unsafe"
)

// Session is a ptrace debugging session controlling one debuggee process.
type Session struct {
	pid      int
	cmd      *exec.Cmd
	stopped  bool
	exited   bool
	codeBase uint64 // base address of the JIT code in the debuggee
}

// Launch starts the debuggee subprocess (gasm debug --target ...) and
// attaches to it via ptrace.  The debuggee assembles the file, maps the
// JIT code, calls PTRACE_TRACEME and raises SIGSTOP; Launch waits for
// that initial stop and returns a ready Session.
func Launch(gasmBin, asmPath, funcName string, args []byte) (*Session, error) {
	sess, _, err := LaunchWithBuffers(gasmBin, asmPath, funcName, args, "")
	return sess, err
}

// LaunchWithBuffers is like Launch but also allocates buffers in the debuggee
// based on the buffer specification.  Returns the Session and the buffer
// addresses (in the order they appear in the spec).
func LaunchWithBuffers(gasmBin, asmPath, funcName string, args []byte, bufSpec string) (*Session, []uint64, error) {
	self, err := os.Executable()
	if err != nil {
		return nil, nil, fmt.Errorf("debug: cannot find gasm binary: %w", err)
	}
	if gasmBin != "" {
		self = gasmBin
	}

	// Write the arg block to a temp file (the child reads it).
	tmpDir, err := os.MkdirTemp("", "gasm-debug-*")
	if err != nil {
		return nil, nil, fmt.Errorf("debug: tempdir: %w", err)
	}
	argsFile := filepath.Join(tmpDir, "args.bin")
	if err := os.WriteFile(argsFile, args, 0o644); err != nil {
		os.RemoveAll(tmpDir)
		return nil, nil, fmt.Errorf("debug: write args: %w", err)
	}

	// Write the buffer spec if present.
	if bufSpec != "" {
		if err := os.WriteFile(filepath.Join(tmpDir, "bufspec"), []byte(bufSpec), 0o644); err != nil {
			os.RemoveAll(tmpDir)
			return nil, nil, fmt.Errorf("debug: write bufspec: %w", err)
		}
	}

	cmd := exec.Command(self, "debug", "--target", "--func", funcName, "--args", argsFile, asmPath)
	cmd.Env = append(os.Environ(), "GASM_DEBUG_TMP="+tmpDir)
	cmd.Stdout = nil // output goes to the debugger, not the terminal
	cmd.Stderr = os.Stderr
	cmd.SysProcAttr = &syscall.SysProcAttr{}

	if err := cmd.Start(); err != nil {
		os.RemoveAll(tmpDir)
		return nil, nil, fmt.Errorf("debug: start debuggee: %w", err)
	}

	s := &Session{pid: cmd.Process.Pid, cmd: cmd}

	// Wait for the child to signal readiness and stop.  The child calls
	// PTRACE_TRACEME then SIGSTOP, so Wait4 with WUNTRACED observes the
	// ptrace-stop directly (no PTRACE_ATTACH needed).
	readyFile := filepath.Join(tmpDir, "ready")
	for i := 0; i < 500; i++ {
		if _, err := os.Stat(readyFile); err == nil {
			break
		}
		time.Sleep(5 * time.Millisecond)
	}
	var ws syscall.WaitStatus
	if _, err := syscall.Wait4(s.pid, &ws, syscall.WUNTRACED, nil); err != nil {
		cmd.Process.Kill()
		os.RemoveAll(tmpDir)
		return nil, nil, fmt.Errorf("debug: wait for stop: %w", err)
	}

	// Wait for the debuggee to reach the function entry point.
	entryFile := filepath.Join(tmpDir, "entry")
	for i := 0; i < 500; i++ {
		if _, err := os.Stat(entryFile); err == nil {
			break
		}
		time.Sleep(5 * time.Millisecond)
	}

	// Continue the debuggee to the entry point.
	if err := s.Continue(); err != nil {
		return nil, nil, fmt.Errorf("debug: continue to entry: %w", err)
	}

	// Wait for the entry stop.
	if _, err := syscall.Wait4(s.pid, &ws, syscall.WUNTRACED, nil); err != nil {
		return nil, nil, fmt.Errorf("debug: wait for entry: %w", err)
	}

	s.stopped = true

	// Read the code base from /proc/pid/maps (find the RWX mapping).
	s.codeBase = findRWXMapping(s.pid)
	if s.codeBase == 0 {
		// Fallback: try the file the child wrote.
		baseFile := filepath.Join(tmpDir, "codebase")
		if data, err := os.ReadFile(baseFile); err == nil {
			fmt.Sscanf(string(data), "%d", &s.codeBase)
		}
	}

	// Read buffer addresses if buffers were allocated.
	var bufAddrs []uint64
	if bufSpec != "" {
		addrFile := filepath.Join(tmpDir, "bufaddrs")
		if data, err := os.ReadFile(addrFile); err == nil {
			for _, line := range strings.Split(strings.TrimSpace(string(data)), "\n") {
				var addr uint64
				if _, err := fmt.Sscanf(line, "%d", &addr); err == nil {
					bufAddrs = append(bufAddrs, addr)
				}
			}
		}
	}

	return s, bufAddrs, nil
}

// wait waits for the debuggee to stop and returns the wait status.
func (s *Session) wait() error {
	var ws syscall.WaitStatus
	_, err := syscall.Wait4(s.pid, &ws, 0, nil)
	if err != nil {
		return err
	}
	if ws.Exited() {
		s.exited = true
		return fmt.Errorf("debuggee exited with status %d", ws.ExitStatus())
	}
	s.stopped = true
	return nil
}

// GetRegs reads the general-purpose registers of the stopped debuggee.
func (s *Session) GetRegs() (Regs, error) {
	var regs Regs
	_, _, errno := syscall.Syscall6(
		syscall.SYS_PTRACE,
		uintptr(syscall.PTRACE_GETREGS),
		uintptr(s.pid),
		0,
		uintptr(unsafe.Pointer(&regs)),
		0, 0,
	)
	if errno != 0 {
		return regs, fmt.Errorf("debug: PTRACE_GETREGS: %w", errno)
	}
	return regs, nil
}

// SetRegs writes the general-purpose registers of the stopped debuggee.
func (s *Session) SetRegs(regs *Regs) error {
	_, _, errno := syscall.Syscall6(
		syscall.SYS_PTRACE,
		uintptr(syscall.PTRACE_SETREGS),
		uintptr(s.pid),
		0,
		uintptr(unsafe.Pointer(regs)),
		0, 0,
	)
	if errno != 0 {
		return fmt.Errorf("debug: PTRACE_SETREGS: %w", errno)
	}
	return nil
}

// FPRegs holds the x87 FPU and SSE (XMM) register state from PTRACE_GETFPREGS.
type FPRegs struct {
	FCW       uint16
	FSW       uint16
	FTW       byte
	FOP       uint16
	FIP       uint64
	FCS       uint16
	FDP       uint64
	FDS       uint16
	MXCSR     uint32
	MXCSRMask uint32
	ST        [8][16]byte  // x87 stack (10 bytes per reg, padded to 16)
	XMM       [16][16]byte // XMM0-15
}

// GetFPRegs retrieves the FPU/SSE register state of the stopped debuggee.
func (s *Session) GetFPRegs() (FPRegs, error) {
	var fp FPRegs
	_, _, errno := syscall.Syscall6(
		syscall.SYS_PTRACE,
		uintptr(syscall.PTRACE_GETFPREGS),
		uintptr(s.pid),
		0,
		uintptr(unsafe.Pointer(&fp)),
		0, 0,
	)
	if errno != 0 {
		return fp, fmt.Errorf("debug: PTRACE_GETFPREGS: %w", errno)
	}
	return fp, nil
}

// VectorRegs holds the YMM register state extracted from XSAVE.
type VectorRegs struct {
	YMM [16][32]byte // YMM0-15 (full 256-bit values)
}

// GetVectorRegs retrieves the YMM registers via PTRACE_GETREGSET + XSAVE.
// Falls back to XMM if XSAVE is unavailable.
func (s *Session) GetVectorRegs() (VectorRegs, error) {
	var v VectorRegs
	fp, err := s.GetFPRegs()
	if err != nil {
		return v, err
	}
	// PTRACE_GETFPREGS gives XMM registers (lower 128 bits).
	// For YMM we'd need XSAVE; for now, copy XMM and zero the upper half.
	for i := 0; i < 16; i++ {
		for j := 0; j < 16; j++ {
			v.YMM[i][j] = fp.XMM[i][j]
		}
		// Upper 128 bits would come from XSAVE, not available via GETFPREGS.
	}
	return v, nil
}

// Peek reads a word (8 bytes) from the debuggee's memory at addr.
// Uses /proc/pid/mem which works reliably with Go's multi-threaded runtime.
func (s *Session) Peek(addr uint64) (uint64, error) {
	mem, err := os.OpenFile(fmt.Sprintf("/proc/%d/mem", s.pid), os.O_RDONLY, 0)
	if err != nil {
		return 0, fmt.Errorf("debug: open /proc/%d/mem: %w", s.pid, err)
	}
	defer mem.Close()
	buf := make([]byte, 8)
	if _, err := mem.ReadAt(buf, int64(addr)); err != nil {
		return 0, fmt.Errorf("debug: read mem %#x: %w", addr, err)
	}
	return uint64(buf[0]) | uint64(buf[1])<<8 | uint64(buf[2])<<16 | uint64(buf[3])<<24 |
		uint64(buf[4])<<32 | uint64(buf[5])<<40 | uint64(buf[6])<<48 | uint64(buf[7])<<56, nil
}

// Poke writes a word (8 bytes) to the debuggee's memory at addr.
func (s *Session) Poke(addr, val uint64) error {
	mem, err := os.OpenFile(fmt.Sprintf("/proc/%d/mem", s.pid), os.O_WRONLY, 0)
	if err != nil {
		return fmt.Errorf("debug: open /proc/%d/mem: %w", s.pid, err)
	}
	defer mem.Close()
	buf := []byte{byte(val), byte(val >> 8), byte(val >> 16), byte(val >> 24),
		byte(val >> 32), byte(val >> 40), byte(val >> 48), byte(val >> 56)}
	if _, err := mem.WriteAt(buf, int64(addr)); err != nil {
		return fmt.Errorf("debug: write mem %#x: %w", addr, err)
	}
	return nil
}

// ReadMemory reads len bytes from the debuggee's memory at addr.
func (s *Session) ReadMemory(addr uint64, length int) ([]byte, error) {
	out := make([]byte, length)
	for i := 0; i < length; i += 8 {
		word, err := s.Peek(addr + uint64(i))
		if err != nil {
			return out[:i], err
		}
		for j := 0; j < 8 && i+j < length; j++ {
			out[i+j] = byte(word >> (8 * j))
		}
	}
	return out, nil
}

// WriteMemory writes bytes to the debuggee's memory at addr.
func (s *Session) WriteMemory(addr uint64, data []byte) error {
	for i := 0; i < len(data); i += 8 {
		end := i + 8
		if end > len(data) {
			end = len(data)
		}
		var word uint64
		for j := 0; j < end-i; j++ {
			word |= uint64(data[i+j]) << (8 * j)
		}
		// For partial writes, read-modify-write the existing word.
		if end-i < 8 {
			existing, err := s.Peek(addr + uint64(i))
			if err != nil {
				return err
			}
			// Clear the bytes we're overwriting and merge.
			mask := ^((uint64(1) << (8 * (end - i))) - 1)
			word = (existing & mask) | word
		}
		if err := s.Poke(addr+uint64(i), word); err != nil {
			return err
		}
	}
	return nil
}

// Step executes a single instruction in the debuggee.
func (s *Session) Step() error {
	if s.exited {
		return fmt.Errorf("debug: debuggee has exited")
	}
	_, _, errno := syscall.Syscall6(
		syscall.SYS_PTRACE,
		uintptr(syscall.PTRACE_SINGLESTEP),
		uintptr(s.pid),
		0, 0, 0, 0,
	)
	if errno != 0 {
		return fmt.Errorf("debug: PTRACE_SINGLESTEP: %w", errno)
	}
	return s.wait()
}

// Continue resumes execution until the next breakpoint or exit.
func (s *Session) Continue() error {
	if s.exited {
		return fmt.Errorf("debug: debuggee has exited")
	}
	_, _, errno := syscall.Syscall6(
		syscall.SYS_PTRACE,
		uintptr(syscall.PTRACE_CONT),
		uintptr(s.pid),
		0, 0, 0, 0,
	)
	if errno != 0 {
		return fmt.Errorf("debug: PTRACE_CONT: %w", errno)
	}
	return s.wait()
}

// Exited returns true if the debuggee has terminated.
func (s *Session) Exited() bool {
	return s.exited
}

// Pid returns the debuggee's process ID.
func (s *Session) Pid() int {
	return s.pid
}

// CodeBase returns the base address of the JIT code in the debuggee.
func (s *Session) CodeBase() uint64 {
	return s.codeBase
}

// Kill terminates the debuggee.
func (s *Session) Kill() {
	if !s.exited {
		syscall.Kill(s.pid, syscall.SIGKILL)
		syscall.Wait4(s.pid, nil, 0, nil)
		s.exited = true
	}
	if s.cmd != nil && s.cmd.Process != nil {
		s.cmd.Wait()
	}
}

// findRWXMapping reads /proc/pid/maps and returns the base address of the
// first read-write-execute mapping (the JIT code region).
func findRWXMapping(pid int) uint64 {
	data, err := os.ReadFile(fmt.Sprintf("/proc/%d/maps", pid))
	if err != nil {
		return 0
	}
	for _, line := range strings.Split(string(data), "\n") {
		// Format: addr-addr perms offset dev inode pathname
		fields := strings.Fields(line)
		if len(fields) < 2 {
			continue
		}
		perms := fields[1]
		if len(perms) >= 3 && perms[0] == 'r' && perms[1] == 'w' && perms[2] == 'x' {
			// Parse the start address.
			var start uint64
			fmt.Sscanf(fields[0], "%x-", &start)
			return start
		}
	}
	return 0
}