Merge pull request #415 from 0xced/SelectedPreferredInstance-Fastest

Select the fastest (Kdc/Kpasswd) instance instead of at random

Author: SteveSyfuhs

Kerberos.NET/Client/Transport/KerberosTransportBase.cs

@@ -6,6 +6,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Net.NetworkInformation;
 using System.Threading;
 using System.Threading.Tasks;
 using Kerberos.NET.Asn1;
@@ -19,15 +20,15 @@ namespace Kerberos.NET.Transport
 {
     public abstract class KerberosTransportBase : IKerberosTransport2, IDisposable
     {
-        private static readonly Random Random = new Random();
-
         protected KerberosTransportBase(ILoggerFactory logger)
         {
             this.ClientRealmService = new ClientDomainService(logger);
         }
 
         private bool disposedValue;
 
+        private DnsRecord fastest;
+
         public virtual bool TransportFailed { get; set; }
 
         public virtual KerberosTransportException LastError { get; set; }
@@ -165,34 +166,58 @@ public void Dispose()
         protected virtual async Task<DnsRecord> LocatePreferredKdc(string domain, string servicePrefix)
         {
             var results = await this.LocateKdc(domain, servicePrefix);
-            return SelectedPreferredInstance(domain, servicePrefix, results, ClientDomainService.DefaultKerberosPort);
+            return await SelectedPreferredInstance(domain, servicePrefix, results, ClientDomainService.DefaultKerberosPort);
         }
 
         protected virtual async Task<DnsRecord> LocatePreferredKpasswd(string domain, string servicePrefix)
         {
             var results = await this.LocateKpasswd(domain, servicePrefix);
-            return SelectedPreferredInstance(domain, servicePrefix, results, ClientDomainService.DefaultKpasswdPort);
+            return await SelectedPreferredInstance(domain, servicePrefix, results, ClientDomainService.DefaultKpasswdPort);
         }
 
-        protected virtual DnsRecord SelectedPreferredInstance(string domain, string servicePrefix, IEnumerable<DnsRecord> results, int defaultPort)
+        protected virtual async Task<DnsRecord> SelectedPreferredInstance(string domain, string servicePrefix, IEnumerable<DnsRecord> results, int defaultPort)
         {
-            results = results.Where(r => r.Name.StartsWith(servicePrefix));
+            if (results.Contains(fastest, DnsRecordComparer.Instance))
+            {
+                return fastest;
+            }
 
-            var rand = Random.Next(0, results?.Count() ?? 0);
+            fastest = await results.Where(r => r.Name.StartsWith(servicePrefix)).GetFastestAsync(PingAsync);
+            return fastest ?? throw new KerberosTransportException($"Cannot locate SRV record for {domain}");
+        }
 
-            var srv = results?.ElementAtOrDefault(rand);
+        private async Task<DnsRecord> PingAsync(DnsRecord record, CancellationToken cancellationToken)
+        {
+            using var ping = new Ping();
+            cancellationToken.Register(() => ping.SendAsyncCancel());
+            var reply = await ping.SendPingAsync(record.Target, Convert.ToInt32(ConnectTimeout.TotalMilliseconds));
+            return reply.Status == IPStatus.Success ? record : throw new PingException($"Ping {record.Target} returned {reply.Status}");
+        }
 
-            if (srv == null)
+        private class DnsRecordComparer : IEqualityComparer<DnsRecord>
+        {
+            public static readonly DnsRecordComparer Instance = new();
+
+            private DnsRecordComparer()
             {
-                throw new KerberosTransportException($"Cannot locate SRV record for {domain}");
             }
 
-            if (srv.Port <= 0)
+            public bool Equals(DnsRecord x, DnsRecord y)
             {
-                srv.Port = defaultPort;
+                if (ReferenceEquals(x, y)) return true;
+                if (x is null) return false;
+                if (y is null) return false;
+                if (x.GetType() != y.GetType()) return false;
+                return x.Target == y.Target && x.Port == y.Port;
             }
 
-            return srv;
+            public int GetHashCode(DnsRecord obj)
+            {
+                unchecked
+                {
+                    return ((obj.Target != null ? obj.Target.GetHashCode() : 0) * 397) ^ obj.Port;
+                }
+            }
         }
     }
 }

Kerberos.NET/TaskExtensions.cs

@@ -0,0 +1,42 @@
+// -----------------------------------------------------------------------
+// Licensed to The .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// -----------------------------------------------------------------------
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+
+internal static class TaskExtensions
+{
+    public static async Task<TResult> GetFastestAsync<TSource, TResult>(this IEnumerable<TSource> source, Func<TSource, CancellationToken, Task<TResult>> task, CancellationToken cancellationToken = default)
+    {
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        var tasks = new HashSet<Task<TResult>>(source.Select(e => task(e, cts.Token)));
+        if (tasks.Count == 0)
+        {
+            return default;
+        }
+
+        var exceptions = new List<Exception>();
+        do
+        {
+            var completedTask = await Task.WhenAny(tasks);
+            if (completedTask.Status == TaskStatus.RanToCompletion)
+            {
+                cts.Cancel();
+                return completedTask.Result;
+            }
+
+            if (completedTask.Exception != null)
+            {
+                exceptions.AddRange(completedTask.Exception.InnerExceptions);
+            }
+            tasks.Remove(completedTask);
+        } while (tasks.Count > 0);
+
+        throw new AggregateException(exceptions);
+    }
+}