From 74a6ea8190b00212053044dc2167d2edadba67ef Mon Sep 17 00:00:00 2001
From: Shreyas Zare <shreyas@technitium.com>
Date: Sat, 11 Sep 2021 15:22:07 +0530
Subject: [PATCH] ResolverPrefetchDnsCache: updated code to support dns app to
 find closest delegation.

---
 DnsServerCore/Dns/ResolverPrefetchDnsCache.cs | 53 +++++++++++++++++--
 1 file changed, 50 insertions(+), 3 deletions(-)

diff --git a/DnsServerCore/Dns/ResolverPrefetchDnsCache.cs b/DnsServerCore/Dns/ResolverPrefetchDnsCache.cs
index 92c44e29..e1017de5 100644
--- a/DnsServerCore/Dns/ResolverPrefetchDnsCache.cs
+++ b/DnsServerCore/Dns/ResolverPrefetchDnsCache.cs
@@ -17,7 +17,11 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 */
 
+using DnsApplicationCommon;
+using DnsServerCore.Dns.Applications;
 using DnsServerCore.Dns.ZoneManagers;
+using System.Net;
+using TechnitiumLibrary.IO;
 using TechnitiumLibrary.Net.Dns;
 
 namespace DnsServerCore.Dns
@@ -32,14 +36,54 @@ namespace DnsServerCore.Dns
 
         #region constructor
 
-        public ResolverPrefetchDnsCache(AuthZoneManager authZoneManager, CacheZoneManager cacheZoneManager, DnsQuestionRecord prefetchQuery)
-            : base(authZoneManager, cacheZoneManager)
+        public ResolverPrefetchDnsCache(DnsApplicationManager dnsApplicationManager, AuthZoneManager authZoneManager, CacheZoneManager cacheZoneManager, DnsQuestionRecord prefetchQuery)
+            : base(dnsApplicationManager, authZoneManager, cacheZoneManager)
         {
             _prefetchQuery = prefetchQuery;
         }
 
         #endregion
 
+        #region private
+
+        private DnsDatagram DnsApplicationQueryClosestDelegation(DnsDatagram request)
+        {
+            if ((_dnsApplicationManager.DnsAuthoritativeRequestHandlers.Count < 1) || (request.Question.Count != 1))
+                return null;
+
+            IPEndPoint localEP = new IPEndPoint(IPAddress.Any, 0);
+            DnsQuestionRecord question = request.Question[0];
+            string currentDomain = question.Name;
+
+            while (true)
+            {
+                DnsDatagram nsRequest = new DnsDatagram(0, false, DnsOpcode.StandardQuery, false, false, true, false, false, false, DnsResponseCode.NoError, new DnsQuestionRecord[] { new DnsQuestionRecord(currentDomain, DnsResourceRecordType.NS, DnsClass.IN) });
+
+                foreach (IDnsAuthoritativeRequestHandler requestHandler in _dnsApplicationManager.DnsAuthoritativeRequestHandlers)
+                {
+                    DnsDatagram nsResponse = requestHandler.ProcessRequestAsync(nsRequest, localEP, DnsTransportProtocol.Tcp, false).Sync();
+                    if (nsResponse is not null)
+                    {
+                        if ((nsResponse.Answer.Count > 0) && (nsResponse.Answer[0].Type == DnsResourceRecordType.NS))
+                            return new DnsDatagram(request.Identifier, true, nsResponse.OPCODE, nsResponse.AuthoritativeAnswer, nsResponse.Truncation, nsResponse.RecursionDesired, nsResponse.RecursionAvailable, nsResponse.AuthenticData, nsResponse.CheckingDisabled, nsResponse.RCODE, request.Question, null, nsResponse.Answer, nsResponse.Additional);
+                        else if ((nsResponse.Authority.Count > 0) && (nsResponse.Authority[0].Type == DnsResourceRecordType.NS))
+                            return new DnsDatagram(request.Identifier, true, nsResponse.OPCODE, nsResponse.AuthoritativeAnswer, nsResponse.Truncation, nsResponse.RecursionDesired, nsResponse.RecursionAvailable, nsResponse.AuthenticData, nsResponse.CheckingDisabled, nsResponse.RCODE, request.Question, null, nsResponse.Authority, nsResponse.Additional);
+                    }
+                }
+
+                //get parent domain
+                int i = currentDomain.IndexOf('.');
+                if (i < 0)
+                    break;
+
+                currentDomain = currentDomain.Substring(i + 1);
+            }
+
+            return null;
+        }
+
+        #endregion
+
         #region public
 
         public override DnsDatagram Query(DnsDatagram request, bool serveStaleAndResetExpiry = false, bool findClosestNameServers = false)
@@ -47,7 +91,10 @@ namespace DnsServerCore.Dns
             if (findClosestNameServers && _prefetchQuery.Equals(request.Question[0]))
             {
                 //return closest name servers so that the recursive resolver queries them to refreshes cache instead of returning response from cache
-                DnsDatagram authResponse = _authZoneManager.QueryClosestDelegation(request);
+                DnsDatagram authResponse = DnsApplicationQueryClosestDelegation(request);
+                if (authResponse is null)
+                    authResponse = _authZoneManager.QueryClosestDelegation(request);
+
                 DnsDatagram cacheResponse = _cacheZoneManager.QueryClosestDelegation(request);
 
                 if ((authResponse is not null) && (authResponse.Authority.Count > 0))