Propagate HTTP trace context in Go

bpf/go_common.h

@@ -78,27 +78,31 @@ static __always_inline u64 find_parent_goroutine(void *goroutine_addr) {
     return 0;
 }
 
-static __always_inline void decode_go_traceparent(unsigned char *buf, unsigned char *trace_id, unsigned char *span_id) {
+static __always_inline void decode_go_traceparent(unsigned char *buf, unsigned char *trace_id, unsigned char *span_id, unsigned char *flags) {
     unsigned char *t_id = buf + 2 + 1; // strlen(ver) + strlen("-")
     unsigned char *s_id = buf + 2 + 1 + 32 + 1; // strlen(ver) + strlen("-") + strlen(trace_id) + strlen("-")
+    unsigned char *f_id = buf + 2 + 1 + 32 + 1 + 16 + 1; // strlen(ver) + strlen("-") + strlen(trace_id) + strlen("-") + strlen(span_id) + strlen("-")
 
     decode_hex(trace_id, t_id, TRACE_ID_CHAR_LEN);
     decode_hex(span_id, s_id, SPAN_ID_CHAR_LEN);
+    decode_hex(flags, f_id, FLAGS_CHAR_LEN);
 } 
 
 static __always_inline void server_trace_parent(void *goroutine_addr, tp_info_t *tp, void *req_header) {
+    // May get overriden when decoding existing traceparent, but otherwise we set sample ON
+    tp->flags = 1;
     // Get traceparent from the Request.Header
     void *traceparent_ptr = extract_traceparent_from_req_headers(req_header);
     if (traceparent_ptr != NULL) {
-        unsigned char buf[W3C_VAL_LENGTH];
+        unsigned char buf[TP_MAX_VAL_LENGTH];
         long res = bpf_probe_read(buf, sizeof(buf), traceparent_ptr);
         if (res < 0) {
             bpf_dbg_printk("can't copy traceparent header");
             urand_bytes(tp->trace_id, TRACE_ID_SIZE_BYTES);
             *((u64 *)tp->parent_id) = 0;
         } else {
             bpf_dbg_printk("Decoding traceparent from headers %s", buf);
-            decode_go_traceparent(buf, tp->trace_id, tp->parent_id);
+            decode_go_traceparent(buf, tp->trace_id, tp->parent_id, &tp->flags);
         }
     } else {
         bpf_dbg_printk("No traceparent in headers, generating");
@@ -110,20 +114,25 @@ static __always_inline void server_trace_parent(void *goroutine_addr, tp_info_t
     bpf_map_update_elem(&go_trace_map, &goroutine_addr, tp, BPF_ANY);
 }
 
-static __always_inline void client_trace_parent(void *goroutine_addr, tp_info_t *tp_i, void *req_header) {
+static __always_inline u8 client_trace_parent(void *goroutine_addr, tp_info_t *tp_i, void *req_header) {
     // Get traceparent from the Request.Header
     u8 found_trace_id = 0;
+    u8 trace_id_exists = 0;
 
+    // May get overriden when decoding existing traceparent or finding a server span, but otherwise we set sample ON
+    tp_i->flags = 1;
+
     if (req_header) {
         void *traceparent_ptr = extract_traceparent_from_req_headers(req_header);
         if (traceparent_ptr != NULL) {
             unsigned char buf[W3C_VAL_LENGTH];
+            trace_id_exists = 1;
             long res = bpf_probe_read(buf, sizeof(buf), traceparent_ptr);
             if (res < 0) {
                 bpf_dbg_printk("can't copy traceparent header");
             } else {
                 found_trace_id = 1;
-                decode_go_traceparent(buf, tp_i->trace_id, tp_i->span_id);
+                decode_go_traceparent(buf, tp_i->trace_id, tp_i->span_id, &tp_i->flags);
             }
         }
     }
@@ -142,12 +151,15 @@ static __always_inline void client_trace_parent(void *goroutine_addr, tp_info_t
             *((u64 *)tp_i->trace_id) = *((u64 *)tp->trace_id);
             *((u64 *)(tp_i->trace_id + 8)) = *((u64 *)(tp->trace_id + 8));
             *((u64 *)tp_i->parent_id) = *((u64 *)tp->span_id);
+            tp_i->flags = tp->flags;
         } else {
             urand_bytes(tp_i->trace_id, TRACE_ID_SIZE_BYTES);    
         }
 
         urand_bytes(tp_i->span_id, SPAN_ID_SIZE_BYTES);
     }
+
+    return trace_id_exists;
 }
 
 

bpf/go_nethttp.c

@@ -61,6 +61,8 @@ int uprobe_ServeHTTP(struct pt_regs *ctx) {
 
     if (req) {
         server_trace_parent(goroutine_addr, &invocation.tp, (void*)(req + req_header_ptr_pos));
+        // TODO: if context propagation is supported, overwrite the header value in the map with the 
+        // new span context and the same thread id.
     }
 
     // Write event
@@ -193,6 +195,16 @@ int uprobe_WriteHeader(struct pt_regs *ctx) {
     return 0;
 }
 
+#ifndef NO_HEADER_PROPAGATION
+struct {
+    __uint(type, BPF_MAP_TYPE_HASH);
+    __type(key, void *); // key: pointer to the request header map
+    __type(value, u64); // the goroutine of the transport request
+    __uint(max_entries, MAX_CONCURRENT_REQUESTS);
+} header_req_map SEC(".maps");
+
+#endif
+
 /* HTTP Client. We expect to see HTTP client in both HTTP server and gRPC server calls.*/
 
 SEC("uprobe/roundTrip")
@@ -210,13 +222,25 @@ int uprobe_roundTrip(struct pt_regs *ctx) {
         .tp = {0}
     };
 
-    client_trace_parent(goroutine_addr, &invocation.tp, (void*)(req + req_header_ptr_pos));
+    __attribute__((__unused__)) u8 existing_tp = client_trace_parent(goroutine_addr, &invocation.tp, (void*)(req + req_header_ptr_pos));
 
     // Write event
     if (bpf_map_update_elem(&ongoing_http_client_requests, &goroutine_addr, &invocation, BPF_ANY)) {
         bpf_dbg_printk("can't update http client map element");
     }
 
+#ifndef NO_HEADER_PROPAGATION
+    if (!existing_tp) {
+        void *headers_ptr = 0;
+        bpf_probe_read(&headers_ptr, sizeof(headers_ptr), (void*)(req + req_header_ptr_pos));
+        bpf_dbg_printk("goroutine_addr %lx, req ptr %llx, headers_ptr %llx", goroutine_addr, req, headers_ptr);
+
+        if (headers_ptr) {
+            bpf_map_update_elem(&header_req_map, &headers_ptr, &goroutine_addr, BPF_ANY);
+        }
+    }
+#endif    
+
     return 0;
 }
 
@@ -290,3 +314,68 @@ int uprobe_roundTripReturn(struct pt_regs *ctx) {
 
     return 0;
 }
+
+#ifndef NO_HEADER_PROPAGATION
+// Context propagation through HTTP headers
+SEC("uprobe/header_writeSubset")
+int uprobe_writeSubset(struct pt_regs *ctx) {
+    bpf_dbg_printk("=== uprobe/proc header writeSubset === ");
+
+    void *header_addr = GO_PARAM1(ctx);
+    void *io_writer_addr = GO_PARAM3(ctx);
+
+    bpf_dbg_printk("goroutine_addr %lx, header ptr %llx", GOROUTINE_PTR(ctx), header_addr);
+
+    u64 *request_goaddr = bpf_map_lookup_elem(&header_req_map, &header_addr);
+
+    if (!request_goaddr) {
+        bpf_dbg_printk("Can't find parent go routine for header %llx", header_addr);
+        return 0;
+    }
+
+    u64 parent_goaddr = *request_goaddr;
+
+    http_func_invocation_t *func_inv = bpf_map_lookup_elem(&ongoing_http_client_requests, &parent_goaddr);
+    if (!func_inv) {
+        bpf_dbg_printk("Can't find client request for goroutine %llx", parent_goaddr);
+        return 0;
+    }
+
+    unsigned char buf[TRACEPARENT_LEN];
+
+    make_tp_string(buf, &func_inv->tp);
+
+    void *buf_ptr = 0;
+    bpf_probe_read(&buf_ptr, sizeof(buf_ptr), (void *)(io_writer_addr + io_writer_buf_ptr_pos));
+    if (!buf_ptr) {
+        return 0;
+    }
+
+    s64 size = 0;
+    bpf_probe_read(&size, sizeof(s64), (void *)(io_writer_addr + io_writer_buf_ptr_pos + 8)); // grab size
+
+    s64 len = 0;
+    bpf_probe_read(&len, sizeof(s64), (void *)(io_writer_addr + io_writer_n_pos)); // grab len
+
+    bpf_dbg_printk("buf_ptr %llx, len=%d, size=%d", (void*)buf_ptr, len, size);
+
+    if (len < (size - TP_MAX_VAL_LENGTH - TP_MAX_KEY_LENGTH - 4)) { // 4 = strlen(":_") + strlen("\r\n")
+        char key[TP_MAX_KEY_LENGTH + 2] = "Traceparent: ";
+        char end[2] = "\r\n";
+        bpf_probe_write_user(buf_ptr + (len & 0x0ffff), key, sizeof(key));
+        len += TP_MAX_KEY_LENGTH + 2;
+        bpf_probe_write_user(buf_ptr + (len & 0x0ffff), buf, sizeof(buf));
+        len += TP_MAX_VAL_LENGTH;
+        bpf_probe_write_user(buf_ptr + (len & 0x0ffff), end, sizeof(end));
+        len += 2;
+        bpf_probe_write_user((void *)(io_writer_addr + io_writer_n_pos), &len, sizeof(len));
+    }
+
+    return 0;
+}
+#else
+SEC("uprobe/header_writeSubset")
+int uprobe_writeSubset(struct pt_regs *ctx) {
+    return 0;
+}
+#endif

bpf/go_nethttp.h

@@ -27,5 +27,7 @@ volatile const u64 host_ptr_pos;
 volatile const u64 content_length_ptr_pos;
 volatile const u64 resp_req_pos;
 volatile const u64 req_header_ptr_pos;
+volatile const u64 io_writer_buf_ptr_pos;
+volatile const u64 io_writer_n_pos;
 
 #endif

bpf/trace_common.h

@@ -97,6 +97,10 @@ static __always_inline unsigned char *extract_span_id(unsigned char *tp_start) {
     return tp_start + 13 + 2 + 1 + 32 + 1; // strlen("Traceparent: ") + strlen(ver) + strlen("-") + strlen(trace_id) + strlen("-")
 }
 
+static __always_inline unsigned char *extract_flags(unsigned char *tp_start) {
+    return tp_start + 13 + 2 + 1 + 32 + 1 + 16 + 1; // strlen("Traceparent: ") + strlen(ver) + strlen("-") + strlen(trace_id) + strlen("-") + strlen(span_id) + strlen("-")
+}
+
 static __always_inline u64 current_epoch() {
     u64 ts = bpf_ktime_get_ns();
     u64 temp = ts / NANOSECONDS_PER_EPOCH;
@@ -122,6 +126,7 @@ static __always_inline void get_or_create_trace_info(http_connection_metadata_t
     }
 
     tp->epoch = current_epoch();
+    tp->flags = 1;
     urand_bytes(tp->span_id, SPAN_ID_SIZE_BYTES);
     bpf_memset(tp->parent_id, 0, sizeof(tp->span_id));
 
@@ -151,8 +156,10 @@ static __always_inline void get_or_create_trace_info(http_connection_metadata_t
             bpf_dbg_printk("Found traceparent %s", res);
             unsigned char *t_id = extract_trace_id(res);
             unsigned char *s_id = extract_span_id(res);
+            unsigned char *f_id = extract_flags(res);
 
             decode_hex(tp->trace_id, t_id, TRACE_ID_CHAR_LEN);
+            decode_hex((unsigned char *)&tp->flags, f_id, FLAGS_CHAR_LEN);
             if (meta && meta->type == EVENT_HTTP_CLIENT) {
                 decode_hex(tp->span_id, s_id, SPAN_ID_CHAR_LEN);
             } else {

bpf/trace_util.h

@@ -1,8 +1,6 @@
 #ifndef TRACE_UTIL_H
 #define TRACE_UTIL_H
 
-#include "utils.h"
-
 // 55+13
 #define TRACE_PARENT_HEADER_LEN 68
 

bpf/tracing.h

@@ -1,18 +1,41 @@
 #ifndef TRACING_H
 #define TRACING_H
 #include "vmlinux.h"
+#include "trace_util.h"
 
 #define TRACE_ID_SIZE_BYTES 16
 #define SPAN_ID_SIZE_BYTES   8
+#define FLAGS_SIZE_BYTES     1
 #define TRACE_ID_CHAR_LEN   32
 #define SPAN_ID_CHAR_LEN    16
+#define FLAGS_CHAR_LEN       2
 #define TP_MAX_VAL_LENGTH   55
+#define TP_MAX_KEY_LENGTH   11
 
 typedef struct tp_info {
     unsigned char trace_id[TRACE_ID_SIZE_BYTES];
     unsigned char span_id[SPAN_ID_SIZE_BYTES];
     unsigned char parent_id[SPAN_ID_SIZE_BYTES];
     u64 epoch;
+    u8  flags;
 } tp_info_t;
 
+static __always_inline void make_tp_string(unsigned char *buf, tp_info_t *tp) {
+    // Version
+    *buf++ = '0'; *buf++ = '0'; *buf++ = '-';
+
+    // TraceID
+    encode_hex(buf, tp->trace_id, TRACE_ID_SIZE_BYTES);
+    buf += TRACE_ID_CHAR_LEN;
+    *buf++ = '-';
+
+    // SpanID
+    encode_hex(buf, tp->span_id, SPAN_ID_SIZE_BYTES);
+    buf += SPAN_ID_CHAR_LEN;
+    *buf++ = '-';
+
+    // Flags
+    *buf++ = '0'; *buf = (tp->flags == 0) ? '0' : '1';
+}
+
 #endif

configs/offsets/tracker_input.json

@@ -27,6 +27,10 @@
       ],
       "context.valueCtx": [
         "val"
+      ],
+      "bufio.Writer": [
+        "buf",
+        "n"
       ]
     }
   },

examples/greeting-apps/docker-compose.yaml

@@ -86,6 +86,7 @@ services:
       - --config=/configs/beyla-config.yml
     volumes:
       - ./configs/:/configs
+      - ./system/sys/kernel/security:/sys/kernel/security
     container_name: demo-javabeyla
     privileged: true 
     network_mode: "service:javatestserver"

examples/greeting-apps/system/sys/kernel/security/lockdown

@@ -0,0 +1 @@
+none [integrity] confidentiality

pkg/internal/ebpf/common/spanner.go

@@ -60,6 +60,7 @@ func HTTPRequestTraceToSpan(trace *HTTPRequestTrace) request.Span {
 		TraceID:       trace2.TraceID(trace.Tp.TraceId),
 		SpanID:        trace2.SpanID(trace.Tp.SpanId),
 		ParentSpanID:  trace2.SpanID(trace.Tp.ParentId),
+		Flags:         trace.Tp.Flags,
 		Pid: request.PidInfo{
 			HostPID:   trace.Pid.HostPid,
 			UserPID:   trace.Pid.UserPid,
@@ -98,6 +99,7 @@ func SQLRequestTraceToSpan(trace *SQLRequestTrace) request.Span {
 		TraceID:       trace2.TraceID(trace.Tp.TraceId),
 		SpanID:        trace2.SpanID(trace.Tp.SpanId),
 		ParentSpanID:  trace2.SpanID(trace.Tp.ParentId),
+		Flags:         trace.Tp.Flags,
 		Pid: request.PidInfo{
 			HostPID:   trace.Pid.HostPid,
 			UserPID:   trace.Pid.UserPid,