/*
 * xdp-util.h -- set of xdp related helpers
 *
 * Copyright (c) 2024, NLnet Labs. All rights reserved.
 *
 * See LICENSE for the license.
 *
 */

#ifndef XDP_UTIL_H
#define XDP_UTIL_H

#include <linux/ip.h>
#include <linux/ipv6.h>
#include <linux/types.h>
#include <linux/udp.h>

/*
 * Get number of combined or rx queues available on ifname.
 */
int ethtool_channels_get(char const *ifname);

/*
 * Set capabilities to only include the ones needed for using AF_XDP/
 */
void set_caps(int unset_setid_caps);

/*
 * Add u16 to checksum value and preserve one's complement sum.
 */
static inline __sum16 csum16_add(__sum16 csum, __be16 addend) {
	uint16_t res = (uint16_t)csum;

	res += (__u16)addend;
	return (__sum16)(res + (res < (__u16)addend));
}

/*
 * Subtract u16 from checksum value and preserve one's complement sum.
 */
static inline __sum16 csum16_sub(__sum16 csum, __be16 addend) {
	return csum16_add(csum, ~addend);
}

/*
 * Replace u16 from checksum value and preserve one's complement sum.
 */
static inline void csum16_replace(__sum16 *sum, __be16 old, __be16 new) {
	*sum = ~csum16_add(csum16_sub(~(*sum), old), new);
}

/*
 * Sum up _data_len amount of 16-bit words in _data and add to result.
 */
static inline void csum_add_data(uint32_t *result,
                          const void *_data,
                          uint32_t len) {
	const uint16_t *data = _data;
	while (len > 1) {
		*result += *data++;
		len -= 2;
	}
	if (len)
		*result += *data & 0xff;
		// *result += *(uint8_t *)data;
}

/*
 * Add single 16-bit words to result.
 */
static inline void csum_add_u16(uint32_t *result, uint16_t x) { *result += x; }

/*
 * Apply one's complement to result.
 */
static inline void csum_reduce(uint32_t *result) {
	while (*result >> 16)
		*result = (*result & 0xffff) + (*result >> 16);
}

/*
 * Calculate UDP checksum with IPv6 pseudo-header
 */
static inline uint16_t calc_csum_udp6(struct udphdr *udp, struct ipv6hdr *ipv6) {
	uint32_t sum = 0;
	sum += udp->len;
	sum += htons(IPPROTO_UDP);
	csum_add_data(&sum, &ipv6->saddr, sizeof(ipv6->saddr));
	csum_add_data(&sum, &ipv6->daddr, sizeof(ipv6->daddr));

	udp->check = 0;
	csum_add_data(&sum, udp, ntohs(udp->len));
	/* maybe restore previous checksum to remove side effects? */

	// reduces sum to 16bit
	csum_reduce(&sum);

	if (sum != 0xffff)
		return (uint16_t) ~sum;
	else
		return (uint16_t) sum;
}

/*
 * Calculate UDP checksum with IPv4 pseudo-header
 */
static inline uint16_t calc_csum_udp4(struct udphdr *udp, struct iphdr *ipv4) {
	uint32_t sum = 0;
	sum += udp->len;
	sum += htons(IPPROTO_UDP);
	csum_add_data(&sum, &ipv4->saddr, sizeof(ipv4->saddr));
	csum_add_data(&sum, &ipv4->daddr, sizeof(ipv4->daddr));

	udp->check = 0;
	csum_add_data(&sum, udp, ntohs(udp->len));
	/* maybe restore previous checksum to remove side effects? */

	// reduces sum to 16bit
	csum_reduce(&sum);

	if (sum != 0xffff)
		return (uint16_t) ~sum;
	else
		return (uint16_t) sum;
}

#endif /* XDP_UTIL_H */