From 80b5eaac6e00db91684b8fdbb4cbe1179fffce33 Mon Sep 17 00:00:00 2001
From: starkpl <46371668+starkpl@users.noreply.github.com>
Date: Tue, 11 May 2021 07:37:20 +0200
Subject: [PATCH] Add AdditionalServiceRegistration action for custom ASP.NET
 Core DI setup (#611)

---
 .../WireMock.Net.Console.Net452.Classic/MainApp.cs    |  1 +
 src/WireMock.Net/Owin/AspNetCoreSelfHost.cs           |  2 ++
 src/WireMock.Net/Owin/IWireMockMiddlewareOptions.cs   |  5 +++++
 src/WireMock.Net/Owin/WireMockMiddlewareOptions.cs    |  5 +++++
 src/WireMock.Net/Server/WireMockServer.cs             |  1 +
 src/WireMock.Net/Settings/IWireMockServerSettings.cs  | 11 +++++++++++
 src/WireMock.Net/Settings/WireMockServerSettings.cs   | 10 ++++++++++
 7 files changed, 35 insertions(+)

diff --git a/examples/WireMock.Net.Console.Net452.Classic/MainApp.cs b/examples/WireMock.Net.Console.Net452.Classic/MainApp.cs
index 867d42d7..e7d2b537 100644
--- a/examples/WireMock.Net.Console.Net452.Classic/MainApp.cs
+++ b/examples/WireMock.Net.Console.Net452.Classic/MainApp.cs
@@ -59,6 +59,7 @@ namespace WireMock.Net.ConsoleApplication
                 //},
                 PreWireMockMiddlewareInit = app => { System.Console.WriteLine($"PreWireMockMiddlewareInit : {app.GetType()}"); },
                 PostWireMockMiddlewareInit = app => { System.Console.WriteLine($"PostWireMockMiddlewareInit : {app.GetType()}"); },
+                AdditionalServiceRegistration = services => { System.Console.WriteLine($"AdditionalServiceRegistration : {services.GetType()}"); },
                 Logger = new WireMockConsoleLogger(),
 
                 HandlebarsRegistrationCallback = (handlebarsContext, fileSystemHandler) =>
diff --git a/src/WireMock.Net/Owin/AspNetCoreSelfHost.cs b/src/WireMock.Net/Owin/AspNetCoreSelfHost.cs
index 29bf4ba8..b70b8bff 100644
--- a/src/WireMock.Net/Owin/AspNetCoreSelfHost.cs
+++ b/src/WireMock.Net/Owin/AspNetCoreSelfHost.cs
@@ -66,6 +66,8 @@ namespace WireMock.Owin
                     services.AddSingleton<IMappingMatcher, MappingMatcher>();
                     services.AddSingleton<IOwinRequestMapper, OwinRequestMapper>();
                     services.AddSingleton<IOwinResponseMapper, OwinResponseMapper>();
+
+                    _wireMockMiddlewareOptions.AdditionalServiceRegistration?.Invoke(services);
                 })
                 .Configure(appBuilder =>
                 {
diff --git a/src/WireMock.Net/Owin/IWireMockMiddlewareOptions.cs b/src/WireMock.Net/Owin/IWireMockMiddlewareOptions.cs
index f56eb877..4586ae8e 100644
--- a/src/WireMock.Net/Owin/IWireMockMiddlewareOptions.cs
+++ b/src/WireMock.Net/Owin/IWireMockMiddlewareOptions.cs
@@ -8,6 +8,7 @@ using WireMock.Util;
 using Owin;
 #else
 using IAppBuilder = Microsoft.AspNetCore.Builder.IApplicationBuilder;
+using Microsoft.Extensions.DependencyInjection;
 #endif
 
 namespace WireMock.Owin
@@ -36,6 +37,10 @@ namespace WireMock.Owin
 
         Action<IAppBuilder> PostWireMockMiddlewareInit { get; set; }
 
+#if USE_ASPNETCORE
+        Action<IServiceCollection> AdditionalServiceRegistration { get; set; }
+#endif
+
         IFileSystemHandler FileSystemHandler { get; set; }
 
         bool? AllowBodyForAllHttpMethods { get; set; }
diff --git a/src/WireMock.Net/Owin/WireMockMiddlewareOptions.cs b/src/WireMock.Net/Owin/WireMockMiddlewareOptions.cs
index 8b9698b1..49c9aa92 100644
--- a/src/WireMock.Net/Owin/WireMockMiddlewareOptions.cs
+++ b/src/WireMock.Net/Owin/WireMockMiddlewareOptions.cs
@@ -8,6 +8,7 @@ using WireMock.Util;
 using Owin;
 #else
 using IAppBuilder = Microsoft.AspNetCore.Builder.IApplicationBuilder;
+using Microsoft.Extensions.DependencyInjection;
 #endif
 
 namespace WireMock.Owin
@@ -36,6 +37,10 @@ namespace WireMock.Owin
 
         public Action<IAppBuilder> PostWireMockMiddlewareInit { get; set; }
 
+#if USE_ASPNETCORE
+        public Action<IServiceCollection> AdditionalServiceRegistration { get; set; }
+#endif
+
         /// <inheritdoc cref="IWireMockMiddlewareOptions.FileSystemHandler"/>
         public IFileSystemHandler FileSystemHandler { get; set; }
 
diff --git a/src/WireMock.Net/Server/WireMockServer.cs b/src/WireMock.Net/Server/WireMockServer.cs
index d5b0faa9..b9c1daea 100644
--- a/src/WireMock.Net/Server/WireMockServer.cs
+++ b/src/WireMock.Net/Server/WireMockServer.cs
@@ -242,6 +242,7 @@ namespace WireMock.Server
             _mappingToFileSaver = new MappingToFileSaver(_settings, _mappingConverter);
 
 #if USE_ASPNETCORE
+            _options.AdditionalServiceRegistration = _settings.AdditionalServiceRegistration;
             _httpServer = new AspNetCoreSelfHost(_options, urlOptions);
 #else
             _httpServer = new OwinSelfHost(_options, urlOptions);
diff --git a/src/WireMock.Net/Settings/IWireMockServerSettings.cs b/src/WireMock.Net/Settings/IWireMockServerSettings.cs
index d70e0411..3de86643 100644
--- a/src/WireMock.Net/Settings/IWireMockServerSettings.cs
+++ b/src/WireMock.Net/Settings/IWireMockServerSettings.cs
@@ -4,6 +4,9 @@ using JetBrains.Annotations;
 using WireMock.Handlers;
 using WireMock.Logging;
 using WireMock.Matchers;
+#if USE_ASPNETCORE
+using Microsoft.Extensions.DependencyInjection;
+#endif
 
 namespace WireMock.Settings
 {
@@ -109,6 +112,14 @@ namespace WireMock.Settings
         [PublicAPI]
         Action<object> PostWireMockMiddlewareInit { get; set; }
 
+#if USE_ASPNETCORE
+        /// <summary>
+        /// Action which is called with IServiceCollection when ASP.NET Core DI is being configured. [Optional]
+        /// </summary>
+        [PublicAPI]
+        Action<IServiceCollection> AdditionalServiceRegistration { get; set; }
+#endif
+
         /// <summary>
         /// The IWireMockLogger which logs Debug, Info, Warning or Error
         /// </summary>
diff --git a/src/WireMock.Net/Settings/WireMockServerSettings.cs b/src/WireMock.Net/Settings/WireMockServerSettings.cs
index 1074a64c..04ac657c 100644
--- a/src/WireMock.Net/Settings/WireMockServerSettings.cs
+++ b/src/WireMock.Net/Settings/WireMockServerSettings.cs
@@ -4,6 +4,9 @@ using JetBrains.Annotations;
 using Newtonsoft.Json;
 using WireMock.Handlers;
 using WireMock.Logging;
+#if USE_ASPNETCORE
+using Microsoft.Extensions.DependencyInjection;
+#endif
 
 namespace WireMock.Settings
 {
@@ -79,6 +82,13 @@ namespace WireMock.Settings
         [JsonIgnore]
         public Action<object> PostWireMockMiddlewareInit { get; set; }
 
+#if USE_ASPNETCORE
+        /// <inheritdoc cref="IWireMockServerSettings.AdditionalServiceRegistration"/>
+        [PublicAPI]
+        [JsonIgnore]
+        public Action<IServiceCollection> AdditionalServiceRegistration { get; set; }
+#endif
+
         /// <inheritdoc cref="IWireMockServerSettings.Logger"/>
         [PublicAPI]
         [JsonIgnore]