mirror of
https://github.com/PluralKit/PluralKit.git
synced 2026-02-04 13:06:50 +00:00
chore: merge avatars service into monorepo
This commit is contained in:
parent
f427d4d727
commit
17f5561293
27 changed files with 1925 additions and 111 deletions
3
.github/workflows/rust.yml
vendored
3
.github/workflows/rust.yml
vendored
|
|
@ -10,6 +10,7 @@ on:
|
|||
- 'lib/libpk/**'
|
||||
- 'services/api/**'
|
||||
- 'services/gateway/**'
|
||||
- 'services/avatars/**'
|
||||
- '.github/workflows/rust.yml'
|
||||
- 'Dockerfile.rust'
|
||||
- 'Dockerfile.bin'
|
||||
|
|
@ -47,7 +48,7 @@ jobs:
|
|||
|
||||
# add more binaries here
|
||||
- run: |
|
||||
for binary in "api" "gateway"; do
|
||||
for binary in "api" "gateway" "avatars"; do
|
||||
for tag in latest ${{ env.BRANCH_NAME }} ${{ github.sha }}; do
|
||||
cat Dockerfile.bin | sed "s/__BINARY__/$binary/g" | docker build -t ghcr.io/pluralkit/$binary:$tag -f - .
|
||||
done
|
||||
|
|
|
|||
746
Cargo.lock
generated
746
Cargo.lock
generated
File diff suppressed because it is too large
Load diff
14
Cargo.toml
14
Cargo.toml
|
|
@ -3,7 +3,8 @@ members = [
|
|||
"./lib/libpk",
|
||||
"./services/api",
|
||||
"./services/dispatch",
|
||||
"./services/gateway"
|
||||
"./services/gateway",
|
||||
"./services/avatars"
|
||||
]
|
||||
|
||||
[workspace.dependencies]
|
||||
|
|
@ -16,13 +17,16 @@ fred = { version = "5.2.0", default-features = false, features = ["tracing", "po
|
|||
futures = "0.3.30"
|
||||
lazy_static = "1.4.0"
|
||||
metrics = "0.23.0"
|
||||
serde = "1.0.152"
|
||||
reqwest = { version = "0.12.7" , default-features = false, features = ["rustls-tls", "trust-dns"]}
|
||||
serde = { version = "1.0.196", features = ["derive"] }
|
||||
serde_json = "1.0.117"
|
||||
signal-hook = "0.3.17"
|
||||
sqlx = { version = "0.7.4", features = ["runtime-tokio", "postgres", "chrono", "macros"] }
|
||||
tokio = { version = "1.25.0", features = ["full"] }
|
||||
tracing = "0.1.37"
|
||||
sqlx = { version = "0.7.4", features = ["runtime-tokio", "postgres", "time", "macros", "uuid"] }
|
||||
time = "0.3.34"
|
||||
tokio = { version = "1.36.0", features = ["full"] }
|
||||
tracing = "0.1.40"
|
||||
tracing-subscriber = { version = "0.3.16", features = ["env-filter", "json"] }
|
||||
uuid = { version = "1.7.0", features = ["serde"] }
|
||||
|
||||
twilight-gateway = { git = "https://github.com/pluralkit/twilight" }
|
||||
twilight-cache-inmemory = { git = "https://github.com/pluralkit/twilight", features = ["permission-calculator"] }
|
||||
|
|
|
|||
|
|
@ -28,11 +28,14 @@ COPY proto/ /build/proto
|
|||
COPY lib/libpk /build/lib/libpk
|
||||
COPY services/api/ /build/services/api
|
||||
COPY services/gateway/ /build/services/gateway
|
||||
COPY services/avatars/ /build/services/avatars
|
||||
|
||||
RUN cargo build --bin api --release --target x86_64-unknown-linux-musl
|
||||
RUN cargo build --bin gateway --release --target x86_64-unknown-linux-musl
|
||||
RUN cargo build --bin avatars --release --target x86_64-unknown-linux-musl
|
||||
|
||||
FROM scratch
|
||||
|
||||
COPY --from=binary-builder /build/target/x86_64-unknown-linux-musl/release/api /api
|
||||
COPY --from=binary-builder /build/target/x86_64-unknown-linux-musl/release/gateway /gateway
|
||||
COPY --from=binary-builder /build/target/x86_64-unknown-linux-musl/release/avatars /avatars
|
||||
|
|
|
|||
|
|
@ -13,11 +13,13 @@ metrics = { workspace = true }
|
|||
metrics-exporter-prometheus = { version = "0.15.3", default-features = false, features = ["tokio", "http-listener", "tracing"] }
|
||||
serde = { workspace = true }
|
||||
sqlx = { workspace = true }
|
||||
time = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
tracing-gelf = "0.7.1"
|
||||
tracing-subscriber = { workspace = true}
|
||||
twilight-model = { workspace = true }
|
||||
uuid = { workspace = true }
|
||||
|
||||
prost = { workspace = true }
|
||||
prost-types = { workspace = true }
|
||||
|
|
|
|||
|
|
@ -20,6 +20,9 @@ pub struct DiscordConfig {
|
|||
pub max_concurrency: u32,
|
||||
pub cluster: Option<ClusterSettings>,
|
||||
pub api_base_url: Option<String>,
|
||||
|
||||
#[serde(default = "_default_api_addr")]
|
||||
pub cache_api_addr: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
|
|
@ -36,7 +39,7 @@ fn _default_api_addr() -> String {
|
|||
"0.0.0.0:5000".to_string()
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
#[derive(Deserialize, Clone, Debug)]
|
||||
pub struct ApiConfig {
|
||||
#[serde(default = "_default_api_addr")]
|
||||
pub addr: String,
|
||||
|
|
@ -50,6 +53,23 @@ pub struct ApiConfig {
|
|||
pub temp_token2: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Clone, Debug)]
|
||||
pub struct AvatarsConfig {
|
||||
pub s3: S3Config,
|
||||
pub cdn_url: String,
|
||||
|
||||
#[serde(default)]
|
||||
pub migrate_worker_count: u32,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Clone, Debug)]
|
||||
pub struct S3Config {
|
||||
pub bucket: String,
|
||||
pub application_id: String,
|
||||
pub application_key: String,
|
||||
pub endpoint: String,
|
||||
}
|
||||
|
||||
fn _metrics_default() -> bool {
|
||||
false
|
||||
}
|
||||
|
|
@ -61,8 +81,9 @@ fn _json_log_default() -> bool {
|
|||
pub struct PKConfig {
|
||||
pub db: DatabaseConfig,
|
||||
|
||||
pub discord: DiscordConfig,
|
||||
pub api: ApiConfig,
|
||||
pub discord: Option<DiscordConfig>,
|
||||
pub api: Option<ApiConfig>,
|
||||
pub avatars: Option<AvatarsConfig>,
|
||||
|
||||
#[serde(default = "_metrics_default")]
|
||||
pub run_metrics_server: bool,
|
||||
|
|
@ -71,6 +92,16 @@ pub struct PKConfig {
|
|||
pub(crate) json_log: bool,
|
||||
}
|
||||
|
||||
impl PKConfig {
|
||||
pub fn api(self) -> ApiConfig {
|
||||
self.api.expect("missing api config")
|
||||
}
|
||||
|
||||
pub fn discord_config(self) -> DiscordConfig {
|
||||
self.discord.expect("missing discord config")
|
||||
}
|
||||
}
|
||||
|
||||
lazy_static! {
|
||||
#[derive(Debug)]
|
||||
pub static ref CONFIG: Arc<PKConfig> = {
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ use std::str::FromStr;
|
|||
use tracing::info;
|
||||
|
||||
pub mod repository;
|
||||
pub mod types;
|
||||
|
||||
pub async fn init_redis() -> anyhow::Result<RedisPool> {
|
||||
info!("connecting to redis");
|
||||
|
|
|
|||
87
lib/libpk/src/db/repository/avatars.rs
Normal file
87
lib/libpk/src/db/repository/avatars.rs
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
use sqlx::{PgPool, Postgres, Transaction};
|
||||
|
||||
use crate::db::types::avatars::*;
|
||||
|
||||
pub async fn get_by_original_url(
|
||||
pool: &PgPool,
|
||||
original_url: &str,
|
||||
) -> anyhow::Result<Option<ImageMeta>> {
|
||||
Ok(
|
||||
sqlx::query_as("select * from images where original_url = $1")
|
||||
.bind(original_url)
|
||||
.fetch_optional(pool)
|
||||
.await?,
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn get_by_attachment_id(
|
||||
pool: &PgPool,
|
||||
attachment_id: u64,
|
||||
) -> anyhow::Result<Option<ImageMeta>> {
|
||||
Ok(
|
||||
sqlx::query_as("select * from images where original_attachment_id = $1")
|
||||
.bind(attachment_id as i64)
|
||||
.fetch_optional(pool)
|
||||
.await?,
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn pop_queue(
|
||||
pool: &PgPool,
|
||||
) -> anyhow::Result<Option<(Transaction<Postgres>, ImageQueueEntry)>> {
|
||||
let mut tx = pool.begin().await?;
|
||||
let res: Option<ImageQueueEntry> = sqlx::query_as("delete from image_queue where itemid = (select itemid from image_queue order by itemid for update skip locked limit 1) returning *")
|
||||
.fetch_optional(&mut *tx).await?;
|
||||
Ok(res.map(|x| (tx, x)))
|
||||
}
|
||||
|
||||
pub async fn get_queue_length(pool: &PgPool) -> anyhow::Result<i64> {
|
||||
Ok(sqlx::query_scalar("select count(*) from image_queue")
|
||||
.fetch_one(pool)
|
||||
.await?)
|
||||
}
|
||||
|
||||
pub async fn get_stats(pool: &PgPool) -> anyhow::Result<Stats> {
|
||||
Ok(sqlx::query_as(
|
||||
"select count(*) as total_images, sum(file_size) as total_file_size from images",
|
||||
)
|
||||
.fetch_one(pool)
|
||||
.await?)
|
||||
}
|
||||
|
||||
pub async fn add_image(pool: &PgPool, meta: ImageMeta) -> anyhow::Result<bool> {
|
||||
let kind_str = match meta.kind {
|
||||
ImageKind::Avatar => "avatar",
|
||||
ImageKind::Banner => "banner",
|
||||
};
|
||||
|
||||
let res = sqlx::query("insert into images (id, url, content_type, original_url, file_size, width, height, original_file_size, original_type, original_attachment_id, kind, uploaded_by_account, uploaded_by_system, uploaded_at) values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, (now() at time zone 'utc')) on conflict (id) do nothing")
|
||||
.bind(meta.id)
|
||||
.bind(meta.url)
|
||||
.bind(meta.content_type)
|
||||
.bind(meta.original_url)
|
||||
.bind(meta.file_size)
|
||||
.bind(meta.width)
|
||||
.bind(meta.height)
|
||||
.bind(meta.original_file_size)
|
||||
.bind(meta.original_type)
|
||||
.bind(meta.original_attachment_id)
|
||||
.bind(kind_str)
|
||||
.bind(meta.uploaded_by_account)
|
||||
.bind(meta.uploaded_by_system)
|
||||
.execute(pool).await?;
|
||||
Ok(res.rows_affected() > 0)
|
||||
}
|
||||
|
||||
pub async fn push_queue(
|
||||
conn: &mut sqlx::PgConnection,
|
||||
url: &str,
|
||||
kind: ImageKind,
|
||||
) -> anyhow::Result<()> {
|
||||
sqlx::query("insert into image_queue (url, kind) values ($1, $2)")
|
||||
.bind(url)
|
||||
.bind(kind)
|
||||
.execute(conn)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -1,5 +1,7 @@
|
|||
mod stats;
|
||||
pub use stats::*;
|
||||
|
||||
pub mod avatars;
|
||||
|
||||
mod auth;
|
||||
pub use auth::*;
|
||||
|
|
|
|||
53
lib/libpk/src/db/types/avatars.rs
Normal file
53
lib/libpk/src/db/types/avatars.rs
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::FromRow;
|
||||
use time::OffsetDateTime;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(FromRow)]
|
||||
pub struct ImageMeta {
|
||||
pub id: String,
|
||||
pub kind: ImageKind,
|
||||
pub content_type: String,
|
||||
pub url: String,
|
||||
pub file_size: i32,
|
||||
pub width: i32,
|
||||
pub height: i32,
|
||||
pub uploaded_at: Option<OffsetDateTime>,
|
||||
|
||||
pub original_url: Option<String>,
|
||||
pub original_attachment_id: Option<i64>,
|
||||
pub original_file_size: Option<i32>,
|
||||
pub original_type: Option<String>,
|
||||
pub uploaded_by_account: Option<i64>,
|
||||
pub uploaded_by_system: Option<Uuid>,
|
||||
}
|
||||
|
||||
#[derive(FromRow, Serialize)]
|
||||
pub struct Stats {
|
||||
pub total_images: i64,
|
||||
pub total_file_size: i64,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Clone, Copy, Debug, sqlx::Type, PartialEq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
#[sqlx(rename_all = "snake_case", type_name = "text")]
|
||||
pub enum ImageKind {
|
||||
Avatar,
|
||||
Banner,
|
||||
}
|
||||
|
||||
impl ImageKind {
|
||||
pub fn size(&self) -> (u32, u32) {
|
||||
match self {
|
||||
Self::Avatar => (512, 512),
|
||||
Self::Banner => (1024, 1024),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(FromRow)]
|
||||
pub struct ImageQueueEntry {
|
||||
pub itemid: i32,
|
||||
pub url: String,
|
||||
pub kind: ImageKind,
|
||||
}
|
||||
1
lib/libpk/src/db/types/mod.rs
Normal file
1
lib/libpk/src/db/types/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod avatars;
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
use metrics_exporter_prometheus::PrometheusBuilder;
|
||||
use tracing_subscriber::{EnvFilter, Registry};
|
||||
use tracing_subscriber::EnvFilter;
|
||||
|
||||
pub mod db;
|
||||
pub mod proto;
|
||||
|
|
|
|||
|
|
@ -63,7 +63,7 @@ async fn main() -> anyhow::Result<()> {
|
|||
let db = libpk::db::init_data_db().await?;
|
||||
let redis = libpk::db::init_redis().await?;
|
||||
|
||||
let rproxy_uri = Uri::from_static(&libpk::config.api.remote_url).to_string();
|
||||
let rproxy_uri = Uri::from_static(&libpk::config.api.as_ref().expect("missing api config").remote_url).to_string();
|
||||
let rproxy_client = hyper_util::client::legacy::Client::<(), ()>::builder(TokioExecutor::new())
|
||||
.build(HttpConnector::new());
|
||||
|
||||
|
|
@ -145,7 +145,7 @@ async fn main() -> anyhow::Result<()> {
|
|||
|
||||
.route("/", get(|| async { axum::response::Redirect::to("https://pluralkit.me/api") }));
|
||||
|
||||
let addr: &str = libpk::config.api.addr.as_ref();
|
||||
let addr: &str = libpk::config.api.as_ref().expect("missing api config").addr.as_ref();
|
||||
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||||
info!("listening on {}", addr);
|
||||
axum::serve(listener, app).await?;
|
||||
|
|
|
|||
|
|
@ -20,32 +20,38 @@ lazy_static::lazy_static! {
|
|||
|
||||
// this is awful but it works
|
||||
pub fn ratelimiter<F, T>(f: F) -> FromFnLayer<F, Option<RedisPool>, T> {
|
||||
let redis = libpk::config.api.ratelimit_redis_addr.as_ref().map(|val| {
|
||||
let r = fred::pool::RedisPool::new(
|
||||
fred::types::RedisConfig::from_url_centralized(val.as_ref())
|
||||
.expect("redis url is invalid"),
|
||||
10,
|
||||
)
|
||||
.expect("failed to connect to redis");
|
||||
let redis = libpk::config
|
||||
.api
|
||||
.as_ref()
|
||||
.expect("missing api config")
|
||||
.ratelimit_redis_addr
|
||||
.as_ref()
|
||||
.map(|val| {
|
||||
let r = fred::pool::RedisPool::new(
|
||||
fred::types::RedisConfig::from_url_centralized(val.as_ref())
|
||||
.expect("redis url is invalid"),
|
||||
10,
|
||||
)
|
||||
.expect("failed to connect to redis");
|
||||
|
||||
let handle = r.connect(Some(ReconnectPolicy::default()));
|
||||
let handle = r.connect(Some(ReconnectPolicy::default()));
|
||||
|
||||
tokio::spawn(async move { handle });
|
||||
tokio::spawn(async move { handle });
|
||||
|
||||
let rscript = r.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Ok(()) = rscript.wait_for_connect().await {
|
||||
match rscript.script_load(LUA_SCRIPT).await {
|
||||
Ok(_) => info!("connected to redis for request rate limiting"),
|
||||
Err(err) => error!("could not load redis script: {}", err),
|
||||
let rscript = r.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Ok(()) = rscript.wait_for_connect().await {
|
||||
match rscript.script_load(LUA_SCRIPT).await {
|
||||
Ok(_) => info!("connected to redis for request rate limiting"),
|
||||
Err(err) => error!("could not load redis script: {}", err),
|
||||
}
|
||||
} else {
|
||||
error!("could not wait for connection to load redis script!");
|
||||
}
|
||||
} else {
|
||||
error!("could not wait for connection to load redis script!");
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
r
|
||||
});
|
||||
r
|
||||
});
|
||||
|
||||
if redis.is_none() {
|
||||
warn!("running without request rate limiting!");
|
||||
|
|
@ -95,7 +101,12 @@ pub async fn do_request_ratelimited(
|
|||
// https://github.com/rust-lang/rust/issues/53667
|
||||
let is_temp_token2 = if let Some(header) = request.headers().clone().get("X-PluralKit-App")
|
||||
{
|
||||
if let Some(token2) = &libpk::config.api.temp_token2 {
|
||||
if let Some(token2) = &libpk::config
|
||||
.api
|
||||
.as_ref()
|
||||
.expect("missing api config")
|
||||
.temp_token2
|
||||
{
|
||||
if header.to_str().unwrap_or("invalid") == token2 {
|
||||
true
|
||||
} else {
|
||||
|
|
|
|||
25
services/avatars/Cargo.toml
Normal file
25
services/avatars/Cargo.toml
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
[package]
|
||||
name = "avatars"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[dependencies]
|
||||
libpk = { path = "../../lib/libpk" }
|
||||
anyhow = { workspace = true }
|
||||
axum = { workspace = true }
|
||||
data-encoding = "2.5.0"
|
||||
form_urlencoded = "1.2.1"
|
||||
futures = { workspace = true }
|
||||
gif = "0.13.1"
|
||||
image = { version = "0.24.8", default-features = false, features = ["gif", "jpeg", "png", "webp", "tiff"] }
|
||||
reqwest = { workspace = true }
|
||||
rust-s3 = { version = "0.33.0", default-features = false, features = ["tokio-rustls-tls"] }
|
||||
sha2 = "0.10.8"
|
||||
serde = { workspace = true }
|
||||
sqlx = { workspace = true }
|
||||
thiserror = "1.0.56"
|
||||
time = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
uuid = { workspace = true }
|
||||
webp = "0.2.6"
|
||||
21
services/avatars/src/hash.rs
Normal file
21
services/avatars/src/hash.rs
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
use std::fmt::Display;
|
||||
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct Hash([u8; 32]);
|
||||
|
||||
impl Hash {
|
||||
pub fn sha256(data: &[u8]) -> Hash {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(data);
|
||||
Hash(hasher.finalize().into())
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for Hash {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
let encoding = data_encoding::BASE32_NOPAD;
|
||||
write!(f, "{}", encoding.encode(&self.0[..16]).to_lowercase())
|
||||
}
|
||||
}
|
||||
24
services/avatars/src/init.sql
Normal file
24
services/avatars/src/init.sql
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
create table if not exists images
|
||||
(
|
||||
id text primary key,
|
||||
url text not null,
|
||||
original_url text,
|
||||
original_file_size int,
|
||||
original_type text,
|
||||
original_attachment_id bigint,
|
||||
file_size int not null,
|
||||
width int not null,
|
||||
height int not null,
|
||||
kind text not null,
|
||||
uploaded_at timestamptz not null,
|
||||
uploaded_by_account bigint
|
||||
);
|
||||
|
||||
create index if not exists images_original_url_idx on images (original_url);
|
||||
create index if not exists images_original_attachment_id_idx on images (original_attachment_id);
|
||||
create index if not exists images_uploaded_by_account_idx on images (uploaded_by_account);
|
||||
|
||||
create table if not exists image_queue (itemid serial primary key, url text not null, kind text not null);
|
||||
|
||||
alter table images add column if not exists uploaded_by_system uuid;
|
||||
alter table images add column if not exists content_type text default 'image/webp';
|
||||
259
services/avatars/src/main.rs
Normal file
259
services/avatars/src/main.rs
Normal file
|
|
@ -0,0 +1,259 @@
|
|||
mod hash;
|
||||
mod migrate;
|
||||
mod process;
|
||||
mod pull;
|
||||
mod store;
|
||||
|
||||
use anyhow::Context;
|
||||
use axum::extract::State;
|
||||
use axum::routing::get;
|
||||
use axum::{
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
routing::post,
|
||||
Json, Router,
|
||||
};
|
||||
use libpk::_config::AvatarsConfig;
|
||||
use libpk::db::repository::avatars as db;
|
||||
use libpk::db::types::avatars::*;
|
||||
use reqwest::{Client, ClientBuilder};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::PgPool;
|
||||
use std::error::Error;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use thiserror::Error;
|
||||
use tracing::{error, info};
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum PKAvatarError {
|
||||
// todo: split off into logical groups (cdn/url error, image format error, etc)
|
||||
#[error("invalid cdn url")]
|
||||
InvalidCdnUrl,
|
||||
|
||||
#[error("discord cdn responded with status code: {0}")]
|
||||
BadCdnResponse(reqwest::StatusCode),
|
||||
|
||||
#[error("network error: {0}")]
|
||||
NetworkError(reqwest::Error),
|
||||
|
||||
#[error("response is missing header: {0}")]
|
||||
MissingHeader(&'static str),
|
||||
|
||||
#[error("unsupported content type: {0}")]
|
||||
UnsupportedContentType(String),
|
||||
|
||||
#[error("image file size too large ({0} > {1})")]
|
||||
ImageFileSizeTooLarge(u64, u64),
|
||||
|
||||
#[error("unsupported image format: {0:?}")]
|
||||
UnsupportedImageFormat(image::ImageFormat),
|
||||
|
||||
#[error("could not detect image format")]
|
||||
UnknownImageFormat,
|
||||
|
||||
#[error("original image dimensions too large: {0:?} > {1:?}")]
|
||||
ImageDimensionsTooLarge((u32, u32), (u32, u32)),
|
||||
|
||||
#[error("could not decode image, is it corrupted?")]
|
||||
ImageFormatError(#[from] image::ImageError),
|
||||
|
||||
#[error("unknown error")]
|
||||
InternalError(#[from] anyhow::Error),
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub struct PullRequest {
|
||||
url: String,
|
||||
kind: ImageKind,
|
||||
uploaded_by: Option<u64>, // should be String? serde makes this hard :/
|
||||
system_id: Option<Uuid>,
|
||||
|
||||
#[serde(default)]
|
||||
force: bool,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct PullResponse {
|
||||
url: String,
|
||||
new: bool,
|
||||
}
|
||||
|
||||
async fn pull(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<PullRequest>,
|
||||
) -> 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? {
|
||||
return Ok(Json(PullResponse {
|
||||
url: existing.url,
|
||||
new: false,
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
let result = crate::pull::pull(state.pull_client, &parsed).await?;
|
||||
|
||||
let original_file_size = result.data.len();
|
||||
let encoded = process::process_async(result.data, req.kind).await?;
|
||||
|
||||
let store_res = crate::store::store(&state.bucket, &encoded).await?;
|
||||
let final_url = format!("{}{}", state.config.cdn_url, store_res.path);
|
||||
let is_new = db::add_image(
|
||||
&state.pool,
|
||||
ImageMeta {
|
||||
id: store_res.id,
|
||||
url: final_url.clone(),
|
||||
content_type: encoded.format.mime_type().to_string(),
|
||||
original_url: Some(parsed.full_url),
|
||||
original_type: Some(result.content_type),
|
||||
original_file_size: Some(original_file_size as i32),
|
||||
original_attachment_id: Some(parsed.attachment_id as i64),
|
||||
file_size: encoded.data.len() as i32,
|
||||
width: encoded.width as i32,
|
||||
height: encoded.height as i32,
|
||||
kind: req.kind,
|
||||
uploaded_at: None,
|
||||
uploaded_by_account: req.uploaded_by.map(|x| x as i64),
|
||||
uploaded_by_system: req.system_id,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(Json(PullResponse {
|
||||
url: final_url,
|
||||
new: is_new,
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn stats(State(state): State<AppState>) -> Result<Json<Stats>, PKAvatarError> {
|
||||
Ok(Json(db::get_stats(&state.pool).await?))
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
bucket: Arc<s3::Bucket>,
|
||||
pull_client: Arc<Client>,
|
||||
pool: PgPool,
|
||||
config: Arc<AvatarsConfig>,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
libpk::init_logging("avatars")?;
|
||||
libpk::init_metrics()?;
|
||||
info!("hello world");
|
||||
|
||||
let config = libpk::config
|
||||
.avatars
|
||||
.as_ref()
|
||||
.expect("missing avatar service config");
|
||||
|
||||
let bucket = {
|
||||
let region = s3::Region::Custom {
|
||||
region: "s3".to_string(),
|
||||
endpoint: config.s3.endpoint.to_string(),
|
||||
};
|
||||
|
||||
let credentials = s3::creds::Credentials::new(
|
||||
Some(&config.s3.application_id),
|
||||
Some(&config.s3.application_key),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let bucket = s3::Bucket::new(&config.s3.bucket, region, credentials)?;
|
||||
|
||||
Arc::new(bucket)
|
||||
};
|
||||
|
||||
let pull_client = Arc::new(
|
||||
ClientBuilder::new()
|
||||
.connect_timeout(Duration::from_secs(3))
|
||||
.timeout(Duration::from_secs(3))
|
||||
.user_agent("PluralKit-Avatars/0.1")
|
||||
.build()
|
||||
.context("error making client")?,
|
||||
);
|
||||
|
||||
let pool = libpk::db::init_data_db().await?;
|
||||
|
||||
let state = AppState {
|
||||
bucket,
|
||||
pull_client,
|
||||
pool,
|
||||
config: Arc::new(config.clone()),
|
||||
};
|
||||
|
||||
// migrations are done, disable this
|
||||
// migrate::spawn_migrate_workers(Arc::new(state.clone()), state.config.migrate_worker_count);
|
||||
|
||||
let app = Router::new()
|
||||
.route("/pull", post(pull))
|
||||
.route("/stats", get(stats))
|
||||
.with_state(state);
|
||||
|
||||
let host = "0.0.0.0:3000";
|
||||
info!("starting server on {}!", host);
|
||||
let listener = tokio::net::TcpListener::bind(host).await.unwrap();
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct AppError(anyhow::Error);
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ErrorResponse {
|
||||
error: String,
|
||||
}
|
||||
|
||||
impl IntoResponse for AppError {
|
||||
fn into_response(self) -> Response {
|
||||
error!("error handling request: {}", self.0);
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ErrorResponse {
|
||||
error: self.0.to_string(),
|
||||
}),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoResponse for PKAvatarError {
|
||||
fn into_response(self) -> Response {
|
||||
let status_code = match self {
|
||||
PKAvatarError::InternalError(_) | PKAvatarError::NetworkError(_) => {
|
||||
StatusCode::INTERNAL_SERVER_ERROR
|
||||
}
|
||||
_ => StatusCode::BAD_REQUEST,
|
||||
};
|
||||
|
||||
// print inner error if otherwise hidden
|
||||
error!("error: {}", self.source().unwrap_or(&self));
|
||||
|
||||
(
|
||||
status_code,
|
||||
Json(ErrorResponse {
|
||||
error: self.to_string(),
|
||||
}),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
|
||||
impl<E> From<E> for AppError
|
||||
where
|
||||
E: Into<anyhow::Error>,
|
||||
{
|
||||
fn from(err: E) -> Self {
|
||||
Self(err.into())
|
||||
}
|
||||
}
|
||||
146
services/avatars/src/migrate.rs
Normal file
146
services/avatars/src/migrate.rs
Normal file
|
|
@ -0,0 +1,146 @@
|
|||
use crate::pull::parse_url;
|
||||
use crate::{db, process, AppState, PKAvatarError};
|
||||
use libpk::db::types::avatars::{ImageMeta, ImageQueueEntry};
|
||||
use reqwest::StatusCode;
|
||||
use std::error::Error;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use time::Instant;
|
||||
use tokio::sync::Semaphore;
|
||||
use tracing::{error, info, instrument, warn};
|
||||
|
||||
static PROCESS_SEMAPHORE: Semaphore = Semaphore::const_new(100);
|
||||
|
||||
pub async fn handle_item_inner(
|
||||
state: &AppState,
|
||||
item: &ImageQueueEntry,
|
||||
) -> Result<(), PKAvatarError> {
|
||||
let parsed = parse_url(&item.url).map_err(|_| PKAvatarError::InvalidCdnUrl)?;
|
||||
|
||||
if let Some(_) = db::get_by_attachment_id(&state.pool, parsed.attachment_id).await? {
|
||||
info!(
|
||||
"attachment {} already migrated, skipping",
|
||||
parsed.attachment_id
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let pulled = crate::pull::pull(state.pull_client.clone(), &parsed).await?;
|
||||
let data_len = pulled.data.len();
|
||||
|
||||
let encoded = {
|
||||
// Trying to reduce CPU load/potentially blocking the worker by adding a bottleneck on parallel encodes
|
||||
// no semaphore on the main api though, that one should ideally be low latency
|
||||
// todo: configurable?
|
||||
let time_before_semaphore = Instant::now();
|
||||
let permit = PROCESS_SEMAPHORE
|
||||
.acquire()
|
||||
.await
|
||||
.map_err(|e| PKAvatarError::InternalError(e.into()))?;
|
||||
let time_after_semaphore = Instant::now();
|
||||
let semaphore_time = time_after_semaphore - time_before_semaphore;
|
||||
if semaphore_time.whole_milliseconds() > 100 {
|
||||
warn!(
|
||||
"waited more than {} ms for process semaphore",
|
||||
semaphore_time.whole_milliseconds()
|
||||
);
|
||||
}
|
||||
|
||||
let encoded = process::process_async(pulled.data, item.kind).await?;
|
||||
drop(permit);
|
||||
encoded
|
||||
};
|
||||
let store_res = crate::store::store(&state.bucket, &encoded).await?;
|
||||
let final_url = format!("{}{}", state.config.cdn_url, store_res.path);
|
||||
|
||||
db::add_image(
|
||||
&state.pool,
|
||||
ImageMeta {
|
||||
id: store_res.id,
|
||||
url: final_url.clone(),
|
||||
content_type: encoded.format.mime_type().to_string(),
|
||||
original_url: Some(parsed.full_url),
|
||||
original_type: Some(pulled.content_type),
|
||||
original_file_size: Some(data_len as i32),
|
||||
original_attachment_id: Some(parsed.attachment_id as i64),
|
||||
file_size: encoded.data.len() as i32,
|
||||
width: encoded.width as i32,
|
||||
height: encoded.height as i32,
|
||||
kind: item.kind,
|
||||
uploaded_at: None,
|
||||
uploaded_by_account: None,
|
||||
uploaded_by_system: None,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
|
||||
info!(
|
||||
"migrated {} ({}k -> {}k)",
|
||||
final_url,
|
||||
data_len,
|
||||
encoded.data.len()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn handle_item(state: &AppState) -> Result<(), PKAvatarError> {
|
||||
// let queue_length = db::get_queue_length(&state.pool).await?;
|
||||
// info!("migrate queue length: {}", queue_length);
|
||||
|
||||
if let Some((mut tx, item)) = db::pop_queue(&state.pool).await? {
|
||||
match handle_item_inner(state, &item).await {
|
||||
Ok(_) => {
|
||||
tx.commit().await.map_err(Into::<anyhow::Error>::into)?;
|
||||
Ok(())
|
||||
}
|
||||
Err(
|
||||
// Errors that mean the image can't be migrated and doesn't need to be retried
|
||||
e @ (PKAvatarError::ImageDimensionsTooLarge(_, _)
|
||||
| PKAvatarError::UnknownImageFormat
|
||||
| PKAvatarError::UnsupportedImageFormat(_)
|
||||
| PKAvatarError::UnsupportedContentType(_)
|
||||
| PKAvatarError::ImageFileSizeTooLarge(_, _)
|
||||
| PKAvatarError::InvalidCdnUrl
|
||||
| PKAvatarError::BadCdnResponse(StatusCode::NOT_FOUND | StatusCode::FORBIDDEN)),
|
||||
) => {
|
||||
warn!("error migrating {}, skipping: {}", item.url, e);
|
||||
tx.commit().await.map_err(Into::<anyhow::Error>::into)?;
|
||||
Ok(())
|
||||
}
|
||||
Err(e @ PKAvatarError::ImageFormatError(_)) => {
|
||||
// will add this item back to the end of the queue
|
||||
db::push_queue(&mut *tx, &item.url, item.kind).await?;
|
||||
tx.commit().await.map_err(Into::<anyhow::Error>::into)?;
|
||||
Err(e)
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
} else {
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[instrument(skip(state))]
|
||||
pub async fn worker(worker_id: u32, state: Arc<AppState>) {
|
||||
info!("spawned migrate worker with id {}", worker_id);
|
||||
loop {
|
||||
match handle_item(&state).await {
|
||||
Ok(()) => {}
|
||||
Err(e) => {
|
||||
error!(
|
||||
"error in migrate worker {}: {}",
|
||||
worker_id,
|
||||
e.source().unwrap_or(&e)
|
||||
);
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn spawn_migrate_workers(state: Arc<AppState>, count: u32) {
|
||||
for i in 0..count {
|
||||
tokio::spawn(worker(i, state.clone()));
|
||||
}
|
||||
}
|
||||
257
services/avatars/src/process.rs
Normal file
257
services/avatars/src/process.rs
Normal file
|
|
@ -0,0 +1,257 @@
|
|||
use image::{DynamicImage, ImageFormat};
|
||||
use std::borrow::Cow;
|
||||
use std::io::Cursor;
|
||||
use time::Instant;
|
||||
use tracing::{debug, error, info, instrument};
|
||||
|
||||
use crate::{hash::Hash, ImageKind, PKAvatarError};
|
||||
|
||||
const MAX_DIMENSION: u32 = 4000;
|
||||
|
||||
pub struct ProcessOutput {
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
pub hash: Hash,
|
||||
pub format: ProcessedFormat,
|
||||
pub data: Vec<u8>,
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub enum ProcessedFormat {
|
||||
Webp,
|
||||
Gif,
|
||||
}
|
||||
|
||||
impl ProcessedFormat {
|
||||
pub fn mime_type(&self) -> &'static str {
|
||||
match self {
|
||||
ProcessedFormat::Gif => "image/gif",
|
||||
ProcessedFormat::Webp => "image/webp",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn extension(&self) -> &'static str {
|
||||
match self {
|
||||
ProcessedFormat::Webp => "webp",
|
||||
ProcessedFormat::Gif => "gif",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Moving Vec<u8> in here since the thread needs ownership of it now, it's fine, don't need it after
|
||||
pub async fn process_async(data: Vec<u8>, kind: ImageKind) -> Result<ProcessOutput, PKAvatarError> {
|
||||
tokio::task::spawn_blocking(move || process(&data, kind))
|
||||
.await
|
||||
.map_err(|je| PKAvatarError::InternalError(je.into()))?
|
||||
}
|
||||
|
||||
#[instrument(skip_all)]
|
||||
pub fn process(data: &[u8], kind: ImageKind) -> Result<ProcessOutput, PKAvatarError> {
|
||||
let time_before = Instant::now();
|
||||
let reader = reader_for(data);
|
||||
match reader.format() {
|
||||
Some(ImageFormat::Png | ImageFormat::WebP | ImageFormat::Jpeg | ImageFormat::Tiff) => {} // ok :)
|
||||
Some(ImageFormat::Gif) => {
|
||||
// animated gifs will need to be handled totally differently
|
||||
// so split off processing here and come back if it's not applicable
|
||||
// (non-banner gifs + 1-frame animated gifs still need to be webp'd)
|
||||
if let Some(output) = process_gif(data, kind)? {
|
||||
return Ok(output);
|
||||
}
|
||||
}
|
||||
Some(other) => return Err(PKAvatarError::UnsupportedImageFormat(other)),
|
||||
None => return Err(PKAvatarError::UnknownImageFormat),
|
||||
}
|
||||
|
||||
// want to check dimensions *before* decoding so we don't accidentally end up with a memory bomb
|
||||
// eg. a 16000x16000 png file is only 31kb and expands to almost a gig of memory
|
||||
let (width, height) = assert_dimensions(reader.into_dimensions()?)?;
|
||||
|
||||
// need to make a new reader??? why can't it just use the same one. reduce duplication?
|
||||
let reader = reader_for(data);
|
||||
|
||||
let time_after_parse = Instant::now();
|
||||
|
||||
// apparently `image` sometimes decodes webp images wrong/weird.
|
||||
// see: https://discord.com/channels/466707357099884544/667795132971614229/1209925940835262464
|
||||
// instead, for webp, we use libwebp itself to decode, as well.
|
||||
// (pls no cve)
|
||||
let image = if reader.format() == Some(ImageFormat::WebP) {
|
||||
let webp_image = webp::Decoder::new(data).decode().ok_or_else(|| {
|
||||
PKAvatarError::InternalError(anyhow::anyhow!("webp decode failed").into())
|
||||
})?;
|
||||
webp_image.to_image()
|
||||
} else {
|
||||
reader.decode().map_err(|e| {
|
||||
// print the ugly error, return the nice error
|
||||
error!("error decoding image: {}", e);
|
||||
PKAvatarError::ImageFormatError(e)
|
||||
})?
|
||||
};
|
||||
|
||||
let time_after_decode = Instant::now();
|
||||
let image = resize(image, kind);
|
||||
let time_after_resize = Instant::now();
|
||||
|
||||
let encoded = encode(image);
|
||||
let time_after = Instant::now();
|
||||
|
||||
info!(
|
||||
"{}: lossy size {}K (parse: {} ms, decode: {} ms, resize: {} ms, encode: {} ms)",
|
||||
encoded.hash,
|
||||
encoded.data.len() / 1024,
|
||||
(time_after_parse - time_before).whole_milliseconds(),
|
||||
(time_after_decode - time_after_parse).whole_milliseconds(),
|
||||
(time_after_resize - time_after_decode).whole_milliseconds(),
|
||||
(time_after - time_after_resize).whole_milliseconds(),
|
||||
);
|
||||
|
||||
debug!(
|
||||
"processed image {}: {} bytes, {}x{} -> {} bytes, {}x{}",
|
||||
encoded.hash,
|
||||
data.len(),
|
||||
width,
|
||||
height,
|
||||
encoded.data.len(),
|
||||
encoded.width,
|
||||
encoded.height
|
||||
);
|
||||
Ok(encoded)
|
||||
}
|
||||
|
||||
fn assert_dimensions((width, height): (u32, u32)) -> Result<(u32, u32), PKAvatarError> {
|
||||
if width > MAX_DIMENSION || height > MAX_DIMENSION {
|
||||
return Err(PKAvatarError::ImageDimensionsTooLarge(
|
||||
(width, height),
|
||||
(MAX_DIMENSION, MAX_DIMENSION),
|
||||
));
|
||||
}
|
||||
return Ok((width, height));
|
||||
}
|
||||
fn process_gif(input_data: &[u8], kind: ImageKind) -> Result<Option<ProcessOutput>, PKAvatarError> {
|
||||
// gifs only supported for banners
|
||||
if kind != ImageKind::Banner {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// and we can't rescale gifs (i tried :/) so the max size is the real limit
|
||||
if kind != ImageKind::Banner {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let reader = gif::Decoder::new(Cursor::new(input_data)).map_err(Into::<anyhow::Error>::into)?;
|
||||
let (max_width, max_height) = kind.size();
|
||||
if reader.width() as u32 > max_width || reader.height() as u32 > max_height {
|
||||
return Err(PKAvatarError::ImageDimensionsTooLarge(
|
||||
(reader.width() as u32, reader.height() as u32),
|
||||
(max_width, max_height),
|
||||
));
|
||||
}
|
||||
Ok(process_gif_inner(reader).map_err(Into::<anyhow::Error>::into)?)
|
||||
}
|
||||
|
||||
fn process_gif_inner(
|
||||
mut reader: gif::Decoder<Cursor<&[u8]>>,
|
||||
) -> Result<Option<ProcessOutput>, anyhow::Error> {
|
||||
let time_before = Instant::now();
|
||||
|
||||
let (width, height) = (reader.width(), reader.height());
|
||||
|
||||
let mut writer = gif::Encoder::new(
|
||||
Vec::new(),
|
||||
width as u16,
|
||||
height as u16,
|
||||
reader.global_palette().unwrap_or(&[]),
|
||||
)?;
|
||||
writer.set_repeat(reader.repeat())?;
|
||||
|
||||
let mut frame_buf = Vec::new();
|
||||
|
||||
let mut frame_count = 0;
|
||||
while let Some(frame) = reader.next_frame_info()? {
|
||||
let mut frame = frame.clone();
|
||||
assert_dimensions((frame.width as u32, frame.height as u32))?;
|
||||
frame_buf.clear();
|
||||
frame_buf.resize(reader.buffer_size(), 0);
|
||||
reader.read_into_buffer(&mut frame_buf)?;
|
||||
frame.buffer = Cow::Borrowed(&frame_buf);
|
||||
|
||||
frame.make_lzw_pre_encoded();
|
||||
writer.write_lzw_pre_encoded_frame(&frame)?;
|
||||
frame_count += 1;
|
||||
}
|
||||
|
||||
if frame_count == 1 {
|
||||
// If there's only one frame, then this doesn't need to be a gif. webp it
|
||||
// (unfortunately we can't tell if there's only one frame until after the first frame's been decoded...)
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let data = writer.into_inner()?;
|
||||
let time_after = Instant::now();
|
||||
|
||||
let hash = Hash::sha256(&data);
|
||||
|
||||
let original_data = reader.into_inner();
|
||||
info!(
|
||||
"processed gif {}: {}K -> {}K ({} ms, frames: {})",
|
||||
hash,
|
||||
original_data.buffer().len() / 1024,
|
||||
data.len() / 1024,
|
||||
(time_after - time_before).whole_milliseconds(),
|
||||
frame_count
|
||||
);
|
||||
|
||||
Ok(Some(ProcessOutput {
|
||||
data,
|
||||
format: ProcessedFormat::Gif,
|
||||
hash,
|
||||
width: width as u32,
|
||||
height: height as u32,
|
||||
}))
|
||||
}
|
||||
|
||||
fn reader_for(data: &[u8]) -> image::io::Reader<Cursor<&[u8]>> {
|
||||
image::io::Reader::new(Cursor::new(data))
|
||||
.with_guessed_format()
|
||||
.expect("cursor i/o is infallible")
|
||||
}
|
||||
|
||||
#[instrument(skip_all)]
|
||||
fn resize(image: DynamicImage, kind: ImageKind) -> DynamicImage {
|
||||
let (target_width, target_height) = kind.size();
|
||||
if image.width() <= target_width && image.height() <= target_height {
|
||||
// don't resize if already smaller
|
||||
return image;
|
||||
}
|
||||
|
||||
// todo: best filter?
|
||||
let resized = image.resize(
|
||||
target_width,
|
||||
target_height,
|
||||
image::imageops::FilterType::Lanczos3,
|
||||
);
|
||||
return resized;
|
||||
}
|
||||
|
||||
#[instrument(skip_all)]
|
||||
// can't believe this is infallible
|
||||
fn encode(image: DynamicImage) -> ProcessOutput {
|
||||
let (width, height) = (image.width(), image.height());
|
||||
let image_buf = image.to_rgba8();
|
||||
|
||||
let encoded_lossy = webp::Encoder::new(&*image_buf, webp::PixelLayout::Rgba, width, height)
|
||||
.encode_simple(false, 90.0)
|
||||
.expect("encode should be infallible")
|
||||
.to_vec();
|
||||
|
||||
let hash = Hash::sha256(&encoded_lossy);
|
||||
|
||||
ProcessOutput {
|
||||
data: encoded_lossy,
|
||||
format: ProcessedFormat::Webp,
|
||||
hash,
|
||||
width,
|
||||
height,
|
||||
}
|
||||
}
|
||||
166
services/avatars/src/pull.rs
Normal file
166
services/avatars/src/pull.rs
Normal file
|
|
@ -0,0 +1,166 @@
|
|||
use std::time::Duration;
|
||||
use std::{str::FromStr, sync::Arc};
|
||||
|
||||
use crate::PKAvatarError;
|
||||
use anyhow::Context;
|
||||
use reqwest::{Client, ClientBuilder, StatusCode, Url};
|
||||
use time::Instant;
|
||||
use tracing::{error, instrument};
|
||||
|
||||
const MAX_SIZE: u64 = 8 * 1024 * 1024;
|
||||
|
||||
pub struct PullResult {
|
||||
pub data: Vec<u8>,
|
||||
pub content_type: String,
|
||||
pub last_modified: Option<String>,
|
||||
}
|
||||
|
||||
#[instrument(skip_all)]
|
||||
pub async fn pull(
|
||||
client: Arc<Client>,
|
||||
parsed_url: &ParsedUrl,
|
||||
) -> Result<PullResult, PKAvatarError> {
|
||||
let time_before = Instant::now();
|
||||
let mut trimmed_url = trim_url_query(&parsed_url.full_url)?;
|
||||
if trimmed_url.host_str() == Some("media.discordapp.net") {
|
||||
trimmed_url
|
||||
.set_host(Some("cdn.discordapp.com"))
|
||||
.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)
|
||||
})?;
|
||||
let time_after_headers = Instant::now();
|
||||
let status = response.status();
|
||||
|
||||
if status != StatusCode::OK {
|
||||
return Err(PKAvatarError::BadCdnResponse(status));
|
||||
}
|
||||
|
||||
let size = match response.content_length() {
|
||||
None => return Err(PKAvatarError::MissingHeader("Content-Length")),
|
||||
Some(size) if size > MAX_SIZE => {
|
||||
return Err(PKAvatarError::ImageFileSizeTooLarge(size, MAX_SIZE))
|
||||
}
|
||||
Some(size) => size,
|
||||
};
|
||||
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_TYPE)
|
||||
.and_then(|x| x.to_str().ok()) // invalid (non-unicode) header = missing, why not
|
||||
.map(|mime| mime.split(';').next().unwrap_or("")) // cut off at ;
|
||||
.ok_or(PKAvatarError::MissingHeader("Content-Type"))?
|
||||
.to_owned();
|
||||
let mime = match content_type.as_str() {
|
||||
mime @ ("image/jpeg" | "image/png" | "image/gif" | "image/webp" | "image/tiff") => mime,
|
||||
_ => return Err(PKAvatarError::UnsupportedContentType(content_type)),
|
||||
};
|
||||
|
||||
let last_modified = response
|
||||
.headers()
|
||||
.get(reqwest::header::LAST_MODIFIED)
|
||||
.and_then(|x| x.to_str().ok())
|
||||
.map(|x| x.to_string());
|
||||
|
||||
let body = response.bytes().await.map_err(|e| {
|
||||
error!("network error for {}: {}", parsed_url.full_url, e);
|
||||
PKAvatarError::NetworkError(e)
|
||||
})?;
|
||||
if body.len() != size as usize {
|
||||
// ???does this ever happen?
|
||||
return Err(PKAvatarError::InternalError(anyhow::anyhow!(
|
||||
"server responded with wrong length"
|
||||
)));
|
||||
}
|
||||
let time_after_body = Instant::now();
|
||||
|
||||
let headers_time = time_after_headers - time_before;
|
||||
let body_time = time_after_body - time_after_headers;
|
||||
|
||||
// can't do dynamic log level lmao
|
||||
if status != StatusCode::OK {
|
||||
tracing::warn!(
|
||||
"{}: {} (headers: {}ms, body: {}ms)",
|
||||
status,
|
||||
&trimmed_url,
|
||||
headers_time.whole_milliseconds(),
|
||||
body_time.whole_milliseconds()
|
||||
);
|
||||
} else {
|
||||
tracing::info!(
|
||||
"{}: {} (headers: {}ms, body: {}ms)",
|
||||
status,
|
||||
&trimmed_url,
|
||||
headers_time.whole_milliseconds(),
|
||||
body_time.whole_milliseconds()
|
||||
);
|
||||
};
|
||||
|
||||
Ok(PullResult {
|
||||
data: body.to_vec(),
|
||||
content_type: mime.to_string(),
|
||||
last_modified,
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ParsedUrl {
|
||||
pub channel_id: u64,
|
||||
pub attachment_id: u64,
|
||||
pub filename: String,
|
||||
pub full_url: String,
|
||||
}
|
||||
|
||||
pub fn parse_url(url: &str) -> anyhow::Result<ParsedUrl> {
|
||||
// todo: should this return PKAvatarError::InvalidCdnUrl?
|
||||
let url = Url::from_str(url).context("invalid url")?;
|
||||
|
||||
match (url.scheme(), url.domain()) {
|
||||
("https", Some("media.discordapp.net" | "cdn.discordapp.com")) => {}
|
||||
_ => anyhow::bail!("not a discord cdn url"),
|
||||
}
|
||||
|
||||
match url
|
||||
.path_segments()
|
||||
.map(|x| x.collect::<Vec<_>>())
|
||||
.as_deref()
|
||||
{
|
||||
Some([_, channel_id, attachment_id, filename]) => {
|
||||
let channel_id = u64::from_str(channel_id).context("invalid channel id")?;
|
||||
let attachment_id = u64::from_str(attachment_id).context("invalid channel id")?;
|
||||
|
||||
Ok(ParsedUrl {
|
||||
channel_id,
|
||||
attachment_id,
|
||||
filename: filename.to_string(),
|
||||
full_url: url.to_string(),
|
||||
})
|
||||
}
|
||||
_ => anyhow::bail!("invaild discord cdn url"),
|
||||
}
|
||||
}
|
||||
|
||||
fn trim_url_query(url: &str) -> anyhow::Result<Url> {
|
||||
let mut parsed = Url::parse(url)?;
|
||||
|
||||
let mut qs = form_urlencoded::Serializer::new(String::new());
|
||||
for (key, value) in parsed.query_pairs() {
|
||||
match key.as_ref() {
|
||||
"ex" | "is" | "hm" => {
|
||||
qs.append_pair(key.as_ref(), value.as_ref());
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
let new_query = qs.finish();
|
||||
parsed.set_query(if new_query.len() > 0 {
|
||||
Some(&new_query)
|
||||
} else {
|
||||
None
|
||||
});
|
||||
|
||||
Ok(parsed)
|
||||
}
|
||||
60
services/avatars/src/store.rs
Normal file
60
services/avatars/src/store.rs
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
use crate::process::ProcessOutput;
|
||||
use tracing::error;
|
||||
|
||||
pub struct StoreResult {
|
||||
pub id: String,
|
||||
pub path: String,
|
||||
}
|
||||
|
||||
pub async fn store(bucket: &s3::Bucket, res: &ProcessOutput) -> anyhow::Result<StoreResult> {
|
||||
// errors here are all going to be internal
|
||||
let encoded_hash = res.hash.to_string();
|
||||
let path = format!(
|
||||
"images/{}/{}.{}",
|
||||
&encoded_hash[..2],
|
||||
&encoded_hash[2..],
|
||||
res.format.extension()
|
||||
);
|
||||
|
||||
// todo: something better than these retries
|
||||
let mut retry_count = 0;
|
||||
loop {
|
||||
if retry_count == 2 {
|
||||
tokio::time::sleep(tokio::time::Duration::new(2, 0)).await;
|
||||
}
|
||||
if retry_count > 2 {
|
||||
anyhow::bail!("error uploading image to cdn, too many retries") // nicer user-facing error?
|
||||
}
|
||||
retry_count += 1;
|
||||
|
||||
let resp = bucket
|
||||
.put_object_with_content_type(&path, &res.data, res.format.mime_type())
|
||||
.await?;
|
||||
match resp.status_code() {
|
||||
200 => {
|
||||
tracing::debug!("uploaded image to {}", &path);
|
||||
|
||||
return Ok(StoreResult {
|
||||
id: encoded_hash,
|
||||
path,
|
||||
});
|
||||
}
|
||||
500 | 503 => {
|
||||
tracing::warn!(
|
||||
"got 503 uploading image to {} ({}), retrying... (try {}/3)",
|
||||
&path,
|
||||
resp.as_str()?,
|
||||
retry_count
|
||||
);
|
||||
continue;
|
||||
}
|
||||
_ => {
|
||||
error!(
|
||||
"storage backend responded status code {}",
|
||||
resp.status_code()
|
||||
);
|
||||
anyhow::bail!("error uploading image to cdn") // nicer user-facing error?
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -13,4 +13,4 @@ tracing = { workspace = true }
|
|||
tracing-subscriber = { workspace = true }
|
||||
|
||||
hickory-client = "0.24.1"
|
||||
reqwest = { version = "0.12.7", default-features = false, features = ["rustls-tls"] }
|
||||
reqwest = { workspace = true }
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ pub async fn run_server(cache: Arc<DiscordCache>) -> anyhow::Result<()> {
|
|||
.route(
|
||||
"/guilds/:guild_id/members/@me",
|
||||
get(|State(cache): State<Arc<DiscordCache>>, Path(guild_id): Path<u64>| async move {
|
||||
match cache.0.member(Id::new(guild_id), libpk::config.discord.client_id) {
|
||||
match cache.0.member(Id::new(guild_id), libpk::config.discord.as_ref().expect("missing discord config").client_id) {
|
||||
Some(member) => status_code(StatusCode::FOUND, to_string(member.value()).unwrap()),
|
||||
None => status_code(StatusCode::NOT_FOUND, "".to_string()),
|
||||
}
|
||||
|
|
@ -45,7 +45,7 @@ pub async fn run_server(cache: Arc<DiscordCache>) -> anyhow::Result<()> {
|
|||
.route(
|
||||
"/guilds/:guild_id/permissions/@me",
|
||||
get(|State(cache): State<Arc<DiscordCache>>, Path(guild_id): Path<u64>| async move {
|
||||
match cache.guild_permissions(Id::new(guild_id), libpk::config.discord.client_id).await {
|
||||
match cache.guild_permissions(Id::new(guild_id), libpk::config.discord.as_ref().expect("missing discord config").client_id).await {
|
||||
Ok(val) => {
|
||||
println!("hh {}", Permissions::all().bits());
|
||||
status_code(StatusCode::FOUND, to_string(&val.bits()).unwrap())
|
||||
|
|
@ -114,7 +114,7 @@ pub async fn run_server(cache: Arc<DiscordCache>) -> anyhow::Result<()> {
|
|||
if guild_id == 0 {
|
||||
return status_code(StatusCode::FOUND, to_string(&*DM_PERMISSIONS).unwrap());
|
||||
}
|
||||
match cache.channel_permissions(Id::new(channel_id), libpk::config.discord.client_id).await {
|
||||
match cache.channel_permissions(Id::new(channel_id), libpk::config.discord.as_ref().expect("missing discord config").client_id).await {
|
||||
Ok(val) => status_code(StatusCode::FOUND, to_string(&val).unwrap()),
|
||||
Err(err) => {
|
||||
error!(?err, ?channel_id, ?guild_id, "failed to get own channelpermissions");
|
||||
|
|
@ -176,7 +176,7 @@ pub async fn run_server(cache: Arc<DiscordCache>) -> anyhow::Result<()> {
|
|||
.layer(axum::middleware::from_fn(crate::logger::logger))
|
||||
.with_state(cache);
|
||||
|
||||
let addr: &str = libpk::config.api.addr.as_ref();
|
||||
let addr: &str = libpk::config.discord.as_ref().expect("missing discord config").cache_api_addr.as_ref();
|
||||
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||||
info!("listening on {}", addr);
|
||||
axum::serve(listener, app).await?;
|
||||
|
|
|
|||
|
|
@ -89,10 +89,22 @@ fn member_to_cached_member(item: Member, id: Id<UserMarker>) -> CachedMember {
|
|||
}
|
||||
|
||||
pub fn new() -> DiscordCache {
|
||||
let mut client_builder =
|
||||
twilight_http::Client::builder().token(libpk::config.discord.bot_token.clone());
|
||||
let mut client_builder = twilight_http::Client::builder().token(
|
||||
libpk::config
|
||||
.discord
|
||||
.as_ref()
|
||||
.expect("missing discord config")
|
||||
.bot_token
|
||||
.clone(),
|
||||
);
|
||||
|
||||
if let Some(base_url) = libpk::config.discord.api_base_url.clone() {
|
||||
if let Some(base_url) = libpk::config
|
||||
.discord
|
||||
.as_ref()
|
||||
.expect("missing discord config")
|
||||
.api_base_url
|
||||
.clone()
|
||||
{
|
||||
client_builder = client_builder.proxy(base_url, true);
|
||||
}
|
||||
|
||||
|
|
@ -136,7 +148,13 @@ impl DiscordCache {
|
|||
return Ok(Permissions::all());
|
||||
}
|
||||
|
||||
let member = if user_id == libpk::config.discord.client_id {
|
||||
let member = if user_id
|
||||
== libpk::config
|
||||
.discord
|
||||
.as_ref()
|
||||
.expect("missing discord config")
|
||||
.client_id
|
||||
{
|
||||
self.0
|
||||
.member(guild_id, user_id)
|
||||
.ok_or(format_err!("self member not found"))?
|
||||
|
|
@ -202,7 +220,13 @@ impl DiscordCache {
|
|||
return Ok(Permissions::all());
|
||||
}
|
||||
|
||||
let member = if user_id == libpk::config.discord.client_id {
|
||||
let member = if user_id
|
||||
== libpk::config
|
||||
.discord
|
||||
.as_ref()
|
||||
.expect("missing discord config")
|
||||
.client_id
|
||||
{
|
||||
self.0
|
||||
.member(guild_id, user_id)
|
||||
.ok_or_else(|| {
|
||||
|
|
|
|||
|
|
@ -17,6 +17,8 @@ use super::{cache::DiscordCache, shard_state::ShardStateManager};
|
|||
pub fn cluster_config() -> ClusterSettings {
|
||||
libpk::config
|
||||
.discord
|
||||
.as_ref()
|
||||
.expect("missing discord config")
|
||||
.cluster
|
||||
.clone()
|
||||
.unwrap_or(libpk::_config::ClusterSettings {
|
||||
|
|
@ -51,10 +53,18 @@ pub fn create_shards(redis: fred::pool::RedisPool) -> anyhow::Result<Vec<Shard<R
|
|||
let shards = create_iterator(
|
||||
start_shard..end_shard + 1,
|
||||
cluster_settings.total_shards,
|
||||
ConfigBuilder::new(libpk::config.discord.bot_token.to_owned(), intents)
|
||||
.presence(presence("pk;help", false))
|
||||
.queue(queue.clone())
|
||||
.build(),
|
||||
ConfigBuilder::new(
|
||||
libpk::config
|
||||
.discord
|
||||
.as_ref()
|
||||
.expect("missing discord config")
|
||||
.bot_token
|
||||
.to_owned(),
|
||||
intents,
|
||||
)
|
||||
.presence(presence("pk;help", false))
|
||||
.queue(queue.clone())
|
||||
.build(),
|
||||
|_, builder| builder.build(),
|
||||
);
|
||||
|
||||
|
|
|
|||
|
|
@ -15,7 +15,11 @@ use libpk::util::redis::RedisErrorExt;
|
|||
pub fn new(redis: RedisPool) -> RedisQueue {
|
||||
RedisQueue {
|
||||
redis,
|
||||
concurrency: libpk::config.discord.max_concurrency,
|
||||
concurrency: libpk::config
|
||||
.discord
|
||||
.as_ref()
|
||||
.expect("missing discord config")
|
||||
.max_concurrency,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue