// bbr_switch_510.bpf.c
// Debian 11 / Linux 5.10
// Logic:
// 1) Initial CC: HTCP
// 2) Observe RTT callbacks
// 3) After >=3s, if RTT>120ms pattern is stable (>=10 of >=12 samples), switch to BBR

#include <linux/bpf.h>
#include <linux/in.h>
#include <linux/tcp.h>
#include <bpf/bpf_helpers.h>

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

#ifndef TCP_CONGESTION
#define TCP_CONGESTION 13
#endif

#define STATE_PROBING   1
#define STATE_SWITCHED  2
#define STATE_HOLD      3

// Tunables
#define RTT_HIGH_US           120000U   // 120ms
#define MIN_SWITCH_DELAY_MS     3000U   // at least 3s before switch
#define EARLY_SAMPLE_NEED         12U
#define EARLY_HIGH_NEED           10U
#define FINAL_DEADLINE_MS      10000U   // stop probing after 10s

struct flow_state {
    __u64 start_ns;
    __u32 sample_cnt;
    __u32 high_rtt_cnt;
    __u8  state;
};

struct {
    __uint(type, BPF_MAP_TYPE_LRU_HASH);
    __uint(max_entries, 262144);
    __type(key, __u64);   // socket cookie
    __type(value, struct flow_state);
} flow_map SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
    __uint(max_entries, 6);
    __type(key, __u32);
    __type(value, __u64);
} stats SEC(".maps");
/*
0 established
1 switched_to_bbr
2 hold_timeout
3 setsockopt_fail
4 rtt_cb_hits
5 too_early_gate
*/

static __always_inline void inc_stat(__u32 idx)
{
    __u64 *v = bpf_map_lookup_elem(&stats, &idx);
    if (v)
        __sync_fetch_and_add(v, 1);
}

static __always_inline int set_cc(struct bpf_sock_ops *skops, const char *name, int len)
{
    return bpf_setsockopt(skops, IPPROTO_TCP, TCP_CONGESTION, (void *)name, len);
}

SEC("sockops")
int bbr_switch_sockops(struct bpf_sock_ops *skops)
{
    __u32 op = skops->op;
    __u64 cookie = bpf_get_socket_cookie(skops);
    __u64 now = bpf_ktime_get_ns();

    if (!cookie)
        return 0;

    if (op == BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB ||
        op == BPF_SOCK_OPS_PASSIVE_ESTABLISHED_CB) {
        struct flow_state st = {};
        const char htcp[] = "htcp";

        st.start_ns = now;
        st.state = STATE_PROBING;
        bpf_map_update_elem(&flow_map, &cookie, &st, BPF_ANY);
        inc_stat(0);

        // Initial CC = HTCP
        if (set_cc(skops, htcp, sizeof(htcp)) != 0)
            inc_stat(3);

        // Ask kernel to call us at RTT events
        bpf_sock_ops_cb_flags_set(skops, BPF_SOCK_OPS_RTT_CB_FLAG);
        return 0;
    }

    if (op == BPF_SOCK_OPS_RTT_CB) {
        struct flow_state *st = bpf_map_lookup_elem(&flow_map, &cookie);
        __u64 age_ms;
        __u32 srtt_us;

        inc_stat(4);

        if (!st || st->state != STATE_PROBING)
            return 0;

        age_ms = (now - st->start_ns) / 1000000ULL;
        if (age_ms > FINAL_DEADLINE_MS) {
            st->state = STATE_HOLD;
            inc_stat(2);
            return 0;
        }

        // On many kernels srtt_us is stored with 3-bit fraction.
        srtt_us = skops->srtt_us >> 3;

        st->sample_cnt++;
        if (srtt_us > RTT_HIGH_US)
            st->high_rtt_cnt++;

        if (age_ms < MIN_SWITCH_DELAY_MS) {
            inc_stat(5);
            return 0;
        }

        if (st->sample_cnt >= EARLY_SAMPLE_NEED &&
            st->high_rtt_cnt >= EARLY_HIGH_NEED) {
            const char bbr[] = "bbr";
            if (set_cc(skops, bbr, sizeof(bbr)) == 0) {
                st->state = STATE_SWITCHED;
                inc_stat(1);
            } else {
                inc_stat(3);
            }
        }
    }

    return 0;
}