Enhance the inject_header function, fix bug in net/http server instru…

CHANGELOG.md

@@ -8,6 +8,10 @@ OpenTelemetry Go Automatic Instrumentation adheres to [Semantic Versioning](http
 
 ## [Unreleased]
 
+### Changed
+- Re-wrtie the inject_header function to be more readable ([#266]https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/266)
+- Fix bug in the net/http server instrumentation which always created a new span context ([#266]https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/266)
+
 ### Deprecated
 
 - The `go.opentelemetry.io/auto/examples/rolldice` module is deprecated.

internal/include/go_types.h

@@ -54,6 +54,17 @@ struct map_bucket {
     void *overflow;
 };
 
+// A header for a Go map. Defined in go/src/runtime/map.go
+struct go_map_header {
+    u64 map_keyvalue_count;
+    u8 flags;
+    u8 log2_bucket_count;
+    u16 noverflow;
+    u32 hash0;
+    void* buckets;
+    // More fiels which we don't care about
+};
+
 static __always_inline struct go_string write_user_go_string(char *str, u32 len)
 {
     // Copy chars to userspace

internal/include/utils.h

@@ -75,3 +75,11 @@ static __always_inline void copy_byte_arrays(unsigned char *src, unsigned char *
         dst[i] = src[i];
     }
 }
+
+static __always_inline void bpf_memset(unsigned char *dst, u32 size, unsigned char value)
+{
+    for (int i = 0; i < size; i++)
+    {
+        dst[i] = value;
+    }
+}

internal/pkg/instrumentors/bpf/net/http/client/bpf/probe.bpf.c

@@ -44,94 +44,79 @@ volatile const u64 headers_ptr_pos;
 volatile const u64 ctx_ptr_pos;
 
 static __always_inline long inject_header(void* headers_ptr, struct span_context* propagated_ctx) {
+    // Read the map header
+    struct go_map_header map_header = {0};
+    long res = bpf_probe_read_user(&map_header, sizeof(struct go_map_header), headers_ptr);
 
-    // Read headers map count
-    u64 map_keyvalue_count = 0;
-    bpf_probe_read(&map_keyvalue_count, sizeof(map_keyvalue_count), headers_ptr);
+    if (res < 0) {
+        bpf_printk("Couldn't read map header from user");
+        return -1;
+    }
+
+    u64 curr_keyvalue_count = map_header.map_keyvalue_count;
 
-    // Currently only maps with less than 8 keys are supported for injection
-    if (map_keyvalue_count >= 8) {
+    if (curr_keyvalue_count >= 8) {
         bpf_printk("Map size is bigger than 8, skipping context propagation");
         return 0;
     }
 
-    long res;
-    if (map_keyvalue_count == 0) {
-        u32 map_id = 0;
-        struct map_bucket *map_value = bpf_map_lookup_elem(&golang_mapbucket_storage_map, &map_id);
-        if (!map_value) {
-            return -1;
-        }
-        void *bucket_ptr = write_target_data(map_value, sizeof(struct map_bucket));
-        res = bpf_probe_write_user(headers_ptr + 16, &bucket_ptr, sizeof(bucket_ptr));
-
-        if(res < 0) {
-            bpf_printk("Failed to write bucket ptr, return code: %d", res);
-            return -1;
-        }
+    // Get pointer to temp bucket struct we store in a map (avoiding large local variable on the stack)
+    // Performing read-modify-write on the bucket and for the header
+    u32 map_id = 0;
+    struct map_bucket *bucket_map_value = bpf_map_lookup_elem(&golang_mapbucket_storage_map, &map_id);
+    if (!bucket_map_value) {
+        return -1;
+    }
 
+    if (curr_keyvalue_count == 0) {
+        // No key-value pairs in the Go map, need to "allocate" memory for the user
+        void *bucket_ptr = write_target_data(bucket_map_value, sizeof(struct map_bucket));
+        map_header.buckets = bucket_ptr;
+    } else {
+        // There is at least 1 key-value pair, hence the bucket pointer from the user is valid
+        // Read the user's bucket to the eBPF map entry
+        bpf_probe_read_user(bucket_map_value, sizeof(struct map_bucket), map_header.buckets);
     }
 
-    void *map_keyvalues_ptr = NULL;
-    bpf_probe_read(&map_keyvalues_ptr, sizeof(map_keyvalues_ptr), headers_ptr + 16);
-    void *injected_key_ptr = map_keyvalues_ptr + 8 + (16 * map_keyvalue_count);
     char traceparent_tophash = 0xee;
-    void *tophashes_ptr = map_keyvalues_ptr +  map_keyvalue_count;
-    res = bpf_probe_write_user(tophashes_ptr, &traceparent_tophash, 1);
-
-    if(res < 0) {
-        bpf_printk("Failed to write tophash, return code: %d", res);
-        return -1;
-    }
+    bucket_map_value->tophash[curr_keyvalue_count] = traceparent_tophash;
 
+    // Prepare the key string for the user
     char key[W3C_KEY_LENGTH] = "traceparent";
     void *ptr = write_target_data(key, W3C_KEY_LENGTH);
+    bucket_map_value->keys[curr_keyvalue_count] = (struct go_string) {.len = W3C_KEY_LENGTH, .str = ptr};
 
-    res = bpf_probe_write_user(injected_key_ptr, &ptr, sizeof(ptr));
-    if(res < 0) {
-        return -1;
-    }
-
-    u64 header_key_length = W3C_KEY_LENGTH;
-    res = bpf_probe_write_user(injected_key_ptr + 8, &header_key_length, sizeof(header_key_length));
-
-    if(res < 0) {
-        return -1;
-    }
-
-    void *injected_value_ptr = injected_key_ptr + (16 * (8 - map_keyvalue_count)) + 24 * map_keyvalue_count;
+    // Prepare the value string slice
+    // First the value string which constains the span context
     char val[W3C_VAL_LENGTH];
     span_context_to_w3c_string(propagated_ctx, val);
-
     ptr = write_target_data(val, sizeof(val));
-    struct go_string header_value = {};
-    header_value.str = ptr;
-    header_value.len = W3C_VAL_LENGTH;
+    if(ptr == NULL) {
+        return -1;
+    }
 
+    // The go string pointing to the above val
+    struct go_string header_value = {.len = W3C_VAL_LENGTH, .str = ptr};
     ptr = write_target_data((void*)&header_value, sizeof(header_value));
-
     if(ptr == NULL) {
         return -1;
     }
 
-    struct go_slice values_slice = {};
-    values_slice.array = ptr;
-    values_slice.len = 1;
-    values_slice.cap = 1;
-
-    res = bpf_probe_write_user(injected_value_ptr, &values_slice, sizeof(values_slice));
+    // Last, go_slice pointing to the above go_string
+    bucket_map_value->values[curr_keyvalue_count] = (struct go_slice) {.array = ptr, .cap = 1, .len = 1};
 
-    if(res < 0) {
+    map_header.map_keyvalue_count++;
+    // Update the map header
+    res = bpf_probe_write_user(headers_ptr, &map_header, sizeof(struct go_map_header));
+    if (res < 0) {
         return -1;
     }
-
-    map_keyvalue_count += 1;
-    res = bpf_probe_write_user(headers_ptr, &map_keyvalue_count, sizeof(map_keyvalue_count));
-
-    if(res < 0) {
+    // Update the bucket
+    res = bpf_probe_write_user(map_header.buckets, bucket_map_value, sizeof(struct map_bucket));
+    if (res < 0) {
         return -1;
     }
-
+    bpf_memset((unsigned char *)bucket_map_value, sizeof(struct map_bucket), 0);
     return 0;
 }
 

internal/pkg/instrumentors/bpf/net/http/server/bpf/probe.bpf.c

@@ -216,7 +216,6 @@ int uprobe_ServerMux_ServeHTTP(struct pt_regs *ctx)
     void *key = get_consistent_key(ctx, (void *)(req_ptr + ctx_ptr_pos));
 
     // Write event
-    httpReq.sc = generate_span_context();
     bpf_map_update_elem(&http_events, &key, &httpReq, 0);
     start_tracking_span(req_ctx_ptr, &httpReq.sc);
     return 0;