Expand other branch pipes (#416)

* Expand other branch pipes

* Expand other branch pipes

* Expand other branch pipes

* optimization Expand other branch pipes ，Add Unit test

*  I hope to add two attributes to IOcelotBuilder for easy extension.

src/Ocelot/DependencyInjection/IOcelotBuilder.cs

@@ -4,11 +4,15 @@ using System;
 using System.Net.Http;
 using IdentityServer4.AccessTokenValidation;
 using Ocelot.Middleware.Multiplexer;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Configuration;
 
 namespace Ocelot.DependencyInjection
 {
     public interface IOcelotBuilder
     {
+        IServiceCollection Services { get; }
+        IConfiguration Configuration { get; }
         IOcelotBuilder AddStoreOcelotConfigurationInConsul();
 
         IOcelotBuilder AddCacheManager(Action<ConfigurationBuilderCachePart> settings);

src/Ocelot/DependencyInjection/OcelotBuilder.cs

@@ -55,6 +55,10 @@ namespace Ocelot.DependencyInjection
         private readonly IServiceCollection _services;
         private readonly IConfiguration _configurationRoot;
 
+        public IServiceCollection Services => _services;
+
+        public IConfiguration Configuration => _configurationRoot;
+
         public OcelotBuilder(IServiceCollection services, IConfiguration configurationRoot)
         {
             _configurationRoot = configurationRoot;

src/Ocelot/Middleware/OcelotMiddlewareExtensions.cs

@@ -29,13 +29,19 @@
             return builder;
         }
 
+        public static async Task<IApplicationBuilder> UseOcelot(this IApplicationBuilder builder, Action<OcelotPipelineConfiguration> pipelineConfiguration)
+        {
+            var config = new OcelotPipelineConfiguration();
+            pipelineConfiguration?.Invoke(config);
+            return await builder.UseOcelot(config);
+        }
         public static async Task<IApplicationBuilder> UseOcelot(this IApplicationBuilder builder, OcelotPipelineConfiguration pipelineConfiguration)
         {
             var configuration = await CreateConfiguration(builder);
 
             CreateAdministrationArea(builder, configuration);
 
-            if(UsingRafty(builder))
+            if (UsingRafty(builder))
             {
                 SetUpRafty(builder);
             }
@@ -78,7 +84,7 @@
         private static bool UsingRafty(IApplicationBuilder builder)
         {
             var possible = builder.ApplicationServices.GetService(typeof(INode)) as INode;
-            if(possible != null)
+            if (possible != null)
             {
                 return true;
             }
@@ -196,7 +202,7 @@
         {
             var ocelotConfiguration = provider.Get();
 
-            if(ocelotConfiguration?.Data == null || ocelotConfiguration.IsError)
+            if (ocelotConfiguration?.Data == null || ocelotConfiguration.IsError)
             {
                 ThrowToStopOcelotStarting(ocelotConfiguration);
             }
@@ -216,7 +222,7 @@
 
         private static void CreateAdministrationArea(IApplicationBuilder builder, IInternalConfiguration configuration)
         {
-            if(!string.IsNullOrEmpty(configuration.AdministrationPath))
+            if (!string.IsNullOrEmpty(configuration.AdministrationPath))
             {
                 builder.Map(configuration.AdministrationPath, app =>
                 {

src/Ocelot/Middleware/OcelotPipelineConfiguration.cs

@@ -1,6 +1,8 @@
 namespace Ocelot.Middleware
 {
+    using Ocelot.Middleware.Pipeline;
     using System;
+    using System.Collections.Generic;
     using System.Threading.Tasks;
 
     public class OcelotPipelineConfiguration
@@ -38,5 +40,9 @@
         /// This allows the user to implement there own query string manipulation logic
         /// </summary>
         public Func<DownstreamContext, Func<Task>, Task> PreQueryStringBuilderMiddleware { get; set; }
+        /// <summary>
+        /// This is an extension that will branch to different pipes
+        /// </summary>
+        public List<Func<IOcelotPipelineBuilder, Func<DownstreamContext, bool>>> MapWhenOcelotPipeline { get; } = new List<Func<IOcelotPipelineBuilder, Func<DownstreamContext, bool>>>();
     }
 }

src/Ocelot/Middleware/Pipeline/OcelotPipelineBuilderExtensions.cs

@@ -80,13 +80,14 @@ namespace Ocelot.Middleware.Pipeline
                     var diagnosticListener = (DiagnosticListener)app.ApplicationServices.GetService(typeof(DiagnosticListener));
                     var middlewareName = ocelotDelegate.Target.GetType().Name;
 
-                    OcelotRequestDelegate wrapped = context => {
+                    OcelotRequestDelegate wrapped = context =>
+                    {
                         try
                         {
                             Write(diagnosticListener, "Ocelot.MiddlewareStarted", middlewareName, context);
                             return ocelotDelegate(context);
                         }
-                        catch(Exception ex)
+                        catch (Exception ex)
                         {
                             WriteException(diagnosticListener, ex, "Ocelot.MiddlewareException", middlewareName, context);
                             throw ex;
@@ -117,7 +118,7 @@ namespace Ocelot.Middleware.Pipeline
 
         private static void Write(DiagnosticListener diagnosticListener, string message, string middlewareName, DownstreamContext context)
         {
-            if(diagnosticListener != null)
+            if (diagnosticListener != null)
             {
                 diagnosticListener.Write(message, new { name = middlewareName, context = context });
             }
@@ -125,7 +126,7 @@ namespace Ocelot.Middleware.Pipeline
 
         private static void WriteException(DiagnosticListener diagnosticListener, Exception exception, string message, string middlewareName, DownstreamContext context)
         {
-            if(diagnosticListener != null)
+            if (diagnosticListener != null)
             {
                 diagnosticListener.Write(message, new { name = middlewareName, context = context, exception = exception });
             }
@@ -160,6 +161,28 @@ namespace Ocelot.Middleware.Pipeline
             return app.Use(next => new MapWhenMiddleware(next, options).Invoke);
         }
 
+        public static IOcelotPipelineBuilder MapWhen(this IOcelotPipelineBuilder app, Func<IOcelotPipelineBuilder, Predicate> pipelineBuilderFunc)
+        {
+            if (app == null)
+            {
+                throw new ArgumentNullException(nameof(app));
+            }
+            if (pipelineBuilderFunc == null)
+            {
+                throw new ArgumentNullException(nameof(pipelineBuilderFunc));
+            }
+            var branchBuilder = app.New();
+            var predicate = pipelineBuilderFunc.Invoke(app);
+            var branch = branchBuilder.Build();
+
+            var options = new MapWhenOptions
+            {
+                Predicate = predicate,
+                Branch = branch
+            };
+            return app.Use(next => new MapWhenMiddleware(next, options).Invoke);
+        }
+
         private static Func<T, DownstreamContext, IServiceProvider, Task> Compile<T>(MethodInfo methodinfo, ParameterInfo[] parameters)
         {
             var middleware = typeof(T);

src/Ocelot/Middleware/Pipeline/OcelotPipelineExtensions.cs

@@ -28,6 +28,15 @@ namespace Ocelot.Middleware.Pipeline
             // It also sets the Request Id if anything is set globally
             builder.UseExceptionHandlerMiddleware();
 
+            //Expand other branch pipes
+            if (pipelineConfiguration.MapWhenOcelotPipeline != null)
+            {
+                foreach (var pipeline in pipelineConfiguration.MapWhenOcelotPipeline)
+                {
+                    builder.MapWhen(pipeline);
+                }
+            }
+
             // If the request is for websockets upgrade we fork into a different pipeline
             builder.MapWhen(context => context.HttpContext.WebSockets.IsWebSocketRequest,
                 app =>

test/Ocelot.UnitTests/Middleware/OcelotPipelineExtensionsTests.cs

@@ -3,8 +3,13 @@ namespace Ocelot.UnitTests.Middleware
     using Microsoft.Extensions.Configuration;
     using Microsoft.Extensions.DependencyInjection;
     using Ocelot.DependencyInjection;
+    using Ocelot.DownstreamRouteFinder.Middleware;
+    using Ocelot.DownstreamUrlCreator.Middleware;
+    using Ocelot.LoadBalancer.Middleware;
     using Ocelot.Middleware;
     using Ocelot.Middleware.Pipeline;
+    using Ocelot.Request.Middleware;
+    using Ocelot.WebSockets.Middleware;
     using Pivotal.Discovery.Client;
     using Shouldly;
     using Steeltoe.Common.Discovery;
@@ -26,6 +31,16 @@ namespace Ocelot.UnitTests.Middleware
                 .BDDfy();
         }
 
+        [Fact]
+        public void should_expand_pipeline()
+        {
+            this.Given(_ => GivenTheDepedenciesAreSetUp())
+                 .When(_ => WhenIExpandBuild())
+                 .Then(_ => ThenThePipelineIsBuilt())
+                 .BDDfy();
+        }
+
+
         private void ThenThePipelineIsBuilt()
         {
             _handlers.ShouldNotBeNull();
@@ -36,6 +51,23 @@ namespace Ocelot.UnitTests.Middleware
             _handlers = _builder.BuildOcelotPipeline(new OcelotPipelineConfiguration());
         }
 
+        private void WhenIExpandBuild()
+        {
+            OcelotPipelineConfiguration configuration = new OcelotPipelineConfiguration();
+            configuration.MapWhenOcelotPipeline.Add((app) =>
+            {
+                app.UseDownstreamRouteFinderMiddleware();
+                app.UseDownstreamRequestInitialiser();
+                app.UseLoadBalancingMiddleware();
+                app.UseDownstreamUrlCreatorMiddleware();
+                app.UseWebSocketsProxyMiddleware();
+
+                return context => context.HttpContext.WebSockets.IsWebSocketRequest;
+            });
+            _handlers = _builder.BuildOcelotPipeline(new OcelotPipelineConfiguration());
+        }
+
+
         private void GivenTheDepedenciesAreSetUp()
         {
             IConfigurationBuilder test = new ConfigurationBuilder();

test/Ocelot.UnitTests/Middleware/OcelotPiplineBuilderTests.cs

@@ -79,6 +79,8 @@ namespace Ocelot.UnitTests.Middleware
             del.Invoke(_downstreamContext);
         }
 
+  
+
         private void ThenTheFuncIsInThePipeline()
         {
             _counter.ShouldBe(1);