/*
 * xdp-server.c -- integration of AF_XDP into nsd
 *
 * Copyright (c) 2024, NLnet Labs. All rights reserved.
 *
 * See LICENSE for the license.
 *
 */

/*
 * Parts inspired by https://github.com/xdp-project/xdp-tutorial
 */

#include "config.h"

#ifdef USE_XDP

#include <assert.h>
#include <errno.h>
#include <netinet/in.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <unistd.h>
#include <linux/limits.h>
#include <sys/mman.h>

#include <sys/poll.h>
#include <sys/resource.h>

/* #include <bpf/bpf.h> */
#include <xdp/xsk.h>
#include <xdp/libxdp.h>
#include <bpf/libbpf.h>

#include <arpa/inet.h>
#include <linux/icmpv6.h>
#include <linux/if_ether.h>
#include <linux/ipv6.h>
#include <linux/ip.h>
#include <linux/udp.h>
#include <net/if.h>

#include <arpa/inet.h>
#include <netdb.h>
#include <ifaddrs.h>
#include <linux/if_link.h>

#include "query.h"
#include "dns.h"
#include "util.h"
#include "xdp-server.h"
#include "xdp-util.h"
#include "nsd.h"

// TODO: make configurable
#define DNS_PORT 53

struct xdp_config {
	__u32 xdp_flags;
	__u32 libxdp_flags;
	__u16 xsk_bind_flags;
};

struct umem_ptr {
	uint64_t addr;
	uint32_t len;
};

static struct umem_ptr umem_ptrs[XDP_RX_BATCH_SIZE];

/*
 * Allocate memory for UMEM and setup rings
 */
static int
xsk_configure_umem(struct xsk_umem_info *umem_info, uint64_t size);

/*
 * Retrieve a UMEM frame address for allocation
 *
 * Returns XDP_INVALID_UMEM_FRAME when there are no free frames available.
 */
static uint64_t xsk_alloc_umem_frame(struct xsk_socket_info *xsk);

/*
 * Bind AF_XDP socket and setup rings
 */
static int xsk_configure_socket(struct xdp_server *xdp,
                                struct xsk_socket_info *xsk_info,
                                struct xsk_umem_info *umem,
                                uint32_t queue_index);

/*
 * Get number of free frames in UMEM
 */
static uint64_t xsk_umem_free_frames(struct xsk_socket_info *xsk);

/*
 * Free a frame in UMEM
 */
static void xsk_free_umem_frame(struct xsk_socket_info *xsk, uint64_t frame);

/*
 * Fill fill ring with as many frames as possible
 */
static void fill_fq(struct xsk_socket_info *xsk);

/*
 * Load eBPF program to forward traffic to our socket
 */
static int load_xdp_program_and_map(struct xdp_server *xdp);

/*
 * Unload eBPF/XDP program
 */
static void unload_xdp_program(struct xdp_server *xdp);

/*
 * Figure out IP addresses to listen to.
 */
static int figure_ip_addresses(struct xdp_server *xdp);

/*
 * Add IP address to allowed destination addresses for incoming packets
 */
static void add_ip_address(struct xdp_server *xdp,
                           struct sockaddr_storage *addr);

/*
 * Check whether destination IPv4 is in allowed IPs list
 */
static int dest_ip_allowed4(struct xdp_server *xdp, struct iphdr *ipv4);

/*
 * Check whether destination IPv6 is in allowed IPs list
 */
static int dest_ip_allowed6(struct xdp_server *xdp, struct ipv6hdr *ipv6);

/*
 * Setup XDP sockets
 */
static int xdp_sockets_init(struct xdp_server *xdp);

/*
 * Cleanup XDP sockets and memory
 */
static void xdp_sockets_cleanup(struct xdp_server *xdp);

/*
 * Allocate a block of shared memory
 */
static void *alloc_shared_mem(size_t len);

/*
 * Collect free frames from completion queue
 */
static void drain_cq(struct xsk_socket_info *xsk);

/*
 * Send outstanding packets and recollect completed frame addresses
 */
static void handle_tx(struct xsk_socket_info *xsk);

/*
 * Process packet and indicate if it should be dropped
 * return 0 or less => drop
 * return greater than 0 => use for tx
 */
static int
process_packet(struct xdp_server *xdp,
               uint8_t *pkt,
               uint32_t *len,
               struct query *query);

