#include "common.h"

enum filter_ret_code {
    // 不处理当前连接
	PASS = 0,
	// 处理当前连接
	STORE = 1,
};


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

static inline
void store_sock_ops(struct bpf_sock_ops *skops)
{
    struct sock_key key = {
        .dip = skops->remote_ip4,
        .sip = skops->local_ip4,
        .sport = bpf_htonl(skops->local_port),  /* convert to network byte order */
        .dport = skops->remote_port,
        .family = skops->family,
    };

    int ret = bpf_sock_hash_update(skops, &sock_ops_map, &key, BPF_NOEXIST);
    if (ret != 0) {
        bpf_printk("sock_hash_update() failed, ret: %d\n", ret);
    }

}

// 判断传入连接是否需要处理
static inline 
__u8 sock_filter(struct bpf_sock_ops *skops)
{
    if (skops->family != 2){ // AF_INET
        return PASS;
    }
    /*
        local traffic only, so check both local and remote IPs
    */
    __u32 remote_ip = skops->remote_ip4;
    __u32 local_ip = skops->local_ip4;

    // 检查 remote_ip 是否在目标 map 中
    if (!bpf_map_lookup_elem(&target_ip_map, &remote_ip)) {
        return PASS;
    }

    // 如果是本地流量，大概率remote ip 和 local ip 是相同的，此时可以省去一次 map 查询
    if (remote_ip == local_ip) {
        return STORE;
    }

    // 检查 local_ip 是否在目标 map 中
    if (!bpf_map_lookup_elem(&target_ip_map, &local_ip)) {
        return PASS;
    }

    return STORE;
}

SEC("sockops")
int bpf_sockops_handler(struct bpf_sock_ops *skops)
{
    // only interested in established events
    if (skops->op != BPF_SOCK_OPS_PASSIVE_ESTABLISHED_CB &&
        skops->op != BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB) {
        return BPF_OK;
    }
    if (sock_filter(skops)){ // AF_INET
        store_sock_ops(skops);  // 将 socket 信息记录到到 sockmap
    }

    return BPF_OK;
}