Add possibility of customizing how discovery metadata are created (#10)

keghub · Dec 17, 2019 · cc040ac · cc040ac
1 parent ef75b12
commit cc040ac
Show file tree

Hide file tree

Showing 6 changed files with 224 additions and 22 deletions.
diff --git a/src/ServiceModel.ECS/ECSMetadataExtensions.cs b/src/ServiceModel.ECS/ECSMetadataExtensions.cs
@@ -1,12 +1,19 @@
 using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.ServiceModel.Description;
+using System.ServiceModel.Discovery;
 using EMG.Extensions.Configuration.Model;
 using EMG.Utilities.ServiceModel.Configuration;
+using EMG.Utilities.ServiceModel.Discovery;
 using Microsoft.Extensions.Configuration;
+using Microsoft.Extensions.DependencyInjection;
 
 namespace EMG.Utilities
 {
     public static class ECSMetadataExtensions
     {
+        [Obsolete("Use AddECSContainerSupport instead.")]
         public static HttpEndpointAddress UseECS(this HttpEndpointAddress endpointAddress, IConfiguration configuration, Action<ECSContainerMetadataOptions> configureOptions = null)
         {
             if (Environment.GetEnvironmentVariable(EMG.Extensions.Configuration.ECSMetadataExtensions.ECSContainerMetadataFileKey) == null)
@@ -34,6 +41,7 @@ public static HttpEndpointAddress UseECS(this HttpEndpointAddress endpointAddres
             return EndpointAddress.ForHttp(containerMetadata.HostPrivateIPv4Address, endpointAddress.Path, portMapping.HostPort, endpointAddress.IsSecure);
         }
 
+        [Obsolete("Use AddECSContainerSupport instead.")]
         public static NetTcpEndpointAddress UseECS(this NetTcpEndpointAddress endpointAddress, IConfiguration configuration, Action<ECSContainerMetadataOptions> configureOptions = null)
         {
             if (Environment.GetEnvironmentVariable(EMG.Extensions.Configuration.ECSMetadataExtensions.ECSContainerMetadataFileKey) == null)
@@ -61,5 +69,59 @@ public static NetTcpEndpointAddress UseECS(this NetTcpEndpointAddress endpointAd
             return EndpointAddress.ForNetTcp(portMapping.HostPort, containerMetadata.HostPrivateIPv4Address, endpointAddress.Path);
 
         }
+
+        public static IServiceCollection AddECSContainerSupport(this IServiceCollection services, IConfiguration configuration)
+        {
+            if (Environment.GetEnvironmentVariable(EMG.Extensions.Configuration.ECSMetadataExtensions.ECSContainerMetadataFileKey) != null)
+            {
+                services.Configure<AnnouncementServiceOptions>(o => o.EndpointDiscoveryMetadata = GetEndpointMetadataFromECS);
+            }
+
+            return services;
+
+            EndpointDiscoveryMetadata GetEndpointMetadataFromECS(ServiceEndpoint endpoint)
+            {
+                if (endpoint.Address is null)
+                {
+                    throw new ArgumentNullException(nameof(endpoint),"endpoint address should not be null");
+                }
+
+                var containerMetadata = configuration.Get<ECSContainerMetadata>();
+
+                var endpointMetadata = EndpointDiscoveryMetadata.FromServiceEndpoint(endpoint);
+
+                if (endpointMetadata is null)
+                {
+                    return null;
+                }
+
+                if (containerMetadata?.HostPrivateIPv4Address is null)
+                {
+                    return endpointMetadata;
+                }
+
+                var uriBuilder = new UriBuilder(endpointMetadata.Address.Uri)
+                {
+                    Host = containerMetadata.HostPrivateIPv4Address,
+                    Port = GetMappedPortOrDefault(containerMetadata.PortMappings, endpointMetadata.Address.Uri.Port)
+                };
+
+                endpointMetadata.Address = new System.ServiceModel.EndpointAddress(uriBuilder.Uri.ToString());
+
+                return endpointMetadata;
+            }
+
+            int GetMappedPortOrDefault(IEnumerable<PortMapping> mappings, int port)
+            {
+                var hostPortByContainerPort = mappings.ToDictionary(k => k.ContainerPort, v => v.HostPort);
+
+                if (hostPortByContainerPort.TryGetValue(port, out var mappedPort))
+                {
+                    return mappedPort;
+                }
+
+                return port;
+            }
+        }
     }
 }
diff --git a/src/ServiceModel.ECS/ServiceModel.ECS.csproj b/src/ServiceModel.ECS/ServiceModel.ECS.csproj
@@ -19,4 +19,9 @@
     <ProjectReference Include="..\ServiceModel\ServiceModel.csproj" />
   </ItemGroup>
 
+  <ItemGroup>
+    <Reference Include="System.ServiceModel" />
+    <Reference Include="System.ServiceModel.Discovery" />
+  </ItemGroup>
+
 </Project>
diff --git a/src/ServiceModel/ServiceModel/Discovery/AnnouncementService.cs b/src/ServiceModel/ServiceModel/Discovery/AnnouncementService.cs
@@ -22,6 +22,8 @@ public class AnnouncementServiceOptions
         public TimeSpan Interval { get; set; }
 
         public Binding Binding { get; set; }
+
+        public Func<ServiceEndpoint, EndpointDiscoveryMetadata> EndpointDiscoveryMetadata { get; set; } = System.ServiceModel.Discovery.EndpointDiscoveryMetadata.FromServiceEndpoint;
     }
 
     public class AnnouncementService : IAnnouncementService
@@ -54,7 +56,7 @@ from endpoint in endpointsToAnnounce
             {
                 using (var client = CreateClient())
                 {
-                    var metadata = EndpointDiscoveryMetadata.FromServiceEndpoint(endpoint);
+                    var metadata = _options.EndpointDiscoveryMetadata(endpoint);
 
                     if (metadata != null)
                     {
@@ -73,7 +75,7 @@ private void UnannounceService(IReadOnlyList<ServiceEndpoint> endpoints)
                 {
                     foreach (var endpoint in endpoints)
                     {
-                        var metadata = EndpointDiscoveryMetadata.FromServiceEndpoint(endpoint);
+                        var metadata = _options.EndpointDiscoveryMetadata(endpoint);
 
                         if (metadata != null)
                         {

diff --git a/tests/Tests.ServiceModel.ECS/CustomAutoDataAttribute.cs b/tests/Tests.ServiceModel.ECS/CustomAutoDataAttribute.cs
@@ -1,10 +1,13 @@
 using System;
 using System.Collections.Generic;
+using System.Linq;
+using System.ServiceModel.Description;
 using AutoFixture;
 using AutoFixture.AutoMoq;
 using AutoFixture.NUnit3;
 using EMG.Extensions.Configuration.Model;
 using Microsoft.Extensions.Configuration;
+using Microsoft.Extensions.DependencyInjection;
 
 namespace Tests
 {
@@ -33,13 +36,28 @@ public static IFixture CreateFixture()
                 GenerateDelegates = true
             });
 
+            fixture.Customize<ECSContainerMetadata>(o => o
+                                                        .With(p => p.PortMappings, (Generator<PortMapping> g) => g.Take(1).ToArray())
+            );
+
             fixture.Customize<IConfiguration>(o => o.FromFactory((ConfigurationBuilder builder, ECSContainerMetadata metadata) =>
             {
                 builder.AddObject(metadata);
 
                 return builder.Build();
             }));
 
+            fixture.Customize<ServiceCollection>(o => o.FromFactory(() =>
+            {
+                var services = new ServiceCollection();
+
+                services.AddOptions();
+
+                return services;
+            }));
+
+            fixture.Customize<ServiceEndpoint>(o => o.Without(p => p.Address));
+
             return fixture;
         }
     }