// SPDX-License-Identifier: GPL-2.0

#include <linux/bpf.h>
#include <bpf/bpf_helpers.h>
#include "bpf_misc.h"

SEC(".maps") struct {
	__uint(type, BPF_MAP_TYPE_SOCKMAP);
	__uint(max_entries, 1);
	__type(key, __u32);
	__type(value, __u64);
} nop_map, sock_map;

SEC(".maps") struct {
	__uint(type, BPF_MAP_TYPE_SOCKHASH);
	__uint(max_entries, 1);
	__type(key, __u32);
	__type(value, __u64);
} nop_hash, sock_hash;

SEC(".maps") struct {
	__uint(type, BPF_MAP_TYPE_ARRAY);
	__uint(max_entries, 2);
	__type(key, int);
	__type(value, unsigned int);
} verdict_map;

/* Set by user space */
int redirect_type;
int redirect_flags;

#define redirect_map(__data)                                                   \
	_Generic((__data),                                                     \
		 struct __sk_buff * : bpf_sk_redirect_map,                     \
		 struct sk_msg_md * : bpf_msg_redirect_map                     \
	)((__data), &sock_map, (__u32){0}, redirect_flags)

#define redirect_hash(__data)                                                  \
	_Generic((__data),                                                     \
		 struct __sk_buff * : bpf_sk_redirect_hash,                    \
		 struct sk_msg_md * : bpf_msg_redirect_hash                    \
	)((__data), &sock_hash, &(__u32){0}, redirect_flags)

#define DEFINE_PROG(__type, __param)                                           \
SEC("sk_" XSTR(__type))                                                        \
int prog_ ## __type ## _verdict(__param data)                                  \
{                                                                              \
	unsigned int *count;                                                   \
	int verdict;                                                           \
									       \
	if (redirect_type == BPF_MAP_TYPE_SOCKMAP)                             \
		verdict = redirect_map(data);                                  \
	else if (redirect_type == BPF_MAP_TYPE_SOCKHASH)                       \
		verdict = redirect_hash(data);                                 \
	else                                                                   \
		verdict = redirect_type - __MAX_BPF_MAP_TYPE;                  \
									       \
	count = bpf_map_lookup_elem(&verdict_map, &verdict);                   \
	if (count)                                                             \
		(*count)++;                                                    \
									       \
	return verdict;                                                        \
}

DEFINE_PROG(skb, struct __sk_buff *);
DEFINE_PROG(msg, struct sk_msg_md *);

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