Implemented DNS Cookie based rate limiter with performance tricks, i.e. drop without responding beyond a threshold

Author: zbalkan

DnsServerCore/Dns/DnsServer.cs

@@ -285,16 +285,27 @@ enum ServiceState
         long _cookieBadcookieSent;
         long _cookieClientOnly;
         long _cookieInvalidDropped;
+        long _cookieRateLimited;
 
         // Hard-coded initial policy for invalid DNS cookie abuse handling.
         // Iteration 2 can make these values configurable.
         const int COOKIE_FAILURE_WINDOW_SECONDS = 60;
         const int COOKIE_FAILURE_SOFT_THRESHOLD = 20;
+        const int COOKIE_FAILURE_BADCOOKIE_SLIP_THRESHOLD = 40;
+        const int COOKIE_FAILURE_BADCOOKIE_SLIP_FACTOR = 4; // send BADCOOKIE for every Nth request after slip threshold
         const int COOKIE_FAILURE_HARD_THRESHOLD = 100;
         const int COOKIE_FAILURE_RETENTION_SECONDS = 300;
+        const int COOKIE_BOOTSTRAP_WINDOW_SECONDS = 10;
+        const int COOKIE_BOOTSTRAP_HARD_THRESHOLD = 30;
+        const int COOKIE_BOOTSTRAP_REFILL_RATE_PER_SECOND = 3; // 30 tokens per 10 seconds
+        const int COOKIE_BOOTSTRAP_IDLE_EVICT_SECONDS = 120;
+        const int COOKIE_BOOTSTRAP_BUCKET_COUNT = 16384; // power-of-two for fast masking
+        const int COOKIE_BOOTSTRAP_MAX_PROBES = 8;
 
         readonly Dictionary<IPAddress, (long WindowStart, int Count)> _cookieFailureByClient = new();
         readonly object _cookieFailureLock = new();
+        readonly object _cookieBootstrapLock = new();
+        readonly CookieBootstrapBucket[] _cookieBootstrapBuckets = new CookieBootstrapBucket[COOKIE_BOOTSTRAP_BUCKET_COUNT];
 
         #endregion
 
@@ -1582,6 +1593,14 @@ private string ConvertToAbsolutePath(string path)
 
         #region cookie
 
+        private struct CookieBootstrapBucket
+        {
+            public ulong ClientHash;
+            public uint LastRefillSecond;
+            public uint LastSeenSecond;
+            public int Tokens;
+        }
+
         private void InitDnsCookies()
         {
             lock (_saveLock)
@@ -1596,6 +1615,10 @@ private void InitDnsCookies()
 
                     lock (_cookieFailureLock)
                         _cookieFailureByClient.Clear();
+                    lock (_cookieBootstrapLock)
+                    {
+                        Array.Clear(_cookieBootstrapBuckets, 0, _cookieBootstrapBuckets.Length);
+                    }
 
                     return;
                 }
@@ -1988,7 +2011,124 @@ private bool ShouldDropInvalidCookie(IPAddress clientAddress, out int failureCou
                 }
             }
 
