Update middleware that assumes UseRouting is called after them, for minimal hosting (#35426)

Author: BrennanConroy

src/Middleware/Diagnostics/src/ExceptionHandler/ExceptionHandlerExtensions.cs

@@ -1,9 +1,11 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System;
+using System.Diagnostics;
 using Microsoft.AspNetCore.Diagnostics;
 using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 
 namespace Microsoft.AspNetCore.Builder
@@ -26,7 +28,7 @@ public static IApplicationBuilder UseExceptionHandler(this IApplicationBuilder a
                 throw new ArgumentNullException(nameof(app));
             }
 
-            return app.UseMiddleware<ExceptionHandlerMiddleware>();
+            return SetExceptionHandlerMiddleware(app, options: null);
         }
 
         /// <summary>
@@ -95,7 +97,50 @@ public static IApplicationBuilder UseExceptionHandler(this IApplicationBuilder a
                 throw new ArgumentNullException(nameof(options));
             }
 
-            return app.UseMiddleware<ExceptionHandlerMiddleware>(Options.Create(options));
+            var iOptions = Options.Create(options);
+            return SetExceptionHandlerMiddleware(app, iOptions);
+        }
+
+        private static IApplicationBuilder SetExceptionHandlerMiddleware(IApplicationBuilder app, IOptions<ExceptionHandlerOptions>? options)
+        {
+            const string globalRouteBuilderKey = "__GlobalEndpointRouteBuilder";
+            // Only use this path if there's a global router (in the 'WebApplication' case).
+            if (app.Properties.TryGetValue(globalRouteBuilderKey, out var routeBuilder) && routeBuilder is not null)
+            {
+                return app.Use(next =>
+                {
+                    var loggerFactory = app.ApplicationServices.GetRequiredService<ILoggerFactory>();
+                    var diagnosticListener = app.ApplicationServices.GetRequiredService<DiagnosticListener>();
+
+                    if (options is null)
+                    {
+                        options = app.ApplicationServices.GetRequiredService<IOptions<ExceptionHandlerOptions>>();
+                    }
+
+                    if (!string.IsNullOrEmpty(options.Value.ExceptionHandlingPath) && options.Value.ExceptionHandler is null)
+                    {
+                        // start a new middleware pipeline
+                        var builder = app.New();
+                        // use the old routing pipeline if it exists so we preserve all the routes and matching logic
+                        // ((IApplicationBuilder)WebApplication).New() does not copy globalRouteBuilderKey automatically like it does for all other properties.
+                        builder.Properties[globalRouteBuilderKey] = routeBuilder;
+                        builder.UseRouting();
+                        // apply the next middleware
+                        builder.Run(next);
+                        // store the pipeline for the error case
+                        options.Value.ExceptionHandler = builder.Build();
+                    }
+
+                    return new ExceptionHandlerMiddleware(next, loggerFactory, options, diagnosticListener).Invoke;
+                });
+            }
+
+            if (options is null)
+            {
+                return app.UseMiddleware<ExceptionHandlerMiddleware>();
+            }
+
+            return app.UseMiddleware<ExceptionHandlerMiddleware>(options);
         }
     }
 }

src/Middleware/Diagnostics/src/StatusCodePage/StatusCodePagesExtensions.cs

@@ -174,7 +174,34 @@ public static IApplicationBuilder UseStatusCodePagesWithReExecute(
                 throw new ArgumentNullException(nameof(app));
             }
 
-            return app.UseStatusCodePages(async context =>
+            const string globalRouteBuilderKey = "__GlobalEndpointRouteBuilder";
+            // Only use this path if there's a global router (in the 'WebApplication' case).
+            if (app.Properties.TryGetValue(globalRouteBuilderKey, out var routeBuilder) && routeBuilder is not null)
+            {
+                return app.Use(next =>
+                {
+                    RequestDelegate? newNext = null;
+                    // start a new middleware pipeline
+                    var builder = app.New();
+                    // use the old routing pipeline if it exists so we preserve all the routes and matching logic
+                    // ((IApplicationBuilder)WebApplication).New() does not copy globalRouteBuilderKey automatically like it does for all other properties.
+                    builder.Properties[globalRouteBuilderKey] = routeBuilder;
+                    builder.UseRouting();
+                    // apply the next middleware
+                    builder.Run(next);
+                    newNext = builder.Build();
+
+                    return new StatusCodePagesMiddleware(next,
+                        Options.Create(new StatusCodePagesOptions() { HandleAsync = CreateHandler(pathFormat, queryFormat, newNext) })).Invoke;
+                });
+            }
+
+            return app.UseStatusCodePages(CreateHandler(pathFormat, queryFormat));
+        }
+
+        private static Func<StatusCodeContext, Task> CreateHandler(string pathFormat, string? queryFormat, RequestDelegate? next = null)
+        {
+            var handler = async (StatusCodeContext context) =>
             {
                 var newPath = new PathString(
                     string.Format(CultureInfo.InvariantCulture, pathFormat, context.HttpContext.Response.StatusCode));
@@ -202,15 +229,24 @@ public static IApplicationBuilder UseStatusCodePagesWithReExecute(
                 context.HttpContext.Request.QueryString = newQueryString;
                 try
                 {
-                    await context.Next(context.HttpContext);
+                    if (next is not null)
+                    {
+                        await next(context.HttpContext);
+                    }
+                    else
+                    {
+                        await context.Next(context.HttpContext);
+                    }
                 }
                 finally
                 {
                     context.HttpContext.Request.QueryString = originalQueryString;
                     context.HttpContext.Request.Path = originalPath;
                     context.HttpContext.Features.Set<IStatusCodeReExecuteFeature?>(null);
                 }
-            });
+            };
+
+            return handler;
         }
     }
 }