/*	$OpenBSD: dt_prov_kprobe.c,v 1.10 2025/08/14 13:04:48 mpi Exp $	*/

/*
 * Copyright (c) 2024 Martin Pieuchot <mpi@openbsd.org>
 * Copyright (c) 2020 Tom Rollet <tom.rollet@epita.fr>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#if defined(DDBPROF) && (defined(__amd64__) || defined(__i386__))

#include <sys/types.h>
#include <sys/systm.h>
#include <sys/param.h>
#include <sys/malloc.h>
#include <sys/exec_elf.h>

#include <ddb/db_elf.h>
#include <machine/db_machdep.h>
#include <ddb/db_extern.h>
#include <ddb/db_access.h>
#include <ddb/db_sym.h>
#include <ddb/db_interface.h>

#include <dev/dt/dtvar.h>

#define KPROBE_ENTRY	0x1
#define KPROBE_RETURN	0x2

extern db_symtab_t	db_symtab;
extern char		__kutext_end[];
extern int		db_prof_on;

extern void	db_prof_count(struct trapframe *);
extern vaddr_t	db_get_probe_addr(struct trapframe *);

/* Lists of probes per ELF symbol. */
SLIST_HEAD(, dt_probe) *dtpf_entry;
SLIST_HEAD(, dt_probe) *dtpf_return;

int	dt_prov_kprobe_alloc(struct dt_probe *, struct dt_softc *,
	    struct dt_pcb_list *, struct dtioc_req *);
int	dt_prov_kprobe_hook(struct dt_provider *, ...);
int	dt_prov_kprobe_dealloc(struct dt_probe *, struct dt_softc *,
	    struct dtioc_req *);

#define DTEVT_PROV_KPROBE (DTEVT_COMMON|DTEVT_FUNCARGS)

struct dt_provider dt_prov_kprobe = {
	.dtpv_name    = "kprobe",
	.dtpv_alloc   = dt_prov_kprobe_alloc,
	.dtpv_enter   = dt_prov_kprobe_hook,
	.dtpv_leave   = NULL,
	.dtpv_dealloc = dt_prov_kprobe_dealloc,
};

/* Bob Jenkin's public domain 32-bit integer hashing function.
 * Original at https://burtleburtle.net/bob/hash/integer.html.
 */
uint32_t
ptr_hash(uint32_t a) {
	a = (a + 0x7ed55d16) + (a<<12);
	a = (a ^ 0xc761c23c) ^ (a>>19);
	a = (a + 0x165667b1) + (a<<5);
	a = (a + 0xd3a2646c) ^ (a<<9);
	a = (a + 0xfd7046c5) + (a<<3);
	a = (a ^ 0xb55a4f09) ^ (a>>16);
	return a;
}

#define	PPTSIZE		(PAGE_SIZE * 30) /* XXX */
#define	PPTMASK		((PPTSIZE / sizeof(struct dt_probe)) - 1)
#define	INSTTOIDX(inst)	(ptr_hash(inst) & PPTMASK)

#if defined(__amd64__)
#define	IBT_SIZE	4
#define	RTGD_MOV_SIZE	7
#define	RTGD_XOR_SIZE	4
#define	FRAME_SIZE	(IBT_SIZE + RTGD_XOR_SIZE + RTGD_XOR_SIZE + 1 + 3)

#define	RET_INST	0xc3
#define	RET_SIZE	1

/*
 * Validate that this prologue respect a well-known layout and return
 * the offset where the symbol can be patched.
 */
