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

// Configurações de MTU para WireGuard
#define WG_OVERHEAD 60         // Overhead do WireGuard (cabeçalhos + criptografia)
#define MAX_WG_MTU 1420        // MTU máximo recomendado para WireGuard (1480 - overhead)
#define ETH_DEFAULT_MTU 1500   // MTU padrão da ethernet
#define MSS_CLAMPING_VALUE (MAX_WG_MTU - 40)  // MSS = MTU - IP(20) - TCP(20)
#define MIN_MSS_VALUE 536      // Valor mínimo de MSS conforme RFC
#define ICMP_NEEDED_HEADROOM 8 // Espaço extra para cabeçalho ICMP
#define IPV4_HEADER_DF_OFFSET 6 // Offset do campo Flags no cabeçalho IP

// Estatísticas de pacotes processados
struct mtu_stats {
    __u64 total_packets;
    __u64 oversized_packets;
    __u64 mss_adjusted;
    __u64 fragmented;
    __u64 pmtu_triggered;
};

// Mapa para armazenar estatísticas
struct {
    __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
    __uint(max_entries, 1);
    __type(key, __u32);
    __type(value, struct mtu_stats);
} mtu_stats_map SEC(".maps");

// Mapa para MTUs conhecidos por destino
struct {
    __uint(type, BPF_MAP_TYPE_LRU_HASH);
    __uint(max_entries, 1024);
    __type(key, __u32);      // IP de destino
    __type(value, __u16);    // MTU conhecido
} path_mtu_map SEC(".maps");

// Atualização incremental de checksum de 16 bits (RFC 1071/1624) — a mesma
// técnica usada pelo csum_replace2() do kernel. Necessária porque XDP não
// tem acesso a helpers de checksum baseados em skb (bpf_l4_csum_replace só
// existe para programas TC/socket, não XDP).
static __always_inline __u16 csum16_replace(__u16 csum, __u16 old_val, __u16 new_val) {
    __u32 sum = (__u16)~csum + (__u16)~old_val + new_val;
    sum = (sum & 0xffff) + (sum >> 16);
    sum = (sum & 0xffff) + (sum >> 16);
    return (__u16)~sum;
}

// Clampa o MSS anunciado em pacotes SYN, reescrevendo o pacote de verdade
// (a versão anterior só incrementava uma estatística e admitia no comentário
// que não mexia no pacote — ou seja, não fazia nada útil). Como não mudamos
// o tamanho do pacote, dá para escrever em cima do buffer existente sem
// bpf_xdp_adjust_head, só recalculando o checksum do TCP.
static __always_inline int clamp_tcp_mss(struct tcphdr *tcp, void *data_end, __u16 max_mss_host) {
    if ((void *)(tcp + 1) > data_end)
        return 0;

    // Ajustamos MSS apenas em pacotes SYN puros (não SYN-ACK): este
    // programa roda no ingress de wg0, ou seja, só vê tráfego entrando
    // pelo túnel em direção à rede interna. O SYN-ACK de volta sai por
    // wg0 (egress), fora do alcance do XDP — cobrir essa direção também
    // exigiria um programa TC de egress, não só XDP.
    if (!tcp->syn || tcp->ack)
        return 0;

    __u8 tcp_hdr_len = tcp->doff * 4;
    if (tcp_hdr_len <= sizeof(struct tcphdr))
        return 0;  // Sem opções TCP

    __u8 *opt = (__u8 *)(tcp + 1);
    __u8 *opt_end = (__u8 *)tcp + tcp_hdr_len;
    if ((void *)opt_end > data_end)
        return 0;

    // Loop limitado (verificador do BPF exige bound estático): cabeçalho
    // TCP tem no máximo 40 bytes de opções, então 10 iterações cobrem até
    // opções de 4 bytes cada — mais que suficiente para chegar no MSS,
    // que por convenção vem logo no início das opções.
    #pragma unroll
    for (int i = 0; i < 10; i++) {
        if ((void *)(opt + 1) > data_end || opt >= opt_end)
            break;

        __u8 kind = opt[0];
        if (kind == 0)          // End of Option List
            break;
        if (kind == 1) {        // NOP
            opt += 1;
            continue;
        }

        if ((void *)(opt + 2) > data_end || opt + 1 >= opt_end)
            break;
        __u8 opt_len = opt[1];
        if (opt_len < 2 || (void *)(opt + opt_len) > data_end || opt + opt_len > opt_end)
            break;

        if (kind == 2 && opt_len == 4) {   // Maximum Segment Size
            if ((void *)(opt + 4) > data_end)
                break;
            __be16 *mss_ptr = (__be16 *)(opt + 2);
            __u16 cur_mss = bpf_ntohs(*mss_ptr);
            if (cur_mss > max_mss_host) {
                __be16 old_val = *mss_ptr;
                __be16 new_val = bpf_htons(max_mss_host);
                *mss_ptr = new_val;
                tcp->check = csum16_replace(tcp->check, old_val, new_val);
                return 1;
            }
            return 0;
        }

        opt += opt_len;
    }
    return 0;
}

