Handle distributed context headers in custom propagator (#1533)

* Handle distributed context headers in custom propagator

* Update src/ReverseProxy/Forwarder/ForwarderHttpClientFactory.cs

Co-authored-by: Miha Zupan <mihazupan.zupan1@gmail.com>

* Addressing review feedback

* Improve docs

* Apply suggestions from code review

Co-authored-by: Miha Zupan <mihazupan.zupan1@gmail.com>

* Apply suggestions from code review

Co-authored-by: Miha Zupan <mihazupan.zupan1@gmail.com>

* Add ActivityHeadersPropagator into Direct.Sample

Co-authored-by: Miha Zupan <mihazupan.zupan1@gmail.com>

docs/docfx/articles/direct-forwarding.md

@@ -53,13 +53,13 @@ public void Configure(IApplicationBuilder app, IHttpForwarder forwarder)
         UseProxy = false,
         AllowAutoRedirect = false,
         AutomaticDecompression = DecompressionMethods.None,
-        UseCookies = false
+        UseCookies = false,
+        ActivityHeadersPropagator = new ReverseProxyPropagator(DistributedContextPropagator.Current)
     });
     var transformer = new CustomTransformer(); // or HttpTransformer.Default;
     var requestConfig = new ForwarderRequestConfig { ActivityTimeout = TimeSpan.FromSeconds(100) };
 
     app.UseRouting();
-    app.UseAuthorization();
     app.UseEndpoints(endpoints =>
     {
         endpoints.Map("/{**catch-all}", async httpContext =>
@@ -83,9 +83,18 @@ private class CustomTransformer : HttpTransformer
     public override async ValueTask TransformRequestAsync(HttpContext httpContext,
         HttpRequestMessage proxyRequest, string destinationPrefix)
     {
-        // Copy headers normally and then remove the original host.
-        // Use the destination host from proxyRequest.RequestUri instead.
+        // Copy all request headers
         await base.TransformRequestAsync(httpContext, proxyRequest, destinationPrefix);
+
+        // Customize the query string:
+        var queryContext = new QueryTransformContext(httpContext.Request);
+        queryContext.Collection.Remove("param1");
+        queryContext.Collection["area"] = "xx2";
+
+        // Assign the custom uri. Be careful about extra slashes when concatenating here. RequestUtilities.MakeDestinationAddress is a safe default.
+        proxyRequest.RequestUri = RequestUtilities.MakeDestinationAddress("https://example.com", httpContext.Request.Path, queryContext.QueryString);
+
+        // Suppress the original request header, use the one from the destination Uri.
         proxyRequest.Headers.Host = null;
     }
 }

docs/docfx/articles/header-guidelines.md

@@ -38,6 +38,11 @@ This response header is used with HTTP/3 upgrades and only applies to the immedi
 ### TraceParent, Request-Id, TraceState, Baggage, Correlation-Context
 
 These headers relate to distributed tracing. They are automatically removed on .NET 6 or later so that the forwarding HttpClient can replace them with updated values.
+You can opt out of modifying these headers by setting `SocketsHttpHandler.ActivityHeadersPropagator ` to `null`:
+```C#
+services.AddReverseProxy()
+    .ConfigureHttpClient((_, handler) => handler.ActivityHeadersPropagator = null);
+```
 
 ## Other Header Guidelines
 

samples/ReverseProxy.Direct.Sample/Startup.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Diagnostics;
 using System.Net;
 using System.Net.Http;
 using System.Threading.Tasks;
@@ -38,7 +39,10 @@ public void Configure(IApplicationBuilder app, IHttpForwarder forwarder)
                 UseProxy = false,
                 AllowAutoRedirect = false,
                 AutomaticDecompression = DecompressionMethods.None,
-                UseCookies = false
+                UseCookies = false,
+#if NET6_0_OR_GREATER
+                ActivityHeadersPropagator = new ReverseProxyPropagator(DistributedContextPropagator.Current)
+#endif
             });
 
             // Setup our own request transform class

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 = new ReverseProxyPropagator(DistributedContextPropagator.Current)
+#endif
 
             // NOTE: MaxResponseHeadersLength = 64, which means up to 64 KB of headers are allowed by default as of .NET Core 3.1.
         };

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,53 @@
+// 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;
+
+/// <summary>
+/// Removes existing headers and then delegates to the inner propagator.
+/// </summary>
+public sealed class ReverseProxyPropagator : DistributedContextPropagator
+{
+    private readonly DistributedContextPropagator _innerPropagator;
+    private readonly string[] _headersToRemove;
+
+    /// <summary>
+    /// ReverseProxyPropagator removes headers pointed out in innerPropagator.
+    /// </summary>
+    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
     };