HttpLoggingMiddleware + Websockets/Upgrade

src/Middleware/HttpLogging/src/HttpLoggingMiddleware.cs

@@ -131,10 +131,25 @@ private async Task InvokeInternal(HttpContext context)
         ResponseBufferingStream? responseBufferingStream = null;
         IHttpResponseBodyFeature? originalBodyFeature = null;
 
+        UpgradeFeatureLoggingDecorator? loggableUpgradeFeature = null;
+        IHttpUpgradeFeature? originalUpgradeFeature = null;
+
         try
         {
             var response = context.Response;
 
+            if (options.LoggingFields.HasFlag(HttpLoggingFields.ResponseStatusCode) || options.LoggingFields.HasFlag(HttpLoggingFields.ResponseHeaders))
+            {
+                originalUpgradeFeature = context.Features.Get<IHttpUpgradeFeature>();
+
+                if (originalUpgradeFeature != null && originalUpgradeFeature.IsUpgradableRequest)
+                {
+                    loggableUpgradeFeature = new UpgradeFeatureLoggingDecorator(originalUpgradeFeature, response, options, _logger);
+
+                    context.Features.Set<IHttpUpgradeFeature>(loggableUpgradeFeature);
+                }
+            }
+
             if (options.LoggingFields.HasFlag(HttpLoggingFields.ResponseBody))
             {
                 originalBodyFeature = context.Features.Get<IHttpResponseBodyFeature>()!;
@@ -159,9 +174,9 @@ private async Task InvokeInternal(HttpContext context)
                 requestBufferingStream.LogRequestBody();
             }
 
-            if (responseBufferingStream == null || responseBufferingStream.FirstWrite == false)
+            if (ResponseHeadersNotYetWritten(responseBufferingStream, loggableUpgradeFeature))
             {
-                // No body, write headers here.
+                // No body, not an upgradable request or request not upgraded, write headers here. 
                 LogResponseHeaders(response, options, _logger);
             }
 
@@ -189,9 +204,29 @@ private async Task InvokeInternal(HttpContext context)
             {
                 context.Request.Body = originalBody;
             }
+
+            if (loggableUpgradeFeature != null)
+            {
+                context.Features.Set(originalUpgradeFeature);
+            }
         }
     }
 
+    private static bool ResponseHeadersNotYetWritten(ResponseBufferingStream? responseBufferingStream, UpgradeFeatureLoggingDecorator? upgradeFeatureLogging)
+    {
+        return BodyNotYetWritten(responseBufferingStream) && NotUpgradeableRequestOrRequestNotUpgraded(upgradeFeatureLogging);
+    }
+
+    private static bool BodyNotYetWritten(ResponseBufferingStream? responseBufferingStream)
+    {
+        return responseBufferingStream == null || responseBufferingStream.FirstWrite == false;
+    }
+
+    private static bool NotUpgradeableRequestOrRequestNotUpgraded(UpgradeFeatureLoggingDecorator? upgradeFeatureLogging)
+    {
+        return upgradeFeatureLogging == null || !upgradeFeatureLogging.IsUpgraded;
+    }
+
     private static void AddToList(List<KeyValuePair<string, object?>> list, string key, string? value)
     {
         list.Add(new KeyValuePair<string, object?>(key, value));

src/Middleware/HttpLogging/src/UpgradeFeatureLoggingDecorator.cs

@@ -0,0 +1,41 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Http.Features;
+using Microsoft.Extensions.Logging;
+
+namespace Microsoft.AspNetCore.HttpLogging;
+
+internal sealed class UpgradeFeatureLoggingDecorator : IHttpUpgradeFeature
+{
+    private readonly IHttpUpgradeFeature _innerUpgradeFeature;
+    private readonly HttpResponse _response;
+    private readonly HttpLoggingOptions _options;
+    private readonly ILogger _logger;
+
+    private bool _isUpgraded;
+
+    public UpgradeFeatureLoggingDecorator(IHttpUpgradeFeature innerUpgradeFeature, HttpResponse response, HttpLoggingOptions options, ILogger logger)
+    {
+        _innerUpgradeFeature = innerUpgradeFeature ?? throw new ArgumentNullException(nameof(innerUpgradeFeature));
+        _response = response ?? throw new ArgumentNullException(nameof(response));
+        _options = options ?? throw new ArgumentNullException(nameof(options));
+        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
+    }
+
+    public bool IsUpgradableRequest => _innerUpgradeFeature.IsUpgradableRequest;
+
+    public bool IsUpgraded { get => _isUpgraded; }
+
+    public async Task<Stream> UpgradeAsync()
+    {
+        var upgradeStream = await _innerUpgradeFeature.UpgradeAsync();
+
+        _isUpgraded = true;
+
+        HttpLoggingMiddleware.LogResponseHeaders(_response, _options, _logger);
+
+        return upgradeStream;
+    }
+}

src/Middleware/HttpLogging/test/HttpLoggingMiddlewareTests.cs

@@ -3,6 +3,7 @@
 
 using System.Text;
 using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Testing;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
@@ -838,6 +839,251 @@ public async Task UnrecognizedMediaType()
         Assert.Contains(TestSink.Writes, w => w.Message.Contains("Unrecognized Content-Type for body."));
     }
 
