fix: preserve outgoing response streams during body capture

Author: georgidhristov

DebugProbe.AspNetCore.Tests/Handlers/DebugProbeHttpClientHandlerBodyTests.cs

@@ -1,4 +1,5 @@
 using System.Net;
+using System.Net.Http.Headers;
 using System.Text;
 using DebugProbe.AspNetCore.Handlers;
 using DebugProbe.AspNetCore.Models;
@@ -94,6 +95,45 @@ public async Task Body_over_limit_is_truncated_and_marker_is_appended()
         Assert.StartsWith(new string('C', 1024), outgoing.ResponseBody);
     }
 
+    [Fact]
+    public async Task Capturing_streamed_response_body_does_not_consume_body_for_caller()
+    {
+        const int limitKb = 1;
+        var bodyText = new string('D', 2048);
+        var responseContent = new StreamContent(new NonSeekableMemoryStream(Encoding.UTF8.GetBytes(bodyText)));
+        responseContent.Headers.ContentType = new MediaTypeHeaderValue("text/plain")
+        {
+            CharSet = Encoding.UTF8.WebName
+        };
+
+        var options = new DebugProbeOptions { MaxBodyCaptureSizeKb = limitKb };
+        var (client, entry) = BuildClient(options, responseContent);
+
+        using var response = await client.GetAsync("https://example.test/");
+
+        Assert.Equal(bodyText, await response.Content.ReadAsStringAsync());
+
+        var outgoing = Assert.Single(entry.OutgoingRequests);
+        Assert.EndsWith("[truncated]", outgoing.ResponseBody);
+    }
+
+    [Fact]
+    public async Task Captured_body_uses_content_type_charset()
+    {
+        const string bodyText = "Hello \u0100";
+        var responseContent = new StringContent(bodyText, Encoding.Unicode, "text/plain");
+
+        var options = new DebugProbeOptions { MaxBodyCaptureSizeKb = 1 };
+        var (client, entry) = BuildClient(options, responseContent);
+
+        using var response = await client.GetAsync("https://example.test/");
+
+        Assert.Equal(bodyText, await response.Content.ReadAsStringAsync());
+
+        var outgoing = Assert.Single(entry.OutgoingRequests);
+        Assert.Contains(bodyText, outgoing.ResponseBody);
+    }
+
     [Fact]
     public async Task Null_content_returns_empty_string()
     {
@@ -132,4 +172,9 @@ protected override Task<HttpResponseMessage> SendAsync(
             HttpRequestMessage request, CancellationToken cancellationToken)
             => Task.FromResult(send(request));
     }
+
+    private sealed class NonSeekableMemoryStream(byte[] buffer) : MemoryStream(buffer)
+    {
+        public override bool CanSeek => false;
+    }
 }

DebugProbe.AspNetCore/Handlers/DebugProbeHttpClientHandler.cs

@@ -1,4 +1,6 @@
 using System.Diagnostics;
+using System.Net;
+using System.Text;
 using DebugProbe.AspNetCore.Internal.Utils;
 using DebugProbe.AspNetCore.Models;
 using DebugProbe.AspNetCore.Options;
@@ -33,13 +35,13 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
         {
             var response = await base.SendAsync(request, cancellationToken);
 
-            await CaptureRequest(request, response, null, started.ElapsedMilliseconds);
+            await CaptureRequest(request, response, null, started.ElapsedMilliseconds, cancellationToken);
 
             return response;
         }
         catch (Exception ex)
         {
-            await CaptureRequest(request, null, ex, started.ElapsedMilliseconds);
+            await CaptureRequest(request, null, ex, started.ElapsedMilliseconds, cancellationToken);
 
             throw;
         }
@@ -48,7 +50,12 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
     /// <summary>
     /// Captures outgoing request details and stores them in the active DebugProbe entry.
     /// </summary>
-    private async Task CaptureRequest(HttpRequestMessage request, HttpResponseMessage? response, Exception? exception, long durationMs)
+    private async Task CaptureRequest(
+        HttpRequestMessage request,
+        HttpResponseMessage? response,
+        Exception? exception,
+        long durationMs,
+        CancellationToken cancellationToken)
     {
         if (!TryGetActiveEntry(out var entry))
         {
@@ -76,9 +83,9 @@ private async Task CaptureRequest(HttpRequestMessage request, HttpResponseMessag
             ResponseHeaders = response != null ? response.Headers.ToDictionary(x => x.Key, x => RedactionUtils.RedactHeader(x.Key, string.Join(", ", x.Value), _options)) : []
         };
 
-        outgoing.RequestBody = await CaptureBodyAsync(request.Content);
+        outgoing.RequestBody = (await CaptureBodyAsync(request.Content, cancellationToken)).Body;
 
-        outgoing.ResponseBody = await CaptureBodyAsync(response?.Content);
+        outgoing.ResponseBody = await CaptureResponseBodyAsync(response, cancellationToken);
 
         entry.OutgoingRequests.Add(outgoing);
     }
@@ -100,7 +107,25 @@ private bool TryGetActiveEntry(out DebugEntry entry)
         return true;
     }
 
