Always add adapter features into context

src/Microsoft.AspNetCore.SystemWebAdapters.CoreServices/BufferResponseStreamAttribute.cs

@@ -8,9 +8,11 @@ namespace Microsoft.AspNetCore.SystemWebAdapters;
 [AttributeUsage(AttributeTargets.Method | AttributeTargets.Class)]
 public sealed class BufferResponseStreamAttribute : Attribute
 {
+    internal const int DefaultMemoryThreshold = 32768; // Same default as FileBufferingWriteStream
+
     public bool IsDisabled { get; set; }
 
-    public int MemoryThreshold { get; set; } = 32768; // Same default as FileBufferingWriteStream
+    public int MemoryThreshold { get; set; } = DefaultMemoryThreshold;
 
     public long? BufferLimit { get; set; }
 }

src/Microsoft.AspNetCore.SystemWebAdapters.CoreServices/BufferResponseStreamMiddleware.cs

@@ -3,7 +3,6 @@
 
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
-using Microsoft.AspNetCore.Http.Features;
 using Microsoft.Extensions.Logging;
 
 namespace Microsoft.AspNetCore.SystemWebAdapters;
@@ -23,30 +22,14 @@ public BufferResponseStreamMiddleware(RequestDelegate next, ILogger<BufferRespon
     }
 
     public Task InvokeAsync(HttpContextCore context)
-        => context.GetEndpoint()?.Metadata.GetMetadata<BufferResponseStreamAttribute>() is { IsDisabled: false } metadata
-            ? BufferResponseStreamAsync(context, metadata)
-            : _next(context);
-
-    private async Task BufferResponseStreamAsync(HttpContextCore context, BufferResponseStreamAttribute metadata)
     {
-        LogBuffering(metadata.BufferLimit, metadata.MemoryThreshold);
-
-        var responseBodyFeature = context.Features.GetRequired<IHttpResponseBodyFeature>();
-
-        await using var bufferedFeature = new HttpResponseAdapterFeature(responseBodyFeature, metadata);
-
-        context.Features.Set<IHttpResponseBodyFeature>(bufferedFeature);
-        context.Features.Set<IHttpResponseAdapterFeature>(bufferedFeature);
-
-        try
+        if (context.GetEndpoint()?.Metadata.GetMetadata<BufferResponseStreamAttribute>() is { IsDisabled: false } metadata)
         {
-            await _next(context);
-            await bufferedFeature.FlushBufferedStreamAsync();
-        }
-        finally
-        {
-            context.Features.Set(responseBodyFeature);
-            context.Features.Set<IHttpResponseAdapterFeature>(null);
+            LogBuffering(metadata.BufferLimit, metadata.MemoryThreshold);
+
+            context.Features.GetRequired<IHttpResponseBufferingFeature>().EnableBuffering(metadata.MemoryThreshold, metadata.BufferLimit);
         }
+
+        return _next(context);
     }
 }

src/Microsoft.AspNetCore.SystemWebAdapters.CoreServices/HttpRequestAdapterFeature.cs → src/Microsoft.AspNetCore.SystemWebAdapters.CoreServices/HttpRequestInputStreamFeature.cs

@@ -14,19 +14,16 @@
 
 namespace Microsoft.AspNetCore.SystemWebAdapters;
 
-internal class HttpRequestAdapterFeature : IHttpRequestAdapterFeature, IHttpRequestFeature, IRequestBodyPipeFeature, IDisposable
+internal class HttpRequestInputStreamFeature : IHttpRequestInputStreamFeature, IHttpRequestFeature, IRequestBodyPipeFeature, IDisposable
 {
-    private readonly int _bufferThreshold;
-    private readonly long? _bufferLimit;
     private readonly IHttpRequestFeature _other;
 
     private PipeReader? _pipeReader;
     private Stream? _bufferedStream;
 
-    public HttpRequestAdapterFeature(IHttpRequestFeature other, int bufferThreshold, long? bufferLimit)
+    public HttpRequestInputStreamFeature(IHttpRequestFeature other)
     {
-        _bufferThreshold = bufferThreshold;
-        _bufferLimit = bufferLimit;
+        BufferThreshold = PreBufferRequestStreamAttribute.DefaultBufferThreshold;
         _other = other;
     }
 
@@ -44,13 +41,17 @@ public Stream GetBufferedInputStream()
         {
             Mode = ReadEntityBodyMode.Buffered;
 
-            return _bufferedStream = new FileBufferingReadStream(_other.Body, _bufferThreshold, _bufferLimit, AspNetCoreTempDirectory.TempDirectoryFactory);
+            return _bufferedStream = new FileBufferingReadStream(_other.Body, BufferThreshold, BufferLimit, AspNetCoreTempDirectory.TempDirectoryFactory);
         }
 
         throw new InvalidOperationException("GetBufferlessInputStream cannot be called after other stream access");
     }
 
