Ensure PreSendHeader event is raised when response is started (#532)

Previously, we were just raining it when a flush occurred, which would be after headers are made readonly.

Fixes #531

Author: twsouthwick

src/Microsoft.AspNetCore.SystemWebAdapters.CoreServices/Features/HttpApplicationPreSendEventsResponseBodyFeature.cs

@@ -7,6 +7,7 @@
 using System.IO.Pipelines;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Web;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
 
@@ -50,7 +51,15 @@ public HttpApplicationPreSendEventsResponseBodyFeature(HttpContextCore context,
     public Task SendFileAsync(string path, long offset, long? count, CancellationToken cancellationToken = default)
         => SendFileFallback.SendFileAsync(Stream, path, offset, count, cancellationToken);
 
-    public Task StartAsync(CancellationToken cancellationToken = default) => _other.StartAsync(cancellationToken);
+    public async Task StartAsync(CancellationToken cancellationToken = default)
+    {
+        if (_context.Features.Get<IHttpApplicationFeature>() is { } httpApplication)
+        {
+            await RaisePreSendRequestHeadersEvent(httpApplication);
+        }
+
+        await _other.StartAsync(cancellationToken);
+    }
 
     public override void Advance(int bytes)
     {
@@ -69,29 +78,43 @@ public override void Advance(int bytes)
 
     public override async ValueTask<FlushResult> FlushAsync(CancellationToken cancellationToken = default)
     {
-        // Only need to raise events if data will be flushed and the feature is available
-        if (_byteCount > 0 && _context.Features.Get<IHttpApplicationFeature>() is { } httpApplication)
+        // Only need to raise events if data will be flushed 
+        if (_byteCount == 0)
+        {
+            return default;
+        }
+
+        if (_context.Features.Get<IHttpApplicationFeature>() is { } httpApplication)
         {
             _byteCount = 0;
 
-            if (_state is State.NotStarted)
-            {
-                _state = State.RaisingPreHeader;
-                await _context.Features.GetRequiredFeature<IHttpApplicationFeature>().RaiseEventAsync(ApplicationEvent.PreSendRequestHeaders);
-                _state = State.ReadyForContent;
-            }
-
-            if (_state is State.ReadyForContent)
-            {
-                _state = State.RaisingPreContent;
-                await httpApplication.RaiseEventAsync(ApplicationEvent.PreSendRequestContent);
-                _state = State.ReadyForContent;
-            }
+            await RaisePreSendRequestHeadersEvent(httpApplication);
+            await RaisePreSendRequestContentEvent(httpApplication);
         }
 
         return await _pipe.FlushAsync(cancellationToken);
     }
 
+    private async ValueTask RaisePreSendRequestHeadersEvent(IHttpApplicationFeature httpApplication)
+    {
+        if (_state is State.NotStarted)
+        {
+            _state = State.RaisingPreHeader;
+            await httpApplication.RaiseEventAsync(ApplicationEvent.PreSendRequestHeaders);
+            _state = State.ReadyForContent;
+        }
+    }
+
+    private async ValueTask RaisePreSendRequestContentEvent(IHttpApplicationFeature httpApplication)
+    {
+        if (_state is State.ReadyForContent)
+        {
+            _state = State.RaisingPreContent;
+            await httpApplication.RaiseEventAsync(ApplicationEvent.PreSendRequestContent);
+            _state = State.ReadyForContent;
+        }
+    }
+
     public override Memory<byte> GetMemory(int sizeHint = 0) => _pipe.GetMemory(sizeHint);
 
     public override Span<byte> GetSpan(int sizeHint = 0) => _pipe.GetSpan(sizeHint);

test/Microsoft.AspNetCore.SystemWebAdapters.CoreServices.Tests/Modules/BufferedModuleTests.cs

@@ -1,12 +1,8 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
 namespace Microsoft.AspNetCore.SystemWebAdapters.CoreServices.Tests;
 
-public class BufferedModuleTests : ModuleTests
+public class BufferedModuleTests : ModuleTests<BufferedModuleTests>
 {
-    public BufferedModuleTests()
-        : base(true)
-    {
-    }
 }

test/Microsoft.AspNetCore.SystemWebAdapters.CoreServices.Tests/Modules/ModuleTests.cs

@@ -9,6 +9,7 @@
 using System.Web;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Hosting;
+using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.SystemWebAdapters.Features;
 using Microsoft.AspNetCore.TestHost;
@@ -20,7 +21,7 @@
 namespace Microsoft.AspNetCore.SystemWebAdapters.CoreServices.Tests;
 
 [Collection(nameof(SelfHostedTests))]
-public abstract class ModuleTests(bool isBuffered)
+public abstract class ModuleTests<T>
 {
     private static readonly ImmutableArray<ApplicationEvent> BeforeHandlerEvents =
     [
@@ -54,7 +55,27 @@ public abstract class ModuleTests(bool isBuffered)
         ApplicationEvent.EndRequest,
     ];
 
-    public static IEnumerable<object[]> GetAllEvents()
+    private static bool IsBuffered
+    {
+        get
+        {
+            if (typeof(T) == typeof(NotBufferedModuleTests))
+            {
+                return false;
+            }
+            else if (typeof(T) == typeof(BufferedModuleTests))
+            {
+                return true;
+            }
+            else
+            {
+                throw new ArgumentOutOfRangeException(typeof(T).FullName);
+            }
+        }
+    }
+
+    [System.Diagnostics.CodeAnalysis.SuppressMessage("Design", "CA1000:Do not declare static members on generic types", Justification = "Needed for xUnit theory tests")]
+    public static TheoryData<ApplicationEvent, RegisterMode> GetAllEvents()
     {
         IEnumerable<ApplicationEvent> all =
         [
@@ -66,14 +87,17 @@ public static IEnumerable<object[]> GetAllEvents()
         ];
 
         var modes = Enum.GetValues<RegisterMode>();
+        var data = new TheoryData<ApplicationEvent, RegisterMode>();
 
         foreach (var notification in all)
         {
             foreach (var mode in modes)
             {
-                yield return new object[] { notification, mode };
+                data.Add(notification, mode);
             }
         }
+
+        return data;
     }
 
     [MemberData(nameof(GetAllEvents))]
@@ -105,7 +129,7 @@ public async Task ModulesThrow(ApplicationEvent notification, RegisterMode mode)
 
         Assert.Equal(expected, result);
 
-        IEnumerable<ApplicationEvent> GetExpected(ApplicationEvent notification)
+        static IEnumerable<ApplicationEvent> GetExpected(ApplicationEvent notification)
         {
             foreach (var item in GetNotificationsUpTo(notification, isThrowing: true))
             {
@@ -119,9 +143,77 @@ IEnumerable<ApplicationEvent> GetExpected(ApplicationEvent notification)
         }
     }
 
-    private IEnumerable<ApplicationEvent> GetNotificationsUpTo(ApplicationEvent notification, bool isThrowing = false)
+    [Fact]
+    public async Task PreSendRequestHeadersAddHeaders()
     {
-        return isBuffered
+        // Arrange
+        const string HeaderName = "name";
+        const string HeaderValue = "value";
+        const string Result = "Hello world!";
+
+        using var host = await new HostBuilder()
+            .ConfigureWebHost(webBuilder => webBuilder
+                .UseTestServer()
+                .ConfigureServices(services =>
+                {
+                    services.AddSystemWebAdapters()
+                        .AddHttpApplication(options =>
+                        {
+                            options.RegisterModule<BufferToggleModule>();
+                            options.RegisterModule<PreSendHeadersAddHeaderModule>();
+                            options.ArePreSendEventsEnabled = true;
+                        });
+
+                })
+                .Configure(app =>
+                {
+                    app.Use((ctx, next) =>
+                    {
+                        ctx.Features.Set(new ResultHeader { { HeaderName, HeaderValue } });
+                        return next(ctx);
+                    });
+                    app.UseSystemWebAdapters();
+
+                    app.Run(ctx => ctx.Response.WriteAsync(Result));
+                })).StartAsync();
+
+        // Act
+        using var response = await host.GetTestClient().GetAsync(new Uri("/", UriKind.Relative));
+
+        // Assert
+        Assert.True(response.Headers.TryGetValues(HeaderName, out var resultHeader));
+        Assert.Equal([HeaderValue], resultHeader);
+        Assert.Equal(Result, await response.Content.ReadAsStringAsync());
+    }
+
+    private sealed class PreSendHeadersAddHeaderModule : IHttpModule
+    {
+        public void Dispose()
+        {
+        }
+
+        public void Init(HttpApplication application)
+        {
+            application.PreSendRequestHeaders += (s, o) =>
+            {
+                if (s is HttpApplication { Context: { } context })
+                {
+                    foreach (var (name, value) in context.AsAspNetCore().Features.GetRequiredFeature<ResultHeader>())
+                    {
+                        context.Response.AddHeader(name, value);
+                    }
+                }
+            };
+        }
+    }
+
+    private sealed class ResultHeader : Dictionary<string, string>
+    {
+    }
+
+    private static IEnumerable<ApplicationEvent> GetNotificationsUpTo(ApplicationEvent notification, bool isThrowing = false)
+    {
+        return IsBuffered
                 ? GetExpectedBufferedNotificationsUntilAction(notification, isThrowing)
                 : GetExpectedUnbufferedNotificationsUntilAction(notification, isThrowing);
 
@@ -206,10 +298,9 @@ static IEnumerable<ApplicationEvent> GetExpectedUnbufferedNotificationsUntilActi
         }
     }
 
-    private async Task<List<ApplicationEvent>> RunAsync(string action, ApplicationEvent @event, RegisterMode mode)
+    private static async Task<List<ApplicationEvent>> RunAsync(string action, ApplicationEvent @event, RegisterMode mode)
     {
         var notifier = new NotificationCollection();
-        var module = isBuffered ? typeof(BufferedTestModule) : typeof(NotBufferedTestModule);
 
         using var host = await new HostBuilder()
             .ConfigureWebHost(webBuilder =>
@@ -226,9 +317,11 @@ private async Task<List<ApplicationEvent>> RunAsync(string action, ApplicationEv
                         services.AddSystemWebAdapters()
                             .AddHttpApplication(options =>
                             {
+                                options.RegisterModule<BufferToggleModule>();
+
                                 if (mode == RegisterMode.Options)
                                 {
-                                    options.RegisterModule(module);
+                                    options.RegisterModule<NotificationTrackingModule>();
                                 }
 
                                 options.ArePreSendEventsEnabled = true;
@@ -239,13 +332,13 @@ private async Task<List<ApplicationEvent>> RunAsync(string action, ApplicationEv
                     {
                         if (mode == RegisterMode.RegisterModule)
                         {
-                            HttpApplication.RegisterModule(module);
+                            HttpApplication.RegisterModule(typeof(NotificationTrackingModule));
                         }
                         else if (mode == RegisterMode.RegisterModuleOnStartup)
                         {
                             app.ApplicationServices.GetRequiredService<IHostApplicationLifetime>().ApplicationStarted.Register(() =>
                             {
-                                HttpApplication.RegisterModule(module);
+                                HttpApplication.RegisterModule(typeof(NotificationTrackingModule));
                             });
                         }
 
@@ -316,16 +409,22 @@ public Action<IApplicationBuilder> Configure(Action<IApplicationBuilder> next)
             };
     }
 
-    private sealed class NotBufferedTestModule : BufferedTestModule
+    /// <summary>
+    /// Module used to toggle buffer output depending on the test suite we're running
+    /// </summary>
+    private sealed class BufferToggleModule : IHttpModule
     {
-        public override void Init(HttpApplication application)
+        public void Dispose()
+        {
+        }
+
+        public void Init(HttpApplication application)
         {
-            application.BeginRequest += (s, o) => ((HttpApplication)s!).Context.Response.BufferOutput = false;
-            base.Init(application);
+            application.BeginRequest += (s, o) => ((HttpApplication)s!).Context.Response.BufferOutput = IsBuffered;
         }
     }
 
-    private class BufferedTestModule : EventsModule
+    private sealed class NotificationTrackingModule : EventsModule
     {
         protected override void InvokeEvent(HttpContext context, string name)
         {

test/Microsoft.AspNetCore.SystemWebAdapters.CoreServices.Tests/Modules/NotBufferedModuleTests.cs

@@ -1,12 +1,8 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
 namespace Microsoft.AspNetCore.SystemWebAdapters.CoreServices.Tests;
 
-public class NotBufferedModuleTests : ModuleTests
+public class NotBufferedModuleTests : ModuleTests<NotBufferedModuleTests>
 {
-    public NotBufferedModuleTests()
-        : base(false)
-    {
-    }
 }