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

61
PluralKit.Bot/Services/LogChannelService.cs
View File

@@ -2,8 +2,9 @@ using System.Threading.Tasks;
using Dapper;
using DSharpPlus;
using DSharpPlus.Entities;
using Myriad.Cache;
using Myriad.Rest;
using Myriad.Types;
using PluralKit.Core;
@@ -15,56 +16,62 @@ namespace PluralKit.Bot {
private readonly IDatabase _db;
private readonly ModelRepository _repo;
private readonly ILogger _logger;
private readonly IDiscordCache _cache;
private readonly DiscordApiClient _rest;
public LogChannelService(EmbedService embed, ILogger logger, IDatabase db, ModelRepository repo)
public LogChannelService(EmbedService embed, ILogger logger, IDatabase db, ModelRepository repo, IDiscordCache cache, DiscordApiClient rest)
{
_embed = embed;
_db = db;
_repo = repo;
_cache = cache;
_rest = rest;
_logger = logger.ForContext<LogChannelService>();
}
public async ValueTask LogMessage(DiscordClient client, MessageContext ctx, ProxyMatch proxy, DiscordMessage trigger, ulong hookMessage)
public async ValueTask LogMessage(MessageContext ctx, ProxyMatch proxy, Message trigger, ulong hookMessage)
{
if (ctx.SystemId == null || ctx.LogChannel == null || ctx.InLogBlacklist) return;
// Find log channel and check if valid
var logChannel = await FindLogChannel(client, trigger.Channel.GuildId, ctx.LogChannel.Value);
if (logChannel == null || logChannel.Type != ChannelType.Text) return;
var logChannel = await FindLogChannel(trigger.GuildId!.Value, ctx.LogChannel.Value);
if (logChannel == null || logChannel.Type != Channel.ChannelType.GuildText) return;
// Check bot permissions
if (!logChannel.BotHasAllPermissions(Permissions.SendMessages | Permissions.EmbedLinks))
{
_logger.Information(
"Does not have permission to proxy log, ignoring (channel: {ChannelId}, guild: {GuildId}, bot permissions: {BotPermissions})",
ctx.LogChannel.Value, trigger.Channel.GuildId, trigger.Channel.BotPermissions());
return;
}
// if (!logChannel.BotHasAllPermissions(Permissions.SendMessages | Permissions.EmbedLinks))
// {
// _logger.Information(
// "Does not have permission to proxy log, ignoring (channel: {ChannelId}, guild: {GuildId}, bot permissions: {BotPermissions})",
// ctx.LogChannel.Value, trigger.GuildId!.Value, trigger.Channel.BotPermissions());
// return;
// }
//
// Send embed!
await using var conn = await _db.Obtain();
var embed = _embed.CreateLoggedMessageEmbed(await _repo.GetSystem(conn, ctx.SystemId.Value),
await _repo.GetMember(conn, proxy.Member.Id), hookMessage, trigger.Id, trigger.Author, proxy.Content,
trigger.Channel);
var url = $"https://discord.com/channels/{trigger.Channel.GuildId}/{trigger.ChannelId}/{hookMessage}";
await logChannel.SendMessageFixedAsync(content: url, embed: embed);
// TODO: fix?
// await using var conn = await _db.Obtain();
// var embed = _embed.CreateLoggedMessageEmbed(await _repo.GetSystem(conn, ctx.SystemId.Value),
// await _repo.GetMember(conn, proxy.Member.Id), hookMessage, trigger.Id, trigger.Author, proxy.Content,
// trigger.Channel);
// var url = $"https://discord.com/channels/{trigger.Channel.GuildId}/{trigger.ChannelId}/{hookMessage}";
// await logChannel.SendMessageFixedAsync(content: url, embed: embed);
}
private async Task<DiscordChannel> FindLogChannel(DiscordClient client, ulong guild, ulong channel)
private async Task<Channel?> FindLogChannel(ulong guildId, ulong channelId)
{
// MUST use this client here, otherwise we get strange cache issues where the guild doesn't exist... >.>
var obj = await client.GetChannel(channel);
// TODO: fetch it directly on cache miss?
var channel = await _cache.GetChannel(channelId);
if (obj == null)
if (channel == null)
{
// Channel doesn't exist or we don't have permission to access it, let's remove it from the database too
_logger.Warning("Attempted to fetch missing log channel {LogChannel} for guild {Guild}, removing from database", channel, guild);
_logger.Warning("Attempted to fetch missing log channel {LogChannel} for guild {Guild}, removing from database", channelId, guildId);
await using var conn = await _db.Obtain();
await conn.ExecuteAsync("update servers set log_channel = null where id = @Guild",
new {Guild = guild});
new {Guild = guildId});
}
return obj;
return channel;
}
}
}