-    Stream IHttpRequestAdapterFeature.GetBufferlessInputStream()
+    public int BufferThreshold { get; set; }
+
+    public long? BufferLimit { get; set; }
+
+    Stream IHttpRequestInputStreamFeature.GetBufferlessInputStream()
     {
         if (Mode is ReadEntityBodyMode.Bufferless or ReadEntityBodyMode.None)
         {
@@ -61,7 +62,7 @@ Stream IHttpRequestAdapterFeature.GetBufferlessInputStream()
         throw new InvalidOperationException("GetBufferlessInputStream cannot be called after other stream access");
     }
 
-    Stream IHttpRequestAdapterFeature.InputStream
+    Stream IHttpRequestInputStreamFeature.InputStream
     {
         get
         {
@@ -74,13 +75,7 @@ Stream IHttpRequestAdapterFeature.InputStream
         }
     }
 
-    async Task<Stream> IHttpRequestAdapterFeature.GetInputStreamAsync(CancellationToken token)
-    {
-        await BufferInputStreamAsync(token);
-        return GetBody();
-    }
-
-    public async Task BufferInputStreamAsync(CancellationToken token)
+    async ValueTask IHttpRequestInputStreamFeature.BufferInputStreamAsync(CancellationToken token)
     {
         if (Mode is ReadEntityBodyMode.Classic)
         {

src/Microsoft.AspNetCore.SystemWebAdapters.CoreServices/HttpResponseAdapterFeature.cs

@@ -2,6 +2,8 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
 using System.IO;
 using System.IO.Pipelines;
 using System.Threading;
@@ -12,17 +14,7 @@
 
 namespace Microsoft.AspNetCore.SystemWebAdapters;
 
-/// <summary>
-/// This feature implements the <see cref="IHttpResponseAdapterFeature"/> to expose functionality for the adapters. As part of that,
-/// it overrides the following features as well:
-/// 
-/// <list>
-///   <item>
-///     <see cref="IHttpResponseBodyFeature"/>: Provide ability to turn off writing to the stream, while also supporting the ability to clear and suppress output
-///   </item>
-/// </list> 
-/// </summary>
-internal class HttpResponseAdapterFeature : Stream, IHttpResponseBodyFeature, IHttpResponseAdapterFeature
+internal class HttpResponseAdapterFeature : Stream, IHttpResponseBodyFeature, IHttpResponseBufferingFeature, IHttpResponseEndFeature, IHttpResponseContentFeature
 {
     private enum StreamState
     {
@@ -33,17 +25,16 @@ private enum StreamState
     }
 
     private readonly IHttpResponseBodyFeature _responseBodyFeature;
-    private readonly BufferResponseStreamAttribute _metadata;
 
     private FileBufferingWriteStream? _bufferedStream;
     private PipeWriter? _pipeWriter;
+    private StreamState _state;
+    private Func<FileBufferingWriteStream>? _factory;
     private bool _suppressContent;
-    private StreamState _state; 
 
-    public HttpResponseAdapterFeature(IHttpResponseBodyFeature httpResponseBody, BufferResponseStreamAttribute metadata)
+    public HttpResponseAdapterFeature(IHttpResponseBodyFeature httpResponseBody)
     {
         _responseBodyFeature = httpResponseBody;
-        _metadata = metadata;
         _state = StreamState.NotStarted;
     }
 
@@ -59,55 +50,107 @@ void IHttpResponseBodyFeature.DisableBuffering()
         }
     }
 
-    Task IHttpResponseBodyFeature.StartAsync(CancellationToken cancellationToken)
+    void IHttpResponseBufferingFeature.EnableBuffering(int memoryThreshold, long? bufferLimit)
     {
         if (_state == StreamState.NotStarted)
         {
+            Debug.Assert(_bufferedStream is null);
+
             _state = StreamState.Buffering;
+            _factory = () => new FileBufferingWriteStream(memoryThreshold, bufferLimit);
+        }
+        else
+        {
+            throw new InvalidOperationException("Cannot enable buffering if writing has begun");
+        }
+    }
+
+    Task IHttpResponseBodyFeature.StartAsync(CancellationToken cancellationToken)
+    {
+        if (_state == StreamState.NotStarted)
+        {
+            _state = StreamState.NotBuffering;
         }
 
         return _responseBodyFeature.StartAsync(cancellationToken);
     }
 
+    private async ValueTask FlushInternalAsync()
+    {
+        if (_state is StreamState.Buffering && _bufferedStream is not null && !SuppressContent)
+        {
+            await _bufferedStream.DrainBufferAsync(_responseBodyFeature.Stream);
+            await _bufferedStream.DisposeAsync();
+            _bufferedStream = null;
+        }
+    }
+
     Stream IHttpResponseBodyFeature.Stream => this;
 
     PipeWriter IHttpResponseBodyFeature.Writer => _pipeWriter ??= PipeWriter.Create(this, new StreamPipeWriterOptions(leaveOpen: true));
 
-    bool IHttpResponseAdapterFeature.SuppressContent
+    public bool SuppressContent
     {
         get => _suppressContent;
-        set => _suppressContent = value;
+        set
+        {
+            if (value)
+            {
+                VerifyBuffering();
+            }
+
+            _suppressContent = value;
+        }
     }
 
-    Task IHttpResponseAdapterFeature.EndAsync() => CompleteAsync();
+    Task IHttpResponseEndFeature.EndAsync() => CompleteAsync();
 
-    bool IHttpResponseAdapterFeature.IsEnded => _state == StreamState.Complete;
+    bool IHttpResponseEndFeature.IsEnded => _state == StreamState.Complete;
 
-    void IHttpResponseAdapterFeature.ClearContent()
+    void IHttpResponseContentFeature.ClearContent()
     {
-        if (_bufferedStream is not null)
+        if (CurrentStream is { CanSeek: true } body)
         {
-            _bufferedStream.Dispose();
-            _bufferedStream = null;
+            body.SetLength(0);
+            return;
         }
+
+        VerifyBuffering();
+
+        _bufferedStream?.Dispose();
+        _bufferedStream = null;
+    }
+
+    [MemberNotNull(nameof(_factory))]
+    private void VerifyBuffering()
+    {
+        if (_state != StreamState.Buffering)
+        {
+            throw new InvalidOperationException("Can only clear content if response is buffered.");
+        }
+
+        Debug.Assert(_factory is not null);
     }
 
+    ValueTask IHttpResponseBufferingFeature.FlushAsync() => FlushInternalAsync();
+
     private Stream CurrentStream
     {
         get
         {
-            if (_state == StreamState.NotBuffering)
-            {
-                return _responseBodyFeature.Stream;
-            }
-            else if (_state == StreamState.Complete)
+            if (_state == StreamState.Buffering)
             {
-                return Null;
+                VerifyBuffering();
+                return _bufferedStream ??= _factory();
             }
             else
             {
-                _state = StreamState.Buffering;
-                return _bufferedStream ??= new FileBufferingWriteStream(_metadata.MemoryThreshold, _metadata.BufferLimit);
+                if (_state != StreamState.Complete)
+                {
+                    _state = StreamState.NotBuffering;
+                }
+
+                return _responseBodyFeature.Stream;
             }
         }
     }
@@ -122,15 +165,7 @@ public override async ValueTask DisposeAsync()
         await base.DisposeAsync();
     }
 
-    public async ValueTask FlushBufferedStreamAsync()
-    {
-        if (_state is StreamState.Buffering && _bufferedStream is not null && !_suppressContent)
-        {
-            await _bufferedStream.DrainBufferAsync(_responseBodyFeature.Stream);
-        }
-    }
-
-    public override bool CanRead => true;
+    public override bool CanRead => false;
 
     public override bool CanSeek => false;
 
@@ -144,12 +179,18 @@ public override long Position
         set => throw new NotSupportedException();
     }
 
