/*	$OpenBSD: kcov.c,v 1.18 2024/08/23 12:56:26 anton Exp $	*/

/*
 * Copyright (c) 2018 Anton Lindqvist <anton@openbsd.org>
 *
 * 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.
 */

#include <sys/types.h>
#include <sys/event.h>
#include <sys/ioctl.h>
#include <sys/kcov.h>
#include <sys/mman.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <sys/wait.h>

#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

struct context {
	int c_fd;
	int c_mode;
	unsigned long c_bufsize;
};

static int test_close(struct context *);
static int test_coverage(struct context *);
static int test_dying(struct context *);
static int test_exec(struct context *);
static int test_fdsend(struct context *);
static int test_fork(struct context *);
static int test_open(struct context *);
static int test_remote(struct context *);
static int test_remote_close(struct context *);
static int test_remote_interrupt(struct context *);
static int test_state(struct context *);

static int check_coverage(const unsigned long *, int, unsigned long, int);
static void do_syscall(void);
static void dump(const unsigned long *, int mode);
static void kcov_disable(int);
static void kcov_enable(int, int);
static int kcov_open(void);
static __dead void usage(void);

static const char *self;

int
main(int argc, char *argv[])
{
	struct {
		const char *name;
		int (*fn)(struct context *);
		int coverage;		/* test must produce coverage */
	} tests[] = {
		{ "close",		test_close,		0 },
		{ "coverage",		test_coverage,		1 },
		{ "dying",		test_dying,		-1 },
		{ "exec",		test_exec,		1 },
		{ "fdsend",		test_fdsend,		-1 },
		{ "fork",		test_fork,		1 },
		{ "open",		test_open,		0 },
		{ "remote",		test_remote,		1 },
		{ "remote-close",	test_remote_close,	0 },
		{ "remote-interrupt",	test_remote_interrupt,	-1 },
		{ "state",		test_state,		1 },
		{ NULL,			NULL,			0 },
	};
	struct context ctx;
	const char *errstr;
	unsigned long *cover, frac;
	int c, i;
	int error = 0;
	int prereq = 0;
	int reexec = 0;
	int verbose = 0;

	self = argv[0];

	memset(&ctx, 0, sizeof(ctx));
	ctx.c_bufsize = 256 << 10;

	while ((c = getopt(argc, argv, "b:Em:pv")) != -1)
		switch (c) {
		case 'b':
			frac = strtonum(optarg, 1, 100, &errstr);
			if (frac == 0)
				errx(1, "buffer size fraction %s", errstr);
			else if (frac > ctx.c_bufsize)
				errx(1, "buffer size fraction too large");
			ctx.c_bufsize /= frac;
			break;
		case 'E':
			reexec = 1;
			break;
		case 'm':
			if (strcmp(optarg, "pc") == 0)
				ctx.c_mode = KCOV_MODE_TRACE_PC;
			else if (strcmp(optarg, "cmp") == 0)
				ctx.c_mode = KCOV_MODE_TRACE_CMP;
			else
				errx(1, "unknown mode %s", optarg);
			break;
		case 'p':
			prereq = 1;
			break;
		case 'v':
			verbose = 1;
			break;
		default:
			usage();
		}
	argc -= optind;
	argv += optind;

	if (prereq) {
		ctx.c_fd = kcov_open();
		close(ctx.c_fd);
		return 0;
	}

	if (reexec) {
		do_syscall();
		return 0;
	}

	if (ctx.c_mode == 0 || argc != 1)
		usage();
	for (i = 0; tests[i].name != NULL; i++)
		if (strcmp(argv[0], tests[i].name) == 0)
			break;
	if (tests[i].name == NULL)
		errx(1, "%s: no such test", argv[0]);

	ctx.c_fd = kcov_open();
	if (ioctl(ctx.c_fd, KIOSETBUFSIZE, &ctx.c_bufsize) == -1)
		err(1, "ioctl: KIOSETBUFSIZE");
	cover = mmap(NULL, ctx.c_bufsize * sizeof(unsigned long),
	    PROT_READ | PROT_WRITE, MAP_SHARED, ctx.c_fd, 0);
	if (cover == MAP_FAILED)
		err(1, "mmap");

	*cover = 0;
	error = tests[i].fn(&ctx);
	if (verbose)
		dump(cover, ctx.c_mode);
	if (check_coverage(cover, ctx.c_mode, ctx.c_bufsize, tests[i].coverage))
		error = 1;

	if (munmap(cover, ctx.c_bufsize * sizeof(unsigned long)) == -1)
		err(1, "munmap");
	if (ctx.c_fd != -1) {
		if (close(ctx.c_fd) == -1)
			err(1, "close");
	}

	return error;
}

