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.ResponsePropertiesAndHeaders))
+            {
+                originalUpgradeFeature = context.Features.Get<IHttpUpgradeFeature>();
+
+                if (originalUpgradeFeature != null && originalUpgradeFeature.IsUpgradableRequest)
+                {
+                    loggableUpgradeFeature = new UpgradeFeatureLoggingDecorator(originalUpgradeFeature, () => LogResponseHeaders(context.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 (ResponseHeadersNotWrittenYet(responseBufferingStream, loggableUpgradeFeature))
             {
-                // No body, write headers here.
+                // No body and not an upgradable request, write headers here.
                 LogResponseHeaders(response, options, _logger);
             }
 
@@ -189,9 +204,20 @@ private async Task InvokeInternal(HttpContext context)
             {
                 context.Request.Body = originalBody;
             }
+
+            if (loggableUpgradeFeature != null)
+            {
+                context.Features.Set(originalUpgradeFeature);
+            }
         }
     }
 
+    private static bool ResponseHeadersNotWrittenYet(ResponseBufferingStream? responseBufferingStream, UpgradeFeatureLoggingDecorator? upgradeFeatureLogging)
+    {
+        return (responseBufferingStream == null || responseBufferingStream.FirstWrite == false) &&
+               (upgradeFeatureLogging == null || !upgradeFeatureLogging.IsUpgradableRequest);
+    }
+
     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,29 @@
+// 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.Features;
+
+namespace Microsoft.AspNetCore.HttpLogging;
+
+internal sealed class UpgradeFeatureLoggingDecorator : IHttpUpgradeFeature
+{
+    private readonly IHttpUpgradeFeature _innerUpgradeFeature;
+    private readonly Action _loggingDelegate;
+
+    public UpgradeFeatureLoggingDecorator(IHttpUpgradeFeature innerUpgradeFeature, Action loggingDelegate)
+    {
+        _innerUpgradeFeature = innerUpgradeFeature ?? throw new ArgumentNullException(nameof(innerUpgradeFeature));
+        _loggingDelegate = loggingDelegate ?? throw new ArgumentNullException(nameof(loggingDelegate));
+    }
+
+    public bool IsUpgradableRequest => _innerUpgradeFeature.IsUpgradableRequest;
+
+    public async Task<Stream> UpgradeAsync()
+    {
+        var upgradeTask = await _innerUpgradeFeature.UpgradeAsync();
+
+        _loggingDelegate();
+
+        return upgradeTask;
+    }
+}

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,125 @@ public async Task UnrecognizedMediaType()
         Assert.Contains(TestSink.Writes, w => w.Message.Contains("Unrecognized Content-Type for body."));
     }
 
+    [Fact]
+    public async Task UpgradeToWebSocketLogsResponsePropertiesAndHeadersWhenResponseIsFlushed()
+    {
+        var options = CreateOptionsAccessor();
+        options.CurrentValue.LoggingFields = 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();
+                await c.Response.CompleteAsync();
+                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"));
+        Assert.Contains(TestSink.Writes, w => w.Message.Contains("Connection: Upgrade"));
+
+        letBodyFinish.SetResult();
+
+        await middlewareTask;
+    }
+
+    [Fact]
+    public async Task UpgradeToWebSocketLogsWrittenOnlyOnce()
+    {
+        var options = CreateOptionsAccessor();
+        options.CurrentValue.LoggingFields = HttpLoggingFields.ResponsePropertiesAndHeaders;
+
+        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();
+                await c.Response.CompleteAsync();
+            },
+            options,
+            LoggerFactory.CreateLogger<HttpLoggingMiddleware>());
+
+        await middleware.Invoke(httpContext);
+
+        Assert.Equal(1, writeCount);
+    }
+
+    [Fact]
+    public async Task OriginalUpgradeFeatureIsRestoredBeforeMiddlewareCompletes()
+    {
+        var options = CreateOptionsAccessor();
+        options.CurrentValue.LoggingFields = HttpLoggingFields.ResponsePropertiesAndHeaders;
+
+        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();

src/Middleware/HttpLogging/test/UpgradeFeatureLoggingDecoratorTests.cs

@@ -0,0 +1,60 @@
+// 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.Features;
+using Moq;
+
+namespace Microsoft.AspNetCore.HttpLogging.Tests;
+public class UpgradeFeatureLoggingDecoratorTests
+{
+    private readonly Mock<IHttpUpgradeFeature> _upgradeFeatureMock;
+    private readonly Mock<Action> _loggingDelegateMock;
+
+    public UpgradeFeatureLoggingDecoratorTests()
+    {
+        _upgradeFeatureMock = new();
+        _loggingDelegateMock = new();
+    }
+
+    [Fact]
+    public void Ctor_ThrowsExceptionsWhenNullArgs()
+    {
+        Assert.Throws<ArgumentNullException>(() => new UpgradeFeatureLoggingDecorator(
+            null,
+            _loggingDelegateMock.Object));
+
+        Assert.Throws<ArgumentNullException>(() => new UpgradeFeatureLoggingDecorator(
+            _upgradeFeatureMock.Object,
+            null));
+    }
+
+    [Fact]
+    public void IsUpgradableRequest_ShouldForwardToWrappedFeature()
+    {
+        var decorator = new UpgradeFeatureLoggingDecorator(_upgradeFeatureMock.Object, _loggingDelegateMock.Object);
+
+        var isUpgradableRequest = decorator.IsUpgradableRequest;
+
+        Assert.Null(
+            Record.Exception(() => _upgradeFeatureMock.Verify(m => m.IsUpgradableRequest, Times.Once)));
+        Assert.Null(
+            Record.Exception(() => _upgradeFeatureMock.VerifyNoOtherCalls()));
+    }
+
+    [Fact]
+    public async Task UpgradeAsync_ShouldForwardToWrappedFeatureAndThenLog()
+    {
+        var decorator = new UpgradeFeatureLoggingDecorator(_upgradeFeatureMock.Object, _loggingDelegateMock.Object);
+
+        await decorator.UpgradeAsync();
+
+        Assert.Null(
+            Record.Exception(() => _upgradeFeatureMock.Verify(m => m.UpgradeAsync(), Times.Once)));
+        Assert.Null(
+            Record.Exception(() => _upgradeFeatureMock.VerifyNoOtherCalls()));
+        Assert.Null(
+            Record.Exception(() => _loggingDelegateMock.Verify(m => m(), Times.Once)));
+        Assert.Null(
+            Record.Exception(() => _loggingDelegateMock.VerifyNoOtherCalls()));
+    }
+}