feat: Clean map when destroy socket (#691)

* feat: Clean map when destroy socket When a socket has been closed in kernel, the corresponding connection info in user space should be cleaned in order to avoid memory leaking. Then, in user space, use another map to keep pid and fd for the socket with socket pointer in kernel space as key. It is because the tuple info in the socket is incorrect when destroy it. e.g. ``` 2024-12-13T13:23:10Z DBG AddConn success fd=5 pid=134386 tuple=192.168.241.133:35348-172.217.194.113:443 2024-12-13T13:23:10Z DBG DestroyConn success fd=5 pid=134386 tuple=192.168.241.133:35348-172.217.194.113:443 2024-12-13T13:23:16Z DBG AddConn success fd=4 pid=134404 tuple=192.168.241.1:65063-192.168.241.133:8080 2024-12-13T13:23:16Z DBG DestroyConn success fd=4 pid=134404 tuple=192.168.241.1:65063-192.168.241.133:8080 ``` Signed-off-by: Leon Hwang <hffilwlqm@gmail.com> * fix: Only handle tcp socket when probe connect and accept It's unnecessary to record excluding-tcp sockets' info. So, when probe connect and accept, we only handle tcp socket by filtering `sk->sk_protocol`. Signed-off-by: Leon Hwang <hffilwlqm@gmail.com> --------- Signed-off-by: Leon Hwang <hffilwlqm@gmail.com>
gojue · Dec 14, 2024 · f0f1752 · f0f1752
1 parent 2a05372
commit f0f1752
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 11 deletions.
diff --git a/kern/openssl.h b/kern/openssl.h
@@ -48,6 +48,9 @@ struct connect_event_t {
     __be32 saddr;
     __be32 daddr;
     char comm[TASK_COMM_LEN];
+    u64 sock;
+    u8 is_destroy;
+    u8 pad[7];
 } __attribute__((packed)); // NOTE: do not leave padding hole in this struct.
 
 struct active_ssl_buf {
@@ -511,6 +514,7 @@ static __inline int kretprobe_connect(struct pt_regs *ctx, int fd, struct sock *
     u64 current_uid_gid = bpf_get_current_uid_gid();
     u32 uid = current_uid_gid;
     u16 address_family = 0;
+    u16 protocol;
     u64 addrs;
     u32 ports;
 
@@ -529,6 +533,11 @@ static __inline int kretprobe_connect(struct pt_regs *ctx, int fd, struct sock *
         return 0;
     }
 
+    bpf_probe_read_kernel(&protocol, sizeof(protocol), &sk->sk_protocol);
+    if (protocol != IPPROTO_TCP) {
+        return 0;
+    }
+
     // if the connection hasn't been established yet, the ports or addrs are 0.
     bpf_probe_read_kernel(&addrs, sizeof(addrs), &sk->__sk_common.skc_addrpair);
     bpf_probe_read_kernel(&ports, sizeof(ports), &sk->__sk_common.skc_portpair);
@@ -556,6 +565,7 @@ static __inline int kretprobe_connect(struct pt_regs *ctx, int fd, struct sock *
         conn.daddr = (__be32)(addrs >> 32);
     }
     bpf_get_current_comm(&conn.comm, sizeof(conn.comm));
+    conn.sock = (u64)sk;
 
     bpf_perf_event_output(ctx, &connect_events, BPF_F_CURRENT_CPU, &conn,
                           sizeof(struct connect_event_t));
@@ -620,6 +630,20 @@ int retprobe_accept4(struct pt_regs* ctx) {
     return 0;
 }
 
+SEC("kprobe/tcp_v4_destroy_sock")
+int probe_tcp_v4_destroy_sock(struct pt_regs* ctx) {
+    struct sock *sk = (struct sock *)PT_REGS_PARM1(ctx);
+
+    struct connect_event_t conn;
+    conn.sock = (u64)sk;
+    conn.is_destroy = 1;
+
+    bpf_perf_event_output(ctx, &connect_events, BPF_F_CURRENT_CPU, &conn,
+                          sizeof(struct connect_event_t));
+
+    return BPF_OK;
+}
+
 
 // int SSL_set_fd(SSL *s, int fd)
 // int SSL_set_rfd(SSL *s, int fd)

diff --git a/user/event/event_openssl.go b/user/event/event_openssl.go
@@ -215,6 +215,9 @@ type connDataEvent struct {
 	Saddr       [4]byte  `json:"saddr"`
 	Daddr       [4]byte  `json:"daddr"`
 	Comm        [16]byte `json:"Comm"`
+	Sock        uint64   `json:"sock"`
+	IsDestroy   uint8    `json:"isDestroy"`
+	Pad         [7]byte  `json:"-"`
 
 	// NOTE: do not leave padding hole in this struct.
 }

diff --git a/user/module/probe_openssl.go b/user/module/probe_openssl.go
@@ -82,9 +82,11 @@ type MOpenSSLProbe struct {
 	eventFuncMaps     map[*ebpf.Map]event.IEventStruct
 	eventMaps         []*ebpf.Map
 
-	// pid[fd:Addr]
-	pidConns  map[uint32]map[uint32]string
-	pidLocker sync.Locker
+	// pid[fd:tuple]
+	pidConns map[uint32]map[uint32]string
+	// sock:[pid,fd], for destroying conn
+	sock2pidFd map[uint64][2]uint32
+	pidLocker  sync.Locker
 
 	keyloggerFilename string
 	keylogger         *os.File
@@ -110,6 +112,7 @@ func (m *MOpenSSLProbe) Init(ctx context.Context, logger *zerolog.Logger, conf c
 	m.eventMaps = make([]*ebpf.Map, 0, 2)
 	m.eventFuncMaps = make(map[*ebpf.Map]event.IEventStruct)
 	m.pidConns = make(map[uint32]map[uint32]string)
+	m.sock2pidFd = make(map[uint64][2]uint32)
 	m.pidLocker = new(sync.Mutex)
 	m.masterKeys = make(map[string]bool)
 	m.sslVersionBpfMap = make(map[string]string)
@@ -392,7 +395,7 @@ func (m *MOpenSSLProbe) Events() []*ebpf.Map {
 	return m.eventMaps
 }
 
-func (m *MOpenSSLProbe) AddConn(pid, fd uint32, tuple string) {
+func (m *MOpenSSLProbe) AddConn(pid, fd uint32, tuple string, sock uint64) {
 	if fd <= 0 {
 		m.logger.Info().Uint32("pid", pid).Uint32("fd", fd).Str("tuple", tuple).Msg("AddConn failed")
 		return
@@ -408,8 +411,38 @@ func (m *MOpenSSLProbe) AddConn(pid, fd uint32, tuple string) {
 	}
 	connMap[fd] = tuple
 	m.pidConns[pid] = connMap
+
+	m.sock2pidFd[sock] = [2]uint32{pid, fd}
+
 	m.logger.Debug().Uint32("pid", pid).Uint32("fd", fd).Str("tuple", tuple).Msg("AddConn success")
-	return
+}
+
+func (m *MOpenSSLProbe) DestroyConn(sock uint64) {
+	m.pidLocker.Lock()
+	defer m.pidLocker.Unlock()
+
+	pidFd, ok := m.sock2pidFd[sock]
+	if !ok {
+		return
+	}
+
+	delete(m.sock2pidFd, sock)
+	pid, fd := pidFd[0], pidFd[1]
+
+	connMap, ok := m.pidConns[pid]
+	if !ok {
+		return
+	}
+
+	tuple, ok := connMap[fd]
+	if ok {
+		delete(connMap, fd)
+		if len(connMap) == 0 {
+			delete(m.pidConns, pid)
+		}
+	}
+
+	m.logger.Debug().Uint32("pid", pid).Uint32("fd", fd).Str("tuple", tuple).Msg("DestroyConn success")
 }
 
 // process exit :fd is 0 , delete all pid map
@@ -705,20 +738,24 @@ func (m *MOpenSSLProbe) mk13NullSecrets(hashLen int,
 
 func (m *MOpenSSLProbe) Dispatcher(eventStruct event.IEventStruct) {
 	// detect eventStruct type
-	switch eventStruct.(type) {
+	switch ev := eventStruct.(type) {
 	case *event.ConnDataEvent:
-		m.AddConn(eventStruct.(*event.ConnDataEvent).Pid, eventStruct.(*event.ConnDataEvent).Fd, eventStruct.(*event.ConnDataEvent).Tuple)
+		if ev.IsDestroy == 0 {
+			m.AddConn(ev.Pid, ev.Fd, ev.Tuple, ev.Sock)
+		} else {
+			m.DestroyConn(ev.Sock)
+		}
 	case *event.MasterSecretEvent:
-		m.saveMasterSecret(eventStruct.(*event.MasterSecretEvent))
+		m.saveMasterSecret(ev)
 	case *event.MasterSecretBSSLEvent:
-		m.saveMasterSecretBSSL(eventStruct.(*event.MasterSecretBSSLEvent))
+		m.saveMasterSecretBSSL(ev)
 	case *event.TcSkbEvent:
-		err := m.dumpTcSkb(eventStruct.(*event.TcSkbEvent))
+		err := m.dumpTcSkb(ev)
 		if err != nil {
 			m.logger.Error().Err(err).Msg("save packet error.")
 		}
 	case *event.SSLDataEvent:
-		m.dumpSslData(eventStruct.(*event.SSLDataEvent))
+		m.dumpSslData(ev)
 	}
 }
 

diff --git a/user/module/probe_openssl_text.go b/user/module/probe_openssl_text.go
@@ -107,6 +107,12 @@ func (m *MOpenSSLProbe) setupManagersText() error {
 				AttachToFuncName: "__sys_accept4",
 				UID:              "kretprobe_sys_accept4",
 			},
+			{
+				Section:          "kprobe/tcp_v4_destroy_sock",
+				EbpfFuncName:     "probe_tcp_v4_destroy_sock",
+				AttachToFuncName: "tcp_v4_destroy_sock",
+				UID:              "kprobe_tcp_v4_destroy_sock",
+			},
 
 			// --------------------------------------------------