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

#266)

* Enhance the inject_header function, fix bug in net/http server instrumentation

* Add changelog

* Fix formating

* Update changelog after rebase

* Apply suggestions from code review

Co-authored-by: Tyler Yahn <MrAlias@users.noreply.github.com>

* Change map_keyvalue_count to be s64 to fit better for the Go defintion

* Add buckets field to injected fields for the http inject header method

---------

Co-authored-by: Tyler Yahn <MrAlias@users.noreply.github.com>

.gitignore

@@ -25,7 +25,7 @@ launcher/
 opentelemetry-helm-charts/
 
 # don't track temp e2e trace json files
-test/**/traces-orig.json
+internal/test/**/traces-orig.json
 
 # ignore db files created by the example or tests
 *.db

CHANGELOG.md

@@ -14,6 +14,7 @@ OpenTelemetry Go Automatic Instrumentation adheres to [Semantic Versioning](http
 
 ### Changed
 
+- 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)
 - Fix runtime panic if OTEL_GO_AUTO_TARGET_EXE is not set. ([#339](https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/339))
 
 ### Deprecated

Makefile

@@ -145,7 +145,7 @@ fixture-databasesql: fixtures/databasesql
 fixtures/%: LIBRARY=$*
 fixtures/%:
 	$(MAKE) docker-build
-	cd test/e2e/$(LIBRARY) && docker build -t sample-app .
+	cd internal/test/e2e/$(LIBRARY) && docker build -t sample-app .
 	kind create cluster
 	kind load docker-image otel-go-instrumentation sample-app
 	helm repo add open-telemetry https://open-telemetry.github.io/opentelemetry-helm-charts
@@ -156,9 +156,9 @@ fixtures/%:
 	kubectl wait --for=condition=Ready --timeout=60s pod/test-opentelemetry-collector-0
 	kubectl -n default create -f .github/workflows/e2e/k8s/sample-job.yml
 	kubectl wait --for=condition=Complete --timeout=60s job/sample-job
-	kubectl cp -c filecp default/test-opentelemetry-collector-0:tmp/trace.json ./test/e2e/$(LIBRARY)/traces-orig.json
-	rm -f ./test/e2e/$(LIBRARY)/traces.json
-	bats ./test/e2e/$(LIBRARY)/verify.bats
+	kubectl cp -c filecp default/test-opentelemetry-collector-0:tmp/trace.json ./internal/test/e2e/$(LIBRARY)/traces-orig.json
+	rm -f ./internal/test/e2e/$(LIBRARY)/traces.json
+	bats ./internal/test/e2e/$(LIBRARY)/verify.bats
 	kind delete cluster
 
 .PHONY: prerelease

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/inject/offset_results.json

@@ -289,6 +289,22 @@
           ]
         }
       ]
+    },
+    "runtime.hmap": {
+      "buckets": [
+        {
+          "versions": {
+            "oldest": "1.12.0",
+            "newest": "1.21.1"
+          },
+          "offsets": [
+            {
+              "offset": 16,
+              "since": "1.12.0"
+            }
+          ]
+        }
+      ]
     }
   }
 }

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

@@ -42,96 +42,96 @@ volatile const u64 url_ptr_pos;
 volatile const u64 path_ptr_pos;
 volatile const u64 headers_ptr_pos;
 volatile const u64 ctx_ptr_pos;
+volatile const u64 buckets_ptr_pos;
 
 static __always_inline long inject_header(void* headers_ptr, struct span_context* propagated_ctx) {
+    // Read the key-value count - this field must be the first one in the hmap struct as documented in src/runtime/map.go
+    u64 curr_keyvalue_count = 0;
+    long res = bpf_probe_read_user(&curr_keyvalue_count, sizeof(curr_keyvalue_count), 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 key-value count from user");
+        return -1;
+    }
 
-    // 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;
+        return -1;
     }
 
-    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));
+    // 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
+    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;
+    }
+
+    void *buckets_ptr_ptr = (void*) (headers_ptr + buckets_ptr_pos);
+    void *bucket_ptr = 0; // The actual pointer to the buckets
 
