/*
 * Minimal MGLRU memcg reparent race PoC.
 *
 * This program expects the companion instrumentation patch to add a short
 * delay before reset_batch_size() for cgroups named /lru_gen_race_* and to log
 * "delay_before_reset".  The program waits for that log line, tears down the
 * target memcg, and lets the stale MGLRU batch commit into the offlined child.
 */

#define _GNU_SOURCE
#include <errno.h>
#include <fcntl.h>
#include <pthread.h>
#include <signal.h>
#include <stdatomic.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/mount.h>
#include <sys/prctl.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <unistd.h>

#define CGROUP_ROOT "/sys/fs/cgroup"
#define LRU_GEN_FILE "/sys/kernel/debug/lru_gen"
#define PAGE_BYTES 4096UL
#define MAX_NODES 8
#define DEFAULT_ITERS 20
#define DEFAULT_FILE_MIB 32
#define WINDOW_TIMEOUT_MS 90000
#define CONFIRM_TIMEOUT_MS 5000

enum {
	PHASE_IDLE = 0,
	PHASE_WINDOW = 1,
	PHASE_DONE = 2,
};

/* LruInfo stores the debugfs MGLRU ids needed to issue aging commands. */
struct lru_info {
	unsigned long memcg_id; /* Memcg id accepted by /sys/kernel/debug/lru_gen. */
	int nr_nodes; /* Number of NUMA nodes parsed for this memcg. */
	int nodes[MAX_NODES]; /* Node ids parsed from /sys/kernel/debug/lru_gen. */
	unsigned long max_seq[MAX_NODES]; /* Latest generation sequence for each node. */
};

/* RaceState is shared by the ager, kmsg reader, and main iteration. */
struct race_state {
	char leaf[1024]; /* Absolute cgroup path for the current leaf. */
	char leaf_rel[1024]; /* Cgroup path as printed by cgroup_path(). */
	char file_path[512]; /* File mapped by the victim process. */
	pid_t victim; /* Victim pid charged to the leaf memcg. */
	atomic_int phase; /* Current synchronization phase. */
	int iter; /* Iteration index used only for concise progress output. */
};

/* Die prints a syscall failure and exits the process. */
static void die(const char *what)
{
	perror(what);
	exit(1);
}

/* WriteFile writes a short string into a sysfs/cgroupfs control file. */
static int write_file(const char *path, const char *value)
{
	int fd = open(path, O_WRONLY | O_CLOEXEC);

	if (fd < 0)
		return -1;

	ssize_t ret = write(fd, value, strlen(value));
	int saved_errno = errno;

	close(fd);
	errno = saved_errno;
	return ret < 0 ? -1 : 0;
}

/* MkdirIfMissing creates a cgroup directory if it is not already present. */
static int mkdir_if_missing(const char *path)
{
	if (!mkdir(path, 0755) || errno == EEXIST)
		return 0;

	return -1;
}

/* EnableMemoryController enables memory accounting below a cgroup. */
static void enable_memory_controller(const char *cg)
{
	char path[640];

	snprintf(path, sizeof(path), "%s/cgroup.subtree_control", cg);
	(void)write_file(path, "+memory");
}

/* MovePid moves a process into the target cgroup. */
static int move_pid(const char *cg, pid_t pid)
{
	char path[640];
	char value[32];

	snprintf(path, sizeof(path), "%s/cgroup.procs", cg);
	snprintf(value, sizeof(value), "%d", (int)pid);
	return write_file(path, value);
}

/* RmdirRetry removes a cgroup after css teardown has made it removable. */
static int rmdir_retry(const char *path)
{
	for (int i = 0; i < 600; i++) {
		if (!rmdir(path))
			return 0;
		if (errno != EBUSY && errno != EINVAL)
			return -1;
		usleep(5000);
	}

	return -1;
}

/* WaitPhase waits until the shared phase reaches the requested value. */
static bool wait_phase(struct race_state *st, int want, int timeout_ms)
{
	for (int i = 0; i < timeout_ms; i++) {
		if (atomic_load(&st->phase) >= want)
			return true;
		usleep(1000);
	}

	return false;
}

