OXYGENHUB
◆ close-proposal-140 oxygenhub / crates / hub / src / web / avatar.rs
11.6 KiB · Rust · last changed Agents get their marks; contributors one row each; an assistant is tested when added 1 month ago source blame history raw
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.
30
pub fn process(input: &[u8]) -> Result<Picture, String> {
31
    if input.len() > MAX_BYTES {
32
        return Err(format!("A picture can be at most {} MiB.", MAX_BYTES >> 20));
33
    }
34
    let format = image::guess_format(input).map_err(|_| "That is not a PNG, JPEG or WebP image.".to_string())?;
35
    if !matches!(format, image::ImageFormat::Png | image::ImageFormat::Jpeg | image::ImageFormat::WebP) {
36
        return Err("Pictures must be PNG, JPEG or WebP.".into());
37
    }
38
    let mut reader = image::ImageReader::with_format(Cursor::new(input), format);
39
    let mut limits = image::Limits::default();
40
    limits.max_image_width = Some(MAX_SIDE);
41
    limits.max_image_height = Some(MAX_SIDE);
42
    limits.max_alloc = Some(512 << 20);
43
    reader.limits(limits);
44
    let img = reader.decode().map_err(|e| format!("The image could not be read: {e}"))?;
45
    let side = img.width().min(img.height());
46
    if side == 0 {
47
        return Err("The image is empty.".into());
48
    }
49
    let square = img.crop_imm((img.width() - side) / 2, (img.height() - side) / 2, side, side);
50
    let scaled = square.resize_exact(SIZE, SIZE, image::imageops::FilterType::Lanczos3);
51
    let mut out = Vec::new();
52
    if scaled.color().has_alpha() {
53
        scaled.to_rgba8().write_to(&mut Cursor::new(&mut out), image::ImageFormat::Png).map_err(|e| e.to_string())?;
54
        Ok(Picture { bytes: out, format: "png" })
55
    } else {
56
        image::codecs::jpeg::JpegEncoder::new_with_quality(&mut out, 88).encode_image(&scaled.to_rgb8()).map_err(|e| e.to_string())?;
57
        Ok(Picture { bytes: out, format: "jpeg" })
58
    }
59
}
60
61
fn file_of(app: &Shared, kind: &str, id: i64) -> std::path::PathBuf {
62
    app.config.data.join("avatars").join(format!("{kind}-{id}"))
63
}
64
65
/// Process and store an upload as the picture of an account (`account`) or
66
/// organisation (`org`).
67
pub async fn store(app: &Shared, kind: &str, id: i64, upload: Vec<u8>) -> Result<(), String> {
68
    let pic = tokio::task::spawn_blocking(move || process(&upload)).await.map_err(|e| e.to_string())??;
69
    let path = file_of(app, kind, id);
70
    let dir = path.parent().unwrap().to_path_buf();
71
    tokio::fs::create_dir_all(&dir).await.map_err(|e| e.to_string())?;
72
    let tmp = dir.join(format!(".{kind}-{id}.tmp"));
73
    tokio::fs::write(&tmp, &pic.bytes).await.map_err(|e| e.to_string())?;
74
    tokio::fs::rename(&tmp, &path).await.map_err(|e| e.to_string())?;
75
    let hash = blake3::hash(&pic.bytes).to_hex();
76
    accounts::set_avatar(&app.db, kind, id, &format!("{}:{}", pic.format, &hash[..16])).await.map_err(|e| e.to_string())
77
}
78
79
/// Go back to the identicon.
80
pub async fn remove(app: &Shared, kind: &str, id: i64) -> Result<(), String> {
81
    accounts::set_avatar(&app.db, kind, id, "").await.map_err(|e| e.to_string())?;
82
    let _ = tokio::fs::remove_file(file_of(app, kind, id)).await;
83
    Ok(())
84
}
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.
88
pub fn identicon(name: &str) -> String {
89
    let h = blake3::hash(name.to_ascii_lowercase().as_bytes());
90
    let b = h.as_bytes();
91
    let hue = u16::from_le_bytes([b[0], b[1]]) % 360;
92
    let bits = u32::from_le_bytes([b[2], b[3], b[4], b[5]]);
93
    let mut cells = String::new();
94
    for row in 0..5u32 {
95
        for col in 0..3u32 {
96
            if bits >> (row * 3 + col) & 1 == 1 {
97
                for x in [col, 4 - col] {
98
                    cells.push_str(&format!(r#"<rect x="{}" y="{}" width="1" height="1"/>"#, x + 1, row + 1));
99
                    if x == 2 {
100
                        break;
101
                    }
102
                }
103
            }
104
        }
105
    }
106
    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
    )
109
}
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
118
pub async fn read_upload(mut multipart: axum::extract::Multipart) -> Upload {
119
    let mut u = Upload { csrf: String::new(), file: Vec::new(), remove: false };
120
    while let Ok(Some(field)) = multipart.next_field().await {
121
        match field.name() {
122
            Some("_csrf") => u.csrf = field.text().await.unwrap_or_default(),
123
            Some("picture") => u.file = field.bytes().await.map(|b| b.to_vec()).unwrap_or_default(),
124
            Some("action") => u.remove = field.text().await.unwrap_or_default() == "remove",
125
            _ => {}
126
        }
127
    }
128
    u
129
}
130
131
/// Apply a picture form to an account or organisation; the message to show.
132
pub async fn apply(app: &Shared, kind: &str, id: i64, u: Upload) -> Result<String, String> {
133
    if u.remove {
134
        remove(app, kind, id).await?;
135
        return Ok("Picture removed; the generated one is back.".into());
136
    }
137
    if u.file.is_empty() {
138
        return Err("Choose a picture to upload.".into());
139
    }
140
    store(app, kind, id, u.file).await?;
141
    Ok("Picture saved.".into())
142
}
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.
147
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
    ];