// XXX use db_get_value();
int
db_prologue_validate(Elf_Sym *symp)
{
	uint8_t *inst = (uint8_t *)symp->st_value;
	int off = symp->st_size - 1;

	if (off < IBT_SIZE + RTGD_MOV_SIZE + RTGD_XOR_SIZE)
		return -1;

	/* Check for IBT */
	if (inst[0] != 0xf3 || inst[1] != 0x0f ||
	    inst[2] != 0x1e || inst[3] != 0xfa)
		return -1;

	/* Check for retguard */
	off = IBT_SIZE;
	if (inst[off] != 0x4c || inst[off + 1] != 0x8b || inst[off + 2] != 0x1d)
	    	return -1;

	/* Check for `xorq off(%rsp), %reg' */
	off += RTGD_MOV_SIZE;
	if (inst[off] != 0x4c || inst[off + 1] != 0x33 || inst[off + 2] != 0x1c)
		return -1;

	/* Check for `pushq %rbp' */
	off += RTGD_XOR_SIZE;
	if (inst[off] != SSF_INST)
		return -1;

	/* Check for `movq %rsp, %rbp'  */
	if (inst[off + 1] != 0x48 || inst[off + 2] != 0x89 ||
	    inst[off + 3] != 0xe5)
	    	return -1;

	return off;
}

/*
 * Insert a breakpoint or restore `pushq %rbp'.
 */
void
db_prologue_patch(vaddr_t addr, int restore)
{
	uint8_t patch;
	size_t size;
	unsigned s;

	CTASSERT(SSF_SIZE == BKPT_SIZE);
	if (restore) {
		patch = SSF_INST;
		size = SSF_SIZE;
	} else {
		patch = BKPT_INST;
		size = BKPT_SIZE;
	}

	s = intr_disable();
	db_write_bytes(addr, size, &patch);
	intr_restore(s);
}

int
db_epilogue_validate(Elf_Sym *symp)
{
	uint8_t *inst = (uint8_t *)symp->st_value;
	int off = symp->st_size - 1;

	while (off > FRAME_SIZE + 2) {
		if (inst[off - 1] == 0xcc && inst[off] == RET_INST)
			return off;
		off--;
	}

	return -1;
}

void
db_epilogue_patch(vaddr_t addr, int restore)
{
	uint8_t patch;
	size_t size;
	unsigned s;

	CTASSERT(RET_SIZE == BKPT_SIZE);
	if (restore) {
		patch = RET_INST;
		size = RET_SIZE;
	} else {
		patch = BKPT_INST;
		size = BKPT_SIZE;
	}

	s = intr_disable();
	db_write_bytes(addr, size, &patch);
	intr_restore(s);
}

#elif defined(__i386__)
#define	POP_RBP_INST	0x5d
#define	POP_RBP_SIZE	1

int
db_prologue_validate(Elf_Sym *symp)
{
	uint8_t *inst = (uint8_t *)symp->st_value;

	/* No retguard or IBT on i386 */
	if (*inst != SSF_INST)
		return -1;

	return 0;
}

/*
 * Insert a breakpoint or restore
 */
void
db_prologue_patch(vaddr_t addr, int restore)
{
	uint8_t patch;
	size_t size;
	unsigned s;

	CTASSERT(SSF_SIZE == BKPT_SIZE);
	if (restore) {
		patch = SSF_INST;
		size = SSF_SIZE;
	} else {
		patch = BKPT_INST;
		size = BKPT_SIZE;
	}

	s = intr_disable();
	db_write_bytes(addr, size, &patch);
	intr_restore(s);
}

int
db_epilogue_validate(Elf_Sym *symp)
{
	vaddr_t limit = symp->st_value + symp->st_size;
#if 0
	/*
	 * Little temporary hack to find some return probe
	 *   => always int3 after 'pop %rpb; ret'
	 */
	while(*((uint8_t *)inst) == 0xcc)
		(*(uint8_t *)inst) -= 1;
#endif
	if (*(uint8_t *)(limit - 2) != POP_RBP)
		return -1;

	return symp->st_size - 2;
}

void
db_epilogue_patch(vaddr_t addr, int restore)
{
	uint8_t patch;
	size_t size;
	unsigned s;

	CTASSERT(SSF_SIZE == BKPT_SIZE);
	if (restore) {
		patch = POP_RBP_INST;
		size = POP_RBP_SIZE;
	} else {
		patch = BKPT_INST;
		size = BKPT_SIZE;
	}

	s = intr_disable();
	db_write_bytes(addr, size, &patch);
	intr_restore(s);
}
#endif /* defined(__i386__) */