static inline void swap_eth(struct ethhdr *eth);
static inline void swap_udp(struct udphdr *udp);
static inline void swap_ipv6(struct ipv6hdr *ipv6);
static inline void swap_ipv4(struct iphdr *ipv4);
static inline void *parse_udp(struct udphdr *udp);
static inline void *parse_ipv6(struct ipv6hdr *ipv6);
static inline void *parse_ipv4(struct iphdr *ipv4);

/*
 * Parse dns message and return new length of dns message
 */
static uint32_t parse_dns(struct nsd* nsd,
                          uint32_t dnslen,
                          struct query *q,
                          sa_family_t ai_family);

/* *************** */
/* Implementations */
/* *************** */

static uint64_t xsk_alloc_umem_frame(struct xsk_socket_info *xsk) {
	uint64_t frame;
	if (xsk->umem->umem_frame_free == 0) {
		return XDP_INVALID_UMEM_FRAME;
	}

	frame = xsk->umem->umem_frame_addr[--xsk->umem->umem_frame_free];
	xsk->umem->umem_frame_addr[xsk->umem->umem_frame_free] =
		XDP_INVALID_UMEM_FRAME;
	return frame;
}

static uint64_t xsk_umem_free_frames(struct xsk_socket_info *xsk) {
	return xsk->umem->umem_frame_free;
}

static void xsk_free_umem_frame(struct xsk_socket_info *xsk, uint64_t frame) {
	assert(xsk->umem->umem_frame_free < XDP_NUM_FRAMES);
	xsk->umem->umem_frame_addr[xsk->umem->umem_frame_free++] = frame;
}

static void fill_fq(struct xsk_socket_info *xsk) {
	uint32_t stock_frames;
	uint32_t idx_fq = 0;
	/* fill the fill ring with as many frames as are available */
	/* get number of spots available in fq */
	stock_frames = xsk_prod_nb_free(&xsk->umem->fq,
	                                (uint32_t) xsk_umem_free_frames(xsk));
	if (stock_frames > 0) {
		/* ignoring prod__reserve return value, because we got stock_frames
		 * from xsk_prod_nb_free(), which are therefore available */
		xsk_ring_prod__reserve(&xsk->umem->fq, stock_frames, &idx_fq);

		for (uint32_t i = 0; i < stock_frames; ++i) {
			/* TODO: handle lack of available frames?
			 * Is not necessary when the total amount of frames exceeds the
			 * total slots available across all queues combined */
			/* uint64_t frame = xsk_alloc_umem_frame(xsk); */
			/* if (frame == XDP_INVALID_UMEM_FRAME) */
			/*     printf("xdp: trying to fill_addr INVALID UMEM FRAME"); */
			*xsk_ring_prod__fill_addr(&xsk->umem->fq, idx_fq++) =
				xsk_alloc_umem_frame(xsk);
		}

		xsk_ring_prod__submit(&xsk->umem->fq, stock_frames);
	}
}

static int load_xdp_program_and_map(struct xdp_server *xdp) {
	struct bpf_map *map;
	char errmsg[512];
	int err, ret;
	/* UNSPEC => let libxdp decide */
	// TODO: put this into a config option as well?
	enum xdp_attach_mode attach_mode = XDP_MODE_UNSPEC;

	DECLARE_LIBXDP_OPTS(bpf_object_open_opts, opts);
	if (xdp->bpf_bpffs_path)
		opts.pin_root_path = xdp->bpf_bpffs_path;

	/* for now our xdp program should contain just one program section */
	// TODO: look at xdp_program__create because it can take a pinned prog
	xdp->bpf_prog = xdp_program__open_file(xdp->bpf_prog_filename, NULL, &opts);

	// conversion should be fine, libxdp errors shouldn't exceed (int),
	// also libxdp_strerr takes int anyway...
	err = (int) libxdp_get_error(xdp->bpf_prog);
	if (err) {
		libxdp_strerror(err, errmsg, sizeof(errmsg));
		log_msg(LOG_ERR, "xdp: could not open xdp program: %s\n", errmsg);
		return err;
	}

	if (xdp->bpf_prog_should_load) {
		/* TODO: I find setting environment variables from within a program
		 * not a good thing to do, but for the meantime this helps... */
		/* This is done to allow unloading the XDP program we load without
		 * needing the SYS_ADMIN capability, and libxdp doesn't allow skipping
		 * the dispatcher through different means. */
		putenv("LIBXDP_SKIP_DISPATCHER=1");
		err = xdp_program__attach(xdp->bpf_prog, (int) xdp->interface_index, attach_mode, 0);
		/* err = xdp_program__attach_single(xdp->bpf_prog, xdp->interface_index, attach_mode); */
		if (err) {
			libxdp_strerror(err, errmsg, sizeof(errmsg));
			log_msg(LOG_ERR, "xdp: could not attach xdp program to interface '%s' : %s\n", 
					xdp->interface_name, errmsg);
			return err;
		}

		xdp->bpf_prog_fd = xdp_program__fd(xdp->bpf_prog);
		xdp->bpf_prog_id = xdp_program__id(xdp->bpf_prog);

		/* We also need to get the file descriptor to the xsks_map */
		map = bpf_object__find_map_by_name(xdp_program__bpf_obj(xdp->bpf_prog), "xsks_map");
		ret = bpf_map__fd(map);
		if (ret < 0) {
			log_msg(LOG_ERR, "xdp: no xsks map found in xdp program: %s\n", strerror(ret));
			return ret;
		}
		xdp->xsk_map_fd = ret;
		xdp->xsk_map = map;
	} else {
		char map_path[PATH_MAX];
		int fd;

		snprintf(map_path, PATH_MAX, "%s/%s", xdp->bpf_bpffs_path, "xsks_map");

		fd = bpf_obj_get(map_path);
		if (fd < 0) {
			log_msg(LOG_ERR, "xdp: could not retrieve xsks_map pin from %s: %s", map_path, strerror(errno));
			return fd;
		}

		map = bpf_object__find_map_by_name(xdp_program__bpf_obj(xdp->bpf_prog), "xsks_map");
		if ((ret = bpf_map__reuse_fd(map, fd))) {
			log_msg(LOG_ERR, "xdp: could not re-use xsks_map: %s\n", strerror(errno));
			return ret;
		}

		xdp->xsk_map_fd = fd;
		xdp->xsk_map = map;
	}

	return 0;
}

