improve associate context for exist connections

.github/workflows/test.yml

@@ -48,9 +48,9 @@ jobs:
         kernel:
           - '5.4-20240515.073534'
           - '5.10-20240515.073534'
-          # - '5.15-20240515.073534'
+          - '5.15-20240515.073534'
           - '6.1-20240515.073534'
-          # - '6.6-20240515.073534'
+          - '6.6-20240515.073534'
           - 'bpf-next-20240515.073534'
     timeout-minutes: 10
     runs-on: ubuntu-latest
@@ -198,14 +198,3 @@ jobs:
             cat /etc/issue
 
             bash /host/testdata/test_nat.sh /host/ptcpdump/ptcpdump
-
-      - name: Test docker
-        uses: cilium/little-vm-helper@3c748d6fc9d6c44a433de85a66f70e8f7043be04 # v0.0.18
-        with:
-          provision: 'false'
-          cmd: |
-            set -ex
-            uname -a
-            cat /etc/issue
-
-            bash /host/testdata/test_docker.sh /host/ptcpdump/ptcpdump

bpf/bpf.go

@@ -22,6 +22,9 @@ import (
 const tcFilterName = "ptcpdump"
 
 type BpfObjectsWithoutCgroup struct {
+	KprobeTcpSendmsg              *ebpf.Program `ebpf:"kprobe__tcp_sendmsg"`
+	KprobeUdpSendmsg              *ebpf.Program `ebpf:"kprobe__udp_sendmsg"`
+	KprobeUdpSendSkb              *ebpf.Program `ebpf:"kprobe__udp_send_skb"`
 	KprobeNfNatManipPkt           *ebpf.Program `ebpf:"kprobe__nf_nat_manip_pkt"`
 	KprobeNfNatPacket             *ebpf.Program `ebpf:"kprobe__nf_nat_packet"`
 	KprobeSecuritySkClassifyFlow  *ebpf.Program `ebpf:"kprobe__security_sk_classify_flow"`
@@ -141,6 +144,9 @@ func (b *BPF) Load(opts Options) error {
 			}); err != nil {
 				return err
 			}
+			b.objs.KprobeTcpSendmsg = objs.KprobeTcpSendmsg
+			b.objs.KprobeUdpSendmsg = objs.KprobeUdpSendmsg
+			b.objs.KprobeUdpSendSkb = objs.KprobeUdpSendSkb
 			b.objs.KprobeNfNatManipPkt = objs.KprobeNfNatManipPkt
 			b.objs.KprobeNfNatPacket = objs.KprobeNfNatPacket
 			b.objs.KprobeSecuritySkClassifyFlow = objs.KprobeSecuritySkClassifyFlow
@@ -223,6 +229,26 @@ func (b *BPF) AttachKprobes() error {
 	}
 	b.links = append(b.links, lk)
 
+	lk, err = link.Kprobe("tcp_sendmsg",
+		b.objs.KprobeTcpSendmsg, &link.KprobeOptions{})
+	if err != nil {
+		return xerrors.Errorf("attach kprobe/tcp_sendmsg: %w", err)
+	}
+	b.links = append(b.links, lk)
+
+	lk, err = link.Kprobe("udp_send_skb", b.objs.KprobeUdpSendSkb, &link.KprobeOptions{})
+	if err != nil {
+		if strings.Contains(err.Error(), "no such file or directory") {
+			lk, err = link.Kprobe("udp_sendmsg", b.objs.KprobeUdpSendmsg, &link.KprobeOptions{})
+			if err != nil {
+				return xerrors.Errorf("attach kprobe/udp_sendmsg: %w", err)
+			}
+		} else {
+			return xerrors.Errorf("attach kprobe/udp_send_skb: %w", err)
+		}
+	}
+	b.links = append(b.links, lk)
+
 	lk, err = link.Kprobe("nf_nat_packet",
 		b.objs.KprobeNfNatPacket, &link.KprobeOptions{})
 	if err != nil {
@@ -254,17 +280,17 @@ func (b *BPF) AttachKprobes() error {
 
 func (b *BPF) AttachTracepoints() error {
 	lk, err := link.AttachRawTracepoint(link.RawTracepointOptions{
-		"sched_process_exec",
-		b.objs.RawTracepointSchedProcessExec,
+		Name:    "sched_process_exec",
+		Program: b.objs.RawTracepointSchedProcessExec,
 	})
 	if err != nil {
 		return xerrors.Errorf("attach raw_tracepoint/sched_process_exec: %w", err)
 	}
 	b.links = append(b.links, lk)
 
 	lk, err = link.AttachRawTracepoint(link.RawTracepointOptions{
-		"sched_process_exit",
-		b.objs.RawTracepointSchedProcessExit,
+		Name:    "sched_process_exit",
+		Program: b.objs.RawTracepointSchedProcessExit,
 	})
 	if err != nil {
 		return xerrors.Errorf("attach raw_tracepoint/sched_process_exit: %w", err)
@@ -273,8 +299,8 @@ func (b *BPF) AttachTracepoints() error {
 
 	if b.opts.attachForks() {
 		lk, err := link.AttachRawTracepoint(link.RawTracepointOptions{
-			"sched_process_fork",
-			b.objs.RawTracepointSchedProcessFork,
+			Name:    "sched_process_fork",
+			Program: b.objs.RawTracepointSchedProcessFork,
 		})
 		if err != nil {
 			return xerrors.Errorf("attach raw_tracepoint/sched_process_fork: %w", err)
@@ -344,14 +370,14 @@ func attachTcHook(ifindex int, prog *ebpf.Program, ingress bool) (func(), error)
 	}
 
 	filter := tc.Object{
-		tc.Msg{
+		Msg: tc.Msg{
 			Family:  unix.AF_UNSPEC,
 			Ifindex: uint32(ifindex),
 			Handle:  0,
 			Parent:  core.BuildHandle(tc.HandleRoot, parent),
 			Info:    1<<16 | uint32(htons(unix.ETH_P_ALL)),
 		},
-		tc.Attribute{
+		Attribute: tc.Attribute{
 			Kind: "bpf",
 			BPF: &tc.Bpf{
 				FD:   &fd,

bpf/ptcpdump.c

@@ -490,6 +490,53 @@ int BPF_KPROBE(kprobe__security_sk_classify_flow, struct sock *sk) {
     return 0;
 }
 
+static __always_inline void handle_sendmsg(struct sock *sk) {
+    struct flow_pid_key_t key = {0};
+    struct process_meta_t value = {0};
+    struct task_struct *task = (struct task_struct *)bpf_get_current_task();
+
+    if (parent_process_filter(task) < 0) {
+        if (process_filter(task) < 0) {
+            return;
+        }
+    }
+
+    fill_sk_meta(sk, &key);
+    if (bpf_map_lookup_elem(&flow_pid_map, &key)) {
+        return;
+    }
+
+    fill_process_meta(task, &value);
+    if (key.sport == 0) {
+        return;
+    }
+    // bpf_printk("[ptcpdump] flow key: %pI4 %d", &key.saddr[0], key.sport);
+    int ret = bpf_map_update_elem(&flow_pid_map, &key, &value, BPF_NOEXIST);
+    if (ret != 0) {
+        bpf_printk("[handle_tcp_sendmsg] bpf_map_update_elem flow_pid_map failed: %d", ret);
+    }
+    return;
+}
+
+SEC("kprobe/tcp_sendmsg")
+int BPF_KPROBE(kprobe__tcp_sendmsg, struct sock *sk) {
+    handle_sendmsg(sk);
+    return 0;
+}
+
+SEC("kprobe/udp_sendmsg")
+int BPF_KPROBE(kprobe__udp_sendmsg, struct sock *sk) {
+    handle_sendmsg(sk);
+    return 0;
+}
+
+SEC("kprobe/udp_send_skb")
+int BPF_KPROBE(kprobe__udp_send_skb, struct sk_buff *skb) {
+    struct sock *sk = BPF_CORE_READ(skb, sk);
+    handle_sendmsg(sk);
+    return 0;
+}
+
 static __noinline bool pcap_filter(void *_skb, void *__skb, void *___skb, void *data, void *data_end) {
     return data != data_end && _skb == __skb && __skb == ___skb;
 }