/* VictimMain faults file-backed pages after it has been moved into the leaf. */
static void victim_main(int start_fd, int ready_fd, const char *path, size_t bytes)
{
	char ch;

	if (read(start_fd, &ch, 1) != 1)
		_exit(10);
	close(start_fd);

	int fd = open(path, O_CREAT | O_TRUNC | O_RDWR | O_CLOEXEC, 0600);

	if (fd < 0)
		_exit(11);
	if (ftruncate(fd, (off_t)bytes))
		_exit(12);

	volatile uint8_t *mapping = mmap(NULL, bytes, PROT_READ | PROT_WRITE,
					 MAP_SHARED, fd, 0);
	if (mapping == MAP_FAILED)
		_exit(13);
	close(fd);

	for (size_t off = 0; off < bytes; off += PAGE_BYTES)
		mapping[off] = (uint8_t)(off >> 12);

	ssize_t ready_ret = write(ready_fd, "R", 1);

	(void)ready_ret;
	close(ready_fd);

	for (uint8_t seed = 1;; seed++) {
		for (size_t off = 0; off < bytes; off += PAGE_BYTES)
			mapping[off] ^= seed;
	}
}

/* ReadLruInfo parses the target memcg section from /sys/kernel/debug/lru_gen. */
static int read_lru_info(const char *leaf_rel, struct lru_info *info)
{
	FILE *file = fopen(LRU_GEN_FILE, "r");
	char *line = NULL;
	size_t cap = 0;
	bool in_target = false;
	int current = -1;
	int ret = -1;

	if (!file)
		return -1;

	memset(info, 0, sizeof(*info));

	while (getline(&line, &cap, file) > 0) {
		unsigned long id;
		unsigned long seq;
		char path[1024];
		int node;

		if (sscanf(line, " memcg %lu %1023s", &id, path) == 2) {
			in_target = !strcmp(path, leaf_rel);
			current = -1;
			if (in_target) {
				info->memcg_id = id;
				ret = 0;
			}
			continue;
		}

		if (!in_target)
			continue;

		if (sscanf(line, " node %d", &node) == 1) {
			if (info->nr_nodes >= MAX_NODES)
				continue;
			current = info->nr_nodes++;
			info->nodes[current] = node;
			continue;
		}

		if (current >= 0 && sscanf(line, " %lu", &seq) == 1 &&
		    seq > info->max_seq[current])
			info->max_seq[current] = seq;
	}

	free(line);
	fclose(file);
	return ret;
}

/* AgerThread repeatedly asks MGLRU debugfs to age the target memcg. */
static void *ager_thread(void *arg)
{
	struct race_state *st = arg;
	int fd;

	prctl(PR_SET_NAME, "lru_ager", 0, 0, 0);

	fd = open(LRU_GEN_FILE, O_WRONLY | O_CLOEXEC);
	if (fd < 0)
		return NULL;

	while (atomic_load(&st->phase) < PHASE_WINDOW) {
		struct lru_info info;

		if (read_lru_info(st->leaf_rel, &info) || !info.memcg_id)
			break;

		for (int i = 0; i < info.nr_nodes; i++) {
			char cmd[128];

			if (atomic_load(&st->phase) >= PHASE_WINDOW)
				break;

			snprintf(cmd, sizeof(cmd), "+ %lu %d %lu 1 1\n",
				 info.memcg_id, info.nodes[i], info.max_seq[i]);
			ssize_t write_ret = write(fd, cmd, strlen(cmd));

			(void)write_ret;
		}
	}

	close(fd);
	return NULL;
}

/* KmsgThread watches for the instrumentation lines used for synchronization. */
static void *kmsg_thread(void *arg)
{
	struct race_state *st = arg;
	char buf[8192];
	int fd = open("/dev/kmsg", O_RDONLY | O_NONBLOCK | O_CLOEXEC);

	if (fd < 0)
		return NULL;

	lseek(fd, 0, SEEK_END);

	while (atomic_load(&st->phase) < PHASE_DONE) {
		ssize_t len = read(fd, buf, sizeof(buf) - 1);
		const char *msg;
		bool ours;

		if (len <= 0) {
			usleep(500);
			continue;
		}

		buf[len] = '\0';
		msg = strchr(buf, ';');
		if (msg)
			msg++;
		else
			msg = buf;

		ours = strstr(msg, st->leaf_rel) != NULL;
		if (ours && strstr(msg, "delay_before_reset")) {
			int idle = PHASE_IDLE;

			atomic_compare_exchange_strong(&st->phase, &idle, PHASE_WINDOW);
		}

		if ((ours && strstr(msg, "exit_nonzero")) ||
		    (atomic_load(&st->phase) >= PHASE_WINDOW &&
		     strstr(msg, "WARNING: mm/vmscan.c")))
			atomic_store(&st->phase, PHASE_DONE);
	}

	close(fd);
	return NULL;
}

