// eBPF program to count TCP retransmissions per flow
#include <uapi/linux/bpf.h>
#include <linux/tcp.h>
#include <linux/ip.h>
#include <linux/in.h>
#include <linux/inet.h>

struct flow_key_t {
    __be32 src_ip;
    __be32 dst_ip;
    __be16 src_port;
    __be16 dst_port;
};

BPF_HASH(rtx_count, struct flow_key_t, u64);

int count_tcp_retransmit(struct pt_regs *ctx) {
    struct sock *sk;
    struct flow_key_t flow_key = {};
    u64 *count, zero = 0;

    sk = (struct sock *)PT_REGS_PARM1(ctx);
    if (!sk)
        return 0;

    // Extract source and destination IPs and ports
    bpf_probe_read_kernel(&flow_key.src_ip, sizeof(flow_key.src_ip), &sk->sk_rcv_saddr);
    bpf_probe_read_kernel(&flow_key.dst_ip, sizeof(flow_key.dst_ip), &sk->sk_daddr);
    bpf_probe_read_kernel(&flow_key.src_port, sizeof(flow_key.src_port), &sk->sk_num);
    bpf_probe_read_kernel(&flow_key.dst_port, sizeof(flow_key.dst_port), &sk->sk_dport);

    // Swap byte order of destination port
    flow_key.dst_port = ntohs(flow_key.dst_port);

    // Update retransmission counter for this flow
    count = rtx_count.lookup_or_try_init(&flow_key, &zero);
    if (count) {
        (*count)++;
    }

    return 0;
}