static int
xsk_configure_umem(struct xsk_umem_info *umem_info, uint64_t size) {
	int ret;
	struct xsk_umem_config umem_config = {
		.fill_size = XSK_RING_PROD__NUM_DESCS,
		.comp_size = XSK_RING_CONS__NUM_DESCS,
		.frame_size = XDP_FRAME_SIZE,
		.frame_headroom = XSK_UMEM_FRAME_HEADROOM,
		.flags = XSK_UMEM_FLAGS,
	};

	ret = xsk_umem__create(&umem_info->umem, umem_info->buffer, size, &umem_info->fq, &umem_info->cq, &umem_config);
	if (ret) {
		errno = -ret;
		return ret;
	}

	return 0;
}

static int
xsk_configure_socket(struct xdp_server *xdp, struct xsk_socket_info *xsk_info,
                     struct xsk_umem_info *umem, uint32_t queue_index) {
	struct xsk_socket_config xsk_cfg;
	uint32_t idx, reserved;
	int ret;

	struct xdp_config cfg = {
		.xdp_flags = 0,
		.xsk_bind_flags = 0,
		.libxdp_flags = XSK_LIBXDP_FLAGS__INHIBIT_PROG_LOAD,
	};

	uint16_t xsk_bind_flags = XDP_USE_NEED_WAKEUP;
	if (xdp->force_copy) {
		xsk_bind_flags |= XDP_COPY;
	}
	cfg.xsk_bind_flags = xsk_bind_flags;

	xsk_info->umem = umem;
	xsk_cfg.rx_size = XSK_RING_CONS__NUM_DESCS;
	xsk_cfg.tx_size = XSK_RING_PROD__NUM_DESCS;
	xsk_cfg.xdp_flags = cfg.xdp_flags;
	xsk_cfg.bind_flags = cfg.xsk_bind_flags;
	xsk_cfg.libxdp_flags = cfg.libxdp_flags;

	ret = xsk_socket__create(&xsk_info->xsk,
	                         xdp->interface_name,
	                         queue_index,
	                         umem->umem,
	                         &xsk_info->rx,
	                         &xsk_info->tx,
	                         &xsk_cfg);
	if (ret) {
		log_msg(LOG_ERR, "xdp: failed to create xsk_socket");
		goto error_exit;
	}

	ret = xsk_socket__update_xskmap(xsk_info->xsk, xdp->xsk_map_fd);
	if (ret) {
		log_msg(LOG_ERR, "xdp: failed to update xskmap");
		goto error_exit;
	}

	/* Initialize umem frame allocation */
	for (uint32_t i = 0; i < XDP_NUM_FRAMES; ++i) {
		xsk_info->umem->umem_frame_addr[i] = i * XDP_FRAME_SIZE;
	}

	xsk_info->umem->umem_frame_free = XDP_NUM_FRAMES;

	reserved = xsk_ring_prod__reserve(&xsk_info->umem->fq,
	                             XSK_RING_PROD__NUM_DESCS,
	                             &idx);

	if (reserved != XSK_RING_PROD__NUM_DESCS) {
		log_msg(LOG_ERR,
		        "xdp: amount of reserved addr not as expected (is %d)", reserved);
		// "ENOMEM 12 Cannot allocate memory" is the closest to the
		// error that not as much memory was reserved as expected
		ret = -12;
		goto error_exit;
	}

	for (uint32_t i = 0; i < XSK_RING_PROD__NUM_DESCS; ++i) {
		*xsk_ring_prod__fill_addr(&xsk_info->umem->fq, idx++) =
			xsk_alloc_umem_frame(xsk_info);
	}

	xsk_ring_prod__submit(&xsk_info->umem->fq, XSK_RING_PROD__NUM_DESCS);

	return 0;

error_exit:
	errno = -ret;
	return ret;
}

