// SPDX-License-Identifier: GPL-2.0
//
// sockmap.bpf.c — Programa eBPF para acelerar comunicação entre containers Docker
//
// Este programa faz duas coisas:
//   1) sock_ops: intercepta conexões TCP sendo estabelecidas e registra
//      os sockets num mapa (sockhash).
//   2) sk_msg:   quando um socket registrado envia dados, redireciona
//      diretamente para o socket do outro lado, PULANDO toda a pilha
//      de rede do kernel (TCP/IP, netfilter, bridge, veth, etc).
//
// Resultado: comunicação entre containers no MESMO host fica muito mais rápida
// porque os dados vão direto de um socket para o outro na memória.

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

// ─────────────────────────────────────────────────────────
// Chave que identifica uma conexão TCP (4-tupla + família)
// ─────────────────────────────────────────────────────────
struct sock_key {
    __u32 sip4;     // IP de origem (source IP)
    __u32 dip4;     // IP de destino (destination IP)
    __u32 sport;    // Porta de origem (source port)
    __u32 dport;    // Porta de destino (destination port)
    __u32 family;   // Família do socket (AF_INET = 2)
};

// ─────────────────────────────────────────────────────────
// Mapa SOCKHASH: armazena sockets indexados pela chave acima
// Capacidade: 65535 entradas (conexões simultâneas)
// ─────────────────────────────────────────────────────────
struct {
    __uint(type, BPF_MAP_TYPE_SOCKHASH);
    __uint(max_entries, 65535);
    __type(key, struct sock_key);
    __type(value, int);
} sock_hash SEC(".maps");

// ─────────────────────────────────────────────────────────
// Mapa para estatísticas
// Index 0 = conexões registradas no sockhash
// Index 1 = mensagens redirecionadas com sucesso (caminho acelerado)
// Index 2 = fallback: par não encontrado (caminho normal)
// Index 3 = conexões ignoradas (IPs externos/filtrados)
// ─────────────────────────────────────────────────────────
struct {
    __uint(type, BPF_MAP_TYPE_ARRAY);
    __uint(max_entries, 4);
    __type(key, __u32);
    __type(value, __u64);
} stats SEC(".maps");

// Incrementa um contador de estatísticas
static __always_inline void stats_inc(__u32 idx)
{
    __u64 *val = bpf_map_lookup_elem(&stats, &idx);
    if (val)
        __sync_fetch_and_add(val, 1);
}

#define SOCKMAP_AF_INET  2
#define SOCKMAP_AF_INET6 10

// ─────────────────────────────────────────────────────────
// BUG ENCONTRADO E CORRIGIDO: o filtro "só IPv4" (family != AF_INET)
// descartava toda conexão aceita por um servidor que escuta em dual-stack
// (bind em "::", padrão do Apache/muitos daemons quando não força IPv4).
// O socket ACEITO aparece como AF_INET6 pro kernel/eBPF mesmo a conexão
// sendo, na prática, IPv4 — o endereço vem mapeado como ::ffff:a.b.c.d.
//
// Evidência que confirmou isso: em produção, o Nextcloud (Apache
// dual-stack) tinha 0 de 184 conexões aceitas registradas no sockhash —
// o lado ATIVO dele (quando ELE disca pra MariaDB/Redis com IP literal,
// então family=AF_INET) funcionava 375/375, mas o lado PASSIVO (aceitando
// conexão do nginx) nunca registrava, porque a mesma conexão TCP aparece
// como AF_INET6 do lado de quem aceitou. Isso deixava o socket do nginx
// "meio-registrado" (só o lado dele no mapa), redirect sempre em
// fallback, e é o cenário suspeito por trás dos travamentos em
// transferências grandes.
//
// Verifica se um endereço IPv6 é, na verdade, um IPv4 mapeado
// (::ffff:a.b.c.d, RFC 4291 §2.5.5.2: 80 bits zero + 16 bits 0xffff +
// 32 bits do IPv4) e extrai o IPv4 equivalente.
//
// Recebe os 4 words POR VALOR (w0..w3), não um ponteiro/array — passar
// o endereço de um campo array do ctx (ex: &ops->local_ip6[0]) pra uma
// função separada é "aritmética de ponteiro modificado no ctx", que o
// verificador do BPF rejeita (bpf_sock_ops/sk_msg_md só permitem acesso
// direto a campos em offsets conhecidos, reescrito pelo compilador;
// tirar o endereço de um array de campo quebra esse reconhecimento).
// Cada chamador lê os 4 elementos diretamente (ops->local_ip6[0] etc)
// e passa os valores, não o array.
static __always_inline bool ipv4_mapped(__u32 w0, __u32 w1, __u32 w2, __u32 w3, __u32 *out_ip4)
{
    if (w0 != 0 || w1 != 0)
        return false;
    if (bpf_ntohl(w2) != 0x0000ffff)
        return false;
    *out_ip4 = w3;
    return true;
}

