Add proper webhook rate limit support

2020-03-27 00:01:42 +01:00
parent 12d29eba44
commit fa70df8f98
4 changed files with 125 additions and 6 deletions
--- a/PluralKit.Bot/Services/WebhookRateLimitService.cs
+++ b/PluralKit.Bot/Services/WebhookRateLimitService.cs
@@ -0,0 +1,104 @@
+using System;
+using System.Collections.Concurrent;
+using System.Globalization;
+using System.Linq;
+using System.Net;
+using System.Net.Http;
+
+using Discord;
+
+using NodaTime;
+
+using Serilog;
+
+namespace PluralKit.Bot
+{
+    // Simplified rate limit handler for webhooks only, disregards most bucket functionality since scope is limited and just denies requests if too fast.
+    public class WebhookRateLimitService
+    {
+        private ILogger _logger;
+        private ConcurrentDictionary<ulong, WebhookRateLimitInfo> _info = new ConcurrentDictionary<ulong, WebhookRateLimitInfo>();
+
+        public WebhookRateLimitService(ILogger logger)
+        {
+            _logger = logger.ForContext<WebhookRateLimitService>();
+        }
+
+        public bool TryExecuteWebhook(IWebhook webhook)
+        {
+            // If we have nothing saved, just allow it (we'll save something once the response returns)
+            if (!_info.TryGetValue(webhook.Id, out var info)) return true;
+            
+            // If we're past the reset time, allow the request and update the bucket limit
+            if (SystemClock.Instance.GetCurrentInstant() > info.resetTime)
+            {
+                if (!info.hasResetTimeExpired)
+                    info.remaining = info.maxLimit;
+                info.hasResetTimeExpired = true;
+                return true;
+            }
+
+            // If we don't have any more requests left, deny the request
+            if (info.remaining == 0)
+            {
+                _logger.Debug("Rate limit bucket for {Webhook} out of requests, denying request", webhook.Id);
+                return false;
+            }
+            
+            // Otherwise, decrement the request count and allow the request
+            info.remaining--;
+            return true;
+        }
+
+        public void UpdateRateLimitInfo(IWebhook webhook, HttpResponseMessage response)
+        {
+            var info = _info.GetOrAdd(webhook.Id, _ => new WebhookRateLimitInfo());
+
+            if (int.TryParse(GetHeader(response, "X-RateLimit-Limit"), out var limit))
+                info.maxLimit = limit;
+
+            // Max "safe" is way above UNIX timestamp values, and we get fractional seconds, hence the double
+            // but need culture/format specifiers to get around Some Locales (cough, my local PC) having different settings for decimal point...
+            // We also use Reset-After to avoid issues with clock desync between us and Discord's server, this way it's all relative (plus latency errors work in our favor)    
+            if (double.TryParse(GetHeader(response, "X-RateLimit-Reset-After"), NumberStyles.Float, CultureInfo.InvariantCulture, out var resetTimestampDelta))
+            {
+                var resetTime = SystemClock.Instance.GetCurrentInstant() + Duration.FromSeconds(resetTimestampDelta);
+                if (resetTime > info.resetTime)
+                {
+                    // Set to the *latest* reset value we have (for safety), since we rely on relative times this can jitter a bit
+                    info.resetTime = resetTime;
+                    info.hasResetTimeExpired = false;
+                }
+            }
+            
+            if (int.TryParse(GetHeader(response, "X-RateLimit-Remaining"), out var remainingRequests))
+                // Overwrite a negative "we don't know" value with whatever we just got
+                // Otherwise, *lower* remaining requests takes precedence
+                if (info.remaining < 0 || remainingRequests < info.remaining)
+                    info.remaining = remainingRequests;
+            
+            _logger.Debug("Updated rate limit information for {Webhook}, bucket has {RequestsRemaining} requests remaining, reset in {ResetTime}", webhook.Id, info.remaining, info.resetTime - SystemClock.Instance.GetCurrentInstant());
+            
+            if (response.StatusCode == HttpStatusCode.TooManyRequests)
+            {
+                // 429, we're *definitely* out of requests
+                info.remaining = 0;
+                _logger.Warning("Got 429 Too Many Requests when invoking webhook {Webhook}, next bucket reset in {ResetTime}", webhook.Id, info.resetTime - SystemClock.Instance.GetCurrentInstant());
+            }
+        }
+
+        private string GetHeader(HttpResponseMessage response, string key)
+        {
+            var firstPair = response.Headers.FirstOrDefault(pair => pair.Key.Equals(key, StringComparison.InvariantCultureIgnoreCase));
+            return firstPair.Value?.FirstOrDefault(); // If key is missing, default value is null
+        }
+
+        private class WebhookRateLimitInfo
+        {
+            public Instant resetTime;
+            public bool hasResetTimeExpired;
+            public int remaining = -1;
+            public int maxLimit = 0;
+        } 
+    }
+}