bpf: implement large buffer action, improve mysql protocol inference (#318)

mmat11 · web-flow · commit 8c997d25995e · 2025-07-21T15:13:59.000-04:00
diff --git a/bpf/common/common.h b/bpf/common/common.h
@@ -47,6 +47,11 @@ enum {
     k_mysql_error_message_max_mask = k_mysql_error_message_max - 1
 };
 
+enum large_buf_action : u8 {
+    k_large_buf_action_init = 0,
+    k_large_buf_action_append = 1,
+};
+
 #define MAX_SPAN_NAME_LEN 64
 #define MAX_STATUS_DESCRIPTION_LEN 64
 
@@ -147,7 +152,8 @@ typedef struct tcp_req {
 typedef struct tcp_large_buffer {
     u8 type; // Must be first
     u8 direction;
-    u8 _pad[2];
+    enum large_buf_action action;
+    u8 _pad;
     u32 len;
     tp_info_t tp;
     u8 buf[];
@@ -223,4 +229,4 @@ typedef struct otel_span {
     pid_info pid;
     otel_attributes_t span_attrs;
     u8 _epad[6];
-} otel_span_t;
+} otel_span_t;
diff --git a/bpf/common/sql.h b/bpf/common/sql.h
@@ -6,7 +6,7 @@
 #include <common/strings.h>
 
 enum {
-    k_max_query_offset = 4,
+    k_max_query_offset = 8,
 
     k_max_sql_op_len = 6, // Maximum length of SQL operation names (e.g., "SELECT")
 };
diff --git a/bpf/generictracer/k_tracer.c b/bpf/generictracer/k_tracer.c
@@ -956,7 +956,7 @@ int obi_handle_buf_with_args(void *ctx) {
                         &args->packet_type,
                         &args->protocol_type)) {
         bpf_dbg_printk("Found mysql connection");
-        bpf_tail_call(ctx, &jump_table, k_tail_protocol_mysql);
+        bpf_tail_call(ctx, &jump_table, k_tail_protocol_tcp);
     } else { // large request tracking and generic TCP
         http_info_t *info = bpf_map_lookup_elem(&ongoing_http, &args->pid_conn);
 
@@ -1013,4 +1013,4 @@ int obi_handle_buf_with_args(void *ctx) {
     }
 
     return 0;
-}
+}
diff --git a/bpf/generictracer/k_tracer_tailcall.h b/bpf/generictracer/k_tracer_tailcall.h
@@ -17,6 +17,5 @@ enum {
     k_tail_protocol_http2_grpc_frames = 3,
     k_tail_protocol_http2_grpc_handle_start_frame = 4,
     k_tail_protocol_http2_grpc_handle_end_frame = 5,
-    k_tail_protocol_mysql = 6,
-    k_tail_handle_buf_with_args = 7,
+    k_tail_handle_buf_with_args = 6,
 };
