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

#include "replica.h"

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

#define FLAG_SSYNC_HAPPENED 0
#define TARGET_PORT 8888

/* ================= BPF Maps ================= */

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

struct {
    __uint(type, BPF_MAP_TYPE_RINGBUF);
    __uint(max_entries, 1 << 26);  // 64MB
} rb SEC(".maps");

/* ================= Helper Functions ================= */

// 无需 process filter，改用 socket port filter

/* ================= Kernel Hooks (TCP Layer) ================= */

/*
 * 使用 kprobe 挂载 tcp_rcv_established
 * 此时 skb 包含完整的 TCP 包（Header + Payload），数据在内核态。
 */
SEC("kprobe/tcp_rcv_established")
int BPF_KPROBE(trace_tcp_rcv, struct sock *sk, struct sk_buff *skb)
{
    // 1. 检查 SSYNC 标志是否已开启 (只在全量同步后开始抓包)
    __u32 flag_key = FLAG_SSYNC_HAPPENED;
    __u32 *flag_val = bpf_map_lookup_elem(&flags, &flag_key);
    if (!flag_val || *flag_val == 0)
        return 0;

    // 2. 过滤端口 8888
    // sk->sk_num 存储的是 Host Byte Order 的本地端口
    __u16 lport = BPF_CORE_READ(sk, __sk_common.skc_num);
    if (lport != TARGET_PORT)
        return 0;

    // 3. 计算数据长度
    // 在 tcp_rcv_established 中，skb->len 是 (TCP Header + Data) 的长度
    // skb->data 指向 TCP Header 的起始位置
    unsigned int skb_len = BPF_CORE_READ(skb, len);
    
    // 读取 TCP Header 长度 (doff 字段，单位是 4字节)
    // 需要读取 skb->data 指向的内存的前几个字节来获取 doff
    unsigned char *skb_data = BPF_CORE_READ(skb, data);
    
    // 读取 TCP Header 的第 12 个字节 (包含 Data Offset)
    // Offset 12: Data Offset (4 bits) | Reserved (3 bits) | NS (1 bit)
    unsigned char doff_byte;
    if (bpf_probe_read_kernel(&doff_byte, 1, skb_data + 12) < 0)
        return 0;
    
    unsigned int tcp_hdr_len = (doff_byte >> 4) * 4;
    
    // 计算 Payload 长度
    if (skb_len <= tcp_hdr_len)
        return 0; // 只有 ACK 没有数据，或者是控制包

    unsigned int payload_len = skb_len - tcp_hdr_len;

    // 4. 准备 RingBuffer 数据
    struct replica_event *e = bpf_ringbuf_reserve(&rb, sizeof(*e), 0);
    if (!e)
        return 0;

    e->type = EVENT_COMPLETED_CMD;
    
    // 截断超长数据
    if (payload_len > MAX_CMD_LEN)
        e->complete.len = MAX_CMD_LEN;
    else
        e->complete.len = payload_len;

    // 5. 核心修改：使用 bpf_probe_read_kernel 读取数据
    // 数据起始位置 = skb->data + tcp_hdr_len
    if (bpf_probe_read_kernel(&e->complete.cmd[0], e->complete.len, skb_data + tcp_hdr_len) < 0) {
        bpf_ringbuf_discard(e, 0);
        return 0;
    }

    bpf_ringbuf_submit(e, 0);
    return 0;
}

/* ================= Uprobe Hooks================= */

SEC("uprobe//home/lian/share/9.1-kvstore/kvstore:__ssync")
int BPF_KPROBE(handle_ssync,
               const __u8 *ip, __u32 ip_len, int port, __u64 seq_unused)
{
    __u32 key = FLAG_SSYNC_HAPPENED;
    __u32 val = 1;
    bpf_map_update_elem(&flags, &key, &val, BPF_ANY);

    struct replica_event *e = bpf_ringbuf_reserve(&rb, sizeof(*e), 0);
    if (!e) return 0;

    e->type = EVENT_SSYNC;
    e->sync.port = port;

    __u32 copy_len = ip_len;
    if (copy_len > sizeof(e->sync.ip)) copy_len = sizeof(e->sync.ip);
    if (ip) bpf_probe_read_user(e->sync.ip, copy_len, ip);

    bpf_ringbuf_submit(e, 0);
    return 0;
}

SEC("uprobe//home/lian/share/9.1-kvstore/kvstore:__sready")
int BPF_KPROBE(handle_sready)
{
    struct replica_event *e = bpf_ringbuf_reserve(&rb, sizeof(*e), 0);
    if (!e) return 0;

    e->type = EVENT_SREADY;
    bpf_ringbuf_submit(e, 0);
    return 0;
}