From 9a6881b0e8d2697e5a6a12dd9ddbdfaf6ce6437f Mon Sep 17 00:00:00 2001
From: "N. Taylor Mullen" <nimullen@microsoft.com>
Date: Thu, 18 Jul 2019 16:38:15 -0700
Subject: [PATCH] Fix endpoint routing statefulness.

- In the case that other middleware change the path of an `HttpContext` and cause middleware to re-invoke we used to short-circuit on second time through the middleware pipeline, now we allow routing to occur.
- Added unit tests to validate the clearing of state.

#11233
---
 .../Routing/src/EndpointRoutingMiddleware.cs  | 15 +++++++--
 .../EndpointRoutingMiddlewareTest.cs          | 33 ++++++++++++++++---
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/src/Http/Routing/src/EndpointRoutingMiddleware.cs b/src/Http/Routing/src/EndpointRoutingMiddleware.cs
index 482e86bb6d..cfaff0befb 100644
--- a/src/Http/Routing/src/EndpointRoutingMiddleware.cs
+++ b/src/Http/Routing/src/EndpointRoutingMiddleware.cs
@@ -51,6 +51,8 @@ namespace Microsoft.AspNetCore.Routing
             if (endpoint != null)
             {
                 Log.MatchSkipped(_logger, endpoint);
+
+                // Someone else set the endpoint, we'll let them handle the unsetting.
                 return _next(httpContext);
             }
 
@@ -87,7 +89,7 @@ namespace Microsoft.AspNetCore.Routing
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private Task SetRoutingAndContinue(HttpContext httpContext)
+        private async Task SetRoutingAndContinue(HttpContext httpContext)
         {
             // If there was no mutation of the endpoint then log failure
             var endpoint = httpContext.GetEndpoint();
@@ -107,7 +109,16 @@ namespace Microsoft.AspNetCore.Routing
                 Log.MatchSuccess(_logger, endpoint);
             }
 
-            return _next(httpContext);
+            try
+            {
+                await _next(httpContext);
+            }
+            finally
+            {
+                // We unset the endpoint after calling through to the next middleware. This enables any future calls into
+                // endpoint routing don't no-op from there already being an endpoint set.
+                httpContext.SetEndpoint(endpoint: null);
+            }
         }
 
         // Initialization is async to avoid blocking threads while reflection and things
diff --git a/src/Http/Routing/test/UnitTests/EndpointRoutingMiddlewareTest.cs b/src/Http/Routing/test/UnitTests/EndpointRoutingMiddlewareTest.cs
index 48e1157195..b2e08d9a2c 100644
--- a/src/Http/Routing/test/UnitTests/EndpointRoutingMiddlewareTest.cs
+++ b/src/Http/Routing/test/UnitTests/EndpointRoutingMiddlewareTest.cs
@@ -21,7 +21,30 @@ namespace Microsoft.AspNetCore.Routing
     public class EndpointRoutingMiddlewareTest
     {
         [Fact]
-        public async Task Invoke_OnCall_SetsEndpointFeature()
+        public async Task Invoke_OnException_ResetsEndpoint()
+        {
+            // Arrange
+            var httpContext = CreateHttpContext();
+
+            var middleware = CreateMiddleware(next: context => throw new Exception());
+
+            // Act
+            try
+            {
+                await middleware.Invoke(httpContext);
+            }
+            catch
+            {
+                // Do nothing, we expect the test to throw.
+            }
+
+            // Assert
+            var endpoint = httpContext.GetEndpoint();
+            Assert.Null(endpoint);
+        }
+
+        [Fact]
+        public async Task Invoke_OnCall_SetsEndpointFeatureAndResetsEndpoint()
         {
             // Arrange
             var httpContext = CreateHttpContext();
@@ -34,14 +57,16 @@ namespace Microsoft.AspNetCore.Routing
             // Assert
             var endpointFeature = httpContext.Features.Get<IEndpointFeature>();
             Assert.NotNull(endpointFeature);
+            Assert.Null(endpointFeature.Endpoint);
         }
 
         [Fact]
-        public async Task Invoke_SkipsRouting_IfEndpointSet()
+        public async Task Invoke_SkipsRoutingAndMaintainsEndpoint_IfEndpointSet()
         {
             // Arrange
             var httpContext = CreateHttpContext();
-            httpContext.SetEndpoint(new Endpoint(c => Task.CompletedTask, new EndpointMetadataCollection(), "myapp"));
+            var expectedEndpoint = new Endpoint(c => Task.CompletedTask, new EndpointMetadataCollection(), "myapp");
+            httpContext.SetEndpoint(expectedEndpoint);
 
             var middleware = CreateMiddleware();
 
@@ -50,7 +75,7 @@ namespace Microsoft.AspNetCore.Routing
 
             // Assert
             var endpoint = httpContext.GetEndpoint();
-            Assert.NotNull(endpoint);
+            Assert.Same(expectedEndpoint, endpoint);
             Assert.Equal("myapp", endpoint.DisplayName);
         }