Support NotFound content rendering for a custom Router (#62635)

* SSR streaming not started with custom router.

* NotFound fragment cannot be conditionally added inside of Router and is not nullable. Use if-else instead.

* Feedback: do not dispatch re-executed POST requests.

Author: ilonatommy

src/Components/Components/src/NavigationManager.cs

@@ -61,7 +61,6 @@ public event EventHandler<NotFoundEventArgs> OnNotFound
     // The URI. Always represented an absolute URI.
     private string? _uri;
     private bool _isInitialized;
-    internal string NotFoundPageRoute { get; set; } = string.Empty;
 
     /// <summary>
     /// Gets or sets the current base URI. The <see cref="BaseUri" /> is always represented as an absolute URI in string form with trailing slash.
@@ -211,7 +210,7 @@ private void NotFoundCore()
         }
         else
         {
-            _notFound.Invoke(this, new NotFoundEventArgs(NotFoundPageRoute));
+            _notFound.Invoke(this, new NotFoundEventArgs());
         }
     }
 

src/Components/Components/src/PublicAPI.Unshipped.txt

@@ -8,8 +8,9 @@ Microsoft.AspNetCore.Components.NavigationManager.OnNotFound -> System.EventHand
 Microsoft.AspNetCore.Components.NavigationManager.NotFound() -> void
 Microsoft.AspNetCore.Components.Routing.IHostEnvironmentNavigationManager.Initialize(string! baseUri, string! uri, System.Func<string!, System.Threading.Tasks.Task!>! onNavigateTo) -> void
 Microsoft.AspNetCore.Components.Routing.NotFoundEventArgs
-Microsoft.AspNetCore.Components.Routing.NotFoundEventArgs.NotFoundEventArgs(string! url) -> void
-Microsoft.AspNetCore.Components.Routing.NotFoundEventArgs.Path.get -> string!
+Microsoft.AspNetCore.Components.Routing.NotFoundEventArgs.NotFoundEventArgs() -> void
+Microsoft.AspNetCore.Components.Routing.NotFoundEventArgs.Path.get -> string?
+Microsoft.AspNetCore.Components.Routing.NotFoundEventArgs.Path.set -> void
 Microsoft.AspNetCore.Components.Infrastructure.ComponentStatePersistenceManager.ComponentStatePersistenceManager(Microsoft.Extensions.Logging.ILogger<Microsoft.AspNetCore.Components.Infrastructure.ComponentStatePersistenceManager!>! logger, System.IServiceProvider! serviceProvider) -> void
 Microsoft.AspNetCore.Components.Infrastructure.ComponentStatePersistenceManager.SetPlatformRenderMode(Microsoft.AspNetCore.Components.IComponentRenderMode! renderMode) -> void
 Microsoft.AspNetCore.Components.Infrastructure.RegisterPersistentComponentStateServiceCollectionExtensions

src/Components/Components/src/Routing/NotFoundEventArgs.cs

@@ -9,16 +9,7 @@ namespace Microsoft.AspNetCore.Components.Routing;
 public sealed class NotFoundEventArgs : EventArgs
 {
     /// <summary>
-    /// Gets the path of NotFoundPage.
+    /// Gets the path of NotFoundPage. If the path is set, it indicates that a subscriber has handled the rendering of the NotFound contents.
     /// </summary>
-    public string Path { get; }
-
-    /// <summary>
-    /// Initializes a new instance of <see cref="NotFoundEventArgs" />.
-    /// </summary>
-    public NotFoundEventArgs(string url)
-    {
-        Path = url;
-    }
-
+    public string? Path { get; set; }
 }

src/Components/Components/src/Routing/Router.cs

@@ -29,6 +29,7 @@ static readonly IReadOnlyDictionary<string, object> _emptyParametersDictionary
     string _locationAbsolute;
     bool _navigationInterceptionEnabled;
     ILogger<Router> _logger;
+    string _notFoundPageRoute;
 
     private string _updateScrollPositionForHashLastLocation;
     private bool _updateScrollPositionForHash;
@@ -166,7 +167,7 @@ public async Task SetParametersAsync(ParameterView parameters)
             var routeAttribute = (RouteAttribute)routeAttributes[0];
             if (routeAttribute.Template != null)
             {
-                NavigationManager.NotFoundPageRoute = routeAttribute.Template;
+                _notFoundPageRoute = routeAttribute.Template;
             }
         }
 
@@ -388,10 +389,12 @@ private void OnLocationChanged(object sender, LocationChangedEventArgs args)
         }
     }
 
