OXYGENHUB

Coverage / crates/hub/src/web/avatar.rs 93.7%

192 of 205 lines run, 42 of 79 functions · run #536 on main · run missed

1
//! Pictures of accounts and organisations (hub spec §6.3): an uploaded
2
//! photo or logo, re-encoded so no uploaded bytes are ever served, or an
3
//! identicon: a five-by-five mirrored pattern and a colour drawn from
4
//! BLAKE3 of the lower-cased name, so everyone has a stable picture.
5
6
use crate::db::accounts;
7
use crate::Shared;
8
use axum::extract::{Path, State};
9
use axum::http::header;
10
use axum::response::{IntoResponse, Response};
11
use std::io::Cursor;
12
13
/// Largest upload accepted.
14
pub const MAX_BYTES: usize = 5 << 20;
15
/// Largest width or height decoded.
16
const MAX_SIDE: u32 = 8000;
17
/// The stored picture is this many pixels square.
18
pub const SIZE: u32 = 400;
19
20
/// A re-encoded picture.
21
pub struct Picture {
22
    pub bytes: Vec<u8>,
23
    /// `png` or `jpeg`.
24
    pub format: &'static str,
25
}
26
27
/// Decode an upload (PNG, JPEG or WebP by content), crop it to the middle
28
/// square, scale it to `SIZE`, and encode it afresh: PNG when it has
29
/// transparency, JPEG otherwise.
3011×
pub fn process(input: &[u8]) -> Result<Picture, String> {
3111×
    if input.len() > MAX_BYTES {
321×
        return Err(format!("A picture can be at most {} MiB.", MAX_BYTES >> 20));
3310×
    }
3410×
    let format = image::guess_format(input).map_err(|_| "That is not a PNG, JPEG or WebP image.".to_string())?;
358×
    if !matches!(format, image::ImageFormat::Png | image::ImageFormat::Jpeg | image::ImageFormat::WebP) {
361×
        return Err("Pictures must be PNG, JPEG or WebP.".into());
377×
    }
387×
    let mut reader = image::ImageReader::with_format(Cursor::new(input), format);
397×
    let mut limits = image::Limits::default();
407×
    limits.max_image_width = Some(MAX_SIDE);
417×
    limits.max_image_height = Some(MAX_SIDE);
427×
    limits.max_alloc = Some(512 << 20);
437×
    reader.limits(limits);
447×
    let img = reader.decode().map_err(|e| format!("The image could not be read: {e}"))?;
455×
    let side = img.width().min(img.height());
465×
    if side == 0 {
470×
        return Err("The image is empty.".into());
485×
    }
495×
    let square = img.crop_imm((img.width() - side) / 2, (img.height() - side) / 2, side, side);
505×
    let scaled = square.resize_exact(SIZE, SIZE, image::imageops::FilterType::Lanczos3);
515×
    let mut out = Vec::new();
525×
    if scaled.color().has_alpha() {
532×
        scaled.to_rgba8().write_to(&mut Cursor::new(&mut out), image::ImageFormat::Png).map_err(|e| e.to_string())?;
542×
        Ok(Picture { bytes: out, format: "png" })
55
    } else {
563×
        image::codecs::jpeg::JpegEncoder::new_with_quality(&mut out, 88).encode_image(&scaled.to_rgb8()).map_err(|e| e.to_string())?;
573×
        Ok(Picture { bytes: out, format: "jpeg" })
58
    }
5911×
}
60
617×
fn file_of(app: &Shared, kind: &str, id: i64) -> std::path::PathBuf {
627×
    app.config.data.join("avatars").join(format!("{kind}-{id}"))
637×
}
64
65
/// Process and store an upload as the picture of an account (`account`) or
66
/// organisation (`org`).
673×
pub async fn store(app: &Shared, kind: &str, id: i64, upload: Vec<u8>) -> Result<(), String> {
683×
    let pic = tokio::task::spawn_blocking(move || process(&upload)).await.map_err(|e| e.to_string())??;
692×
    let path = file_of(app, kind, id);
702×
    let dir = path.parent().unwrap().to_path_buf();
712×
    tokio::fs::create_dir_all(&dir).await.map_err(|e| e.to_string())?;
722×
    let tmp = dir.join(format!(".{kind}-{id}.tmp"));
732×
    tokio::fs::write(&tmp, &pic.bytes).await.map_err(|e| e.to_string())?;
742×
    tokio::fs::rename(&tmp, &path).await.map_err(|e| e.to_string())?;
752×
    let hash = blake3::hash(&pic.bytes).to_hex();
762×
    accounts::set_avatar(&app.db, kind, id, &format!("{}:{}", pic.format, &hash[..16])).await.map_err(|e| e.to_string())
773×
}
78
79
/// Go back to the identicon.
801×
pub async fn remove(app: &Shared, kind: &str, id: i64) -> Result<(), String> {
811×
    accounts::set_avatar(&app.db, kind, id, "").await.map_err(|e| e.to_string())?;
821×
    let _ = tokio::fs::remove_file(file_of(app, kind, id)).await;
831×
    Ok(())
841×
}
85
86
/// The SVG for `name`: 15 bits choose the cells of the left three columns
87
/// (mirrored to five), two bytes choose the hue.
888×
pub fn identicon(name: &str) -> String {
898×
    let h = blake3::hash(name.to_ascii_lowercase().as_bytes());
908×
    let b = h.as_bytes();
918×
    let hue = u16::from_le_bytes([b[0], b[1]]) % 360;
928×
    let bits = u32::from_le_bytes([b[2], b[3], b[4], b[5]]);
938×
    let mut cells = String::new();
9440×
    for row in 0..5u32 {
95120×
        for col in 0..3u32 {
96120×
            if bits >> (row * 3 + col) & 1 == 1 {
9795×
                for x in [col, 4 - col] {
9895×
                    cells.push_str(&format!(r#"<rect x="{}" y="{}" width="1" height="1"/>"#, x + 1, row + 1));
9995×
                    if x == 2 {
10021×
                        break;
10174×
                    }
102
                }
10362×
            }
104
        }
105
    }
1068×
    format!(
107
        r#"<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 7 7" shape-rendering="crispEdges"><rect width="7" height="7" fill="hsl({hue},45%,92%)"/><g fill="hsl({hue},65%,45%)">{cells}</g></svg>"#
108
    )
1098×
}
110
111
/// A picture form: the CSRF token, the file, and `action=remove` to clear it.
112
pub struct Upload {
113
    pub csrf: String,
114
    pub file: Vec<u8>,
115
    pub remove: bool,
116
}
117
1185×
pub async fn read_upload(mut multipart: axum::extract::Multipart) -> Upload {
1195×
    let mut u = Upload { csrf: String::new(), file: Vec::new(), remove: false };
12020×
    while let Ok(Some(field)) = multipart.next_field().await {
12115×
        match field.name() {
12215×
            Some("_csrf") => u.csrf = field.text().await.unwrap_or_default(),
12310×
            Some("picture") => u.file = field.bytes().await.map(|b| b.to_vec()).unwrap_or_default(),
1245×
            Some("action") => u.remove = field.text().await.unwrap_or_default() == "remove",
1250×
            _ => {}
126
        }
127
    }
1285×
    u
1295×
}
130
131
/// Apply a picture form to an account or organisation; the message to show.
1324×
pub async fn apply(app: &Shared, kind: &str, id: i64, u: Upload) -> Result<String, String> {
1334×
    if u.remove {
1341×
        remove(app, kind, id).await?;
1351×
        return Ok("Picture removed; the generated one is back.".into());
1363×
    }
1373×
    if u.file.is_empty() {
1380×
        return Err("Choose a picture to upload.".into());
1393×
    }
1403×
    store(app, kind, id, u.file).await?;
1412×
    Ok("Picture saved.".into())
1424×
}
143
144
/// The brand mark for an AI agent, by the name its identity carries
145
/// (hub spec §6): `Claude Code` → Claude's, `ChatGPT` → OpenAI's, and so
146
/// on; `None` for a name the hub has no mark for.
1477×
pub fn agent_mark(name: &str) -> Option<&'static str> {
148
    const MARKS: &[(&str, &str)] = &[
149
        ("claude", include_str!("../../../../assets/agents/claude.svg")),
150
        ("anthropic", include_str!("../../../../assets/agents/claude.svg")),
151
        ("chatgpt", include_str!("../../../../assets/agents/chatgpt.svg")),
152
        ("openai", include_str!("../../../../assets/agents/chatgpt.svg")),
153
        ("codex", include_str!("../../../../assets/agents/chatgpt.svg")),
154
        ("perplexity", include_str!("../../../../assets/agents/perplexity.svg")),
155
        ("deepseek", include_str!("../../../../assets/agents/deepseek.svg")),
156
        ("mistral", include_str!("../../../../assets/agents/mistral.svg")),
157
        ("antigravity", include_str!("../../../../assets/agents/antigravity.svg")),
158
        ("gemini", include_str!("../../../../assets/agents/antigravity.svg")),
159
        ("kimi", include_str!("../../../../assets/agents/kimi.svg")),
160
        ("moonshot", include_str!("../../../../assets/agents/kimi.svg")),
161
    ];
1627×
    let lower = name.to_lowercase();
16369×
    MARKS.iter().find(|(k, _)| lower.contains(k)).map(|(_, svg)| *svg)
1647×
}
165
166
/// `GET /avatar/agent/{name}`: an AI agent's brand mark, or its identicon.
1672×
pub async fn agent(Path(name): Path<String>) -> Response {
1682×
    let name = name.trim_end_matches(".svg");
1692×
    let svg = agent_mark(name).map(String::from).unwrap_or_else(|| identicon(name));
1702×
    ([(header::CONTENT_TYPE, "image/svg+xml"), (header::CACHE_CONTROL, "public, max-age=86400")], svg).into_response()
1712×
}
172
173
/// `GET /avatar/{name}`: the uploaded picture, or the identicon.
1748×
pub async fn avatar(State(app): State<Shared>, Path(name): Path<String>) -> Response {
1758×
    let name = name.trim_end_matches(".svg");
1768×
    let uploaded = match accounts::account_by_name(&app.db, name).await.ok().flatten() {
1776×
        Some(a) => Some(("account", a.id, a.avatar)),
1782×
        None => accounts::org_by_name(&app.db, name).await.ok().flatten().map(|o| ("org", o.id, o.avatar)),
179
    };
1808×
    if let Some((kind, id, value)) = uploaded.filter(|(_, _, v)| !v.is_empty()) {
1814×
        if let (Some((format, hash)), Ok(bytes)) = (value.split_once(':'), tokio::fs::read(file_of(&app, kind, id)).await) {
1824×
            let mime = if format == "png" { "image/png" } else { "image/jpeg" };
1834×
            return (
1844×
                [
1854×
                    (header::CONTENT_TYPE, mime.to_string()),
1864×
                    (header::CACHE_CONTROL, "public, max-age=300".to_string()),
1874×
                    (header::ETAG, format!("\"{hash}\"")),
1884×
                    (header::X_CONTENT_TYPE_OPTIONS, "nosniff".to_string()),
1894×
                ],
1904×
                bytes,
1914×
            )
1924×
                .into_response();
1930×
        }
1944×
    }
195
    // The hub's own accounts have their marks until someone uploads another.
1964×
    let svg = system_mark(name).map(String::from).unwrap_or_else(|| identicon(name));
1974×
    ([(header::CONTENT_TYPE, "image/svg+xml"), (header::CACHE_CONTROL, "public, max-age=300")], svg).into_response()
1988×
}
199
200
/// The mark of an account the hub runs itself: Diffusion's ripples from a
201
/// source out to its consumers.
2027×
pub fn system_mark(name: &str) -> Option<&'static str> {
2037×
    (name == crate::diffusion::ACCOUNT).then_some(include_str!("../../../../assets/icons/diffusion.svg"))
2047×
}
205
206
#[cfg(test)]
207
mod tests {
208
    use super::{agent_mark, identicon, process, system_mark, SIZE};
209
2104×
    fn encoded(img: image::DynamicImage, format: image::ImageFormat) -> Vec<u8> {
2114×
        let mut out = Vec::new();
2124×
        img.write_to(&mut std::io::Cursor::new(&mut out), format).unwrap();
2134×
        out
2144×
    }
215
216
    #[test]
2171×
    fn uploads_are_cropped_scaled_and_reencoded() {
218
        // A wide photo becomes a square JPEG from its middle.
219270000×
        let photo = image::DynamicImage::ImageRgb8(image::RgbImage::from_fn(900, 300, |x, _| image::Rgb([(x / 4) as u8, 80, 160])));
2201×
        let pic = process(&encoded(photo, image::ImageFormat::Png)).unwrap();
2211×
        assert_eq!(pic.format, "jpeg");
2221×
        let back = image::load_from_memory(&pic.bytes).unwrap();
2231×
        assert_eq!((back.width(), back.height()), (SIZE, SIZE));
224
        // A logo with transparency stays PNG with its alpha.
2254096×
        let logo = image::DynamicImage::ImageRgba8(image::RgbaImage::from_fn(64, 64, |x, _| {
2264096×
            image::Rgba([0, 114, 255, if x < 32 { 0 } else { 255 }])
2274096×
        }));
2281×
        let pic = process(&encoded(logo, image::ImageFormat::Png)).unwrap();
2291×
        assert_eq!(pic.format, "png");
2301×
        assert!(image::load_from_memory(&pic.bytes).unwrap().color().has_alpha());
231
        // JPEG in, too.
2321×
        let jpeg = encoded(image::DynamicImage::ImageRgb8(image::RgbImage::new(50, 80)), image::ImageFormat::Jpeg);
2331×
        assert_eq!(process(&jpeg).unwrap().format, "jpeg");
2341×
    }
235
236
    #[test]
2371×
    fn anything_else_is_refused() {
2381×
        assert!(process(b"<svg xmlns='http://www.w3.org/2000/svg'><script>alert(1)</script></svg>").is_err(), "no SVG");
2391×
        assert!(process(b"GIF89a\x01\0\x01\0").is_err(), "no GIF");
2401×
        assert!(process(b"\x89PNG\r\n\x1a\n broken").is_err(), "a truncated PNG");
2411×
        assert!(process(&vec![0u8; super::MAX_BYTES + 1]).err().unwrap().contains("at most 5 MiB"));
242
        // Larger than the side limit: refused before allocating it.
2431×
        let big = encoded(image::DynamicImage::ImageLuma8(image::GrayImage::new(8001, 1)), image::ImageFormat::Png);
2441×
        assert!(process(&big).is_err(), "8001 pixels wide");
2451×
    }
246
247
    #[test]
2481×
    fn kimi_has_its_mark() {
2494×
        for name in ["kimi", "Kimi", "agent:kimi@andymccall", "Moonshot Kimi"] {
2504×
            let svg = agent_mark(name).unwrap_or_else(|| panic!("{name} has no mark"));
2514×
            assert!(svg.contains("aria-label=\"Kimi\""), "{name}");
252
        }
2531×
        assert!(agent_mark("kim").is_none(), "only the name itself matches");
2541×
    }
255
256
    #[test]
2571×
    fn the_diffusion_account_has_its_mark() {
2581×
        let m = system_mark("diffusion").expect("a mark");
2591×
        assert!(m.starts_with("<svg") && m.contains("<title>Diffusion</title>"));
2601×
        assert!(system_mark("andy").is_none() && system_mark("Diffusion").is_none(), "only the hub's own account, by its exact name");
2611×
    }
262
263
    #[test]
2641×
    fn identicons_are_stable_symmetric_and_case_blind() {
2651×
        let a = identicon("andy");
2661×
        assert_eq!(a, identicon("Andy"), "names are case-insensitive");
2671×
        assert_ne!(a, identicon("bob"));
2681×
        assert!(a.starts_with("<svg") && a.ends_with("</svg>"));
269
        // Every cell has its mirror image.
2701×
        let cells: Vec<(u32, u32)> = a
2711×
            .split(r#"<rect x=""#)
2721×
            .skip(1)
27311×
            .filter_map(|r| {
27411×
                let x: u32 = r.split('"').next()?.parse().ok()?;
27511×
                let y: u32 = r.split(r#"y=""#).nth(1)?.split('"').next()?.parse().ok()?;
27611×
                Some((x, y))
27711×
            })
27811×
            .filter(|(x, _)| (1..=5).contains(x))
2791×
            .collect();
28011×
        for (x, y) in &cells {
28111×
            assert!(cells.contains(&(6 - x, *y)), "({x},{y}) has no mirror in {a}");
282
        }
2831×
    }
284
}