diff --git a/bpf/generictracer/protocol_mysql.h b/bpf/generictracer/protocol_mysql.h
@@ -56,6 +56,9 @@ enum {
     // Large buffer
     k_large_buf_max_size = 1 << 14, // 16K
     k_large_buf_max_size_mask = k_large_buf_max_size - 1,
+
+    // Sanity checks
+    k_mysql_payload_length_max = 1 << 13, // 8K
 };
 
 struct {
@@ -67,43 +70,50 @@ struct {
 
 SCRATCH_MEM_SIZED(mysql_large_buffers, k_large_buf_max_size);
 
+// This function is used to store the MySQL header if it comes in split packets
+// from double send.
+// Given the fact that we need to store this for the duration of the full request
+// (split in potentially multiple packets), we will **not** process or preserve
+// any actual payloads that are exactly 4 bytes long — they are intentionally
+// dropped in favor of state storage.
 static __always_inline int mysql_store_state_data(const connection_info_t *conn_info,
                                                   const unsigned char *data,
                                                   size_t data_len) {
     if (data_len != k_mysql_hdr_without_command_size) {
         return 0;
     }
 
-    struct mysql_state_data *state_data = bpf_map_lookup_elem(&mysql_state, conn_info);
-    if (state_data == NULL) {
-        // State data not found, treat this data as a header.
-        struct mysql_state_data new_state_data = {};
-        bpf_probe_read(&new_state_data, k_mysql_hdr_without_command_size, (const void *)data);
-        bpf_map_update_elem(&mysql_state, conn_info, &new_state_data, BPF_ANY);
-        return -1;
-    }
+    struct mysql_state_data new_state_data = {};
+    bpf_probe_read(&new_state_data, k_mysql_hdr_without_command_size, (const void *)data);
+    bpf_map_update_elem(&mysql_state, conn_info, &new_state_data, BPF_ANY);
 
-    // This is a payload.
-    return 0;
+    return -1;
 }
 
 static __always_inline int mysql_parse_fixup_header(const connection_info_t *conn_info,
                                                     struct mysql_hdr *hdr,
                                                     const unsigned char *data,
                                                     size_t data_len) {
+    // Try to parse and validate the header first.
+    bpf_probe_read(hdr, k_mysql_hdr_size, (const void *)data);
+    if (mysql_payload_length(hdr->payload_length) ==
+        (data_len - k_mysql_hdr_without_command_size)) {
+        // Header is valid and we have the full data, we can proceed.
+        hdr->hdr_arrived = false;
+        return 0;
+    }
+
+    // Prepend the header from state data.
     struct mysql_state_data *state_data = bpf_map_lookup_elem(&mysql_state, conn_info);
     if (state_data != NULL) {
         __builtin_memcpy(hdr, state_data, k_mysql_hdr_without_command_size);
         bpf_probe_read(&hdr->command_id, k_mysql_hdr_command_id_size, (const void *)data);
         hdr->hdr_arrived = true;
-    } else {
-        if (data_len < k_mysql_hdr_size) {
-            bpf_dbg_printk("mysql_parse_fixup_header: data_len is too short: %d", data_len);
-            return -1;
-        }
-        bpf_probe_read(hdr, k_mysql_hdr_size, (const void *)data);
+        return 0;
     }
-    return 0;
+
+    bpf_dbg_printk("mysql_parse_fixup_header: failed to parse mysql header");
+    return -1;
 }
 
 // This is an alternative version of mysql_parse_fixup_header that fills the buffer
@@ -140,43 +150,44 @@ static __always_inline int mysql_read_fixup_buffer(const connection_info_t *conn
     return *buf_len;
 }
 
-static __always_inline void mysql_send_large_buffer(tcp_req_t *req,
-                                                    pid_connection_info_t *pid_conn,
-                                                    const void *u_buf,
-                                                    u32 bytes_len,
-                                                    u8 direction) {
+// Emit a large buffer event for MySQL protocol.
+// The return value is used to control the flow for this specific protocol.
+// -1: wait additional data; 0: continue, regardless of errors.
+static __always_inline int mysql_send_large_buffer(tcp_req_t *req,
+                                                   pid_connection_info_t *pid_conn,
+                                                   const void *u_buf,
+                                                   u32 bytes_len,
+                                                   u8 direction,
+                                                   enum large_buf_action action) {
     if (mysql_store_state_data(&pid_conn->conn, u_buf, bytes_len) < 0) {
         bpf_dbg_printk("mysql_send_large_buffer: 4 bytes packet, storing state data");
-        return;
-    }
-
-    if (bytes_len < (k_mysql_hdr_size + 1)) {
-        bpf_dbg_printk("mysql_send_large_buffer: bytes_len is too short: %d", bytes_len);
-        return;
+        return -1;
     }
 
     tcp_large_buffer_t *large_buf = (tcp_large_buffer_t *)mysql_large_buffers_mem();
     if (!large_buf) {
         bpf_dbg_printk("mysql_send_large_buffer: failed to reserve space for MySQL large buffer");
-        return;
+        return 0;
     }
 
     large_buf->type = EVENT_TCP_LARGE_BUFFER;
     large_buf->direction = direction;
+    large_buf->action = action;
     __builtin_memcpy((void *)&large_buf->tp, (void *)&req->tp, sizeof(tp_info_t));
 
     int written =
         mysql_read_fixup_buffer(&pid_conn->conn, large_buf->buf, &large_buf->len, u_buf, bytes_len);
     if (written < 0) {
         bpf_dbg_printk("mysql_send_large_buffer: failed to read buffer, not sending large buffer");
-        return;
+        return 0;
     }
 
     req->has_large_buffers = true;
     bpf_ringbuf_output(&events,
                        large_buf,
                        (sizeof(tcp_large_buffer_t) + written) & k_large_buf_max_size_mask,
                        get_flags());
+    return 0;
 }
 
 static __always_inline u32 data_offset(struct mysql_hdr *hdr) {
@@ -188,32 +199,6 @@ static __always_inline u32 mysql_command_offset(struct mysql_hdr *hdr) {
     return data_offset(hdr) - k_mysql_hdr_command_id_size;
 }
 
-// k_tail_protocol_mysql
-SEC("kprobe/mysql")
-int obi_protocol_mysql(void *ctx) {
-    call_protocol_args_t *args = protocol_args();
-    if (!args) {
-        return 0;
-    }
-
-    bpf_dbg_printk("=== tcp_mysql_event len=%d pid=%d ===",
-                   args->bytes_len,
-                   pid_from_pid_tgid(bpf_get_current_pid_tgid()));
-
-    if (mysql_store_state_data(
-            &args->pid_conn.conn, (const unsigned char *)args->u_buf, args->bytes_len) < 0) {
-        bpf_dbg_printk("mysql: 4 bytes packet, storing state data");
-        return 0;
-    }
-
-    // Tail call back into generic TCP handler.
-    // Once the header is fixed up, we can use the generic TCP handling code
-    // in order to reuse all the common logic.
-    bpf_tail_call(ctx, &jump_table, k_tail_protocol_tcp);
-
-    return 0;
-}
-
 static __always_inline u8 is_mysql(connection_info_t *conn_info,
                                    const unsigned char *data,
                                    u32 data_len,
@@ -224,19 +209,20 @@ static __always_inline u8 is_mysql(connection_info_t *conn_info,
         return 0;
     }
 
-    if (data_len < (k_mysql_hdr_size + 1)) {
-        bpf_dbg_printk("is_mysql: data_len is too short: %d", data_len);
-        return 0;
-    }
-
     struct mysql_hdr hdr = {};
     if (mysql_parse_fixup_header(conn_info, &hdr, data, data_len) != 0) {
         bpf_dbg_printk("is_mysql: failed to parse mysql header");
         return 0;
     }
+    const u32 payload_len = mysql_payload_length(hdr.payload_length);
+
+    if (payload_len > k_mysql_payload_length_max) {
+        bpf_dbg_printk("is_mysql: payload length is too large: %d", payload_len);
+        return 0;
+    }
 
     bpf_dbg_printk("is_mysql: payload_length=%d sequence_id=%d command_id=%d",
-                   mysql_payload_length(hdr.payload_length),
+                   payload_len,
                    hdr.sequence_id,
                    hdr.command_id);
 
@@ -274,8 +260,9 @@ static __always_inline u8 is_mysql(connection_info_t *conn_info,
             // NOTE: Trying to classify the connection based on this command
             // would be unreliable, as the check is too shallow.
             *packet_type = PACKET_TYPE_REQUEST;
+            break;
         }
-        break;
+        return 0;
     default:
         if (*protocol_type == k_protocol_type_mysql) {
             // Check sequence ID and make sure we are processing a response.
diff --git a/bpf/generictracer/protocol_tcp.h b/bpf/generictracer/protocol_tcp.h
@@ -96,21 +96,26 @@ static __always_inline void cleanup_tcp_trace_info_if_needed(pid_connection_info
     }
 }
 
-static __always_inline void tcp_send_large_buffer(tcp_req_t *req,
-                                                  pid_connection_info_t *pid_conn,
-                                                  void *u_buf,
-                                                  int bytes_len,
-                                                  u8 direction,
-                                                  enum protocol_type protocol_type) {
+static __always_inline int tcp_send_large_buffer(tcp_req_t *req,
+                                                 pid_connection_info_t *pid_conn,
+                                                 void *u_buf,
+                                                 int bytes_len,
+                                                 u8 direction,
+                                                 enum protocol_type protocol_type,
+                                                 enum large_buf_action action) {
+    int ret = 0;
+
     switch (protocol_type) {
     case k_protocol_type_mysql:
         if (mysql_buffer_size > 0) {
-            mysql_send_large_buffer(req, pid_conn, u_buf, bytes_len, direction);
+            ret = mysql_send_large_buffer(req, pid_conn, u_buf, bytes_len, direction, action);
         }
         break;
     case k_protocol_type_unknown:
         break;
     }
+
+    return ret;
 }
 
 static __always_inline void handle_unknown_tcp_connection(pid_connection_info_t *pid_conn,
@@ -183,12 +188,22 @@ static __always_inline void handle_unknown_tcp_connection(pid_connection_info_t
 
             tcp_get_or_set_trace_info(req, pid_conn, ssl, orig_dport);
 
-            tcp_send_large_buffer(req, pid_conn, u_buf, bytes_len, direction, protocol_type);
+            tcp_send_large_buffer(
+                req, pid_conn, u_buf, bytes_len, direction, protocol_type, k_large_buf_action_init);
 
             bpf_map_update_elem(&ongoing_tcp_req, pid_conn, req, BPF_ANY);
         }
     } else if (existing->direction != direction) {
-        tcp_send_large_buffer(existing, pid_conn, u_buf, bytes_len, direction, protocol_type);
+        if (tcp_send_large_buffer(existing,
+                                  pid_conn,
+                                  u_buf,
+                                  bytes_len,
+                                  direction,
+                                  protocol_type,
+                                  k_large_buf_action_init) < 0) {
+            bpf_dbg_printk("handle_unknown_tcp_connection: waiting additional response data");
+            return;
+        }
 
         if (existing->end_monotime_ns == 0) {
             bpf_clamp_umax(bytes_len, K_TCP_RES_LEN);
@@ -220,7 +235,14 @@ static __always_inline void handle_unknown_tcp_connection(pid_connection_info_t
         existing->len += bytes_len;
         existing->req_len = existing->len;
         existing->protocol_type = protocol_type;
-        tcp_send_large_buffer(existing, pid_conn, u_buf, bytes_len, direction, protocol_type);
+
+        tcp_send_large_buffer(existing,
+                              pid_conn,
+                              u_buf,
+                              bytes_len,
+                              direction,
+                              protocol_type,
+                              k_large_buf_action_append);
     } else {
         existing->req_len += bytes_len;
     }
diff --git a/pkg/components/ebpf/common/common.go b/pkg/components/ebpf/common/common.go
@@ -118,7 +118,7 @@ type MisclassifiedEvent struct {
 type EBPFParseContext struct {
 	h2c               *lru.Cache[uint64, h2Connection]
 	redisDBCache      *simplelru.LRU[BpfConnectionInfoT, int]
-	largeBuffers      *expirable.LRU[largeBufferKey, largeBuffer]
+	largeBuffers      *expirable.LRU[largeBufferKey, *largeBuffer]
 	mongoRequestCache *PendingMongoDBRequests
 }
 
@@ -138,7 +138,7 @@ func ptlog() *slog.Logger { return slog.With("component", "ebpf.ProcessTracer")
 func NewEBPFParseContext(cfg *config.EBPFTracer) *EBPFParseContext {
 	var redisDBCache *simplelru.LRU[BpfConnectionInfoT, int]
 	h2c, _ := lru.New[uint64, h2Connection](1024 * 10)
-	largeBuffers := expirable.NewLRU[largeBufferKey, largeBuffer](1024, nil, 5*time.Minute)
+	largeBuffers := expirable.NewLRU[largeBufferKey, *largeBuffer](1024, nil, 5*time.Minute)
 
 	if cfg != nil && cfg.RedisDBCache.Enabled {
 		var err error
@@ -189,7 +189,7 @@ func ReadBPFTraceAsSpan(parseCtx *EBPFParseContext, cfg *config.EBPFTracer, reco
 	case EventTypeGoKafkaGo:
 		return ReadGoKafkaGoRequestIntoSpan(record)
 	case EventTypeTCPLargeBuffer:
-		return setTCPLargeBuffer(parseCtx, record)
+		return appendTCPLargeBuffer(parseCtx, record)
 	case EventOTelSDKGo:
 		return ReadGoOTelEventIntoSpan(record)
 	}
diff --git a/pkg/components/ebpf/common/tcp_detect_transform.go b/pkg/components/ebpf/common/tcp_detect_transform.go
@@ -41,10 +41,10 @@ func ReadTCPRequestIntoSpan(parseCtx *EBPFParseContext, cfg *config.EBPFTracer,
 	responseBuffer = event.Rbuf[:l]
 
 	if event.HasLargeBuffers == 1 {
-		if b, ok := getTCPLargeBuffer(parseCtx, event.Tp.TraceId, event.Tp.SpanId, 0); ok {
+		if b, ok := extractTCPLargeBuffer(parseCtx, event.Tp.TraceId, event.Tp.SpanId, 0); ok {
 			requestBuffer = b
 		}
-		if b, ok := getTCPLargeBuffer(parseCtx, event.Tp.TraceId, event.Tp.SpanId, 1); ok {
+		if b, ok := extractTCPLargeBuffer(parseCtx, event.Tp.TraceId, event.Tp.SpanId, 1); ok {
 			responseBuffer = b
 		}
 	}
diff --git a/pkg/components/ebpf/common/tcp_large_buffer.go b/pkg/components/ebpf/common/tcp_large_buffer.go
diff --git a/pkg/components/ebpf/common/tcp_large_buffer_test.go b/pkg/components/ebpf/common/tcp_large_buffer_test.go
diff --git a/pkg/components/ebpf/generictracer/generictracer.go b/pkg/components/ebpf/generictracer/generictracer.go
diff --git a/pkg/config/ebpf_tracer.go b/pkg/config/ebpf_tracer.go

Original file line number	Diff line number	Diff line change
`@@ -41,10 +41,10 @@ func ReadTCPRequestIntoSpan(parseCtx EBPFParseContext, cfg config.EBPFTracer,`
`41`	`41`	`responseBuffer = event.Rbuf[:l]`
`42`	`42`
`43`	`43`	`if event.HasLargeBuffers == 1 {`
`44`		`- if b, ok := getTCPLargeBuffer(parseCtx, event.Tp.TraceId, event.Tp.SpanId, 0); ok {`
	`44`	`+ if b, ok := extractTCPLargeBuffer(parseCtx, event.Tp.TraceId, event.Tp.SpanId, 0); ok {`
`45`	`45`	`requestBuffer = b`
`46`	`46`	`}`
`47`		`- if b, ok := getTCPLargeBuffer(parseCtx, event.Tp.TraceId, event.Tp.SpanId, 1); ok {`
	`47`	`+ if b, ok := extractTCPLargeBuffer(parseCtx, event.Tp.TraceId, event.Tp.SpanId, 1); ok {`
`48`	`48`	`responseBuffer = b`
`49`	`49`	`}`
`50`	`50`	`}`