static void *alloc_shared_mem(size_t len) {
	/* MAP_ANONYMOUS memory is initialized with zero */
	return mmap(NULL, len, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, -1, 0);
}

static int xdp_sockets_init(struct xdp_server *xdp) {
	size_t umems_len = sizeof(struct xsk_umem_info) * xdp->queue_count;
	size_t xsks_len = sizeof(struct xsk_socket_info) * xdp->queue_count;

	xdp->umems = (struct xsk_umem_info *) alloc_shared_mem(umems_len);
	if (xdp->umems == MAP_FAILED) {
		log_msg(LOG_ERR,
		        "xdp: failed to allocate shared memory for umem info: %s",
		        strerror(errno));
		return -1;
	}

	xdp->xsks = (struct xsk_socket_info *) alloc_shared_mem(xsks_len);
	if (xdp->xsks == MAP_FAILED) {
		log_msg(LOG_ERR,
		        "xdp: failed to allocate shared memory for xsk info: %s",
		        strerror(errno));
		return -1;
	}

	for (uint32_t q_idx = 0; q_idx < xdp->queue_count; ++q_idx) {
		/* mmap is supposedly page-aligned, so should be fine */
		xdp->umems[q_idx].buffer = alloc_shared_mem(XDP_BUFFER_SIZE);

		if (xsk_configure_umem(&xdp->umems[q_idx], XDP_BUFFER_SIZE)) {
			log_msg(LOG_ERR, "xdp: cannot create umem: %s", strerror(errno));
			goto out_err_umem;
		}

		if (xsk_configure_socket(xdp, &xdp->xsks[q_idx], &xdp->umems[q_idx],
		                         q_idx)) {
			log_msg(LOG_ERR,
			        "xdp: cannot create AF_XDP socket: %s",
			        strerror(errno));
			goto out_err_xsk;
		}
	}

	return 0;

out_err_xsk:
	for (uint32_t i = 0; i < xdp->queue_count; ++i)
		xsk_umem__delete(xdp->umems[i].umem);

out_err_umem:
	return -1;
}

static void xdp_sockets_cleanup(struct xdp_server *xdp) {
	for (uint32_t i = 0; i < xdp->queue_count; ++i) {
		xsk_socket__delete(xdp->xsks[i].xsk);
		xsk_umem__delete(xdp->umems[i].umem);
	}
}

int xdp_server_init(struct xdp_server *xdp) {
	struct rlimit rlim = {RLIM_INFINITY, RLIM_INFINITY};

	/* check if interface name exists */
	xdp->interface_index = if_nametoindex(xdp->interface_name);
	if (xdp->interface_index == 0) {
		log_msg(LOG_ERR, "xdp: configured xdp-interface (%s) is unknown: %s",
		        xdp->interface_name, strerror(errno));
		return -1;
	}

	/* (optionally) load xdp program and (definitely) set xsks_map_fd */
	if (load_xdp_program_and_map(xdp)) {
		log_msg(LOG_ERR, "xdp: failed to load/pin xdp program/map");
		return -1;
	}

	/* if we don't do set rlimit, libbpf does it */
	/* this either has to be done before privilege drop or
	 * requires CAP_SYS_RESOURCE */
	if (setrlimit(RLIMIT_MEMLOCK, &rlim)) {
		log_msg(LOG_ERR, "xdp: cannot adjust rlimit (RLIMIT_MEMLOCK): \"%s\"\n",
			strerror(errno));
		return -1;
	}

	if (xdp_sockets_init(xdp))
		return -1;

	for (int i = 0; i < XDP_RX_BATCH_SIZE; ++i) {
		umem_ptrs[i].addr = XDP_INVALID_UMEM_FRAME;
		umem_ptrs[i].len = 0;
	}

	if (!xdp->ip_addresses)
		figure_ip_addresses(xdp);

	return 0;
}

