xdp/xdp_accelerated_nat.c

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

// Definir arquitetura ARM64 explicitamente
#ifndef __TARGET_ARCH_arm64
#define __TARGET_ARCH_arm64
#endif

// Definições para contornar erros de compilação
#ifndef __section
#define __section(NAME) \
   __attribute__((section(NAME), used))
#endif

// Substituir o SEC original pela versão compatível com ARM64
#undef SEC
#define SEC(NAME) __section(NAME)

/* Estrutura para conexões NAT */
struct nat_conn {
    __u32 orig_src_ip;    // IP de origem original
    __u16 orig_src_port;  // Porta de origem original
    __u32 orig_dst_ip;    // IP de destino original
    __u16 orig_dst_port;  // Porta de destino original
    __u32 new_src_ip;     // Novo IP de origem (para SNAT)
    __u16 new_src_port;   // Nova porta de origem (para SNAT)
    __u32 new_dst_ip;     // Novo IP de destino (para DNAT)
    __u16 new_dst_port;   // Nova porta de destino (para DNAT)
    __u8 proto;           // Protocolo (TCP/UDP)
    __u8 state;           // Estado da conexão (1=ativo)
    __u64 last_update;    // Timestamp da última atualização
};

/* Chave para consulta rápida de DNAT em pacotes de entrada */
struct dnat_key {
    __u32 dst_ip;
    __u16 dst_port;
    __u8 proto;
    __u8 padding;         // Adicionado para alinhamento em ARM64
} __attribute__((packed));

/* Chave para busca de conexões de saída */
struct conn_key {
    __u32 src_ip;
    __u16 src_port;
    __u32 dst_ip;
    __u16 dst_port;
    __u8 proto;
    __u8 padding;         // Adicionado para alinhamento em ARM64
} __attribute__((packed));

/* Mapas BPF para NAT */
struct {
    __uint(type, BPF_MAP_TYPE_HASH); // Alterado de LRU_HASH para HASH para maior compatibilidade
    __uint(max_entries, 65536);
    __type(key, struct conn_key);
    __type(value, struct nat_conn);
} conn_track_map SEC(".maps");

/* Mapa para regras DNAT estáticas */
struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 1024);
    __type(key, struct dnat_key);
    __type(value, struct nat_conn);
} dnat_rules SEC(".maps");

/* Mapa para redes que devem ter SNAT aplicado */
struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 256);
    __type(key, __u32);   // Prefixo de rede (/16)
    __type(value, __u8);  // Flag (1=ativo)
} snat_networks SEC(".maps");

/* Mapa para estatísticas */
struct nat_stats {
    __u64 processed;   // Pacotes processados
    __u64 snat_hits;   // Acertos de SNAT
    __u64 dnat_hits;   // Acertos de DNAT
    __u64 new_conns;   // Novas conexões
    __u64 errors;      // Erros
};

struct {
    __uint(type, BPF_MAP_TYPE_ARRAY);
    __uint(max_entries, 1);
    __type(key, __u32);
    __type(value, struct nat_stats);
} stats_map SEC(".maps");

/* Configuração para IP do gateway (usado para SNAT) */
struct {
    __uint(type, BPF_MAP_TYPE_ARRAY);
    __uint(max_entries, 1);
    __type(key, __u32);
    __type(value, __u32);
} gateway_ip_map SEC(".maps");

/* Função para calcular IPv4 checksum - simplificada para evitar loops */
static __always_inline __u16 ip_checksum(struct iphdr *ip) {
    __u32 sum = 0;
    
    // Para ARM64, calcular o checksum sem depender de loops
    __u16 *iph = (__u16 *)ip;
    
    // Zeramos o campo checksum do IP header
    ip->check = 0;
    
    // Calculamos a soma de palavras de 16 bits do header IP
    // O tamanho padrão do header IP é 20 bytes (5 palavras de 32 bits ou 10 de 16 bits)
    sum += iph[0];
    sum += iph[1];
    sum += iph[2];
    sum += iph[3];
    sum += iph[4];
    sum += iph[5];
    sum += iph[6];
    sum += iph[7];
    sum += iph[8];
    sum += iph[9];
    
    // Dobrar 32-bit para 16-bit para obter o checksum final
    sum = (sum & 0xFFFF) + (sum >> 16);
    sum = (sum & 0xFFFF) + (sum >> 16);
    
    return ~sum; // Complemento de 1
}

/* Função simplificada para atualizar checksum */
static __always_inline __u16 update_checksum(__u16 old_checksum, __u32 old_value, __u32 new_value) {
    __u32 sum = ~((uint32_t)old_checksum) & 0xFFFF;
    sum += ~(old_value >> 16) & 0xFFFF;
    sum += ~old_value & 0xFFFF;
    sum += (new_value >> 16) & 0xFFFF;
    sum += new_value & 0xFFFF;
    sum = (sum & 0xFFFF) + (sum >> 16);
    sum = (sum & 0xFFFF) + (sum >> 16);
    return ~sum & 0xFFFF;
}

/* Função para atualizar estatísticas */
static __always_inline void update_stats(__u32 stat_type) {
    __u32 key = 0;
    struct nat_stats *stats;
    
    stats = bpf_map_lookup_elem(&stats_map, &key);
    if (stats) {
        if (stat_type == 0) stats->processed++;
        else if (stat_type == 1) stats->snat_hits++;
        else if (stat_type == 2) stats->dnat_hits++;
        else if (stat_type == 3) stats->new_conns++;
        else if (stat_type == 4) stats->errors++;
    }
}

/* Função para verificar e aplicar SNAT para redes internas */
static __always_inline int should_apply_snat(__u32 src_ip) {
    // Extrair o prefixo /16
    __u32 prefix = src_ip & 0xFFFF0000;
    __u8 *flag;
    
    flag = bpf_map_lookup_elem(&snat_networks, &prefix);
    return (flag && *flag == 1) ? 1 : 0;
}

/* Programa principal XDP para aceleração de NAT */
SEC("xdp")
int xdp_accelerate_nat(struct xdp_md *ctx) {
    void *data = (void *)(long)ctx->data;
    void *data_end = (void *)(long)ctx->data_end;
    
    // Verificação básica de limites
    if (data + sizeof(struct ethhdr) > data_end)
        return XDP_PASS;
    
    struct ethhdr *eth = data;
    
    // Só processa IPv4
    if (eth->h_proto != bpf_htons(ETH_P_IP))
        return XDP_PASS;
    
    struct iphdr *ip = (void *)(eth + 1);
    if ((void *)(ip + 1) > data_end)
        return XDP_PASS;
    
    // Ignorar pacotes não-TCP/UDP por enquanto
    if (ip->protocol != IPPROTO_TCP && ip->protocol != IPPROTO_UDP)
        return XDP_PASS;
    
    // Incrementar contador de pacotes processados
    update_stats(0);
    
    // Restante do código simplificado para compatibilidade ARM64
    // No caso de XDP modo genérico, vamos deixar a maioria do processamento
    // NAT para o TC/iptables padrão para garantir compatibilidade
    
    // Aplicar SNAT básico para rede interna
    if (should_apply_snat(ip->saddr)) {
        update_stats(1);  // SNAT hit
        // Permitir que o TC/iptables faça o trabalho real de SNAT
        return XDP_PASS;
    }
    
    // Apenas passar pacotes normalmente
    return XDP_PASS;
}

/* Simplificação extrema para compatibilidade */
char _license[] SEC("license") = "GPL";