eightball/PluralKit
1
0
Fork 0
mirror of https://github.com/PluralKit/PluralKit.git synced 2026-02-04 04:56:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

feat(bot): use avater service for image verify

Browse source
This commit is contained in:
alyssa 2025-02-23 09:33:32 +00:00
parent 5f6c8c0d14
commit d537f05b23
9 changed files with 94 additions and 63 deletions
Download patch file Download diff file Expand all files Collapse all files

4
PluralKit.Bot/Commands/Groups.cs
View file

@ -291,7 +291,7 @@ public class Groups
ctx.CheckOwnGroup(target);
img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Avatar, ctx.Author.Id, ctx.System);
await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url);
await _avatarHosting.VerifyAvatarOrThrow(img.Url);
await ctx.Repository.UpdateGroup(target.Id, new GroupPatch { Icon = img.CleanUrl ?? img.Url });
@ -366,7 +366,7 @@ public class Groups
ctx.CheckOwnGroup(target);
img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Banner, ctx.Author.Id, ctx.System);
await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url, true);
await _avatarHosting.VerifyAvatarOrThrow(img.Url, true);
await ctx.Repository.UpdateGroup(target.Id, new GroupPatch { BannerImage = img.CleanUrl ?? img.Url });

2
PluralKit.Bot/Commands/Member.cs
View file

@ -83,7 +83,7 @@ public class Member
img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Avatar, ctx.Author.Id, ctx.System);
await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url);
await _avatarHosting.VerifyAvatarOrThrow(img.Url);
await ctx.Repository.UpdateMember(member.Id, new MemberPatch { AvatarUrl = img.CleanUrl ?? img.Url }, conn);
dispatchData.Add("avatar_url", img.CleanUrl);

2
PluralKit.Bot/Commands/MemberAvatar.cs
View file

@ -159,7 +159,7 @@ public class MemberAvatar
ctx.CheckSystem().CheckOwnMember(target);
avatarArg = await _avatarHosting.TryRehostImage(avatarArg.Value, AvatarHostingService.RehostedImageType.Avatar, ctx.Author.Id, ctx.System);
await AvatarUtils.VerifyAvatarOrThrow(_client, avatarArg.Value.Url);
await _avatarHosting.VerifyAvatarOrThrow(avatarArg.Value.Url);
await UpdateAvatar(location, ctx, target, avatarArg.Value.CleanUrl ?? avatarArg.Value.Url);
await PrintResponse(location, ctx, target, avatarArg.Value, guildData);
}

2
PluralKit.Bot/Commands/MemberEdit.cs
View file

@ -231,7 +231,7 @@ public class MemberEdit
{
ctx.CheckOwnMember(target);
img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Banner, ctx.Author.Id, ctx.System);
await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url, true);
await _avatarHosting.VerifyAvatarOrThrow(img.Url, true);
await ctx.Repository.UpdateMember(target.Id, new MemberPatch { BannerImage = img.CleanUrl ?? img.Url });

6
PluralKit.Bot/Commands/SystemEdit.cs
View file

@ -572,7 +572,7 @@ public class SystemEdit
ctx.CheckOwnSystem(target);
img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Avatar, ctx.Author.Id, ctx.System);
await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url);
await _avatarHosting.VerifyAvatarOrThrow(img.Url);
await ctx.Repository.UpdateSystem(target.Id, new SystemPatch { AvatarUrl = img.CleanUrl ?? img.Url });
@ -659,7 +659,7 @@ public class SystemEdit
ctx.CheckOwnSystem(target);
img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Avatar, ctx.Author.Id, ctx.System);
await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url);
await _avatarHosting.VerifyAvatarOrThrow(img.Url);
await ctx.Repository.UpdateSystemGuild(target.Id, ctx.Guild.Id, new SystemGuildPatch { AvatarUrl = img.CleanUrl ?? img.Url });
@ -781,7 +781,7 @@ public class SystemEdit
else if (await ctx.MatchImage() is { } img)
{
img = await _avatarHosting.TryRehostImage(img, AvatarHostingService.RehostedImageType.Banner, ctx.Author.Id, ctx.System);
await AvatarUtils.VerifyAvatarOrThrow(_client, img.Url, true);
await _avatarHosting.VerifyAvatarOrThrow(img.Url, true);
await ctx.Repository.UpdateSystem(target.Id, new SystemPatch { BannerImage = img.CleanUrl ?? img.Url });

38
PluralKit.Bot/Services/AvatarHostingService.cs
View file

@ -18,6 +18,44 @@ public class AvatarHostingService
};
}
public async Task VerifyAvatarOrThrow(string url, bool isBanner = false)
{
if (url.Length > Limits.MaxUriLength)
throw Errors.UrlTooLong(url);
if (!PluralKit.Core.MiscUtils.TryMatchUri(url, out var uri))
throw Errors.InvalidUrl;
if (uri.Host.Contains("toyhou.se"))
throw new PKError("Due to server issues, PluralKit is unable to read images hosted on toyhou.se.");
if (uri.Host == "cdn.pluralkit.me") return;
if (_config.AvatarServiceUrl == null)
return;
var kind = isBanner ? "banner" : "avatar";
try
{
var response = await _client.PostAsJsonAsync(_config.AvatarServiceUrl + "/verify",
new { url, kind });
if (response.StatusCode != HttpStatusCode.OK)
{
var error = await response.Content.ReadFromJsonAsync<ErrorResponse>();
throw new PKError($"{error.Error}");
}
}
catch (TaskCanceledException e)
{
// don't show an internal error to users
if (e.Message.Contains("HttpClient.Timeout"))
throw new PKError("Temporary error setting image, please try again later");
throw;
}
}
public async Task<ParsedImage> TryRehostImage(ParsedImage input, RehostedImageType type, ulong userId, PKSystem? system)
{
try

51
PluralKit.Bot/Utils/AvatarUtils.cs
View file

@ -8,57 +8,6 @@ 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", "image/webp"
};
if (!PluralKit.Core.MiscUtils.TryMatchUri(url, out var uri))
throw Errors.InvalidUrl;
if (uri.Host.Contains("toyhou.se"))
throw new PKError("Due to server issues, PluralKit is unable to read images hosted on toyhou.se.");
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);
try
{
if (!acceptableMimeTypes.Contains(response.Content.Headers.ContentType.MediaType)) // Check MIME type
throw Errors.AvatarNotAnImage(response.Content.Headers.ContentType.MediaType);
}
catch (NullReferenceException)
{
throw new PKError("Could not verify avatar is an image. This can happen when the server sends a malformed response."
+ "\nPlease join the support server for help: <https://discord.gg/PczBt78>");
}
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!

