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>
 
@@ -1415,6 +1416,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 (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 (is_kafka(&args->pid_conn.conn,
                         (const unsigned char *)args->u_buf,
                         args->bytes_len,

bpf/generictracer/protocol_mssql.h

@@ -0,0 +1,176 @@
+// 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 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.
+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) {
+    u32 expected_len = 0;
+
+    if (packet_type == PACKET_TYPE_RESPONSE) {
+        // Response Phase
+        if (req->resp_len > 0) {
+            // The first part of the response (including the header) is already in req->rbuf
+            struct mssql_hdr hdr = mssql_parse_hdr(req->rbuf);
+            expected_len = hdr.length;
+        } else if (bytes_len >= k_mssql_hdr_size) {
+            struct mssql_hdr hdr = mssql_parse_hdr(u_buf);
+            expected_len = hdr.length;
+            bpf_dbg_printk("mssql response: expected_len=%d", hdr.length);
+
+            // Copy the first response packet, which contains the TDS header and error token for userspace parsing
+            bpf_probe_read(req->rbuf, k_tcp_res_len, u_buf);
+        }
+    }
+
+    // these are the bytes already sent so far
+    u32 *bytes_sent_ptr =
+        (packet_type == PACKET_TYPE_REQUEST) ? &req->lb_req_bytes : &req->lb_res_bytes;
+    const u32 bytes_sent = *bytes_sent_ptr;
+
+    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) {
+            if (packet_type == PACKET_TYPE_RESPONSE) {
+                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);
+
+            u32 available_bytes = min(bytes_len, max_available_bytes);
+            const u32 consumed_bytes = large_buf_emit_chunks(large_buf, u_buf, available_bytes);
+
+            if (consumed_bytes > 0) {
+                req->has_large_buffers = true;
+            }
+
+            *bytes_sent_ptr += consumed_bytes;
+        }
+    }
+
+    if (packet_type == PACKET_TYPE_RESPONSE) {
+        req->resp_len += bytes_len;
+        if (expected_len > 0 && req->resp_len < expected_len) {
+            bpf_dbg_printk("mssql response: partial, acc=%d exp=%d", req->resp_len, expected_len);
+            return -1;
+        }
+    }
+
+    return 0;
+}

bpf/generictracer/protocol_tcp.h

@@ -22,6 +22,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>
 
@@ -142,6 +143,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:
@@ -266,13 +269,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"]

internal/test/integration/components/mssqldb/setup.sql

@@ -0,0 +1,12 @@
+CREATE DATABASE testdb;
+GO
+USE testdb;
+GO
+CREATE TABLE actor (
+    actor_id INT PRIMARY KEY,
+    first_name VARCHAR(50),
+    last_name VARCHAR(50)
+);
+GO
+INSERT INTO actor (actor_id, first_name, last_name) VALUES (1, 'TOM', 'CRUISE');
+GO

internal/test/integration/components/pythonsql/Dockerfile_mssql

@@ -0,0 +1,7 @@
+# Dockerfile that will build a container that runs python with FastAPI and uvicorn on port 8080
+FROM python:3.14@sha256:61346539f7b26521a230e72c11da5ebd872924745074b19736e7d65ba748c366
+EXPOSE 8080
+COPY requirements.txt /requirements.txt
+RUN pip install --require-hashes -r /requirements.txt
+COPY main_mssql.py /main.py
+CMD ["uvicorn", "--port", "8080", "--host", "0.0.0.0", "main:app"]