From 35e25eeaf0c2e104ada3c55faa9e0c687deb8ca1 Mon Sep 17 00:00:00 2001
From: John Luo <jt.luo@mail.utoronto.ca>
Date: Mon, 9 Nov 2015 10:07:35 -0800
Subject: [PATCH] Adding option to configure services when exposing the ASP.NET
 5 pipeline via OWIN #398

---
 src/Microsoft.AspNet.Owin/OwinExtensions.cs   | 21 +++--
 .../OwinExtensionTests.cs                     | 80 +++++++++++++++++++
 .../OwinFeatureCollectionTests.cs             |  1 -
 test/Microsoft.AspNet.Owin.Tests/project.json |  1 +
 4 files changed, 97 insertions(+), 6 deletions(-)
 create mode 100644 test/Microsoft.AspNet.Owin.Tests/OwinExtensionTests.cs

diff --git a/src/Microsoft.AspNet.Owin/OwinExtensions.cs b/src/Microsoft.AspNet.Owin/OwinExtensions.cs
index 3f31ce59..3f688265 100644
--- a/src/Microsoft.AspNet.Owin/OwinExtensions.cs
+++ b/src/Microsoft.AspNet.Owin/OwinExtensions.cs
@@ -66,22 +66,27 @@ public static IApplicationBuilder UseOwin(this IApplicationBuilder builder, Acti
         }
 
         public static IApplicationBuilder UseBuilder(this AddMiddleware app)
+        {
+            return app.UseBuilder(serviceProvider: null);
+        }
+
+        public static IApplicationBuilder UseBuilder(this AddMiddleware app, IServiceProvider serviceProvider)
         {
             // Adapt WebSockets by default.
             app(OwinWebSocketAcceptAdapter.AdaptWebSockets);
-            var builder = new ApplicationBuilder(serviceProvider: null);
+            var builder = new ApplicationBuilder(serviceProvider: serviceProvider);
 
-            CreateMiddleware middleware = CreateMiddlewareFactory(exit =>
+            var middleware = CreateMiddlewareFactory(exit =>
             {
                 builder.Use(ignored => exit);
                 return builder.Build();
-            });
+            }, builder.ApplicationServices);
 
             app(middleware);
             return builder;
         }
 