void xdp_server_cleanup(struct xdp_server *xdp) {
	xdp_sockets_cleanup(xdp);

	/* only unpin if we loaded the program */
	if (xdp->bpf_prog_should_load) {
		if (xdp->xsk_map && bpf_map__is_pinned(xdp->xsk_map)) {
			if (bpf_map__unpin(xdp->xsk_map, NULL)) {
				/* We currently ship an XDP program that doesn't pin the map. So
				 * if this error happens, it is because the user specified their
				 * custom XDP program to load by NSD. Therefore they should know
				 * about the pinned map and be able to unlink it themselves.
				 */
				log_msg(LOG_ERR, "xdp: failed to unpin bpf map during cleanup: \"%s\". "
				        "This is usually ok, but you need to unpin the map yourself. "
				        "This can usually be fixed by executing chmod o+wx %s\n",
				        strerror(errno), xdp->bpf_bpffs_path);
			}
		}

		unload_xdp_program(xdp);
	}
}

static void unload_xdp_program(struct xdp_server *xdp) {
	DECLARE_LIBBPF_OPTS(bpf_xdp_attach_opts, bpf_opts,
	                    .old_prog_fd = xdp->bpf_prog_fd);

	log_msg(LOG_INFO, "xdp: detaching xdp program %u from %s\n",
			xdp->bpf_prog_id, xdp->interface_name);

	if (bpf_xdp_detach((int) xdp->interface_index, 0, &bpf_opts))
		log_msg(LOG_ERR, "xdp: failed to detach xdp program: %s\n",
		        strerror(errno));
}

static int dest_ip_allowed6(struct xdp_server *xdp, struct ipv6hdr *ipv6) {
	struct xdp_ip_address *ip = xdp->ip_addresses;
	if (!ip)
		// no IPs available, allowing all
		return 1;

	while (ip) {
		if (ip->addr.ss_family == AF_INET6 &&
		    !memcmp(&(((struct sockaddr_in6 *) &ip->addr)->sin6_addr),
		            &ipv6->daddr,
		            sizeof(struct in6_addr)))
			return 1;
		ip = ip->next;
	}

	return 0;
}

static int dest_ip_allowed4(struct xdp_server *xdp, struct iphdr *ipv4) {
	struct xdp_ip_address *ip = xdp->ip_addresses;
	if (!ip)
		// no IPs available, allowing all
		return 1;

	while (ip) {
		if (ip->addr.ss_family == AF_INET &&
		    ipv4->daddr == ((struct sockaddr_in *) &ip->addr)->sin_addr.s_addr)
			return 1;
		ip = ip->next;
	}

	return 0;
}

static void
add_ip_address(struct xdp_server *xdp, struct sockaddr_storage *addr) {
	struct xdp_ip_address *ip = xdp->ip_addresses;
	if (!ip) {
		xdp->ip_addresses = region_alloc_zero(xdp->region,
		                                      sizeof(struct xdp_ip_address));
		ip = xdp->ip_addresses;
	} else {
		while (ip->next)
			ip = ip->next;

		ip->next = region_alloc_zero(xdp->region,
		                             sizeof(struct xdp_ip_address));
		ip = ip->next;
	}

	memcpy(&ip->addr, addr, sizeof(struct sockaddr_storage));
}

static int figure_ip_addresses(struct xdp_server *xdp) {
	// TODO: if using VLANs, also find appropriate IP addresses?
	struct ifaddrs *ifaddr;
	int family, ret = 0;

	if (getifaddrs(&ifaddr) == -1) {
		log_msg(LOG_ERR, "xdp: couldn't determine local IP addresses. "
		                 "Serving all IP addresses now");
		return -1;
	}

	for (struct ifaddrs *ifa = ifaddr; ifa != NULL; ifa = ifa->ifa_next) {
		if (ifa->ifa_addr == NULL)
			continue;

		if (strcmp(ifa->ifa_name, xdp->interface_name))
			continue;

		family = ifa->ifa_addr->sa_family;

		switch (family) {
			default:
				continue;
			case AF_INET:
			case AF_INET6:
				add_ip_address(xdp, (struct sockaddr_storage *) ifa->ifa_addr);
		}
	}

	freeifaddrs(ifaddr);
	return ret;
}

