// SPDX-License-Identifier: GPL-2.0
// XDP CIDR Blocker - High-performance packet filtering with LPM trie
// Supports whitelist (allowlist) and blocklist (denylist)
// Part of xdp-defense: chained via libxdp dispatcher (priority 10)
#include <linux/bpf.h>
#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/ipv6.h>
#include <linux/in.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_endian.h>
#include <xdp/xdp_helpers.h>

// LPM trie key for IPv4 CIDR matching
struct ipv4_lpm_key {
    __u32 prefixlen;
    __u32 addr;
};

// LPM trie key for IPv6 CIDR matching
struct ipv6_lpm_key {
    __u32 prefixlen;
    __u8 addr[16];
};

// VLAN header (802.1Q)
struct vlan_hdr {
    __be16 h_vlan_TCI;
    __be16 h_vlan_encapsulated_proto;
};

// Statistics structure
struct stats {
    __u64 packets;
    __u64 bytes;
};

// IPv4 WHITELIST - checked first, allows traffic even if in blocklist
// Pinned for sharing with xdp_ddos program
struct {
    __uint(type, BPF_MAP_TYPE_LPM_TRIE);
    __type(key, struct ipv4_lpm_key);
    __type(value, __u64);
    __uint(max_entries, 4096);
    __uint(map_flags, BPF_F_NO_PREALLOC);
    __uint(pinning, LIBBPF_PIN_BY_NAME);
} whitelist_v4 SEC(".maps");

// IPv4 BLOCKLIST - blocks traffic unless in whitelist
struct {
    __uint(type, BPF_MAP_TYPE_LPM_TRIE);
    __type(key, struct ipv4_lpm_key);
    __type(value, __u64);
    __uint(max_entries, 262144);
    __uint(map_flags, BPF_F_NO_PREALLOC);
} blocklist_v4 SEC(".maps");

// IPv6 whitelist - pinned for sharing with xdp_ddos program
struct {
    __uint(type, BPF_MAP_TYPE_LPM_TRIE);
    __type(key, struct ipv6_lpm_key);
    __type(value, __u64);
    __uint(max_entries, 4096);
    __uint(map_flags, BPF_F_NO_PREALLOC);
    __uint(pinning, LIBBPF_PIN_BY_NAME);
} whitelist_v6 SEC(".maps");

// IPv6 blocklist
struct {
    __uint(type, BPF_MAP_TYPE_LPM_TRIE);
    __type(key, struct ipv6_lpm_key);
    __type(value, __u64);
    __uint(max_entries, 262144);
    __uint(map_flags, BPF_F_NO_PREALLOC);
} blocklist_v6 SEC(".maps");

// Per-CPU statistics: 0=passed, 1=dropped, 2=whitelisted, 3=errors
struct {
    __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
    __type(key, __u32);
    __type(value, struct stats);
    __uint(max_entries, 4);
} stats_map SEC(".maps");

// Configuration map
struct {
    __uint(type, BPF_MAP_TYPE_ARRAY);
    __type(key, __u32);
    __type(value, __u32);
    __uint(max_entries, 1);  // 0=enabled/disabled
} config SEC(".maps");

static __always_inline void update_stats(__u32 idx, __u64 bytes) {
    struct stats *s = bpf_map_lookup_elem(&stats_map, &idx);
    if (s) {
        s->packets++;
        s->bytes += bytes;
    }
}

SEC("xdp")
int xdp_blocker(struct xdp_md *ctx) {
    void *data_end = (void *)(long)ctx->data_end;
    void *data = (void *)(long)ctx->data;

    // Check if enabled
    __u32 cfg_key = 0;
    __u32 *enabled = bpf_map_lookup_elem(&config, &cfg_key);
    if (enabled && *enabled == 0) {
        return XDP_PASS;
    }

    struct ethhdr *eth = data;
    if ((void *)(eth + 1) > data_end) {
        update_stats(3, 0);
        return XDP_PASS;
    }

    __u16 eth_proto = bpf_ntohs(eth->h_proto);
    __u64 pkt_len = data_end - data;
    void *l3_hdr = (void *)(eth + 1);

    // Handle VLAN tags (802.1Q and QinQ)
    if (eth_proto == ETH_P_8021Q || eth_proto == ETH_P_8021AD) {
        struct vlan_hdr *vhdr = l3_hdr;
        if ((void *)(vhdr + 1) > data_end) {
            update_stats(3, pkt_len);
            return XDP_PASS;
        }
        eth_proto = bpf_ntohs(vhdr->h_vlan_encapsulated_proto);
        l3_hdr = (void *)(vhdr + 1);

        // Handle QinQ (double VLAN)
        if (eth_proto == ETH_P_8021Q || eth_proto == ETH_P_8021AD) {
            vhdr = l3_hdr;
            if ((void *)(vhdr + 1) > data_end) {
                update_stats(3, pkt_len);
                return XDP_PASS;
            }
            eth_proto = bpf_ntohs(vhdr->h_vlan_encapsulated_proto);
            l3_hdr = (void *)(vhdr + 1);
        }
    }

    // Handle IPv4
    if (eth_proto == ETH_P_IP) {
        struct iphdr *iph = l3_hdr;
        if ((void *)(iph + 1) > data_end) {
            update_stats(3, pkt_len);
            return XDP_PASS;
        }

        struct ipv4_lpm_key key = {
            .prefixlen = 32,
            .addr = iph->saddr,
        };

        // Check WHITELIST first - if whitelisted, always allow
        if (bpf_map_lookup_elem(&whitelist_v4, &key)) {
            update_stats(2, pkt_len);  // whitelisted
            return XDP_PASS;
        }

        // Check BLOCKLIST - if blocked and not whitelisted, drop
        if (bpf_map_lookup_elem(&blocklist_v4, &key)) {
            update_stats(1, pkt_len);  // dropped
            return XDP_DROP;
        }
    }
    // Handle IPv6
    else if (eth_proto == ETH_P_IPV6) {
        struct ipv6hdr *ip6h = l3_hdr;
        if ((void *)(ip6h + 1) > data_end) {
            update_stats(3, pkt_len);
            return XDP_PASS;
        }

        struct ipv6_lpm_key key = {
            .prefixlen = 128,
        };
        __builtin_memcpy(key.addr, &ip6h->saddr, 16);

        // Check WHITELIST first
        if (bpf_map_lookup_elem(&whitelist_v6, &key)) {
            update_stats(2, pkt_len);
            return XDP_PASS;
        }

        // Check BLOCKLIST
        if (bpf_map_lookup_elem(&blocklist_v6, &key)) {
            update_stats(1, pkt_len);
            return XDP_DROP;
        }
    }

    update_stats(0, pkt_len);
    return XDP_PASS;
}

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

// libxdp dispatcher configuration: priority 10, chain on XDP_PASS
struct {
    __uint(priority, 10);
    __uint(XDP_PASS, 1);
} XDP_RUN_CONFIG(xdp_blocker);