// SPDX-License-Identifier: GPL-2.0-only
/*
 * BastionGuard™ Secure Connection — eBPF probe kernel-side (esteso)
 * Copyright (C) 2025–2026 Calogero Scarnà
 *
 * COPERTURA EVENTI
 * ================
 * Il programma aggancia tre punti nel kernel per catturare praticamente
 * tutte le connessioni uscenti di ogni protocollo supportato:
 *
 *  1) tracepoint/sock/inet_sock_set_state  →  transizioni TCP a ESTABLISHED.
 *     Copre sia connessioni attive (dopo connect()) sia passive (dopo
 *     accept()). Avevamo gia' questo: resta per retro-compatibilita' e
 *     perche' e' il punto piu' "pulito" per attribuire il processo
 *     di un socket appena handshake-ato.
 *
 *  2) kprobe/tcp_v4_connect  +  kprobe/tcp_v6_connect
 *     Cattura ogni tentativo di connect() TCP — anche quelli che falliscono
 *     per network unreachable, timeout SYN, ecc. Era il gap principale:
 *     prima vedevamo solo le connessioni che si stabilivano correttamente.
 *
 *  3) kprobe/udp_sendmsg  +  kprobe/udpv6_sendmsg
 *     Cattura ogni datagramma UDP in uscita con destinazione esplicita.
 *     Il tracepoint inet_sock_set_state NON copre UDP senza connect(),
 *     e la stragrande maggioranza del traffico UDP (DNS, QUIC, mDNS,
 *     NTP, DHCP...) usa sendto()/sendmsg() diretti.
 *
 * DEDUPLICAZIONE UDP
 * ==================
 * Senza dedup, un client QUIC spara migliaia di sendmsg() al secondo e
 * la ring buffer si riempirebbe di eventi ripetuti. Usiamo una LRU hash
 * map (BPF_MAP_TYPE_LRU_HASH, 8192 entry) con chiave <pid, dst_ip, dst_port>
 * e valore timestamp ns: emettiamo un evento solo se non ne abbiamo emesso
 * uno identico negli ultimi UDP_DEDUP_WINDOW_NS nanosecondi (default 5s).
 * La LRU evita la fuga di memoria: entry vecchie vengono sfrattate in
 * automatico.
 *
 * Lo stesso meccanismo (con finestra piu' corta) si applica a TCP connect
 * per gestire il retry: Linux riprova SYN 3-5 volte, non ci interessa
 * vedere 5 righe "Firefox connect a 1.2.3.4:443".
 *
 * REQUISITI
 * =========
 *   kernel >= 5.8   (bpf_ringbuf_*, LRU_HASH, fentry a volonta')
 *   BTF abilitato   (CONFIG_DEBUG_INFO_BTF=y)
 *   CAP_BPF + CAP_PERFMON  (o CAP_SYS_ADMIN in fallback)
 */

#include "vmlinux.h"
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>
#include <bpf/bpf_core_read.h>
#include <bpf/bpf_endian.h>

/* ── Costanti kernel ──────────────────────────────────────────────── */
#define AF_INET        2
#define AF_INET6      10
#define TCP_ESTABLISHED 1

#define IPPROTO_TCP    6
#define IPPROTO_UDP   17

/* Dimensione ring buffer: 8 MB — piu' grande dell'originale perche'
 * adesso copriamo anche UDP/connect e il volume totale cresce. */
#define RINGBUF_SIZE (8 * 1024 * 1024)

/* Finestra dedup in nanosecondi. */
#define UDP_DEDUP_WINDOW_NS (5ULL * 1000ULL * 1000ULL * 1000ULL)  /* 5 s  */
#define TCP_DEDUP_WINDOW_NS (2ULL * 1000ULL * 1000ULL * 1000ULL)  /* 2 s  */

