Add tests for manual injection/extraction and guard invalid implementations (#6998)

## Summary of changes

Guard against incorrect inject/extract propagation implementations

## Reason for change

We saw an error where an injector/extractor was returning `null` during
context extraction, which it should not do.
#6991 introduced a `null`
check to avoid throwing. This PR adds additional safety to the checks,
explicitly related to the manual-instrumentation implementation, and
adds extra tests.

## Implementation details

- Wrap each of the functions passed in the manual API with a "safe"
extractor/injector to guard against exceptions/nulls
- We just swallow the errors, but that's what would happen _anyway_, and
this way we don't unactionable logs.
- It might be nice to have a way to _actually_ bubble-up specific
non-custom exceptions someway, but not critical, and would require a
bunch of changes, so meh.

## Test coverage

- Add unit tests for the integrations to confirm the wrapping handles
errors
- Add integration tests for the instrumentation

## Other details

Related to #6991

Author: andrewlock

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/ManualInstrumentation/Propagators/SafeExtractor.cs

@@ -0,0 +1,39 @@
+// <copyright file="SafeExtractor.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+#nullable enable
+using System;
+using System.Collections.Generic;
+
+namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.ManualInstrumentation.Propagators;
+
+/// <summary>
+/// An extractor we use to wrap customer extraction functions
+/// to force nullable constraints and handle exceptions.
+/// </summary>
+internal readonly struct SafeExtractor<TCarrier>
+{
+    private readonly Func<TCarrier, string, IEnumerable<string?>> _extractor;
+
+    public SafeExtractor(Func<TCarrier, string, IEnumerable<string?>> extractor)
+    {
+        _extractor = extractor;
+    }
+
+    public IEnumerable<string?> SafeExtract(TCarrier carrier, string key)
+    {
+        try
+        {
+            return _extractor(carrier, key) ?? [];
+        }
+        catch (Exception)
+        {
+            // This is a customer error, so we don't want to log it
+            // There's a question of how/if we should throw _sometihng_ here to propagate the error
+            // to the customer, but for now, just swallow it;
+            return [];
+        }
+    }
+}

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/ManualInstrumentation/Propagators/SafeInjector.cs

@@ -0,0 +1,38 @@
+// <copyright file="SafeInjector.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+#nullable enable
+using System;
+using System.Collections.Generic;
+
+namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.ManualInstrumentation.Propagators;
+
+/// <summary>
+/// An injector we use to wrap customer extraction functions
+/// to force nullable constraints and handle exceptions.
+/// </summary>
+internal readonly struct SafeInjector<TCarrier>
+{
+    private readonly Action<TCarrier, string, string> _injector;
+
+    public SafeInjector(Action<TCarrier, string, string> injector)
+    {
+        _injector = injector;
+    }
+
+    public void SafeInject(TCarrier carrier, string key, string value)
+    {
+        try
+        {
+            _injector(carrier, key, value);
+        }
+        catch (Exception)
+        {
+            // This is a customer error, so we don't want to log it
+            // There's a question of how/if we should throw _sometihng_ here to propagate the error
+            // to the customer, but for now, just swallow it;
+        }
+    }
+}

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/ManualInstrumentation/Propagators/SpanContextExtractorExtractIncludingDsmIntegration.cs

@@ -35,8 +35,13 @@ public class SpanContextExtractorExtractIncludingDsmIntegration
     internal static CallTargetState OnMethodBegin<TTarget, TCarrier, TAction>(TTarget instance, in TCarrier carrier, in TAction getter, string messageType, string source)
     {
         TelemetryFactory.Metrics.Record(PublicApiUsage.SpanContextExtractor_ExtractIncludingDsm);
+
+        // we can't trust that the customer honours the nullable reference types here,
+        // so we wrap the method in a try/catch and ensure we always return non-null
         var extract = (Func<TCarrier, string, IEnumerable<string?>>)(object)getter!;
-        var extracted = SpanContextExtractor.ExtractInternal(carrier, extract, messageType, source);
+        var extractor = new SafeExtractor<TCarrier>(extract);
+
+        var extracted = SpanContextExtractor.ExtractInternal(carrier, extractor.SafeExtract, messageType, source);
         return new CallTargetState(scope: null, extracted);
     }
 

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/ManualInstrumentation/Propagators/SpanContextExtractorExtractIntegration.cs

@@ -34,8 +34,13 @@ public class SpanContextExtractorExtractIntegration
     internal static CallTargetState OnMethodBegin<TTarget, TCarrier, TAction>(TTarget instance, in TCarrier carrier, in TAction getter)
     {
         TelemetryFactory.Metrics.Record(PublicApiUsage.SpanContextExtractor_Extract);
+
+        // we can't trust that the customer honours the nullable reference types here,
+        // so we wrap the method in a try/catch and ensure we always return non-null
         var extract = (Func<TCarrier, string, IEnumerable<string?>>)(object)getter!;
-        var extracted = SpanContextExtractor.ExtractInternal(carrier, extract);
+        var extractor = new SafeExtractor<TCarrier>(extract);
+
+        var extracted = SpanContextExtractor.ExtractInternal(carrier, extractor.SafeExtract);
         return new CallTargetState(scope: null, state: extracted);
     }
 

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/ManualInstrumentation/Propagators/SpanContextInjectorInjectIncludingDsmIntegration.cs

@@ -33,11 +33,14 @@ internal static CallTargetState OnMethodBegin<TTarget, TCarrier, TAction, TSpanC
         // The Injector.Inject method currently _only_ works with SpanContext objects
         // Therefore, there's no point calling inject unless we can remap it to a SpanContext
         TelemetryFactory.Metrics.Record(PublicApiUsage.SpanContextInjector_Inject);
-        var inject = (Action<TCarrier, string, string>)(object)setter!;
 
         if (SpanContextHelper.GetContext(context) is { } spanContext)
         {
-            SpanContextInjector.InjectInternal(carrier, inject, spanContext, messageType, target);
+            // We can't trust that the customer function is safe,
+            // so we wrap the method in a try/catch to ensure we don't throw
+            var inject = (Action<TCarrier, string, string>)(object)setter!;
+            var injector = new SafeInjector<TCarrier>(inject);
+            SpanContextInjector.InjectInternal(carrier, injector.SafeInject, spanContext, messageType, target);
         }
 
         return CallTargetState.GetDefault();

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/ManualInstrumentation/Propagators/SpanContextInjectorInjectIntegration.cs

@@ -33,11 +33,14 @@ internal static CallTargetState OnMethodBegin<TTarget, TCarrier, TAction, TSpanC
         // The Injector.Inject method currently _only_ works with SpanContext objects
         // Therefore, there's no point calling inject unless we can remap it to a SpanContext
         TelemetryFactory.Metrics.Record(PublicApiUsage.SpanContextInjector_Inject);
-        var inject = (Action<TCarrier, string, string>)(object)setter!;
 
         if (SpanContextHelper.GetContext(context) is { } spanContext)
         {
-            SpanContextInjector.InjectInternal(carrier, inject, spanContext);
+            // We can't trust that the customer function is safe,
+            // so we wrap the method in a try/catch to ensure we don't throw
+            var inject = (Action<TCarrier, string, string>)(object)setter!;
+            var injector = new SafeInjector<TCarrier>(inject);
+            SpanContextInjector.InjectInternal(carrier, injector.SafeInject, spanContext);
         }
 
         return CallTargetState.GetDefault();

tracer/test/Datadog.Trace.Tests/ManualInstrumentation/SpanContextPropagationTests.cs

@@ -0,0 +1,89 @@
+// <copyright file="SpanContextPropagationTests.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+using System;
+using System.Collections.Generic;
+using Datadog.Trace.ClrProfiler.AutoInstrumentation.ManualInstrumentation.Propagators;
+using FluentAssertions;
+using Xunit;
+
+namespace Datadog.Trace.Tests.ManualInstrumentation;
+
+public class SpanContextPropagationTests
+{
+    [Fact]
+    public void SpanContextExtractorExtractIntegration_WhenUserProvidedGetterReturnsNull_IntegrationCatchesException()
+    {
+        var headers = new Dictionary<string, string[]>();
+        Func<Dictionary<string, string[]>, string, IEnumerable<string>> getter = (dict, key) => null;
+
+        // instance isn't used, so just passing this to avoid generic gymnastics
+        var state = SpanContextExtractorExtractIntegration.OnMethodBegin(this, headers, getter);
+        var result = state.State as SpanContext;
+        result.Should().BeNull();
+    }
+
+    [Fact]
+    public void SpanContextExtractorExtractIntegration_WhenUserProvidedGetterThrowsException_IntegrationCatchesException()
+    {
+        var headers = new Dictionary<string, string[]>();
+        Func<Dictionary<string, string[]>, string, IEnumerable<string>> getter = (dict, key) => throw new Exception("oops!");
+
+        // instance isn't used, so just passing this to avoid generic gymnastics
+        var state = SpanContextExtractorExtractIntegration.OnMethodBegin(this, headers, getter);
+        var result = state.State as SpanContext;
+        result.Should().BeNull();
+    }
+
+    [Fact]
+    public void SpanContextExtractorExtractIncludingDsmIntegration_WhenUserProvidedGetterReturnsNull_IntegrationCatchesException()
+    {
+        var headers = new Dictionary<string, string[]>();
+        Func<Dictionary<string, string[]>, string, IEnumerable<string>> getter = (dict, key) => null;
+
+        // instance isn't used, so just passing this to avoid generic gymnastics
+        var state = SpanContextExtractorExtractIncludingDsmIntegration.OnMethodBegin(this, headers, getter, "messageType", "source");
+        var result = state.State as SpanContext;
+        result.Should().BeNull();
+    }
+
+    [Fact]
+    public void SpanContextExtractorExtractIncludingDsmIntegration_WhenUserProvidedGetterThrowsException_IntegrationCatchesException()
+    {
+        var headers = new Dictionary<string, string[]>();
+        Func<Dictionary<string, string[]>, string, IEnumerable<string>> getter = (dict, key) => throw new Exception("oops!");
+
+        // instance isn't used, so just passing this to avoid generic gymnastics
+        var state = SpanContextExtractorExtractIncludingDsmIntegration.OnMethodBegin(this, headers, getter, "messageType", "source");
+        var result = state.State as SpanContext;
+        result.Should().BeNull();
+    }
+
+    [Fact]
+    public void SpanContextInjectorInjectIntegration_WhenUserProvidedGetterThrowsException_IntegrationCatchesException()
+    {
+        var spanContext = new SpanContext(123, 123, SamplingPriority.UserKeep);
+        var headers = new Dictionary<string, string[]>();
+        Action<Dictionary<string, string[]>, string, string> setter = (dict, key, value) => throw new Exception("oops!");
+
+        // instance isn't used, so just passing this to avoid generic gymnastics
+        var state = SpanContextInjectorInjectIntegration.OnMethodBegin(this, headers, setter, spanContext);
+        var result = state.State as SpanContext;
+        result.Should().BeNull();
+    }
+
+    [Fact]
+    public void SpanContextInjectorInjectIncludingDsmIntegration_WhenUserProvidedGetterThrowsException_IntegrationCatchesException()
+    {
+        var headers = new Dictionary<string, string[]>();
+        var spanContext = new SpanContext(123, 123, SamplingPriority.UserKeep);
+        Action<Dictionary<string, string[]>, string, string> setter = (dict, key, value) => throw new Exception("oops!");
+
+        // instance isn't used, so just passing this to avoid generic gymnastics
+        var state = SpanContextInjectorInjectIncludingDsmIntegration.OnMethodBegin(this, headers, setter, spanContext, "messageType", "source");
+        var result = state.State as SpanContext;
+        result.Should().BeNull();
+    }
+}

tracer/test/test-applications/integrations/Samples.ManualInstrumentation/Program.cs

@@ -255,6 +255,64 @@ async Task OtherStuff()
             {
                 s1.Span.ServiceName = Tracer.Instance.DefaultServiceName;
                 await SendHttpRequest("CustomContext");
+
+                // Test injection
+                Dictionary<string, List<string>> headers = new();
+                new SpanContextInjector().Inject(
+                    headers,
+                    setter: (dict, key, value) => headers[key] = new List<string> { value },
+                    s1.Span.Context);
+                var context = new SpanContextExtractor().Extract(
+                    headers,
+                    getter: (dict, key) => dict.TryGetValue(key, out var values) ? values : Enumerable.Empty<string>());
+
+                Expect(context is not null, "Extracted context should not be null");
+                Expect(s1.Span.Context.SpanId == context?.SpanId, "SpanId should be extracted");
+                Expect(s1.Span.Context.TraceId == context?.TraceId, "TraceId should be extracted");
+
+                // Test DSM injection
+                Dictionary<string, List<string>> dsmHeaders = new();
+                new SpanContextInjector().InjectIncludingDsm(
+                    dsmHeaders,
+                    setter: (dict, key, value) => dsmHeaders[key] = new List<string> { value },
+                    s1.Span.Context,
+                    "messageType",
+                    "messageId");
+                var dsmContext = new SpanContextExtractor().ExtractIncludingDsm(
+                    dsmHeaders,
+                    getter: (dict, key) => dict.TryGetValue(key, out var values) ? values : Enumerable.Empty<string>(),
+                    "messageType",
+                    "messageId");
+                Expect(dsmContext is not null, "Extracted context should not be null");
+                Expect(s1.Span.Context.SpanId == dsmContext?.SpanId, "SpanId should be extracted");
+                Expect(s1.Span.Context.TraceId == dsmContext?.TraceId, "TraceId should be extracted");
+
+                // Test that we handle incorrect (null returning) implementations
+                var nullContext1 = new SpanContextExtractor().Extract(
+                    headers,
+                    getter: (dict, key) => null); // Always return null
+                Expect(nullContext1 is null, "Extracted context should be null");
+
+                var nullContext2 = new SpanContextExtractor().ExtractIncludingDsm(
+                    dsmHeaders,
+                    getter: (dict, key) => null, // Always return null
+                    "messageType",
+                    "messageId");
+                Expect(nullContext2 is null, "Extracted context should be null");
+
+                // Exceptions thrown in the faulty injector/extractor will not bubble up
+                // as they're caught in the integration, but they will write error logs if we don't handle them
+                // so will be caught in the CheckLogsForErrors stage in CI.
+                new SpanContextInjector().Inject(
+                    headers,
+                    setter: (dict, key, value) => throw new Exception("Throwing exception that should be ignored"),
+                    s1.Span.Context);
+                new SpanContextInjector().InjectIncludingDsm(
+                    headers,
+                    setter: (dict, key, value) => throw new Exception("Throwing exception that should be ignored"),
+                    s1.Span.Context,
+                    "messageType",
+                    "messageId");
             }
 
             // Manually disable debug logs