// SPDX-License-Identifier: GPL-2.0
#define BPF_NO_KFUNC_PROTOTYPES
#include <vmlinux.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_endian.h>

#define ETH_P_IP	0x0800
#define ETH_P_IPV6	0x86dd
#define IP_MF		0x2000	/* "More Fragments" */
#define IP_OFFSET	0x1fff	/* "Fragment Offset" */
#define AF_INET		2
#define AF_INET6	10

struct bpf_flowtable_opts___local {
	s32 error;
};

struct flow_offload_tuple_rhash *
bpf_xdp_flow_lookup(struct xdp_md *, struct bpf_fib_lookup *,
		    struct bpf_flowtable_opts___local *, u32) __ksym;

struct {
	__uint(type, BPF_MAP_TYPE_ARRAY);
	__type(key, __u32);
	__type(value, __u32);
	__uint(max_entries, 1);
} stats SEC(".maps");

static bool xdp_flowtable_offload_check_iphdr(struct iphdr *iph)
{
	/* ip fragmented traffic */
	if (iph->frag_off & bpf_htons(IP_MF | IP_OFFSET))
		return false;

	/* ip options */
	if (iph->ihl * 4 != sizeof(*iph))
		return false;

	if (iph->ttl <= 1)
		return false;

	return true;
}

static bool xdp_flowtable_offload_check_tcp_state(void *ports, void *data_end,
						  u8 proto)
{
	if (proto == IPPROTO_TCP) {
		struct tcphdr *tcph = ports;

		if (tcph + 1 > data_end)
			return false;

		if (tcph->fin || tcph->rst)
			return false;
	}

	return true;
}

struct flow_ports___local {
	__be16 source, dest;
} __attribute__((preserve_access_index));

SEC("xdp.frags")
int xdp_flowtable_do_lookup(struct xdp_md *ctx)
{
	void *data_end = (void *)(long)ctx->data_end;
	struct bpf_flowtable_opts___local opts = {};
	struct flow_offload_tuple_rhash *tuplehash;
	struct bpf_fib_lookup tuple = {
		.ifindex = ctx->ingress_ifindex,
	};
	void *data = (void *)(long)ctx->data;
	struct ethhdr *eth = data;
	struct flow_ports___local *ports;
	__u32 *val, key = 0;

	if (eth + 1 > data_end)
		return XDP_DROP;

	switch (eth->h_proto) {
	case bpf_htons(ETH_P_IP): {
		struct iphdr *iph = data + sizeof(*eth);

		ports = (struct flow_ports___local *)(iph + 1);
		if (ports + 1 > data_end)
			return XDP_PASS;

		/* sanity check on ip header */
		if (!xdp_flowtable_offload_check_iphdr(iph))
			return XDP_PASS;

		if (!xdp_flowtable_offload_check_tcp_state(ports, data_end,
							   iph->protocol))
			return XDP_PASS;

		tuple.family		= AF_INET;
		tuple.tos		= iph->tos;
		tuple.l4_protocol	= iph->protocol;
		tuple.tot_len		= bpf_ntohs(iph->tot_len);
		tuple.ipv4_src		= iph->saddr;
		tuple.ipv4_dst		= iph->daddr;
		tuple.sport		= ports->source;
		tuple.dport		= ports->dest;
		break;
	}
	case bpf_htons(ETH_P_IPV6): {
		struct in6_addr *src = (struct in6_addr *)tuple.ipv6_src;
		struct in6_addr *dst = (struct in6_addr *)tuple.ipv6_dst;
		struct ipv6hdr *ip6h = data + sizeof(*eth);

		ports = (struct flow_ports___local *)(ip6h + 1);
		if (ports + 1 > data_end)
			return XDP_PASS;

		if (ip6h->hop_limit <= 1)
			return XDP_PASS;

		if (!xdp_flowtable_offload_check_tcp_state(ports, data_end,
							   ip6h->nexthdr))
			return XDP_PASS;

		tuple.family		= AF_INET6;
		tuple.l4_protocol	= ip6h->nexthdr;
		tuple.tot_len		= bpf_ntohs(ip6h->payload_len);
		*src			= ip6h->saddr;
		*dst			= ip6h->daddr;
		tuple.sport		= ports->source;
		tuple.dport		= ports->dest;
		break;
	}
	default:
		return XDP_PASS;
	}

	tuplehash = bpf_xdp_flow_lookup(ctx, &tuple, &opts, sizeof(opts));
	if (!tuplehash)
		return XDP_PASS;

	val = bpf_map_lookup_elem(&stats, &key);
	if (val)
		__sync_add_and_fetch(val, 1);

	return XDP_PASS;
}

char _license[] SEC("license") = "GPL";