/* ioctl_trace.c - log the NVIDIA driver ioctls of a child process.
 *
 * libcuda issues raw syscall instructions, so LD_PRELOAD/ltrace never see
 * them; ptrace does. Logs every ioctl whose request has the NVIDIA 'F' magic
 * (0xc03046xx / 0xc01846xx / ...) with its arguments and return value, so the
 * last driver call before a crash is visible.
 *
 * Build (guest): gcc -O1 -w -o /tmp/iot ioctl_trace.c
 * Run:           /tmp/iot /tmp/cdt
 */
#define _GNU_SOURCE
#include <errno.h>
#include <signal.h>
#include <stdio.h>
#include <string.h>
#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#define SYS_IOCTL_X86_64 16

/* Read memory from the traced child. */
static int read_child(pid_t pid, unsigned long addr, void *buf, size_t len)
{
	size_t done = 0;
	unsigned char *p = buf;

	while (done < len) {
		long word = ptrace(PTRACE_PEEKDATA, pid, addr + done, 0);

		if (word == -1 && errno)
			return -1;
		memcpy(p + done, &word, sizeof(word) < len - done ?
		       sizeof(word) : len - done);
		done += sizeof(word);
	}
	return 0;
}

static void dump_alloc(pid_t pid, unsigned long arg, unsigned int size)
{
	unsigned int v[12];
	unsigned int st;

	if (read_child(pid, arg, v, sizeof(v)))
		return;
	/* NVOS21_PARAMETERS (0x20): ... paramsSize@0x14, status@0x10
	 * NVOS64_PARAMETERS (0x30): ... status@0x28 */
	st = (size == 0x30) ? v[10] : v[4];
	fprintf(stderr,
		"[iot]   alloc root=0x%x parent=0x%x new=0x%x class=0x%x st=0x%x\n",
		v[0], v[1], v[2], v[3], st);
}

static void dump_free(pid_t pid, unsigned long arg)
{
	unsigned int v[4];

	if (read_child(pid, arg, v, sizeof(v)))
		return;
	/* NVOS00_PARAMETERS: hRoot, hObjectParent, hObjectOld, status */
	fprintf(stderr, "[iot]   free old=0x%x st=0x%x\n", v[2], v[3]);
}

/* Decode the interesting RM control replies (params buffer read at syscall
 * exit, so for OUT params this is the answer the client received). */
static void decode_ctrl_reply(pid_t pid, unsigned int cmd,
			      unsigned long long pParams, unsigned int size)
{
	unsigned int w[8];
	unsigned long long listPtr;
	unsigned int i, n;

	if (!pParams || !size)
		return;

	switch (cmd) {
	case 0x20800123: /* GPU_GET_ENGINES (V1): {count, ptr} */
		if (size < 16 || read_child(pid, pParams, w, 16))
			return;
		memcpy(&listPtr, &w[2], sizeof(listPtr));
		n = w[0];
		fprintf(stderr, "[iot]     V1 engineCount=%u list@0x%llx:", n, listPtr);
		if (n > 0x54)
			n = 0x54;
		for (i = 0; i < n; i++) {
			unsigned int e = 0;
			if (listPtr && read_child(pid, listPtr + i * 4, &e, 4))
				break;
			fprintf(stderr, " 0x%x", e);
		}
		fprintf(stderr, "\n");
		break;
	case 0x20800170: /* GPU_GET_ENGINES_V2: {count, list[]} */
		if (size < 4 || read_child(pid, pParams, w, 8))
			return;
		n = w[0];
		fprintf(stderr, "[iot]     V2 engineCount=%u:", n);
		if (n > 0x54)
			n = 0x54;
		for (i = 0; i < n; i++) {
			unsigned int e = 0;
			if (read_child(pid, pParams + 4 + i * 4, &e, 4))
				break;
			fprintf(stderr, " 0x%x", e);
		}
		fprintf(stderr, "\n");
		break;
	case 0x20800124: /* GPU_GET_ENGINE_CLASSLIST: {engineType, n, ptr} */
		if (size < 16 || read_child(pid, pParams, w, 16))
			return;
		memcpy(&listPtr, &w[4], sizeof(listPtr));
		n = w[2];
		fprintf(stderr, "[iot]     classlist eng=0x%x n=%u:", w[0], n);
		if (n > 64)
			n = 64;
		for (i = 0; i < n; i++) {
			unsigned int c = 0;
			if (listPtr && read_child(pid, listPtr + i * 4, &c, 4))
				break;
			fprintf(stderr, " 0x%x", c);
		}
		fprintf(stderr, "\n");
		break;
	case 0x20802a07:  /* CE_GET_PHYSICAL_CAPS: {type, caps[8]} */
	case 0x20802a01:  /* CE_GET_CAPS: {type, capsTblSize, capsTbl ptr...} */
	case 0x20802a02: { /* CE_GET_CE_PCE_MASK: {type, pceMask} */
		if (size < 8 || read_child(pid, pParams, w, 8))
			return;
		fprintf(stderr, "[iot]     ce ctrl 0x%x type=0x%x w4=0x%x w5=0x%x\n",
			cmd, w[0], w[4], w[5]);
		break;
	}
	case 0x20802a0a:  /* CE_GET_ALL_CAPS */
	case 0x20802a0b: { /* CE_GET_ALL_PHYSICAL_CAPS: 128 caps + 8 present */
		unsigned int present_lo = 0, present_hi = 0;
		if (size < 136)
			return;
		if (read_child(pid, pParams + 128, &present_lo, 4))
			return;
		if (read_child(pid, pParams + 132, &present_hi, 4))
			return;
		fprintf(stderr, "[iot]     ce all-caps 0x%x present=0x%x%08x\n",
			cmd, present_hi, present_lo);
		break;
	}
	default:
		break;
	}
}