static inline void swap_eth(struct ethhdr *eth) {
	uint8_t tmp_mac[ETH_ALEN];
	memcpy(tmp_mac, eth->h_dest, ETH_ALEN);
	memcpy(eth->h_dest, eth->h_source, ETH_ALEN);
	memcpy(eth->h_source, tmp_mac, ETH_ALEN);
}

static inline void swap_udp(struct udphdr *udp) {
	uint16_t tmp_port; /* not touching endianness */
	tmp_port = udp->source;
	udp->source = udp->dest;
	udp->dest = tmp_port;
}

static inline void swap_ipv6(struct ipv6hdr *ipv6) {
	struct in6_addr tmp_ip;
	memcpy(&tmp_ip, &ipv6->saddr, sizeof(tmp_ip));
	memcpy(&ipv6->saddr, &ipv6->daddr, sizeof(tmp_ip));
	memcpy(&ipv6->daddr, &tmp_ip, sizeof(tmp_ip));
}

static inline void swap_ipv4(struct iphdr *ipv4) {
	struct in_addr tmp_ip;
	memcpy(&tmp_ip, &ipv4->saddr, sizeof(tmp_ip));
	memcpy(&ipv4->saddr, &ipv4->daddr, sizeof(tmp_ip));
	memcpy(&ipv4->daddr, &tmp_ip, sizeof(tmp_ip));
}

static inline void *parse_udp(struct udphdr *udp) {
	if (ntohs(udp->dest) != DNS_PORT)
		return NULL;

	return (void *)(udp + 1);
}

static inline void *parse_ipv6(struct ipv6hdr *ipv6) {
	if (ipv6->nexthdr != IPPROTO_UDP)
		return NULL;

	return (void *)(ipv6 + 1);
}

static inline void *parse_ipv4(struct iphdr *ipv4) {
	if (ipv4->protocol != IPPROTO_UDP)
		return NULL;

	return (void *)(ipv4 + 1);
}

static uint32_t parse_dns(struct nsd* nsd, uint32_t dnslen,
                          struct query *q, sa_family_t ai_family) {
	/* TODO: implement DNSTAP, PROXYv2, ...? */
	uint32_t now = 0;

	/* set the size of the dns message and move position to start */
	buffer_skip(q->packet, dnslen);
	buffer_flip(q->packet);

	if (query_process(q, nsd, &now) != QUERY_DISCARDED) {
		if (RCODE(q->packet) == RCODE_OK && !AA(q->packet)) {
			STATUP(nsd, nona);
			ZTATUP(nsd, q->zone, nona);
		}

#ifdef USE_ZONE_STATS
		if (ai_family == AF_INET) {
			ZTATUP(nsd, q->zone, qudp);
		} else if (ai_family == AF_INET6) {
			ZTATUP(nsd, q->zone, qudp6);
		}
#endif /* USE_ZONE_STATS */

		query_add_optional(q, nsd, &now);

		buffer_flip(q->packet);

#ifdef BIND8_STATS
			/* Account the rcode & TC... */
			STATUP2(nsd, rcode, RCODE(q->packet));
			ZTATUP2(nsd, q->zone, rcode, RCODE(q->packet));
			if (TC(q->packet)) {
				STATUP(nsd, truncated);
				ZTATUP(nsd, q->zone, truncated);
			}
#endif /* BIND8_STATS */

		/* return new dns message length */
		return (uint32_t) buffer_remaining(q->packet);
	} else {
		query_reset(q, UDP_MAX_MESSAGE_LEN, 0);
		STATUP(nsd, dropped);
		ZTATUP(nsd, q->zone, dropped);
		return 0;
	}
}

