Pass the connection address from write to gzip through a common address between gzip and dial

tap/tlstapper/bpf/golang_uprobes.c

@@ -57,6 +57,20 @@ static __always_inline int golang_crypto_tls_write_uprobe(struct pt_regs *ctx) {
     b->is_request = true;
     b->is_gzip_chunk = false;
 
+    struct socket x = {
+        .pid = s->pid,
+        .fd = ctx->rax,
+        .key_dial = s->key_dial,
+        .key_gzip = s->key_gzip,
+        .conn_addr = b->conn_addr,
+    };
+
+    __u64 key_gzip_full = (pid << 32) + s->key_gzip;
+    status = bpf_map_update_elem(&golang_write_to_gzip, &key_gzip_full, &x, BPF_ANY);
+    if (status != 0) {
+        bpf_printk("[golang_crypto_tls_write_uprobe] error updating gzip conn addr: %d", status);
+    }
+
     status = bpf_probe_read_str(&b->data, sizeof(b->data), (void*)ctx->rbx);
     if (status < 0) {
         bpf_printk("[golang_crypto_tls_write_uprobe] error reading data: %d", status);
@@ -114,23 +128,38 @@ static __always_inline int golang_crypto_tls_read_uprobe(struct pt_regs *ctx) {
 
 SEC("uprobe/golang_net_http_gzipreader_read")
 static __always_inline int golang_net_http_gzipreader_read_uprobe(struct pt_regs *ctx) {
+    void* stack_addr = (void*)ctx->rsp;
+    __u64 pid_tgid = bpf_get_current_pid_tgid();
+    __u64 pid = pid_tgid >> 32;
+    __u32 key_gzip;
+    // Address at ctx->rsp + 0x150 is common between golang_net_http_gzipreader_read_uprobe and golang_net_http_dialconn_uprobe
+    __u32 status = bpf_probe_read(&key_gzip, sizeof(key_gzip), stack_addr + 0x150);
+    if (status < 0) {
+        bpf_printk("[golang_net_http_gzipreader_read_uprobe] error reading key_gzip: %d", status);
+        return 0;
+    }
+
+    __u64 key_gzip_full = (pid << 32) + key_gzip;
+    struct socket *s = bpf_map_lookup_elem(&golang_write_to_gzip, &key_gzip_full);
+    if (s == NULL) {
+        bpf_printk("[golang_net_http_gzipreader_read_uprobe] error getting socket");
+        return 0;
+    }
+
     struct golang_read_write *b = NULL;
     b = bpf_ringbuf_reserve(&golang_read_writes, sizeof(struct golang_read_write), 0);
     if (!b) {
         return 0;
     }
 
-    // __u64 pid_tgid = bpf_get_current_pid_tgid();
-    // b->pid = pid_tgid >> 32;
-    // // ctx->rsi is common between golang_crypto_tls_write_uprobe and golang_crypto_tls_read_uprobe
-    // b->conn_addr = ctx->rsi; // go.itab.*net.TCPConn,net.Conn address
+    b->pid = pid;
+    b->conn_addr = s->conn_addr;
     b->is_request = false;
     b->is_gzip_chunk = true;
 
-    void* stack_addr = (void*)ctx->rsp;
     __u64 data_p;
     // Address at ctx->rsp + 0x8 holds the data
-    __u32 status = bpf_probe_read(&data_p, sizeof(data_p), stack_addr + 0x8);
+    status = bpf_probe_read(&data_p, sizeof(data_p), stack_addr + 0x8);
     if (status < 0) {
         bpf_printk("[golang_net_http_gzipreader_read_uprobe] error reading data pointer: %d", status);
         bpf_ringbuf_discard(b, BPF_RB_FORCE_WAKEUP);
@@ -160,7 +189,13 @@ static __always_inline int golang_net_socket_uprobe(struct pt_regs *ctx) {
         return 0;
     }
 
-    struct socket b = { .pid = s->pid, .fd = ctx->rax, .key_dial = s->key_dial };
+    struct socket b = {
+        .pid = s->pid,
+        .fd = ctx->rax,
+        .key_dial = s->key_dial,
+        .key_gzip = s->key_gzip,
+        .conn_addr = 0,
+    };
 
     __u64 key_dial_full = (pid << 32) + s->key_dial;
     __u32 status = bpf_map_update_elem(&golang_socket_to_write, &key_dial_full, &b, BPF_ANY);
@@ -182,8 +217,22 @@ static __always_inline int golang_net_http_dialconn_uprobe(struct pt_regs *ctx)
         return 0;
     }
 
+    __u32 key_gzip;
+    // Address at ctx->rsp + 0x1d0 is common between golang_net_http_gzipreader_read_uprobe and golang_net_http_dialconn_uprobe
+    status = bpf_probe_read(&key_gzip, sizeof(key_gzip), stack_addr + 0x1d0);
+    if (status < 0) {
+        bpf_printk("[golang_net_http_dialconn_uprobe] error reading key_gzip: %d", status);
+        return 0;
+    }
+
     __u64 pid_tgid = bpf_get_current_pid_tgid();
-    struct socket b = { .pid = pid_tgid >> 32, .fd = 0, .key_dial = key_dial };
+    struct socket b = {
+        .pid = pid_tgid >> 32,
+        .fd = 0,
+        .key_dial = key_dial,
+        .key_gzip = key_gzip,
+        .conn_addr = 0,
+    };
 
     __u64 pid = b.pid;
     // ctx->r14 is common between golang_net_socket_uprobe and golang_net_http_dialconn_uprobe

tap/tlstapper/bpf/include/maps.h

@@ -57,6 +57,8 @@ struct socket {
     __u32 pid;
     __u32 fd;
     __u64 key_dial;
+    __u64 key_gzip;
+    __u64 conn_addr;
 };
 
 #define BPF_MAP(_name, _type, _key_type, _value_type, _max_entries)     \
@@ -91,6 +93,7 @@ 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_LRU_HASH(golang_write_to_gzip, __u64, struct socket);
 BPF_RINGBUF(golang_read_writes);
 
 #endif /* __MAPS__ */

tap/tlstapper/golang_connection.go

@@ -4,10 +4,9 @@ type golangConnection struct {
 	Pid         uint32
 	ConnAddr    uint32
 	AddressPair addressPair
-	Request     []byte
-	Response    []byte
-	GotRequest  bool
-	GotResponse bool
+	Requests    [][]byte
+	Responses   [][]byte
+	Gzipped     bool
 }
 
 func NewGolangConnection(pid uint32, connAddr uint32) *golangConnection {

tap/tlstapper/tls_poller.go

@@ -22,7 +22,6 @@ import (
 	"github.com/up9inc/mizu/logger"
 	"github.com/up9inc/mizu/tap/api"
 	orderedmap "github.com/wk8/go-ordered-map"
-	"golang.org/x/sys/unix"
 )
 
 const (
@@ -138,12 +137,6 @@ func (p *tlsPoller) pollGolangReadWrite(rd *ringbuf.Reader, emitter api.Emitter,
 			continue
 		}
 
-		if b.IsGzipChunk {
-			chunk := unix.ByteSliceToString(b.Data[:])
-			fmt.Printf("chunk: %v\n", chunk)
-			continue
-		}
-
 		if p.golangReadWriteMap.Len()+1 > golangMapLimit {
 			pair := p.golangReadWriteMap.Oldest()
 			p.golangReadWriteMap.Delete(pair.Key)
@@ -162,6 +155,10 @@ func (p *tlsPoller) pollGolangReadWrite(rd *ringbuf.Reader, emitter api.Emitter,
 			connection = _connection.(*golangConnection)
 		}
 
+		if b.IsGzipChunk {
+			connection.Gzipped = true
+		}
+
 		if b.IsRequest {
 			err := connection.setAddressBySockfd(p.procfs, b.Pid, b.Fd)
 			if err != nil {
@@ -169,71 +166,77 @@ func (p *tlsPoller) pollGolangReadWrite(rd *ringbuf.Reader, emitter api.Emitter,
 				continue
 			}
 
-			connection.Request = make([]byte, len(b.Data[:]))
-			copy(connection.Request, b.Data[:])
-			connection.GotRequest = true
+			request := make([]byte, len(b.Data[:]))
+			copy(request, b.Data[:])
+			connection.Requests = append(connection.Requests, request)
 		} else {
-			connection.Response = make([]byte, len(b.Data[:]))
-			copy(connection.Response, b.Data[:])
-			connection.GotResponse = true
-		}
+			response := make([]byte, len(b.Data[:]))
+			copy(response, b.Data[:])
+			connection.Responses = append(connection.Responses, response)
 
-		if connection.GotRequest && connection.GotResponse {
-			// TODO: Remove these comments
-			// fmt.Printf("\n\nconnection.Pid: %v\n", connection.Pid)
-			// fmt.Printf("connection.ConnAddr: 0x%x\n", connection.ConnAddr)
-			// fmt.Printf("connection.AddressPair.srcIp: %v\n", connection.AddressPair.srcIp)
-			// fmt.Printf("connection.AddressPair.srcPort: %v\n", connection.AddressPair.srcPort)
-			// fmt.Printf("connection.AddressPair.dstIp: %v\n", connection.AddressPair.dstIp)
-			// fmt.Printf("connection.AddressPair.dstPort: %v\n", connection.AddressPair.dstPort)
-			// fmt.Printf("connection.Request:\n%v\n", unix.ByteSliceToString(connection.Request))
-			// fmt.Printf("connection.Response:\n%v\n", unix.ByteSliceToString(connection.Response))
+			if !b.IsGzipChunk {
+				// TODO: Remove these comments
+				// fmt.Printf("\n\nidentifier: %v\n", identifier)
+				// fmt.Printf("connection.Pid: %v\n", connection.Pid)
+				// fmt.Printf("connection.ConnAddr: 0x%x\n", connection.ConnAddr)
+				// fmt.Printf("connection.AddressPair.srcIp: %v\n", connection.AddressPair.srcIp)
+				// fmt.Printf("connection.AddressPair.srcPort: %v\n", connection.AddressPair.srcPort)
+				// fmt.Printf("connection.AddressPair.dstIp: %v\n", connection.AddressPair.dstIp)
+				// fmt.Printf("connection.AddressPair.dstPort: %v\n", connection.AddressPair.dstPort)
+				// fmt.Printf("connection.Gzipped: %v\n", connection.Gzipped)
+				// for i, x := range connection.Requests {
+				// 	fmt.Printf("connection.Request[%d]:\n%v\n", i, unix.ByteSliceToString(x))
+				// }
+				// for i, y := range connection.Responses {
+				// 	fmt.Printf("connection.Response[%d]:\n%v\n", i, unix.ByteSliceToString(y))
+				// }
 
-			tcpid := p.buildTcpId(&connection.AddressPair)
+				// tcpid := p.buildTcpId(&connection.AddressPair)
 
-			tlsEmitter := &tlsEmitter{
-				delegate:  emitter,
-				namespace: p.getNamespace(b.Pid),
-			}
+				// tlsEmitter := &tlsEmitter{
+				// 	delegate:  emitter,
+				// 	namespace: p.getNamespace(b.Pid),
+				// }
 
-			reader := &tlsReader{
-				chunks:        make(chan *tlsChunk, 1),
-				progress:      &api.ReadProgress{},
-				tcpID:         &tcpid,
-				isClient:      true,
-				captureTime:   time.Now(),
-				extension:     p.extension,
-				emitter:       tlsEmitter,
-				counterPair:   &api.CounterPair{},
-				reqResMatcher: p.extension.Dissector.NewResponseRequestMatcher(),
-			}
+				// reader := &tlsReader{
+				// 	chunks:        make(chan *tlsChunk, 1),
+				// 	progress:      &api.ReadProgress{},
+				// 	tcpID:         &tcpid,
+				// 	isClient:      true,
+				// 	captureTime:   time.Now(),
+				// 	extension:     p.extension,
+				// 	emitter:       tlsEmitter,
+				// 	counterPair:   &api.CounterPair{},
+				// 	reqResMatcher: p.extension.Dissector.NewResponseRequestMatcher(),
+				// }
 
-			stream := &tlsStream{
-				reader: reader,
-			}
-			streamsMap.Store(streamsMap.NextId(), stream)
+				// stream := &tlsStream{
+				// 	reader: reader,
+				// }
+				// streamsMap.Store(streamsMap.NextId(), stream)
 
-			reader.parent = stream
+				// reader.parent = stream
 
-			err := p.extension.Dissector.Dissect(bufio.NewReader(bytes.NewReader(connection.Request)), reader, options)
+				// err := p.extension.Dissector.Dissect(bufio.NewReader(bytes.NewReader(connection.Requests[0])), reader, options)
 
-			if err != nil {
-				logger.Log.Warningf("Error dissecting TLS %v - %v", reader.GetTcpID(), err)
-			}
+				// if err != nil {
+				// 	logger.Log.Warningf("Error dissecting TLS %v - %v", reader.GetTcpID(), err)
+				// }
 
-			reader.isClient = false
-			reader.tcpID = &api.TcpID{
-				SrcIP:   reader.tcpID.DstIP,
-				DstIP:   reader.tcpID.SrcIP,
-				SrcPort: reader.tcpID.DstPort,
-				DstPort: reader.tcpID.SrcPort,
-			}
-			reader.progress = &api.ReadProgress{}
+				// reader.isClient = false
+				// reader.tcpID = &api.TcpID{
+				// 	SrcIP:   reader.tcpID.DstIP,
+				// 	DstIP:   reader.tcpID.SrcIP,
+				// 	SrcPort: reader.tcpID.DstPort,
+				// 	DstPort: reader.tcpID.SrcPort,
+				// }
+				// reader.progress = &api.ReadProgress{}
 
-			err = p.extension.Dissector.Dissect(bufio.NewReader(bytes.NewReader(connection.Response)), reader, options)
+				// err = p.extension.Dissector.Dissect(bufio.NewReader(bytes.NewReader(connection.Responses[0])), reader, options)
 
-			if err != nil {
-				logger.Log.Warningf("Error dissecting TLS %v - %v", reader.GetTcpID(), err)
+				// if err != nil {
+				// 	logger.Log.Warningf("Error dissecting TLS %v - %v", reader.GetTcpID(), err)
+				// }
 			}
 		}
 	}

tap/tlstapper/tlstapper_bpfeb.go

@@ -108,6 +108,7 @@ type tlsTapperMapSpecs struct {
 	GolangDialToSocket   *ebpf.MapSpec `ebpf:"golang_dial_to_socket"`
 	GolangReadWrites     *ebpf.MapSpec `ebpf:"golang_read_writes"`
 	GolangSocketToWrite  *ebpf.MapSpec `ebpf:"golang_socket_to_write"`
+	GolangWriteToGzip    *ebpf.MapSpec `ebpf:"golang_write_to_gzip"`
 	Heap                 *ebpf.MapSpec `ebpf:"heap"`
 	LogBuffer            *ebpf.MapSpec `ebpf:"log_buffer"`
 	PidsMap              *ebpf.MapSpec `ebpf:"pids_map"`
@@ -141,6 +142,7 @@ type tlsTapperMaps struct {
 	GolangDialToSocket   *ebpf.Map `ebpf:"golang_dial_to_socket"`
 	GolangReadWrites     *ebpf.Map `ebpf:"golang_read_writes"`
 	GolangSocketToWrite  *ebpf.Map `ebpf:"golang_socket_to_write"`
+	GolangWriteToGzip    *ebpf.Map `ebpf:"golang_write_to_gzip"`
 	Heap                 *ebpf.Map `ebpf:"heap"`
 	LogBuffer            *ebpf.Map `ebpf:"log_buffer"`
 	PidsMap              *ebpf.Map `ebpf:"pids_map"`
@@ -157,6 +159,7 @@ func (m *tlsTapperMaps) Close() error {
 		m.GolangDialToSocket,
 		m.GolangReadWrites,
 		m.GolangSocketToWrite,
+		m.GolangWriteToGzip,
 		m.Heap,
 		m.LogBuffer,
 		m.PidsMap,

tap/tlstapper/tlstapper_bpfel.go

@@ -108,6 +108,7 @@ type tlsTapperMapSpecs struct {
 	GolangDialToSocket   *ebpf.MapSpec `ebpf:"golang_dial_to_socket"`
 	GolangReadWrites     *ebpf.MapSpec `ebpf:"golang_read_writes"`
 	GolangSocketToWrite  *ebpf.MapSpec `ebpf:"golang_socket_to_write"`
+	GolangWriteToGzip    *ebpf.MapSpec `ebpf:"golang_write_to_gzip"`
 	Heap                 *ebpf.MapSpec `ebpf:"heap"`
 	LogBuffer            *ebpf.MapSpec `ebpf:"log_buffer"`
 	PidsMap              *ebpf.MapSpec `ebpf:"pids_map"`
@@ -141,6 +142,7 @@ type tlsTapperMaps struct {
 	GolangDialToSocket   *ebpf.Map `ebpf:"golang_dial_to_socket"`
 	GolangReadWrites     *ebpf.Map `ebpf:"golang_read_writes"`
 	GolangSocketToWrite  *ebpf.Map `ebpf:"golang_socket_to_write"`
+	GolangWriteToGzip    *ebpf.Map `ebpf:"golang_write_to_gzip"`
 	Heap                 *ebpf.Map `ebpf:"heap"`
 	LogBuffer            *ebpf.Map `ebpf:"log_buffer"`
 	PidsMap              *ebpf.Map `ebpf:"pids_map"`
@@ -157,6 +159,7 @@ func (m *tlsTapperMaps) Close() error {
 		m.GolangDialToSocket,
 		m.GolangReadWrites,
 		m.GolangSocketToWrite,
+		m.GolangWriteToGzip,
 		m.Heap,
 		m.LogBuffer,
 		m.PidsMap,