From 7ec6df807e9765aee4caf9b2d3d725fb5fe2ec87 Mon Sep 17 00:00:00 2001
From: Shreyas Zare <shreyas@technitium.com>
Date: Sat, 6 Jun 2020 15:50:15 +0530
Subject: [PATCH] DnsResourceRecordExtension: Fixed issue in SetGlueRecords().
 Code refactoring done.

---
 .../DnsResourceRecordExtension.cs             | 92 +++++++++++++++----
 1 file changed, 76 insertions(+), 16 deletions(-)
 rename DnsServerCore/Dns/{Zones => ResourceRecords}/DnsResourceRecordExtension.cs (53%)

diff --git a/DnsServerCore/Dns/Zones/DnsResourceRecordExtension.cs b/DnsServerCore/Dns/ResourceRecords/DnsResourceRecordExtension.cs
similarity index 53%
rename from DnsServerCore/Dns/Zones/DnsResourceRecordExtension.cs
rename to DnsServerCore/Dns/ResourceRecords/DnsResourceRecordExtension.cs
index 265470c0..bc0fd564 100644
--- a/DnsServerCore/Dns/Zones/DnsResourceRecordExtension.cs
+++ b/DnsServerCore/Dns/ResourceRecords/DnsResourceRecordExtension.cs
@@ -24,27 +24,43 @@ using System.Net.Sockets;
 using TechnitiumLibrary.Net.Dns;
 using TechnitiumLibrary.Net.Dns.ResourceRecords;
 