// Programa XDP principal - simplificado
SEC("xdp")
int xdp_wg_optimizer(struct xdp_md *ctx) {
    void *data_end = (void *)(long)ctx->data_end;
    void *data = (void *)(long)ctx->data;
    __u32 key = 0;
    
    // Incrementar estatística de pacotes totais
    struct mtu_stats *stats = bpf_map_lookup_elem(&mtu_stats_map, &key);
    if (stats)
        stats->total_packets++;
    
    // Verificação básica de limites
    if (data + sizeof(struct ethhdr) > data_end)
        return XDP_PASS;
    
    struct ethhdr *eth = data;
    
    // Filtrar apenas pacotes IPv4
    if (eth->h_proto != bpf_htons(ETH_P_IP))
        return XDP_PASS;
    
    // Verificar se o pacote contém um cabeçalho IP completo
    if ((void *)(eth + 1) + sizeof(struct iphdr) > data_end)
        return XDP_PASS;
    
    struct iphdr *iph = (void *)(eth + 1);
    
    // Verificar cabeçalho IP
    if ((void *)(iph + 1) > data_end || iph->version != 4)
        return XDP_PASS;
    
    // Verificação de MTU
    __u16 tot_len = bpf_ntohs(iph->tot_len);
    if (tot_len > MAX_WG_MTU) {
        // Atualizar estatísticas
        if (stats)
            stats->oversized_packets++;
        
        // Verificar flag DF
        __u16 frag_off = bpf_ntohs(iph->frag_off);
        if (frag_off & 0x4000) {
            // Don't Fragment está ativo, registrar o MTU para este destino
            __u32 dest_ip = iph->daddr;
            __u16 mtu = MAX_WG_MTU;
            bpf_map_update_elem(&path_mtu_map, &dest_ip, &mtu, BPF_ANY);
            
            // Atualizar estatísticas
            if (stats)
                stats->pmtu_triggered++;
                
            return XDP_ABORTED; // Sinalizar para enviar ICMP Fragmentation Needed
        } else {
            // Pacote pode ser fragmentado
            if (stats)
                stats->fragmented++;
        }
    }
    
    // Ajustar MSS para pacotes TCP
    if (iph->protocol == IPPROTO_TCP) {
        __u8 ihl = iph->ihl * 4; // IP header length
        void *tcp_ptr = (void *)iph + ihl;
        
        // Verificar se temos espaço para o cabeçalho TCP
        if (tcp_ptr + sizeof(struct tcphdr) <= data_end) {
            struct tcphdr *tcp = tcp_ptr;
            if (clamp_tcp_mss(tcp, data_end, MSS_CLAMPING_VALUE) && stats)
                stats->mss_adjusted++;
        }
    }
    
    // Verificar MTUs conhecidos
    __u16 *known_mtu = bpf_map_lookup_elem(&path_mtu_map, &iph->daddr);
    if (known_mtu && tot_len > *known_mtu) {
        __u16 frag_off = bpf_ntohs(iph->frag_off);
        if (frag_off & 0x4000) {
            // Don't Fragment está ativo, não pode passar pelo tunnel
            if (stats)
                stats->pmtu_triggered++;
            return XDP_ABORTED;
        }
    }
    
    // Permitir o pacote
    return XDP_PASS;
}

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