Fix context propagation on old kernels (#1576)

bpf/http_types.h

@@ -25,6 +25,9 @@
 #define K_TCP_MAX_LEN 256
 #define K_TCP_RES_LEN 128
 
+// Max of HTTP, HTTP2/GRPC and TCP buffers. Used in sk_msg
+#define MAX_PROTOCOL_BUF_SIZE 256
+
 #define CONN_INFO_FLAG_TRACE 0x1
 
 #define TRACE_ID_SIZE_BYTES 16

bpf/k_tracer.h

@@ -224,6 +224,20 @@ int BPF_KRETPROBE(beyla_kretprobe_sys_connect, int fd) {
     return 0;
 }
 
+static __always_inline void
+tcp_send_ssl_check(u64 id, void *ssl, pid_connection_info_t *p_conn, u16 orig_dport) {
+    bpf_dbg_printk("=== kprobe SSL tcp_sendmsg=%d ssl=%llx ===", id, ssl);
+    ssl_pid_connection_info_t *s_conn = bpf_map_lookup_elem(&ssl_to_conn, &ssl);
+    if (s_conn) {
+        finish_possible_delayed_tls_http_request(&s_conn->p_conn, ssl);
+    }
+    ssl_pid_connection_info_t ssl_conn = {
+        .orig_dport = orig_dport,
+    };
+    __builtin_memcpy(&ssl_conn.p_conn, p_conn, sizeof(pid_connection_info_t));
+    bpf_map_update_elem(&ssl_to_conn, &ssl, &ssl_conn, BPF_ANY);
+}
+
 // Main HTTP read and write operations are handled with tcp_sendmsg and tcp_recvmsg
 
 // The size argument here will be always the total response size.
@@ -311,16 +325,67 @@ int BPF_KPROBE(beyla_kprobe_tcp_sendmsg, struct sock *sk, struct msghdr *msg, si
             return 0;
         }
 
-        bpf_dbg_printk("=== kprobe SSL tcp_sendmsg=%d sock=%llx ssl=%llx ===", id, sk, ssl);
-        ssl_pid_connection_info_t *s_conn = bpf_map_lookup_elem(&ssl_to_conn, &ssl);
-        if (s_conn) {
-            finish_possible_delayed_tls_http_request(&s_conn->p_conn, ssl);
-        }
-        ssl_pid_connection_info_t ssl_conn = {
-            .orig_dport = orig_dport,
+        tcp_send_ssl_check(id, ssl, &s_args.p_conn, orig_dport);
+    }
+
+    return 0;
+}
+
+// This is a backup path kprobe in case tcp_sendmsg doesn't fire, which
+// happens on certain kernels if sk_msg is attached.
+SEC("kprobe/tcp_rate_check_app_limited")
+int BPF_KPROBE(beyla_kprobe_tcp_rate_check_app_limited, struct sock *sk) {
+    u64 id = bpf_get_current_pid_tgid();
+
+    if (!valid_pid(id)) {
+        return 0;
+    }
+
+    bpf_dbg_printk("=== kprobe tcp_rate_check_app_limited=%d sock=%llx ===", id, sk);
+
+    send_args_t s_args = {};
+
+    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 = {
+            .d_port = s_args.p_conn.conn.d_port,
+            .s_port = s_args.p_conn.conn.s_port,
         };
-        __builtin_memcpy(&ssl_conn.p_conn, &s_args.p_conn, sizeof(pid_connection_info_t));
-        bpf_map_update_elem(&ssl_to_conn, &ssl, &ssl_conn, BPF_ANY);
+
+        sort_connection_info(&s_args.p_conn.conn);
+        s_args.p_conn.pid = pid_from_pid_tgid(id);
+
+        msg_buffer_t *m_buf = bpf_map_lookup_elem(&msg_buffers, &e_key);
+        if (m_buf) {
+            u8 *buf = m_buf->buf;
+            // The buffer setup for us by a sock_msg program is always the
+            // full buffer, but when we extend a packet to be able to inject
+            // a Traceparent field, it will actually be split in 3 chunks:
+            // [before the injected header],[70 bytes for 'Traceparent...'],[the rest].
+            // We don't want the handle_buf_with_connection logic to run more than
+            // once on the same data, so if we find a buf we send all of it to the
+            // handle_buf_with_connection logic and then mark it as seen by making
+            // m_buf->pos be the size of the buffer.
+            if (!m_buf->pos) {
+                u16 size = sizeof(m_buf->buf);
+                m_buf->pos = size;
+                s_args.size = size;
+                bpf_dbg_printk("msg_buffer: size %d, buf[%s]", size, buf);
+                u64 sock_p = (u64)sk;
+                bpf_map_update_elem(&active_send_args, &id, &s_args, BPF_ANY);
+                bpf_map_update_elem(&active_send_sock_args, &sock_p, &s_args, BPF_ANY);
+
+                // Logically last for !ssl.
+                handle_buf_with_connection(
+                    ctx, &s_args.p_conn, buf, size, NO_SSL, TCP_SEND, orig_dport);
+            }
+        }
+
+        void *ssl = is_ssl_connection(id, &s_args.p_conn);
+        if (ssl) {
+            tcp_send_ssl_check(id, ssl, &s_args.p_conn, orig_dport);
+        }
     }
 
     return 0;