-namespace DnsServerCore.Dns.Zones
+namespace DnsServerCore.Dns.ResourceRecords
 {
     public static class DnsResourceRecordExtension
     {
-        public static void SetGlueRecords(this DnsResourceRecord nsRecord, string glueAddresses)
+        public static void SetGlueRecords(this DnsResourceRecord record, string glueAddresses)
         {
             List<IPAddress> addresses = new List<IPAddress>();
 
             foreach (string address in glueAddresses.Split(new char[] { ',' }, StringSplitOptions.RemoveEmptyEntries))
                 addresses.Add(IPAddress.Parse(address.Trim()));
 
-            SetGlueRecords(nsRecord, addresses);
+            SetGlueRecords(record, addresses);
         }
 
-        public static void SetGlueRecords(this DnsResourceRecord nsRecord, IReadOnlyList<IPAddress> glueAddresses)
+        public static void SetGlueRecords(this DnsResourceRecord record, IReadOnlyList<IPAddress> glueAddresses)
         {
-            DnsResourceRecordInfo rrInfo = nsRecord.Tag as DnsResourceRecordInfo;
+            DnsResourceRecordInfo rrInfo = record.Tag as DnsResourceRecordInfo;
             if (rrInfo == null)
             {
                 rrInfo = new DnsResourceRecordInfo();
-                nsRecord.Tag = rrInfo;
+                record.Tag = rrInfo;
+            }
+
+            string domain;
+
+            switch (record.Type)
+            {
+                case DnsResourceRecordType.NS:
+                    domain = (record.RDATA as DnsNSRecord).NSDomainName;
+                    break;
+
+                case DnsResourceRecordType.SOA:
+                    domain = (record.RDATA as DnsSOARecord).MasterNameServer;
+                    break;
+
+                default:
+                    throw new NotSupportedException();
             }
 
             DnsResourceRecord[] glueRecords = new DnsResourceRecord[glueAddresses.Count];
@@ -54,11 +70,11 @@ namespace DnsServerCore.Dns.Zones
                 switch (glueAddresses[i].AddressFamily)
                 {
                     case AddressFamily.InterNetwork:
-                        glueRecords[i] = new DnsResourceRecord((nsRecord.RDATA as DnsNSRecord).NSDomainName, DnsResourceRecordType.A, DnsClass.IN, nsRecord.TtlValue, new DnsARecord(glueAddresses[i]));
+                        glueRecords[i] = new DnsResourceRecord(domain, DnsResourceRecordType.A, DnsClass.IN, record.TtlValue, new DnsARecord(glueAddresses[i]));
                         break;
 
                     case AddressFamily.InterNetworkV6:
-                        glueRecords[i] = new DnsResourceRecord((nsRecord.RDATA as DnsNSRecord).NSDomainName, DnsResourceRecordType.AAAA, DnsClass.IN, nsRecord.TtlValue, new DnsAAAARecord(glueAddresses[i]));
+                        glueRecords[i] = new DnsResourceRecord(domain, DnsResourceRecordType.AAAA, DnsClass.IN, record.TtlValue, new DnsAAAARecord(glueAddresses[i]));
                         break;
                 }
             }
@@ -66,27 +82,71 @@ namespace DnsServerCore.Dns.Zones
             rrInfo.GlueRecords = glueRecords;
         }
 
-        public static IReadOnlyList<DnsResourceRecord> GetGlueRecords(this DnsResourceRecord nsRecord)
+        public static void SetGlueRecords(this DnsResourceRecord record, IReadOnlyList<DnsResourceRecord> glueRecords)
         {
-            DnsResourceRecordInfo rrInfo = nsRecord.Tag as DnsResourceRecordInfo;
+            DnsResourceRecordInfo rrInfo = record.Tag as DnsResourceRecordInfo;
+            if (rrInfo == null)
+            {
+                rrInfo = new DnsResourceRecordInfo();
+                record.Tag = rrInfo;
+            }
+
+            string domain;
+
+            switch (record.Type)
+            {
+                case DnsResourceRecordType.NS:
+                    domain = (record.RDATA as DnsNSRecord).NSDomainName;
+                    break;
+
+                case DnsResourceRecordType.SOA:
+                    domain = (record.RDATA as DnsSOARecord).MasterNameServer;
+                    break;
+
+                default:
+                    throw new NotSupportedException();
+            }
+
+            List<DnsResourceRecord> foundGlueRecords = new List<DnsResourceRecord>();
+
+            foreach (DnsResourceRecord glueRecord in glueRecords)
+            {
+                switch (glueRecord.Type)
+                {
+                    case DnsResourceRecordType.A:
+                    case DnsResourceRecordType.AAAA:
+                        if (glueRecord.Name.Equals(domain, StringComparison.OrdinalIgnoreCase))
+                            foundGlueRecords.Add(glueRecord);
+
+                        break;
+                }
+            }
+
+            if (foundGlueRecords.Count > 0)
+                rrInfo.GlueRecords = foundGlueRecords;
+        }
+
+        public static IReadOnlyList<DnsResourceRecord> GetGlueRecords(this DnsResourceRecord record)
+        {
+            DnsResourceRecordInfo rrInfo = record.Tag as DnsResourceRecordInfo;
             if (rrInfo == null)
                 return Array.Empty<DnsResourceRecord>();
 
             IReadOnlyList<DnsResourceRecord> glueRecords = rrInfo.GlueRecords;
-            if ((glueRecords == null) || (glueRecords.Count == 0))
+            if (glueRecords == null)
                 return Array.Empty<DnsResourceRecord>();
 
             return glueRecords;
         }
 
-        public static IReadOnlyList<DnsResourceRecord> GetGlueRecords(this IReadOnlyList<DnsResourceRecord> nsRecords)
+        public static IReadOnlyList<DnsResourceRecord> GetGlueRecords(this IReadOnlyList<DnsResourceRecord> records)
         {
-            if (nsRecords.Count == 1)
-                return GetGlueRecords(nsRecords[0]);
+            if (records.Count == 1)
+                return GetGlueRecords(records[0]);
 
-            List<DnsResourceRecord> glueRecords = new List<DnsResourceRecord>(nsRecords.Count * 2);
+            List<DnsResourceRecord> glueRecords = new List<DnsResourceRecord>(records.Count * 2);
 
-            foreach (DnsResourceRecord nsRecord in nsRecords)
+            foreach (DnsResourceRecord nsRecord in records)
                 glueRecords.AddRange(GetGlueRecords(nsRecord));
 
             return glueRecords;