Add a tracepoint for sys_enter_close to detect socket closes

tap/tlstapper/bpf/fd_tracepoints.c

@@ -90,3 +90,23 @@ void sys_enter_write(struct sys_enter_write_ctx *ctx) {
 		log_error(ctx, LOG_ERROR_PUTTING_FILE_DESCRIPTOR, id, err, ORIGIN_SYS_ENTER_WRITE_CODE);
 	}
 }
+
+struct sys_enter_close_ctx {
+	__u64 __unused_syscall_header;
+	__u32 __unused_syscall_nr;
+
+	__u64 fd;
+};
+
+SEC("tracepoint/syscalls/sys_enter_close")
+void sys_enter_close(struct sys_enter_close_ctx *ctx) {
+	__u64 id = bpf_get_current_pid_tgid();
+
+	if (!should_tap(id >> 32)) {
+		return;
+	}
+
+	struct sys_close event;
+    event.fd = ctx->fd;
+    bpf_perf_event_output(ctx, &sys_closes, BPF_F_CURRENT_CPU, &event, sizeof(event));
+}

tap/tlstapper/bpf/golang_uprobes.c

@@ -4,25 +4,10 @@ SPDX-License-Identifier: GPL-2.0
 Copyright (C) UP9 Inc.
 */
 
-#include <stdbool.h>
-
 #include "include/headers.h"
 #include "include/maps.h"
 
 
