Skip to content

Commit

Permalink
Handle distributed context headers in custom propagator
Browse files Browse the repository at this point in the history
  • Loading branch information
Katya Sokolova committed Jan 26, 2022
1 parent a323abe commit 58a7e18
Show file tree
Hide file tree
Showing 5 changed files with 58 additions and 24 deletions.
15 changes: 14 additions & 1 deletion src/ReverseProxy/Forwarder/ForwarderHttpClientFactory.cs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
// Licensed under the MIT License.

using System;
using System.Diagnostics;
using System.Net;
using System.Net.Http;
using System.Text;
Expand Down Expand Up @@ -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.
};
Expand All @@ -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
Expand Down
9 changes: 0 additions & 9 deletions src/ReverseProxy/Forwarder/RequestUtilities.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
44 changes: 44 additions & 0 deletions src/ReverseProxy/Forwarder/ReverseProxyPropagator.cs
Original file line number Diff line number Diff line change
@@ -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
7 changes: 0 additions & 7 deletions test/ReverseProxy.Tests/Forwarder/HttpForwarderTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
7 changes: 0 additions & 7 deletions test/ReverseProxy.Tests/Forwarder/HttpTransformerTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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
};

Expand Down

0 comments on commit 58a7e18

Please sign in to comment.