/* Differential + guard-page rig for the audited file itself:
 *   /usr/src/sys/libkern/strncpy.c  (compiled UNMODIFIED, -Dstrncpy=kstrncpy)
 *
 * A. Semantics differential vs libc strncpy: byte-exact dst image + return
 *    value over the full (n, srclen) matrix 0..64 plus 100k random cases
 *    (random content incl. embedded NULs).
 * B. Pad-count math: for every (n, srclen), exactly n bytes are written:
 *    [0,k)=src, [k,n)=NUL-pad, canary after n untouched, 0<n.
 * C. Guard-page src: with a NUL inside the window the impl must stop
 *    reading src (pad loop may not re-read src); with no NUL and n equal
 *    to the readable window it must read exactly n bytes and no more.
 * D. Guard-page dst: n == space -> exactly n writes, no fault;
 *    n == space+1 -> fault (bound is exactly n; over-n is caller misuse).
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <sys/mman.h>
#include <sys/wait.h>
#include <unistd.h>

char *kstrncpy(char *dst, const char *src, size_t n); /* kernel source */

static int failures;

#define CHECK(cond, ...) do {						\
	if (!(cond)) { failures++; printf("FAIL: " __VA_ARGS__); printf("\n"); } \
} while (0)

#define MAXN 80
#define PAGESZ 4096

static void
part_a(void)
{
	size_t n, slen, i;
	char src[2 * MAXN], d1[MAXN], d2[MAXN], *r1, *r2;

	for (n = 0; n <= 64; n++) {
		for (slen = 0; slen <= 64; slen++) {
			memset(src, 'S', sizeof(src));
			src[slen] = 0;			/* NUL at slen */
			memset(d1, 0x5A, sizeof(d1));
			memset(d2, 0x5A, sizeof(d2));
			r1 = kstrncpy(d1, src, n);
			r2 = strncpy(d2, src, n);
			CHECK(r1 == d1 && r2 == d2, "return n=%zu slen=%zu", n, slen);
			CHECK(memcmp(d1, d2, sizeof(d1)) == 0,
			    "image n=%zu slen=%zu", n, slen);
		}
	}
	/* random fuzz incl. embedded NULs */
	srandom(4242);
	for (i = 0; i < 100000; i++) {
		size_t len = (size_t)(random() % (2 * MAXN));
		n = (size_t)(random() % 65);
		for (slen = 0; slen < len; slen++)
			src[slen] = (char)(1 + random() % 255); /* no NULs */
		src[len] = 0;
		if (len > n && (random() & 1))	/* sometimes embed a NUL early */
			src[random() % len] = 0;
		memset(d1, 0x5A, sizeof(d1));
		memset(d2, 0x5A, sizeof(d2));
		kstrncpy(d1, src, n);
		strncpy(d2, src, n);
		CHECK(memcmp(d1, d2, sizeof(d1)) == 0, "fuzz i=%zu n=%zu", i, n);
	}
	printf("part A (differential vs libc, matrix+fuzz): %s\n",
	    failures ? "FAIL" : "identical");
}

static void
part_b(void)
{
	size_t n, slen, k;
	char src[3 * MAXN], dst[MAXN + 16], ref[MAXN + 16];

	for (n = 1; n <= 64; n++) {
		for (slen = 0; slen <= 72; slen++) {
			size_t kmax = slen < n ? slen : n;
			memset(src, 'S', sizeof(src));
			src[slen] = 0;
			memset(dst, 0xA5, sizeof(dst));
			memcpy(ref, dst, sizeof(ref));
			kstrncpy(dst, src, n);
			/* expected: [0,kmax)=src bytes, [kmax,n)=0, rest canary */
			for (k = 0; k < n; k++) {
				char want = k < kmax ? src[k] : 0;
				CHECK(dst[k] == want,
				    "pad-math n=%zu slen=%zu k=%zu got %02x want %02x",
				    n, slen, k, (unsigned char)dst[k],
				    (unsigned char)want);
			}
			for (k = n; k < sizeof(dst); k++)
				CHECK((unsigned char)dst[k] == 0xA5,
				    "canary n=%zu slen=%zu k=%zu (wrote past n)",
				    n, slen, k);
			(void)ref;
		}
	}
	printf("part B (pad-count: exactly n bytes, canary intact): %s\n",
	    failures ? "FAIL" : "exact");
}

/* run fn in a forked child; return 0 clean exit, 1 SIGSEGV, -1 other */
static int
fork_run(void (*fn)(void *), void *arg)
{
	pid_t pid = fork();
	int st;

	if (pid == 0) {
		fn(arg);
		_exit(0);
	}
	waitpid(pid, &st, 0);
	if (WIFSIGNALED(st) && WTERMSIG(st) == SIGSEGV)
		return 1;
	if (WIFEXITED(st) && WEXITSTATUS(st) == 0)
		return 0;
	return -1;
}

