// udp_flood_blocker.bpf.c
// Compile with: clang -O2 -target bpf -c udp_flood_blocker.bpf.c -o udp_flood_blocker.bpf.o

#define KBUILD_MODNAME "udp_flood_blocker"
#include <linux/bpf.h>
#include <bpf/bpf_helpers.h>
#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/udp.h>

/*
  Map #1: udp_count
  - Type: PERCPU_ARRAY
  - Holds a single key (0) storing a 64-bit counter of UDP packets.
*/
struct {
    __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
    __uint(max_entries, 1);
    __type(key, __u32);
    __type(value, __u64);
} udp_count SEC(".maps");

/*
  Map #2: block_state
  - Type: ARRAY
  - Holds a single key (0) storing a 32-bit "block" flag (0 = allow, 1 = drop).
*/
struct {
    __uint(type, BPF_MAP_TYPE_ARRAY);
    __uint(max_entries, 1);
    __type(key, __u32);
    __type(value, __u32);
} block_state SEC(".maps");

// XDP hook function: detect UDP and optionally block
SEC("xdp")
int detect_and_block_udp(struct xdp_md *ctx)
{
    void *data_end = (void *)(long)ctx->data_end;
    void *data     = (void *)(long)ctx->data;

    // parse Ethernet header
    struct ethhdr *eth = data;
    if ((void *)(eth + 1) > data_end) {
        return XDP_PASS; // packet too small/malformed
    }

    // check if it's IPv4
    if (bpf_ntohs(eth->h_proto) == ETH_P_IP) {
        struct iphdr *iph = (struct iphdr *)(eth + 1);
        if ((void *)(iph + 1) > data_end) {
            return XDP_PASS;
        }

        // check if it's UDP
        if (iph->protocol == IPPROTO_UDP) {
            // 1) Check block_state
            __u32 zero_key = 0;
            __u32 *block_val = bpf_map_lookup_elem(&block_state, &zero_key);
            if (block_val && *block_val == 1) {
                // If we are in block mode, drop all UDP
                return XDP_DROP;
            }

            // 2) Not blocked => increment the global UDP counter
            __u64 *cnt_val = bpf_map_lookup_elem(&udp_count, &zero_key);
            if (cnt_val) {
                __sync_fetch_and_add(cnt_val, 1);
            }
        }
    }

    // pass non-UDP or non-IPv4 packets
    return XDP_PASS;
}

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