capture core dumps locally before systemd eats them

coreshim

#!/usr/bin/env python3
"""Capture core dumps via ptrace, bypassing the kernels core pattern, including systemd-coredump.

# TODO might want shibboleth instead of polyphemos
Usage: coreshim [--track-refault-sigreturn] [--track-refault-polyphemos] <command> [args...]

Writes <name>-<timestamp>-<pid>.core to cwd on crash.
If the shim dies, the child is SIGKILLed (to prevent stale core files being submitted to systemd).
x86_64 and aarch64.

--track-refault-sigreturn enables PTRACE_SYSCALL so we observe rt_sigreturn and reset
the "already delivered" flag on handler return (~2x ptrace stops per syscall).
--track-refault-polyphemos compares siginfo si_addr across SIGSEGV/SIGBUS/SIGILL/SIGFPE
faults (skipping si_code == SI_USER); same address is treated as a re-fault.
The flags may be combined.
"""

import ctypes
import ctypes.util
import errno
import os
import platform
import signal
import struct
import sys
import time

_libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
_libc.ptrace.restype = ctypes.c_long
_libc.ptrace.argtypes = [ctypes.c_ulong] * 4
_libc.tgkill.restype = ctypes.c_int
_libc.tgkill.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_int]

PTRACE_TRACEME         = 0
PTRACE_CONT            = 7
PTRACE_SYSCALL         = 24
PTRACE_GETEVENTMSG     = 0x4201
PTRACE_GETSIGINFO      = 0x4202
PTRACE_GETREGSET       = 0x4204
PTRACE_SETOPTIONS      = 0x4200

SI_USER = 0
PTRACE_O_TRACESYSGOOD  = 0x1
PTRACE_O_TRACEFORK     = 0x2
PTRACE_O_TRACEVFORK    = 0x4
PTRACE_O_TRACECLONE    = 0x8
PTRACE_O_TRACEEXEC     = 0x10
PTRACE_O_EXITKILL      = 0x100000

PTRACE_EVENT_FORK      = 1
PTRACE_EVENT_VFORK     = 2
PTRACE_EVENT_CLONE     = 3
PTRACE_EVENT_EXEC      = 4

# wait for any descendant including thread-group siblings (CLONE_THREAD)
__WALL = 0x40000000

NT_ARM_SYSTEM_CALL = 0x404

PAGE_SIZE = os.sysconf("SC_PAGE_SIZE")

NT_PRSTATUS = 1;  NT_FPREGSET = 2;  NT_PRPSINFO = 3

_ARCH = platform.machine()
if _ARCH == "x86_64":
    PTRACE_GETREGS   = 12
    PTRACE_GETFPREGS = 14
    _EM = 62   # EM_X86_64
    _GP_REG_SIZE = 216
    _FP_REG_SIZE = 512
    _NR_RT_SIGRETURN = 15
elif _ARCH == "aarch64":
    _EM = 183  # EM_AARCH64
    _GP_REG_SIZE = 272  # 34 * 8  (x0-x30, sp, pc, pstate)
    _FP_REG_SIZE = 528  # 32 * 16 + 4 + 4  (v0-v31, fpsr, fpcr)
    _NR_RT_SIGRETURN = 139
else:
    sys.exit(f"coreshim: unsupported arch {_ARCH}")


def ptrace(req, pid, addr=0, data=0):
    ctypes.set_errno(0)
    r = _libc.ptrace(req, pid, addr, data)
    if r == -1 and ctypes.get_errno():
        raise OSError(ctypes.get_errno(), os.strerror(ctypes.get_errno()))
    return r


def ptrace_getbuf(req, pid, addr, size):
    buf = ctypes.create_string_buffer(size)
    ctypes.set_errno(0)
    _libc.ptrace(req, pid, addr, ctypes.addressof(buf))
    e = ctypes.get_errno()
    if e:
        raise OSError(e, os.strerror(e))
    return buf.raw


def ptrace_getregset(pid, nt_type, size):
    """Fetch registers via PTRACE_GETREGSET (iovec-based, works on all arches)."""
    buf = ctypes.create_string_buffer(size)
    # struct iovec { void *iov_base; size_t iov_len; }
    iov_buf = ctypes.create_string_buffer(
        struct.pack("@PN", ctypes.addressof(buf), size))
    ctypes.set_errno(0)
    _libc.ptrace(PTRACE_GETREGSET, pid, nt_type, ctypes.addressof(iov_buf))
    e = ctypes.get_errno()
    if e:
        raise OSError(e, os.strerror(e))
    return buf.raw