struct gp_args {
	int mode;	/* 0: src-NUL-inside, 1: src-full-window, 2: dst-exact,
			 * 3: dst-overrun(self-check), 4: src-overrun(self-check) */
	size_t n;
};

static char *
gp_map(size_t pages)
{
	char *m = mmap(NULL, pages * PAGESZ, PROT_READ | PROT_WRITE,
	    MAP_PRIVATE | MAP_ANON, -1, 0);
	if (m == MAP_FAILED) { perror("mmap"); exit(2); }
	if (mprotect(m + (pages - 1) * PAGESZ, PAGESZ, PROT_NONE) != 0) {
		perror("mprotect"); exit(2);
	}
	return m;
}

static void
gp_fn(void *v)
{
	struct gp_args *a = v;
	char *m = gp_map(2);
	size_t room = PAGESZ;

	if (a->mode == 0) {		/* NUL inside window; pad must not read src */
		char *src = m + room - 24;
		char dst[64];
		memset(src, 'Q', 24);
		src[4] = 0;			/* NUL at src[4], 19 bytes to guard */
		memset(dst, 0x5A, sizeof(dst));
		kstrncpy(dst, src, 40);		/* n=40 >> bytes-to-guard(19) */
		_exit(memcmp(dst, "QQQQ", 4) == 0 && dst[5] == 0 ? 0 : 3);
	}
	if (a->mode == 1) {		/* no NUL, n == readable bytes exactly */
		char *src = m;		/* fill whole page non-NUL */
		char dst[PAGESZ];
		memset(src, 'Q', room);
		kstrncpy(dst, src, room);	/* reads exactly n, must not fault */
		_exit(dst[room - 1] == 'Q' ? 0 : 3);
	}
	if (a->mode == 2) {		/* dst exactly n bytes before guard */
		char *dst = m + room - 40;
		const char *src = "hello world";
		kstrncpy(dst, src, 40);		/* writes exactly 40, no fault */
		_exit(dst[39] == 0 ? 0 : 3);
	}
	if (a->mode == 3) {		/* SELF-CHECK: dst n = room+1 -> MUST fault */
		char *dst = m + room - 40;
		kstrncpy(dst, "x", 41);		/* 41 > 40 bytes to guard */
		_exit(0);			/* reaching here = no fault = WRONG */
	}
	if (a->mode == 4) {		/* SELF-CHECK: src read past n -> fault */
		char *src;			/* only 4 readable bytes, no NUL */
		char dst[64];
		memset(m, 'Q', room);		/* fresh anon mem is ZEROED: fill! */
		src = m + room - 4;
		kstrncpy(dst, src, 8);		/* reads 8 > 4 -> fault expected */
		_exit(0);
	}
	_exit(9);
}

static void
part_cd(void)
{
	struct gp_args a;
	int r;

	a.n = 0;
	a.mode = 0; r = fork_run(gp_fn, &a);
	CHECK(r == 0, "C0 src-guard NUL-inside n=40 (r=%d)", r);
	printf("part C0 src-guard, NUL in window (pad must not read src): %s\n",
	    r == 0 ? "no fault (correct)" : "FAIL");

	a.mode = 1; r = fork_run(gp_fn, &a);
	CHECK(r == 0, "C1 src-guard full-window n=4096 (r=%d)", r);
	printf("part C1 src-guard, no NUL, n==window: %s\n",
	    r == 0 ? "no fault (reads exactly n)" : "FAIL");

	a.mode = 2; r = fork_run(gp_fn, &a);
	CHECK(r == 0, "D0 dst-guard n==space (r=%d)", r);
	printf("part D0 dst-guard, n==space: %s\n",
	    r == 0 ? "no fault (writes exactly n)" : "FAIL");

	a.mode = 3; r = fork_run(gp_fn, &a);
	CHECK(r == 1, "D1 dst-guard n==space+1 must fault (r=%d)", r);
	printf("part D1 dst-guard, n==space+1 (caller misuse): %s\n",
	    r == 1 ? "SIGSEGV (bound is exactly n)" : "NO FAULT (wrong)");

	a.mode = 4; r = fork_run(gp_fn, &a);
	CHECK(r == 1, "D2 src-guard n>window must fault (r=%d)", r);
	printf("part D2 src-guard, n==window+1 (caller misuse): %s\n",
	    r == 1 ? "SIGSEGV (reads at most n)" : "NO FAULT (wrong)");
}

int
main(void)
{
	printf("rig: kernel %s compiled from /usr/src/sys/libkern/strncpy.c\n",
	    "strncpy");
	part_a();
	part_b();
	part_cd();
	printf("\n%s (%d failures)\n", failures ? "RIG FAILURES" :
	    "ALL PARTS PASS: implementation is libc-identical, reads<=n, "
	    "writes==n, pad-count exact", failures);
	return failures ? 1 : 0;
}