aven_core/attachments/
storage.rs

1#![allow(dead_code)]
2
3use std::fs;
4use std::path::{Path, PathBuf};
5
6use anyhow::{Context, Result, bail};
7use sha2::{Digest, Sha256};
8use sqlx::SqliteConnection;
9
10use crate::ids::now;
11use crate::types::BlobInventoryRow;
12
13use super::decode::{ImageFacts, ValidatedImage, validate_image};
14use super::validation::{validate_media_type, validate_sha256};
15
16pub struct StoredBlob {
17    pub sha256: String,
18    pub byte_size: i64,
19    pub facts: ImageFacts,
20}
21
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct StagedBlob {
24    pub sha256: String,
25    pub byte_size: i64,
26    pub created: bool,
27}
28
29pub fn sha256_hex(bytes: &[u8]) -> String {
30    let mut hasher = Sha256::new();
31    hasher.update(bytes);
32    hex::encode(hasher.finalize())
33}
34
35pub fn object_path(blob_dir: &Path, sha256: &str) -> Result<PathBuf> {
36    validate_sha256(sha256)?;
37    Ok(blob_dir.join("objects").join("sha256").join(sha256))
38}
39
40pub async fn store_blob(
41    conn: &mut SqliteConnection,
42    blob_dir: &Path,
43    media_type: &str,
44    bytes: &[u8],
45) -> Result<StoredBlob> {
46    let validated = validate_image(bytes.to_vec(), Some(media_type.to_string())).await?;
47    store_validated_blob(conn, blob_dir, validated).await
48}
49
50pub async fn store_validated_blob(
51    conn: &mut SqliteConnection,
52    blob_dir: &Path,
53    validated: ValidatedImage,
54) -> Result<StoredBlob> {
55    let sha256 = sha256_hex(&validated.bytes);
56    let bytes = validated.bytes;
57    let staged = stage_blob(blob_dir, &sha256, &bytes).await?;
58    if let Err(error) = upsert_inventory_available(
59        conn,
60        &staged.sha256,
61        staged.byte_size,
62        &validated.facts.media_type,
63    )
64    .await
65    {
66        if staged.created {
67            remove_staged_blob_if_unreferenced(conn, blob_dir, &staged.sha256).await;
68        }
69        return Err(error);
70    }
71    Ok(StoredBlob {
72        sha256: staged.sha256,
73        byte_size: staged.byte_size,
74        facts: validated.facts,
75    })
76}
77
78pub async fn stage_blob(blob_dir: &Path, sha256: &str, bytes: &[u8]) -> Result<StagedBlob> {
79    validate_sha256(sha256)?;
80    if sha256_hex(bytes) != sha256 {
81        bail!("error attachment-prepared-hash-mismatch");
82    }
83    let path = object_path(blob_dir, sha256)?;
84    let bytes = bytes.to_vec();
85    let write_path = path.clone();
86    let created =
87        super::blocking::run(move || write_object_atomically(&write_path, &bytes)).await?;
88    let byte_size = fs::metadata(&path)
89        .context("could not inspect attachment object")?
90        .len();
91    let byte_size = i64::try_from(byte_size).context("attachment bytes exceed i64")?;
92    Ok(StagedBlob {
93        sha256: sha256.to_string(),
94        byte_size,
95        created,
96    })
97}
98
99pub async fn remove_staged_blob_if_unreferenced(
100    conn: &mut SqliteConnection,
101    blob_dir: &Path,
102    sha256: &str,
103) {
104    let Ok(mut tx) = crate::db::begin_immediate(conn).await else {
105        return;
106    };
107    let protected = sqlx::query_scalar::<_, bool>(
108        "SELECT
109             EXISTS(SELECT 1 FROM task_attachments WHERE sha256 = ?)
110             OR EXISTS(
111                 SELECT 1 FROM changes
112                 WHERE server_seq IS NULL AND op_type = 'attachment_add'
113                   AND json_extract(payload, '$.sha256') = ?
114             )
115             OR EXISTS(
116                 SELECT 1 FROM blob_leases WHERE sha256 = ? AND expires_at > ?
117             )
118             OR EXISTS(
119                 SELECT 1 FROM blob_upload_reservations
120                 WHERE sha256 = ? AND expires_at > ?
121             )
122             OR EXISTS(SELECT 1 FROM local_shared_capture_pins WHERE sha256 = ?)",
123    )
124    .bind(sha256)
125    .bind(sha256)
126    .bind(sha256)
127    .bind(now())
128    .bind(sha256)
129    .bind(now())
130    .bind(sha256)
131    .fetch_one(&mut *tx)
132    .await;
133    if !matches!(protected, Ok(false)) {
134        let _ = tx.rollback().await;
135        return;
136    }
137    let Ok(path) = object_path(blob_dir, sha256) else {
138        let _ = tx.rollback().await;
139        return;
140    };
141    match fs::remove_file(path) {
142        Ok(()) => {}
143        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
144        Err(_) => {
145            let _ = tx.rollback().await;
146            return;
147        }
148    }
149    if sqlx::query("DELETE FROM blob_inventory WHERE sha256 = ?")
150        .bind(sha256)
151        .execute(&mut *tx)
152        .await
153        .is_err()
154    {
155        let _ = tx.rollback().await;
156        return;
157    }
158    let _ = tx.commit().await;
159}
160
161fn write_object_atomically(path: &Path, bytes: &[u8]) -> Result<bool> {
162    if let Some(parent) = path.parent() {
163        fs::create_dir_all(parent).context("could not create attachment object directory")?;
164        if path.exists() {
165            let existing = fs::read(path).context("could not read attachment object")?;
166            if sha256_hex(&existing) != sha256_hex(bytes) {
167                bail!("error attachment-object-content-mismatch");
168            }
169            return Ok(false);
170        }
171        let mut temp = tempfile::Builder::new()
172            .prefix(".aven-stage-")
173            .tempfile_in(parent)
174            .context("could not create attachment staging file")?;
175        std::io::Write::write_all(&mut temp, bytes)?;
176        temp.as_file().sync_all()?;
177        match temp.persist_noclobber(path) {
178            Ok(_) => {
179                if let Err(error) =
180                    fs::File::open(parent).and_then(|directory| directory.sync_all())
181                {
182                    let _ = fs::remove_file(path);
183                    return Err(error.into());
184                }
185                return Ok(true);
186            }
187            Err(error) if error.error.kind() == std::io::ErrorKind::AlreadyExists => {
188                let existing = fs::read(path).context("could not read attachment object")?;
189                if sha256_hex(&existing) == sha256_hex(bytes) {
190                    return Ok(false);
191                }
192                bail!("error attachment-object-content-mismatch");
193            }
194            Err(error) => return Err(error.error.into()),
195        }
196    }
197    bail!("error attachment-object-path-invalid")
198}
199
200pub async fn upsert_inventory_available(
201    conn: &mut SqliteConnection,
202    sha256: &str,
203    byte_size: i64,
204    media_type: &str,
205) -> Result<()> {
206    validate_sha256(sha256)?;
207    validate_media_type(media_type)?;
208    let existing = blob_inventory_row(conn, sha256).await?;
209    let timestamp = now();
210    if let Some(row) = existing {
211        if row.byte_size != byte_size || row.media_type != media_type {
212            bail!("error blob-inventory-metadata-mismatch");
213        }
214        sqlx::query(
215            "UPDATE blob_inventory
216             SET available = 1, last_verified_at = ?
217             WHERE sha256 = ?",
218        )
219        .bind(&timestamp)
220        .bind(sha256)
221        .execute(&mut *conn)
222        .await?;
223        return Ok(());
224    }
225    sqlx::query(
226        "INSERT INTO blob_inventory(sha256, byte_size, media_type, available, first_seen_at, last_verified_at)
227         VALUES (?, ?, ?, 1, ?, ?)",
228    )
229    .bind(sha256)
230    .bind(byte_size)
231    .bind(media_type)
232    .bind(&timestamp)
233    .bind(&timestamp)
234    .execute(&mut *conn)
235    .await?;
236    Ok(())
237}
238
239pub async fn blob_inventory_row(
240    conn: &mut SqliteConnection,
241    sha256: &str,
242) -> Result<Option<BlobInventoryRow>> {
243    validate_sha256(sha256)?;
244    let row = sqlx::query!(
245        "SELECT sha256, byte_size, media_type, available, first_seen_at, last_verified_at
246         FROM blob_inventory WHERE sha256 = ?",
247        sha256
248    )
249    .fetch_optional(&mut *conn)
250    .await?;
251    Ok(row.map(|row| BlobInventoryRow {
252        sha256: row.sha256.unwrap_or_default(),
253        byte_size: row.byte_size,
254        media_type: row.media_type,
255        available: row.available != 0,
256        first_seen_at: row.first_seen_at,
257        last_verified_at: row.last_verified_at,
258    }))
259}
260
261#[cfg(test)]
262mod tests {
263    use std::io::Cursor;
264    use std::path::PathBuf;
265
266    use image::{DynamicImage, ImageFormat, RgbaImage};
267
268    use crate::db::open_db;
269
270    use super::*;
271
272    fn png_bytes() -> Vec<u8> {
273        let mut bytes = Cursor::new(Vec::new());
274        DynamicImage::ImageRgba8(RgbaImage::new(2, 1))
275            .write_to(&mut bytes, ImageFormat::Png)
276            .unwrap();
277        bytes.into_inner()
278    }
279
280    #[tokio::test]
281    async fn stores_and_retrieves_blob() {
282        let temp = tempfile::tempdir().unwrap();
283        let db_path = temp.path().join("test.sqlite");
284        let pool = open_db(&db_path).await.unwrap();
285        let mut conn = pool.acquire().await.unwrap();
286
287        let blob_dir = temp.path().join("blobs");
288
289        let bytes = png_bytes();
290        let stored = store_blob(&mut conn, &blob_dir, "image/png", &bytes)
291            .await
292            .unwrap();
293
294        assert_eq!(stored.sha256, sha256_hex(&bytes));
295        assert!(stored.byte_size > 0);
296
297        let obj_path = object_path(&blob_dir, &stored.sha256).unwrap();
298        assert!(obj_path.exists(), "blob file should exist on disk");
299
300        let on_disk = std::fs::read(&obj_path).unwrap();
301        assert_eq!(on_disk, bytes);
302
303        let row = blob_inventory_row(&mut conn, &stored.sha256)
304            .await
305            .unwrap()
306            .expect("inventory row should exist");
307        assert!(row.available);
308        assert_eq!(row.sha256, stored.sha256);
309        assert_eq!(row.byte_size, stored.byte_size);
310    }
311
312    #[tokio::test]
313    async fn stores_blob_idempotently() {
314        let temp = tempfile::tempdir().unwrap();
315        let db_path = temp.path().join("test.sqlite");
316        let pool = open_db(&db_path).await.unwrap();
317        let mut conn = pool.acquire().await.unwrap();
318
319        let blob_dir = temp.path().join("blobs");
320
321        let bytes = png_bytes();
322        let stored1 = store_blob(&mut conn, &blob_dir, "image/png", &bytes)
323            .await
324            .unwrap();
325        let stored2 = store_blob(&mut conn, &blob_dir, "image/png", &bytes)
326            .await
327            .unwrap();
328
329        assert_eq!(stored1.sha256, stored2.sha256);
330        let row = blob_inventory_row(&mut conn, &stored1.sha256)
331            .await
332            .unwrap()
333            .unwrap();
334        assert!(row.available);
335    }
336
337    #[tokio::test]
338    async fn rejects_inventory_metadata_mismatch() {
339        let temp = tempfile::tempdir().unwrap();
340        let db_path = temp.path().join("test.sqlite");
341        let pool = open_db(&db_path).await.unwrap();
342        let mut conn = pool.acquire().await.unwrap();
343
344        let sha256 = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
345        upsert_inventory_available(&mut conn, sha256, 12, "image/png")
346            .await
347            .unwrap();
348
349        let error = upsert_inventory_available(&mut conn, sha256, 13, "image/png")
350            .await
351            .unwrap_err()
352            .to_string();
353        assert_eq!(error, "error blob-inventory-metadata-mismatch");
354
355        let error = upsert_inventory_available(&mut conn, sha256, 12, "image/jpeg")
356            .await
357            .unwrap_err()
358            .to_string();
359        assert_eq!(error, "error blob-inventory-metadata-mismatch");
360    }
361
362    #[tokio::test]
363    async fn inventory_failure_removes_only_newly_staged_blob() {
364        let temp = tempfile::tempdir().unwrap();
365        let db_path = temp.path().join("test.sqlite");
366        let pool = open_db(&db_path).await.unwrap();
367        let mut conn = pool.acquire().await.unwrap();
368        let blob_dir = temp.path().join("blobs");
369        let bytes = png_bytes();
370        let sha256 = sha256_hex(&bytes);
371        upsert_inventory_available(&mut conn, &sha256, 12, "image/png")
372            .await
373            .unwrap();
374
375        assert!(
376            store_blob(&mut conn, &blob_dir, "image/png", &bytes)
377                .await
378                .is_err()
379        );
380        assert!(!object_path(&blob_dir, &sha256).unwrap().exists());
381    }
382
383    #[test]
384    fn computes_sha256_hex() {
385        let result = sha256_hex(b"hello");
386        assert_eq!(result.len(), 64);
387        assert!(result.bytes().all(|b| b.is_ascii_hexdigit()));
388    }
389
390    #[test]
391    fn computes_object_path() {
392        let blob_dir = Path::new("/tmp/aven");
393        let hash = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
394        let path = object_path(blob_dir, hash).unwrap();
395        assert_eq!(
396            path,
397            PathBuf::from(
398                "/tmp/aven/objects/sha256/abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789"
399            )
400        );
401    }
402
403    #[test]
404    fn rejects_invalid_hash_in_object_path() {
405        assert!(object_path(Path::new("/tmp"), "short").is_err());
406    }
407}
aven_core/attachments/storage.rs

aven_core/attachments/
storage.rs
