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

The implementation that chooses at random is fine as long as all kdcs advertised are working. But in case some of them are not reachable then choosing at random will randomly fail.

The new implementation pings all the kdcs and selects the server that replies first.

Author: 0xced

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);
+    }
+}