/* ── Struttura evento condivisa con userspace ─────────────────────── */
/*
 * ATTENZIONE: deve combaciare byte-per-byte con bsc::ebpf::ConnEvent
 * in daemon/include/bsc/ebpf/EBPFMon.hpp. Campo `event_kind` aggiunto
 * per distinguere la sorgente (tracepoint/kprobe); i vecchi loader che
 * non lo conoscono leggeranno 0 nel padding, quindi retro-compat ok.
 */
struct bsc_conn_event {
    /* processo */
    __u32 pid;
    __u32 tid;
    __u32 uid;
    __u32 gid;
    char  comm[16];       /* task_comm_len */

    /* rete */
    __u8  family;         /* AF_INET o AF_INET6 */
    __u8  proto;          /* IPPROTO_TCP=6 o IPPROTO_UDP=17 */
    __u8  event_kind;     /* 0=tp_set_state, 1=tcp_connect, 2=udp_sendmsg */
    __u8  _pad[1];
    __u16 src_port;       /* host byte order */
    __u16 dst_port;       /* host byte order */

    /* indirizzi: primi 4 byte validi per IPv4, tutti 16 per IPv6 */
    __u8  src_addr[16];
    __u8  dst_addr[16];

    /* timestamp monotono (ns) */
    __u64 ts_ns;
};

/* Chiave di dedup: pid + dst_ip + dst_port sufficiente.
 * Non includiamo src_port perche' per UDP e' effimero e farebbe saltare
 * il dedup. */
struct dedup_key {
    __u32 pid;
    __u16 dst_port;
    __u8  family;
    __u8  proto;
    __u8  dst_addr[16];
};

/* ── Maps ─────────────────────────────────────────────────────────── */
struct {
    __uint(type, BPF_MAP_TYPE_RINGBUF);
    __uint(max_entries, RINGBUF_SIZE);
} bsc_events SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_LRU_HASH);
    __uint(max_entries, 8192);
    __type(key,   struct dedup_key);
    __type(value, __u64);              /* ts_ns ultimo emit */
} bsc_dedup SEC(".maps");

/* ── Helper: riempie i campi "processo" dell'evento ───────────────── */
static __always_inline void fill_proc(struct bsc_conn_event* ev)
{
    const __u64 pid_tgid = bpf_get_current_pid_tgid();
    ev->pid = (__u32)(pid_tgid >> 32);
    ev->tid = (__u32)(pid_tgid & 0xffffffffull);
    const __u64 uid_gid = bpf_get_current_uid_gid();
    ev->uid = (__u32)(uid_gid & 0xffffffffull);
    ev->gid = (__u32)(uid_gid >> 32);
    bpf_get_current_comm(ev->comm, sizeof(ev->comm));
}

/* ── Helper: dedup check. Ritorna true se dobbiamo emettere. ──────── */
static __always_inline bool should_emit(const struct dedup_key* k,
                                        __u64 now_ns,
                                        __u64 window_ns)
{
    __u64* last = bpf_map_lookup_elem(&bsc_dedup, k);
    if (last) {
        if (now_ns - *last < window_ns) return false;  /* scarto: duplicato */
        *last = now_ns;                                /* update in-place */
        return true;
    }
    bpf_map_update_elem(&bsc_dedup, k, &now_ns, BPF_ANY);
    return true;
}

/* =================================================================== *
 * 1) tracepoint TCP state change (ESTABLISHED)
 * =================================================================== */