static __dead void
usage(void)
{
	fprintf(stderr, "usage: kcov [-Epv] [-b fraction] -m mode test\n");
	exit(1);
}

static void
do_syscall(void)
{
	getpid();
}

static int
check_coverage(const unsigned long *cover, int mode, unsigned long maxsize,
    int nonzero)
{
	unsigned long arg1, arg2, exp, i, pc, type;
	int error = 0;

	if (nonzero == -1) {
		return 0;
	} else if (nonzero && cover[0] == 0) {
		warnx("coverage empty (count=0)");
		return 1;
	} else if (!nonzero && cover[0] != 0) {
		warnx("coverage not empty (count=%lu)", *cover);
		return 1;
	} else if (cover[0] >= maxsize) {
		warnx("coverage overflow (count=%lu, max=%lu)",
		    *cover, maxsize);
		return 1;
	}

	if (mode == KCOV_MODE_TRACE_CMP) {
		if (*cover * 4 >= maxsize) {
			warnx("coverage cmp overflow (count=%lu, max=%lu)",
			    *cover * 4, maxsize);
			return 1;
		}

		for (i = 0; i < cover[0]; i++) {
			type = cover[i * 4 + 1];
			arg1 = cover[i * 4 + 2];
			arg2 = cover[i * 4 + 3];
			pc = cover[i * 4 + 4];

			exp = type >> 1;
			if (exp <= 3)
				continue;

			warnx("coverage cmp invalid size (i=%lu, exp=%lx, "
			    "const=%ld, arg1=%lu, arg2=%lu, pc=%p)\n",
			    i, exp, type & 0x1, arg1, arg2, (void *)pc);
			error = 1;
		}
	}

	return error;
}

static void
dump(const unsigned long *cover, int mode)
{
	unsigned long i;
	int stride = 1;

	if (mode == KCOV_MODE_TRACE_CMP)
		stride = 4;

	for (i = 0; i < cover[0]; i++)
		printf("%p\n", (void *)cover[i * stride + stride]);
}

static int
kcov_open(void)
{
	int fd;

	fd = open("/dev/kcov", O_RDWR);
	if (fd == -1)
		err(1, "open: /dev/kcov");
	return fd;
}

static void
kcov_enable(int fd, int mode)
{
	if (ioctl(fd, KIOENABLE, &mode) == -1)
		err(1, "ioctl: KIOENABLE");
}

static void
kcov_disable(int fd)
{
	if (ioctl(fd, KIODISABLE) == -1)
		err(1, "ioctl: KIODISABLE");
}

/*
 * Close before mmap.
 */
static int
test_close(struct context *ctx)
{
	int fd;

	fd = kcov_open();
	close(fd);
	return 0;
}

/*
 * Coverage of current thread.
 */
static int
test_coverage(struct context *ctx)
{
	kcov_enable(ctx->c_fd, ctx->c_mode);
	do_syscall();
	kcov_disable(ctx->c_fd);
	return 0;
}

static void *
closer(void *arg)
{
	struct context *ctx = arg;

	close(ctx->c_fd);
	return NULL;
}

/*
 * Close kcov descriptor in another thread during tracing.
 */
static int
test_dying(struct context *ctx)
{
	pthread_t th;
	int error;

	kcov_enable(ctx->c_fd, ctx->c_mode);

	if ((error = pthread_create(&th, NULL, closer, (void *)ctx)))
		errc(1, error, "pthread_create");
	if ((error = pthread_join(th, NULL)))
		errc(1, error, "pthread_join");

	error = 0;
	if (close(ctx->c_fd) == -1) {
		if (errno != EBADF)
			err(1, "close");
	} else {
		warnx("expected kcov descriptor to be closed");
		error = 1;
	}
	ctx->c_fd = -1;

	return error;
}

