HeaderPropagationV2.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Net.Http;
using System.Threading;
using Microsoft.AspNetCore;
using Microsoft.AspNetCore.Hosting;
using Microsoft.AspNetCore.Http;

namespace header_propagation
{
    public class Program
    {
        public static void Main(string[] args)
        {
            using (DiagnosticListener.AllListeners.Subscribe(new HeaderInjectionObserver()))
            {
                CreateWebHostBuilder(args).Build().Run();
            }
        }

        public static IWebHostBuilder CreateWebHostBuilder(string[] args) =>
            WebHost.CreateDefaultBuilder(args)
                .UseStartup<Startup>();
    }

    public class HeaderInjectionObserver : IObserver<DiagnosticListener>, IObserver<KeyValuePair<string, object>>
    {
        private static AsyncLocal<HttpContext> Context = new AsyncLocal<HttpContext>();

        public void OnCompleted()
        {

        }

        public void OnError(Exception error)
        {

        }

        public void OnNext(DiagnosticListener value)
        {
            if (value.Name == "HttpHandlerDiagnosticListener")
            {
                value.Subscribe(this);
            }

            if (value.Name == "Microsoft.AspNetCore")
            {
                value.Subscribe(this);
            }
        }

        public void OnNext(KeyValuePair<string, object> value)
        {
            // Incoming requests
            if (value.Key == "Microsoft.AspNetCore.Hosting.HttpRequestIn.Start")
            {
                var property = value.Value.GetType().GetProperty("HttpContext");
                HeaderInjectionObserver.Context.Value = (HttpContext)property.GetValue(value.Value);
            }

            if (value.Key == "Microsoft.AspNetCore.Hosting.HttpRequestOut.Start")
            {
                // Clear the async local
                HeaderInjectionObserver.Context.Value = null;
            }

            // Outgoing requests
            if (value.Key == "System.Net.Http.HttpRequestOut.Start")
            {
                var property = value.Value.GetType().GetProperty("Request");

                var req = (HttpRequestMessage)property.GetValue(value.Value);

                var httpContext = HeaderInjectionObserver.Context.Value;
                if (httpContext != null)
                {
                    var incomingHeader = httpContext.Request.Headers["incoming-header"];
                    // This mutation is EVIL!
                    req.Headers.TryAddWithoutValidation("outgoing-header", incomingHeader.ToArray());
                }
            }
        }
    }
}