feat: Clean map when destroy socket

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)

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.
 }

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)
 	}
 }
 

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",
+			},
 
 			// --------------------------------------------------