liz/PluralKit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial commit, basic proxying working

Browse Source
This commit is contained in:
Ske
2020-12-22 13:15:26 +01:00
parent c3f6becea4
commit a6fbd869be
109 changed files with 3539 additions and 359 deletions
Download Patch File Download Diff File Expand all files Collapse all files

95
PluralKit.Bot/Services/WebhookCacheService.cs
View File

@@ -1,14 +1,15 @@
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Net.Http;
using System.Threading.Tasks;
using App.Metrics;
using DSharpPlus;
using DSharpPlus.Entities;
using Myriad.Gateway;
using Myriad.Rest;
using Myriad.Rest.Types.Requests;
using Myriad.Types;
using Serilog;
@@ -17,38 +18,32 @@ namespace PluralKit.Bot
public class WebhookCacheService
{
public static readonly string WebhookName = "PluralKit Proxy Webhook";
private readonly DiscordShardedClient _client;
private readonly ConcurrentDictionary<ulong, Lazy<Task<DiscordWebhook>>> _webhooks;
private readonly DiscordApiClient _rest;
private readonly ConcurrentDictionary<ulong, Lazy<Task<Webhook>>> _webhooks;
private readonly IMetrics _metrics;
private readonly ILogger _logger;
private readonly Cluster _cluster;
public WebhookCacheService(DiscordShardedClient client, ILogger logger, IMetrics metrics)
public WebhookCacheService(ILogger logger, IMetrics metrics, DiscordApiClient rest, Cluster cluster)
{
_client = client;
_metrics = metrics;
_rest = rest;
_cluster = cluster;
_logger = logger.ForContext<WebhookCacheService>();
_webhooks = new ConcurrentDictionary<ulong, Lazy<Task<DiscordWebhook>>>();
_webhooks = new ConcurrentDictionary<ulong, Lazy<Task<Webhook>>>();
}
public async Task<DiscordWebhook> GetWebhook(DiscordClient client, ulong channelId)
{
var channel = await client.GetChannel(channelId);
if (channel == null) return null;
if (channel.Type == ChannelType.Text) return null;
return await GetWebhook(channel);
}
public async Task<DiscordWebhook> GetWebhook(DiscordChannel channel)
public async Task<Webhook> GetWebhook(ulong channelId)
{
// We cache the webhook through a Lazy<Task<T>>, this way we make sure to only create one webhook per channel
// If the webhook is requested twice before it's actually been found, the Lazy<T> wrapper will stop the
// webhook from being created twice.
Lazy<Task<DiscordWebhook>> GetWebhookTaskInner()
Lazy<Task<Webhook>> GetWebhookTaskInner()
{
Task<DiscordWebhook> Factory() => GetOrCreateWebhook(channel);
return _webhooks.GetOrAdd(channel.Id, new Lazy<Task<DiscordWebhook>>(Factory));
Task<Webhook> Factory() => GetOrCreateWebhook(channelId);
return _webhooks.GetOrAdd(channelId, new Lazy<Task<Webhook>>(Factory));
}
var lazyWebhookValue = GetWebhookTaskInner();
@@ -57,36 +52,38 @@ namespace PluralKit.Bot
// although, keep in mind this block gets hit the call *after* the task failed (since we only await it below)
if (lazyWebhookValue.IsValueCreated && lazyWebhookValue.Value.IsFaulted)
{
_logger.Warning(lazyWebhookValue.Value.Exception, "Cached webhook task for {Channel} faulted with below exception", channel.Id);
_logger.Warning(lazyWebhookValue.Value.Exception, "Cached webhook task for {Channel} faulted with below exception", channelId);
// Specifically don't recurse here so we don't infinite-loop - if this one errors too, it'll "stick"
// until next time this function gets hit (which is okay, probably).
_webhooks.TryRemove(channel.Id, out _);
_webhooks.TryRemove(channelId, out _);
lazyWebhookValue = GetWebhookTaskInner();
}
// It's possible to "move" a webhook to a different channel after creation
// Here, we ensure it's actually still pointing towards the proper channel, and if not, wipe and refetch one.
var webhook = await lazyWebhookValue.Value;
if (webhook.ChannelId != channel.Id && webhook.ChannelId != 0) return await InvalidateAndRefreshWebhook(channel, webhook);
if (webhook.ChannelId != channelId && webhook.ChannelId != 0)
return await InvalidateAndRefreshWebhook(channelId, webhook);
return webhook;
}
public async Task<DiscordWebhook> InvalidateAndRefreshWebhook(DiscordChannel channel, DiscordWebhook webhook)
public async Task<Webhook> InvalidateAndRefreshWebhook(ulong channelId, Webhook webhook)
{
// note: webhook.ChannelId may not be the same as channelId >.>
_logger.Information("Refreshing webhook for channel {Channel}", webhook.ChannelId);
_webhooks.TryRemove(webhook.ChannelId, out _);
return await GetWebhook(channel);
return await GetWebhook(channelId);
}
private async Task<DiscordWebhook> GetOrCreateWebhook(DiscordChannel channel)
private async Task<Webhook?> GetOrCreateWebhook(ulong channelId)
{
_logger.Debug("Webhook for channel {Channel} not found in cache, trying to fetch", channel.Id);
_logger.Debug("Webhook for channel {Channel} not found in cache, trying to fetch", channelId);
_metrics.Measure.Meter.Mark(BotMetrics.WebhookCacheMisses);
_logger.Debug("Finding webhook for channel {Channel}", channel.Id);
var webhooks = await FetchChannelWebhooks(channel);
_logger.Debug("Finding webhook for channel {Channel}", channelId);
var webhooks = await FetchChannelWebhooks(channelId);
// If the channel has a webhook created by PK, just return that one
var ourWebhook = webhooks.FirstOrDefault(IsWebhookMine);
@@ -95,17 +92,17 @@ namespace PluralKit.Bot
// We don't have one, so we gotta create a new one
// but first, make sure we haven't hit the webhook cap yet...
if (webhooks.Count >= 10)
if (webhooks.Length >= 10)
throw new PKError("This channel has the maximum amount of possible webhooks (10) already created. A server admin must delete one or more webhooks so PluralKit can create one for proxying.");
return await DoCreateWebhook(channel);
return await DoCreateWebhook(channelId);
}
private async Task<IReadOnlyList<DiscordWebhook>> FetchChannelWebhooks(DiscordChannel channel)
private async Task<Webhook[]> FetchChannelWebhooks(ulong channelId)
{
try
{
return await channel.GetWebhooksAsync();
return await _rest.GetChannelWebhooks(channelId);
}
catch (HttpRequestException e)
{
@@ -113,33 +110,17 @@ namespace PluralKit.Bot
// This happens sometimes when Discord returns a malformed request for the webhook list
// Nothing we can do than just assume that none exist.
return new DiscordWebhook[0];
return new Webhook[0];
}
}
private async Task<DiscordWebhook> FindExistingWebhook(DiscordChannel channel)
private async Task<Webhook> DoCreateWebhook(ulong channelId)
{
_logger.Debug("Finding webhook for channel {Channel}", channel.Id);
try
{
return (await channel.GetWebhooksAsync()).FirstOrDefault(IsWebhookMine);
}
catch (HttpRequestException e)
{
_logger.Warning(e, "Error occurred while fetching webhook list");
// This happens sometimes when Discord returns a malformed request for the webhook list
// Nothing we can do than just assume that none exist and return null.
return null;
}
_logger.Information("Creating new webhook for channel {Channel}", channelId);
return await _rest.CreateWebhook(channelId, new CreateWebhookRequest(WebhookName));
}
private Task<DiscordWebhook> DoCreateWebhook(DiscordChannel channel)
{
_logger.Information("Creating new webhook for channel {Channel}", channel.Id);
return channel.CreateWebhookAsync(WebhookName);
}
private bool IsWebhookMine(DiscordWebhook arg) => arg.User.Id == _client.CurrentUser.Id && arg.Name == WebhookName;
private bool IsWebhookMine(Webhook arg) => arg.User?.Id == _cluster.User?.Id && arg.Name == WebhookName;
public int CacheSize => _webhooks.Count;
}