def get_regs(pid):
    """Return (gp_regs, fp_regs) bytes for the current architecture."""
    if _ARCH == "x86_64":
        gp = ptrace_getbuf(PTRACE_GETREGS, pid, 0, _GP_REG_SIZE)
        fp = ptrace_getbuf(PTRACE_GETFPREGS, pid, 0, _FP_REG_SIZE)
    else:  # aarch64
        gp = ptrace_getregset(pid, NT_PRSTATUS, _GP_REG_SIZE)
        fp = ptrace_getregset(pid, NT_FPREGSET, _FP_REG_SIZE)
    return gp, fp


def syscall_no(pid):
    """Current syscall number at a syscall-stop."""
    if _ARCH == "x86_64":
        # orig_rax is the 16th u64 in user_regs_struct
        regs = ptrace_getbuf(PTRACE_GETREGS, pid, 0, _GP_REG_SIZE)
        return struct.unpack_from("<q", regs, 15 * 8)[0]
    return struct.unpack("<i",
        ptrace_getregset(pid, NT_ARM_SYSTEM_CALL, 4))[0]


def ptrace_getsiginfo(pid):
    """(si_code, si_addr) from the signal-stop's siginfo (LP64 layout)."""
    buf = ctypes.create_string_buffer(128)
    ctypes.set_errno(0)
    _libc.ptrace(PTRACE_GETSIGINFO, pid, 0, ctypes.addressof(buf))
    e = ctypes.get_errno()
    if e:
        raise OSError(e, os.strerror(e))
    si_code = struct.unpack_from("<i", buf.raw, 8)[0]
    si_addr = struct.unpack_from("<P", buf.raw, 16)[0]
    return si_code, si_addr


ELFCLASS64 = 2;  ELFDATA2LSB = 1;  EV_CURRENT = 1
ET_CORE = 4
PT_NOTE = 4;  PT_LOAD   = 1
PF_R = 4;     PF_W = 2;     PF_X = 1
NT_AUXV = 6;  NT_FILE = 0x46494c45

CORE_SIGS = frozenset({
    signal.SIGQUIT, signal.SIGILL, signal.SIGTRAP, signal.SIGABRT,
    signal.SIGBUS,  signal.SIGFPE, signal.SIGSEGV, signal.SIGXCPU,
    signal.SIGXFSZ, signal.SIGSYS,
})

# fault-bearing signals: siginfo.si_addr identifies the offending insn/address
POLYPHEMOS_SIGS = frozenset({
    signal.SIGSEGV, signal.SIGBUS, signal.SIGILL, signal.SIGFPE,
})


def elf_note(name, ntype, data):
    nb = name.encode() + b'\x00'
    hdr = struct.pack("<III", len(nb), len(data), ntype)
    return hdr + nb.ljust(-(-len(nb) & ~3), b'\x00') \
               + data.ljust(-(-len(data) & ~3), b'\x00')


def prstatus(sig_num, pid, regs):
    """elf_prstatus — layout is arch-independent, only pr_reg size varies."""
    try:
        pgrp, sid = os.getpgid(pid), os.getsid(pid)
    except OSError:
        pgrp = sid = pid
    d  = struct.pack("<iii", sig_num, 0, 0)           # pr_info
    d += struct.pack("<hxx", sig_num)                 # pr_cursig + pad
    d += struct.pack("<QQ", 0, 0)                     # sigpend, sighold
    d += struct.pack("<iiii", pid, os.getpid(), pgrp, sid)
    d += b'\x00' * 64                                 # 4 timevals
    d += regs                                         # pr_reg
    d += struct.pack("<i4x", 1)                       # fpvalid + pad
    return elf_note("CORE", NT_PRSTATUS, d)


def prpsinfo(pid, name, args):
    """elf_prpsinfo (136 bytes, same for all 64-bit arches)."""
    try:
        pgrp, sid = os.getpgid(pid), os.getsid(pid)
    except OSError:
        pgrp = sid = pid
    # pr_state, pr_sname, pr_zomb, pr_nice + pad + pr_flag
    d  = struct.pack("<bbbb4xQ", 0, ord('R'), 0, 0, 0)
    d += struct.pack("<II", os.getuid(), os.getgid())
    d += struct.pack("<iiii", pid, os.getpid(), pgrp, sid)
    d += name[:15].encode().ljust(16, b'\x00')
    d += args[:79].encode().ljust(80, b'\x00')
    return elf_note("CORE", NT_PRPSINFO, d)


