From cb4f4c07a4873d026a6de35d367a086b1bbfc877 Mon Sep 17 00:00:00 2001
From: Wangjunqi123 <wangjunqi@kylinos.cn>
Date: Tue, 19 Nov 2024 15:16:08 +0800
Subject: [PATCH] agent/probes: define probe tcp_cleanup_rbuf

---
 agent/probes/network/src/tcp_netflow.bpf.c | 60 ++++++++++++++++++++++
 1 file changed, 60 insertions(+)

diff --git a/agent/probes/network/src/tcp_netflow.bpf.c b/agent/probes/network/src/tcp_netflow.bpf.c
index 8b3e6d9..65fd731 100644
--- a/agent/probes/network/src/tcp_netflow.bpf.c
+++ b/agent/probes/network/src/tcp_netflow.bpf.c
@@ -91,4 +91,64 @@ int BPF_KRETPROBE(tcp_sendmsg_exit, int ret) {
     // ttcode
     bpf_printk("(tcp_sendmsg_exit) pid: %u", pid);
     return 0;
+}
+
+SEC("kprobe/tcp_cleanup_rbuf")
+int BPF_KPROBE(tcp_cleanup_rbuf, struct sock *sk, int copied) {
+    // ttcode
+    u8 comm[16], _comm[TARGET_NUM][16] = TARGET_PROC;
+    (void)bpf_get_current_comm(&comm, sizeof(comm));
+    if (strcmp(comm, _comm[0]) == 1 && strcmp(comm, _comm[1]) == 1 && strcmp(comm, _comm[2]) == 1) {
+        return 0;
+    }
+
+    u32 pid = bpf_get_current_pid_tgid() >> INT_LEN;
+    if (copied < 0) {
+        bpf_printk("(tcp_cleanup_rbuf) pid: %u errorcode: %d", pid, copied);
+        return 0;
+    }
+
+    // ttcode
+    bpf_printk("(tcp_cleanup_rbuf) pid: %u", pid);
+
+    struct tcp_metrics *metrics = bpf_map_lookup_elem(&tcp_link_map, &pid);
+    // struct tcp_metrics *metrics = bpf_map_lookup_elem(&tcp_link_map, &sk);
+    if (!metrics) {
+        return 0;
+    }
+
+    if (copied <= 0) {
+        return 0;
+    }
+
+    metrics->family = _(sk->sk_family);
+    if (metrics->role == LINK_ROLE_CLIENT) {
+        if (metrics->family == AF_INET) {
+            metrics->c_ip = _(sk->sk_rcv_saddr);
+            metrics->s_ip = _(sk->sk_daddr);
+        } else {
+            BPF_CORE_READ_INTO(metrics->c_ip6, sk, sk_v6_rcv_saddr);
+            BPF_CORE_READ_INTO(metrics->s_ip6, sk, sk_v6_daddr);
+        }
+        metrics->s_port = bpf_ntohs(_(sk->sk_dport));
+        metrics->c_port = _(sk->sk_num);
+    } else {
+        if (metrics->family == AF_INET) {
+            metrics->s_ip = _(sk->sk_rcv_saddr);
+            metrics->c_ip = _(sk->sk_daddr);
+        } else {
+            BPF_CORE_READ_INTO(metrics->s_ip6, sk, sk_v6_rcv_saddr);
+            BPF_CORE_READ_INTO(metrics->c_ip6, sk, sk_v6_daddr);
+        }
+        metrics->s_port = _(sk->sk_num);
+        metrics->c_port = bpf_ntohs(_(sk->sk_dport));
+    }
+    metrics->pid = pid;
+    (void)bpf_get_current_comm(metrics->comm, sizeof(metrics->comm));
+
+    metrics->rx = (u64)copied;
+    // __sync_fetch_and_add(&(metrics->rx), (u64)(copied));
+
+    bpf_ringbuf_output(&tcp_output, metrics, sizeof(struct tcp_metrics), 0);
+    return 0;
 }
\ No newline at end of file
-- 
Gitee