+    [Fact]
+    public async Task UpgradeToWebSocketLogsResponseStatusCodeWhenResponseIsFlushed()
+    {
+        var options = CreateOptionsAccessor();
+        options.CurrentValue.LoggingFields = HttpLoggingFields.ResponseStatusCode;
+
+        var writtenHeaders = new TaskCompletionSource();
+        var letBodyFinish = new TaskCompletionSource();
+
+        var httpContext = new DefaultHttpContext();
+
+        var upgradeFeatureMock = new Mock<IHttpUpgradeFeature>();
+        upgradeFeatureMock.Setup(m => m.IsUpgradableRequest).Returns(true);
+        upgradeFeatureMock
+            .Setup(m => m.UpgradeAsync())
+            .Callback(() =>
+            {
+                httpContext.Response.StatusCode = StatusCodes.Status101SwitchingProtocols;
+                httpContext.Response.Headers[HeaderNames.Connection] = HeaderNames.Upgrade;
+            })
+            .ReturnsAsync(Stream.Null);
+        httpContext.Features.Set<IHttpUpgradeFeature>(upgradeFeatureMock.Object);
+
+        var middleware = new HttpLoggingMiddleware(
+            async c =>
+            {
+                await c.Features.Get<IHttpUpgradeFeature>().UpgradeAsync();
+                writtenHeaders.SetResult();
+                await letBodyFinish.Task;
+            },
+            options,
+            LoggerFactory.CreateLogger<HttpLoggingMiddleware>());
+
+        var middlewareTask = middleware.Invoke(httpContext);
+
+        await writtenHeaders.Task;
+
+        Assert.Contains(TestSink.Writes, w => w.Message.Contains("StatusCode: 101"));
+
+        letBodyFinish.SetResult();
+
+        await middlewareTask;
+    }
+
+    [Fact]
+    public async Task UpgradeToWebSocketLogsResponseHeadersWhenResponseIsFlushed()
+    {
+        var options = CreateOptionsAccessor();
+        options.CurrentValue.LoggingFields = HttpLoggingFields.ResponseHeaders;
+
+        var writtenHeaders = new TaskCompletionSource();
+        var letBodyFinish = new TaskCompletionSource();
+
+        var httpContext = new DefaultHttpContext();
+
+        var upgradeFeatureMock = new Mock<IHttpUpgradeFeature>();
+        upgradeFeatureMock.Setup(m => m.IsUpgradableRequest).Returns(true);
+        upgradeFeatureMock
+            .Setup(m => m.UpgradeAsync())
+            .Callback(() =>
+            {
+                httpContext.Response.StatusCode = StatusCodes.Status101SwitchingProtocols;
+                httpContext.Response.Headers[HeaderNames.Connection] = HeaderNames.Upgrade;
+            })
+            .ReturnsAsync(Stream.Null);
+        httpContext.Features.Set<IHttpUpgradeFeature>(upgradeFeatureMock.Object);
+
+        var middleware = new HttpLoggingMiddleware(
+            async c =>
+            {
+                await c.Features.Get<IHttpUpgradeFeature>().UpgradeAsync();
+                writtenHeaders.SetResult();
+                await letBodyFinish.Task;
+            },
+            options,
+            LoggerFactory.CreateLogger<HttpLoggingMiddleware>());
+
+        var middlewareTask = middleware.Invoke(httpContext);
+
+        await writtenHeaders.Task;
+
+        Assert.Contains(TestSink.Writes, w => w.Message.Contains("Connection: Upgrade"));
+
+        letBodyFinish.SetResult();
+
+        await middlewareTask;
+    }
+
+    [Fact]
+    public async Task UpgradeToWebSocketDoesNotLogWhenResponseIsFlushedIfLoggingOptionsAreOtherThanResponseStatusCodeOrResponseHeaders()
+    {
+        var options = CreateOptionsAccessor();
+        options.CurrentValue.LoggingFields = HttpLoggingFields.All ^ HttpLoggingFields.ResponsePropertiesAndHeaders;
+
+        var writtenHeaders = new TaskCompletionSource();
+        var letBodyFinish = new TaskCompletionSource();
+
+        var httpContext = new DefaultHttpContext();
+
+        var upgradeFeatureMock = new Mock<IHttpUpgradeFeature>();
+        upgradeFeatureMock.Setup(m => m.IsUpgradableRequest).Returns(true);
+        upgradeFeatureMock
+            .Setup(m => m.UpgradeAsync())
+            .Callback(() =>
+            {
+                httpContext.Response.StatusCode = StatusCodes.Status101SwitchingProtocols;
+                httpContext.Response.Headers[HeaderNames.Connection] = HeaderNames.Upgrade;
+            })
+            .ReturnsAsync(Stream.Null);
+        httpContext.Features.Set<IHttpUpgradeFeature>(upgradeFeatureMock.Object);
+
+        var middleware = new HttpLoggingMiddleware(
+            async c =>
+            {
+                await c.Features.Get<IHttpUpgradeFeature>().UpgradeAsync();
+                writtenHeaders.SetResult();
+                await letBodyFinish.Task;
+            },
+            options,
+            LoggerFactory.CreateLogger<HttpLoggingMiddleware>());
+
+        var middlewareTask = middleware.Invoke(httpContext);
+
+        await writtenHeaders.Task;
+
+        Assert.DoesNotContain(TestSink.Writes, w => w.Message.Contains("StatusCode: 101"));
+        Assert.DoesNotContain(TestSink.Writes, w => w.Message.Contains("Connection: Upgrade"));
+
+        letBodyFinish.SetResult();
+
+        await middlewareTask;
+    }
+
+    [Theory]
+    [InlineData(false)]
+    [InlineData(true)]
+    public async Task LogsWrittenOutsideUpgradeWrapperIfUpgradeDoesNotOccur(bool isUpgradableRequest)
+    {
+        var options = CreateOptionsAccessor();
+        options.CurrentValue.LoggingFields = HttpLoggingFields.ResponsePropertiesAndHeaders;
+
+        var httpContext = new DefaultHttpContext();
+
+        var upgradeFeatureMock = new Mock<IHttpUpgradeFeature>();
+        upgradeFeatureMock.Setup(m => m.IsUpgradableRequest).Returns(isUpgradableRequest);
+        httpContext.Features.Set<IHttpUpgradeFeature>(upgradeFeatureMock.Object);
+
+        var middleware = new HttpLoggingMiddleware(
+            async c =>
+            {
+                c.Response.StatusCode = 200;
+                c.Response.Headers[HeaderNames.TransferEncoding] = "test";
+                c.Response.ContentType = "text/plain";
+                await c.Response.StartAsync();
+            },
+            options,
+            LoggerFactory.CreateLogger<HttpLoggingMiddleware>());
+
+        var middlewareTask = middleware.Invoke(httpContext);
+
+        Assert.Contains(TestSink.Writes, w => w.Message.Contains("StatusCode: 200"));
+        Assert.Contains(TestSink.Writes, w => w.Message.Contains("Transfer-Encoding: test"));
+        Assert.Contains(TestSink.Writes, w => w.Message.Contains("Content-Type: text/plain"));
+        Assert.Null(
+            Record.Exception(() => upgradeFeatureMock.Verify(m => m.UpgradeAsync(), Times.Never)));
+
+        await middlewareTask;
+    }
+
+    [Theory]
+    [InlineData(HttpLoggingFields.ResponseStatusCode)]
+    [InlineData(HttpLoggingFields.ResponseHeaders)]
+    public async Task UpgradeToWebSocketLogsWrittenOnlyOnce(HttpLoggingFields loggingFields)
+    {
+        var options = CreateOptionsAccessor();
+        options.CurrentValue.LoggingFields = loggingFields;
+
+        var httpContext = new DefaultHttpContext();
+
+        var upgradeFeatureMock = new Mock<IHttpUpgradeFeature>();
+        upgradeFeatureMock.Setup(m => m.IsUpgradableRequest).Returns(true);
+        upgradeFeatureMock
+            .Setup(m => m.UpgradeAsync())
+            .Callback(() =>
+            {
+                httpContext.Response.StatusCode = StatusCodes.Status101SwitchingProtocols;
+                httpContext.Response.Headers[HeaderNames.Connection] = HeaderNames.Upgrade;
+            })
+            .ReturnsAsync(Stream.Null);
+        httpContext.Features.Set<IHttpUpgradeFeature>(upgradeFeatureMock.Object);
+
+        var writeCount = 0;
+        TestSink.MessageLogged += (context) => { writeCount++; };
+
+        var middleware = new HttpLoggingMiddleware(
+            async c =>
+            {
+                await c.Features.Get<IHttpUpgradeFeature>().UpgradeAsync();
+            },
+            options,
+            LoggerFactory.CreateLogger<HttpLoggingMiddleware>());
+
+        await middleware.Invoke(httpContext);
+
+        Assert.Equal(1, writeCount);
+    }
+
+    [Theory]
+    [InlineData(HttpLoggingFields.ResponseStatusCode)]
+    [InlineData(HttpLoggingFields.ResponseHeaders)]
+    public async Task OriginalUpgradeFeatureIsRestoredBeforeMiddlewareCompletes(HttpLoggingFields loggingFields)
+    {
+        var options = CreateOptionsAccessor();
+        options.CurrentValue.LoggingFields = loggingFields;
+
+        var letBodyFinish = new TaskCompletionSource();
+
+        var httpContext = new DefaultHttpContext();
+
+        var upgradeFeatureMock = new Mock<IHttpUpgradeFeature>();
+        upgradeFeatureMock.Setup(m => m.IsUpgradableRequest).Returns(true);
+        upgradeFeatureMock.Setup(m => m.UpgradeAsync()).ReturnsAsync(Stream.Null);
+        httpContext.Features.Set<IHttpUpgradeFeature>(upgradeFeatureMock.Object);
+
+        IHttpUpgradeFeature upgradeFeature = null;
+
+        var middleware = new HttpLoggingMiddleware(
+            async c =>
+            {
+                upgradeFeature = c.Features.Get<IHttpUpgradeFeature>();
+                await letBodyFinish.Task;
+            },
+            options,
+            LoggerFactory.CreateLogger<HttpLoggingMiddleware>());
+
+        var middlewareTask = middleware.Invoke(httpContext);
+
+        Assert.True(upgradeFeature is UpgradeFeatureLoggingDecorator);
+
+        letBodyFinish.SetResult();
+        await middlewareTask;
+
+        Assert.False(httpContext.Features.Get<IHttpUpgradeFeature>() is UpgradeFeatureLoggingDecorator);
+    }
+
     private IOptionsMonitor<HttpLoggingOptions> CreateOptionsAccessor()
     {
         var options = new HttpLoggingOptions();