def auxv(pid):
    try:
        with open(f"/proc/{pid}/auxv", "rb") as f:
            return elf_note("CORE", NT_AUXV, f.read())
    except OSError:
        return b''


def file_note(mappings):
    entries = [(s, e, off, p) for s, e, _, off, p in mappings if p]
    if not entries:
        return b''
    d = struct.pack("<QQ", len(entries), PAGE_SIZE)
    d += b''.join(struct.pack("<QQQ", s, e, off) for s, e, off, _ in entries)
    d += b''.join(p.encode() + b'\x00' for _, _, _, p in entries)
    return elf_note("CORE", NT_FILE, d)


def parse_maps(pid):
    """→ [(start, end, elf_flags, pgoff_pages, path|'')]"""
    result = []
    with open(f"/proc/{pid}/maps") as f:
        for line in f:
            parts = line.split(None, 5)
            lo, hi = (int(x, 16) for x in parts[0].split('-'))
            p = parts[1]
            flags = (PF_R if 'r' in p else 0) | \
                    (PF_W if 'w' in p else 0) | \
                    (PF_X if 'x' in p else 0)
            pgoff = int(parts[2], 16) // PAGE_SIZE
            path = parts[5].strip() if len(parts) > 5 else ''
            if not path or path.startswith('['):
                path = ''
            result.append((lo, hi, flags, pgoff, path))
    return result