static int
process_packet(struct xdp_server *xdp, uint8_t *pkt,
               uint32_t *len, struct query *query) {
	/* log_msg(LOG_INFO, "xdp: received packet with len %d", *len); */

	uint32_t dnslen = *len;
	uint32_t data_before_dnshdr_len = 0;

	struct ethhdr *eth = (struct ethhdr *)pkt;
	struct ipv6hdr *ipv6 = NULL;
	struct iphdr *ipv4 = NULL;
	struct udphdr *udp = NULL;
	void *dnshdr = NULL;

	/* doing the check here, so that the packet/frame is large enough to contain
	 * at least an ethernet header, an ipv4 header (ipv6 header is larger), and
	 * a udp header.
	 */
	if (*len < (sizeof(*eth) + sizeof(struct iphdr) + sizeof(*udp)))
		return -1;

	data_before_dnshdr_len = sizeof(*eth) + sizeof(*udp);

	switch (ntohs(eth->h_proto)) {
	case ETH_P_IPV6: {
		ipv6 = (struct ipv6hdr *)(eth + 1);

		if (*len < (sizeof(*eth) + sizeof(*ipv6) + sizeof(*udp)))
			return -2;
		if (!(udp = parse_ipv6(ipv6)))
			return -3;

		dnslen -= (uint32_t) (sizeof(*eth) + sizeof(*ipv6) + sizeof(*udp));
		data_before_dnshdr_len += sizeof(*ipv6);

		if (!dest_ip_allowed6(xdp, ipv6))
			return -4;

		break;
	} case ETH_P_IP: {
		ipv4 = (struct iphdr *)(eth + 1);

		if (!(udp = parse_ipv4(ipv4)))
			return -5;

		dnslen -= (uint32_t) (sizeof(*eth) + sizeof(*ipv4) + sizeof(*udp));
		data_before_dnshdr_len += sizeof(*ipv4);

		if (!dest_ip_allowed4(xdp, ipv4))
			return -6;

		break;
	}

	/* TODO: vlan? */
	/* case ETH_P_8021AD: case ETH_P_8021Q: */
	/*     if (*len < (sizeof(*eth) + sizeof(*vlan))) */
	/*         break; */
	default:
		return -7;
	}

	if (!(dnshdr = parse_udp(udp)))
		return -8;

	query_set_buffer_data(query, dnshdr, XDP_FRAME_SIZE - data_before_dnshdr_len);

	if(ipv6) {
#ifdef INET6
		struct sockaddr_in6* sock6 = (struct sockaddr_in6*)&query->remote_addr;
		sock6->sin6_family = AF_INET6;
		sock6->sin6_port = udp->dest;
		sock6->sin6_flowinfo = 0;
		sock6->sin6_scope_id = 0;
		memcpy(&sock6->sin6_addr, &ipv6->saddr, sizeof(ipv6->saddr));
		query->remote_addrlen = (socklen_t)sizeof(struct sockaddr_in6);
#else
		return 0; /* no inet6 no network */
#endif /* INET6 */
#ifdef BIND8_STATS
		STATUP(xdp->nsd, qudp6);
#endif /* BIND8_STATS */
	} else {
		struct sockaddr_in* sock4 = (struct sockaddr_in*)&query->remote_addr;
		sock4->sin_family = AF_INET;
		sock4->sin_port = udp->dest;
		sock4->sin_addr.s_addr = ipv4->saddr;
		query->remote_addrlen = (socklen_t)sizeof(struct sockaddr_in);
#ifdef BIND8_STATS
		STATUP(xdp->nsd, qudp);
#endif /* BIND8_STATS */
	}

	query->client_addr    = query->remote_addr;
	query->client_addrlen = query->remote_addrlen;
	query->is_proxied = 0;

	dnslen = parse_dns(xdp->nsd, dnslen, query, query->remote_addr.ss_family);
	if (!dnslen) {
		return -9;
	}

	// Not verifying the packet length (that it fits in an IP packet), as
	// parse_dns truncates too long response messages.
	udp->len = htons((uint16_t) (sizeof(*udp) + dnslen));

	swap_eth(eth);
	swap_udp(udp);

	if (ipv4) {
		__be16 ipv4_old_len = ipv4->tot_len;
		swap_ipv4(ipv4);
		ipv4->tot_len = htons(sizeof(*ipv4)) + udp->len;
		csum16_replace(&ipv4->check, ipv4_old_len, ipv4->tot_len);
		udp->check = calc_csum_udp4(udp, ipv4);
	} else if (ipv6) {
		swap_ipv6(ipv6);
		ipv6->payload_len = udp->len;
		udp->check = calc_csum_udp6(udp, ipv6);
	} else {
		log_msg(LOG_ERR, "xdp: we forgot to implement something... oops");
		return 0;
	}

	/* log_msg(LOG_INFO, "xdp: done with processing the packet"); */

	*len = data_before_dnshdr_len + dnslen;
	return 1;
}