-    private async Task<string> CaptureBodyAsync(HttpContent? content)
+    private async Task<string> CaptureResponseBodyAsync(HttpResponseMessage? response, CancellationToken cancellationToken)
+    {
+        var content = response?.Content;
+        if (content == null)
+        {
+            return string.Empty;
+        }
+
+        var result = await CaptureBodyAsync(content, cancellationToken);
+
+        if (result.BytesRead.Length > 0)
+        {
+            response!.Content = new PrefixReplayHttpContent(content, result.Stream, result.BytesRead);
+        }
+
+        return result.Body;
+    }
+
+    private async Task<BodyCaptureResult> CaptureBodyAsync(HttpContent? content, CancellationToken cancellationToken)
     {
         if (content == null ||
             !HttpContentUtils.IsTextContent(content.Headers.ContentType?.MediaType))
@@ -115,24 +140,100 @@ private async Task<string> CaptureBodyAsync(HttpContent? content)
         var buffer = new byte[limit + 1];
         var totalRead = 0;
 
-        var stream = await content.ReadAsStreamAsync();
+        var stream = await content.ReadAsStreamAsync(cancellationToken);
 
         int bytesRead;
         while (totalRead < buffer.Length &&
-               (bytesRead = await stream.ReadAsync(buffer.AsMemory(totalRead, buffer.Length - totalRead))) > 0)
+               (bytesRead = await stream.ReadAsync(buffer.AsMemory(totalRead, buffer.Length - totalRead), cancellationToken)) > 0)
         {
             totalRead += bytesRead;
         }
 
         var truncated = totalRead > limit;
-        var encoding = System.Text.Encoding.UTF8;
+        var encoding = GetEncoding(content);
         var body = encoding.GetString(buffer, 0, Math.Min(totalRead, limit));
 
         if (truncated)
         {
             body += "[truncated]";
         }
 
-        return JsonUtils.Format(RedactionUtils.RedactJsonFields(body, _options));
+        return new BodyCaptureResult(
+            JsonUtils.Format(RedactionUtils.RedactJsonFields(body, _options)),
+            stream,
+            buffer[..totalRead]);
+    }
+
+    private static Encoding GetEncoding(HttpContent content)
+    {
+        var charset = content.Headers.ContentType?.CharSet;
+        if (string.IsNullOrWhiteSpace(charset))
+        {
+            return Encoding.UTF8;
+        }
+
+        try
+        {
+            return Encoding.GetEncoding(charset.Trim('"'));
+        }
+        catch
+        {
+            return Encoding.UTF8;
+        }
+    }
+
+    private readonly record struct BodyCaptureResult(string Body, Stream Stream, byte[] BytesRead)
+    {
+        public static implicit operator BodyCaptureResult(string body) => new(body, Stream.Null, []);
+    }
+
+    private sealed class PrefixReplayHttpContent : HttpContent
+    {
+        private readonly HttpContent _innerContent;
+
+        private readonly Stream _remainingStream;
+
+        private readonly byte[] _prefix;
+
+        public PrefixReplayHttpContent(HttpContent innerContent, Stream remainingStream, byte[] prefix)
+        {
+            _innerContent = innerContent;
+            _remainingStream = remainingStream;
+            _prefix = prefix;
+
+            foreach (var header in innerContent.Headers)
+            {
+                Headers.TryAddWithoutValidation(header.Key, header.Value);
+            }
+        }
+
+        protected override async Task SerializeToStreamAsync(Stream stream, TransportContext? context)
+        {
+            await stream.WriteAsync(_prefix);
+            await _remainingStream.CopyToAsync(stream);
+        }
+
+        protected override bool TryComputeLength(out long length)
+        {
+            if (_innerContent.Headers.ContentLength is { } contentLength)
+            {
+                length = contentLength;
+                return true;
+            }
+
+            length = 0;
+            return false;
+        }
+
+        protected override void Dispose(bool disposing)
+        {
+            if (disposing)
+            {
+                _remainingStream.Dispose();
+                _innerContent.Dispose();
+            }
+
+            base.Dispose(disposing);
+        }
     }
 }