From 6bb1e9c4b62c4f10fdd9634ec862b49678f5902f Mon Sep 17 00:00:00 2001
From: David Jackman <djackman@microsoft.com>
Date: Fri, 24 May 2019 11:23:02 -0600
Subject: [PATCH] Fix issue #176.  Old service provider should be restored to
 the HttpContext after the branch scope is finished.

---
 .../ParallelApplicationPipelinesExtensions.cs   |  2 ++
 .../ParallelPipelinesStartup.cs                 | 17 +++++++++++++++++
 .../ParallelPipelinesTests.cs                   |  9 +++++++++
 3 files changed, 28 insertions(+)

diff --git a/src/WebApiContrib.Core/ParallelApplicationPipelinesExtensions.cs b/src/WebApiContrib.Core/ParallelApplicationPipelinesExtensions.cs
index 3c037a7..ca48045 100644
--- a/src/WebApiContrib.Core/ParallelApplicationPipelinesExtensions.cs
+++ b/src/WebApiContrib.Core/ParallelApplicationPipelinesExtensions.cs
@@ -55,6 +55,7 @@ public static IApplicationBuilder UseBranchWithServices(this IApplicationBuilder
 
             branchBuilder.Use(async (context, next) =>
             {
+                var oldServiceProvider = context.RequestServices;
                 using (var scope = factory.CreateScope())
                 {
                     context.RequestServices = scope.ServiceProvider;
@@ -67,6 +68,7 @@ public static IApplicationBuilder UseBranchWithServices(this IApplicationBuilder
 
                     await next();
                 }
+                context.RequestServices = oldServiceProvider;
             });
 
             appBuilderConfiguration(branchBuilder);
diff --git a/tests/WebApiContrib.Core.Tests/ParallelPipelinesStartup.cs b/tests/WebApiContrib.Core.Tests/ParallelPipelinesStartup.cs
index 9c86c41..7e41538 100644
--- a/tests/WebApiContrib.Core.Tests/ParallelPipelinesStartup.cs
+++ b/tests/WebApiContrib.Core.Tests/ParallelPipelinesStartup.cs
@@ -3,6 +3,7 @@
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Logging;
 
 namespace WebApiContrib.Core.Tests
 {
@@ -87,6 +88,22 @@ public void Configure(IApplicationBuilder app)
                     await c.Response.WriteAsync(service.Hi());
                 });
             });
+
+            // Add middleware before the next branch to verify state of the context coming out of the branch
+            app.Use(async (c, next) =>
+            {
+                // Call into the branch
+                await next();
+
+                // Verify the service provider is working
+                var service = c.RequestServices.GetService<ILogger>();
+            });
+
+            app.UseBranchWithServices("/baz", s =>
+            {
+            }, a =>
+            {
+            });
         }
     }
 }
diff --git a/tests/WebApiContrib.Core.Tests/ParallelPipelinesTests.cs b/tests/WebApiContrib.Core.Tests/ParallelPipelinesTests.cs
index 824cccb..b15e200 100644
--- a/tests/WebApiContrib.Core.Tests/ParallelPipelinesTests.cs
+++ b/tests/WebApiContrib.Core.Tests/ParallelPipelinesTests.cs
@@ -50,5 +50,14 @@ public async Task CanReachBranchWithMultipleEntryPoints(string path)
 
             Assert.Equal($"hi {path.Substring(0, path.Length-1)}", stringResult);
         }
+
+        [Fact]
+        public async Task PipelineBeforeTheBranchStillHasWorkingServiceProvider()
+        {
+            var client = _server.CreateClient();
+
+            var request = new HttpRequestMessage(HttpMethod.Get, $"/baz");
+            await client.SendAsync(request);
+        }
     }
 }