Expand other branch pipes (ThreeMammals#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/Configuration/Builder/QoSOptionsBuilder.cs

@@ -6,7 +6,7 @@ public class QoSOptionsBuilder
 
         private int _durationOfBreak;
 
-        private int _timeoutValue;
+        private int _timeoutValue;
 
         private string _key;
 

src/Ocelot/DependencyInjection/IOcelotBuilder.cs

@@ -4,11 +4,15 @@
 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 @@ public class OcelotBuilder : IOcelotBuilder
         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/Errors/Middleware/ExceptionHandlerMiddleware.cs

@@ -34,7 +34,7 @@ public ExceptionHandlerMiddleware(OcelotRequestDelegate next,
         public async Task Invoke(DownstreamContext context)
         {
             try
-            {
+            {
                 //try and get the global request id and set it for logs...
                 //should this basically be immutable per request...i guess it should!
                 //first thing is get config
@@ -44,7 +44,7 @@ public async Task Invoke(DownstreamContext context)
                 {
                     throw new Exception($"{MiddlewareName} setting pipeline errors. IOcelotConfigurationProvider returned {configuration.Errors.ToErrorString()}");
                 }
-
+
                 TrySetGlobalRequestId(context, configuration.Data);
 
                 context.Configuration = configuration.Data;

src/Ocelot/Middleware/OcelotMiddlewareExtensions.cs

@@ -27,15 +27,21 @@ public static async Task<IApplicationBuilder> UseOcelot(this IApplicationBuilder
             await builder.UseOcelot(new OcelotPipelineConfiguration());
 
             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);
-            
+            var configuration = await CreateConfiguration(builder);
+
             CreateAdministrationArea(builder, configuration);
 
-            if(UsingRafty(builder))
+            if (UsingRafty(builder))
             {
                 SetUpRafty(builder);
             }
@@ -78,7 +84,7 @@ private static bool UsingEurekaServiceDiscoveryProvider(IInternalConfiguration c
         private static bool UsingRafty(IApplicationBuilder builder)
         {
             var possible = builder.ApplicationServices.GetService(typeof(INode)) as INode;
-            if(possible != null)
+            if (possible != null)
             {
                 return true;
             }
@@ -99,9 +105,9 @@ private static async Task<IInternalConfiguration> CreateConfiguration(IApplicati
         {
             // make configuration from file system?
             // earlier user needed to add ocelot files in startup configuration stuff, asp.net will map it to this
-            var fileConfig = (IOptions<FileConfiguration>)builder.ApplicationServices.GetService(typeof(IOptions<FileConfiguration>));
-            
-            // now create the config
+            var fileConfig = (IOptions<FileConfiguration>)builder.ApplicationServices.GetService(typeof(IOptions<FileConfiguration>));
+
+            // now create the config
             var internalConfigCreator = (IInternalConfigurationCreator)builder.ApplicationServices.GetService(typeof(IInternalConfigurationCreator));
             var internalConfig = await internalConfigCreator.Create(fileConfig.Value);
 
@@ -196,7 +202,7 @@ private static IInternalConfiguration GetOcelotConfigAndReturn(IInternalConfigur
         {
             var ocelotConfiguration = provider.Get();
 
-            if(ocelotConfiguration?.Data == null || ocelotConfiguration.IsError)
+            if (ocelotConfiguration?.Data == null || ocelotConfiguration.IsError)
             {
                 ThrowToStopOcelotStarting(ocelotConfiguration);
             }
@@ -216,7 +222,7 @@ private static bool UsingConsul(IFileConfigurationRepository fileConfigRepo)
 
         private static void CreateAdministrationArea(IApplicationBuilder builder, IInternalConfiguration configuration)
         {
-            if(!string.IsNullOrEmpty(configuration.AdministrationPath))
+            if (!string.IsNullOrEmpty(configuration.AdministrationPath))
             {
                 builder.Map(configuration.AdministrationPath, app =>
                 {
@@ -239,8 +245,8 @@ private static void ConfigureDiagnosticListener(IApplicationBuilder builder)
             var listener = (OcelotDiagnosticListener)builder.ApplicationServices.GetService(typeof(OcelotDiagnosticListener));
             var diagnosticListener = (DiagnosticListener)builder.ApplicationServices.GetService(typeof(DiagnosticListener));
             diagnosticListener.SubscribeWithAdapter(listener);
-        }
-        
+        }
+
         private static void OnShutdown(IApplicationBuilder app)
         {
             var node = (INode)app.ApplicationServices.GetService(typeof(INode));

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
@@ -37,6 +39,10 @@ public class OcelotPipelineConfiguration
         /// <summary>
         /// This allows the user to implement there own query string manipulation logic
         /// </summary>
-        public Func<DownstreamContext, Func<Task>, Task> PreQueryStringBuilderMiddleware { get; set; }        
+        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 @@ public static IOcelotPipelineBuilder UseMiddleware(this IOcelotPipelineBuilder a
                     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,15 +118,15 @@ public static IOcelotPipelineBuilder UseMiddleware(this IOcelotPipelineBuilder a
 
         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 });
             }
         }
 
         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 @@ public static IOcelotPipelineBuilder MapWhen(this IOcelotPipelineBuilder app, Pr
             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 @@ public static OcelotRequestDelegate BuildOcelotPipeline(this IOcelotPipelineBuil
             // 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 @@ public void should_set_up_pipeline()
                 .BDDfy();
         }
 
+        [Fact]
+        public void should_expand_pipeline()
+        {
+            this.Given(_ => GivenTheDepedenciesAreSetUp())
+                 .When(_ => WhenIExpandBuild())
+                 .Then(_ => ThenThePipelineIsBuilt())
+                 .BDDfy();
+        }
+
+
         private void ThenThePipelineIsBuilt()
         {
             _handlers.ShouldNotBeNull();
@@ -36,6 +51,23 @@ private void WhenIBuild()
             _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 @@ private void WhenIUseAFunc()
             del.Invoke(_downstreamContext);
         }
 
+
+
         private void ThenTheFuncIsInThePipeline()
         {
             _counter.ShouldBe(1);