Extend packet at the sock msg level (#1740)

Author: rafaelroquetto

bpf/go_nethttp.h

@@ -1153,6 +1153,7 @@ int beyla_uprobe_persistConnRoundTrip(struct pt_regs *ctx) {
                 tp_info_pid_t tp_p = {
                     .pid = pid,
                     .valid = 1,
+                    .written = 0,
                 };
 
                 tp_clone(&tp_p.tp, &invocation->tp);

bpf/http_ssl_defs.h

@@ -204,11 +204,12 @@ handle_ssl_buf(void *ctx, u64 id, ssl_args_t *args, int bytes_len, u8 direction)
     }
 }
 
-static __always_inline void set_active_ssl_connection(pid_connection_info_t *conn, void *ssl) {
+static __always_inline void set_active_ssl_connection(const pid_connection_info_t *conn,
+                                                      void *ssl) {
     bpf_map_update_elem(&active_ssl_connections, conn, &ssl, BPF_ANY);
 }
 
-static __always_inline void *is_ssl_connection(u64 id, pid_connection_info_t *conn) {
+static __always_inline void *is_ssl_connection(u64 id, const pid_connection_info_t *conn) {
     void *ssl = 0;
     // Checks if it's sandwitched between read or write uprobe/uretprobe
     ssl_args_t *ssl_args = bpf_map_lookup_elem(&active_ssl_read_args, &id);

bpf/http_types.h

@@ -99,6 +99,7 @@ typedef struct tp_info_pid {
     tp_info_t tp;
     u32 pid;
     u8 valid;
+    u8 written;
     u8 req_type;
 } tp_info_pid_t;
 
@@ -300,4 +301,4 @@ static __always_inline u8 is_http_request_buf(const unsigned char *p) {
     );
 }
 
-#endif
+#endif

bpf/k_tracer.h

@@ -264,7 +264,7 @@ int BPF_KPROBE(beyla_kprobe_tcp_sendmsg, struct sock *sk, struct msghdr *msg, si
         dbg_print_http_connection_info(
             &s_args.p_conn.conn); // commented out since GitHub CI doesn't like this call
         // Create the egress key before we sort the connection info.
-        egress_key_t e_key = {
+        const egress_key_t e_key = {
             .d_port = s_args.p_conn.conn.d_port,
             .s_port = s_args.p_conn.conn.s_port,
         };
@@ -348,7 +348,7 @@ int BPF_KPROBE(beyla_kprobe_tcp_rate_check_app_limited, struct sock *sk) {
     if (parse_sock_info(sk, &s_args.p_conn.conn)) {
         u16 orig_dport = s_args.p_conn.conn.d_port;
         dbg_print_http_connection_info(&s_args.p_conn.conn);
-        egress_key_t e_key = {
+        const egress_key_t e_key = {
             .d_port = s_args.p_conn.conn.d_port,
             .s_port = s_args.p_conn.conn.s_port,
         };
@@ -505,7 +505,7 @@ static __always_inline int return_recvmsg(void *ctx, u64 id, int copied_len) {
     bpf_map_delete_elem(&active_recv_args, &id);
 
     if (parse_sock_info((struct sock *)sock_ptr, &info.conn)) {
-        u16 orig_dport = info.conn.d_port;
+        const u16 orig_dport = info.conn.d_port;
         //dbg_print_http_connection_info(&info.conn);
         sort_connection_info(&info.conn);
         info.pid = pid_from_pid_tgid(id);

bpf/protocol_http.h

@@ -49,7 +49,30 @@ static __always_inline void http_get_or_create_trace_info(http_connection_metada
                                                           int bytes_len,
                                                           s32 capture_header_buffer,
                                                           u8 ssl) {
-    tp_info_pid_t *tp_p = tp_buf();
+    //TODO use make_key
+    egress_key_t e_key = {
+        .d_port = conn->d_port,
+        .s_port = conn->s_port,
+    };
+
+    sort_egress_key(&e_key);
+
+    tp_info_pid_t *tp_p = bpf_map_lookup_elem(&outgoing_trace_map, &e_key);
+
+    // TODO move this to sock msg
+    if (tp_p && tp_p->req_type == EVENT_HTTP_CLIENT && tp_p->written && tp_p->pid == pid) {
+        bpf_dbg_printk("found tp info previously set by sock msg");
+        // we've already got a tp_info_pid_t setup by the sockmsg program, use
+        // that instead
+
+        set_trace_info_for_connection(conn, TRACE_TYPE_CLIENT, tp_p);
+
+        // clean up so that TC does not pick it up
+        bpf_map_delete_elem(&outgoing_trace_map, &e_key);
+        return;
+    }
+
+    tp_p = tp_buf();
 
     if (!tp_p) {
         return;
@@ -58,6 +81,7 @@ static __always_inline void http_get_or_create_trace_info(http_connection_metada
     tp_p->tp.ts = bpf_ktime_get_ns();
     tp_p->tp.flags = 1;
     tp_p->valid = 1;
+    tp_p->written = 0;
     tp_p->pid = pid; // used for avoiding finding stale server requests with client port reuse
     tp_p->req_type = (meta) ? meta->type : 0;
 
@@ -82,7 +106,7 @@ static __always_inline void http_get_or_create_trace_info(http_connection_metada
     if (!found_tp) {
         bpf_dbg_printk("Generating new traceparent id");
         new_trace_id(&tp_p->tp);
-        __builtin_memset(tp_p->tp.parent_id, 0, sizeof(tp_p->tp.span_id));
+        __builtin_memset(tp_p->tp.parent_id, 0, sizeof(tp_p->tp.parent_id));
     } else {
         bpf_dbg_printk("Using old traceparent id");
     }

bpf/tc_common.h

@@ -11,6 +11,7 @@ const char INV_TP[] = "W3C-BeylaID: 00-00000000000000000000000000000000-00000000
 const u32 EXTEND_SIZE = sizeof(TP) - 1;
 const char TP_PREFIX[] = "Traceparent: ";
 const u32 TP_PREFIX_SIZE = sizeof(TP_PREFIX) - 1;
+const u32 INVALID_POS = 0xffffffff;
 
 static __always_inline unsigned char *
 memchar(unsigned char *haystack, char needle, const unsigned char *end, u32 size) {
@@ -32,22 +33,24 @@ find_first_of(unsigned char *begin, unsigned char *end, char ch) {
     return memchar(begin, ch, end, MAX_INLINE_LEN);
 }
 
-static __always_inline int
-memchar_pos(unsigned const char *haystack, char needle, const unsigned char *end, u32 size) {
+static __always_inline u32 memchar_pos(unsigned char *haystack,
+                                       char needle,
+                                       const unsigned char *end,
+                                       u32 size) {
     for (u32 i = 0; i < size; ++i) {
-        if (&haystack[i] >= end) {
-            break;
-        }
+        unsigned char *ptr = haystack + i;
 
-        if (haystack[i] == needle) {
+        if (ptr + 1 >= end) {
+            break;
+        } else if (ptr && *ptr == needle) {
             return i;
         }
     }
 
-    return -1;
+    return INVALID_POS;
 }
 
-static __always_inline int find_first_pos_of(unsigned char *begin, unsigned char *end, char ch) {
+static __always_inline u32 find_first_pos_of(unsigned char *begin, unsigned char *end, char ch) {
     return memchar_pos(begin, ch, end, MAX_INLINE_LEN);
 }
 
@@ -71,6 +74,26 @@ static __always_inline void *ctx_data_end(struct __sk_buff *ctx) {
     return data_end;
 }
 
+static __always_inline void *ctx_msg_data(struct sk_msg_md *ctx) {
+    void *data;
+
+    asm("%[res] = *(u64 *)(%[base] + %[offset])"
+        : [res] "=r"(data)
+        : [base] "r"(ctx), [offset] "i"(offsetof(struct sk_msg_md, data)), "m"(*ctx));
+
+    return data;
+}
+
+static __always_inline void *ctx_msg_data_end(struct sk_msg_md *ctx) {
+    void *data_end;
+
+    asm("%[res] = *(u64 *)(%[base] + %[offset])"
+        : [res] "=r"(data_end)
+        : [base] "r"(ctx), [offset] "i"(offsetof(struct sk_msg_md, data_end)), "m"(*ctx));
+
+    return data_end;
+}
+
 static __always_inline void
 sk_msg_read_remote_ip6(struct sk_msg_md *ctx, u32 *res) { //NOLINT(readability-non-const-parameter)
     asm("%[res0] = *(u32 *)(%[base] + %[offset] + 0)\n"