/* 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 v[12];

	if (read_child(pid, arg, v, sizeof(v)))
		return;
	/* NVOS21_PARAMETERS: hRoot, hObjectParent, hObjectNew, hClass,
	 * pAllocParms, pRightsRequested, paramsSize, flags, status */
	fprintf(stderr,
		"[iot]   alloc root=0x%x parent=0x%x new=0x%x class=0x%x params=%u flags=0x%x\n",
		v[0], v[1], v[2], v[3], v[6], v[7]);
}

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, ... */
	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\n",
		v[0], v[1], 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 == 0xc030462bULL)
						dump_alloc(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;
}