/*
 * Coverage of thread after exec.
 */
static int
test_exec(struct context *ctx)
{
	pid_t pid;
	int status;

	pid = fork();
	if (pid == -1)
		err(1, "fork");
	if (pid == 0) {
		kcov_enable(ctx->c_fd, ctx->c_mode);
		execlp(self, self, "-E", NULL);
		_exit(1);
	}

	if (waitpid(pid, &status, 0) == -1)
		err(1, "waitpid");
	if (WIFSIGNALED(status)) {
		warnx("terminated by signal (%d)", WTERMSIG(status));
		return 1;
	} else if (WEXITSTATUS(status) != 0) {
		warnx("non-zero exit (%d)", WEXITSTATUS(status));
		return 1;
	}

	/* Upon exit, the kcov descriptor must be reusable again. */
	kcov_enable(ctx->c_fd, ctx->c_mode);
	kcov_disable(ctx->c_fd);

	return 0;
}

/*
 * File descriptor send/receive is not allowed since remote coverage is tied to
 * the current process.
 */
static int
test_fdsend(struct context *ctx)
{
	struct msghdr msg;
	union {
		struct cmsghdr hdr;
		unsigned char buf[CMSG_SPACE(sizeof(int))];
	} cmsgbuf;
	struct cmsghdr *cmsg;
	int pair[2];

	if (socketpair(AF_UNIX, SOCK_DGRAM, 0, pair) == -1)
		err(1, "socketpair");

	memset(&msg, 0, sizeof(msg));
	msg.msg_control = &cmsgbuf.buf;
	msg.msg_controllen = sizeof(cmsgbuf.buf);
	cmsg = CMSG_FIRSTHDR(&msg);
	cmsg->cmsg_len = CMSG_LEN(sizeof(int));
	cmsg->cmsg_level = SOL_SOCKET;
	cmsg->cmsg_type = SCM_RIGHTS;
	*(int *)CMSG_DATA(cmsg) = ctx->c_fd;
	if (sendmsg(pair[1], &msg, 0) != -1)
		errx(1, "sendmsg: expected error");

	close(pair[0]);
	close(pair[1]);
	return 0;
}

/*
 * Coverage of thread after fork.
 */
static int
test_fork(struct context *ctx)
{
	pid_t pid;
	int status;

	pid = fork();
	if (pid == -1)
		err(1, "fork");
	if (pid == 0) {
		kcov_enable(ctx->c_fd, ctx->c_mode);
		do_syscall();
		_exit(0);
	}

	if (waitpid(pid, &status, 0) == -1)
		err(1, "waitpid");
	if (WIFSIGNALED(status)) {
		warnx("terminated by signal (%d)", WTERMSIG(status));
		return 1;
	} else if (WEXITSTATUS(status) != 0) {
		warnx("non-zero exit (%d)", WEXITSTATUS(status));
		return 1;
	}

	/* Upon exit, the kcov descriptor must be reusable again. */
	kcov_enable(ctx->c_fd, ctx->c_mode);
	kcov_disable(ctx->c_fd);

	return 0;
}

/*
 * Open /dev/kcov more than once.
 */