SEC("tracepoint/sock/inet_sock_set_state")
int bsc_inet_sock_set_state(struct trace_event_raw_inet_sock_set_state *ctx)
{
    if (ctx->newstate != TCP_ESTABLISHED)
        return 0;

    const __u16 family = ctx->family;
    if (family != AF_INET && family != AF_INET6)
        return 0;

    const __u8 proto = (__u8)ctx->protocol;
    if (proto != IPPROTO_TCP && proto != IPPROTO_UDP)
        return 0;

    struct bsc_conn_event *ev = bpf_ringbuf_reserve(&bsc_events,
                                                    sizeof(*ev), 0);
    if (!ev) return 0;

    fill_proc(ev);
    ev->family     = (__u8)family;
    ev->proto      = proto;
    ev->event_kind = 0;
    ev->_pad[0]    = 0;
    ev->src_port   = bpf_ntohs(ctx->sport);
    ev->dst_port   = bpf_ntohs(ctx->dport);

    if (family == AF_INET) {
        __builtin_memset(ev->src_addr, 0, sizeof(ev->src_addr));
        __builtin_memset(ev->dst_addr, 0, sizeof(ev->dst_addr));
        bpf_probe_read_kernel(ev->src_addr, 4, &ctx->saddr);
        bpf_probe_read_kernel(ev->dst_addr, 4, &ctx->daddr);
    } else {
        BPF_CORE_READ_INTO(ev->src_addr, ctx, saddr_v6);
        BPF_CORE_READ_INTO(ev->dst_addr, ctx, daddr_v6);
    }

    ev->ts_ns = bpf_ktime_get_ns();
    bpf_ringbuf_submit(ev, 0);
    return 0;
}

/* =================================================================== *
 * 2) TCP connect (IPv4 + IPv6)
 * =================================================================== *
 *
 * Firma:  int tcp_v4_connect(struct sock* sk, struct sockaddr* uaddr, int len)
 *         int tcp_v6_connect(struct sock* sk, struct sockaddr* uaddr, int len)
 *
 * Emettiamo l'evento subito all'ingresso — abbiamo gia' tutti i dati
 * (sk contiene src/dst). Il connect puo' ancora fallire ma per il
 * firewall conta il tentativo.
 */
SEC("kprobe/tcp_v4_connect")
int BPF_KPROBE(bsc_tcp_v4_connect, struct sock* sk)
{
    if (!sk) return 0;

    const __u64 now = bpf_ktime_get_ns();

    /* Leggi destinazione dalla struct sock (inet_sock → skc_daddr/dport).
     * NOTA: in tcp_v4_connect la destinazione viene scritta prima
     * della chiamata effettiva al connect — quindi e' gia' nel sk. */
    __u32 daddr  = BPF_CORE_READ(sk, __sk_common.skc_daddr);
    __u32 saddr  = BPF_CORE_READ(sk, __sk_common.skc_rcv_saddr);
    __u16 dport  = BPF_CORE_READ(sk, __sk_common.skc_dport);   /* net byte order */
    __u16 sport  = BPF_CORE_READ(sk, __sk_common.skc_num);     /* host byte order */
    dport = bpf_ntohs(dport);

    /* saddr puo' essere 0 (kernel non ha ancora scelto src): non e' errore */

    struct dedup_key k = {};
    const __u64 pid_tgid = bpf_get_current_pid_tgid();
    k.pid      = (__u32)(pid_tgid >> 32);
    k.dst_port = dport;
    k.family   = AF_INET;
    k.proto    = IPPROTO_TCP;
    __builtin_memcpy(k.dst_addr, &daddr, 4);
    if (!should_emit(&k, now, TCP_DEDUP_WINDOW_NS)) return 0;

    struct bsc_conn_event *ev = bpf_ringbuf_reserve(&bsc_events,
                                                    sizeof(*ev), 0);
    if (!ev) return 0;

    fill_proc(ev);
    ev->family     = AF_INET;
    ev->proto      = IPPROTO_TCP;
    ev->event_kind = 1;
    ev->_pad[0]    = 0;
    ev->src_port   = sport;
    ev->dst_port   = dport;

    __builtin_memset(ev->src_addr, 0, sizeof(ev->src_addr));
    __builtin_memset(ev->dst_addr, 0, sizeof(ev->dst_addr));
    __builtin_memcpy(ev->src_addr, &saddr, 4);
    __builtin_memcpy(ev->dst_addr, &daddr, 4);

    ev->ts_ns = now;
    bpf_ringbuf_submit(ev, 0);
    return 0;
}

