/*
 * Linuxulator litmus for x86 memory protection keys:
 * pkey_alloc(2), pkey_free(2), pkey_mprotect(2), and the PKRU
 * register (RDPKRU/WRPKRU via glibc pkey_get/pkey_set).
 *
 * Verdicts:
 *   - pkey_alloc -> ENOSYS: kernel lacks the pkey syscalls (stubs);
 *     apply the linux-pkru patch, reload the linux modules, re-run.
 *   - pkey_alloc -> ENOSPC with no keys allocated: CPU or kernel does
 *     not support PKU (check: Structured Extended Features2 OSPKE in
 *     dmesg); this is also what Linux returns on such hardware.
 *   - Otherwise all tests should PASS.
 *
 * Tests:
 *   0. initial PKRU denies access to unallocated keys (the Linux
 *      init_pkru exec default, 0x55555554)
 *   1. pkey_alloc(0, 0) returns a key in 1..15
 *   2. pkey_mprotect tags a page; access allowed with rights 0
 *   3. pkey_set(key, PKEY_DISABLE_ACCESS): read faults (SIGSEGV)
 *   4. pkey_set(key, 0): access restored
 *   5. pkey_set(key, PKEY_DISABLE_WRITE): read OK, write faults
 *   6. pkey_alloc(0, PKEY_DISABLE_WRITE): pkey_get reflects the
 *      initial rights (kernel must set the caller's PKRU)
 *   7. fork: child inherits the allocation map (pkey_free succeeds)
 *   8. pkey_mprotect with pkey = -1 acts as plain mprotect
 *   9. pkey_free; double free fails EINVAL
 *  10. allocation exhausts at 15 keys total, then ENOSPC
 *
 * Build (linux-rl9-devtools):
 *   /compat/linux/usr/bin/gcc -O2 -o ~/linux-pkey-litmus \
 *       ~/linux-pkey-litmus.c
 */

#define _GNU_SOURCE
#include <errno.h>
#include <setjmp.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/wait.h>
#include <unistd.h>

static sigjmp_buf fault_env;
static volatile sig_atomic_t fault_code;
static int fault_key = -1;

static void
segv_handler(int sig, siginfo_t *si, void *uc)
{
	(void)sig;
	(void)uc;
	fault_code = si->si_code;
	if (fault_key >= 0)
		pkey_set(fault_key, 0);	/* re-allow so we can continue */
	siglongjmp(fault_env, 1);
}

static int failed;

static void
check(int ok, const char *what)
{
	printf("%s: %s\n", what, ok ? "OK" : "FAIL");
	if (!ok)
		failed = 1;
}

