Handle distributed context headers in custom propagator

src/ReverseProxy/Forwarder/ForwarderHttpClientFactory.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Diagnostics;
 using System.Net;
 using System.Net.Http;
 using System.Text;
@@ -45,7 +46,10 @@ public HttpMessageInvoker CreateClient(ForwarderHttpClientContext context)
             UseProxy = false,
             AllowAutoRedirect = false,
             AutomaticDecompression = DecompressionMethods.None,
-            UseCookies = false
+            UseCookies = false,
+#if NET6_0_OR_GREATER
+            ActivityHeadersPropagator = Propagator
+#endif
 
             // NOTE: MaxResponseHeadersLength = 64, which means up to 64 KB of headers are allowed by default as of .NET Core 3.1.
         };
@@ -67,6 +71,15 @@ protected virtual bool CanReuseOldClient(ForwarderHttpClientContext context)
     {
         return context.OldClient != null && context.NewConfig == context.OldConfig;
     }
+#if NET6_0_OR_GREATER
+    /// <summary>
+    /// Wraps custom propagator and keeps header removal logic by default.
+    /// </summary>
+    protected virtual DistributedContextPropagator? Propagator
+    {
+        get => new ReverseProxyPropagator(DistributedContextPropagator.CreateDefaultPropagator());
+    }
+#endif
 
     /// <summary>
     /// Allows configuring the <see cref="SocketsHttpHandler"/> instance. The base implementation

src/ReverseProxy/Forwarder/RequestUtilities.cs

@@ -83,15 +83,6 @@ internal static bool ShouldSkipResponseHeader(string headerName)
 #else
         "Alt-Svc",
 #endif
-
-#if NET6_0_OR_GREATER
-        // Distributed context headers
-        HeaderNames.TraceParent,
-        HeaderNames.RequestId,
-        HeaderNames.TraceState,
-        HeaderNames.Baggage,
-        HeaderNames.CorrelationContext,
-#endif
     };
 
     // Headers marked as HttpHeaderType.Content in

src/ReverseProxy/Forwarder/ReverseProxyPropagator.cs

@@ -0,0 +1,47 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#if NET6_0_OR_GREATER
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.Linq;
+using System.Net.Http;
+
+namespace Yarp.ReverseProxy.Forwarder;
+
+public sealed class ReverseProxyPropagator : DistributedContextPropagator
+{
+    private readonly DistributedContextPropagator _innerPropagator;
+    private readonly string[] _headersToRemove;
+
+    public ReverseProxyPropagator(DistributedContextPropagator innerPropagator)
+    {
+        _innerPropagator = innerPropagator ?? throw new ArgumentNullException(nameof(innerPropagator));
+        _headersToRemove = _innerPropagator.Fields.ToArray();
+    }
+
+    public override void Inject(Activity? activity, object? carrier, PropagatorSetterCallback? setter)
+    {
+        if (carrier is HttpRequestMessage message)
+        {
+            var headers = message.Headers;
+
+            foreach (var header in _headersToRemove)
+            {
+                headers.Remove(header);
+            }
+        }
+
+        _innerPropagator.Inject(activity, carrier, setter);
+    }
+
+    public override void ExtractTraceIdAndState(object? carrier, PropagatorGetterCallback? getter, out string? traceId, out string? traceState) =>
+        _innerPropagator.ExtractTraceIdAndState(carrier, getter, out traceId, out traceState);
+
+    public override IEnumerable<KeyValuePair<string, string?>>? ExtractBaggage(object? carrier, PropagatorGetterCallback? getter) =>
+        _innerPropagator.ExtractBaggage(carrier, getter);
+
+    public override IReadOnlyCollection<string> Fields => _innerPropagator.Fields;
+}
+#endif

test/ReverseProxy.Tests/Forwarder/HttpForwarderTests.cs

@@ -2432,13 +2432,6 @@ public static IEnumerable<object[]> GetProhibitedHeaders()
             "HTTP2-Settings: value",
             "Upgrade-Insecure-Requests: value",
             "Alt-Svc: value",
-#if NET6_0_OR_GREATER
-            "traceparent: value",
-            "Request-Id: value",
-            "tracestate: value",
-            "baggage: value",
-            "Correlation-Context: value",
-#endif
         };
 
         foreach (var header in headers)

test/ReverseProxy.Tests/Forwarder/HttpTransformerTests.cs

@@ -35,13 +35,6 @@ public class HttpTransformerTests
         HeaderNames.AltSvc,
 #else
         "Alt-Svc",
-#endif
-#if NET6_0_OR_GREATER
-        HeaderNames.TraceParent,
-        HeaderNames.RequestId,
-        HeaderNames.TraceState,
-        HeaderNames.Baggage,
-        HeaderNames.CorrelationContext,
 #endif
     };