-        if(res < 0) {
-            bpf_printk("Failed to write bucket ptr, return code: %d", res);
+    if (curr_keyvalue_count == 0) {
+        // No key-value pairs in the Go map, need to "allocate" memory for the user
+        bucket_ptr = write_target_data(bucket_map_value, sizeof(struct map_bucket));
+        // Update the buckets pointer in the hmap struct to point to newly allocated bucket
+        res = bpf_probe_write_user(buckets_ptr_ptr, &bucket_ptr, sizeof(bucket_ptr));
+        if (res < 0) {
+            bpf_printk("Failed to update the map bucket pointer for the user");
             return -1;
         }
-
+    } else {
+        // There is at least 1 key-value pair, hence the bucket pointer from the user is valid
+        // Read the bucket pointer
+        res = bpf_probe_read_user(&bucket_ptr, sizeof(bucket_ptr), buckets_ptr_ptr);
+        // Read the user's bucket to the eBPF map entry
+        bpf_probe_read_user(bucket_map_value, sizeof(struct map_bucket), bucket_ptr);
     }
 
-    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);
+    u8 bucket_index = curr_keyvalue_count & 0x7;
 
-    if(res < 0) {
-        bpf_printk("Failed to write tophash, return code: %d", res);
-        return -1;
-    }
+    char traceparent_tophash = 0xee;
+    bucket_map_value->tophash[bucket_index] = 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[bucket_index] = (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;
+    // Last, go_slice pointing to the above go_string
+    bucket_map_value->values[bucket_index] = (struct go_slice) {.array = ptr, .cap = 1, .len = 1};
 
-    res = bpf_probe_write_user(injected_value_ptr, &values_slice, sizeof(values_slice));
-
-    if(res < 0) {
+    // Update the map header count field
+    curr_keyvalue_count += 1;
+    res = bpf_probe_write_user(headers_ptr, &curr_keyvalue_count, sizeof(curr_keyvalue_count));
+    if (res < 0) {
+        bpf_printk("Failed to update key-value count in map header");
         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(bucket_ptr, bucket_map_value, sizeof(struct map_bucket));
+    if (res < 0) {
+        bpf_printk("Failed to update bucket content");
         return -1;
     }
 
+    bpf_memset((unsigned char *)bucket_map_value, sizeof(struct map_bucket), 0);
     return 0;
 }
 

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

@@ -98,6 +98,11 @@ func (h *Instrumentor) Load(ctx *context.InstrumentorContext) error {
 			StructName: "net/http.Request",
 			Field:      "ctx",
 		},
+		{
+			VarName:    "buckets_ptr_pos",
+			StructName: "runtime.hmap",
+			Field:      "buckets",
+		},
 	}, nil, true)
 	if err != nil {
 		return err

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

@@ -69,6 +69,7 @@ volatile const u64 url_ptr_pos;
 volatile const u64 path_ptr_pos;
 volatile const u64 ctx_ptr_pos;
 volatile const u64 headers_ptr_pos;
+volatile const u64 buckets_ptr_pos;
 
 static __always_inline struct span_context *extract_context_from_req_headers(void *headers_ptr_ptr)
 {
@@ -97,7 +98,7 @@ static __always_inline struct span_context *extract_context_from_req_headers(voi
     }
     u64 bucket_count = 1 << log_2_bucket_count;
     void *header_buckets;
-    res = bpf_probe_read(&header_buckets, sizeof(header_buckets), headers_ptr + 16);
+    res = bpf_probe_read(&header_buckets, sizeof(header_buckets), (void*)(headers_ptr + buckets_ptr_pos));
     if (res < 0)
     {
         return NULL;
@@ -216,7 +217,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;

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

@@ -106,6 +106,11 @@ func (h *Instrumentor) Load(ctx *context.InstrumentorContext) error {
 			StructName: "net/http.Request",
 			Field:      "Header",
 		},
+		{
+			VarName:    "buckets_ptr_pos",
+			StructName: "runtime.hmap",
+			Field:      "buckets",
+		},
 	}, nil, false)
 	if err != nil {
 		return err

internal/tools/offsets/main.go

@@ -53,6 +53,10 @@ func main() {
 				StructName: "runtime.g",
 				Field:      "goid",
 			},
+			{
+				StructName: "runtime.hmap",
+				Field:      "buckets",
+			},
 		})
 
 	if err != nil {