void xdp_handle_recv_and_send(struct xdp_server *xdp) {
	struct xsk_socket_info *xsk = &xdp->xsks[xdp->queue_index];
	unsigned int recvd, i, reserved, to_send = 0;
	uint32_t idx_rx = 0;
	uint32_t tx_idx = 0;
	int ret;

	recvd = xsk_ring_cons__peek(&xsk->rx, XDP_RX_BATCH_SIZE, &idx_rx);
	if (!recvd) {
		/* no data available */
		return;
	}

	fill_fq(xsk);

	/* Process received packets */
	for (i = 0; i < recvd; ++i) {
		uint64_t addr = xsk_ring_cons__rx_desc(&xsk->rx, idx_rx)->addr;
		uint32_t len = xsk_ring_cons__rx_desc(&xsk->rx, idx_rx++)->len;

		uint8_t *pkt = xsk_umem__get_data(xsk->umem->buffer, addr);
		if ((ret = process_packet(xdp, pkt, &len, xdp->queries[i])) <= 0) {
			/* drop packet */
			xsk_free_umem_frame(xsk, addr);
		} else {
			umem_ptrs[to_send].addr = addr;
			umem_ptrs[to_send].len = len;
			++to_send;
		}
		/* we can reset the query directly after each packet processing,
		 * because the reset does not delete the underlying buffer/data.
		 * However, if we, in future, need to access data from the query
		 * struct when sending the answer, this needs to change.
		 * This also means, that currently a single query instance (and
		 * not an array) would suffice for this implementation. */
		query_reset(xdp->queries[i], UDP_MAX_MESSAGE_LEN, 0);

		/* xsk->stats.rx_bytes += len; */
	}

	xsk_ring_cons__release(&xsk->rx, recvd);
	/* xsk->stats.rx_packets += rcvd; */

	/* Process sending packets */

	/* TODO: at least send as many packets as slots are available */
	reserved = xsk_ring_prod__reserve(&xsk->tx, to_send, &tx_idx);
	// if we can't reserve to_send frames, we'll get 0 frames, so
	// no need to "un-reserve"
	if (reserved != to_send) {
		// not enough tx slots available, drop packets
		log_msg(LOG_ERR, "xdp: not enough TX frames available, dropping "
		        "whole batch");
		for (i = 0; i < to_send; ++i) {
			xsk_free_umem_frame(xsk, umem_ptrs[i].addr);
			umem_ptrs[i].addr = XDP_INVALID_UMEM_FRAME;
			umem_ptrs[i].len = 0;
		}
#ifdef BIND8_STATS
		xdp->nsd->st->txerr += to_send;
#endif /* BIND8_STATS */
		to_send = 0;
	}

	for (i = 0; i < to_send; ++i) {
		uint64_t addr = umem_ptrs[i].addr;
		uint32_t len = umem_ptrs[i].len;
		xsk_ring_prod__tx_desc(&xsk->tx, tx_idx)->addr = addr;
		xsk_ring_prod__tx_desc(&xsk->tx, tx_idx)->len = len;
		tx_idx++;
		xsk->outstanding_tx++;
		umem_ptrs[i].addr = XDP_INVALID_UMEM_FRAME;
		umem_ptrs[i].len = 0;
	}

	xsk_ring_prod__submit(&xsk->tx, to_send);

	/* wake up kernel for tx if needed and collect completed tx buffers */
	handle_tx(xsk);
	/* TODO: maybe call fill_fq(xsk) here as well? */
}

static void drain_cq(struct xsk_socket_info *xsk) {
	uint32_t completed, idx_cq;

	/* free completed TX buffers */
	completed = xsk_ring_cons__peek(&xsk->umem->cq,
	                                XSK_RING_CONS__NUM_DESCS,
	                                &idx_cq);

	if (completed > 0) {
		for (uint32_t i = 0; i < completed; i++) {
			xsk_free_umem_frame(xsk, *xsk_ring_cons__comp_addr(&xsk->umem->cq,
			                                                   idx_cq++));
		}

		xsk_ring_cons__release(&xsk->umem->cq, completed);
		xsk->outstanding_tx -= completed < xsk->outstanding_tx ?
		                       completed : xsk->outstanding_tx;
	}
}

static void handle_tx(struct xsk_socket_info *xsk) {
	if (!xsk->outstanding_tx)
		return;

	if (xsk_ring_prod__needs_wakeup(&xsk->tx))
		sendto(xsk_socket__fd(xsk->xsk), NULL, 0, MSG_DONTWAIT, NULL, 0);

	drain_cq(xsk);

	// Update TX-queue pointers
	// This is not needed, because prod__reserve calls this function too,
	// and therefore, if not enough frames are free on the cached pointers,
	// it will update the real pointers.
	/* xsk_prod_nb_free(&xsk->tx, XSK_RING_PROD__NUM_DESCS/4); */
}

#endif /* USE_XDP */