/* Initialize all entry and return probes and store them in global arrays */
int
dt_prov_kprobe_init(void)
{
	struct dt_probe *dtp;
	Elf_Sym *symp, *symtab_start, *symtab_end;
	const char *strtab, *name;
	vaddr_t inst;
	int off, nb_sym, nb_probes = 0;

	nb_sym = (db_symtab.end - db_symtab.start) / sizeof (Elf_Sym);

	dtpf_entry = malloc(PPTSIZE, M_DT, M_NOWAIT|M_ZERO);
	if (dtpf_entry == NULL)
		return 0;

	dtpf_return = malloc(PPTSIZE, M_DT, M_NOWAIT|M_ZERO);
	if (dtpf_return == NULL) {
		free(dtpf_entry, M_DT, PPTSIZE);
		return 0;
	}

	db_symtab_t *stab = &db_symtab;

	symtab_start = STAB_TO_SYMSTART(stab);
	symtab_end = STAB_TO_SYMEND(stab);
	strtab = db_elf_find_strtab(stab);

	for (symp = symtab_start; symp < symtab_end; symp++) {
		if (ELF_ST_TYPE(symp->st_info) != STT_FUNC)
			continue;

		inst = symp->st_value;
		name = strtab + symp->st_name;

		/* Filter function that are not mapped in memory */
		if (inst < KERNBASE || inst >= (vaddr_t)&__kutext_end)
			continue;

		/* Remove some function to avoid recursive tracing */
		if (strncmp(name, "dt_", 3) == 0 ||
		    strncmp(name, "trap", 4) == 0 ||
		    strncmp(name, "db_", 3) == 0)
			continue;

		off = db_prologue_validate(symp);
		if (off < 0)
			continue;

		dtp = dt_dev_alloc_probe(name, "entry", &dt_prov_kprobe);
		if (dtp == NULL)
			break;

		dtp->dtp_addr = inst + off;
		dtp->dtp_type = KPROBE_ENTRY;
		dtp->dtp_nargs = db_ctf_func_numargs(symp);
		SLIST_INSERT_HEAD(&dtpf_entry[INSTTOIDX(dtp->dtp_addr)],
		    dtp, dtp_knext);
		dt_dev_register_probe(dtp);
		nb_probes++;

		off = db_epilogue_validate(symp);
		if (off < 0)
			continue;

		dtp = dt_dev_alloc_probe(name, "return", &dt_prov_kprobe);
		if (dtp == NULL)
			break;

		dtp->dtp_addr = inst + off;
		dtp->dtp_type = KPROBE_RETURN;
		SLIST_INSERT_HEAD(&dtpf_return[INSTTOIDX(dtp->dtp_addr)],
		    dtp, dtp_knext);
		dt_dev_register_probe(dtp);
		nb_probes++;
	}

	return nb_probes;
}

int
dt_prov_kprobe_alloc(struct dt_probe *dtp, struct dt_softc *sc,
    struct dt_pcb_list *plist, struct dtioc_req *dtrq)
{
	struct dt_pcb *dp;

	dp = dt_pcb_alloc(dtp, sc);
	if (dp == NULL)
		return ENOMEM;

	dtp->dtp_ref++;
	if (dtp->dtp_ref == 1) {
		switch (dtp->dtp_type) {
		case KPROBE_ENTRY:
			db_prologue_patch(dtp->dtp_addr, 0);
			break;
		case KPROBE_RETURN:
			db_epilogue_patch(dtp->dtp_addr, 0);
			break;
		default:
			panic("unknown probe type %d", dtp->dtp_type);
		}
	}

	dp->dp_evtflags = dtrq->dtrq_evtflags & DTEVT_PROV_KPROBE;
	TAILQ_INSERT_HEAD(plist, dp, dp_snext);
	return 0;
}

int
dt_prov_kprobe_dealloc(struct dt_probe *dtp, struct dt_softc *sc,
   struct dtioc_req *dtrq)
{
	dtp->dtp_ref--;
	if (dtp->dtp_ref > 0)
		return 0;

	switch (dtp->dtp_type) {
	case KPROBE_ENTRY:
		db_prologue_patch(dtp->dtp_addr, 1);
		break;
	case KPROBE_RETURN:
		db_epilogue_patch(dtp->dtp_addr, 1);
		break;
	default:
		panic("unknown probe type %d", dtp->dtp_type);
	}

