// Based on the tcptop BCC tool by Brendan Gregg:
// https://github.com/iovisor/bcc/blob/master/tools/tcptop.py
//
// Please refer to https://github.com/iovisor/bcc/blob/master/tools/tcptop_example.txt
// for overhead measurements.

#include <uapi/linux/ptrace.h>
#include <net/sock.h>

struct netstats {
    u64 tcp_send_calls;
    u64 tcp_send_bytes;
    u64 tcp_recv_calls;
    u64 tcp_recv_bytes;

    u64 udp_send_calls;
    u64 udp_send_bytes;
    u64 udp_recv_calls;
    u64 udp_recv_bytes;
};

BPF_TABLE("lru_hash", u32, struct netstats, netstats_per_pid, 10240);

int kprobe__tcp_sendmsg(struct pt_regs *ctx, struct sock *sk,
    struct msghdr *msg, size_t size)
{
    u32 pid = bpf_get_current_pid_tgid();

    struct netstats *statp, zero = {};
    statp = netstats_per_pid.lookup_or_try_init(&pid, &zero);
    if (statp) {
        statp->tcp_send_calls++;
        statp->tcp_send_bytes += size;
    }

    return 0;
}

/*
 * tcp_recvmsg() would be obvious to trace, but is less suitable because:
 * - we'd need to trace both entry and return, to have both sock and size
 * - misses tcp_read_sock() traffic
 * we'd much prefer tracepoints once they are available.
 */
int kprobe__tcp_cleanup_rbuf(struct pt_regs *ctx, struct sock *sk, int copied)
{
    if (copied <= 0)
        return 0;

    u32 pid = bpf_get_current_pid_tgid();
    struct netstats *statp, zero = {};
    statp = netstats_per_pid.lookup_or_try_init(&pid, &zero);
    if (statp) {
        statp->tcp_recv_calls++;
        statp->tcp_recv_bytes += copied;
    }
    return 0;
}

int kprobe__udp_sendmsg(struct pt_regs *ctx, struct sock *sk,
    struct msghdr *msg, size_t len)
{
    u32 pid = bpf_get_current_pid_tgid();
    struct netstats *statp, zero = {};
    statp = netstats_per_pid.lookup_or_try_init(&pid, &zero);
    if (statp) {
        statp->udp_send_calls++;
        statp->udp_send_bytes += len;
    }
    return 0;
}

// similar reason as above with udp_recvmsg()
int kprobe__skb_consume_udp(struct pt_regs *ctx, struct sock *sk,
    struct sk_buff *skb, int len)
{
    if (len <= 0)
        return 0;

    u32 pid = bpf_get_current_pid_tgid();
    struct netstats *statp, zero = {};
    statp = netstats_per_pid.lookup_or_try_init(&pid, &zero);
    if (statp) {
        statp->udp_recv_calls++;
        statp->udp_recv_bytes += len;
    }
    return 0;
}