SEC("kprobe/tcp_v6_connect")
int BPF_KPROBE(bsc_tcp_v6_connect, struct sock* sk)
{
    if (!sk) return 0;

    const __u64 now = bpf_ktime_get_ns();

    __u16 dport = BPF_CORE_READ(sk, __sk_common.skc_dport);
    __u16 sport = BPF_CORE_READ(sk, __sk_common.skc_num);
    dport = bpf_ntohs(dport);

    struct dedup_key k = {};
    const __u64 pid_tgid = bpf_get_current_pid_tgid();
    k.pid      = (__u32)(pid_tgid >> 32);
    k.dst_port = dport;
    k.family   = AF_INET6;
    k.proto    = IPPROTO_TCP;
    BPF_CORE_READ_INTO(k.dst_addr, sk, __sk_common.skc_v6_daddr);
    if (!should_emit(&k, now, TCP_DEDUP_WINDOW_NS)) return 0;

    struct bsc_conn_event *ev = bpf_ringbuf_reserve(&bsc_events,
                                                    sizeof(*ev), 0);
    if (!ev) return 0;

    fill_proc(ev);
    ev->family     = AF_INET6;
    ev->proto      = IPPROTO_TCP;
    ev->event_kind = 1;
    ev->_pad[0]    = 0;
    ev->src_port   = sport;
    ev->dst_port   = dport;

    BPF_CORE_READ_INTO(ev->src_addr, sk, __sk_common.skc_v6_rcv_saddr);
    BPF_CORE_READ_INTO(ev->dst_addr, sk, __sk_common.skc_v6_daddr);

    ev->ts_ns = now;
    bpf_ringbuf_submit(ev, 0);
    return 0;
}

/* =================================================================== *
 * 3) UDP sendmsg (IPv4 + IPv6)
 * =================================================================== *
 *
 * Firma: int udp_sendmsg(struct sock* sk, struct msghdr* msg, size_t len)
 *        int udpv6_sendmsg(struct sock* sk, struct msghdr* msg, size_t len)
 *
 * Se il socket e' stato "connesso" con connect(), skc_daddr/skc_dport
 * contengono la destinazione. Altrimenti la destinazione e' in
 * msg->msg_name (sockaddr_in/6). Gestiamo entrambi i casi.
 */

/* Estrae dst_ip e dst_port v4 dalla msghdr se presente.
 * Ritorna true se ha letto con successo. */
static __always_inline bool extract_msg_v4(struct msghdr* msg,
                                           __u32* out_daddr,
                                           __u16* out_dport)
{
    if (!msg) return false;
    void* name = BPF_CORE_READ(msg, msg_name);
    int   nlen = BPF_CORE_READ(msg, msg_namelen);
    if (!name || nlen < (int)sizeof(struct sockaddr_in)) return false;

    struct sockaddr_in sa = {};
    if (bpf_probe_read_kernel(&sa, sizeof(sa), name) != 0) return false;
    if (sa.sin_family != AF_INET) return false;

    *out_daddr = sa.sin_addr.s_addr;         /* network byte order */
    *out_dport = bpf_ntohs(sa.sin_port);
    return true;
}

static __always_inline bool extract_msg_v6(struct msghdr* msg,
                                           __u8 out_daddr[16],
                                           __u16* out_dport)
{
    if (!msg) return false;
    void* name = BPF_CORE_READ(msg, msg_name);
    int   nlen = BPF_CORE_READ(msg, msg_namelen);
    if (!name || nlen < (int)sizeof(struct sockaddr_in6)) return false;

    struct sockaddr_in6 sa = {};
    if (bpf_probe_read_kernel(&sa, sizeof(sa), name) != 0) return false;
    if (sa.sin6_family != AF_INET6) return false;

    __builtin_memcpy(out_daddr, &sa.sin6_addr, 16);
    *out_dport = bpf_ntohs(sa.sin6_port);
    return true;
}