160
    let lower = name.to_lowercase();
161
    MARKS.iter().find(|(k, _)| lower.contains(k)).map(|(_, svg)| *svg)
162
}
163
164
/// `GET /avatar/agent/{name}`: an AI agent's brand mark, or its identicon.
165
pub async fn agent(Path(name): Path<String>) -> Response {
166
    let name = name.trim_end_matches(".svg");
167
    let svg = agent_mark(name).map(String::from).unwrap_or_else(|| identicon(name));
168
    ([(header::CONTENT_TYPE, "image/svg+xml"), (header::CACHE_CONTROL, "public, max-age=86400")], svg).into_response()
169
}
170
171
/// `GET /avatar/{name}`: the uploaded picture, or the identicon.
172
pub async fn avatar(State(app): State<Shared>, Path(name): Path<String>) -> Response {
173
    let name = name.trim_end_matches(".svg");
174
    let uploaded = match accounts::account_by_name(&app.db, name).await.ok().flatten() {
175
        Some(a) => Some(("account", a.id, a.avatar)),
176
        None => accounts::org_by_name(&app.db, name).await.ok().flatten().map(|o| ("org", o.id, o.avatar)),
177
    };
178
    if let Some((kind, id, value)) = uploaded.filter(|(_, _, v)| !v.is_empty()) {
179
        if let (Some((format, hash)), Ok(bytes)) = (value.split_once(':'), tokio::fs::read(file_of(&app, kind, id)).await) {
180
            let mime = if format == "png" { "image/png" } else { "image/jpeg" };
181
            return (
182
                [
183
                    (header::CONTENT_TYPE, mime.to_string()),
184
                    (header::CACHE_CONTROL, "public, max-age=300".to_string()),
185
                    (header::ETAG, format!("\"{hash}\"")),
186
                    (header::X_CONTENT_TYPE_OPTIONS, "nosniff".to_string()),
187
                ],
188
                bytes,
189
            )
190
                .into_response();
191
        }
192
    }
193
    ([(header::CONTENT_TYPE, "image/svg+xml"), (header::CACHE_CONTROL, "public, max-age=300")], identicon(name)).into_response()
194
}
195
196
#[cfg(test)]
197
mod tests {
198
    use super::{identicon, process, SIZE};
199
200
    fn encoded(img: image::DynamicImage, format: image::ImageFormat) -> Vec<u8> {
201
        let mut out = Vec::new();
202
        img.write_to(&mut std::io::Cursor::new(&mut out), format).unwrap();
203
        out
204
    }
205
206
    #[test]
207
    fn uploads_are_cropped_scaled_and_reencoded() {
208
        // A wide photo becomes a square JPEG from its middle.
209
        let photo = image::DynamicImage::ImageRgb8(image::RgbImage::from_fn(900, 300, |x, _| image::Rgb([(x / 4) as u8, 80, 160])));
210
        let pic = process(&encoded(photo, image::ImageFormat::Png)).unwrap();
211
        assert_eq!(pic.format, "jpeg");
212
        let back = image::load_from_memory(&pic.bytes).unwrap();
213
        assert_eq!((back.width(), back.height()), (SIZE, SIZE));
214
        // A logo with transparency stays PNG with its alpha.
215
        let logo = image::DynamicImage::ImageRgba8(image::RgbaImage::from_fn(64, 64, |x, _| {
216
            image::Rgba([0, 114, 255, if x < 32 { 0 } else { 255 }])
217
        }));
218
        let pic = process(&encoded(logo, image::ImageFormat::Png)).unwrap();
219
        assert_eq!(pic.format, "png");
220
        assert!(image::load_from_memory(&pic.bytes).unwrap().color().has_alpha());
221
        // JPEG in, too.
222
        let jpeg = encoded(image::DynamicImage::ImageRgb8(image::RgbImage::new(50, 80)), image::ImageFormat::Jpeg);
223
        assert_eq!(process(&jpeg).unwrap().format, "jpeg");
224
    }
225
226
    #[test]
227
    fn anything_else_is_refused() {
228
        assert!(process(b"<svg xmlns='http://www.w3.org/2000/svg'><script>alert(1)</script></svg>").is_err(), "no SVG");
229
        assert!(process(b"GIF89a\x01\0\x01\0").is_err(), "no GIF");
230
        assert!(process(b"\x89PNG\r\n\x1a\n broken").is_err(), "a truncated PNG");
231
        assert!(process(&vec![0u8; super::MAX_BYTES + 1]).err().unwrap().contains("at most 5 MiB"));
232
        // Larger than the side limit: refused before allocating it.
233
        let big = encoded(image::DynamicImage::ImageLuma8(image::GrayImage::new(8001, 1)), image::ImageFormat::Png);
234
        assert!(process(&big).is_err(), "8001 pixels wide");
235
    }
236
237
    #[test]
238
    fn identicons_are_stable_symmetric_and_case_blind() {
239
        let a = identicon("andy");
240
        assert_eq!(a, identicon("Andy"), "names are case-insensitive");
241
        assert_ne!(a, identicon("bob"));
242
        assert!(a.starts_with("<svg") && a.ends_with("</svg>"));
243
        // Every cell has its mirror image.
244
        let cells: Vec<(u32, u32)> = a
245
            .split(r#"<rect x=""#)
246
            .skip(1)
247
            .filter_map(|r| {
248
                let x: u32 = r.split('"').next()?.parse().ok()?;
249
                let y: u32 = r.split(r#"y=""#).nth(1)?.split('"').next()?.parse().ok()?;
250
                Some((x, y))
251
            })
252
            .filter(|(x, _)| (1..=5).contains(x))
253
            .collect();
254
        for (x, y) in &cells {
255
            assert!(cells.contains(&(6 - x, *y)), "({x},{y}) has no mirror in {a}");
256
        }
257
    }
258
}