// ─────────────────────────────────────────────────────────
// Verifica se um IP (em network byte order / big-endian) é
// endereço privado RFC-1918 ou loopback.
//
// Por que isso importa:
//   O sock_ops roda para TODAS as conexões TCP do cgroup raiz,
//   incluindo tráfego externo (HTTPS do mundo, metadata AWS…).
//   Registrar IPs públicos no sockhash é inútil e perigoso:
//   o par nunca será encontrado, o sk_msg vai ao fallback e
//   entradas mortas poluem o mapa.
//   Só registramos conexões onde ambos os lados são privados
//   (i.e., comunicação intra-host / intra-container).
// ─────────────────────────────────────────────────────────
static __always_inline bool is_local_ip(__u32 ip_be)
{
    // local_ip4 / remote_ip4 são big-endian (network byte order).
    // Em x86 (little-endian) o valor numérico da u32 aparece
    // byte-invertido: para 10.15.0.45 (0A 0F 00 2D em memória),
    // a u32 lida vale 0x2D000F0A — ou seja, ip_be >> 24 = 45 (ERRADO!).
    // bpf_ntohl() converte para host order, colocando o primeiro
    // octeto (10) no byte mais significativo, onde >> 24 extrai
    // corretamente 10.
    __u32 ip = bpf_ntohl(ip_be);
    __u8 o1  = (__u8)(ip >> 24);
    __u8 o2  = (__u8)(ip >> 16);

    if (o1 == 10)                              return true; // 10.0.0.0/8
    if (o1 == 172 && o2 >= 16 && o2 <= 31)    return true; // 172.16.0.0/12
    if (o1 == 192 && o2 == 168)               return true; // 192.168.0.0/16
    if (o1 == 127)                             return true; // 127.0.0.0/8 loopback
    return false;
}

// ─────────────────────────────────────────────────────────
// Extrai a chave de conexão a partir do contexto sock_ops
// (usado quando uma conexão TCP é estabelecida)
// ─────────────────────────────────────────────────────────
// BUG CORRIGIDO: o código antigo misturava dois domínios de codificação
// de porta sem converter os dois pro mesmo formato.
//   - ops->local_port  vem em HOST byte order (puro, ex: 6379)
//   - ops->remote_port vem com o valor em network byte order alojado nos
//     16 bits SUPERIORES de um campo de 32 bits — não é um u32 comum em
//     network order, então bpf_ntohl() sozinho faz um byte-swap completo
//     de 32 bits e produz um número sem relação com a porta real.
// A comparação `sport == local_port do outro lado` só funciona se AMBOS
// os valores estiverem no MESMO domínio. Em vez de tentar decifrar o
// valor "verdadeiro" da porta (o que exigiria confiar numa convenção que
// não dá pra validar sem testar em kernel real), a correção normaliza os
// dois pro domínio de remote_port: extrai remote_port com ">> 16" (sem
// ntohl) e converte local_port PRA esse mesmo domínio com
// "bpf_htonl(local_port) >> 16". Assim os dois lados de uma conexão
// calculam exatamente o mesmo valor pro mesmo par de portas, não importa
// qual seja a codificação exata do kernel — o que importa é que os dois
// membros do par usem a MESMA transformação simétrica.
// Retorna false se a conexão não puder ser reduzida a um par IPv4 (IPv6
// nativo de verdade, fora do escopo deste programa) — quem chama deve
// tratar isso como "não é pra registrar", igual ao antigo filtro de
// family, só que agora sem descartar dual-stack por engano.
static __always_inline bool
extract_key_from_ops(struct bpf_sock_ops *ops, struct sock_key *key)
{
    __u32 local_ip4, remote_ip4;

    if (ops->family == SOCKMAP_AF_INET) {
        local_ip4  = ops->local_ip4;
        remote_ip4 = ops->remote_ip4;
    } else if (ops->family == SOCKMAP_AF_INET6) {
        if (!ipv4_mapped(ops->local_ip6[0], ops->local_ip6[1], ops->local_ip6[2], ops->local_ip6[3], &local_ip4) ||
            !ipv4_mapped(ops->remote_ip6[0], ops->remote_ip6[1], ops->remote_ip6[2], ops->remote_ip6[3], &remote_ip4))
            return false;
    } else {
        return false;
    }

    key->sip4   = local_ip4;
    key->dip4   = remote_ip4;
    key->sport  = bpf_htonl(ops->local_port) >> 16;
    key->dport  = ops->remote_port >> 16;
    // family SEMPRE normalizado pra AF_INET: os dois lados da MESMA
    // conexão TCP podem reportar families diferentes entre si (um vê
    // AF_INET literal, o outro vê AF_INET6 mapeado por causa de bind
    // dual-stack) — sem normalizar, as chaves dos dois lados nunca
    // bateriam e o par nunca seria encontrado no sockhash. Foi
    // exatamente isso que deixava o Nextcloud (Apache em dual-stack)
    // com 0 registros passivos em 184 conexões aceitas.
    key->family = SOCKMAP_AF_INET;
    return true;
}