/* RunIteration creates one child memcg and races its teardown against aging. */
static bool run_iteration(const char *base, int iter, size_t file_mib)
{
	struct race_state st;
	int start_pipe[2];
	int ready_pipe[2];
	pthread_t ager;
	pthread_t kmsg;
	bool got_window;
	bool confirmed = false;

	memset(&st, 0, sizeof(st));
	st.iter = iter;
	snprintf(st.leaf, sizeof(st.leaf), "%s/leaf_%03d", base, iter);
	snprintf(st.leaf_rel, sizeof(st.leaf_rel), "%s/leaf_%03d",
		 base + strlen(CGROUP_ROOT), iter);
	snprintf(st.file_path, sizeof(st.file_path), "/root/lru_race_%d_%03d.dat",
		 getpid(), iter);
	atomic_store(&st.phase, PHASE_IDLE);

	if (mkdir_if_missing(st.leaf))
		return false;
	if (pipe(start_pipe) || pipe(ready_pipe))
		die("pipe");

	st.victim = fork();
	if (!st.victim) {
		close(start_pipe[1]);
		close(ready_pipe[0]);
		victim_main(start_pipe[0], ready_pipe[1], st.file_path,
			    file_mib << 20);
		_exit(0);
	}

	close(start_pipe[0]);
	close(ready_pipe[1]);

	if (move_pid(st.leaf, st.victim)) {
		ssize_t start_ret = write(start_pipe[1], "g", 1);

		(void)start_ret;
		close(start_pipe[1]);
		kill(st.victim, SIGKILL);
		waitpid(st.victim, NULL, 0);
		rmdir_retry(st.leaf);
		return false;
	}

	ssize_t start_ret = write(start_pipe[1], "g", 1);

	(void)start_ret;
	close(start_pipe[1]);

	char ready;
	ssize_t ready_ret = read(ready_pipe[0], &ready, 1);

	(void)ready_ret;
	close(ready_pipe[0]);

	for (int retry = 0; retry < 400; retry++) {
		struct lru_info info;

		if (!read_lru_info(st.leaf_rel, &info) && info.memcg_id &&
		    info.nr_nodes > 0)
			break;
		usleep(5000);
	}

	pthread_create(&kmsg, NULL, kmsg_thread, &st);
	pthread_create(&ager, NULL, ager_thread, &st);

	got_window = wait_phase(&st, PHASE_WINDOW, WINDOW_TIMEOUT_MS);
	if (got_window) {
		kill(st.victim, SIGKILL);
		waitpid(st.victim, NULL, 0);
		st.victim = 0;
		rmdir_retry(st.leaf);
		confirmed = wait_phase(&st, PHASE_DONE, CONFIRM_TIMEOUT_MS);
	}

	atomic_store(&st.phase, PHASE_DONE);
	pthread_join(ager, NULL);
	pthread_join(kmsg, NULL);

	if (st.victim) {
		kill(st.victim, SIGKILL);
		waitpid(st.victim, NULL, 0);
	}
	rmdir_retry(st.leaf);
	unlink(st.file_path);

	printf("iter %d: %s\n", iter, confirmed ? "confirmed" :
	       got_window ? "window-only" : "miss");
	return confirmed;
}

/* Main prepares cgroup/debugfs state and runs bounded race attempts. */
int main(int argc, char **argv)
{
	int iters = argc > 1 ? atoi(argv[1]) : DEFAULT_ITERS;
	size_t file_mib = argc > 2 ? strtoul(argv[2], NULL, 0) : DEFAULT_FILE_MIB;
	char base[512];
	int confirmed = 0;

	if (geteuid()) {
		fprintf(stderr, "must run as root\n");
		return 1;
	}

	if (mount("debugfs", "/sys/kernel/debug", "debugfs", 0, NULL) && errno != EBUSY)
		perror("mount debugfs");

	enable_memory_controller(CGROUP_ROOT);
	snprintf(base, sizeof(base), CGROUP_ROOT "/lru_gen_race_%d", getpid());
	if (mkdir_if_missing(base))
		die("mkdir base cgroup");
	enable_memory_controller(base);

	for (int i = 0; i < iters; i++) {
		if (run_iteration(base, i, file_mib))
			confirmed++;
	}

	rmdir_retry(base);
	printf("confirmed=%d/%d\n", confirmed, iters);
	return confirmed ? 0 : 1;
}