-        private static CreateMiddleware CreateMiddlewareFactory(Func<RequestDelegate, RequestDelegate> middleware)
+        private static CreateMiddleware CreateMiddlewareFactory(Func<RequestDelegate, RequestDelegate> middleware, IServiceProvider applicationServices)
         {
             return next =>
             {
@@ -105,6 +110,7 @@ private static CreateMiddleware CreateMiddlewareFactory(Func<RequestDelegate, Re
                         context = new DefaultHttpContext(
                                     new FeatureCollection(
                                         new OwinFeatureCollection(env)));
+                        context.ApplicationServices = applicationServices;
                     }
 
                     return app.Invoke(context);
@@ -114,7 +120,12 @@ private static CreateMiddleware CreateMiddlewareFactory(Func<RequestDelegate, Re
 
         public static AddMiddleware UseBuilder(this AddMiddleware app, Action<IApplicationBuilder> pipeline)
         {
-            var builder = app.UseBuilder();
+            return app.UseBuilder(pipeline, serviceProvider: null);
+        }
+
+        public static AddMiddleware UseBuilder(this AddMiddleware app, Action<IApplicationBuilder> pipeline, IServiceProvider serviceProvider)
+        {
+            var builder = app.UseBuilder(serviceProvider);
             pipeline(builder);
             return app;
         }
diff --git a/test/Microsoft.AspNet.Owin.Tests/OwinExtensionTests.cs b/test/Microsoft.AspNet.Owin.Tests/OwinExtensionTests.cs
new file mode 100644
index 00000000..67d3683b
--- /dev/null
+++ b/test/Microsoft.AspNet.Owin.Tests/OwinExtensionTests.cs
@@ -0,0 +1,80 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+using Microsoft.AspNet.Builder;
+using Microsoft.Extensions.DependencyInjection;
+using Xunit;
+
+
+namespace Microsoft.AspNet.Owin
+{
+    using AppFunc = Func<IDictionary<string, object>, Task>;
+    using CreateMiddleware = Func<
+          Func<IDictionary<string, object>, Task>,
+          Func<IDictionary<string, object>, Task>
+        >;
+    using AddMiddleware = Action<Func<
+          Func<IDictionary<string, object>, Task>,
+          Func<IDictionary<string, object>, Task>
+        >>;
+
+    public class OwinExtensionTests
+    {
+        static AppFunc notFound = async env => env["owin.ResponseStatusCode"] = 404;
+
+        [Fact]
+        public void OwinConfigureServiceProviderAddsServices()
+        {
+            IList<CreateMiddleware> list = new List<CreateMiddleware>();
+            AddMiddleware build = list.Add;
+            IServiceProvider serviceProvider = null;
+            FakeService fakeService = null;
+
+            var builder = build.UseBuilder(applicationBuilder =>
+            {
+                serviceProvider = applicationBuilder.ApplicationServices;
+                applicationBuilder.Run(async context =>
+                {
+                    fakeService = context.ApplicationServices.GetService<FakeService>();
+                });
+            }, new ServiceCollection().AddSingleton(new FakeService()).BuildServiceProvider());
+
+            list.Reverse().Aggregate(notFound, (next, middleware) => middleware(next)).Invoke(new Dictionary<string, object>());
+
+            Assert.NotNull(fakeService);
+            Assert.NotNull(serviceProvider?.GetService<FakeService>());
+        }
+
+        [Fact]
+        public void OwinDefaultNoServices()
+        {
+            IList<CreateMiddleware> list = new List<CreateMiddleware>();
+            AddMiddleware build = list.Add;
+            IServiceProvider serviceProvider = null;
+            FakeService fakeService = null;
+
+            var builder = build.UseBuilder(applicationBuilder =>
+            {
+                serviceProvider = applicationBuilder.ApplicationServices;
+                applicationBuilder.Run(async context =>
+                {
+                    fakeService = context.ApplicationServices.GetService<FakeService>();
+                });
+            });
+
+            list.Reverse().Aggregate(notFound, (next, middleware) => middleware(next)).Invoke(new Dictionary<string, object>());
+
+            Assert.Null(fakeService);
+            Assert.Null(serviceProvider);
+        }
+
+        public class FakeService
+        {
+
+        }
+    }
+}
diff --git a/test/Microsoft.AspNet.Owin.Tests/OwinFeatureCollectionTests.cs b/test/Microsoft.AspNet.Owin.Tests/OwinFeatureCollectionTests.cs
index b4958629..b945af47 100644
--- a/test/Microsoft.AspNet.Owin.Tests/OwinFeatureCollectionTests.cs
+++ b/test/Microsoft.AspNet.Owin.Tests/OwinFeatureCollectionTests.cs
@@ -2,7 +2,6 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System.Collections.Generic;
-using System.Linq;
 using Microsoft.AspNet.Http.Features;
 using Xunit;
 
diff --git a/test/Microsoft.AspNet.Owin.Tests/project.json b/test/Microsoft.AspNet.Owin.Tests/project.json
index 05c31c8a..9fc7f68b 100644
--- a/test/Microsoft.AspNet.Owin.Tests/project.json
+++ b/test/Microsoft.AspNet.Owin.Tests/project.json
@@ -2,6 +2,7 @@
   "dependencies": {
     "Microsoft.AspNet.Http": "1.0.0-*",
     "Microsoft.AspNet.Owin": "1.0.0-*",
+    "Microsoft.Extensions.DependencyInjection": "1.0.0-*",
     "xunit.runner.aspnet": "2.0.0-aspnet-*"
   },
   "commands": {