Skip to main content

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 referenced = sqlx::query_scalar::<_, bool>(
105        "SELECT EXISTS(SELECT 1 FROM task_attachments WHERE sha256 = ?)",
106    )
107    .bind(sha256)
108    .fetch_one(&mut *conn)
109    .await;
110    let pending_upload = sqlx::query_scalar::<_, bool>(
111        "SELECT EXISTS(
112            SELECT 1 FROM changes
113            WHERE server_seq IS NULL AND op_type = 'attachment_add'
114              AND json_extract(payload, '$.sha256') = ?
115         )",
116    )
117    .bind(sha256)
118    .fetch_one(&mut *conn)
119    .await;
120    let active_lease = sqlx::query_scalar::<_, bool>(
121        "SELECT EXISTS(SELECT 1 FROM blob_leases WHERE sha256 = ? AND expires_at > ?)",
122    )
123    .bind(sha256)
124    .bind(now())
125    .fetch_one(&mut *conn)
126    .await;
127    if matches!(referenced, Ok(false))
128        && matches!(pending_upload, Ok(false))
129        && matches!(active_lease, Ok(false))
130        && let Ok(path) = object_path(blob_dir, sha256)
131    {
132        match fs::remove_file(path) {
133            Ok(()) => {
134                let _ = sqlx::query("DELETE FROM blob_inventory WHERE sha256 = ?")
135                    .bind(sha256)
136                    .execute(&mut *conn)
137                    .await;
138            }
139            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
140                let _ = sqlx::query("DELETE FROM blob_inventory WHERE sha256 = ?")
141                    .bind(sha256)
142                    .execute(&mut *conn)
143                    .await;
144            }
145            Err(_) => {}
146        }
147    }
148}
149
150fn write_object_atomically(path: &Path, bytes: &[u8]) -> Result<bool> {
151    if let Some(parent) = path.parent() {
152        fs::create_dir_all(parent).context("could not create attachment object directory")?;
153        if path.exists() {
154            let existing = fs::read(path).context("could not read attachment object")?;
155            if sha256_hex(&existing) != sha256_hex(bytes) {
156                bail!("error attachment-object-content-mismatch");
157            }
158            return Ok(false);
159        }
160        let mut temp = tempfile::Builder::new()
161            .prefix(".aven-stage-")
162            .tempfile_in(parent)
163            .context("could not create attachment staging file")?;
164        std::io::Write::write_all(&mut temp, bytes)?;
165        temp.as_file().sync_all()?;
166        match temp.persist_noclobber(path) {
167            Ok(_) => {
168                if let Err(error) =
169                    fs::File::open(parent).and_then(|directory| directory.sync_all())
170                {
171                    let _ = fs::remove_file(path);
172                    return Err(error.into());
173                }
174                return Ok(true);
175            }
176            Err(error) if error.error.kind() == std::io::ErrorKind::AlreadyExists => {
177                let existing = fs::read(path).context("could not read attachment object")?;
178                if sha256_hex(&existing) == sha256_hex(bytes) {
179                    return Ok(false);
180                }
181                bail!("error attachment-object-content-mismatch");
182            }
183            Err(error) => return Err(error.error.into()),
184        }
185    }
186    bail!("error attachment-object-path-invalid")
187}
188
189pub async fn upsert_inventory_available(
190    conn: &mut SqliteConnection,
191    sha256: &str,
192    byte_size: i64,
193    media_type: &str,
194) -> Result<()> {
195    validate_sha256(sha256)?;
196    validate_media_type(media_type)?;
197    let existing = blob_inventory_row(conn, sha256).await?;
198    let timestamp = now();
199    if let Some(row) = existing {
200        if row.byte_size != byte_size || row.media_type != media_type {
201            bail!("error blob-inventory-metadata-mismatch");
202        }
203        sqlx::query(
204            "UPDATE blob_inventory
205             SET available = 1, last_verified_at = ?
206             WHERE sha256 = ?",
207        )
208        .bind(&timestamp)
209        .bind(sha256)
210        .execute(&mut *conn)
211        .await?;
212        return Ok(());
213    }
214    sqlx::query(
215        "INSERT INTO blob_inventory(sha256, byte_size, media_type, available, first_seen_at, last_verified_at)
216         VALUES (?, ?, ?, 1, ?, ?)",
217    )
218    .bind(sha256)
219    .bind(byte_size)
220    .bind(media_type)
221    .bind(&timestamp)
222    .bind(&timestamp)
223    .execute(&mut *conn)
224    .await?;
225    Ok(())
226}
227
228pub async fn blob_inventory_row(
229    conn: &mut SqliteConnection,
230    sha256: &str,
231) -> Result<Option<BlobInventoryRow>> {
232    validate_sha256(sha256)?;
233    let row = sqlx::query!(
234        "SELECT sha256, byte_size, media_type, available, first_seen_at, last_verified_at
235         FROM blob_inventory WHERE sha256 = ?",
236        sha256
237    )
238    .fetch_optional(&mut *conn)
239    .await?;
240    Ok(row.map(|row| BlobInventoryRow {
241        sha256: row.sha256.unwrap_or_default(),
242        byte_size: row.byte_size,
243        media_type: row.media_type,
244        available: row.available != 0,
245        first_seen_at: row.first_seen_at,
246        last_verified_at: row.last_verified_at,
247    }))
248}
249
250#[cfg(test)]
251mod tests {
252    use std::io::Cursor;
253    use std::path::PathBuf;
254
255    use image::{DynamicImage, ImageFormat, RgbaImage};
256
257    use crate::db::open_db;
258
259    use super::*;
260
261    fn png_bytes() -> Vec<u8> {
262        let mut bytes = Cursor::new(Vec::new());
263        DynamicImage::ImageRgba8(RgbaImage::new(2, 1))
264            .write_to(&mut bytes, ImageFormat::Png)
265            .unwrap();
266        bytes.into_inner()
267    }
268
269    #[tokio::test]
270    async fn stores_and_retrieves_blob() {
271        let temp = tempfile::tempdir().unwrap();
272        let db_path = temp.path().join("test.sqlite");
273        let pool = open_db(&db_path).await.unwrap();
274        let mut conn = pool.acquire().await.unwrap();
275
276        let blob_dir = temp.path().join("blobs");
277
278        let bytes = png_bytes();
279        let stored = store_blob(&mut conn, &blob_dir, "image/png", &bytes)
280            .await
281            .unwrap();
282
283        assert_eq!(stored.sha256, sha256_hex(&bytes));
284        assert!(stored.byte_size > 0);
285
286        let obj_path = object_path(&blob_dir, &stored.sha256).unwrap();
287        assert!(obj_path.exists(), "blob file should exist on disk");
288
289        let on_disk = std::fs::read(&obj_path).unwrap();
290        assert_eq!(on_disk, bytes);
291
292        let row = blob_inventory_row(&mut conn, &stored.sha256)
293            .await
294            .unwrap()
295            .expect("inventory row should exist");
296        assert!(row.available);
297        assert_eq!(row.sha256, stored.sha256);
298        assert_eq!(row.byte_size, stored.byte_size);
299    }
300
301    #[tokio::test]
302    async fn stores_blob_idempotently() {
303        let temp = tempfile::tempdir().unwrap();
304        let db_path = temp.path().join("test.sqlite");
305        let pool = open_db(&db_path).await.unwrap();
306        let mut conn = pool.acquire().await.unwrap();
307
308        let blob_dir = temp.path().join("blobs");
309
310        let bytes = png_bytes();
311        let stored1 = store_blob(&mut conn, &blob_dir, "image/png", &bytes)
312            .await
313            .unwrap();
314        let stored2 = store_blob(&mut conn, &blob_dir, "image/png", &bytes)
315            .await
316            .unwrap();
317
318        assert_eq!(stored1.sha256, stored2.sha256);
319        let row = blob_inventory_row(&mut conn, &stored1.sha256)
320            .await
321            .unwrap()
322            .unwrap();
323        assert!(row.available);
324    }
325
326    #[tokio::test]
327    async fn rejects_inventory_metadata_mismatch() {
328        let temp = tempfile::tempdir().unwrap();
329        let db_path = temp.path().join("test.sqlite");
330        let pool = open_db(&db_path).await.unwrap();
331        let mut conn = pool.acquire().await.unwrap();
332
333        let sha256 = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
334        upsert_inventory_available(&mut conn, sha256, 12, "image/png")
335            .await
336            .unwrap();
337
338        let error = upsert_inventory_available(&mut conn, sha256, 13, "image/png")
339            .await
340            .unwrap_err()
341            .to_string();
342        assert_eq!(error, "error blob-inventory-metadata-mismatch");
343
344        let error = upsert_inventory_available(&mut conn, sha256, 12, "image/jpeg")
345            .await
346            .unwrap_err()
347            .to_string();
348        assert_eq!(error, "error blob-inventory-metadata-mismatch");
349    }
350
351    #[tokio::test]
352    async fn inventory_failure_removes_only_newly_staged_blob() {
353        let temp = tempfile::tempdir().unwrap();
354        let db_path = temp.path().join("test.sqlite");
355        let pool = open_db(&db_path).await.unwrap();
356        let mut conn = pool.acquire().await.unwrap();
357        let blob_dir = temp.path().join("blobs");
358        let bytes = png_bytes();
359        let sha256 = sha256_hex(&bytes);
360        upsert_inventory_available(&mut conn, &sha256, 12, "image/png")
361            .await
362            .unwrap();
363
364        assert!(
365            store_blob(&mut conn, &blob_dir, "image/png", &bytes)
366                .await
367                .is_err()
368        );
369        assert!(!object_path(&blob_dir, &sha256).unwrap().exists());
370    }
371
372    #[test]
373    fn computes_sha256_hex() {
374        let result = sha256_hex(b"hello");
375        assert_eq!(result.len(), 64);
376        assert!(result.bytes().all(|b| b.is_ascii_hexdigit()));
377    }
378
379    #[test]
380    fn computes_object_path() {
381        let blob_dir = Path::new("/tmp/aven");
382        let hash = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
383        let path = object_path(blob_dir, hash).unwrap();
384        assert_eq!(
385            path,
386            PathBuf::from(
387                "/tmp/aven/objects/sha256/abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789"
388            )
389        );
390    }
391
392    #[test]
393    fn rejects_invalid_hash_in_object_path() {
394        assert!(object_path(Path::new("/tmp"), "short").is_err());
395    }
396}