/*
 * pfilrace_feeder.c -- set up IPv4 multicast routing with a rate-limited
 * vif and flood multicast so ip_mroute's tbf_reprocess_q callout
 * (a SOFTCLOCK-thread callback) continuously drains the tbf queue via
 * ip_output() -> pfil_run_hooks(PFIL_OUT).
 *
 * usage: pfilrace_feeder <vif1-lcl-addr>   (e.g. 10.0.2.15)
 * run as root.
 */
#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 <errno.h>

#ifndef MRT_INIT
#define MRT_INIT	100
#define MRT_ADD_VIF	102
#define MRT_ADD_MFC	104
#endif

struct vifctl {
	u_short	vifc_vifi;
	u_char	vifc_flags;
	u_char	vifc_threshold;
	u_int	vifc_rate_limit;
	struct in_addr vifc_lcl_addr;
	struct in_addr vifc_rmt_addr;
};

struct mfcctl {
	struct in_addr mfcc_origin;
	struct in_addr mfcc_mcastgrp;
	u_short mfcc_parent;
	u_char mfcc_ttls[32];
};

#define GROUP	"239.200.200.200"
#define PORT	4711

int
main(int argc, char **argv)
{
	struct sockaddr_in dst;
	struct vifctl vc;
	struct mfcctl mc;
	int rate = 4000;
	struct ip_mreq imr;
	char buf[256];
	int s, u, i, rc, n = 0;

	if (argc < 2) {
		fprintf(stderr, "usage: %s <vif1 lcl addr> [rate]\n", argv[0]);
		exit(1);
	}
	if (argc > 2)
		rate = atoi(argv[2]);

	s = socket(AF_INET, SOCK_RAW, IPPROTO_IGMP);
	if (s < 0) { perror("raw igmp socket"); exit(1); }

	i = 1;
	rc = setsockopt(s, IPPROTO_IP, MRT_INIT, &i, sizeof(i));
	if (rc < 0) { perror("MRT_INIT"); exit(1); }

	/* vif0: phyint on lo0 (incoming) */
	memset(&vc, 0, sizeof(vc));
	vc.vifc_vifi = 0;
	vc.vifc_flags = 0;
	vc.vifc_threshold = 1;
	vc.vifc_rate_limit = 0;
	vc.vifc_lcl_addr.s_addr = inet_addr("127.0.0.1");
	rc = setsockopt(s, IPPROTO_IP, MRT_ADD_VIF, &vc, sizeof(vc));
	if (rc < 0) { perror("MRT_ADD_VIF vif0"); exit(1); }

	/* vif1: phyint on the real interface, RATE LIMITED (outgoing) */
	vc.vifc_vifi = 1;
	vc.vifc_rate_limit = rate;	/* token refill, argv[2] */
	vc.vifc_lcl_addr.s_addr = inet_addr(argv[1]);
	rc = setsockopt(s, IPPROTO_IP, MRT_ADD_VIF, &vc, sizeof(vc));
	if (rc < 0) { perror("MRT_ADD_VIF vif1"); exit(1); }

	/* MFC: 127.0.0.1 -> GROUP, incoming vif0, forward on vif1 */
	memset(&mc, 0, sizeof(mc));
	mc.mfcc_origin.s_addr = inet_addr("127.0.0.1");
	mc.mfcc_mcastgrp.s_addr = inet_addr(GROUP);
	mc.mfcc_parent = 0;
	mc.mfcc_ttls[1] = 1;
	rc = setsockopt(s, IPPROTO_IP, MRT_ADD_MFC, &mc, sizeof(mc));
	if (rc < 0) { perror("MRT_ADD_MFC"); exit(1); }
	printf("feeder: mroute up (vif0 lo0, vif1 %s rate 4000, mfc set)\n",
	    argv[1]);

#define NPORT 6
	int uu[NPORT], k;
	char portbuf[16];
	struct sockaddr_in bnd;
	for (k = 0; k < NPORT; k++) {
		uu[k] = socket(AF_INET, SOCK_DGRAM, 0);
		if (uu[k] < 0) { perror("udp socket"); exit(1); }
		memset(&bnd, 0, sizeof(bnd));
		bnd.sin_family = AF_INET;
		bnd.sin_port = htons(41000 + k);
		bnd.sin_addr.s_addr = inet_addr("127.0.0.1");
		if (bind(uu[k], (struct sockaddr *)&bnd, sizeof(bnd)) < 0) {
			perror("bind"); exit(1);
		}
		{
			in_addr_t lo = inet_addr("127.0.0.1");
			struct ip_mreq m2;
			setsockopt(uu[k], IPPROTO_IP, IP_MULTICAST_IF, &lo,
			    sizeof(lo));
			/* join on every socket so each flow loops back
			 * through ip_input and is hashed by its 5-tuple */
			m2.imr_multiaddr.s_addr = inet_addr(GROUP);
			m2.imr_interface.s_addr = inet_addr("127.0.0.1");
			setsockopt(uu[k], IPPROTO_IP, IP_ADD_MEMBERSHIP,
			    &m2, sizeof(m2));
		}
	}
	u = uu[0];

	/* join the group on lo0 so our own sends loop back into ip_input */
	imr.imr_multiaddr.s_addr = inet_addr(GROUP);
	imr.imr_interface.s_addr = inet_addr("127.0.0.1");
	rc = setsockopt(u, IPPROTO_IP, IP_ADD_MEMBERSHIP, &imr, sizeof(imr));
	if (rc < 0) { perror("IP_ADD_MEMBERSHIP"); exit(1); }

	i = 64;
	setsockopt(u, IPPROTO_IP, IP_MULTICAST_TTL, &i, sizeof(i));
	/* loop defaults to enabled; force it on */
	i = 1;
	setsockopt(u, IPPROTO_IP, IP_MULTICAST_LOOP, &i, sizeof(i));
	imr.imr_interface.s_addr = inet_addr("127.0.0.1");
	setsockopt(u, IPPROTO_IP, IP_MULTICAST_IF, &imr.imr_interface,
	    sizeof(imr.imr_interface));

	memset(&dst, 0, sizeof(dst));
	dst.sin_family = AF_INET;
	dst.sin_port = htons(PORT);
	dst.sin_addr.s_addr = inet_addr(GROUP);

	memset(buf, 0x5a, sizeof(buf));
	printf("feeder: flooding UDP %s:%d via lo0 (mcast loop)\n", GROUP, PORT);
	fflush(stdout);

	/* argv[3]: if set, send only from port 41000+argv[3] (flow pinning) */
	if (argc > 3)
		u = uu[atoi(argv[3])];

	for (;;) {
		if (argc <= 3) {
			static int rr;
			u = uu[rr];
			rr = (rr + 1) % NPORT;
		}
		rc = sendto(u, buf, sizeof(buf), 0,
		    (struct sockaddr *)&dst, sizeof(dst));
		if (rc < 0) {
			if (errno == ENOBUFS || errno == EINTR)
				usleep(2000);
			else {
				perror("sendto");
				usleep(10000);
			}
		}
		if (++n >= 1) {	/* throttle: ~2kpps keeps netisr/mbufs sane */
			n = 0;
			usleep(300);
		}
	}
	/* NOTREACHED */
	return (0);
}