-    private void OnNotFound(object sender, EventArgs args)
+    private void OnNotFound(object sender, NotFoundEventArgs args)
     {
-        if (_renderHandle.IsInitialized)
+        if (_renderHandle.IsInitialized && NotFoundPage != null)
         {
+            // setting the path signals to the endpoint renderer that router handled rendering
+            args.Path = _notFoundPageRoute;
             Log.DisplayingNotFound(_logger);
             RenderNotFound();
         }

src/Components/Endpoints/src/RazorComponentEndpointInvoker.cs

@@ -106,7 +106,7 @@ await _renderer.InitializeStandardComponentServicesAsync(
             waitForQuiescence: result.IsPost || isErrorHandlerOrReExecuted);
 
         Task quiesceTask;
-        if (!result.IsPost)
+        if (!result.IsPost || isReExecuted)
         {
             quiesceTask = htmlContent.QuiescenceTask;
         }
@@ -130,6 +130,11 @@ await _renderer.InitializeStandardComponentServicesAsync(
             }
         }
 
+        if (_renderer.NotFoundEventArgs != null)
+        {
+            _renderer.SetNotFoundWhenResponseNotStarted();
+        }
+
         if (!quiesceTask.IsCompleted)
         {
             // An incomplete QuiescenceTask indicates there may be streaming rendering updates.
@@ -158,6 +163,10 @@ await _renderer.InitializeStandardComponentServicesAsync(
         if (!quiesceTask.IsCompletedSuccessfully)
         {
             await _renderer.SendStreamingUpdatesAsync(context, quiesceTask, bufferWriter);
+            if (_renderer.NotFoundEventArgs != null)
+            {
+                await _renderer.SetNotFoundWhenResponseHasStarted();
+            }
         }
         else
         {
@@ -171,6 +180,17 @@ await _renderer.InitializeStandardComponentServicesAsync(
             componentStateHtmlContent.WriteTo(bufferWriter, HtmlEncoder.Default);
         }
 
+        if (context.Response.StatusCode == StatusCodes.Status404NotFound &&
+            !isReExecuted &&
+            string.IsNullOrEmpty(_renderer.NotFoundEventArgs?.Path))
+        {
+            // Router did not handle the NotFound event, otherwise this would not be empty.
+            // Don't flush the response if we have an unhandled 404 rendering
+            // This will allow the StatusCodePages middleware to re-execute the request
+            context.Response.ContentType = null;
+            return;
+        }
+
         // Invoke FlushAsync to ensure any buffered content is asynchronously written to the underlying
         // response asynchronously. In the absence of this line, the buffer gets synchronously written to the
         // response as part of the Dispose which has a perf impact.

src/Components/Endpoints/src/Rendering/EndpointHtmlRenderer.EventDispatch.cs

@@ -79,43 +79,44 @@ private Task ReturnErrorResponse(string detailedMessage)
             : Task.CompletedTask;
     }
 
-    internal async Task SetNotFoundResponseAsync(string baseUri, NotFoundEventArgs args)
+    internal void SetNotFoundWhenResponseNotStarted()
     {
-        if (_httpContext.Response.HasStarted ||
-            // POST waits for quiescence -> rendering the NotFoundPage would be queued for the next batch
-            // but we want to send the signal to the renderer to stop rendering future batches -> use client rendering
-            HttpMethods.IsPost(_httpContext.Request.Method))
-        {
-            if (string.IsNullOrEmpty(_notFoundUrl))
-            {
-                _notFoundUrl = GetNotFoundUrl(baseUri, args);
-            }
-            var defaultBufferSize = 16 * 1024;
-            await using var writer = new HttpResponseStreamWriter(_httpContext.Response.Body, Encoding.UTF8, defaultBufferSize, ArrayPool<byte>.Shared, ArrayPool<char>.Shared);
-            using var bufferWriter = new BufferedTextWriter(writer);
-            HandleNotFoundAfterResponseStarted(bufferWriter, _httpContext, _notFoundUrl);
-            await bufferWriter.FlushAsync();
-        }
-        else
+        _httpContext.Response.StatusCode = StatusCodes.Status404NotFound;
+
+        // When the application triggers a NotFound event, we continue rendering the current batch.
+        // However, after completing this batch, we do not want to process any further UI updates,
+        // as we are going to return a 404 status and discard the UI updates generated so far.
+        SignalRendererToFinishRendering();
+    }
+
+    internal async Task SetNotFoundWhenResponseHasStarted()
+    {
+        if (string.IsNullOrEmpty(_notFoundUrl))
         {
-            _httpContext.Response.StatusCode = StatusCodes.Status404NotFound;
+            var baseUri = $"{_httpContext.Request.Scheme}://{_httpContext.Request.Host}{_httpContext.Request.PathBase}/";
+            _notFoundUrl = GetNotFoundUrl(baseUri, NotFoundEventArgs);
         }
+        var defaultBufferSize = 16 * 1024;
+        await using var writer = new HttpResponseStreamWriter(_httpContext.Response.Body, Encoding.UTF8, defaultBufferSize, ArrayPool<byte>.Shared, ArrayPool<char>.Shared);
+        using var bufferWriter = new BufferedTextWriter(writer);
+        HandleNotFoundAfterResponseStarted(bufferWriter, _httpContext, _notFoundUrl);
+        await bufferWriter.FlushAsync();
 
         // When the application triggers a NotFound event, we continue rendering the current batch.
         // However, after completing this batch, we do not want to process any further UI updates,
         // as we are going to return a 404 status and discard the UI updates generated so far.
         SignalRendererToFinishRendering();
     }
 
-    private string GetNotFoundUrl(string baseUri, NotFoundEventArgs args)
+    private string GetNotFoundUrl(string baseUri, NotFoundEventArgs? args)
     {
-        string path = args.Path;
+        string? path = args?.Path;
         if (string.IsNullOrEmpty(path))
         {
             var pathFormat = _httpContext.Items[nameof(StatusCodePagesOptions)] as string;
             if (string.IsNullOrEmpty(pathFormat))
             {
-                throw new InvalidOperationException("The NotFoundPage route must be specified or re-execution middleware has to be set to render NotFoundPage when the response has started.");
+                throw new InvalidOperationException($"The {nameof(Router.NotFoundPage)} route must be specified or re-execution middleware has to be set to render not found content.");
             }
 
             path = pathFormat;

src/Components/Endpoints/src/Rendering/EndpointHtmlRenderer.cs

@@ -63,6 +63,7 @@ public EndpointHtmlRenderer(IServiceProvider serviceProvider, ILoggerFactory log
     }
 
     internal HttpContext? HttpContext => _httpContext;
+    internal NotFoundEventArgs? NotFoundEventArgs { get; private set; }
 
     internal void SetHttpContext(HttpContext httpContext)
     {
@@ -85,10 +86,7 @@ internal async Task InitializeStandardComponentServicesAsync(
         var navigationManager = httpContext.RequestServices.GetRequiredService<NavigationManager>();
         ((IHostEnvironmentNavigationManager)navigationManager)?.Initialize(GetContextBaseUri(httpContext.Request), GetFullUri(httpContext.Request), OnNavigateTo);
 
-        navigationManager?.OnNotFound += (sender, args) =>
-        {
-            _ = GetErrorHandledTask(SetNotFoundResponseAsync(navigationManager.BaseUri, args));
-        };
+        navigationManager?.OnNotFound += (sender, args) => NotFoundEventArgs = args;
 
         var authenticationStateProvider = httpContext.RequestServices.GetService<AuthenticationStateProvider>();
         if (authenticationStateProvider is IHostEnvironmentAuthenticationStateProvider hostEnvironmentAuthenticationStateProvider)

src/Components/Endpoints/test/EndpointHtmlRendererTest.cs

@@ -944,9 +944,9 @@ public async Task Renderer_WhenNoNotFoundPathProvided_Throws()
         httpContext.Items[nameof(StatusCodePagesOptions)] = null; // simulate missing re-execution route
 
         var exception = await Assert.ThrowsAsync<InvalidOperationException>(async () =>
-            await renderer.SetNotFoundResponseAsync(httpContext, new NotFoundEventArgs(""))
+            await renderer.SetNotFoundResponseAsync(httpContext, new NotFoundEventArgs())
         );
-        string expectedError = "The NotFoundPage route must be specified or re-execution middleware has to be set to render NotFoundPage when the response has started.";
+        string expectedError = $"The {nameof(Router.NotFoundPage)} route must be specified or re-execution middleware has to be set to render not found content.";
 
         Assert.Equal(expectedError, exception.Message);
     }
@@ -1823,7 +1823,7 @@ protected override void ProcessPendingRender()
         public async Task SetNotFoundResponseAsync(HttpContext httpContext, NotFoundEventArgs args)
         {
             SetHttpContext(httpContext);
-            await SetNotFoundResponseAsync(httpContext.Request.PathBase, args);
+            await SetNotFoundWhenResponseHasStarted();
         }
     }