From 51b545b3836e2910d7172fcb1a2ae2f1302b276a Mon Sep 17 00:00:00 2001
From: Shreyas Zare <shreyas@technitium.com>
Date: Sun, 20 Nov 2022 15:24:21 +0530
Subject: [PATCH] GeoCountryApp: refactored and fixed ECS implementation.

---
 Apps/GeoCountryApp/Address.cs | 43 +++++++++++++++--------------------
 Apps/GeoCountryApp/CNAME.cs   | 43 +++++++++++++++--------------------
 2 files changed, 36 insertions(+), 50 deletions(-)

diff --git a/Apps/GeoCountryApp/Address.cs b/Apps/GeoCountryApp/Address.cs
index 81ec52d0..517a0c74 100644
--- a/Apps/GeoCountryApp/Address.cs
+++ b/Apps/GeoCountryApp/Address.cs
@@ -85,27 +85,16 @@ namespace GeoCountry
                     dynamic jsonAppRecordData = JsonConvert.DeserializeObject(appRecordData);
                     dynamic jsonCountry = null;
 
-                    EDnsClientSubnetOptionData clientSubnet = null;
-
-                    if (request.EDNS is not null)
+                    bool ecsUsed = false;
+                    EDnsClientSubnetOptionData requestECS = request.GetEDnsClientSubnetOption();
+                    if (requestECS is not null)
                     {
-                        foreach (EDnsOption option in request.EDNS.Options)
+                        if (_maxMind.DatabaseReader.TryCountry(requestECS.Address, out CountryResponse csResponse))
                         {
-                            if (option.Code == EDnsOptionCode.EDNS_CLIENT_SUBNET)
-                            {
-                                EDnsClientSubnetOptionData cs = option.Data as EDnsClientSubnetOptionData;
-
-                                if (_maxMind.DatabaseReader.TryCountry(cs.Address, out CountryResponse csResponse))
-                                {
-                                    jsonCountry = jsonAppRecordData[csResponse.Country.IsoCode];
-                                    if (jsonCountry is null)
-                                        jsonCountry = jsonAppRecordData["default"];
-                                    else
-                                        clientSubnet = cs;
-                                }
-
-                                break;
-                            }
+                            ecsUsed = true;
+                            jsonCountry = jsonAppRecordData[csResponse.Country.IsoCode];
+                            if (jsonCountry is null)
+                                jsonCountry = jsonAppRecordData["default"];
                         }
                     }
 
@@ -157,14 +146,18 @@ namespace GeoCountry
                     if (answers.Count > 1)
                         answers.Shuffle();
 
-                    EDnsOption[] options = null;
+                    EDnsOption[] options;
 
-                    if (clientSubnet is not null)
+                    if (requestECS is null)
                     {
-                        options = new EDnsOption[]
-                        {
-                            new EDnsOption(EDnsOptionCode.EDNS_CLIENT_SUBNET, new EDnsClientSubnetOptionData(clientSubnet.SourcePrefixLength, clientSubnet.SourcePrefixLength, clientSubnet.Address))
-                        };
+                        options = null;
+                    }
+                    else
+                    {
+                        if (ecsUsed)
+                            options = EDnsClientSubnetOptionData.GetEDnsClientSubnetOption(requestECS.SourcePrefixLength, requestECS.SourcePrefixLength, requestECS.Address);
+                        else
+                            options = EDnsClientSubnetOptionData.GetEDnsClientSubnetOption(requestECS.SourcePrefixLength, 0, requestECS.Address);
                     }
 
                     return Task.FromResult(new DnsDatagram(request.Identifier, true, request.OPCODE, true, false, request.RecursionDesired, isRecursionAllowed, false, false, DnsResponseCode.NoError, request.Question, answers, null, null, _dnsServer.UdpPayloadSize, EDnsHeaderFlags.None, options));
diff --git a/Apps/GeoCountryApp/CNAME.cs b/Apps/GeoCountryApp/CNAME.cs
index a6af618c..63f16326 100644
--- a/Apps/GeoCountryApp/CNAME.cs
+++ b/Apps/GeoCountryApp/CNAME.cs
@@ -79,27 +79,16 @@ namespace GeoCountry
             dynamic jsonAppRecordData = JsonConvert.DeserializeObject(appRecordData);
             dynamic jsonCountry = null;
 
-            EDnsClientSubnetOptionData clientSubnet = null;
-
-            if (request.EDNS is not null)
+            bool ecsUsed = false;
+            EDnsClientSubnetOptionData requestECS = request.GetEDnsClientSubnetOption();
+            if (requestECS is not null)
             {
-                foreach (EDnsOption option in request.EDNS.Options)
+                if (_maxMind.DatabaseReader.TryCountry(requestECS.Address, out CountryResponse csResponse))
                 {
-                    if (option.Code == EDnsOptionCode.EDNS_CLIENT_SUBNET)
-                    {
-                        EDnsClientSubnetOptionData cs = option.Data as EDnsClientSubnetOptionData;
-
-                        if (_maxMind.DatabaseReader.TryCountry(cs.Address, out CountryResponse csResponse))
-                        {
-                            jsonCountry = jsonAppRecordData[csResponse.Country.IsoCode];
-                            if (jsonCountry is null)
-                                jsonCountry = jsonAppRecordData["default"];
-                            else
-                                clientSubnet = cs;
-                        }
-
-                        break;
-                    }
+                    ecsUsed = true;
+                    jsonCountry = jsonAppRecordData[csResponse.Country.IsoCode];
+                    if (jsonCountry is null)
+                        jsonCountry = jsonAppRecordData["default"];
                 }
             }
 
@@ -131,14 +120,18 @@ namespace GeoCountry
             else
                 answers = new DnsResourceRecord[] { new DnsResourceRecord(request.Question[0].Name, DnsResourceRecordType.CNAME, DnsClass.IN, appRecordTtl, new DnsCNAMERecordData(cname)) };
 
-            EDnsOption[] options = null;
+            EDnsOption[] options;
 
-            if (clientSubnet is not null)
+            if (requestECS is null)
             {
-                options = new EDnsOption[]
-                {
-                    new EDnsOption(EDnsOptionCode.EDNS_CLIENT_SUBNET, new EDnsClientSubnetOptionData(clientSubnet.SourcePrefixLength, clientSubnet.SourcePrefixLength, clientSubnet.Address))
-                };
+                options = null;
+            }
+            else
+            {
+                if (ecsUsed)
+                    options = EDnsClientSubnetOptionData.GetEDnsClientSubnetOption(requestECS.SourcePrefixLength, requestECS.SourcePrefixLength, requestECS.Address);
+                else
+                    options = EDnsClientSubnetOptionData.GetEDnsClientSubnetOption(requestECS.SourcePrefixLength, 0, requestECS.Address);
             }
 
             return Task.FromResult(new DnsDatagram(request.Identifier, true, request.OPCODE, true, false, request.RecursionDesired, isRecursionAllowed, false, false, DnsResponseCode.NoError, request.Question, answers, null, null, _dnsServer.UdpPayloadSize, EDnsHeaderFlags.None, options));