// ─────────────────────────────────────────────────────────
// PROGRAMA 1: sock_ops
//
// Anexado ao cgroup, é chamado em eventos de socket TCP.
// Quando uma conexão é ESTABELECIDA (ativa ou passiva),
// registramos o socket no sockhash.
// ─────────────────────────────────────────────────────────
SEC("sockops")
int bpf_sockops_handler(struct bpf_sock_ops *skops)
{
    __u32 op = skops->op;

    // Aceita AF_INET e AF_INET6 aqui (o filtro real de família acontece
    // dentro de extract_key_from_ops, que resolve IPv6 mapeado pra IPv4
    // ou retorna false se for IPv6 de verdade). Cortar por family==AF_INET
    // logo no topo, como a versão antiga fazia, foi o que descartava
    // silenciosamente toda conexão aceita por um listener dual-stack.
    if (skops->family != SOCKMAP_AF_INET && skops->family != SOCKMAP_AF_INET6)
        return 0;

    switch (op) {
    // Conexão passiva: servidor aceitou uma conexão de um cliente
    case BPF_SOCK_OPS_PASSIVE_ESTABLISHED_CB:
    // Conexão ativa: cliente conectou a um servidor
    case BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB: {
        struct sock_key key = {};
        if (!extract_key_from_ops(skops, &key)) {
            stats_inc(3); // IPv6 de verdade (não mapeado) — fora do escopo
            return 0;
        }

        // ── Filtro: só conexões entre IPs privados RFC-1918 ───────
        // Ignora tráfego externo (HTTPS público, metadata AWS 169.254.x.x).
        if (!is_local_ip(key.sip4) || !is_local_ip(key.dip4)) {
            stats_inc(3); // conta conexões externas ignoradas
            return 0;
        }

        if (skops->family == SOCKMAP_AF_INET6)
            bpf_printk("REG via IPv6 mapeado: op=%u sip=%x dip=%x", op, key.sip4, key.dip4);

        // BPF_ANY: atualiza mesmo se já existir (cobre reconexões keepalive)
        int ret = bpf_sock_hash_update(skops, &sock_hash, &key, BPF_ANY);
        if (ret == 0) {
            stats_inc(0);

            // ── Hints de performance TCP ──────────────────────────
            // Ativa todos os callbacks relevantes de uma vez:
            //   STATE_CB_FLAG  : notifica mudanças de estado TCP
            //                    (necessário para limpar sockets fechados)
            //   RTT_CB_FLAG    : callback a cada RTT (permite ajuste dinâmico)
            __u32 flags = skops->bpf_sock_ops_cb_flags
                          | BPF_SOCK_OPS_STATE_CB_FLAG
                          | BPF_SOCK_OPS_RTT_CB_FLAG;
            bpf_sock_ops_cb_flags_set(skops, flags);
        }
        break;
    }

    // Mudança de estado TCP: limpa o mapa em qualquer transição
    // que indique encerramento, para não deixar sockets mortos.
    case BPF_SOCK_OPS_STATE_CB: {
        __u32 new_state = skops->args[1];
        if (new_state == BPF_TCP_CLOSE      ||
            new_state == BPF_TCP_CLOSE_WAIT ||
            new_state == BPF_TCP_FIN_WAIT1  ||
            new_state == BPF_TCP_FIN_WAIT2  ||
            new_state == BPF_TCP_TIME_WAIT) {
            struct sock_key key = {};
            if (extract_key_from_ops(skops, &key))
                bpf_map_delete_elem(&sock_hash, &key);
        }
        break;
    }

    // RTT callback: pode ser usado no futuro para métricas por conexão.
    // Por ora não fazemos nada aqui — o callback foi ativado para
    // manter o kernel notificando o sock_ops sobre a saúde da conexão.
    case BPF_SOCK_OPS_RTT_CB:
        break;

    default:
        break;
    }

    return 0;
}

