WIP: revise avatars service

crates/avatars/src/cleanup.rs

@@ -1,8 +1,8 @@
 use anyhow::Context;
-use reqwest::{ClientBuilder, StatusCode};
-use sqlx::prelude::FromRow;
-use std::{sync::Arc, time::Duration};
+use reqwest::{ClientBuilder, StatusCode, Url};
+use std::{path::Path, sync::Arc, time::Duration};
 use tracing::{error, info};
+use uuid::Uuid;
 
 #[libpk::main]
 async fn main() -> anyhow::Result<()> {
@@ -36,109 +36,196 @@ async fn main() -> anyhow::Result<()> {
     loop {
         // no infinite loops
         tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
-        match cleanup_job(pool.clone(), bucket.clone()).await {
+        match cleanup_job(pool.clone()).await {
             Ok(()) => {}
             Err(error) => {
                 error!(?error, "failed to run avatar cleanup job");
                 // sentry
             }
         }
+        match cleanup_hash_job(pool.clone(), bucket.clone()).await {
+            Ok(()) => {}
+            Err(error) => {
+                error!(?error, "failed to run hash cleanup job");
+                // sentry
+            }
+        }
     }
 }
 
-#[derive(FromRow)]
+#[derive(sqlx::FromRow)]
 struct CleanupJobEntry {
-    id: String,
+    id: Uuid,
+    system_uuid: Uuid,
 }
 
-async fn cleanup_job(pool: sqlx::PgPool, bucket: Arc<s3::Bucket>) -> anyhow::Result<()> {
+async fn cleanup_job(pool: sqlx::PgPool) -> anyhow::Result<()> {
     let mut tx = pool.begin().await?;
 
-    let image_id: Option<CleanupJobEntry> = sqlx::query_as(
+    let entry: Option<CleanupJobEntry> = sqlx::query_as(
         // no timestamp checking here
         // images are only added to the table after 24h
         r#"
-                select id from image_cleanup_jobs
+                select id, system_uuid from image_cleanup_jobs
                 for update skip locked limit 1;"#,
     )
     .fetch_optional(&mut *tx)
     .await?;
-    if image_id.is_none() {
+    if entry.is_none() {
         info!("no job to run, sleeping for 1 minute");
         tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
         return Ok(());
     }
-    let image_id = image_id.unwrap().id;
+    let entry = entry.unwrap();
+    let image_id = entry.id;
+    let system_uuid = entry.system_uuid;
     info!("got image {image_id}, cleaning up...");
 
-    let image_data = libpk::db::repository::avatars::get_by_id(&pool, image_id.clone()).await?;
-    if image_data.is_none() {
+    let image =
+        libpk::db::repository::avatars::get_by_id(&pool, system_uuid.clone(), image_id.clone())
+            .await?;
+    if image.is_none() {
         // unsure how this can happen? there is a FK reference
         info!("image {image_id} was already deleted, skipping");
-        sqlx::query("delete from image_cleanup_jobs where id = $1")
+        sqlx::query("delete from image_cleanup_jobs where id = $1 and system_uuid = $2")
             .bind(image_id)
+            .bind(system_uuid)
             .execute(&mut *tx)
             .await?;
         return Ok(());
     }
-    let image_data = image_data.unwrap();
+    let image = image.unwrap();
 
     let config = libpk::config
         .avatars
         .as_ref()
         .expect("missing avatar service config");
 
-    let path = image_data
-        .url
-        .strip_prefix(config.cdn_url.as_str())
-        .unwrap();
-
-    let s3_resp = bucket.delete_object(path).await?;
-    match s3_resp.status_code() {
-        204 => {
-            info!("successfully deleted image {image_id} from s3");
-        }
-        _ => {
-            anyhow::bail!("s3 returned bad error code {}", s3_resp.status_code());
-        }
-    }
-
-    if let Some(zone_id) = config.cloudflare_zone_id.as_ref() {
+    if let Some(store_id) = config.fastly_store_id.as_ref() {
         let client = ClientBuilder::new()
             .connect_timeout(Duration::from_secs(3))
             .timeout(Duration::from_secs(3))
             .build()
             .context("error making client")?;
 
-        let cf_resp = client
-            .post(format!(
-                "https://api.cloudflare.com/client/v4/zones/{zone_id}/purge_cache"
+        let url = Url::parse(&image.data.url).expect("invalid url");
+        let extension = Path::new(url.path())
+            .extension()
+            .and_then(|s| s.to_str())
+            .unwrap_or("");
+        let key = format!("{system_uuid}:{image_id}.{extension}");
+
+        let kv_resp = client
+            .delete(format!(
+                "https://api.fastly.com/resources/stores/kv/{store_id}/keys/{key}"
             ))
-            .header(
-                "Authorization",
-                format!("Bearer {}", config.cloudflare_token.as_ref().unwrap()),
-            )
-            .body(format!(r#"{{"files":["{}"]}}"#, image_data.url))
+            .header("Fastly-Key", config.fastly_token.as_ref().unwrap())
             .send()
             .await?;
 
-        match cf_resp.status() {
+        match kv_resp.status() {
             StatusCode::OK => {
                 info!(
-                    "successfully purged url {} from cloudflare cache",
-                    image_data.url
+                    "successfully purged image {}:{}.{} from fastly kv",
+                    system_uuid, image_id, extension
                 );
             }
             _ => {
-                let status = cf_resp.status();
-                tracing::info!("raw response from cloudflare: {:#?}", cf_resp.text().await?);
-                anyhow::bail!("cloudflare returned bad error code {}", status);
+                let status = kv_resp.status();
+                tracing::info!("raw response from fastly: {:#?}", kv_resp.text().await?);
+                tracing::warn!("fastly returned bad error code {}", status);
+            }
+        }
+
+        let cdn_url_parsed = Url::parse(config.cdn_url.as_str())?;
+        let cdn_host = cdn_url_parsed.host_str().unwrap_or(config.cdn_url.as_str());
+
+        let cache_resp = client
+            .post(format!(
+                "https://api.fastly.com/purge/{}/{}/{}.{}",
+                cdn_host, system_uuid, image_id, extension
+            ))
+            .header("Fastly-Key", config.fastly_token.as_ref().unwrap())
+            .send()
+            .await?;
+
+        match cache_resp.status() {
+            StatusCode::OK => {
+                info!(
+                    "successfully purged image {}/{}.{} from fastly cache",
+                    system_uuid, image_id, extension
+                );
+            }
+            _ => {
+                let status = cache_resp.status();
+                tracing::info!("raw response from fastly: {:#?}", cache_resp.text().await?);
+                tracing::warn!("fastly returned bad error code {}", status);
             }
         }
     }
 
-    sqlx::query("delete from images where id = $1")
+    sqlx::query("delete from images_assets where id = $1 and system_uuid = $2")
         .bind(image_id.clone())
+        .bind(system_uuid.clone())
+        .execute(&mut *tx)
+        .await?;
+
+    tx.commit().await?;
+
+    Ok(())
+}
+
+#[derive(sqlx::FromRow)]
+struct HashCleanupJobEntry {
+    hash: String,
+}
+
+async fn cleanup_hash_job(pool: sqlx::PgPool, bucket: Arc<s3::Bucket>) -> anyhow::Result<()> {
+    let mut tx = pool.begin().await?;
+
+    let config = libpk::config
+        .avatars
+        .as_ref()
+        .expect("missing avatar service config");
+
+    let entry: Option<HashCleanupJobEntry> = sqlx::query_as(
+        // no timestamp checking here
+        // images are only added to the table after 24h
+        r#"
+                select hash from image_hash_cleanup_jobs
+                for update skip locked limit 1;"#,
+    )
+    .fetch_optional(&mut *tx)
+    .await?;
+    if entry.is_none() {
+        info!("no hash job to run, sleeping for 1 minute");
+        tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
+        return Ok(());
+    }
+    let entry = entry.unwrap();
+    let hash = entry.hash;
+    info!("got orphaned hash {hash}, cleaning up...");
+
+    let url: Option<String> = sqlx::query_scalar("select url from images_hashes where hash = $1")
+        .bind(&hash)
+        .fetch_optional(&mut *tx)
+        .await?;
+
+    if let Some(url) = url {
+        let path = url.strip_prefix(config.cdn_url.as_str()).unwrap();
+        let s3_resp = bucket.delete_object(path).await?;
+        match s3_resp.status_code() {
+            204 => {
+                info!("successfully deleted image {hash} from s3");
+            }
+            _ => {
+                anyhow::bail!("s3 returned bad error code {}", s3_resp.status_code());
+            }
+        }
+    }
+
+    sqlx::query("delete from images_hashes where hash = $1")
+        .bind(&hash)
         .execute(&mut *tx)
         .await?;
 

crates/avatars/src/main.rs

@@ -5,7 +5,8 @@ mod pull;
 mod store;
 
 use anyhow::Context;
-use axum::extract::State;
+use axum::extract::{DefaultBodyLimit, Multipart, State};
+use axum::http::HeaderMap;
 use axum::routing::get;
 use axum::{
     Json, Router,
@@ -21,12 +22,18 @@ use reqwest::{Client, ClientBuilder};
 use serde::{Deserialize, Serialize};
 use sqlx::PgPool;
 use std::error::Error;
+use std::net::IpAddr;
+use std::str::FromStr;
 use std::sync::Arc;
 use std::time::Duration;
 use thiserror::Error;
 use tracing::{error, info, warn};
 use uuid::Uuid;
 
+const NORMAL_HARD_LIMIT: usize = 8 * 1024 * 1024;
+const PREMIUM_SOFT_LIMIT: usize = 30 * 1024 * 1024;
+const PREMIUM_HARD_LIMIT: usize = 50 * 1024 * 1024;
+
 #[derive(Error, Debug)]
 pub enum PKAvatarError {
     // todo: split off into logical groups (cdn/url error, image format error, etc)
@@ -82,60 +89,130 @@ pub struct PullRequest {
 }
 
 #[derive(Serialize)]
-pub struct PullResponse {
+pub struct ImageResponse {
     url: String,
     new: bool,
 }
 
+async fn gen_proxy_image(state: &AppState, data: Vec<u8>) -> Result<Option<String>, PKAvatarError> {
+    let encoded_proxy = process::process_async(data, ImageKind::Avatar).await?;
+    let store_proxy_res = crate::store::store(&state.bucket, &encoded_proxy).await?;
+    let proxy_url = format!("{}{}", state.config.cdn_url, store_proxy_res.path);
+    db::add_image_data(
+        &state.pool,
+        &ImageData {
+            hash: encoded_proxy.hash.to_string(),
+            url: proxy_url,
+            file_size: encoded_proxy.data.len() as i32,
+            width: encoded_proxy.width as i32,
+            height: encoded_proxy.height as i32,
+            content_type: encoded_proxy.format.to_mime_type().to_string(),
+            created_at: None,
+        },
+    )
+    .await?;
+    Ok(Some(encoded_proxy.hash.to_string()))
+}
+
+async fn handle_image(
+    state: &AppState,
+    data: Vec<u8>,
+    mut meta: ImageMeta,
+) -> Result<ImageResponse, PKAvatarError> {
+    let original_file_size = data.len();
+    let system_uuid = meta.system_uuid;
+
+    if meta.kind.is_premium() && original_file_size > NORMAL_HARD_LIMIT {
+        meta.proxy_image = gen_proxy_image(&state, data.clone()).await?;
+    }
+
+    let encoded = process::process_async(data, meta.kind).await?;
+    let store_res = crate::store::store(&state.bucket, &encoded).await?;
+    meta.image = store_res.id.clone();
+    let storage_url = format!("{}{}", state.config.cdn_url, store_res.path);
+
+    let res = db::add_image(
+        &state.pool,
+        Image {
+            meta: meta,
+            data: ImageData {
+                hash: store_res.id,
+                url: storage_url,
+                file_size: encoded.data.len() as i32,
+                width: encoded.width as i32,
+                height: encoded.height as i32,
+                content_type: encoded.format.to_mime_type().to_string(),
+                created_at: None,
+            },
+        },
+    )
+    .await?;
+
+    if original_file_size >= PREMIUM_SOFT_LIMIT {
+        warn!(
+            "large image {} of size {} uploaded",
+            res.uuid, original_file_size
+        )
+    }
+
+    let final_url = format!(
+        "{}images/{}/{}.{}",
+        state.config.edge_url,
+        system_uuid,
+        res.uuid,
+        encoded
+            .format
+            .extensions_str()
+            .first()
+            .expect("expected valid extension")
+    );
+
+    Ok(ImageResponse {
+        url: final_url,
+        new: res.is_new,
+    })
+}
+
 async fn pull(
     State(state): State<AppState>,
     Json(req): Json<PullRequest>,
-) -> Result<Json<PullResponse>, PKAvatarError> {
+) -> Result<Json<ImageResponse>, PKAvatarError> {
     let parsed = pull::parse_url(&req.url) // parsing beforehand to "normalize"
         .map_err(|_| PKAvatarError::InvalidCdnUrl)?;
     if !(req.force || req.url.contains("https://serve.apparyllis.com/")) {
         if let Some(existing) = db::get_by_attachment_id(&state.pool, parsed.attachment_id).await? {
             // remove any pending image cleanup
             db::remove_deletion_queue(&state.pool, parsed.attachment_id).await?;
-            return Ok(Json(PullResponse {
-                url: existing.url,
+            return Ok(Json(ImageResponse {
+                url: existing.data.url,
                 new: false,
             }));
         }
     }
-
-    let result = crate::pull::pull(state.pull_client, &parsed).await?;
-
+    let result = crate::pull::pull(&state.pull_client, &parsed, req.kind.is_premium()).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,
-    }))
+    Ok(Json(
+        handle_image(
+            &state,
+            result.data,
+            ImageMeta {
+                id: Uuid::default(),
+                system_uuid: req.system_id.expect("expected system id"),
+                image: "".to_string(),
+                proxy_image: None,
+                kind: req.kind,
+                original_url: Some(parsed.full_url),
+                original_file_size: Some(original_file_size as i32),
+                original_type: Some(result.content_type),
+                original_attachment_id: Some(parsed.attachment_id as i64),
+                uploaded_by_account: req.uploaded_by.map(|x| x as i64),
+                uploaded_by_ip: None,
+                uploaded_at: None,
+            },
+        )
+        .await?,
+    ))
 }
 
 async fn verify(
@@ -143,13 +220,14 @@ async fn verify(
     Json(req): Json<PullRequest>,
 ) -> Result<(), PKAvatarError> {
     let result = crate::pull::pull(
-        state.pull_client,
+        &state.pull_client,
         &ParsedUrl {
             full_url: req.url.clone(),
             channel_id: 0,
             attachment_id: 0,
             filename: "".to_string(),
         },
+        req.kind.is_premium(),
     )
     .await?;
 
@@ -158,6 +236,81 @@ async fn verify(
     Ok(())
 }
 
+async fn upload(
+    State(state): State<AppState>,
+    headers: HeaderMap,
+    mut multipart: Multipart,
+) -> Result<Json<ImageResponse>, PKAvatarError> {
+    let mut data: Option<Vec<u8>> = None;
+    let mut kind: Option<ImageKind> = None;
+    let mut system_id: Option<Uuid> = None;
+    let mut upload_ip: Option<IpAddr> = None;
+
+    if let Some(val) = headers.get("x-pluralkit-systemuuid")
+        && let Ok(s) = val.to_str()
+    {
+        system_id = Uuid::parse_str(s).ok();
+    }
+    if let Some(val) = headers.get("x-pluralkit-client-ip")
+        && let Ok(s) = val.to_str()
+    {
+        upload_ip = IpAddr::from_str(s).ok();
+    }
+
+    while let Some(field) = multipart
+        .next_field()
+        .await
+        .map_err(|e| PKAvatarError::InternalError(e.into()))?
+    {
+        let name = field.name().unwrap_or("").to_string();
+
+        match name.as_str() {
+            "file" => {
+                let bytes = field
+                    .bytes()
+                    .await
+                    .map_err(|e| PKAvatarError::InternalError(e.into()))?;
+                data = Some(bytes.to_vec());
+            }
+            "kind" => {
+                let txt = field
+                    .text()
+                    .await
+                    .map_err(|e| PKAvatarError::InternalError(e.into()))?;
+                kind = ImageKind::from_string(&txt);
+            }
+            _ => {}
+        }
+    }
+
+    let data = data.ok_or(PKAvatarError::MissingHeader("file"))?;
+    let kind = kind.ok_or(PKAvatarError::MissingHeader("kind"))?;
+    let system_id = system_id.ok_or(PKAvatarError::MissingHeader("x-pluralkit-systemuuid"))?;
+    let upload_ip = upload_ip.ok_or(PKAvatarError::MissingHeader("x-pluralkit-client-ip"))?;
+
+    Ok(Json(
+        handle_image(
+            &state,
+            data,
+            ImageMeta {
+                id: Uuid::default(),
+                system_uuid: system_id,
+                image: "".to_string(),
+                proxy_image: None,
+                kind: kind,
+                original_url: None,
+                original_file_size: None,
+                original_type: None,
+                original_attachment_id: None,
+                uploaded_by_account: None,
+                uploaded_by_ip: Some(upload_ip),
+                uploaded_at: None,
+            },
+        )
+        .await?,
+    ))
+}
+
 pub async fn stats(State(state): State<AppState>) -> Result<Json<Stats>, PKAvatarError> {
     Ok(Json(db::get_stats(&state.pool).await?))
 }
@@ -221,7 +374,9 @@ async fn main() -> anyhow::Result<()> {
     let app = Router::new()
         .route("/verify", post(verify))
         .route("/pull", post(pull))
+        .route("/upload", post(upload))
         .route("/stats", get(stats))
+        .layer(DefaultBodyLimit::max(PREMIUM_HARD_LIMIT))
         .with_state(state);
 
     let host = &config.bind_addr;

crates/avatars/src/new.sql

@@ -0,0 +1,97 @@
+create table images_hashes (
+    hash text primary key,
+    url text not null,
+    file_size int not null,
+    width int not null,
+    height int not null,
+    content_type text not null,
+    created_at timestamptz not null default now()
+);
+
+create table images_assets (
+    id uuid primary key default gen_random_uuid(),
+    system_uuid uuid not null,
+    image text references images_hashes(hash),
+    proxy_image text references images_hashes(hash),
+    kind text not null,
+
+    original_url text,
+    original_file_size int,
+    original_type text,
+    original_attachment_id bigint,
+
+    uploaded_by_account bigint,
+    uploaded_by_ip inet,
+    uploaded_at timestamptz not null default now()
+
+    unique (id, system_uuid)
+);
+
+insert into images_hashes (
+    hash, 
+    url, 
+    file_size, 
+    width, 
+    height, 
+    content_type, 
+    created_at
+)
+select 
+    id,
+    url, 
+    file_size, 
+    width, 
+    height, 
+    coalesce(content_type, 'image/webp'),
+    uploaded_at
+from images;
+
+alter table images rename to images_legacy;
+
+create index if not exists images_original_url_idx on images_assets (original_url);
+create index if not exists images_original_attachment_id_idx on images_assets (original_attachment_id);
+create index if not exists images_uploaded_by_account_idx on images_assets (uploaded_by_account);
+
+create index if not exists images_system_id_idx on images_assets (system_uuid);
+create index if not exists images_proxy_hash_idx on images_assets (image);
+
+-- image cleanup stuffs
+alter table image_cleanup_jobs rename to image_cleanup_jobs_legacy;
+
+create table image_cleanup_jobs (
+    id uuid primary key,
+    system_uuid uuid not null,
+    
+    foreign key (id, system_uuid) 
+        references images_assets(id, system_uuid) 
+        on delete cascade
+);
+
+alter table image_cleanup_pending_jobs rename to image_cleanup_pending_jobs_legacy;
+
+create table image_cleanup_pending_jobs (
+    id uuid primary key,
+    system_uuid uuid not null,
+    ts timestamp not null default now(),
+
+    foreign key (id, system_uuid)
+        references images_assets(id, system_uuid)
+        on delete cascade
+);
+
+create table image_hash_cleanup_jobs (
+    hash text primary key
+
+    foreign key (hash)
+        references images_hashes(hash)
+        on delete cascade
+);
+
+create table image_hash_cleanup_pending_jobs (
+    hash text primary key,
+    ts timestamp not null default now()
+
+    foreign key (hash)
+        references images_hashes(hash)
+        on delete cascade
+);

crates/avatars/src/process.rs

@@ -12,32 +12,10 @@ pub struct ProcessOutput {
     pub width: u32,
     pub height: u32,
     pub hash: Hash,
-    pub format: ProcessedFormat,
+    pub format: ImageFormat,
     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))
@@ -49,13 +27,16 @@ pub async fn process_async(data: Vec<u8>, kind: ImageKind) -> Result<ProcessOutp
 pub fn process(data: &[u8], kind: ImageKind) -> Result<ProcessOutput, PKAvatarError> {
     let time_before = Instant::now();
     let reader = reader_for(data);
-    match reader.format() {
+    let format = reader.format();
+    match 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)? {
+            if !kind.is_premium()
+                && let Some(output) = process_gif(data, kind)?
+            {
                 return Ok(output);
             }
         }
@@ -70,6 +51,18 @@ pub fn process(data: &[u8], kind: ImageKind) -> Result<ProcessOutput, PKAvatarEr
     // need to make a new reader??? why can't it just use the same one. reduce duplication?
     let reader = reader_for(data);
 
+    //if it's a 'premium' image, skip encoding
+    if kind.is_premium() {
+        let hash = Hash::sha256(&data);
+        return Ok(ProcessOutput {
+            width: width,
+            height: height,
+            hash: hash,
+            format: reader.format().expect("expected supported format"),
+            data: data.to_vec(),
+        });
+    }
+
     let time_after_parse = Instant::now();
 
     // apparently `image` sometimes decodes webp images wrong/weird.
@@ -204,7 +197,7 @@ fn process_gif_inner(
 
     Ok(Some(ProcessOutput {
         data,
-        format: ProcessedFormat::Gif,
+        format: ImageFormat::Gif,
         hash,
         width: width as u32,
         height: height as u32,
@@ -249,7 +242,7 @@ fn encode(image: DynamicImage) -> ProcessOutput {
 
     ProcessOutput {
         data: encoded_lossy,
-        format: ProcessedFormat::Webp,
+        format: ImageFormat::WebP,
         hash,
         width,
         height,

crates/avatars/src/pull.rs

@@ -8,8 +8,6 @@ use std::fmt::Write;
 use std::time::Instant;
 use tracing::{error, instrument};
 
-const MAX_SIZE: u64 = 8 * 1024 * 1024;
-
 #[allow(dead_code)]
 pub struct PullResult {
     pub data: Vec<u8>,
@@ -19,8 +17,9 @@ pub struct PullResult {
 
 #[instrument(skip_all)]
 pub async fn pull(
-    client: Arc<Client>,
+    client: &Arc<Client>,
     parsed_url: &ParsedUrl,
+    premium: bool,
 ) -> Result<PullResult, PKAvatarError> {
     let time_before = Instant::now();
     let mut trimmed_url = trim_url_query(&parsed_url.full_url)?;
@@ -59,10 +58,14 @@ pub async fn pull(
         }
     }
 
+    let max_size = match premium {
+        true => super::PREMIUM_HARD_LIMIT as u64,
+        false => super::NORMAL_HARD_LIMIT as u64,
+    };
     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) if size > max_size => {
+            return Err(PKAvatarError::ImageFileSizeTooLarge(size, max_size));
         }
         Some(size) => size,
     };

crates/avatars/src/store.rs

@@ -13,7 +13,10 @@ pub async fn store(bucket: &s3::Bucket, res: &ProcessOutput) -> anyhow::Result<S
         "images/{}/{}.{}",
         &encoded_hash[..2],
         &encoded_hash[2..],
-        res.format.extension()
+        res.format
+            .extensions_str()
+            .first()
+            .expect("expected valid extension")
     );
 
     // todo: something better than these retries
@@ -28,7 +31,7 @@ pub async fn store(bucket: &s3::Bucket, res: &ProcessOutput) -> anyhow::Result<S
         retry_count += 1;
 
         let resp = bucket
-            .put_object_with_content_type(&path, &res.data, res.format.mime_type())
+            .put_object_with_content_type(&path, &res.data, res.format.to_mime_type())
             .await?;
         match resp.status_code() {
             200 => {