From 8157c6932e5d24b4134c7bee64c70f63ffa6e355 Mon Sep 17 00:00:00 2001 From: Astrid Date: Sun, 11 Feb 2024 03:53:46 +0100 Subject: [PATCH] feat: add support for external avatar hosting service (#614) --- PluralKit.Bot/BotConfig.cs | 1 + PluralKit.Bot/Commands/Groups.cs | 6 +- PluralKit.Bot/Commands/Member.cs | 6 +- PluralKit.Bot/Commands/MemberAvatar.cs | 6 +- PluralKit.Bot/Commands/MemberEdit.cs | 5 +- PluralKit.Bot/Commands/SystemEdit.cs | 7 +- PluralKit.Bot/Modules.cs | 1 + .../Services/AvatarHostingService.cs | 65 +++++++++++++++++++ PluralKit.Bot/Utils/AvatarUtils.cs | 2 + 9 files changed, 94 insertions(+), 5 deletions(-) create mode 100644 PluralKit.Bot/Services/AvatarHostingService.cs diff --git a/PluralKit.Bot/BotConfig.cs b/PluralKit.Bot/BotConfig.cs index 18c44b28..4621717e 100644 --- a/PluralKit.Bot/BotConfig.cs +++ b/PluralKit.Bot/BotConfig.cs @@ -25,6 +25,7 @@ public class BotConfig public string? RedisGatewayUrl { get; set; } public string? DiscordBaseUrl { get; set; } + public string? AvatarServiceUrl { get; set; } public bool DisableErrorReporting { get; set; } = false; diff --git a/PluralKit.Bot/Commands/Groups.cs b/PluralKit.Bot/Commands/Groups.cs index e3f539a2..dad4190c 100644 --- a/PluralKit.Bot/Commands/Groups.cs +++ b/PluralKit.Bot/Commands/Groups.cs @@ -21,13 +21,15 @@ public class Groups private readonly HttpClient _client; private readonly DispatchService _dispatch; private readonly EmbedService _embeds; + private readonly AvatarHostingService _avatarHosting; public Groups(EmbedService embeds, HttpClient client, - DispatchService dispatch) + DispatchService dispatch, AvatarHostingService avatarHosting) { _embeds = embeds; _client = client; _dispatch = dispatch; + _avatarHosting = avatarHosting; } public async Task CreateGroup(Context ctx) @@ -261,6 +263,7 @@ public class Groups { ctx.CheckOwnGroup(target); + img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Avatar, ctx.Author.Id); await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url); await ctx.Repository.UpdateGroup(target.Id, new GroupPatch { Icon = img.CleanUrl ?? img.Url }); @@ -326,6 +329,7 @@ public class Groups { ctx.CheckOwnGroup(target); + img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Banner, ctx.Author.Id); await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url, true); await ctx.Repository.UpdateGroup(target.Id, new GroupPatch { BannerImage = img.CleanUrl ?? img.Url }); diff --git a/PluralKit.Bot/Commands/Member.cs b/PluralKit.Bot/Commands/Member.cs index a34087c1..dad5e1cc 100644 --- a/PluralKit.Bot/Commands/Member.cs +++ b/PluralKit.Bot/Commands/Member.cs @@ -16,13 +16,15 @@ public class Member private readonly HttpClient _client; private readonly DispatchService _dispatch; private readonly EmbedService _embeds; + private readonly AvatarHostingService _avatarHosting; public Member(EmbedService embeds, HttpClient client, - DispatchService dispatch) + DispatchService dispatch, AvatarHostingService avatarHosting) { _embeds = embeds; _client = client; _dispatch = dispatch; + _avatarHosting = avatarHosting; } public async Task NewMember(Context ctx) @@ -78,6 +80,8 @@ public class Member uriBuilder.Query = ""; img.CleanUrl = uriBuilder.Uri.AbsoluteUri; + img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Avatar, ctx.Author.Id); + await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url); await ctx.Repository.UpdateMember(member.Id, new MemberPatch { AvatarUrl = img.CleanUrl ?? img.Url }, conn); diff --git a/PluralKit.Bot/Commands/MemberAvatar.cs b/PluralKit.Bot/Commands/MemberAvatar.cs index 7986e9d4..a0654a25 100644 --- a/PluralKit.Bot/Commands/MemberAvatar.cs +++ b/PluralKit.Bot/Commands/MemberAvatar.cs @@ -9,10 +9,12 @@ namespace PluralKit.Bot; public class MemberAvatar { private readonly HttpClient _client; + private readonly AvatarHostingService _avatarHosting; - public MemberAvatar(HttpClient client) + public MemberAvatar(HttpClient client, AvatarHostingService avatarHosting) { _client = client; + _avatarHosting = avatarHosting; } private async Task AvatarClear(MemberAvatarLocation location, Context ctx, PKMember target, MemberGuildSettings? mgs) @@ -138,6 +140,8 @@ public class MemberAvatar } ctx.CheckSystem().CheckOwnMember(target); + + avatarArg = await _avatarHosting.TryRehostImage(avatarArg.Value, AvatarHostingService.RehostedImageType.Avatar, ctx.Author.Id); await AvatarUtils.VerifyAvatarOrThrow(_client, avatarArg.Value.Url); await UpdateAvatar(location, ctx, target, avatarArg.Value.CleanUrl ?? avatarArg.Value.Url); await PrintResponse(location, ctx, target, avatarArg.Value, guildData); diff --git a/PluralKit.Bot/Commands/MemberEdit.cs b/PluralKit.Bot/Commands/MemberEdit.cs index b6193249..d2e943b3 100644 --- a/PluralKit.Bot/Commands/MemberEdit.cs +++ b/PluralKit.Bot/Commands/MemberEdit.cs @@ -13,10 +13,12 @@ namespace PluralKit.Bot; public class MemberEdit { private readonly HttpClient _client; + private readonly AvatarHostingService _avatarHosting; - public MemberEdit(HttpClient client) + public MemberEdit(HttpClient client, AvatarHostingService avatarHosting) { _client = client; + _avatarHosting = avatarHosting; } public async Task Name(Context ctx, PKMember target) @@ -180,6 +182,7 @@ public class MemberEdit async Task SetBannerImage(ParsedImage img) { + img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Banner, ctx.Author.Id); await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url, true); await ctx.Repository.UpdateMember(target.Id, new MemberPatch { BannerImage = img.CleanUrl ?? img.Url }); diff --git a/PluralKit.Bot/Commands/SystemEdit.cs b/PluralKit.Bot/Commands/SystemEdit.cs index b1a73e74..243ccaa6 100644 --- a/PluralKit.Bot/Commands/SystemEdit.cs +++ b/PluralKit.Bot/Commands/SystemEdit.cs @@ -18,12 +18,14 @@ public class SystemEdit private readonly HttpClient _client; private readonly DataFileService _dataFiles; private readonly PrivateChannelService _dmCache; + private readonly AvatarHostingService _avatarHosting; - public SystemEdit(DataFileService dataFiles, HttpClient client, PrivateChannelService dmCache) + public SystemEdit(DataFileService dataFiles, HttpClient client, PrivateChannelService dmCache, AvatarHostingService avatarHosting) { _dataFiles = dataFiles; _client = client; _dmCache = dmCache; + _avatarHosting = avatarHosting; } public async Task Name(Context ctx, PKSystem target) @@ -473,6 +475,7 @@ public class SystemEdit { ctx.CheckOwnSystem(target); + img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Avatar, ctx.Author.Id); await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url); await ctx.Repository.UpdateSystem(target.Id, new SystemPatch { AvatarUrl = img.CleanUrl ?? img.Url }); @@ -541,6 +544,7 @@ public class SystemEdit { ctx.CheckOwnSystem(target); + img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Avatar, ctx.Author.Id); await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url); await ctx.Repository.UpdateSystemGuild(target.Id, ctx.Guild.Id, new SystemGuildPatch { AvatarUrl = img.CleanUrl ?? img.Url }); @@ -638,6 +642,7 @@ public class SystemEdit else if (await ctx.MatchImage() is { } img) { + img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Banner, ctx.Author.Id); await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url, true); await ctx.Repository.UpdateSystem(target.Id, new SystemPatch { BannerImage = img.CleanUrl ?? img.Url }); diff --git a/PluralKit.Bot/Modules.cs b/PluralKit.Bot/Modules.cs index 2c4c0841..730908a9 100644 --- a/PluralKit.Bot/Modules.cs +++ b/PluralKit.Bot/Modules.cs @@ -136,6 +136,7 @@ public class BotModule: Module builder.RegisterType().AsSelf().SingleInstance(); builder.RegisterType().AsSelf().SingleInstance(); builder.RegisterType().AsSelf().SingleInstance(); + builder.RegisterType().AsSelf().SingleInstance(); // Sentry stuff builder.Register(_ => new Scope(null)).AsSelf().InstancePerLifetimeScope(); diff --git a/PluralKit.Bot/Services/AvatarHostingService.cs b/PluralKit.Bot/Services/AvatarHostingService.cs new file mode 100644 index 00000000..68a512ee --- /dev/null +++ b/PluralKit.Bot/Services/AvatarHostingService.cs @@ -0,0 +1,65 @@ +using System.Net; +using System.Net.Http.Json; + +namespace PluralKit.Bot; + +public class AvatarHostingService +{ + private readonly BotConfig _config; + private readonly HttpClient _client; + + public AvatarHostingService(BotConfig config, HttpClient client) + { + _config = config; + _client = client; + } + + public async Task TryRehostImage(ParsedImage input, RehostedImageType type, ulong userId) + { + var uploaded = await TryUploadAvatar(input.Url, type, userId); + if (uploaded != null) + { + // todo: make new image type called Cdn? + return new ParsedImage { Url = uploaded, Source = AvatarSource.Url }; + } + + return input; + } + + public async Task TryUploadAvatar(string? avatarUrl, RehostedImageType type, ulong userId) + { + if (!AvatarUtils.IsDiscordCdnUrl(avatarUrl)) + return null; + + if (_config.AvatarServiceUrl == null) + return null; + + var kind = type switch + { + RehostedImageType.Avatar => "avatar", + RehostedImageType.Banner => "banner", + _ => throw new ArgumentOutOfRangeException(nameof(type), type, null) + }; + + var response = await _client.PostAsJsonAsync(_config.AvatarServiceUrl + "/pull", + new { url = avatarUrl, kind, uploaded_by = userId }); + if (response.StatusCode != HttpStatusCode.OK) + { + var error = await response.Content.ReadFromJsonAsync(); + throw new PKError($"Error uploading image to CDN: {error.Error}"); + } + + var success = await response.Content.ReadFromJsonAsync(); + return success.Url; + } + + public record ErrorResponse(string Error); + + public record SuccessResponse(string Url, bool New); + + public enum RehostedImageType + { + Avatar, + Banner, + } +} \ No newline at end of file diff --git a/PluralKit.Bot/Utils/AvatarUtils.cs b/PluralKit.Bot/Utils/AvatarUtils.cs index 00f3059f..70206593 100644 --- a/PluralKit.Bot/Utils/AvatarUtils.cs +++ b/PluralKit.Bot/Utils/AvatarUtils.cs @@ -80,4 +80,6 @@ public static class AvatarUtils return newUrl; } + + public static bool IsDiscordCdnUrl(string? url) => url != null && DiscordCdnUrl.Match(url).Success; } \ No newline at end of file