10
PluralKit.Bot/Services/LoggerCleanService.cs
View File

@@ -4,11 +4,10 @@ using System.Linq;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
using Dapper;
using DSharpPlus;
using DSharpPlus.Entities;
using DSharpPlus.Exceptions;
using Myriad.Types;
using PluralKit.Core;
@@ -68,8 +67,10 @@ namespace PluralKit.Bot
public ICollection<LoggerBot> Bots => _bots.Values;
public async ValueTask HandleLoggerBotCleanup(DiscordMessage msg)
public async ValueTask HandleLoggerBotCleanup(Message msg)
{
// TODO: fix!!
/*
if (msg.Channel.Type != ChannelType.Text) return;
if (!msg.Channel.BotHasAllPermissions(Permissions.ManageMessages)) return;
@@ -130,6 +131,7 @@ namespace PluralKit.Bot
// The only thing I can think of that'd cause this are the DeleteAsync() calls which 404 when
// the message doesn't exist anyway - so should be safe to just ignore it, right?
}
*/
}
private static ulong? ExtractAuttaja(DiscordMessage msg)

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;
}

123
PluralKit.Bot/Services/WebhookExecutorService.cs
View File

@@ -1,19 +1,20 @@
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.Http;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
using App.Metrics;
using DSharpPlus.Entities;
using DSharpPlus.Exceptions;
using Humanizer;
using Myriad.Cache;
using Myriad.Rest;
using Myriad.Rest.Types;
using Myriad.Rest.Types.Requests;
using Myriad.Types;
using Newtonsoft.Json;
using Newtonsoft.Json.Serialization;
using Serilog;
@@ -26,64 +27,84 @@ namespace PluralKit.Bot
// Exceptions for control flow? don't mind if I do
// TODO: rewrite both of these as a normal exceptional return value (0?) in case of error to be discarded by caller
}
public record ProxyRequest
{
public ulong GuildId { get; init; }
public ulong ChannelId { get; init; }
public string Name { get; init; }
public string? AvatarUrl { get; init; }
public string? Content { get; init; }
public Message.Attachment[] Attachments { get; init; }
public Embed[] Embeds { get; init; }
public bool AllowEveryone { get; init; }
}
public class WebhookExecutorService
{
private readonly IDiscordCache _cache;
private readonly WebhookCacheService _webhookCache;
private readonly DiscordApiClient _rest;
private readonly ILogger _logger;
private readonly IMetrics _metrics;
private readonly HttpClient _client;
public WebhookExecutorService(IMetrics metrics, WebhookCacheService webhookCache, ILogger logger, HttpClient client)
public WebhookExecutorService(IMetrics metrics, WebhookCacheService webhookCache, ILogger logger, HttpClient client, IDiscordCache cache, DiscordApiClient rest)
{
_metrics = metrics;
_webhookCache = webhookCache;
_client = client;
_cache = cache;
_rest = rest;
_logger = logger.ForContext<WebhookExecutorService>();
}
public async Task<DiscordMessage> ExecuteWebhook(DiscordChannel channel, string name, string avatarUrl, string content, IReadOnlyList<DiscordAttachment> attachments, IReadOnlyList<DiscordEmbed> embeds, bool allowEveryone)
public async Task<Message> ExecuteWebhook(ProxyRequest req)
{
_logger.Verbose("Invoking webhook in channel {Channel}", channel.Id);
_logger.Verbose("Invoking webhook in channel {Channel}", req.ChannelId);
// Get a webhook, execute it
var webhook = await _webhookCache.GetWebhook(channel);
var webhookMessage = await ExecuteWebhookInner(channel, webhook, name, avatarUrl, content, attachments, embeds, allowEveryone);
var webhook = await _webhookCache.GetWebhook(req.ChannelId);
var webhookMessage = await ExecuteWebhookInner(webhook, req);
// Log the relevant metrics
_metrics.Measure.Meter.Mark(BotMetrics.MessagesProxied);
_logger.Information("Invoked webhook {Webhook} in channel {Channel}", webhook.Id,
channel.Id);
req.ChannelId);
return webhookMessage;
}
private async Task<DiscordMessage> ExecuteWebhookInner(
DiscordChannel channel, DiscordWebhook webhook, string name, string avatarUrl, string content,
IReadOnlyList<DiscordAttachment> attachments, IReadOnlyList<DiscordEmbed> embeds, bool allowEveryone, bool hasRetried = false)
private async Task<Message> ExecuteWebhookInner(Webhook webhook, ProxyRequest req, bool hasRetried = false)
{
content = content.Truncate(2000);
var guild = await _cache.GetGuild(req.GuildId)!;
var content = req.Content.Truncate(2000);
var dwb = new DiscordWebhookBuilder();
dwb.WithUsername(FixClyde(name).Truncate(80));
dwb.WithContent(content);
dwb.AddMentions(content.ParseAllMentions(allowEveryone, channel.Guild));
if (!string.IsNullOrWhiteSpace(avatarUrl))
dwb.WithAvatarUrl(avatarUrl);
dwb.AddEmbeds(embeds);
var webhookReq = new ExecuteWebhookRequest
{
Username = FixClyde(req.Name).Truncate(80),
Content = content,
AllowedMentions = null, // todo
AvatarUrl = !string.IsNullOrWhiteSpace(req.AvatarUrl) ? req.AvatarUrl : null,
Embeds = req.Embeds
};
var attachmentChunks = ChunkAttachmentsOrThrow(attachments, 8 * 1024 * 1024);
// dwb.AddMentions(content.ParseAllMentions(guild, req.AllowEveryone));
MultipartFile[] files = null;
var attachmentChunks = ChunkAttachmentsOrThrow(req.Attachments, 8 * 1024 * 1024);
if (attachmentChunks.Count > 0)
{
_logger.Information("Invoking webhook with {AttachmentCount} attachments totalling {AttachmentSize} MiB in {AttachmentChunks} chunks", attachments.Count, attachments.Select(a => a.FileSize).Sum() / 1024 / 1024, attachmentChunks.Count);
await AddAttachmentsToBuilder(dwb, attachmentChunks[0]);
_logger.Information("Invoking webhook with {AttachmentCount} attachments totalling {AttachmentSize} MiB in {AttachmentChunks} chunks",
req.Attachments.Length, req.Attachments.Select(a => a.Size).Sum() / 1024 / 1024, attachmentChunks.Count);
files = await GetAttachmentFiles(attachmentChunks[0]);
}
DiscordMessage webhookMessage;
Message webhookMessage;
using (_metrics.Measure.Timer.Time(BotMetrics.WebhookResponseTime)) {
try
{
webhookMessage = await webhook.ExecuteAsync(dwb);
webhookMessage = await _rest.ExecuteWebhook(webhook.Id, webhook.Token, webhookReq, files);
}
catch (JsonReaderException)
{
@@ -91,17 +112,16 @@ namespace PluralKit.Bot
// Nothing we can do about this - happens sometimes under server load, so just drop the message and give up
throw new WebhookExecutionErrorOnDiscordsEnd();
}
catch (NotFoundException e)
catch (Myriad.Rest.Exceptions.NotFoundException e)
{
var errorText = e.WebResponse?.Response;
if (errorText != null && errorText.Contains("10015") && !hasRetried)
if (e.ErrorCode == 10015 && !hasRetried)
{
// Error 10015 = "Unknown Webhook" - this likely means the webhook was deleted
// but is still in our cache. Invalidate, refresh, try again
_logger.Warning("Error invoking webhook {Webhook} in channel {Channel}", webhook.Id, webhook.ChannelId);
var newWebhook = await _webhookCache.InvalidateAndRefreshWebhook(channel, webhook);
return await ExecuteWebhookInner(channel, newWebhook, name, avatarUrl, content, attachments, embeds, allowEveryone, hasRetried: true);
var newWebhook = await _webhookCache.InvalidateAndRefreshWebhook(req.ChannelId, webhook);
return await ExecuteWebhookInner(newWebhook, req, hasRetried: true);
}
throw;
@@ -109,53 +129,50 @@ namespace PluralKit.Bot
}
// We don't care about whether the sending succeeds, and we don't want to *wait* for it, so we just fork it off
var _ = TrySendRemainingAttachments(webhook, name, avatarUrl, attachmentChunks);
var _ = TrySendRemainingAttachments(webhook, req.Name, req.AvatarUrl, attachmentChunks);
return webhookMessage;
}
private async Task TrySendRemainingAttachments(DiscordWebhook webhook, string name, string avatarUrl, IReadOnlyList<IReadOnlyCollection<DiscordAttachment>> attachmentChunks)
private async Task TrySendRemainingAttachments(Webhook webhook, string name, string avatarUrl, IReadOnlyList<IReadOnlyCollection<Message.Attachment>> attachmentChunks)
{
if (attachmentChunks.Count <= 1) return;
for (var i = 1; i < attachmentChunks.Count; i++)
{
var dwb = new DiscordWebhookBuilder();
if (avatarUrl != null) dwb.WithAvatarUrl(avatarUrl);
dwb.WithUsername(name);
await AddAttachmentsToBuilder(dwb, attachmentChunks[i]);
await webhook.ExecuteAsync(dwb);
var files = await GetAttachmentFiles(attachmentChunks[i]);
var req = new ExecuteWebhookRequest {Username = name, AvatarUrl = avatarUrl};
await _rest.ExecuteWebhook(webhook.Id, webhook.Token!, req, files);
}
}
private async Task AddAttachmentsToBuilder(DiscordWebhookBuilder dwb, IReadOnlyCollection<DiscordAttachment> attachments)
private async Task<MultipartFile[]> GetAttachmentFiles(IReadOnlyCollection<Message.Attachment> attachments)
{
async Task<(DiscordAttachment, Stream)> GetStream(DiscordAttachment attachment)
async Task<MultipartFile> GetStream(Message.Attachment attachment)
{
var attachmentResponse = await _client.GetAsync(attachment.Url, HttpCompletionOption.ResponseHeadersRead);
return (attachment, await attachmentResponse.Content.ReadAsStreamAsync());
return new(attachment.Filename, await attachmentResponse.Content.ReadAsStreamAsync());
}
foreach (var (attachment, attachmentStream) in await Task.WhenAll(attachments.Select(GetStream)))
dwb.AddFile(attachment.FileName, attachmentStream);
return await Task.WhenAll(attachments.Select(GetStream));
}
private IReadOnlyList<IReadOnlyCollection<DiscordAttachment>> ChunkAttachmentsOrThrow(
IReadOnlyList<DiscordAttachment> attachments, int sizeThreshold)
private IReadOnlyList<IReadOnlyCollection<Message.Attachment>> ChunkAttachmentsOrThrow(
IReadOnlyList<Message.Attachment> attachments, int sizeThreshold)
{
// Splits a list of attachments into "chunks" of at most 8MB each
// If any individual attachment is larger than 8MB, will throw an error
var chunks = new List<List<DiscordAttachment>>();
var list = new List<DiscordAttachment>();
var chunks = new List<List<Message.Attachment>>();
var list = new List<Message.Attachment>();
foreach (var attachment in attachments)
{
if (attachment.FileSize >= sizeThreshold) throw Errors.AttachmentTooLarge;
if (attachment.Size >= sizeThreshold) throw Errors.AttachmentTooLarge;
if (list.Sum(a => a.FileSize) + attachment.FileSize >= sizeThreshold)
if (list.Sum(a => a.Size) + attachment.Size >= sizeThreshold)
{
chunks.Add(list);
list = new List<DiscordAttachment>();
list = new List<Message.Attachment>();
}
list.Add(attachment);