29
crates/avatars/src/main.rs
View file

@ -16,6 +16,7 @@ use axum::{
use libpk::_config::AvatarsConfig;
use libpk::db::repository::avatars as db;
use libpk::db::types::avatars::*;
use pull::ParsedUrl;
use reqwest::{Client, ClientBuilder};
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
@ -35,9 +36,15 @@ pub enum PKAvatarError {
#[error("discord cdn responded with status code: {0}")]
BadCdnResponse(reqwest::StatusCode),
#[error("server responded with status code: {0}")]
BadServerResponse(reqwest::StatusCode),
#[error("network error: {0}")]
NetworkError(reqwest::Error),
#[error("network error: {0}")]
NetworkErrorString(String),
#[error("response is missing header: {0}")]
MissingHeader(&'static str),
@ -86,7 +93,6 @@ async fn pull(
) -> Result<Json<PullResponse>, PKAvatarError> {
let parsed = pull::parse_url(&req.url) // parsing beforehand to "normalize"
.map_err(|_| PKAvatarError::InvalidCdnUrl)?;
if !req.force {
if let Some(existing) = db::get_by_attachment_id(&state.pool, parsed.attachment_id).await? {
// remove any pending image cleanup
@ -132,6 +138,26 @@ async fn pull(
}))
}
async fn verify(
State(state): State<AppState>,
Json(req): Json<PullRequest>,
) -> Result<(), PKAvatarError> {
let result = crate::pull::pull(
state.pull_client,
&ParsedUrl {
full_url: req.url.clone(),
channel_id: 0,
attachment_id: 0,
filename: "".to_string(),
},
)
.await?;
let encoded = process::process_async(result.data, req.kind).await?;
Ok(())
}
pub async fn stats(State(state): State<AppState>) -> Result<Json<Stats>, PKAvatarError> {
Ok(Json(db::get_stats(&state.pool).await?))
}
@ -193,6 +219,7 @@ async fn real_main() -> anyhow::Result<()> {
// migrate::spawn_migrate_workers(Arc::new(state.clone()), state.config.migrate_worker_count);
let app = Router::new()
.route("/verify", post(verify))
.route("/pull", post(pull))
.route("/stats", get(stats))
.with_state(state);

23
crates/avatars/src/pull.rs
View file

@ -3,6 +3,8 @@ use std::{str::FromStr, sync::Arc};
use crate::PKAvatarError;
use anyhow::Context;
use reqwest::{Client, StatusCode, Url};
use std::error::Error;
use std::fmt::Write;
use std::time::Instant;
use tracing::{error, instrument};
@ -28,14 +30,29 @@ pub async fn pull(
.expect("set_host should not fail");
}
let response = client.get(trimmed_url.clone()).send().await.map_err(|e| {
error!("network error for {}: {}", parsed_url.full_url, e);
PKAvatarError::NetworkError(e)
// terrible
let mut s = format!("{}", e);
if let Some(src) = e.source() {
let _ = write!(s, "\n\nCaused by: {}", src);
let mut err = src;
while let Some(src) = err.source() {
let _ = write!(s, "\n\nCaused by: {}", src);
err = src;
}
}
error!("network error for {}: {}", parsed_url.full_url, s);
PKAvatarError::NetworkErrorString(s)
})?;
let time_after_headers = Instant::now();
let status = response.status();
if status != StatusCode::OK {
return Err(PKAvatarError::BadCdnResponse(status));
if trimmed_url.host_str() == Some("cdn.discordapp.com") {
return Err(PKAvatarError::BadCdnResponse(status));
} else {
return Err(PKAvatarError::BadServerResponse(status));
}
}
let size = match response.content_length() {