/* ============================================================================
 * uprobe_ssl.c — Uprobes SSL_read/SSL_set_fd et kprobes accept4
 *
 * Intercepte les appels OpenSSL pour capturer le trafic déchiffré,
 * et corrige l'association socket ↔ SSL* via accept4.
 * ============================================================================ */

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

/* Taille maximale de données SSL à copier par événement */
#define MAX_SSL_DATA 4096

/* ---------------------------------------------------------------------------
 * Map temporaire : pid_tgid → upeer_sockaddr (sauvegardé à l'entrée d'accept4)
 * ---------------------------------------------------------------------------*/
struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 10240);
    __type(key, __u64);
    __type(value, __u64);   /* pointeur userspace vers sockaddr_in */
} accept_args_map SEC(".maps");

/* ===========================================================================
 * uprobe_ssl_set_fd — Intercept SSL_set_fd(SSL *s, int fd)
 *
 * Associe un ssl_ptr à ses informations de connexion (fd, src_ip, src_port)
 * en consultant fd_conn_map.
 * ===========================================================================*/
SEC("uprobe/SSL_set_fd")
int uprobe_ssl_set_fd(struct pt_regs *ctx)
{
    __u64 ssl_ptr = ((__u64)PT_REGS_PARM1(ctx));
    __u32 fd      = ((__u32)PT_REGS_PARM2(ctx));

    /* Rechercher les infos de connexion via le fd */
    struct ssl_conn_info *conn = bpf_map_lookup_elem(&fd_conn_map, &fd);
    if (!conn)
        return 0;

    /* Enregistrer l'association ssl_ptr → conn_info */
    struct ssl_conn_info new_conn = *conn;
    new_conn.fd = fd;
    bpf_map_update_elem(&ssl_conn_map, &ssl_ptr, &new_conn, BPF_ANY);

    return 0;
}

/* ===========================================================================
 * uprobe_ssl_read_entry — Entrée de SSL_read(SSL *ssl, void *buf, int num)
 *
 * Sauvegarde les arguments pour l'uretprobe correspondant.
 * ===========================================================================*/
SEC("uprobe/SSL_read")
int uprobe_ssl_read_entry(struct pt_regs *ctx)
{
    __u64 pid_tgid = bpf_get_current_pid_tgid();

    struct ssl_read_args args = {};
    args.ssl_ptr = (__u64)PT_REGS_PARM1(ctx);
    args.buf_ptr = (__u64)PT_REGS_PARM2(ctx);
    args.num     = (__u32)PT_REGS_PARM3(ctx);

    bpf_map_update_elem(&ssl_args_map, &pid_tgid, &args, BPF_ANY);
    return 0;
}

/* ===========================================================================
 * uretprobe_ssl_read_exit — Retour de SSL_read
 *
 * Lit le buffer déchiffré et l'émet dans rb_ssl_data.
 * ===========================================================================*/
SEC("uretprobe/SSL_read")
int uretprobe_ssl_read_exit(struct pt_regs *ctx)
{
    __u64 pid_tgid = bpf_get_current_pid_tgid();

    /* Récupérer les arguments sauvegardés à l'entrée */
    struct ssl_read_args *args = bpf_map_lookup_elem(&ssl_args_map, &pid_tgid);
    if (!args)
        return 0;

    /* Vérifier que la lecture a réussi (valeur de retour > 0) */
    long retval = PT_REGS_RC(ctx);
    if (retval <= 0) {
        bpf_map_delete_elem(&ssl_args_map, &pid_tgid);
        return 0;
    }

    /* Allouer un slot dans le ring buffer */
    struct ssl_data_event *evt = bpf_ringbuf_reserve(&rb_ssl_data, sizeof(*evt), 0);
    if (!evt) {
        bpf_map_delete_elem(&ssl_args_map, &pid_tgid);
        return 0;
    }

    evt->pid_tgid    = pid_tgid;
    evt->direction   = 0;  /* lecture = client vers serveur */
    evt->timestamp_ns = bpf_ktime_get_ns();

    /* Limiter la copie à MAX_SSL_DATA octets */
    __u32 data_len = (retval > MAX_SSL_DATA) ? MAX_SSL_DATA : (__u32)retval;
    evt->data_len = data_len;

    /* Copier depuis l'espace utilisateur */
    bpf_probe_read_user(evt->data, data_len & (MAX_SSL_DATA - 1), (void *)args->buf_ptr);

    /* Retrouver les infos de connexion via ssl_ptr */
    struct ssl_conn_info *conn = bpf_map_lookup_elem(&ssl_conn_map, &args->ssl_ptr);
    if (conn) {
        evt->fd       = conn->fd;
        evt->src_ip   = conn->src_ip;
        evt->src_port = conn->src_port;
    } else {
        evt->fd       = 0;
        evt->src_ip   = 0;
        evt->src_port = 0;
    }

    bpf_ringbuf_submit(evt, 0);
    bpf_map_delete_elem(&ssl_args_map, &pid_tgid);

    return 0;
}

