From 58a7e18f4cca323d6087b40b16c5694c73568c3e Mon Sep 17 00:00:00 2001
From: Katya Sokolova <esokolov@microsoft.com>
Date: Wed, 26 Jan 2022 19:01:05 +0100
Subject: [PATCH] Handle distributed context headers in custom propagator

---
 .../Forwarder/ForwarderHttpClientFactory.cs   | 15 ++++++-
 .../Forwarder/RequestUtilities.cs             |  9 ----
 .../Forwarder/ReverseProxyPropagator.cs       | 44 +++++++++++++++++++
 .../Forwarder/HttpForwarderTests.cs           |  7 ---
 .../Forwarder/HttpTransformerTests.cs         |  7 ---
 5 files changed, 58 insertions(+), 24 deletions(-)
 create mode 100644 src/ReverseProxy/Forwarder/ReverseProxyPropagator.cs

diff --git a/src/ReverseProxy/Forwarder/ForwarderHttpClientFactory.cs b/src/ReverseProxy/Forwarder/ForwarderHttpClientFactory.cs
index 81dad1a85c..0aa828e720 100644
--- a/src/ReverseProxy/Forwarder/ForwarderHttpClientFactory.cs
+++ b/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
diff --git a/src/ReverseProxy/Forwarder/RequestUtilities.cs b/src/ReverseProxy/Forwarder/RequestUtilities.cs
index 1a99e7e85d..82ca843a80 100644
--- a/src/ReverseProxy/Forwarder/RequestUtilities.cs
+++ b/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
diff --git a/src/ReverseProxy/Forwarder/ReverseProxyPropagator.cs b/src/ReverseProxy/Forwarder/ReverseProxyPropagator.cs
new file mode 100644
index 0000000000..23029d880f
--- /dev/null
+++ b/src/ReverseProxy/Forwarder/ReverseProxyPropagator.cs
@@ -0,0 +1,44 @@
+#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
diff --git a/test/ReverseProxy.Tests/Forwarder/HttpForwarderTests.cs b/test/ReverseProxy.Tests/Forwarder/HttpForwarderTests.cs
index d6367e461d..65713d2522 100644
--- a/test/ReverseProxy.Tests/Forwarder/HttpForwarderTests.cs
+++ b/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)
diff --git a/test/ReverseProxy.Tests/Forwarder/HttpTransformerTests.cs b/test/ReverseProxy.Tests/Forwarder/HttpTransformerTests.cs
index 4cbf881727..17d4af5bb1 100644
--- a/test/ReverseProxy.Tests/Forwarder/HttpTransformerTests.cs
+++ b/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
     };