128 lines
5.8 KiB
C#
128 lines
5.8 KiB
C#
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 int CacheSize => _info.Count;
|
|
|
|
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;
|
|
|
|
// We can hit this multiple times if many requests are in flight before a real one gets "back", so we still
|
|
// decrement the remaining request count, this basically "blacklists" the channel given continuous spam until *one* of the requests come back with new rate limit headers
|
|
info.remaining--;
|
|
|
|
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());
|
|
}
|
|
}
|
|
|
|
public void GarbageCollect()
|
|
{
|
|
_logger.Information("Garbage-collecting webhook rate limit buckets...");
|
|
|
|
var collected = 0;
|
|
foreach (var channel in _info.Keys)
|
|
{
|
|
if (!_info.TryGetValue(channel, out var info)) continue;
|
|
|
|
// Remove all keys that expired more than an hour ago (and of course, haven't been reset)
|
|
if (info.resetTime < SystemClock.Instance.GetCurrentInstant() - Duration.FromHours(1))
|
|
if (_info.TryRemove(channel, out _)) collected++;
|
|
}
|
|
|
|
_logger.Information("Garbage-collected {ChannelCount} channels from the webhook rate limit buckets.", collected);
|
|
}
|
|
|
|
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;
|
|
}
|
|
}
|
|
} |