	/* Deallocation of PCB is done by dt_pcb_purge when closing the dev */
	return 0;
}

int
dt_prov_kprobe_hook(struct dt_provider *dtpv, ...)
{
	struct dt_probe *dtp;
	struct dt_pcb *dp;
	struct trapframe *tf;
	va_list ap;
	int is_dt_bkpt = 0;
	int error;	/* Return values for return probes*/
	vaddr_t *args, addr;
	size_t argsize;
	register_t retval[2];

	KASSERT(dtpv == &dt_prov_kprobe);

	va_start(ap, dtpv);
	tf = va_arg(ap, struct trapframe*);
	va_end(ap);

	addr = db_get_probe_addr(tf);

	SLIST_FOREACH(dtp, &dtpf_entry[INSTTOIDX(addr)], dtp_knext) {
		if (dtp->dtp_addr != addr)
			continue;

		is_dt_bkpt = 1;
		if (db_prof_on)
			db_prof_count(tf);

		if (!dtp->dtp_recording)
			continue;

		smr_read_enter();
		SMR_SLIST_FOREACH(dp, &dtp->dtp_pcbs, dp_pnext) {
			struct dt_evt *dtev;

			dtev = dt_pcb_ring_get(dp, 0);
			if (dtev == NULL)
				continue;

#if defined(__amd64__)
			args = (vaddr_t *)tf->tf_rdi;
			/* XXX: use CTF to get the number of arguments. */
			argsize = 6;
#elif defined(__i386__)
			/* All args on stack */
			args = (vaddr_t *)(tf->tf_esp + 4);
			argsize = 10;
#endif

			if (ISSET(dp->dp_evtflags, DTEVT_FUNCARGS))
				memcpy(dtev->dtev_args, args, argsize);

			dt_pcb_ring_consume(dp, dtev);
		}
		smr_read_leave();
	}

	if (is_dt_bkpt)
		return is_dt_bkpt;

	SLIST_FOREACH(dtp, &dtpf_return[INSTTOIDX(addr)], dtp_knext) {
		if (dtp->dtp_addr != addr)
			continue;

		is_dt_bkpt = 2;

		if (!dtp->dtp_recording)
			continue;

		smr_read_enter();
		SMR_SLIST_FOREACH(dp, &dtp->dtp_pcbs, dp_pnext) {
			struct dt_evt *dtev;

			dtev = dt_pcb_ring_get(dp, 0);
			if (dtev == NULL)
				continue;

#if defined(__amd64__)
			retval[0] = tf->tf_rax;
			retval[1] = 0;
			error = 0;
#elif defined(__i386)
			retval[0] = tf->tf_eax;
			retval[1] = 0;
			error = 0;
#endif

			dtev->dtev_retval[0] = retval[0];
			dtev->dtev_retval[1] = retval[1];
			dtev->dtev_error = error;

			dt_pcb_ring_consume(dp, dtev);
		}
		smr_read_leave();
	}
	return is_dt_bkpt;
}

/* Called by ddb to patch all functions without allocating 1 pcb per probe */
void
dt_prov_kprobe_patch_all_entry(void)
{
	struct dt_probe *dtp;
	size_t i;

	for (i = 0; i < PPTMASK; ++i) {
		SLIST_FOREACH(dtp, &dtpf_entry[i], dtp_knext) {
			dtp->dtp_ref++;
			if (dtp->dtp_ref != 1)
				continue;

			db_prologue_patch(dtp->dtp_addr, 0);
		}
	}
}

/* Called by ddb to patch all functions without allocating 1 pcb per probe */
void
dt_prov_kprobe_depatch_all_entry(void)
{
	struct dt_probe *dtp;
	size_t i;

	for (i = 0; i < PPTMASK; ++i) {
		SLIST_FOREACH(dtp, &dtpf_entry[i], dtp_knext) {
			dtp->dtp_ref--;
			if (dtp->dtp_ref != 0)
				continue;

			db_prologue_patch(dtp->dtp_addr, 1);
		}

	}
}
#endif /* __amd64__ || __i386__ */