PluralKit/PluralKit.Bot/Utils/AvatarUtils.cs

using System;
using System.Linq;
using System.Net.Http;
using System.Text.RegularExpressions;
using System.Threading.Tasks;

using PluralKit.Core;

using SixLabors.ImageSharp;

namespace PluralKit.Bot
{
    public static class AvatarUtils
    {
        public static async Task VerifyAvatarOrThrow(HttpClient client, string url, bool isFullSizeImage = false)
        {
            if (url.Length > Limits.MaxUriLength)
                throw Errors.UrlTooLong(url);

            // List of MIME types we consider acceptable
            var acceptableMimeTypes = new[]
            {
                "image/jpeg",
                "image/gif",
                "image/png"
                // TODO: add image/webp once ImageSharp supports this
            };

            if (!PluralKit.Core.MiscUtils.TryMatchUri(url, out var uri))
                throw Errors.InvalidUrl(url);

            url = TryRewriteCdnUrl(url);

            var response = await client.GetAsync(url);
            if (!response.IsSuccessStatusCode) // Check status code
                throw Errors.AvatarServerError(response.StatusCode);
            if (response.Content.Headers.ContentLength == null) // Check presence of content length
                throw Errors.AvatarNotAnImage(null);
            if (!acceptableMimeTypes.Contains(response.Content.Headers.ContentType.MediaType)) // Check MIME type
                throw Errors.AvatarNotAnImage(response.Content.Headers.ContentType.MediaType);

            if (isFullSizeImage)
                // no need to do size checking on banners
                return;

            if (response.Content.Headers.ContentLength > Limits.AvatarFileSizeLimit) // Check content length
                throw Errors.AvatarFileSizeLimit(response.Content.Headers.ContentLength.Value);

            // Parse the image header in a worker
            var stream = await response.Content.ReadAsStreamAsync();
            var image = await Task.Run(() => Image.Identify(stream));
            if (image == null) throw Errors.AvatarInvalid;
            if (image.Width > Limits.AvatarDimensionLimit || image.Height > Limits.AvatarDimensionLimit) // Check image size
                throw Errors.AvatarDimensionsTooLarge(image.Width, image.Height);
        }

        // Rewrite cdn.discordapp.com URLs to media.discordapp.net for jpg/png files
        // This lets us add resizing parameters to "borrow" their media proxy server to downsize the image
        // which in turn makes it more likely to be underneath the size limit!
        private static readonly Regex DiscordCdnUrl = new Regex(@"^https?://(?:cdn\.discordapp\.com|media\.discordapp\.net)/attachments/(\d{17,19})/(\d{17,19})/([^/\\&\?]+)\.(png|jpg|jpeg|webp)(\?.*)?$");
        private static readonly string DiscordMediaUrlReplacement = "https://media.discordapp.net/attachments/$1/$2/$3.$4?width=256&height=256";
        public static string? TryRewriteCdnUrl(string? url)
        {
            return url == null ? null : DiscordCdnUrl.Replace(url, DiscordMediaUrlReplacement);
        }
    }
}