-
     private async Task CompleteAsync()
     {
-        await FlushBufferedStreamAsync();
-        await _responseBodyFeature.CompleteAsync();
+        if (_state == StreamState.Complete)
+        {
+            return;
+        }
+
+        await FlushInternalAsync();
+
         _state = StreamState.Complete;
+
+        await _responseBodyFeature.CompleteAsync();
     }
 
     public override void Flush() => CurrentStream.Flush();

src/Microsoft.AspNetCore.SystemWebAdapters.CoreServices/Mvc/ResponseEndFilter.cs

@@ -20,7 +20,7 @@ public ResponseEndFilter(ILogger<ResponseEndFilter> logger)
 
     public void OnActionExecuted(ActionExecutedContext context)
     {
-        if (context.Result is not null && context.HttpContext.Features.Get<IHttpResponseAdapterFeature>() is { IsEnded: true })
+        if (context.Result is not null && context.HttpContext.Features.Get<IHttpResponseEndFeature>() is { IsEnded: true })
         {
             LogClearingResult();
             context.Result = null;

src/Microsoft.AspNetCore.SystemWebAdapters.CoreServices/PreBufferRequestStreamAttribute.cs

@@ -9,7 +9,7 @@ namespace Microsoft.AspNetCore.SystemWebAdapters;
 public sealed class PreBufferRequestStreamAttribute : Attribute
 {
     // Same limit as the default: https://source.dot.net/#Microsoft.AspNetCore.Http/Internal/BufferingHelper.cs,47b7015acb14f2a4
-    private const int DefaultBufferThreshold = 1024 * 30;
+    internal const int DefaultBufferThreshold = 1024 * 30;
 
     public bool IsDisabled { get; set; }