Make golang_dial_writes key __u64 and include the PID

tap/tlstapper/bpf/golang_uprobes.c

@@ -26,6 +26,8 @@ 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;
+    __u64 pid_tgid = bpf_get_current_pid_tgid();
+    __u64 pid = pid_tgid >> 32;
     __u32 key_dial;
     // Address at ctx->rsp + 0x20 is common between golang_crypto_tls_write_uprobe and golang_net_http_dialconn_uprobe
     __u32 status = bpf_probe_read(&key_dial, sizeof(key_dial), stack_addr + 0x20);
@@ -34,7 +36,8 @@ static __always_inline int golang_crypto_tls_write_uprobe(struct pt_regs *ctx) {
         return 0;
     }
 
-    struct socket *s = bpf_map_lookup_elem(&golang_dial_writes, &key_dial);
+    __u64 key_dial_full = (pid << 32) + key_dial;
+    struct socket *s = bpf_map_lookup_elem(&golang_dial_writes, &key_dial_full);
     if (s == NULL) {
         bpf_printk("[golang_crypto_tls_write_uprobe] error getting socket");
         return 0;
@@ -46,8 +49,7 @@ static __always_inline int golang_crypto_tls_write_uprobe(struct pt_regs *ctx) {
         return 0;
     }
 
-    __u64 pid_tgid = bpf_get_current_pid_tgid();
-    b->pid = pid_tgid >> 32;
+    b->pid = pid;
     b->fd = s->fd;
     // 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
@@ -105,7 +107,8 @@ static __always_inline int golang_net_socket_uprobe(struct pt_regs *ctx) {
 
     struct socket b = { .pid = s->pid, .fd = ctx->rax, .key_dial = s->key_dial };
 
-    __u32 status = bpf_map_update_elem(&golang_dial_writes, &s->key_dial, &b, BPF_ANY);
+    __u64 key_dial_full = (pid << 32) + s->key_dial;
+    __u32 status = bpf_map_update_elem(&golang_dial_writes, &key_dial_full, &b, BPF_ANY);
     if (status != 0) {
         bpf_printk("[golang_net_socket_uprobe] error updating socket file descriptor: %d", status);
     }

tap/tlstapper/bpf/include/maps.h

@@ -56,7 +56,7 @@ struct fd_info {
 struct socket {
     __u32 pid;
     __u32 fd;
-    __u32 key_dial;
+    __u64 key_dial;
 };
 
 #define BPF_MAP(_name, _type, _key_type, _value_type, _max_entries)     \
@@ -90,7 +90,7 @@ BPF_PERF_OUTPUT(chunks_buffer);
 BPF_PERF_OUTPUT(log_buffer);
 
 BPF_LRU_HASH(golang_socket_dials, __u64, struct socket);
-BPF_LRU_HASH(golang_dial_writes, __u32, struct socket);
+BPF_LRU_HASH(golang_dial_writes, __u64, struct socket);
 BPF_RINGBUF(golang_read_writes);
 
 #endif /* __MAPS__ */