SEC("kprobe/udp_sendmsg")
int BPF_KPROBE(bsc_udp_sendmsg, struct sock* sk, struct msghdr* msg)
{
    if (!sk) return 0;

    __u32 daddr = 0;
    __u16 dport = 0;

    /* 1) prova dalla msghdr (UDP unconnected) */
    if (!extract_msg_v4(msg, &daddr, &dport)) {
        /* 2) fallback: dalla struct sock (UDP connected) */
        daddr = BPF_CORE_READ(sk, __sk_common.skc_daddr);
        __u16 d = BPF_CORE_READ(sk, __sk_common.skc_dport);
        dport   = bpf_ntohs(d);
    }
    if (dport == 0) return 0;  /* socket senza destinazione risolta */

    const __u64 now = bpf_ktime_get_ns();
    __u32 saddr = BPF_CORE_READ(sk, __sk_common.skc_rcv_saddr);
    __u16 sport = BPF_CORE_READ(sk, __sk_common.skc_num);

    struct dedup_key k = {};
    const __u64 pid_tgid = bpf_get_current_pid_tgid();
    k.pid      = (__u32)(pid_tgid >> 32);
    k.dst_port = dport;
    k.family   = AF_INET;
    k.proto    = IPPROTO_UDP;
    __builtin_memcpy(k.dst_addr, &daddr, 4);
    if (!should_emit(&k, now, UDP_DEDUP_WINDOW_NS)) return 0;

    struct bsc_conn_event *ev = bpf_ringbuf_reserve(&bsc_events,
                                                    sizeof(*ev), 0);
    if (!ev) return 0;

    fill_proc(ev);
    ev->family     = AF_INET;
    ev->proto      = IPPROTO_UDP;
    ev->event_kind = 2;
    ev->_pad[0]    = 0;
    ev->src_port   = sport;
    ev->dst_port   = dport;

    __builtin_memset(ev->src_addr, 0, sizeof(ev->src_addr));
    __builtin_memset(ev->dst_addr, 0, sizeof(ev->dst_addr));
    __builtin_memcpy(ev->src_addr, &saddr, 4);
    __builtin_memcpy(ev->dst_addr, &daddr, 4);

    ev->ts_ns = now;
    bpf_ringbuf_submit(ev, 0);
    return 0;
}

SEC("kprobe/udpv6_sendmsg")
int BPF_KPROBE(bsc_udpv6_sendmsg, struct sock* sk, struct msghdr* msg)
{
    if (!sk) return 0;

    __u8  daddr[16] = {};
    __u16 dport = 0;

    if (!extract_msg_v6(msg, daddr, &dport)) {
        /* fallback: socket connected */
        BPF_CORE_READ_INTO(daddr, sk, __sk_common.skc_v6_daddr);
        __u16 d = BPF_CORE_READ(sk, __sk_common.skc_dport);
        dport   = bpf_ntohs(d);
    }
    if (dport == 0) return 0;

    const __u64 now = bpf_ktime_get_ns();
    __u16 sport = BPF_CORE_READ(sk, __sk_common.skc_num);

    struct dedup_key k = {};
    const __u64 pid_tgid = bpf_get_current_pid_tgid();
    k.pid      = (__u32)(pid_tgid >> 32);
    k.dst_port = dport;
    k.family   = AF_INET6;
    k.proto    = IPPROTO_UDP;
    __builtin_memcpy(k.dst_addr, daddr, 16);
    if (!should_emit(&k, now, UDP_DEDUP_WINDOW_NS)) return 0;

    struct bsc_conn_event *ev = bpf_ringbuf_reserve(&bsc_events,
                                                    sizeof(*ev), 0);
    if (!ev) return 0;

    fill_proc(ev);
    ev->family     = AF_INET6;
    ev->proto      = IPPROTO_UDP;
    ev->event_kind = 2;
    ev->_pad[0]    = 0;
    ev->src_port   = sport;
    ev->dst_port   = dport;

    BPF_CORE_READ_INTO(ev->src_addr, sk, __sk_common.skc_v6_rcv_saddr);
    __builtin_memcpy(ev->dst_addr, daddr, 16);

    ev->ts_ns = now;
    bpf_ringbuf_submit(ev, 0);
    return 0;
}

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