-struct golang_read_write {
-    __u32 pid;
-    __u32 fd;
-    __u32 conn_addr;
-    bool is_request;
-    __u32 len;
-    __u32 cap;
-    __u8 data[CHUNK_SIZE];
-};
-
-const struct golang_read_write *unused __attribute__((unused));
-
-
 SEC("uprobe/golang_crypto_tls_write")
 static __always_inline int golang_crypto_tls_write_uprobe(struct pt_regs *ctx) {
     void* stack_addr = (void*)ctx->rsp;

tap/tlstapper/bpf/include/headers.h

@@ -8,6 +8,7 @@ Copyright (C) UP9 Inc.
 #define __HEADERS__
 
 #include <stddef.h>
+#include <stdbool.h>
 #include <linux/bpf.h>
 #include <linux/ptrace.h>
 #include <bpf/bpf_helpers.h>

tap/tlstapper/bpf/include/maps.h

@@ -53,6 +53,10 @@ struct fd_info {
     __u8 flags;
 };
 
+struct sys_close {
+    __u32 fd;
+};
+
 struct socket {
     __u32 pid;
     __u32 fd;
@@ -60,6 +64,20 @@ struct socket {
     __u64 conn_addr;
 };
 
+struct golang_read_write {
+    __u32 pid;
+    __u32 fd;
+    __u32 conn_addr;
+    bool is_request;
+    __u32 len;
+    __u32 cap;
+    __u8 data[CHUNK_SIZE];
+};
+
+const struct golang_read_write *unused1 __attribute__((unused));
+const struct sys_close *unused2 __attribute__((unused));
+
+
 #define BPF_MAP(_name, _type, _key_type, _value_type, _max_entries)     \
     struct bpf_map_def SEC("maps") _name = {                            \
         .type = _type,                                                  \
@@ -93,5 +111,6 @@ BPF_PERF_OUTPUT(log_buffer);
 BPF_LRU_HASH(golang_dial_to_socket, __u64, struct socket);
 BPF_LRU_HASH(golang_socket_to_write, __u64, struct socket);
 BPF_RINGBUF(golang_read_writes);
+BPF_PERF_OUTPUT(sys_closes);
 
 #endif /* __MAPS__ */

tap/tlstapper/golang_connection.go

@@ -4,8 +4,10 @@ import "github.com/up9inc/mizu/tap/api"
 
 type golangConnection struct {
 	pid          uint32
+	fd           uint32
 	connAddr     uint32
 	addressPair  addressPair
+	addressIsSet bool
 	stream       *tlsStream
 	clientReader *golangReader
 	serverReader *golangReader
@@ -28,10 +30,24 @@ func NewGolangConnection(pid uint32, connAddr uint32, extension *api.Extension,
 }
 
 func (c *golangConnection) setAddressBySockfd(procfs string, pid uint32, fd uint32) error {
+	if c.addressIsSet {
+		return nil
+	}
+
 	addrPair, err := getAddressBySockfd(procfs, pid, fd)
 	if err != nil {
 		return err
 	}
 	c.addressPair = addrPair
+	c.addressIsSet = true
 	return nil
 }
+
+func (c *golangConnection) close() {
+	if c.clientReader != nil {
+		c.clientReader.close()
+	}
+	if c.serverReader != nil {
+		c.serverReader.close()
+	}
+}

tap/tlstapper/syscall_hooks.go

@@ -8,6 +8,7 @@ import (
 type syscallHooks struct {
 	sysEnterRead    link.Link
 	sysEnterWrite   link.Link
+	sysEnterClose   link.Link
 	sysEnterAccept4 link.Link
 	sysExitAccept4  link.Link
 	sysEnterConnect link.Link
@@ -29,6 +30,12 @@ func (s *syscallHooks) installSyscallHooks(bpfObjects *tlsTapperObjects) error {
 		return errors.Wrap(err, 0)
 	}
 
+	s.sysEnterClose, err = link.Tracepoint("syscalls", "sys_enter_close", bpfObjects.SysEnterClose)
+
+	if err != nil {
+		return errors.Wrap(err, 0)
+	}
+
 	s.sysEnterAccept4, err = link.Tracepoint("syscalls", "sys_enter_accept4", bpfObjects.SysEnterAccept4)
 
 	if err != nil {
@@ -67,6 +74,10 @@ func (s *syscallHooks) close() []error {
 		errors = append(errors, err)
 	}
 
+	if err := s.sysEnterClose.Close(); err != nil {
+		errors = append(errors, err)
+	}
+
 	if err := s.sysEnterAccept4.Close(); err != nil {
 		errors = append(errors, err)
 	}

tap/tlstapper/tls_poller.go

@@ -38,6 +38,7 @@ type tlsPoller struct {
 	chunksReader       *perf.Reader
 	golangReader       *ringbuf.Reader
 	golangReadWriteMap *orderedmap.OrderedMap
+	sysCloses          *perf.Reader
 	extension          *api.Extension
 	procfs             string
 	pidToNamespace     sync.Map
@@ -81,6 +82,12 @@ func (p *tlsPoller) init(bpfObjects *tlsTapperObjects, bufferSize int) error {
 		return errors.Wrap(err, 0)
 	}
 
+	p.sysCloses, err = perf.NewReader(bpfObjects.SysCloses, bufferSize)
+
+	if err != nil {
+		return errors.Wrap(err, 0)
+	}
+
 	p.golangReadWriteMap = orderedmap.New()
 
 	return nil
@@ -113,6 +120,7 @@ func (p *tlsPoller) pollSsllib(emitter api.Emitter, options *api.TrafficFilterin
 
 func (p *tlsPoller) pollGolang(emitter api.Emitter, options *api.TrafficFilteringOptions, streamsMap api.TcpStreamMap) {
 	go p.pollGolangReadWrite(p.golangReader, emitter, options, streamsMap)
+	go p.pollSysClose(p.sysCloses)
 }
 
 func (p *tlsPoller) pollGolangReadWrite(rd *ringbuf.Reader, emitter api.Emitter, options *api.TrafficFilteringOptions,
@@ -139,15 +147,7 @@ func (p *tlsPoller) pollGolangReadWrite(rd *ringbuf.Reader, emitter api.Emitter,
 
 		if p.golangReadWriteMap.Len()+1 > golangMapLimit {
 			pair := p.golangReadWriteMap.Oldest()
-			c := pair.Value.(*golangConnection)
+			pair.Value.(*golangConnection).close()
-			clientReader := c.clientReader
-			if clientReader != nil {
-				clientReader.close()
-			}
-			serverReader := c.serverReader
-			if serverReader != nil {
-				serverReader.close()
-			}
 			p.golangReadWriteMap.Delete(pair.Key)
 		}
 
@@ -171,6 +171,8 @@ func (p *tlsPoller) pollGolangReadWrite(rd *ringbuf.Reader, emitter api.Emitter,
 		}
 
 		if b.IsRequest {
+			connection.fd = b.Fd
+
 			err := connection.setAddressBySockfd(p.procfs, b.Pid, b.Fd)
 			if err != nil {
 				log.Printf("Error resolving address pair from fd: %s", err)
@@ -200,6 +202,41 @@ func (p *tlsPoller) pollGolangReadWrite(rd *ringbuf.Reader, emitter api.Emitter,
 	}
 }
 
+func (p *tlsPoller) pollSysClose(rd *perf.Reader) {
+	nativeEndian := p.getByteOrder()
+	// tlsTapperSysClose is generated by bpf2go.
+	var b tlsTapperSysClose
+	for {
+		record, err := rd.Read()
+		if err != nil {
+			if errors.Is(err, perf.ErrClosed) {
+				return
+			}
+			log.Printf("reading from perf event reader: %s", err)
+			continue
+		}
+
+		if record.LostSamples != 0 {
+			log.Printf("perf event ring buffer full, dropped %d samples", record.LostSamples)
+			continue
+		}
+
+		if err := binary.Read(bytes.NewBuffer(record.RawSample), nativeEndian, &b); err != nil {
+			log.Printf("parsing perf event: %s", err)
+			continue
+		}
+
+		// Close and remove the connection from map if its socket file descriptor is closed.
+		for pair := p.golangReadWriteMap.Oldest(); pair != nil; pair = pair.Next() {
+			connection := pair.Value.(*golangConnection)
+			if connection.fd == b.Fd {
+				connection.close()
+				p.golangReadWriteMap.Delete(pair.Key)
+			}
+		}
+	}
+}
+
 func (p *tlsPoller) pollChunksPerfBuffer(chunks chan<- *tlsChunk) {
 	logger.Log.Infof("Start polling for tls events")
 

tap/tlstapper/tls_tapper.go

@@ -14,7 +14,7 @@ import (
 
 const GLOABL_TAP_PID = 0
 
-//go:generate go run github.com/cilium/ebpf/cmd/bpf2go@0d0727ef53e2f53b1731c73f4c61e0f58693083a -type golang_read_write tlsTapper bpf/tls_tapper.c -- -O2 -g -D__TARGET_ARCH_x86
+//go:generate go run github.com/cilium/ebpf/cmd/bpf2go@0d0727ef53e2f53b1731c73f4c61e0f58693083a -type golang_read_write -type sys_close tlsTapper bpf/tls_tapper.c -- -O2 -g -D__TARGET_ARCH_x86
 
 type TlsTapper struct {
 	bpfObjects         tlsTapperObjects

tap/tlstapper/tlstapper_bpfeb.go

@@ -24,6 +24,8 @@ type tlsTapperGolangReadWrite struct {
 	Data      [4096]uint8
 }
 
+type tlsTapperSysClose struct{ Fd uint32 }
+
 type tlsTapperTlsChunk struct {
 	Pid      uint32
 	Tgid     uint32
@@ -90,6 +92,7 @@ type tlsTapperProgramSpecs struct {
 	SslWrite                    *ebpf.ProgramSpec `ebpf:"ssl_write"`
 	SslWriteEx                  *ebpf.ProgramSpec `ebpf:"ssl_write_ex"`
 	SysEnterAccept4             *ebpf.ProgramSpec `ebpf:"sys_enter_accept4"`
+	SysEnterClose               *ebpf.ProgramSpec `ebpf:"sys_enter_close"`
 	SysEnterConnect             *ebpf.ProgramSpec `ebpf:"sys_enter_connect"`
 	SysEnterRead                *ebpf.ProgramSpec `ebpf:"sys_enter_read"`
 	SysEnterWrite               *ebpf.ProgramSpec `ebpf:"sys_enter_write"`
@@ -113,6 +116,7 @@ type tlsTapperMapSpecs struct {
 	PidsMap              *ebpf.MapSpec `ebpf:"pids_map"`
 	SslReadContext       *ebpf.MapSpec `ebpf:"ssl_read_context"`
 	SslWriteContext      *ebpf.MapSpec `ebpf:"ssl_write_context"`
+	SysCloses            *ebpf.MapSpec `ebpf:"sys_closes"`
 }
 
 // tlsTapperObjects contains all objects after they have been loaded into the kernel.
@@ -146,6 +150,7 @@ type tlsTapperMaps struct {
 	PidsMap              *ebpf.Map `ebpf:"pids_map"`
 	SslReadContext       *ebpf.Map `ebpf:"ssl_read_context"`
 	SslWriteContext      *ebpf.Map `ebpf:"ssl_write_context"`
+	SysCloses            *ebpf.Map `ebpf:"sys_closes"`
 }
 
 func (m *tlsTapperMaps) Close() error {
@@ -162,6 +167,7 @@ func (m *tlsTapperMaps) Close() error {
 		m.PidsMap,
 		m.SslReadContext,
 		m.SslWriteContext,
+		m.SysCloses,
 	)
 }
 
@@ -182,6 +188,7 @@ type tlsTapperPrograms struct {
 	SslWrite                    *ebpf.Program `ebpf:"ssl_write"`
 	SslWriteEx                  *ebpf.Program `ebpf:"ssl_write_ex"`
 	SysEnterAccept4             *ebpf.Program `ebpf:"sys_enter_accept4"`
+	SysEnterClose               *ebpf.Program `ebpf:"sys_enter_close"`
 	SysEnterConnect             *ebpf.Program `ebpf:"sys_enter_connect"`
 	SysEnterRead                *ebpf.Program `ebpf:"sys_enter_read"`
 	SysEnterWrite               *ebpf.Program `ebpf:"sys_enter_write"`
@@ -204,6 +211,7 @@ func (p *tlsTapperPrograms) Close() error {
 		p.SslWrite,
 		p.SslWriteEx,
 		p.SysEnterAccept4,
+		p.SysEnterClose,
 		p.SysEnterConnect,
 		p.SysEnterRead,
 		p.SysEnterWrite,

tap/tlstapper/tlstapper_bpfel.go

@@ -24,6 +24,8 @@ type tlsTapperGolangReadWrite struct {
 	Data      [4096]uint8
 }
 
+type tlsTapperSysClose struct{ Fd uint32 }
+
 type tlsTapperTlsChunk struct {
 	Pid      uint32
 	Tgid     uint32
@@ -90,6 +92,7 @@ type tlsTapperProgramSpecs struct {
 	SslWrite                    *ebpf.ProgramSpec `ebpf:"ssl_write"`
 	SslWriteEx                  *ebpf.ProgramSpec `ebpf:"ssl_write_ex"`
 	SysEnterAccept4             *ebpf.ProgramSpec `ebpf:"sys_enter_accept4"`
+	SysEnterClose               *ebpf.ProgramSpec `ebpf:"sys_enter_close"`
 	SysEnterConnect             *ebpf.ProgramSpec `ebpf:"sys_enter_connect"`
 	SysEnterRead                *ebpf.ProgramSpec `ebpf:"sys_enter_read"`
 	SysEnterWrite               *ebpf.ProgramSpec `ebpf:"sys_enter_write"`
@@ -113,6 +116,7 @@ type tlsTapperMapSpecs struct {
 	PidsMap              *ebpf.MapSpec `ebpf:"pids_map"`
 	SslReadContext       *ebpf.MapSpec `ebpf:"ssl_read_context"`
 	SslWriteContext      *ebpf.MapSpec `ebpf:"ssl_write_context"`
+	SysCloses            *ebpf.MapSpec `ebpf:"sys_closes"`
 }
 
 // tlsTapperObjects contains all objects after they have been loaded into the kernel.
@@ -146,6 +150,7 @@ type tlsTapperMaps struct {
 	PidsMap              *ebpf.Map `ebpf:"pids_map"`
 	SslReadContext       *ebpf.Map `ebpf:"ssl_read_context"`
 	SslWriteContext      *ebpf.Map `ebpf:"ssl_write_context"`
+	SysCloses            *ebpf.Map `ebpf:"sys_closes"`
 }
 
 func (m *tlsTapperMaps) Close() error {
@@ -162,6 +167,7 @@ func (m *tlsTapperMaps) Close() error {
 		m.PidsMap,
 		m.SslReadContext,
 		m.SslWriteContext,
+		m.SysCloses,
 	)
 }
 
@@ -182,6 +188,7 @@ type tlsTapperPrograms struct {
 	SslWrite                    *ebpf.Program `ebpf:"ssl_write"`
 	SslWriteEx                  *ebpf.Program `ebpf:"ssl_write_ex"`
 	SysEnterAccept4             *ebpf.Program `ebpf:"sys_enter_accept4"`
+	SysEnterClose               *ebpf.Program `ebpf:"sys_enter_close"`
 	SysEnterConnect             *ebpf.Program `ebpf:"sys_enter_connect"`
 	SysEnterRead                *ebpf.Program `ebpf:"sys_enter_read"`
 	SysEnterWrite               *ebpf.Program `ebpf:"sys_enter_write"`
@@ -204,6 +211,7 @@ func (p *tlsTapperPrograms) Close() error {
 		p.SslWrite,
 		p.SslWriteEx,
 		p.SysEnterAccept4,
+		p.SysEnterClose,
 		p.SysEnterConnect,
 		p.SysEnterRead,
 		p.SysEnterWrite,