// ─────────────────────────────────────────────────────────
// PROGRAMA 2: sk_msg
//
// Anexado ao sockhash, é chamado quando qualquer socket
// registrado envia dados (sendmsg/sendfile).
// Redireciona os dados diretamente para o socket par,
// pulando toda a pilha de rede.
// ─────────────────────────────────────────────────────────
SEC("sk_msg")
int bpf_msg_redirect(struct sk_msg_md *msg)
{
    // Construímos a chave do socket PAR (invertendo src/dst).
    // Se A→B enviou dados, queremos encontrar no mapa a entrada
    // que B registrou: {sip=B, dip=A, sport=portaB, dport=portaA}
    //
    // Mesma normalização de porta usada em extract_key_from_ops (mesmo
    // domínio dos dois lados): remote_port só com ">> 16", local_port
    // convertido pra esse domínio via "bpf_htonl(...) >> 16". Mesma
    // resolução de IPv6-mapeado-pra-IPv4 também — precisa bater com o
    // que extract_key_from_ops grava no REGISTRO, senão o par nunca é
    // encontrado (era exatamente o bug do Nextcloud dual-stack).
    __u32 sip4, dip4;
    if (msg->family == SOCKMAP_AF_INET) {
        sip4 = msg->remote_ip4;
        dip4 = msg->local_ip4;
    } else if (msg->family == SOCKMAP_AF_INET6) {
        if (!ipv4_mapped(msg->remote_ip6[0], msg->remote_ip6[1], msg->remote_ip6[2], msg->remote_ip6[3], &sip4) ||
            !ipv4_mapped(msg->local_ip6[0], msg->local_ip6[1], msg->local_ip6[2], msg->local_ip6[3], &dip4)) {
            stats_inc(2); // IPv6 de verdade — nunca deveria estar no mapa, mas por segurança cai no fallback
            return SK_PASS;
        }
    } else {
        stats_inc(2);
        return SK_PASS;
    }

    struct sock_key key = {};
    key.sip4   = sip4;                            // IP remoto = IP do par
    key.dip4   = dip4;                            // IP local  = nosso IP
    key.sport  = msg->remote_port >> 16;
    key.dport  = bpf_htonl(msg->local_port) >> 16;
    key.family = SOCKMAP_AF_INET;                 // normalizado, ver extract_key_from_ops

    // Tenta redirecionar diretamente para o socket par.
    int ret = bpf_msg_redirect_hash(msg, &sock_hash, &key, BPF_F_INGRESS);
    if (ret == SK_PASS) {
        stats_inc(1);  // Redirecionamento OK — caminho acelerado
        return SK_PASS;
    }

    // FALLBACK CRÍTICO: se o par não está no mapa (socket ainda não
    // registrado, já fechou, ou conexão externa ao cgroup), deixamos
    // os dados fluírem pelo caminho normal de rede em vez de
    // descartar (SK_DROP), o que causaria ERR_CONNECTION_CLOSED.
    stats_inc(2);
    return SK_PASS;
}

// Licença GPL obrigatória para usar helpers BPF
char _license[] SEC("license") = "GPL";