int
main(void)
{
	struct sigaction sa;
	volatile char *page;
	char probe;
	int key, key2, keys[16], nkeys, i;
	pid_t pid;
	int status;

	memset(&sa, 0, sizeof(sa));
	sa.sa_sigaction = segv_handler;
	sa.sa_flags = SA_SIGINFO;
	if (sigaction(SIGSEGV, &sa, NULL) != 0) {
		perror("sigaction");
		return (1);
	}

	/* 1. Allocate a key. */
	key = pkey_alloc(0, 0);
	if (key < 0) {
		if (errno == ENOSYS) {
			fprintf(stderr,
			    "pkey_alloc: ENOSYS\n"
			    "VERDICT: kernel lacks the pkey syscalls; apply\n"
			    "the linux-pkru patch, reload the linux modules\n"
			    "(linux_common, linux64, linux), and re-run.\n");
			return (1);
		}
		if (errno == ENOSPC) {
			fprintf(stderr,
			    "pkey_alloc: ENOSPC\n"
			    "VERDICT: no protection keys available; this\n"
			    "CPU/kernel has no PKU (OSPKE).  Matches Linux\n"
			    "behavior on such hardware; nothing to fix.\n");
			return (1);
		}
		perror("pkey_alloc");
		return (1);
	}
	printf("pkey_alloc(0, 0) = %d\n", key);
	check(key >= 1 && key <= 15, "key in 1..15");

	/*
	 * 0. Initial PKRU (checked via a key that is still unallocated;
	 * RDPKRU is only safe once the alloc probe above confirmed PKU).
	 */
	check(pkey_get(15) == PKEY_DISABLE_ACCESS,
	    "unallocated key denied by initial PKRU (init_pkru)");

	page = mmap(NULL, 4096, PROT_READ | PROT_WRITE,
	    MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
	if (page == MAP_FAILED) {
		perror("mmap");
		return (1);
	}
	page[0] = 'x';

	/* 2. Tag the page. */
	check(pkey_mprotect((void *)page, 4096, PROT_READ | PROT_WRITE,
	    key) == 0, "pkey_mprotect(key)");
	probe = page[0];
	check(probe == 'x', "read with rights 0");

	/* 3. Deny access; read must fault. */
	fault_key = key;
	fault_code = 0;
	if (sigsetjmp(fault_env, 1) == 0) {
		pkey_set(key, PKEY_DISABLE_ACCESS);
		probe = page[0];
		check(0, "read with PKEY_DISABLE_ACCESS faults");
	} else {
		check(1, "read with PKEY_DISABLE_ACCESS faults");
		printf("  (si_code = %d%s)\n", (int)fault_code,
		    fault_code == SEGV_PKUERR ? " = SEGV_PKUERR" : "");
	}

	/* 4. Restore access. */
	pkey_set(key, 0);
	probe = page[0];
	check(probe == 'x', "read after rights restored");

	/* 5. Write-disable: read OK, write faults. */
	if (sigsetjmp(fault_env, 1) == 0) {
		pkey_set(key, PKEY_DISABLE_WRITE);
		probe = page[0];
		check(probe == 'x', "read with PKEY_DISABLE_WRITE");
		page[0] = 'y';
		check(0, "write with PKEY_DISABLE_WRITE faults");
	} else {
		check(1, "write with PKEY_DISABLE_WRITE faults");
	}
	pkey_set(key, 0);

	/* 6. Initial rights applied by the kernel at alloc time. */
	key2 = pkey_alloc(0, PKEY_DISABLE_WRITE);
	check(key2 >= 1, "second pkey_alloc");
	check(pkey_get(key2) == PKEY_DISABLE_WRITE,
	    "pkey_alloc initial rights visible in PKRU");
	check(pkey_free(key2) == 0, "pkey_free(second key)");

	/* 7. Fork inherits the allocation map. */
	pid = fork();
	if (pid == 0) {
		_exit(pkey_free(key) == 0 ? 0 : 1);
	}
	if (waitpid(pid, &status, 0) < 0) {
		perror("waitpid");
		return (1);
	}
	check(WIFEXITED(status) && WEXITSTATUS(status) == 0,
	    "child inherits allocation map (pkey_free in child)");

	/* 8. pkey = -1 behaves as plain mprotect. */
	check(pkey_mprotect((void *)page, 4096, PROT_READ, -1) == 0,
	    "pkey_mprotect(-1) as plain mprotect");

	/* 9. Free; double free must fail. */
	check(pkey_free(key) == 0, "pkey_free");
	check(pkey_free(key) == -1 && errno == EINVAL,
	    "double pkey_free fails EINVAL");

	/* 10. Exhaustion: 15 keys total, then ENOSPC. */
	nkeys = 0;
	for (i = 0; i < 16; i++) {
		keys[nkeys] = pkey_alloc(0, 0);
		if (keys[nkeys] < 0)
			break;
		nkeys++;
	}
	check(nkeys == 15 && errno == ENOSPC,
	    "15 keys allocatable, then ENOSPC");
	for (i = 0; i < nkeys; i++)
		pkey_free(keys[i]);

	printf("VERDICT: %s\n", failed ? "FAIL" : "PASS");
	return (failed);
}