@@ -540,7 +605,7 @@ int beyla_socket__http_filter(struct __sk_buff *skb) {
             }
             read_skb_bytes(skb, tcp.hdr_len, info->buf, full_len);
             u64 cookie = bpf_get_socket_cookie(skb);
-            //bpf_dbg_printk("=== http_filter cookie = %llx, tcp_seq=%d len=%d %s ===", cookie, tcp.seq, len, buf);
+            //bpf_printk("=== http_filter cookie = %llx, len=%d %s ===", cookie, len, buf);
             //dbg_print_http_connection_info(&conn);
             set_fallback_http_info(info, &conn, skb->len - tcp.hdr_len);
 

bpf/tc_sock.h

@@ -205,18 +205,18 @@ static __always_inline u8 protocol_detector(struct sk_msg_md *msg,
             msg_buffer_t msg_buf = {
                 .pos = 0,
             };
-            bpf_probe_read_kernel(msg_buf.buf, FULL_BUF_SIZE, msg->data);
+            bpf_probe_read_kernel(msg_buf.buf, MAX_PROTOCOL_BUF_SIZE, msg->data);
             // We setup any call that looks like HTTP request to be extended.
             // This must match exactly to what the decision will be for
             // the kprobe program on tcp_sendmsg, which sets up the
             // outgoing_trace_map data used by Traffic Control to write the
             // actual 'Traceparent:...' string.
+            if (bpf_map_update_elem(&msg_buffers, &e_key, &msg_buf, BPF_ANY)) {
+                // fail if we can't setup a msg buffer
+                return 0;
+            }
             if (is_http_request_buf((const unsigned char *)msg_buf.buf)) {
                 bpf_dbg_printk("Setting up request to be extended");
-                if (bpf_map_update_elem(&msg_buffers, &e_key, &msg_buf, BPF_ANY)) {
-                    // fail if we can't setup a msg buffer
-                    return 0;
-                }
 
                 return 1;
             }
@@ -258,6 +258,8 @@ int beyla_packet_extender(struct sk_msg_md *msg) {
     }
     bpf_msg_pull_data(msg, 0, 1024, 0);
 
+    // TODO: execute the protocol handlers here with tail calls, don't
+    // rely on tcp_sendmsg to do it and record these message buffers.
     if (!tracked) {
         // If we didn't have metadata (sock_msg runs before the kprobe),
         // we ensure to mark it for any packet we want to extend.
@@ -300,6 +302,7 @@ int beyla_packet_extender(struct sk_msg_md *msg) {
                             bad_ctx->seen = bad_ctx->offset;
                         }
                     }
+                } else {
                     bpf_dbg_printk("offset to split %d", newline_pos);
                 }
             }