xdp/xdp_wg_mtu_optimizer.c

#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>

// 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");

// Função para atualizar estatísticas
static __always_inline void update_stats(__u64 *stat_counter) {
    __u32 key = 0;
    struct mtu_stats *stats = bpf_map_lookup_elem(&mtu_stats_map, &key);
    if (stats) {
        stats->total_packets++;
        (*stat_counter)++;
    }
}

// Função para ajustar TCP MSS
static __always_inline int adjust_tcp_mss(struct tcphdr *tcp, void *data_end, __u16 max_mss) {
    void *tcp_end = (void *)tcp + sizeof(struct tcphdr);
    if (tcp_end > data_end)
        return 0;
    
    if (!tcp->syn)
        return 0; // Ajustamos MSS apenas em pacotes SYN
    
    __u8 tcp_hdr_len = tcp->doff * 4;
    __u8 *tcp_opt_ptr = (__u8 *)tcp + sizeof(struct tcphdr);
    __u8 *tcp_opt_end = (__u8 *)tcp + tcp_hdr_len;
    
    if (tcp_opt_end > data_end)
        return 0;
    
    // Procurar opção MSS (kind=2, len=4) nas opções TCP
    #pragma unroll
    for (int i = 0; i < 40; i += tcp_opt_ptr[1]) { // máximo de 40 bytes de opções
        if (tcp_opt_ptr + 2 > tcp_opt_end || tcp_opt_ptr + 2 > data_end)
            break;
        
        if (tcp_opt_ptr[0] == 2 && tcp_opt_ptr[1] == 4) { // MSS option
            if (tcp_opt_ptr + 4 > tcp_opt_end || tcp_opt_ptr + 4 > data_end)
                break;
            
            // Extrair valor MSS atual (network byte order)
            __u16 *mss_ptr = (__u16 *)(tcp_opt_ptr + 2);
            __u16 current_mss = bpf_ntohs(*mss_ptr);
            
            // Ajustar MSS se necessário
            if (current_mss > max_mss) {
                *mss_ptr = bpf_htons(max_mss);
                update_stats(&((struct mtu_stats *)bpf_map_lookup_elem(&mtu_stats_map, &(__u32){0}))->mss_adjusted);
                return 1;
            }
            break;
        }
        
        if (tcp_opt_ptr[0] == 0) // End of options
            break;
        
        if (tcp_opt_ptr[1] < 2) // Opções malformadas, prevenir loop infinito
            break;
        
        tcp_opt_ptr += tcp_opt_ptr[1];
    }
    
    return 0;
}

// Função para enviar ICMP "Fragmentation needed" para indicar Path MTU Discovery
static __always_inline int send_icmp_frag_needed(struct xdp_md *ctx, struct ethhdr *eth,
                                               struct iphdr *iph, __u16 mtu) {
    // Enviar "Pacote muito grande" / "Fragmentação necessária mas DF está ativo"
    // Aqui sinalizamos ao kernel para enviar o ICMP, já que é difícil construí-lo no XDP
    
    // Verificar se o pacote tem DF (Don't Fragment) ativo
    __u16 frag_off = bpf_ntohs(iph->frag_off);
    if ((frag_off & 0x4000) == 0) // Sem bit DF, pode fragmentar
        return 0;
    
    // Registrar MTU conhecido para este destino
    __u32 dest_ip = iph->daddr;
    __u16 recommended_mtu = MAX_WG_MTU;
    bpf_map_update_elem(&path_mtu_map, &dest_ip, &recommended_mtu, BPF_ANY);
    
    // Incrementar contador de PMTU
    update_stats(&((struct mtu_stats *)bpf_map_lookup_elem(&mtu_stats_map, &(__u32){0}))->pmtu_triggered);
    
    // Sinalizar ao kernel para dropar e enviar ICMP
    return XDP_ABORTED;
}

// Programa XDP principal para otimização de pacotes WireGuard
SEC("xdp")
int xdp_wg_optimizer(struct xdp_md *ctx) {
    void *data_end = (void *)(long)ctx->data_end;
    void *data = (void *)(long)ctx->data;
    struct ethhdr *eth;
    struct iphdr *iph;
    __u32 key = 0;
    
    // Inicializar contadores
    struct mtu_stats *stats = bpf_map_lookup_elem(&mtu_stats_map, &key);
    if (stats) {
        stats->total_packets++;
    }
    
    // Validação básica dos cabeçalhos
    eth = data;
    if ((void*)(eth + 1) > data_end)
        return XDP_PASS;
        
    // Processamos apenas tráfego IPv4
    if (eth->h_proto != bpf_htons(ETH_P_IP))
        return XDP_PASS;
        
    iph = (void*)(eth + 1);
    if ((void*)(iph + 1) > data_end)
        return XDP_PASS;
    
    __u16 tot_len = bpf_ntohs(iph->tot_len);
    __u16 ip_hdr_len = iph->ihl * 4;
    
    // 1. Verificação de MTU - pacote maior que o MTU do WireGuard?
    if (tot_len > MAX_WG_MTU) {
        update_stats(&((struct mtu_stats *)bpf_map_lookup_elem(&mtu_stats_map, &key))->oversized_packets);
        
        // Verificar se podemos fragmentar (sem flag DF) ou se precisamos sinalizar PMTUD
        __u16 frag_off = bpf_ntohs(iph->frag_off);
        
        if (frag_off & 0x4000) { // Don't Fragment está ativo
            return send_icmp_frag_needed(ctx, eth, iph, MAX_WG_MTU);
        } else {
            // Pacote pode ser fragmentado, deixar para o kernel
            update_stats(&((struct mtu_stats *)bpf_map_lookup_elem(&mtu_stats_map, &key))->fragmented);
            return XDP_PASS;
        }
    }
    
    // 2. Otimização TCP MSS para sessões TCP novas
    if (iph->protocol == IPPROTO_TCP) {
        struct tcphdr *tcp = (void*)(iph) + ip_hdr_len;
        if ((void*)(tcp + 1) <= data_end) {
            if (adjust_tcp_mss(tcp, data_end, MSS_CLAMPING_VALUE)) {
                // MSS foi ajustado
            }
        }
    }
    
    // 3. Para pacotes destinados a hosts com MTU conhecidos, verificar se precisa sinalizar PMTU
    __u32 dest_ip = iph->daddr;
    __u16 *known_mtu = bpf_map_lookup_elem(&path_mtu_map, &dest_ip);
    
    if (known_mtu && tot_len > *known_mtu) {
        // Pacote excede MTU já conhecido para este destino
        if (iph->frag_off & bpf_htons(0x4000)) { // DF bit set
            return send_icmp_frag_needed(ctx, eth, iph, *known_mtu);
        }
    }
    
    // O pacote está dentro dos limites, deixar prosseguir
    return XDP_PASS;
}

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