fix: outgoing body capture stream preservation

georgidhristov · web-flow · commit 368b5f699f7b · 2026-06-26T07:24:58.000+03:00
diff --git a/DebugProbe.AspNetCore.Tests/Handlers/DebugProbeHttpClientHandlerBodyTests.cs b/DebugProbe.AspNetCore.Tests/Handlers/DebugProbeHttpClientHandlerBodyTests.cs
@@ -0,0 +1,180 @@
+using System.Net;
+using System.Net.Http.Headers;
+using System.Text;
+using DebugProbe.AspNetCore.Handlers;
+using DebugProbe.AspNetCore.Models;
+using DebugProbe.AspNetCore.Options;
+
+namespace DebugProbe.AspNetCore.Tests.Handlers;
+
+public class DebugProbeHttpClientHandlerBodyTests
+{
+    // ---------------------------------------------------------------------------
+    // Helpers
+    // ---------------------------------------------------------------------------
+
+    /// <summary>
+    /// Builds a handler wired to the given DebugEntry and options, with a stub
+    /// inner handler that echoes <paramref name="responseContent"/> back.
+    /// </summary>
+    private static (HttpClient client, DebugEntry entry) BuildClient(
+        DebugProbeOptions options,
+        HttpContent? responseContent = null)
+    {
+        var entry = new DebugEntry();
+        var context = new DefaultHttpContext();
+        context.Items["DebugProbeEntry"] = entry;
+
+        var handler = new DebugProbeHttpClientHandler(
+            new HttpContextAccessor { HttpContext = context },
+            options)
+        {
+            InnerHandler = new StubHandler(_ =>
+                new HttpResponseMessage(HttpStatusCode.OK)
+                {
+                    Content = responseContent
+                })
+        };
+
+        return (new HttpClient(handler), entry);
+    }
+
+    // ---------------------------------------------------------------------------
+    // Test cases
+    // ---------------------------------------------------------------------------
+
+    [Fact]
+    public async Task Body_under_limit_is_captured_without_truncation_marker()
+    {
+        const int limitKb = 1;                    // 1 KB limit
+        var bodyText = new string('A', 100);      // 100 bytes — well under 1 KB
+
+        var options = new DebugProbeOptions { MaxBodyCaptureSizeKb = limitKb };
+        var (client, entry) = BuildClient(options,
+            responseContent: new StringContent(bodyText, Encoding.UTF8, "text/plain"));
+
+        await client.GetAsync("https://example.test/");
+
+        var outgoing = Assert.Single(entry.OutgoingRequests);
+        Assert.DoesNotContain("[truncated]", outgoing.ResponseBody);
+        Assert.Contains(bodyText, outgoing.ResponseBody);
+    }
+
+    [Fact]
+    public async Task Body_exactly_at_limit_is_captured_without_truncation_marker()
+    {
+        const int limitKb = 1;
+        var bodyText = new string('B', 1024);     // exactly 1 KB
+
+        var options = new DebugProbeOptions { MaxBodyCaptureSizeKb = limitKb };
+        var (client, entry) = BuildClient(options,
+            responseContent: new StringContent(bodyText, Encoding.UTF8, "text/plain"));
+
+        await client.GetAsync("https://example.test/");
+
+        var outgoing = Assert.Single(entry.OutgoingRequests);
+        Assert.DoesNotContain("[truncated]", outgoing.ResponseBody);
+        Assert.Contains(bodyText, outgoing.ResponseBody);
+    }
+
+    [Fact]
+    public async Task Body_over_limit_is_truncated_and_marker_is_appended()
+    {
+        const int limitKb = 1;
+        var bodyText = new string('C', 2048);     // 2 KB — double the limit
+
+        var options = new DebugProbeOptions { MaxBodyCaptureSizeKb = limitKb };
+        var (client, entry) = BuildClient(options,
+            responseContent: new StringContent(bodyText, Encoding.UTF8, "text/plain"));
+
+        await client.GetAsync("https://example.test/");
+
+        var outgoing = Assert.Single(entry.OutgoingRequests);
+        Assert.EndsWith("[truncated]", outgoing.ResponseBody);
+        // Captured prefix must be exactly the limit (1024 'C' chars)
+        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()
+    {
+        var options = new DebugProbeOptions();
+        var (client, entry) = BuildClient(options, responseContent: null);
+
+        await client.GetAsync("https://example.test/");
+
+        var outgoing = Assert.Single(entry.OutgoingRequests);
+        Assert.Equal(string.Empty, outgoing.ResponseBody);
+    }
+
+    [Fact]
+    public async Task Non_text_content_returns_empty_string()
+    {
+        var binaryContent = new ByteArrayContent([0x89, 0x50, 0x4E, 0x47]); // PNG header bytes
+        binaryContent.Headers.ContentType =
+            new System.Net.Http.Headers.MediaTypeHeaderValue("image/png");
+
+        var options = new DebugProbeOptions();
+        var (client, entry) = BuildClient(options, responseContent: binaryContent);
+
+        await client.GetAsync("https://example.test/");
+
+        var outgoing = Assert.Single(entry.OutgoingRequests);
+        Assert.Equal(string.Empty, outgoing.ResponseBody);
+    }
+
+    // ---------------------------------------------------------------------------
+    // Stub inner handler
+    // ---------------------------------------------------------------------------
+
+    private sealed class StubHandler(Func<HttpRequestMessage, HttpResponseMessage> send) : HttpMessageHandler
+    {
+        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;
+    }
+}
diff --git a/DebugProbe.AspNetCore/Handlers/DebugProbeHttpClientHandler.cs b/DebugProbe.AspNetCore/Handlers/DebugProbeHttpClientHandler.cs
@@ -1,4 +1,6 @@
-﻿using System.Diagnostics;
+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,18 +107,133 @@ 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))
         {
             return string.Empty;
         }
 
-        var body = await content.ReadAsStringAsync();
+        if (_options.MaxBodyCaptureSizeBytes <= 0)
+            return string.Empty;
+
+        var limit = _options.MaxBodyCaptureSizeBytes;
+        var buffer = new byte[limit + 1];
+        var totalRead = 0;
+
+        var stream = await content.ReadAsStreamAsync(cancellationToken);
+
+        int bytesRead;
+        while (totalRead < buffer.Length &&
+               (bytesRead = await stream.ReadAsync(buffer.AsMemory(totalRead, buffer.Length - totalRead), cancellationToken)) > 0)
+        {
+            totalRead += bytesRead;
+        }
+
+        var truncated = totalRead > limit;
+        var encoding = GetEncoding(content);
+        var body = encoding.GetString(buffer, 0, Math.Min(totalRead, limit));
+
+        if (truncated)
+        {
+            body += "[truncated]";
+        }
+
+        return new BodyCaptureResult(
+            JsonUtils.Format(RedactionUtils.RedactJsonFields(body, _options)),
+            stream,
+            buffer[..totalRead]);
+    }
 
-        return JsonUtils.Format(RedactionUtils.RedactJsonFields(
-            HttpContentUtils.Trim(body, _options.MaxBodyCaptureSizeBytes),
-            _options));
+    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);
+        }
     }
 }