static void dump_control(pid_t pid, unsigned long arg)
{
	unsigned int v[8];
	unsigned long long pParams;
	unsigned int cmd;

	if (read_child(pid, arg, v, sizeof(v)))
		return;
	/* NVOS54_PARAMETERS: hClient, hObject, cmd, flags, params, paramsSize,
	 * status@0x1c */
	cmd = v[2];
	memcpy(&pParams, &v[4], sizeof(pParams));
	fprintf(stderr,
		"[iot]   control client=0x%x object=0x%x cmd=0x%x params=0x%llx size=%u st=0x%x\n",
		v[0], v[1], cmd, pParams, v[5], v[7]);
	decode_ctrl_reply(pid, cmd, pParams, v[5]);
}

int main(int argc, char **argv)
{
	pid_t pid;
	int status;
	unsigned long long nr = 0, a0 = 0, a1 = 0, a2 = 0, a3 = 0;
	int entry_args_valid = 0;

	if (argc < 2) {
		fprintf(stderr, "usage: %s <prog> [args...]\n", argv[0]);
		return 1;
	}

	pid = fork();
	if (pid == 0) {
		ptrace(PTRACE_TRACEME, 0, 0, 0);
		raise(SIGSTOP);
		execv(argv[1], &argv[1]);
		perror("execv");
		_exit(127);
	}

	waitpid(pid, &status, 0);
	ptrace(PTRACE_SETOPTIONS, pid, 0, PTRACE_O_TRACESYSGOOD);
	ptrace(PTRACE_SYSCALL, pid, 0, 0);

	while (1) {
		if (waitpid(pid, &status, 0) < 0)
			break;

		if (WIFEXITED(status)) {
			fprintf(stderr, "[iot] child exited rc=%d\n",
				WEXITSTATUS(status));
			break;
		}
		if (WIFSIGNALED(status)) {
			fprintf(stderr, "[iot] child killed by signal %d\n",
				WTERMSIG(status));
			break;
		}
		if (!WIFSTOPPED(status))
			continue;
		if (WSTOPSIG(status) != (SIGTRAP | 0x80))
			goto cont;

		if (!entry_args_valid) {
			/* syscall entry */
			struct user_regs_struct {
				unsigned long long r15, r14, r13, r12, rbp,
					rbx, r11, r10, r9, r8, rax, rcx,
					rdx, rsi, rdi, orig_rax, rip, cs,
					eflags, rsp, ss;
			} regs;

			if (ptrace(PTRACE_GETREGS, pid, 0, &regs) == 0 &&
			    regs.orig_rax == SYS_IOCTL_X86_64) {
				nr = regs.orig_rax;
				a0 = regs.rdi;
				a1 = regs.rsi;
				a2 = regs.rdx;
				a3 = regs.r10;
				entry_args_valid = 1;
			}
		} else {
			/* syscall exit */
			struct user_regs_struct {
				unsigned long long r15, r14, r13, r12, rbp,
					rbx, r11, r10, r9, r8, rax, rcx,
					rdx, rsi, rdi, orig_rax, rip, cs,
					eflags, rsp, ss;
			} regs;

			if (entry_args_valid &&
			    ptrace(PTRACE_GETREGS, pid, 0, &regs) == 0) {
				unsigned long long req = a1;
				unsigned long long ret =
					(unsigned long long)regs.rax;

				if (((req >> 8) & 0xff) == 0x46) {
					fprintf(stderr,
						"[iot] req=0x%llx arg=0x%llx ret=0x%llx\n",
						req, a2, ret);
					if ((req & 0xff) == 0x2b)
						dump_alloc(pid, a2, (req >> 16) & 0x3fff);
					else if ((req & 0xff) == 0x29)
						dump_free(pid, a2);
					else if ((req & 0xff) == 0x2a)
						dump_control(pid, a2);
				}
			}
			entry_args_valid = 0;
			(void)nr; (void)a0; (void)a3;
		}
cont:
		ptrace(PTRACE_SYSCALL, pid, 0, 0);
	}
	return 0;
}