def read_mapping(mem_fd, start, size):
    """Read [start, start+size) from /proc/<pid>/mem, zero-filling holes.

    Short reads in the middle of an otherwise readable mapping are retried
    instead of being silently zero-padded. On a hard error or zero-byte
    return we skip to the next page boundary and continue.
    """
    buf = bytearray(size)
    off = 0
    while off < size:
        try:
            os.lseek(mem_fd, start + off, os.SEEK_SET)
            chunk = os.read(mem_fd, size - off)
        except OSError:
            chunk = b''
        if chunk:
            buf[off:off + len(chunk)] = chunk
            off += len(chunk)
            continue
        # unreadable byte: skip ahead to the next page boundary
        next_page = ((start + off) // PAGE_SIZE + 1) * PAGE_SIZE
        off = min(next_page - start, size)
    return bytes(buf)


def caught_signals(pid):
    """SigCgt bitmask from /proc/<pid>/status; 0 if unreadable."""
    try:
        with open(f"/proc/{pid}/status") as f:
            for line in f:
                if line.startswith("SigCgt:"):
                    return int(line.split()[1], 16)
    except OSError:
        pass
    return 0


def thread_group_id(tid):
    """Tgid (process pid) for a tid; falls back to tid."""
    try:
        with open(f"/proc/{tid}/status") as f:
            for line in f:
                if line.startswith("Tgid:"):
                    return int(line.split()[1])
    except OSError:
        pass
    return tid


def list_thread_tids(tgid):
    try:
        return sorted(int(x) for x in os.listdir(f"/proc/{tgid}/task"))
    except OSError:
        return []


def ptrace_geteventmsg(pid):
    msg = ctypes.c_ulong(0)
    ctypes.set_errno(0)
    _libc.ptrace(PTRACE_GETEVENTMSG, pid, 0, ctypes.addressof(msg))
    e = ctypes.get_errno()
    if e:
        raise OSError(e, os.strerror(e))
    return msg.value


def stop_thread(tgid, tid):
    """Ensure tid is in ptrace-stop so its registers can be read.

    Returns True if the thread is stopped (or was already), False otherwise.
    """
    try:
        get_regs(tid)
        return True
    except OSError as e:
        if e.errno != errno.ESRCH:
            return False
    if _libc.tgkill(tgid, tid, signal.SIGSTOP) != 0:
        return False
    deadline = time.monotonic() + 1.0
    while time.monotonic() < deadline:
        try:
            wpid, st = os.waitpid(tid, __WALL | os.WNOHANG)
        except OSError:
            return False
        if wpid == tid:
            if os.WIFSTOPPED(st):
                return True
            if os.WIFEXITED(st) or os.WIFSIGNALED(st):
                return False
        time.sleep(0.005)
    return False


def write_core(path, tid, sig_num, name, args):
    tgid = thread_group_id(tid)
    siblings = [t for t in list_thread_tids(tgid) if t != tid]
    threads = []
    try:
        regs, fpregs = get_regs(tid)
        threads.append((tid, sig_num, regs, fpregs))
    except OSError:
        threads.append((tid, sig_num,
                        b'\x00' * _GP_REG_SIZE, b'\x00' * _FP_REG_SIZE))
    for sib in siblings:
        if not stop_thread(tgid, sib):
            continue
        try:
            regs, fpregs = get_regs(sib)
        except OSError:
            continue
        threads.append((sib, 0, regs, fpregs))

    mappings = parse_maps(tgid)

    notes = b''
    for t_tid, t_sig, t_regs, t_fpregs in threads:
        notes += prstatus(t_sig, t_tid, t_regs)
        notes += elf_note("CORE", NT_FPREGSET, t_fpregs)
    notes += prpsinfo(tgid, name, args)
    notes += auxv(tgid)
    notes += file_note(mappings)

    n_phdr = 1 + len(mappings)
    ehdr_sz = 64
    phdr_sz = 56
    notes_off = ehdr_sz + n_phdr * phdr_sz
    data_off = notes_off + len(notes)

    # compute file offset per segment
    seg_off = data_off
    offsets = []
    for start, end, flags, _, _ in mappings:
        offsets.append(seg_off if (flags & PF_R) else 0)
        if flags & PF_R:
            seg_off += end - start

    ident = b'\x7fELF' + bytes([ELFCLASS64, ELFDATA2LSB, EV_CURRENT, 0]) + b'\x00' * 8

    with open(path, "wb") as f:
        f.write(struct.pack("<16sHHIQQQIHHHHHH",
            ident, ET_CORE, _EM, EV_CURRENT,
            0, ehdr_sz, 0, 0,
            ehdr_sz, phdr_sz, n_phdr, 0, 0, 0))

        f.write(struct.pack("<IIQQQQQQ",
            PT_NOTE, 0, notes_off, 0, 0, len(notes), 0, 4))

        for i, (start, end, flags, _, _) in enumerate(mappings):
            size = end - start
            filesz = size if (flags & PF_R) else 0
            f.write(struct.pack("<IIQQQQQQ",
                PT_LOAD, flags, offsets[i], start, 0, filesz, size, PAGE_SIZE))

        f.write(notes)

        mem_fd = os.open(f"/proc/{tgid}/mem", os.O_RDONLY)
        try:
            for start, end, flags, _, _ in mappings:
                if flags & PF_R:
                    f.write(read_mapping(mem_fd, start, end - start))
        finally:
            os.close(mem_fd)


def main():
    args = sys.argv[1:]
    track_sigreturn = False
    track_polyphemos = False
    while args and args[0].startswith("--"):
        if args[0] == "--track-refault-sigreturn":
            track_sigreturn = True
        elif args[0] == "--track-refault-polyphemos":
            track_polyphemos = True
        else:
            break
        args = args[1:]

    if not args:
        sys.exit(f"Usage: {os.path.basename(sys.argv[0])} "
                 "[--track-refault-sigreturn] [--track-refault-polyphemos] "
                 "<command> [args...]")

    cmd = args
    bin_name = os.path.basename(cmd[0])

    pid = os.fork()
    if pid == 0:
        try:
            ptrace(PTRACE_TRACEME, 0)
            os.execvp(cmd[0], cmd)
        except Exception as e:
            print(f"coreshim: {e}", file=sys.stderr)
        os._exit(127)

    # wait for exec-stop
    _, status = os.waitpid(pid, 0)
    if not os.WIFSTOPPED(status):
        sys.exit(os.WEXITSTATUS(status) if os.WIFEXITED(status)
                 else 128 + os.WTERMSIG(status))

    options = (PTRACE_O_TRACEEXEC | PTRACE_O_EXITKILL |
               PTRACE_O_TRACEFORK | PTRACE_O_TRACEVFORK | PTRACE_O_TRACECLONE)
    if track_sigreturn:
        options |= PTRACE_O_TRACESYSGOOD
    ptrace(PTRACE_SETOPTIONS, pid, 0, options)

    cont_req = PTRACE_SYSCALL if track_sigreturn else PTRACE_CONT
    syscall_stop_sig = signal.SIGTRAP | 0x80
    ptrace(cont_req, pid, 0, 0)

    # child gets SIGINT via process group; ignore in parent so we don't die
    signal.signal(signal.SIGINT, signal.SIG_IGN)

    tracees = {pid}        # tids whose initial stop has already been consumed
    pending_init = set()   # tids known via FORK/CLONE event, awaiting initial SIGSTOP
    delivered = {}         # tgid -> {sig: si_addr | None}; sigreturn checks
                           # membership, polyphemos compares the stored addr

    while True:
        try:
            wpid, status = os.waitpid(-1, __WALL)
        except ChildProcessError:
            sys.exit(0)

        if os.WIFEXITED(status) or os.WIFSIGNALED(status):
            tracees.discard(wpid)
            pending_init.discard(wpid)
            delivered.pop(wpid, None)
            if wpid == pid:
                if os.WIFEXITED(status):
                    sys.exit(os.WEXITSTATUS(status))
                sys.exit(128 + os.WTERMSIG(status))
            continue

        if not os.WIFSTOPPED(status):
            continue

        sig = os.WSTOPSIG(status)
        event = (status >> 16) & 0xff

        # New tracee from fork/vfork/clone: suppress its initial SIGSTOP.
        # Two arrival orders are possible (parent's event first, or child's
        # SIGSTOP first); both end up here.
        if wpid in pending_init or wpid not in tracees:
            pending_init.discard(wpid)
            tracees.add(wpid)
            ptrace(cont_req, wpid, 0, 0)
            continue

        if track_sigreturn and sig == syscall_stop_sig:
            try:
                if syscall_no(wpid) == _NR_RT_SIGRETURN:
                    delivered.pop(thread_group_id(wpid), None)
            except OSError:
                pass
            ptrace(cont_req, wpid, 0, 0)
            continue

        if event in (PTRACE_EVENT_FORK, PTRACE_EVENT_VFORK, PTRACE_EVENT_CLONE):
            try:
                new_tid = ptrace_geteventmsg(wpid)
                if new_tid not in tracees:
                    pending_init.add(new_tid)
            except OSError:
                pass
            ptrace(cont_req, wpid, 0, 0)
            continue

        if event == PTRACE_EVENT_EXEC:
            ptrace(cont_req, wpid, 0, 0)
            continue
        if event:
            sys.exit(f"coreshim: unhandled ptrace event {event}")

        if sig in CORE_SIGS:
            tgid = thread_group_id(wpid)
            if caught_signals(wpid) & (1 << (sig - 1)):
                if not (track_sigreturn or track_polyphemos):
                    # trust the application's handler; stay out of the way
                    ptrace(cont_req, wpid, 0, sig)
                    continue
                si_addr = None
                if track_polyphemos and sig in POLYPHEMOS_SIGS:
                    try:
                        si_code, addr = ptrace_getsiginfo(wpid)
                        if si_code != SI_USER:
                            si_addr = addr
                    except OSError:
                        pass
                seen = delivered.get(tgid, {})
                if si_addr is not None:
                    # polyphemos has a verdict; different addr means fresh fault
                    refault = sig in seen and seen[sig] == si_addr
                else:
                    # fall back to sigreturn's any-second-occurrence check
                    refault = track_sigreturn and sig in seen
                if not refault:
                    if track_sigreturn or si_addr is not None:
                        delivered.setdefault(tgid, {})[sig] = si_addr
                    ptrace(cont_req, wpid, 0, sig)
                    continue

            args_str = ' '.join(cmd)
            core_path = f"{bin_name}-{int(time.time())}-{tgid}.core"
            try:
                write_core(core_path, wpid, sig, bin_name, args_str)
                sig_name = signal.Signals(sig).name
                print(f"Killed by {sig_name}, core dumped to {core_path}",
                      file=sys.stderr)
            except Exception as e:
                print(f"coreshim: core dump failed: {e}", file=sys.stderr)
            os.kill(tgid, signal.SIGKILL)
            if tgid == pid:
                # original child died; reap any remaining descendants and exit
                while True:
                    try:
                        os.waitpid(-1, __WALL)
                    except ChildProcessError:
                        break
                sys.exit(128 + sig)
            continue

        # non-fatal signal: deliver it
        ptrace(cont_req, wpid, 0, sig)


if __name__ == "__main__":
    main()