Fix HttpLoggingMiddleware Request/Response bodies logging in case of stream being closed by a subsequent middleware (#61489)

Author: ExtraClock

src/Middleware/HttpLogging/src/RequestBufferingStream.cs

@@ -13,6 +13,7 @@ internal sealed class RequestBufferingStream : BufferingStream
     private readonly bool _logOnFinish;
     private readonly int _limit;
     private BodyStatus _status = BodyStatus.None;
+    private string? _bodyBeforeClose;
 
     public bool HasLogged { get; private set; }
 
@@ -116,15 +117,15 @@ public void LogRequestBody()
         if (!HasLogged && _logOnFinish)
         {
             HasLogged = true;
-            _logger.RequestBody(GetString(_encoding), GetStatus(showCompleted: false));
+            _logger.RequestBody(GetStringInternal(), GetStatus(showCompleted: false));
         }
     }
 
     public void LogRequestBody(HttpLoggingInterceptorContext logContext)
     {
         if (logContext.IsAnyEnabled(HttpLoggingFields.RequestBody))
         {
-            logContext.AddParameter("RequestBody", GetString(_encoding));
+            logContext.AddParameter("RequestBody", GetStringInternal());
             logContext.AddParameter("RequestBodyStatus", GetStatus(showCompleted: true));
         }
     }
@@ -138,6 +139,25 @@ public void LogRequestBody(HttpLoggingInterceptorContext logContext)
         _ => throw new NotImplementedException(_status.ToString()),
     };
 
+    private string GetStringInternal()
+    {
+        var result = _bodyBeforeClose ?? GetString(_encoding);
+        // Reset the value after its consumption to preserve GetString(encoding) behavior
+        _bodyBeforeClose = null;
+        return result;
+    }
+
+    public override void Close()
+    {
+        if (!HasLogged)
+        {
+            // Subsequent middleware can close the request stream after reading enough bytes (guided by ContentLength).
+            // Preserving the body for the final GetStringInternal() call.
+            _bodyBeforeClose = GetString(_encoding);
+        }
+        base.Close();
+    }
+
     public override IAsyncResult BeginRead(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state)
     {
         return TaskToApm.Begin(ReadAsync(buffer, offset, count), callback, state);

src/Middleware/HttpLogging/src/ResponseBufferingStream.cs

@@ -20,6 +20,7 @@ internal sealed class ResponseBufferingStream : BufferingStream, IHttpResponseBo
     private readonly IHttpLoggingInterceptor[] _interceptors;
     private bool _logBody;
     private Encoding? _encoding;
+    private string? _bodyBeforeClose;
 
     private static readonly StreamPipeWriterOptions _pipeWriterOptions = new StreamPipeWriterOptions(leaveOpen: true);
 
@@ -179,7 +180,7 @@ public void LogResponseBody()
     {
         if (_logBody)
         {
-            var responseBody = GetString(_encoding!);
+            var responseBody = GetStringInternal();
             _logger.ResponseBody(responseBody);
         }
     }
@@ -188,7 +189,27 @@ public void LogResponseBody(HttpLoggingInterceptorContext logContext)
     {
         if (_logBody)
         {
-            logContext.AddParameter("ResponseBody", GetString(_encoding!));
+            logContext.AddParameter("ResponseBody", GetStringInternal());
         }
     }
+
+    private string GetStringInternal()
+    {
+        var result = _bodyBeforeClose ?? GetString(_encoding!);
+        // Reset the value after its consumption to preserve GetString(encoding) behavior
+        _bodyBeforeClose = null;
+        return result;
+    }
+
+    public override void Close()
+    {
+        if (_logBody)
+        {
+            // Subsequent middleware can close the response stream after writing its body
+            // Preserving the body for the final GetStringInternal() call.
+            _bodyBeforeClose = GetString(_encoding!);
+        }
+
+        base.Close();
+    }
 }

src/Middleware/HttpLogging/test/HttpLoggingMiddlewareTests.cs

@@ -390,6 +390,46 @@ public async Task RequestBodyCopyToAsyncWorks(string expected)
         Assert.Contains(TestSink.Writes, w => w.Message.Contains(expected));
     }
 
+    [Theory]
+    [MemberData(nameof(BodyData))]
+    public async Task RequestBodyWithStreamCloseWorks(string expected)
+    {
+        var options = CreateOptionsAccessor();
+        options.CurrentValue.LoggingFields = HttpLoggingFields.RequestBody;
+
+        var middleware = CreateMiddleware(
+            async c =>
+            {
+                var arr = new byte[4096];
+                var contentLengthBytesLeft = c.Request.Body.Length;
+
+                // (1) The subsequent middleware reads right up to the buffer size (guided by the ContentLength header)
+                while (contentLengthBytesLeft > 0)
+                {
+                    var res = await c.Request.Body.ReadAsync(arr, 0, (int)Math.Min(arr.Length, contentLengthBytesLeft));
+                    contentLengthBytesLeft -= res;
+                    if (res == 0)
+                    {
+                        break;
+                    }
+                }
+
+                // (2) The subsequent middleware closes the request stream after its consumption
+                c.Request.Body.Close();
+            },
+            options);
+
+        var httpContext = new DefaultHttpContext();
+        httpContext.Request.ContentType = "text/plain";
+        var buffer = Encoding.UTF8.GetBytes(expected);
+        httpContext.Request.Body = new MemoryStream(buffer);
+        httpContext.Request.ContentLength = buffer.Length;
+
+        await middleware.Invoke(httpContext);
+
+        Assert.Contains(TestSink.Writes, w => w.Message.Contains(expected));
+    }
+
     [Fact]
     public async Task RequestBodyReadingLimitLongCharactersWorks()
     {
@@ -1155,6 +1195,32 @@ public async Task StartAsyncResponseHeadersLogged()
         await middlewareTask;
     }
 
+    [Theory]
+    [MemberData(nameof(BodyData))]
+    public async Task ResponseBodyWithStreamCloseWorks(string expected)
+    {
+        var options = CreateOptionsAccessor();
+        options.CurrentValue.LoggingFields = HttpLoggingFields.ResponseBody;
+        var middleware = CreateMiddleware(
+            async c =>
+            {
+                c.Response.ContentType = "text/plain";
+
+                // (1) The subsequent middleware writes its response
+                await c.Response.WriteAsync(expected);
+
+                // (2) The subsequent middleware closes the response stream after it has completed writing to it
+                c.Response.Body.Close();
+            },
+            options);
+
+        var httpContext = new DefaultHttpContext();
+
+        await middleware.Invoke(httpContext);
+
+        Assert.Contains(TestSink.Writes, w => w.Message.Contains(expected));
+    }
+
     [Fact]
     public async Task UnrecognizedMediaType()
     {