mail4agent_server/
media.rs1use rand::RngCore;
9use rusqlite::{params, Connection, OptionalExtension};
10
11pub const MAX_UPLOAD_BYTES: usize = 25 * 1024 * 1024;
13
14pub fn create_media_schema(conn: &Connection) -> rusqlite::Result<()> {
16 conn.execute_batch(
17 "CREATE TABLE IF NOT EXISTS media (
18 media_id TEXT PRIMARY KEY,
19 owner_user_id INTEGER NOT NULL,
20 content_type TEXT NOT NULL,
21 filename TEXT,
22 created_ms INTEGER NOT NULL,
23 data BLOB NOT NULL
24 );
25 CREATE INDEX IF NOT EXISTS idx_media_created ON media(created_ms);",
26 )
27}
28
29pub struct Blob {
31 pub content_type: String,
33 pub filename: Option<String>,
35 pub data: Vec<u8>,
37}
38
39pub fn put(conn: &Connection, owner: i64, content_type: &str, filename: Option<&str>, data: &[u8], now_ms: i64) -> rusqlite::Result<String> {
41 let mut raw = [0u8; 18];
42 rand::thread_rng().fill_bytes(&mut raw);
43 let id = base64::Engine::encode(&base64::engine::general_purpose::URL_SAFE_NO_PAD, raw);
44 conn.execute(
45 "INSERT INTO media (media_id, owner_user_id, content_type, filename, created_ms, data) VALUES (?1,?2,?3,?4,?5,?6)",
46 params![id, owner, content_type, filename, now_ms, data],
47 )?;
48 Ok(id)
49}
50
51pub fn get(conn: &Connection, media_id: &str) -> rusqlite::Result<Option<Blob>> {
53 conn.query_row("SELECT content_type, filename, data FROM media WHERE media_id = ?1", params![media_id], |r| {
54 Ok(Blob { content_type: r.get(0)?, filename: r.get(1)?, data: r.get(2)? })
55 })
56 .optional()
57}
58
59pub fn purge_expired(conn: &Connection, now_ms: i64, ttl_ms: i64) -> rusqlite::Result<usize> {
61 conn.execute("DELETE FROM media WHERE created_ms < ?1", params![now_ms - ttl_ms])
62}
63
64#[cfg(test)]
65mod tests {
66 use super::*;
67
68 #[test]
69 fn put_get_and_ttl() {
70 let c = Connection::open_in_memory().unwrap();
71 create_media_schema(&c).unwrap();
72 let id = put(&c, 1, "application/octet-stream", Some("a.bin"), &[1, 2, 3], 1_000).unwrap();
73 assert_eq!(get(&c, &id).unwrap().unwrap().data, vec![1, 2, 3]);
74 assert_eq!(purge_expired(&c, 1_500, 1_000).unwrap(), 0);
75 assert_eq!(purge_expired(&c, 5_000, 1_000).unwrap(), 1);
76 assert!(get(&c, &id).unwrap().is_none());
77 }
78}