/* ===========================================================================
 * kprobe_accept4_entry — Entrée de accept4(fd, upeer_sockaddr, upeer_addrlen, flags)
 *
 * Sauvegarde le pointeur vers la sockaddr pour la récupérer à la sortie.
 * ===========================================================================*/
SEC("kprobe/accept4")
int kprobe_accept4_entry(struct pt_regs *ctx)
{
    __u64 pid_tgid     = bpf_get_current_pid_tgid();
    /* Deuxième argument : pointeur userspace vers struct sockaddr_in */
    __u64 sockaddr_ptr = (__u64)PT_REGS_PARM2(ctx);

    bpf_map_update_elem(&accept_args_map, &pid_tgid, &sockaddr_ptr, BPF_ANY);
    return 0;
}

/* ===========================================================================
 * kretprobe_accept4_exit — Retour de accept4
 *
 * Lit la sockaddr_in pour extraire src_ip:src_port du client,
 * peuple accept_map et fd_conn_map, et émet dans rb_accept.
 * ===========================================================================*/
SEC("kretprobe/accept4")
int kretprobe_accept4_exit(struct pt_regs *ctx)
{
    __u64 pid_tgid = bpf_get_current_pid_tgid();

    /* Vérifier que accept4 a réussi (fd ≥ 0) */
    long new_fd = PT_REGS_RC(ctx);
    if (new_fd < 0) {
        bpf_map_delete_elem(&accept_args_map, &pid_tgid);
        return 0;
    }

    /* Récupérer le pointeur vers sockaddr_in */
    __u64 *sockaddr_ptr_p = bpf_map_lookup_elem(&accept_args_map, &pid_tgid);
    if (!sockaddr_ptr_p) {
        return 0;
    }
    __u64 sockaddr_ptr = *sockaddr_ptr_p;
    bpf_map_delete_elem(&accept_args_map, &pid_tgid);

    if (!sockaddr_ptr)
        return 0;

    /* Lire la structure sockaddr_in depuis l'espace utilisateur */
    /* struct sockaddr_in: sin_family(2) + sin_port(2) + sin_addr(4) */
    __u8 sa_buf[8] = {};
    bpf_probe_read_user(sa_buf, sizeof(sa_buf), (void *)sockaddr_ptr);

    /* Extraire port (octets 2-3) et adresse IP (octets 4-7) */
    __u16 sin_port = (__u16)(sa_buf[2] << 8) | sa_buf[3];   /* network byte order */
    __u32 sin_addr = *(__u32 *)(sa_buf + 4);                 /* network byte order */

    __u32 src_ip   = __builtin_bswap32(sin_addr);            /* host byte order */
    __u16 src_port = __builtin_bswap16(sin_port);            /* host byte order */
    __u32 fd       = (__u32)new_fd;

    /* Peupler accept_map[{pid_tgid, fd}] */
    struct accept_key akey = { .pid_tgid = pid_tgid, .fd = fd };
    struct accept_event aevt = {
        .pid_tgid    = pid_tgid,
        .fd          = fd,
        .src_ip      = src_ip,
        .src_port    = src_port,
        .timestamp_ns = bpf_ktime_get_ns(),
    };
    bpf_map_update_elem(&accept_map, &akey, &aevt, BPF_ANY);

    /* Peupler fd_conn_map[fd] pour accès rapide par SSL_set_fd */
    struct ssl_conn_info conn_info = {
        .fd       = fd,
        .src_ip   = src_ip,
        .src_port = src_port,
    };
    bpf_map_update_elem(&fd_conn_map, &fd, &conn_info, BPF_ANY);

    /* Émettre dans rb_accept */
    struct accept_event *out = bpf_ringbuf_reserve(&rb_accept, sizeof(*out), 0);
    if (!out)
        return 0;

    out->pid_tgid    = pid_tgid;
    out->fd          = fd;
    out->src_ip      = src_ip;
    out->src_port    = src_port;
    out->timestamp_ns = aevt.timestamp_ns;

    bpf_ringbuf_submit(out, 0);

    return 0;
}

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