feat: Support MSSQL packets

bpf/common/connection_info.h

@@ -21,6 +21,7 @@ enum protocol_type : u8 {
     k_protocol_type_http = 3,
     k_protocol_type_kafka = 4,
     k_protocol_type_mqtt = 5,
+    k_protocol_type_mssql = 6,
 };
 
 // Struct to keep information on the connections in flight

bpf/common/large_buffers.h

@@ -9,6 +9,7 @@ volatile const u32 http_max_captured_bytes = 0;
 volatile const u32 mysql_max_captured_bytes = 0;
 volatile const u32 postgres_max_captured_bytes = 0;
 volatile const u32 kafka_max_captured_bytes = 0;
+volatile const u32 mssql_max_captured_bytes = 0;
 
 enum {
     // Maximum payload size per ring buffer chunk.
@@ -26,6 +27,7 @@ enum {
     k_large_buf_max_mysql_captured_bytes = 1 << 16,
     k_large_buf_max_postgres_captured_bytes = 1 << 16,
     k_large_buf_max_kafka_captured_bytes = 1 << 16,
+    k_large_buf_max_mssql_captured_bytes = 1 << 16,
 };
 
 SCRATCH_MEM_SIZED(tcp_large_buffers, k_large_buf_max_size);

bpf/generictracer/k_tracer.c

@@ -36,6 +36,7 @@
 #include <generictracer/protocol_http2.h>
 #include <generictracer/protocol_mysql.h>
 #include <generictracer/protocol_postgres.h>
+#include <generictracer/protocol_mssql.h>
 #include <generictracer/protocol_tcp.h>
 #include <generictracer/ssl_defs.h>
 

bpf/generictracer/protocol_handler.c

@@ -74,6 +74,12 @@ int obi_handle_buf_with_args(void *ctx) {
                                                   &args->protocol_type)) {
         bpf_dbg_printk("Found postgres connection");
         bpf_tail_call(ctx, &jump_table, k_tail_protocol_tcp);
+    } else if (args->protocols.tcp && is_mssql(&args->pid_conn.conn,
+                                               (const unsigned char *)args->u_buf,
+                                               args->bytes_len,
+                                               &args->protocol_type)) {
+        bpf_dbg_printk("Found mssql connection");
+        bpf_tail_call(ctx, &jump_table, k_tail_protocol_tcp);
     } else if (args->protocols.tcp && is_kafka(&args->pid_conn.conn,
                                                (const unsigned char *)args->u_buf,
                                                args->bytes_len,

bpf/generictracer/protocol_mssql.h

@@ -0,0 +1,219 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <logger/bpf_dbg.h>
+#include <bpfcore/vmlinux.h>
+#include <bpfcore/bpf_endian.h>
+#include <bpfcore/bpf_helpers.h>
+#include <bpfcore/utils.h>
+
+#include <common/common.h>
+#include <common/connection_info.h>
+#include <common/large_buffers.h>
+#include <common/ringbuf.h>
+
+#include <generictracer/maps/protocol_cache.h>
+#include <generictracer/protocol_common.h>
+
+// TDS Packet Header
+// https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-tds/7af53667-1b72-4703-8258-7984e838f746
+struct mssql_hdr {
+    u8 type;
+    u8 status;
+    u16 length; // big-endian
+    u16 spid;   // big-endian
+    u8 packet_id;
+    u8 window;
+};
+
+enum {
+    // TDS header
+    k_mssql_hdr_size = 8,
+    k_mssql_messages_in_packet_max = 10,
+
+    // TDS status bits
+    k_mssql_status_eom = 0x01, // End Of Message
+
+    // TDS message types
+    k_mssql_msg_sql_batch = 0x01,
+    k_mssql_msg_rpc = 0x03,
+    k_mssql_msg_response = 0x04,
+    k_mssql_msg_login7 = 0x10,
+    k_mssql_msg_prelogin = 0x12,
+};
+
+static __always_inline struct mssql_hdr mssql_parse_hdr(const unsigned char *data) {
+    struct mssql_hdr hdr = {};
+
+    bpf_probe_read(&hdr, sizeof(hdr), (const void *)data);
+
+    // Length and SPID are big-endian
+    hdr.length = bpf_ntohs(hdr.length);
+    hdr.spid = bpf_ntohs(hdr.spid);
+
+    return hdr;
+}
+
+static __always_inline u8 is_mssql(connection_info_t *conn_info,
+                                   const unsigned char *data,
+                                   u32 data_len,
+                                   enum protocol_type *protocol_type) {
+    if (*protocol_type != k_protocol_type_mssql && *protocol_type != k_protocol_type_unknown) {
+        // Already classified, not mssql.
+        return 0;
+    }
+
+    if (data_len < k_mssql_hdr_size) {
+        return 0;
+    }
+
+    size_t offset = 0;
+    bool includes_known_command = false;
+
+    for (u8 i = 0; i < k_mssql_messages_in_packet_max; i++) {
+        if (offset + k_mssql_hdr_size > data_len) {
+            break;
+        }
+
+        struct mssql_hdr hdr = mssql_parse_hdr(data + offset);
+
+        if (hdr.length < k_mssql_hdr_size || hdr.length > data_len - offset) {
+            return 0;
+        }
+
+        switch (hdr.type) {
+        case k_mssql_msg_sql_batch:
+        case k_mssql_msg_rpc:
+        case k_mssql_msg_response:
+        case k_mssql_msg_login7:
+        case k_mssql_msg_prelogin:
+            includes_known_command = true;
+            break;
+        default:
+            break;
+        }
+
+        offset += hdr.length;
+    }
+
+    if (offset != data_len || !includes_known_command) {
+        return 0;
+    }
+
+    *protocol_type = k_protocol_type_mssql;
+    bpf_map_update_elem(&protocol_cache, conn_info, protocol_type, BPF_ANY);
+
+    return 1;
+}
+
+// Emit a large buffer event for MSSQL 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 mssql_send_large_buffer(tcp_req_t *req,
+                                                   const void *u_buf,
+                                                   u32 bytes_len,
+                                                   u8 packet_type,
+                                                   u8 direction,
+                                                   enum large_buf_action action) {
+    if (mssql_max_captured_bytes > k_large_buf_max_mssql_captured_bytes) {
+        bpf_dbg_printk("BUG: mssql_max_captured_bytes exceeds maximum allowed value.");
+    }
+
+    if (packet_type == PACKET_TYPE_RESPONSE && req->resp_len == 0 &&
+        bytes_len >= k_mssql_hdr_size) {
+        bpf_probe_read(req->rbuf, k_mssql_hdr_size, u_buf);
+    }
+
+    const u32 bytes_sent =
+        packet_type == PACKET_TYPE_REQUEST ? req->lb_req_bytes : req->lb_res_bytes;
+
+    if (mssql_max_captured_bytes > 0 && bytes_sent < mssql_max_captured_bytes && bytes_len > 0) {
+        tcp_large_buffer_t *large_buf = (tcp_large_buffer_t *)tcp_large_buffers_mem();
+        if (!large_buf) {
+            bpf_dbg_printk(
+                "mssql_send_large_buffer: failed to reserve space for MSSQL large buffer");
+        } else {
+            large_buf->type = EVENT_TCP_LARGE_BUFFER;
+            large_buf->packet_type = packet_type;
+            large_buf->action = action;
+            large_buf->direction = direction;
+            large_buf->conn_info = req->conn_info;
+            large_buf->tp = req->tp;
+
+            u32 max_available_bytes = mssql_max_captured_bytes - bytes_sent;
+            bpf_clamp_umax(max_available_bytes, k_large_buf_max_mssql_captured_bytes);
+
+            const u32 available_bytes = min(bytes_len, max_available_bytes);
+            const u32 consumed_bytes = large_buf_emit_chunks(large_buf, u_buf, available_bytes);
+
+            if (packet_type == PACKET_TYPE_REQUEST) {
+                req->lb_req_bytes += consumed_bytes;
+            } else {
+                req->lb_res_bytes += consumed_bytes;
+            }
+
+            if (consumed_bytes > 0) {
+                req->has_large_buffers = true;
+            }
+        }
+    }
+
+    if (packet_type == PACKET_TYPE_RESPONSE) {
+        req->resp_len += bytes_len;
+
+        // Scan complete TDS packets in the current recv buffer for the EOM bit.
+        // A response may arrive as: (a) one or more complete packets in a single
+        // recv, or (b) a single packet split across multiple recvs (header first,
+        // then payload). Handle both by falling back to accumulated-length tracking
+        // when the current buffer does not start at a TDS packet boundary.
+        bool eom = false;
+        bool found_complete_packet = false;
+        u32 offset = 0;
+        for (u8 i = 0; i < k_mssql_messages_in_packet_max; i++) {
+            if (offset + k_mssql_hdr_size > bytes_len) {
+                break;
+            }
+            const struct mssql_hdr hdr = mssql_parse_hdr((const unsigned char *)u_buf + offset);
+            if (hdr.length < k_mssql_hdr_size || offset + hdr.length > bytes_len) {
+                break;
+            }
+            found_complete_packet = true;
+            if (hdr.status & k_mssql_status_eom) {
+                eom = true;
+            }
+            offset += hdr.length;
+        }
+
+        if (eom) {
+            // EOM found: all response packets received.
+        } else if (found_complete_packet) {
+            // Complete packets present but no EOM: more packets expected.
+            bpf_dbg_printk("mssql response: waiting for EOM, acc=%d", req->resp_len);
+            return -1;
+        } else {
+            // Could not parse a complete TDS packet (partial recv or payload
+            // continuation). Fall back to length tracking using the saved header.
+            if (req->resp_len < k_mssql_hdr_size) {
+                return -1;
+            }
+            const struct mssql_hdr first_hdr = mssql_parse_hdr(req->rbuf);
+            if (first_hdr.length >= k_mssql_hdr_size) {
+                const u32 prev_resp_len = req->resp_len - bytes_len;
+                // Wait while packet 1 is still completing, or if it just finished
+                // in this recv with no EOM (more packets follow; the next recv
+                // will start at a TDS boundary where the main loop detects EOM).
+                if (prev_resp_len < first_hdr.length &&
+                    (req->resp_len < first_hdr.length ||
+                     !(first_hdr.status & k_mssql_status_eom))) {
+                    bpf_dbg_printk(
+                        "mssql response: partial, acc=%d exp=%d", req->resp_len, first_hdr.length);
+                    return -1;
+                }
+            }
+        }
+    }
+
+    return 0;
+}

bpf/generictracer/protocol_tcp.h

@@ -23,6 +23,7 @@
 #include <generictracer/protocol_kafka.h>
 #include <generictracer/protocol_mysql.h>
 #include <generictracer/protocol_postgres.h>
+#include <generictracer/protocol_mssql.h>
 
 #include <generictracer/maps/tcp_req_mem.h>
 
@@ -143,6 +144,8 @@ static __always_inline int tcp_send_large_buffer(tcp_req_t *req,
         return postgres_send_large_buffer(req, u_buf, bytes_len, packet_type, direction, action);
     case k_protocol_type_kafka:
         return kafka_send_large_buffer(req, pid_conn, u_buf, bytes_len, direction, action);
+    case k_protocol_type_mssql:
+        return mssql_send_large_buffer(req, u_buf, bytes_len, packet_type, direction, action);
     case k_protocol_type_http:
     case k_protocol_type_mqtt:
     case k_protocol_type_unknown:
@@ -269,13 +272,11 @@ static __always_inline void handle_unknown_tcp_connection(pid_connection_info_t
         }
     } else if (existing->direction != direction) {
         existing->is_server = is_server;
-        if (tcp_send_large_buffer(existing,
-                                  pid_conn,
-                                  u_buf,
-                                  bytes_len,
-                                  direction,
-                                  protocol_type,
-                                  k_large_buf_action_init) < 0) {
+        const enum large_buf_action response_action =
+            (existing->lb_res_bytes > 0) ? k_large_buf_action_append : k_large_buf_action_init;
+        if (tcp_send_large_buffer(
+                existing, pid_conn, u_buf, bytes_len, direction, protocol_type, response_action) <
+            0) {
             bpf_dbg_printk("waiting additional response data");
             return;
         }

devdocs/features.md

@@ -6,6 +6,7 @@
 | gRPC          |    All    |        1.0+ | All                                                                                      |  Yes   |                 No |                                      Can't get method for long living connections before OBI started, will mark method with `*`
 | MySQL         |    All    |         All | All                                                                                      |  Yes   |                 No |             In the case of prepared statements, if the statement was prepared before OBI started then the query might be missed
 | PostgreSQL    |    All    |         All | All                                                                                      |  Yes   |                 No |             In the case of prepared statements, if the statement was prepared before OBI started then the query might be missed
+| MSSQL         |    All    |         All | All                                                                                      |  Yes   |                 No |             In the case of prepared statements, if the statement was prepared before OBI started then the query might be missed
 | Redis         |    All    |         All | All                                                                                      |  Yes   |                 No |             For already started connections, can't infer the number of the database, and won't add the `db.namespace` attribute
 | MongoDB       |    All    |        5.0+ | insert, update, find, delete, findAndModify, aggregate, count, distinct, mapReduce       |  Yes   |                 No |                                                                                              no support for compressed payloads
 | Couchbase     |    All    |         All | All                                                                                      |  Yes   |                 No | Bucket unknown if SELECT_BUCKET occurred before OBI started; Collection unknown if GET_COLLECTION_ID occurred before OBI started
@@ -61,8 +62,9 @@ Large payloads are streamed to userspace across multiple ring-buffer events and
 | `OTEL_EBPF_BPF_BUFFER_SIZE_MYSQL`  | MySQL      | 65535   | 0 (disabled) |
 | `OTEL_EBPF_BPF_BUFFER_SIZE_KAFKA`  | Kafka      | 65535   | 0 (disabled) |
 | `OTEL_EBPF_BPF_BUFFER_SIZE_POSTGRES` | PostgreSQL | 65535 | 0 (disabled) |
+| `OTEL_EBPF_BPF_BUFFER_SIZE_MSSQL`  | MSSQL      | 65535   | 0 (disabled) |
 
-Equivalent YAML keys live under `ebpf.buffer_sizes.{http,mysql,kafka,postgres}`.
+Equivalent YAML keys live under `ebpf.buffer_sizes.{http,mysql,kafka,postgres,mssql}`.
 
 ## GPU Instrumentation
 

docs/config-schema.json

@@ -309,6 +309,10 @@
           "type": "integer",
           "x-env-var": "OTEL_EBPF_BPF_BUFFER_SIZE_KAFKA"
         },
+        "mssql": {
+          "type": "integer",
+          "x-env-var": "OTEL_EBPF_BPF_BUFFER_SIZE_MSSQL"
+        },
         "mysql": {
           "type": "integer",
           "x-env-var": "OTEL_EBPF_BPF_BUFFER_SIZE_MYSQL"
@@ -443,6 +447,11 @@
           "description": "MongoDB requests cache size.",
           "x-env-var": "OTEL_EBPF_BPF_MONGO_REQUESTS_CACHE_SIZE"
         },
+        "mssql_prepared_statements_cache_size": {
+          "type": "integer",
+          "description": "MSSQL prepared statements cache size.",
+          "x-env-var": "OTEL_EBPF_BPF_MSSQL_PREPARED_STATEMENTS_CACHE_SIZE"
+        },
         "mysql_prepared_statements_cache_size": {
           "type": "integer",
           "description": "MySQL prepared statements cache size.",

internal/test/integration/components/mssqldb/Dockerfile

@@ -0,0 +1,24 @@
+FROM mcr.microsoft.com/mssql/server@sha256:49b45a911dc535e9345fbfd7101a1bd8a1e190a5f29b877ef75387a061e5fcf0
+
+USER root
+
+# Create a directory for our initialization
+COPY setup.sql /setup.sql
+RUN chown mssql /setup.sql
+
+# Set required environment variables for the build-time run
+ENV ACCEPT_EULA=Y
+ENV MSSQL_SA_PASSWORD=p_ssW0rd
+
+# Switch to mssql user for the initialization run
+USER mssql
+
+# Start SQL Server, wait for it to be ready, run the setup script, and then shut down.
+RUN /opt/mssql/bin/sqlservr & \
+    bash -c "until /opt/mssql-tools18/bin/sqlcmd -S localhost -U sa -P \"$MSSQL_SA_PASSWORD\" -C -Q 'SELECT 1' &> /dev/null; do sleep 2; echo 'Waiting for SQL Server...'; done" && \
+    /opt/mssql-tools18/bin/sqlcmd -S localhost -U sa -P "$MSSQL_SA_PASSWORD" -C -i /setup.sql && \
+    pkill sqlservr && \
+    sleep 5
+
+# Use the default entrypoint
+ENTRYPOINT ["/opt/mssql/bin/sqlservr"]