1 files changed, 941 insertions, 0 deletions

diff --git a/xdp-trafficgen/xdp-trafficgen.c b/xdp-trafficgen/xdp-trafficgen.c
new file mode 100644
index 0000000..293e416
--- /dev/null
+++ b/xdp-trafficgen/xdp-trafficgen.c
@@ -0,0 +1,941 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+
+#define _GNU_SOURCE
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <signal.h>
+#include <errno.h>
+#include <stdbool.h>
+#include <unistd.h>
+#include <fcntl.h>
+#include <pthread.h>
+#include <sched.h>
+#include <net/if.h>
+#include <sys/ioctl.h>
+#include <sys/types.h>
+#include <sys/socket.h>
+#include <netdb.h>
+
+#include <bpf/bpf.h>
+#include <bpf/bpf_endian.h>
+#include <bpf/libbpf.h>
+#include <xdp/libxdp.h>
+#include <xdp/xdp_stats_kern_user.h>
+#include <linux/bpf.h>
+#include <linux/err.h>
+#include <linux/if_link.h>
+#include <linux/if_ether.h>
+#include <linux/if_packet.h>
+#include <linux/ipv6.h>
+#include <linux/in6.h>
+#include <linux/udp.h>
+#include <linux/tcp.h>
+
+#include "params.h"
+#include "logging.h"
+#include "util.h"
+#include "xdp_sample.h"
+#include "xdp-trafficgen.h"
+
+#include "xdp_trafficgen.skel.h"
+
+#define PROG_NAME "xdp-trafficgen"
+
+#ifndef BPF_F_TEST_XDP_LIVE_FRAMES
+#define BPF_F_TEST_XDP_LIVE_FRAMES	(1U << 1)
+#endif
+
+#define IFINDEX_LO 1
+
+static int mask = SAMPLE_DEVMAP_XMIT_CNT_MULTI | SAMPLE_DROP_OK;
+
+DEFINE_SAMPLE_INIT(xdp_trafficgen);
+
+static bool status_exited = false;
+static bool runners_exited = false;
+
+struct udp_packet {
+	struct ethhdr eth;
+	struct ipv6hdr iph;
+	struct udphdr udp;
+	__u8 payload[64 - sizeof(struct udphdr)
+		     - sizeof(struct ethhdr) - sizeof(struct ipv6hdr)];
+} __attribute__((__packed__));
+
+static struct udp_packet pkt_udp = {
+	.eth.h_proto = __bpf_constant_htons(ETH_P_IPV6),
+	.iph.version = 6,
+	.iph.nexthdr = IPPROTO_UDP,
+	.iph.payload_len = bpf_htons(sizeof(struct udp_packet)
+				     - offsetof(struct udp_packet, udp)),
+	.iph.hop_limit = 1,
+	.iph.saddr.s6_addr16 = {bpf_htons(0xfe80), 0, 0, 0, 0, 0, 0, bpf_htons(1)},
+	.iph.daddr.s6_addr16 = {bpf_htons(0xfe80), 0, 0, 0, 0, 0, 0, bpf_htons(2)},
+	.udp.source = bpf_htons(1),
+	.udp.dest = bpf_htons(1),
+	.udp.len = bpf_htons(sizeof(struct udp_packet)
+			     - offsetof(struct udp_packet, udp)),
+};
+
+struct thread_config {
+	void *pkt;
+	size_t pkt_size;
+	__u32 cpu_core_id;
+	__u32 num_pkts;
+	__u32 batch_size;
+	struct xdp_program *prog;
+};
+
+static int run_prog(const struct thread_config *cfg, bool *status_var)
+{
+#ifdef HAVE_LIBBPF_BPF_PROG_TEST_RUN_OPTS
+	struct xdp_md ctx_in = {
+		.data_end = cfg->pkt_size,
+	};
+	DECLARE_LIBBPF_OPTS(bpf_test_run_opts, opts,
+			    .data_in = cfg->pkt,
+			    .data_size_in = cfg->pkt_size,
+			    .ctx_in = &ctx_in,
+			    .ctx_size_in = sizeof(ctx_in),
+			    .repeat = cfg->num_pkts ?: 1 << 20,
+			    .flags = BPF_F_TEST_XDP_LIVE_FRAMES,
+			    .batch_size = cfg->batch_size,
+		);
+	__u64 iterations = 0;
+	cpu_set_t cpu_cores;
+	int err;
+
+	CPU_ZERO(&cpu_cores);
+	CPU_SET(cfg->cpu_core_id, &cpu_cores);
+	pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpu_cores);
+	do {
+		err = xdp_program__test_run(cfg->prog, &opts, 0);
+		if (err)
+			return -errno;
+		iterations += opts.repeat;
+	} while (!*status_var && (!cfg->num_pkts || cfg->num_pkts > iterations));
+
+	return 0;
+#else
+	__unused const void *c = cfg, *s = status_var;
+	return -EOPNOTSUPP;
+#endif
+}
+
+static void *run_traffic(void *arg)
+{
+	const struct thread_config *cfg = arg;
+	int err;
+
+	err = run_prog(cfg, &status_exited);
+	if (err)
+		pr_warn("Couldn't run trafficgen program: %s\n", strerror(-err));
+
+	runners_exited = true;
+	return NULL;
+}
+
+static int probe_kernel_support(void)
+{
+	DECLARE_LIBXDP_OPTS(xdp_program_opts, opts);
+	struct xdp_trafficgen *skel;
+	struct xdp_program *prog;
+	int data = 0, err;
+	bool status = 0;
+
+	skel = xdp_trafficgen__open();
+	if (!skel) {
+		err = -errno;
+		pr_warn("Couldn't open XDP program: %s\n", strerror(-err));
+		return err;
+	}
+
+	err = sample_init_pre_load(skel, "lo");
+	if (err < 0) {
+		pr_warn("Failed to sample_init_pre_load: %s\n", strerror(-err));
+		goto out;
+	}
+
+	opts.obj = skel->obj;
+	opts.prog_name = "xdp_drop";
+
+	prog = xdp_program__create(&opts);
+	if (!prog) {
+		err = -errno;
+		pr_warn("Couldn't load XDP program: %s\n", strerror(-err));
+		goto out;
+	}
+
+	const struct thread_config cfg = {
+		.pkt = &data,
+		.pkt_size = sizeof(data),
+		.num_pkts = 1,
+		.batch_size = 1,
+		.prog = prog
+	};
+	err = run_prog(&cfg, &status);
+	if (err == -EOPNOTSUPP) {
+		pr_warn("BPF_PROG_RUN with batch size support is missing from libbpf.\n");
+	}  else if (err == -EINVAL) {
+		err = -EOPNOTSUPP;
+		pr_warn("Kernel doesn't support live packet mode for XDP BPF_PROG_RUN.\n");
+	} else if (err) {
+		pr_warn("Error probing kernel support: %s\n", strerror(-err));
+	}
+
+	xdp_program__close(prog);
+out:
+	xdp_trafficgen__destroy(skel);
+	return err;
+}
+
+static int create_runners(pthread_t **runner_threads, struct thread_config **thread_configs, int num_threads,
+			  struct thread_config *tcfg, struct xdp_program *prog)
+{
+	struct thread_config *t;
+	pthread_t *threads;
+	int i, err;
+
+	threads = calloc(sizeof(pthread_t), num_threads);
+	if (!threads) {
+		pr_warn("Couldn't allocate memory\n");
+		return -ENOMEM;
+	}
+
+	t = calloc(sizeof(struct thread_config), num_threads);
+	if (!t) {
+		pr_warn("Couldn't allocate memory\n");
+		free(threads);
+		return -ENOMEM;
+	}
+
+	for (i = 0; i < num_threads; i++) {
+		memcpy(&t[i], tcfg, sizeof(*tcfg));
+		tcfg->cpu_core_id++;
+
+		t[i].prog = xdp_program__clone(prog, 0);
+		err = libxdp_get_error(t[i].prog);
+		if (err) {
+			pr_warn("Failed to clone xdp_program: %s\n", strerror(-err));
+			t[i].prog = NULL;
+			goto err;
+		}
+
+		err = pthread_create(&threads[i], NULL, run_traffic, &t[i]);
+		if (err < 0) {
+			pr_warn("Failed to create traffic thread: %s\n", strerror(-err));
+			goto err;
+		}
+	}
+
+	*runner_threads = threads;
+	*thread_configs = t;
+
+	return 0;
+
+err:
+	for (i = 0; i < num_threads; i++) {
+		pthread_cancel(threads[i]);
+		xdp_program__close(t[i].prog);
+	}
+	free(t);
+	free(threads);
+
+	return err;
+}
+
+
+static __be16 calc_udp_cksum(const struct udp_packet *pkt)
+{
+	__u32 chksum = pkt->iph.nexthdr + bpf_ntohs(pkt->iph.payload_len);
+	int i;
+
+	for (i = 0; i < 8; i++) {
+		chksum += bpf_ntohs(pkt->iph.saddr.s6_addr16[i]);
+		chksum += bpf_ntohs(pkt->iph.daddr.s6_addr16[i]);
+	}
+	chksum += bpf_ntohs(pkt->udp.source);
+	chksum += bpf_ntohs(pkt->udp.dest);
+	chksum += bpf_ntohs(pkt->udp.len);
+
+	while (chksum >> 16)
+		chksum = (chksum & 0xFFFF) + (chksum >> 16);
+	return bpf_htons(~chksum);
+}
+
+static const struct udpopt {
+	__u32 num_pkts;
+	struct iface iface;
+	struct mac_addr dst_mac;
+	struct mac_addr src_mac;
+	struct ip_addr dst_ip;
+	struct ip_addr src_ip;
+	__u16 dst_port;
+	__u16 src_port;
+	__u16 dyn_ports;
+	__u16 threads;
+	__u16 interval;
+} defaults_udp = {
+	.interval = 1,
+	.threads = 1,
+};
+
+static int prepare_udp_pkt(const struct udpopt *cfg)
+{
+	struct mac_addr src_mac = cfg->src_mac;
+	int err;
+
+	if (macaddr_is_null(&src_mac)) {
+		err = get_mac_addr(cfg->iface.ifindex, &src_mac);
+		if (err)
+			return err;
+	}
+	memcpy(pkt_udp.eth.h_source, &src_mac, sizeof(src_mac));
+	if (!macaddr_is_null(&cfg->dst_mac))
+		memcpy(pkt_udp.eth.h_dest, &cfg->dst_mac, sizeof(cfg->dst_mac));
+
+	if (!ipaddr_is_null(&cfg->src_ip)) {
+		if (cfg->src_ip.af != AF_INET6) {
+			pr_warn("Only IPv6 is supported\n");
+			return 1;
+		}
+		pkt_udp.iph.saddr = cfg->src_ip.addr.addr6;
+	}
+
+	if (!ipaddr_is_null(&cfg->dst_ip)) {
+		if (cfg->dst_ip.af != AF_INET6) {
+			pr_warn("Only IPv6 is supported\n");
+			return 1;
+		}
+		pkt_udp.iph.daddr = cfg->dst_ip.addr.addr6;
+	}
+
+	if (cfg->src_port)
+		pkt_udp.udp.source = bpf_htons(cfg->src_port);
+	if (cfg->dst_port)
+		pkt_udp.udp.dest = bpf_htons(cfg->dst_port);
+	pkt_udp.udp.check = calc_udp_cksum(&pkt_udp);
+	return 0;
+}
+
+static struct prog_option udp_options[] = {
+	DEFINE_OPTION("dst-mac", OPT_MACADDR, struct udpopt, dst_mac,
+		      .short_opt = 'm',
+		      .metavar = "<mac addr>",
+		      .help = "Destination MAC address of generated packets"),
+	DEFINE_OPTION("src-mac", OPT_MACADDR, struct udpopt, src_mac,
+		      .short_opt = 'M',
+		      .metavar = "<mac addr>",
+		      .help = "Source MAC address of generated packets"),
+	DEFINE_OPTION("dst-addr", OPT_IPADDR, struct udpopt, dst_ip,
+		      .short_opt = 'a',
+		      .metavar = "<addr>",
+		      .help = "Destination IP address of generated packets"),
+	DEFINE_OPTION("src-addr", OPT_IPADDR, struct udpopt, src_ip,
+		      .short_opt = 'A',
+		      .metavar = "<addr>",
+		      .help = "Source IP address of generated packets"),
+	DEFINE_OPTION("dst-port", OPT_U16, struct udpopt, dst_port,
+		      .short_opt = 'p',
+		      .metavar = "<port>",
+		      .help = "Destination port of generated packets"),
+	DEFINE_OPTION("src-port", OPT_U16, struct udpopt, src_port,
+		      .short_opt = 'P',
+		      .metavar = "<port>",
+		      .help = "Source port of generated packets"),
+	DEFINE_OPTION("dyn-ports", OPT_U16, struct udpopt, dyn_ports,
+		      .short_opt = 'd',
+		      .metavar = "<num ports>",
+		      .help = "Dynamically vary destination port over a range of <num ports>"),
+	DEFINE_OPTION("num-packets", OPT_U32, struct udpopt, num_pkts,
+		      .short_opt = 'n',
+		      .metavar = "<port>",
+		      .help = "Number of packets to send"),
+	DEFINE_OPTION("threads", OPT_U16, struct udpopt, threads,
+		      .short_opt = 't',
+		      .metavar = "<threads>",
+		      .help = "Number of simultaneous threads to transmit from"),
+	DEFINE_OPTION("interval", OPT_U16, struct udpopt, interval,
+		      .short_opt = 'I',
+		      .metavar = "<s>",
+		      .help = "Output statistics with this interval"),
+	DEFINE_OPTION("interface", OPT_IFNAME, struct udpopt, iface,
+		      .positional = true,
+		      .metavar = "<ifname>",
+		      .required = true,
+		      .help = "Load on device <ifname>"),
+	END_OPTIONS
+};
+
+int do_udp(const void *opt, __unused const char *pin_root_path)
+{
+	const struct udpopt *cfg = opt;
+
+	DECLARE_LIBXDP_OPTS(xdp_program_opts, opts);
+	struct thread_config *t = NULL, tcfg = {
+		.pkt = &pkt_udp,
+		.pkt_size = sizeof(pkt_udp),
+		.num_pkts = cfg->num_pkts,
+	};
+	struct trafficgen_state bpf_state = {};
+	struct xdp_trafficgen *skel = NULL;
+	pthread_t *runner_threads = NULL;
+	struct xdp_program *prog = NULL;
+	int err = 0, i;
+	char buf[100];
+	__u32 key = 0;
+
+	err = probe_kernel_support();
+	if (err)
+		return err;
+
+	err = prepare_udp_pkt(cfg);
+	if (err)
+		goto out;
+
+	skel = xdp_trafficgen__open();
+	if (!skel) {
+		err = -errno;
+		pr_warn("Couldn't open XDP program: %s\n", strerror(-err));
+		goto out;
+	}
+
+	err = sample_init_pre_load(skel, cfg->iface.ifname);
+	if (err < 0) {
+		pr_warn("Failed to sample_init_pre_load: %s\n", strerror(-err));
+		goto out;
+	}
+
+	skel->rodata->config.port_start = cfg->dst_port;
+	skel->rodata->config.port_range = cfg->dyn_ports;
+	skel->rodata->config.ifindex_out = cfg->iface.ifindex;
+	bpf_state.next_port = cfg->dst_port;
+
+	if (cfg->dyn_ports)
+		opts.prog_name = "xdp_redirect_update_port";
+	else
+		opts.prog_name = "xdp_redirect_notouch";
+	opts.obj = skel->obj;
+
+	prog = xdp_program__create(&opts);
+	if (!prog) {
+		err = -errno;
+		libxdp_strerror(err, buf, sizeof(buf));
+		pr_warn("Couldn't open BPF file: %s\n", buf);
+		goto out;
+	}
+
+	err = xdp_trafficgen__load(skel);
+	if (err)
+		goto out;
+
+	err = bpf_map_update_elem(bpf_map__fd(skel->maps.state_map),
+				  &key, &bpf_state, BPF_EXIST);
+	if (err) {
+		err = -errno;
+		pr_warn("Couldn't set initial state map value: %s\n", strerror(-err));
+		goto out;
+	}
+
+	err = sample_init(skel, mask, IFINDEX_LO, cfg->iface.ifindex);
+	if (err < 0) {
+		pr_warn("Failed to initialize sample: %s\n", strerror(-err));
+		goto out;
+	}
+
+	err = create_runners(&runner_threads, &t, cfg->threads, &tcfg, prog);
+	if (err)
+		goto out;
+
+	pr_info("Transmitting on %s (ifindex %d)\n",
+	       cfg->iface.ifname, cfg->iface.ifindex);
+
+	err = sample_run(cfg->interval, NULL, NULL);
+	status_exited = true;
+
+	for (i = 0; i < cfg->threads; i++) {
+		pthread_join(runner_threads[i], NULL);
+		xdp_program__close(t[i].prog);
+	}
+
+out:
+	xdp_program__close(prog);
+	xdp_trafficgen__destroy(skel);
+	free(runner_threads);
+	free(t);
+        return err;
+}
+
+struct tcp_packet {
+	struct ethhdr eth;
+	struct ipv6hdr iph;
+	struct tcphdr tcp;
+	__u8 payload[1500 - sizeof(struct tcphdr)
+		     - sizeof(struct ethhdr) - sizeof(struct ipv6hdr)];
+} __attribute__((__packed__));
+
+static __unused struct tcp_packet pkt_tcp = {
+	.eth.h_proto = __bpf_constant_htons(ETH_P_IPV6),
+	.iph.version = 6,
+	.iph.nexthdr = IPPROTO_TCP,
+	.iph.payload_len = bpf_htons(sizeof(struct tcp_packet)
+				     - offsetof(struct tcp_packet, tcp)),
+	.iph.hop_limit = 64,
+	.iph.saddr.s6_addr16 = {bpf_htons(0xfe80), 0, 0, 0, 0, 0, 0, bpf_htons(1)},
+	.iph.daddr.s6_addr16 = {bpf_htons(0xfe80), 0, 0, 0, 0, 0, 0, bpf_htons(2)},
+	.tcp.source = bpf_htons(1),
+	.tcp.dest = bpf_htons(1),
+	.tcp.window = bpf_htons(0x100),
+	.tcp.doff = 5,
+	.tcp.ack = 1,
+};
+
+static void hexdump_data(void *data, int size)
+{
+	unsigned char *ptr = data;
+	int i;
+	for (i = 0; i < size; i++) {
+		if (i % 16 == 0)
+			pr_debug("\n%06X: ", i);
+		else if (i % 2 == 0)
+			pr_debug(" ");
+		pr_debug("%02X", *ptr++);
+	}
+	pr_debug("\n");
+}
+
+static __be16 calc_tcp_cksum(const struct tcp_packet *pkt)
+{
+	__u32 chksum = bpf_htons(pkt->iph.nexthdr) + pkt->iph.payload_len;
+	int payload_len = sizeof(pkt->payload);
+	struct tcphdr tcph_ = pkt->tcp;
+	__u16 *ptr = (void *)&tcph_;
+	int i;
+
+	tcph_.check = 0;
+
+	for (i = 0; i < 8; i++) {
+		chksum += pkt->iph.saddr.s6_addr16[i];
+		chksum += pkt->iph.daddr.s6_addr16[i];
+	}
+	for (i = 0; i < 10; i++)
+		chksum += *(ptr++);
+
+	ptr = (void *)&pkt->payload;
+	for (i = 0; i < payload_len / 2; i++)
+		chksum += *(ptr++);
+
+	if (payload_len % 2)
+		chksum += (*((__u8 *)ptr)) << 8;
+
+	while (chksum >> 16)
+		chksum = (chksum & 0xFFFF) + (chksum >> 16);
+
+	return ~chksum;
+}
+
+static void prepare_tcp_pkt(const struct tcp_flowkey *fkey,
+			    const struct tcp_flowstate *fstate)
+{
+	memcpy(pkt_tcp.eth.h_source, fstate->src_mac, ETH_ALEN);
+	memcpy(pkt_tcp.eth.h_dest, fstate->dst_mac, ETH_ALEN);
+
+	pkt_tcp.iph.saddr = fkey->src_ip;
+	pkt_tcp.iph.daddr = fkey->dst_ip;
+	pkt_tcp.tcp.source = fkey->src_port;
+	pkt_tcp.tcp.dest = fkey->dst_port;
+	pkt_tcp.tcp.seq = bpf_htonl(fstate->seq);
+	pkt_tcp.tcp.ack_seq = bpf_htonl(fstate->rcv_seq);
+
+	pkt_tcp.tcp.check = calc_tcp_cksum(&pkt_tcp);
+	pr_debug("TCP packet:\n");
+	hexdump_data(&pkt_tcp, sizeof(pkt_tcp));
+}
+
+struct enum_val xdp_modes[] = {
+       {"native", XDP_MODE_NATIVE},
+       {"skb", XDP_MODE_SKB},
+       {"hw", XDP_MODE_HW},
+       {NULL, 0}
+};
+
+static const struct tcpopt {
+	__u32 num_pkts;
+	struct iface iface;
+	char *dst_addr;
+	__u16 dst_port;
+	__u16 interval;
+	__u16 timeout;
+	enum xdp_attach_mode mode;
+} defaults_tcp = {
+	.interval = 1,
+	.dst_port = 10000,
+	.timeout = 2,
+	.mode = XDP_MODE_NATIVE,
+};
+
+static struct prog_option tcp_options[] = {
+	DEFINE_OPTION("dst-port", OPT_U16, struct tcpopt, dst_port,
+		      .short_opt = 'p',
+		      .metavar = "<port>",
+		      .help = "Connect to destination <port>. Default 10000"),
+	DEFINE_OPTION("num-packets", OPT_U32, struct tcpopt, num_pkts,
+		      .short_opt = 'n',
+		      .metavar = "<port>",
+		      .help = "Number of packets to send"),
+	DEFINE_OPTION("interval", OPT_U16, struct tcpopt, interval,
+		      .short_opt = 'I',
+		      .metavar = "<s>",
+		      .help = "Output statistics with this interval"),
+	DEFINE_OPTION("timeout", OPT_U16, struct tcpopt, timeout,
+		      .short_opt = 't',
+		      .metavar = "<s>",
+		      .help = "TCP connect timeout (default 2 seconds)."),
+	DEFINE_OPTION("interface", OPT_IFNAME, struct tcpopt, iface,
+		      .metavar = "<ifname>",
+		      .required = true,
+		      .short_opt = 'i',
+		      .help = "Connect through device <ifname>"),
+	DEFINE_OPTION("mode", OPT_ENUM, struct tcpopt, mode,
+		      .short_opt = 'm',
+		      .typearg = xdp_modes,
+		      .metavar = "<mode>",
+		      .help = "Load ingress XDP program in <mode>; default native"),
+	DEFINE_OPTION("dst-addr", OPT_STRING, struct tcpopt, dst_addr,
+		      .positional = true,
+		      .required = true,
+		      .metavar = "<hostname>",
+		      .help = "Destination host of generated stream"),
+	END_OPTIONS
+};
+
+int do_tcp(const void *opt, __unused const char *pin_root_path)
+{
+	const struct tcpopt *cfg = opt;
+
+	struct addrinfo *ai = NULL, hints = {
+		.ai_family = AF_INET6,
+		.ai_socktype = SOCK_STREAM,
+		.ai_protocol = IPPROTO_TCP,
+	};
+	struct ip_addr local_addr = { .af = AF_INET6 }, remote_addr = { .af = AF_INET6 };
+	struct bpf_map *state_map = NULL, *fstate_map;
+	DECLARE_LIBXDP_OPTS(xdp_program_opts, opts,
+			    .prog_name = "xdp_handle_tcp_recv");
+	struct xdp_program *ifindex_prog = NULL, *test_prog = NULL;
+	struct sockaddr_in6 local_saddr = {}, *addr6;
+	struct thread_config *t = NULL, tcfg = {
+		.pkt = &pkt_tcp,
+		.pkt_size = sizeof(pkt_tcp),
+		.num_pkts = cfg->num_pkts,
+	};
+	struct trafficgen_state bpf_state = {};
+	struct xdp_trafficgen *skel = NULL;
+	char buf_local[50], buf_remote[50];
+	pthread_t *runner_threads = NULL;
+	socklen_t sockaddr_sz, tcpi_sz;
+	__u16 local_port, remote_port;
+	int sock = -1, err = -EINVAL;
+	struct tcp_flowstate fstate;
+	struct timeval timeout = {
+		.tv_sec = cfg->timeout,
+	};
+	struct tcp_info tcpi = {};
+	bool attached = false;
+	__u16 num_threads = 1;
+	__u32 key = 0;
+	char port[6];
+	int i, sopt;
+
+	err = probe_kernel_support();
+	if (err)
+		return err;
+
+	skel = xdp_trafficgen__open();
+	if (!skel) {
+		err = -errno;
+		pr_warn("Couldn't open XDP program: %s\n", strerror(-err));
+		goto out;
+	}
+
+	err = sample_init_pre_load(skel, cfg->iface.ifname);
+	if (err < 0) {
+		pr_warn("Failed to sample_init_pre_load: %s\n", strerror(-err));
+		goto out;
+	}
+
+	opts.obj = skel->obj;
+	skel->rodata->config.ifindex_out = cfg->iface.ifindex;
+
+	snprintf(port, sizeof(port), "%d", cfg->dst_port);
+
+	err = getaddrinfo(cfg->dst_addr, port, &hints, &ai);
+	if (err) {
+		pr_warn("Couldn't resolve hostname: %s\n", gai_strerror(err));
+		goto out;
+	}
+
+	addr6 = (struct sockaddr_in6* )ai->ai_addr;
+	remote_addr.addr.addr6 = addr6->sin6_addr;
+	remote_port = bpf_ntohs(addr6->sin6_port);
+
+	bpf_state.flow_key.dst_port = addr6->sin6_port;
+	bpf_state.flow_key.dst_ip = addr6->sin6_addr;
+
+	print_addr(buf_remote, sizeof(buf_remote), &remote_addr);
+
+	ifindex_prog = xdp_program__create(&opts);
+	if (!ifindex_prog) {
+		err = -errno;
+		pr_warn("Couldn't open XDP program: %s\n", strerror(-err));
+		goto out;
+	}
+
+	opts.prog_name = "xdp_redirect_send_tcp";
+	test_prog = xdp_program__create(&opts);
+	if (!test_prog) {
+		err = -errno;
+		pr_warn("Couldn't find test program: %s\n", strerror(-err));
+		goto out;
+	}
+
+	state_map = skel->maps.state_map;
+	fstate_map = skel->maps.flow_state_map;
+
+	if (!fstate_map) {
+		pr_warn("Couldn't find BPF maps\n");
+		goto out;
+	}
+
+	err = xdp_program__attach(ifindex_prog, cfg->iface.ifindex, cfg->mode, 0);
+	if (err) {
+		err = -errno;
+		pr_warn("Couldn't attach XDP program to iface '%s': %s\n",
+			cfg->iface.ifname, strerror(-err));
+		goto out;
+	}
+	attached = true;
+
+	err = bpf_map_update_elem(bpf_map__fd(state_map),
+				  &key, &bpf_state, BPF_EXIST);
+
+	if (err) {
+		err = -errno;
+		pr_warn("Couldn't set initial state map value: %s\n", strerror(-err));
+		goto out;
+	}
+
+	err = sample_init(skel, mask, IFINDEX_LO, cfg->iface.ifindex);
+	if (err < 0) {
+		pr_warn("Failed to initialize sample: %s\n", strerror(-err));
+		goto out;
+	}
+
+	sock = socket(AF_INET6, SOCK_STREAM, IPPROTO_TCP);
+	if (sock < 0) {
+		err = -errno;
+		pr_warn("Couldn't open TCP socket: %s\n", strerror(-err));
+		goto out;
+	}
+
+	err = setsockopt(sock, SOL_SOCKET, SO_BINDTOIFINDEX,
+			 &cfg->iface.ifindex, sizeof(cfg->iface.ifindex));
+	if (err) {
+		err = -errno;
+		pr_warn("Couldn't bind to device '%s': %s\n", cfg->iface.ifname, strerror(-err));
+		goto out;
+	}
+
+	sopt = fcntl(sock, F_GETFL, NULL);
+	if (sopt < 0) {
+		err = -errno;
+		pr_warn("Couldn't get socket opts: %s\n", strerror(-err));
+		goto out;
+	}
+
+	err = fcntl(sock, F_SETFL, sopt | O_NONBLOCK);
+	if (err) {
+		err = -errno;
+		pr_warn("Couldn't set socket non-blocking: %s\n", strerror(-err));
+		goto out;
+	}
+
+	err = connect(sock, ai->ai_addr, ai->ai_addrlen);
+	if (err && errno == EINPROGRESS) {
+		fd_set wait;
+
+		FD_ZERO(&wait);
+		FD_SET(sock, &wait);
+
+		err = select(sock + 1, NULL, &wait, NULL, &timeout);
+		if (!err) {
+			err = -1;
+			errno = ETIMEDOUT;
+		} else if (err > 0) {
+			err = 0;
+		}
+	}
+	if (err) {
+		err = -errno;
+		pr_warn("Couldn't connect to destination: %s\n", strerror(-err));
+		goto out;
+	}
+
+	err = fcntl(sock, F_SETFL, sopt);
+	if (err) {
+		err = -errno;
+		pr_warn("Couldn't reset socket opts: %s\n", strerror(-err));
+		goto out;
+	}
+
+	sockaddr_sz = sizeof(local_saddr);
+	err = getsockname(sock, &local_saddr, &sockaddr_sz);
+	if (err) {
+		err = -errno;
+		pr_warn("Couldn't get local address: %s\n", strerror(-err));
+		goto out;
+	}
+
+	local_addr.addr.addr6 = local_saddr.sin6_addr;
+	local_port = bpf_htons(local_saddr.sin6_port);
+	print_addr(buf_local, sizeof(buf_local), &local_addr);
+
+	printf("Connected to %s port %d from %s port %d\n",
+	       buf_remote, remote_port, buf_local, local_port);
+
+	bpf_state.flow_key.src_port = local_saddr.sin6_port;
+	bpf_state.flow_key.src_ip = local_saddr.sin6_addr;
+
+	tcpi_sz = sizeof(tcpi);
+	err = getsockopt(sock, IPPROTO_TCP, TCP_INFO, &tcpi, &tcpi_sz);
+	if (err) {
+		err = -errno;
+		pr_warn("Couldn't get TCP_INFO for socket: %s\n", strerror(-err));
+		goto out;
+	}
+
+	err = bpf_map_lookup_elem(bpf_map__fd(fstate_map),
+				  &bpf_state.flow_key, &fstate);
+	if (err) {
+		err = -errno;
+		pr_warn("Couldn't find flow state in map: %s\n", strerror(-err));
+		goto out;
+	}
+
+	if (tcpi.tcpi_snd_wnd != fstate.window) {
+		pr_warn("TCP_INFO and packet data disagree on window (%u != %u)\n",
+			tcpi.tcpi_snd_wnd, fstate.window);
+	}
+
+	fstate.wscale = tcpi.tcpi_rcv_wscale;
+	fstate.flow_state = FLOW_STATE_RUNNING;
+	err = bpf_map_update_elem(bpf_map__fd(fstate_map),
+				  &bpf_state.flow_key, &fstate, BPF_EXIST);
+	if (err) {
+		err = -errno;
+		pr_warn("Couldn't update flow state map: %s\n", strerror(-err));
+		goto out;
+	}
+
+	err = bpf_map_update_elem(bpf_map__fd(state_map),
+				  &key, &bpf_state, BPF_EXIST);
+	if (err) {
+		err = -errno;
+		pr_warn("Couldn't update program state map: %s\n", strerror(-err));
+		goto out;
+	}
+
+	prepare_tcp_pkt(&bpf_state.flow_key, &fstate);
+
+	err = create_runners(&runner_threads, &t, num_threads, &tcfg, test_prog);
+	if (err)
+		goto out;
+
+	err = sample_run(cfg->interval, NULL, NULL);
+	status_exited = true;
+	for (i = 0; i < num_threads; i++) {
+		pthread_join(runner_threads[i], NULL);
+		xdp_program__close(t[i].prog);
+	}
+
+	/* send 3 RSTs with 200ms interval to kill the other side of the connection */
+	for (i = 0; i < 3; i++) {
+		usleep(200000);
+
+		pkt_tcp.tcp.rst = 1;
+		pkt_tcp.iph.payload_len = bpf_htons(sizeof(struct tcphdr));
+		pkt_tcp.tcp.check = calc_tcp_cksum(&pkt_tcp);
+		tcfg.cpu_core_id = 0;
+		tcfg.num_pkts = 1;
+		tcfg.pkt_size = offsetof(struct tcp_packet, payload);
+		tcfg.prog = test_prog;
+		run_traffic(&tcfg);
+	}
+
+out:
+	if (ai)
+		freeaddrinfo(ai);
+	if (sock >= 0)
+		close(sock);
+	if (attached)
+		xdp_program__detach(ifindex_prog, cfg->iface.ifindex, cfg->mode, 0);
+
+	xdp_program__close(ifindex_prog);
+	xdp_program__close(test_prog);
+
+	xdp_trafficgen__destroy(skel);
+
+	free(runner_threads);
+	free(t);
+	return err;
+}
+
+static const struct probeopt {
+} defaults_probe = {};
+
+static struct prog_option probe_options[] = {};
+
+int do_probe(__unused const void *cfg, __unused const char *pin_root_path)
+{
+	int err = probe_kernel_support();
+
+	if (!err)
+		pr_info("Kernel supports live packet mode for XDP BPF_PROG_RUN.\n");
+	return err;
+}
+
+int do_help(__unused const void *cfg, __unused const char *pin_root_path)
+{
+	fprintf(stderr,
+		"Usage: xdp-trafficgen COMMAND [options]\n"
+		"\n"
+		"COMMAND can be one of:\n"
+		"       udp         - run in UDP mode\n"
+		"       tcp         - run in TCP mode\n"
+		"       help        - show this help message\n"
+		"\n"
+		"Use 'xdp-trafficgen COMMAND --help' to see options for each command\n");
+	return -1;
+}
+
+static const struct prog_command cmds[] = {
+	DEFINE_COMMAND(udp, "Run in UDP mode"),
+	DEFINE_COMMAND(tcp, "Run in TCP mode"),
+	DEFINE_COMMAND(probe, "Probe kernel support"),
+	{ .name = "help", .func = do_help, .no_cfg = true },
+	END_COMMANDS
+};
+
+union all_opts {
+	struct udpopt udp;
+};
+
+int main(int argc, char **argv)
+{
+	if (argc > 1)
+		return dispatch_commands(argv[1], argc - 1, argv + 1, cmds,
+					 sizeof(union all_opts), PROG_NAME, false);
+
+	return do_help(NULL, NULL);
+}