/*=====================================================================
 * hellgate.c - THE ROUTER FROM HELL - XDP early-drop gate
 *
 * Attaches to the WAN interface (enp16s0f0np0) in native/driver mode
 * (i40e). Drops garbage before skb allocation:
 *
 *   1. IPv4 packets whose SOURCE address falls in a prefix listed in
 *      the pinned LPM map "hellgate_drop4" (bogons/RFC1918 by default,
 *      extendable at runtime with bpftool - no recompile, no reload).
 *   2. Optionally, anything hitting a destination port listed in the
 *      pinned hash map "hellgate_dports" (e.g. scanner bait ports).
 *      Empty by default = feature off.
 *
 * Everything else: XDP_PASS. WireGuard, OSPF unicast, your SSH, DNS,
 * established flows - all untouched, all still governed by nftables.
 *
 * Per-CPU stats are kept in "hellgate_stats":
 *   idx 0 = passed, 1 = dropped_src, 2 = dropped_dport,
 *   idx 3 = passed_non_ip (arp etc), 4 = malformed (dropped)
 *
 * Design notes:
 *  - LPM_TRIE lookup is the same trie walk the kernel FIB does, but
 *    against a tiny table that lives entirely in cache.
 *  - No VLAN parsing: WAN is untagged. If you ever attach this to a
 *    tagged interface, add 802.1Q handling or it will PASS tagged
 *    frames unexamined (fail-open, which is the correct default).
 *  - IPv6 is passed untouched (project policy: no v6 until Socket
 *    offers it natively).
 *
 * Build:   see Makefile (clang -O2 -g -target bpf)
 * Attach:  ip link set dev enp16s0f0np0 xdpdrv obj hellgate.o sec xdp
 *          (or use the hellgate OpenRC service)
 *=====================================================================*/

#include <linux/bpf.h>
#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/udp.h>
#include <linux/tcp.h>
#include <linux/in.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_endian.h>

/* ----- stats indices ------------------------------------------------*/
enum {
    ST_PASS = 0,
    ST_DROP_SRC,
    ST_DROP_DPORT,
    ST_PASS_NONIP,
    ST_DROP_MALFORMED,
    ST_MAX
};

/* ----- maps (pinned under /sys/fs/bpf/hellgate/) --------------------*/

struct lpm_v4_key {
    __u32 prefixlen;
    __u32 addr;          /* network byte order */
};

struct {
    __uint(type, BPF_MAP_TYPE_LPM_TRIE);
    __uint(max_entries, 1024);
    __type(key, struct lpm_v4_key);
    __type(value, __u8);
    __uint(map_flags, BPF_F_NO_PREALLOC);
    __uint(pinning, LIBBPF_PIN_BY_NAME);
} hellgate_drop4 SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 64);
    __type(key, __u16);   /* dest port, host byte order */
    __type(value, __u8);
    __uint(pinning, LIBBPF_PIN_BY_NAME);
} hellgate_dports SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
    __uint(max_entries, ST_MAX);
    __type(key, __u32);
    __type(value, __u64);
    __uint(pinning, LIBBPF_PIN_BY_NAME);
} hellgate_stats SEC(".maps");

static __always_inline void count(__u32 idx)
{
    __u64 *v = bpf_map_lookup_elem(&hellgate_stats, &idx);
    if (v)
        __sync_fetch_and_add(v, 1);
}

/* ----- the gate -----------------------------------------------------*/

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

    /* Ethernet */
    struct ethhdr *eth = data;
    if ((void *)(eth + 1) > data_end) {
        count(ST_DROP_MALFORMED);
        return XDP_DROP;
    }

    /* Only inspect IPv4. ARP, IPv6, LLDP, whatever else: pass. */
    if (eth->h_proto != bpf_htons(ETH_P_IP)) {
        count(ST_PASS_NONIP);
        return XDP_PASS;
    }

    /* IPv4 header */
    struct iphdr *iph = (void *)(eth + 1);
    if ((void *)(iph + 1) > data_end) {
        count(ST_DROP_MALFORMED);
        return XDP_DROP;
    }

    /* 1: source-prefix gate (bogons, RFC1918, runtime blocklist) */
    struct lpm_v4_key key = {
        .prefixlen = 32,
        .addr      = iph->saddr,
    };
    if (bpf_map_lookup_elem(&hellgate_drop4, &key)) {
        count(ST_DROP_SRC);
        return XDP_DROP;
    }

    /* 2: dest-port gate (only if the map has entries you added) */
    __u8 ihl_bytes = iph->ihl * 4;
    if (ihl_bytes < sizeof(*iph)) {
        count(ST_DROP_MALFORMED);
        return XDP_DROP;
    }

    void *l4 = (void *)iph + ihl_bytes;
    __u16 dport_h = 0;

    if (iph->protocol == IPPROTO_TCP) {
        struct tcphdr *th = l4;
        if ((void *)(th + 1) > data_end)
            goto pass;                    /* fragment/short: let stack judge */
        dport_h = bpf_ntohs(th->dest);
    } else if (iph->protocol == IPPROTO_UDP) {
        struct udphdr *uh = l4;
        if ((void *)(uh + 1) > data_end)
            goto pass;
        dport_h = bpf_ntohs(uh->dest);
    } else {
        goto pass;                        /* ICMP, ESP, OSPF... not ours */
    }

    if (bpf_map_lookup_elem(&hellgate_dports, &dport_h)) {
        count(ST_DROP_DPORT);
        return XDP_DROP;
    }

pass:
    count(ST_PASS);
    return XDP_PASS;
}

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