-            return failureCount > COOKIE_FAILURE_HARD_THRESHOLD;
+            if (failureCount > COOKIE_FAILURE_HARD_THRESHOLD)
+                return true;
+
+            if (failureCount > COOKIE_FAILURE_BADCOOKIE_SLIP_THRESHOLD)
+                return (failureCount % COOKIE_FAILURE_BADCOOKIE_SLIP_FACTOR) != 0;
+
+            return false;
+        }
+
+        private static ulong GetCookieBootstrapClientHash(IPAddress clientAddress)
+        {
+            Span<byte> addressBytes = stackalloc byte[16];
+            if (!clientAddress.TryWriteBytes(addressBytes, out int written))
+                written = 0;
+
+            int offset = 0;
+            if (written == 16)
+            {
+                bool isV4Mapped =
+                    addressBytes[0] == 0 && addressBytes[1] == 0 && addressBytes[2] == 0 && addressBytes[3] == 0 &&
+                    addressBytes[4] == 0 && addressBytes[5] == 0 && addressBytes[6] == 0 && addressBytes[7] == 0 &&
+                    addressBytes[8] == 0 && addressBytes[9] == 0 &&
+                    addressBytes[10] == 0xff && addressBytes[11] == 0xff;
+
+                if (isV4Mapped)
+                {
+                    offset = 12;
+                    written = 4;
+                }
+            }
+
+            // FNV-1a 64-bit.
+            ulong hash = 1469598103934665603UL;
+            for (int i = 0; i < written; i++)
+            {
+                hash ^= addressBytes[offset + i];
+                hash *= 1099511628211UL;
+            }
+
+            // Reserve zero as "empty bucket" marker.
+            return hash == 0 ? 1UL : hash;
+        }
+
+        private bool ShouldRateLimitCookieBootstrap(IPAddress clientAddress, out int requestCount)
+        {
+            uint now = unchecked((uint)DateTimeOffset.UtcNow.ToUnixTimeSeconds());
+            ulong clientHash = GetCookieBootstrapClientHash(clientAddress);
+            int mask = COOKIE_BOOTSTRAP_BUCKET_COUNT - 1;
+            int start = (int)(clientHash & (uint)mask);
+
+            lock (_cookieBootstrapLock)
+            {
+                int selectedIndex = -1;
+                int staleIndex = -1;
+
+                for (int probe = 0; probe < COOKIE_BOOTSTRAP_MAX_PROBES; probe++)
+                {
+                    int index = (start + probe) & mask;
+                    ref CookieBootstrapBucket bucket = ref _cookieBootstrapBuckets[index];
+
+                    if (bucket.ClientHash == clientHash)
+                    {
+                        selectedIndex = index;
+                        break;
+                    }
+
+                    if (bucket.ClientHash == 0)
+                    {
+                        selectedIndex = index;
+                        break;
+                    }
+
+                    if ((now - bucket.LastSeenSecond) >= COOKIE_BOOTSTRAP_IDLE_EVICT_SECONDS && staleIndex < 0)
+                        staleIndex = index;
+                }
+
+                if (selectedIndex < 0)
+                    selectedIndex = staleIndex;
+
+                if (selectedIndex < 0)
+                {
+                    requestCount = int.MaxValue; // table pressure fail-closed
+                    return true;
+                }
+
+                ref CookieBootstrapBucket state = ref _cookieBootstrapBuckets[selectedIndex];
+
+                if (state.ClientHash != clientHash)
+                {
+                    state.ClientHash = clientHash;
+                    state.LastRefillSecond = now;
+                    state.LastSeenSecond = now;
+                    state.Tokens = COOKIE_BOOTSTRAP_HARD_THRESHOLD;
+                }
+                else
+                {
+                    uint elapsed = now - state.LastRefillSecond;
+                    if (elapsed > 0)
+                    {
+                        int refill = (int)Math.Min(int.MaxValue, elapsed * (uint)COOKIE_BOOTSTRAP_REFILL_RATE_PER_SECOND);
+                        state.Tokens = Math.Min(COOKIE_BOOTSTRAP_HARD_THRESHOLD, state.Tokens + refill);
+                        state.LastRefillSecond = now;
+                    }
+
+                    state.LastSeenSecond = now;
+                }
+
+                if (state.Tokens <= 0)
+                {
+                    requestCount = COOKIE_BOOTSTRAP_HARD_THRESHOLD + 1;
+                    return true;
+                }
+
+                state.Tokens--;
+            }
+
+            requestCount = 0;
+            return false;
         }
 
         private DnsDatagram HandleInvalidCookieRequest(
@@ -2011,6 +2151,11 @@ private DnsDatagram HandleInvalidCookieRequest(
             {
                 Interlocked.Increment(ref _cookieInvalidDropped);
 
+                if (failuresInWindow == (COOKIE_FAILURE_BADCOOKIE_SLIP_THRESHOLD + 1))
+                {
+                    _log.Write($"Client '{remoteEP.Address}' crossed DNS cookie invalid slip threshold ({COOKIE_FAILURE_BADCOOKIE_SLIP_THRESHOLD}/{COOKIE_FAILURE_WINDOW_SECONDS}s); BADCOOKIE responses are now throttled to 1/{COOKIE_FAILURE_BADCOOKIE_SLIP_FACTOR}.");
+                }
+
                 if (failuresInWindow == (COOKIE_FAILURE_HARD_THRESHOLD + 1))
                 {
                     _log.Write($"Client '{remoteEP.Address}' exceeded DNS cookie invalid hard threshold ({COOKIE_FAILURE_HARD_THRESHOLD}/{COOKIE_FAILURE_WINDOW_SECONDS}s); dropping invalid-cookie requests.");
@@ -2956,6 +3101,18 @@ private async Task<DnsDatagram> ProcessRequestAsync(DnsDatagram request, IPEndPo
                 if (requestCookie == null)
                 {
                     Interlocked.Increment(ref _cookieMissing);
+
+                    if (ShouldRateLimitCookieBootstrap(remoteEP.Address, out int requestsInWindow))
+                    {
+                        Interlocked.Increment(ref _cookieRateLimited);
+
+                        if (requestsInWindow == (COOKIE_BOOTSTRAP_HARD_THRESHOLD + 1))
+                        {
+                            _log.Write($"Client '{remoteEP.Address}' exceeded DNS cookie bootstrap threshold ({COOKIE_BOOTSTRAP_HARD_THRESHOLD}/{COOKIE_BOOTSTRAP_WINDOW_SECONDS}s); dropping UDP requests without cookie.");
+                        }
+
+                        return null;
+                    }
                 }
                 else
                 {
@@ -3000,6 +3157,18 @@ private async Task<DnsDatagram> ProcessRequestAsync(DnsDatagram request, IPEndPo
                     if (scLen == 0)
                     {
                         Interlocked.Increment(ref _cookieClientOnly);
+
+                        if (ShouldRateLimitCookieBootstrap(remoteEP.Address, out int requestsInWindow))
+                        {
+                            Interlocked.Increment(ref _cookieRateLimited);
+
+                            if (requestsInWindow == (COOKIE_BOOTSTRAP_HARD_THRESHOLD + 1))
+                            {
+                                _log.Write($"Client '{remoteEP.Address}' exceeded DNS cookie bootstrap threshold ({COOKIE_BOOTSTRAP_HARD_THRESHOLD}/{COOKIE_BOOTSTRAP_WINDOW_SECONDS}s); dropping client-cookie-only requests.");
+                            }
+
+                            return null;
+                        }
                     }
                     else
                     {