static int
test_open(struct context *ctx)
{
	unsigned long *cover;
	int fd;
	int error = 0;

	fd = kcov_open();
	if (ioctl(fd, KIOSETBUFSIZE, &ctx->c_bufsize) == -1)
		err(1, "ioctl: KIOSETBUFSIZE");
	cover = mmap(NULL, ctx->c_bufsize * sizeof(unsigned long),
	    PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
	if (cover == MAP_FAILED)
		err(1, "mmap");

	kcov_enable(fd, ctx->c_mode);
	do_syscall();
	kcov_disable(fd);

	error = check_coverage(cover, ctx->c_mode, ctx->c_bufsize, 1);

	if (munmap(cover, ctx->c_bufsize * sizeof(unsigned long)))
		err(1, "munmap");
	close(fd);

	return error;
}

/*
 * Remote taskq coverage. One reliable way to trigger a task on behalf of the
 * running process is to monitor a kqueue file descriptor using kqueue.
 */
static int
test_remote(struct context *ctx)
{
	struct kio_remote_attach remote = {
		.subsystem	= KCOV_REMOTE_COMMON,
		.id		= 0,
	};
	struct kevent kev;
	int kq1, kq2, pip[2];
	int x = 0;

	if (ioctl(ctx->c_fd, KIOREMOTEATTACH, &remote) == -1)
		err(1, "ioctl: KIOREMOTEATTACH");
	kcov_enable(ctx->c_fd, ctx->c_mode);

	kq1 = kqueue();
	if (kq1 == -1)
		err(1, "kqueue");
	kq2 = kqueue();
	if (kq1 == -1)
		err(1, "kqueue");
	EV_SET(&kev, kq2, EVFILT_READ, EV_ADD | EV_ENABLE, 0, 0, NULL);
	if (kevent(kq1, &kev, 1, NULL, 0, NULL) == -1)
		err(1, "kqueue");

	if (pipe(pip) == -1)
		err(1, "pipe");

	EV_SET(&kev, pip[0], EVFILT_READ, EV_ADD | EV_ENABLE, 0, 0, NULL);
	if (kevent(kq2, &kev, 1, NULL, 0, NULL) == -1)
		err(1, "kqueue");
	(void)write(pip[1], &x, sizeof(x));

	if (kevent(kq1, NULL, 0, &kev, 1, NULL) == -1)
		err(1, "kevent");

	kcov_disable(ctx->c_fd);

	return 0;
}

/*
 * Close with remote coverage enabled.
 */
static int
test_remote_close(struct context *ctx)
{
	struct kio_remote_attach remote = {
		.subsystem	= KCOV_REMOTE_COMMON,
		.id		= 0,
	};

	if (ioctl(ctx->c_fd, KIOREMOTEATTACH, &remote) == -1)
		err(1, "ioctl: KIOREMOTEATTACH");
	kcov_enable(ctx->c_fd, ctx->c_mode);
	if (close(ctx->c_fd) == -1)
		err(1, "close");
	ctx->c_fd = kcov_open();
	return 0;
}

/*
 * Remote interrupt coverage. There's no reliable way to enter a remote section
 * in interrupt context. This test can however by used to examine the coverage
 * collected in interrupt context:
 *
 *     $ until [ -s cov ]; do kcov -v -m pc remote-interrupt >cov; done
 */
static int
test_remote_interrupt(struct context *ctx)
{
	struct kio_remote_attach remote = {
		.subsystem	= KCOV_REMOTE_COMMON,
		.id		= 0,
	};
	int i;

	if (ioctl(ctx->c_fd, KIOREMOTEATTACH, &remote) == -1)
		err(1, "ioctl: KIOREMOTEATTACH");
	kcov_enable(ctx->c_fd, ctx->c_mode);

	for (i = 0; i < 100; i++)
		(void)getpid();

	kcov_disable(ctx->c_fd);

	return 0;
}

/*
 * State transitions.
 */
static int
test_state(struct context *ctx)
{
	if (ioctl(ctx->c_fd, KIOENABLE, &ctx->c_mode) == -1) {
		warn("KIOSETBUFSIZE -> KIOENABLE");
		return 1;
	}
	if (ioctl(ctx->c_fd, KIODISABLE) == -1) {
		warn("KIOENABLE -> KIODISABLE");
		return 1;
	}
	if (ioctl(ctx->c_fd, KIOSETBUFSIZE, 0) != -1) {
		warnx("KIOSETBUFSIZE -> KIOSETBUFSIZE");
		return 1;
	}
	if (ioctl(ctx->c_fd, KIODISABLE) != -1) {
		warnx("KIOSETBUFSIZE -> KIODISABLE");
		return 1;
	}

	kcov_enable(ctx->c_fd, ctx->c_mode);
	if (ioctl(ctx->c_fd, KIOENABLE, &ctx->c_mode) != -1) {
		warnx("KIOENABLE -> KIOENABLE");
		return 1;
	}
	if (ioctl(ctx->c_fd, KIOSETBUFSIZE, 0) != -1) {
		warnx("KIOENABLE -> KIOSETBUFSIZE");
		return 1;
	}
	kcov_disable(ctx->c_fd);

	return 0;
}