From 30a4ee1ce1d8686f9666361db36492dfaa7897a9 Mon Sep 17 00:00:00 2001
From: Hao Kung <haok@microsoft.com>
Date: Tue, 14 Oct 2014 21:03:16 -0700
Subject: [PATCH] React to EnsureRoutingServices fix

---
 src/Microsoft.AspNet.Routing/BuilderExtensions.cs  |  2 +-
 src/Microsoft.AspNet.Routing/RouterMiddleware.cs   | 12 ++++++++++--
 .../RouterMiddlewareTest.cs                        | 14 ++++++++++----
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/src/Microsoft.AspNet.Routing/BuilderExtensions.cs b/src/Microsoft.AspNet.Routing/BuilderExtensions.cs
index 8fbf5b2bd9..53c3151979 100644
--- a/src/Microsoft.AspNet.Routing/BuilderExtensions.cs
+++ b/src/Microsoft.AspNet.Routing/BuilderExtensions.cs
@@ -9,7 +9,7 @@ namespace Microsoft.AspNet.Builder
     {
         public static IApplicationBuilder UseRouter([NotNull] this IApplicationBuilder builder, [NotNull] IRouter router)
         {
-            return builder.Use((next) => new RouterMiddleware(next, router).Invoke);
+            return builder.UseMiddleware<RouterMiddleware>(router);
         }
     }
 }
\ No newline at end of file
diff --git a/src/Microsoft.AspNet.Routing/RouterMiddleware.cs b/src/Microsoft.AspNet.Routing/RouterMiddleware.cs
index 00fecccb58..6a131760ed 100644
--- a/src/Microsoft.AspNet.Routing/RouterMiddleware.cs
+++ b/src/Microsoft.AspNet.Routing/RouterMiddleware.cs
@@ -8,6 +8,7 @@ using Microsoft.AspNet.Routing;
 using Microsoft.AspNet.Routing.Logging;
 using Microsoft.Framework.DependencyInjection;
 using Microsoft.Framework.Logging;
+using System;
 
 namespace Microsoft.AspNet.Builder
 {
@@ -15,10 +16,11 @@ namespace Microsoft.AspNet.Builder
     {
         private ILogger _logger;
 
-        public RouterMiddleware(RequestDelegate next, IRouter router)
+        public RouterMiddleware(RequestDelegate next, IServiceProvider services, IRouter router)
         {
             Next = next;
             Router = router;
+            Services = services;
         }
 
         private IRouter Router
@@ -33,9 +35,15 @@ namespace Microsoft.AspNet.Builder
             set;
         }
 
+        private IServiceProvider Services
+        {
+            get;
+            set;
+        }
+
         public async Task Invoke(HttpContext httpContext)
         {
-            using (RequestServicesContainer.EnsureRequestServices(httpContext))
+            using (RequestServicesContainer.EnsureRequestServices(httpContext, Services))
             {
                 EnsureLogger(httpContext);
                 using (_logger.BeginScope("RouterMiddleware.Invoke"))
diff --git a/test/Microsoft.AspNet.Routing.Tests/RouterMiddlewareTest.cs b/test/Microsoft.AspNet.Routing.Tests/RouterMiddlewareTest.cs
index d36984efb0..65c54882ca 100644
--- a/test/Microsoft.AspNet.Routing.Tests/RouterMiddlewareTest.cs
+++ b/test/Microsoft.AspNet.Routing.Tests/RouterMiddlewareTest.cs
@@ -10,6 +10,7 @@ using Microsoft.Framework.Logging;
 using Moq;
 #endif
 using Xunit;
+using System;
 
 namespace Microsoft.AspNet.Routing
 {
@@ -31,13 +32,15 @@ namespace Microsoft.AspNet.Routing
             mockContext.Setup(m => m.RequestServices.GetService(typeof(ILoggerFactory)))
                 .Returns(loggerFactory);
 
+            var mockServiceProvider = new Mock<IServiceProvider>();
+
             RequestDelegate next = (c) =>
             {
                 return Task.FromResult<object>(null);
             };
 
             var router = new TestRouter(isHandled);
-            var middleware = new RouterMiddleware(next, router);
+            var middleware = new RouterMiddleware(next, mockServiceProvider.Object, router);
 
             // Act
             await middleware.Invoke(mockContext.Object);
@@ -79,7 +82,8 @@ namespace Microsoft.AspNet.Routing
             };
 
             var router = new TestRouter(isHandled);
-            var middleware = new RouterMiddleware(next, router);
+            var mockServiceProvider = new Mock<IServiceProvider>();
+            var middleware = new RouterMiddleware(next, mockServiceProvider.Object, router);
 
             // Act
             await middleware.Invoke(mockContext.Object);
@@ -114,7 +118,8 @@ namespace Microsoft.AspNet.Routing
             };
 
             var router = new TestRouter(isHandled);
-            var middleware = new RouterMiddleware(next, router);
+            var mockServiceProvider = new Mock<IServiceProvider>();
+            var middleware = new RouterMiddleware(next, mockServiceProvider.Object, router);
 
             // Act
             await middleware.Invoke(mockContext.Object);
@@ -158,7 +163,8 @@ namespace Microsoft.AspNet.Routing
             };
 
             var router = new TestRouter(isHandled);
-            var middleware = new RouterMiddleware(next, router);
+            var mockServiceProvider = new Mock<IServiceProvider>();
+            var middleware = new RouterMiddleware(next, mockServiceProvider.Object, router);
 
             // Act
             await middleware.Invoke(mockContext.Object);