use crate::storage_error::StorageError; use crate::util::{db, sync}; use rand::RngCore; use rusqlite::{OptionalExtension, Row, params}; use sha2::{Digest, Sha256}; use std::collections::HashSet; use std::fs::{self, File, OpenOptions}; use std::io::{Read, Seek, SeekFrom, Write}; use std::path::PathBuf; use std::sync::atomic::{AtomicI64, Ordering}; pub const MAX_CHUNK_BYTES: usize = 8 * 1024 * 1024; const ABANDONED_UPLOAD_MS: i64 = 24 * 60 * 60 * 1_000; const CLEANUP_INTERVAL_MS: i64 = 15 * 60 * 1_000; static LAST_CLEANUP: AtomicI64 = AtomicI64::new(0); #[derive(Debug, Clone)] pub struct UserAssetMetadata { pub id: i64, pub user_id: i64, pub asset_id: String, pub size: i64, pub sha256: Vec, pub mime_type: String, pub revision: i64, pub created_at: i64, pub updated_at: i64, } #[derive(Debug, Clone)] pub struct UploadStatus { pub upload_id: String, pub asset_id: String, pub size: i64, pub sha256: Vec, pub mime_type: String, pub offset: i64, pub state: String, pub revision: Option, pub created_at: i64, pub updated_at: i64, pub committed_at: Option, } fn asset_dir() -> Result { let path = iota_util::file_util::storage_directory().join("assets"); fs::create_dir_all(&path)?; Ok(path) } fn random_file_name() -> String { let mut bytes = [0_u8; 24]; rand::thread_rng().fill_bytes(&mut bytes); bytes.iter().map(|byte| format!("{byte:02x}")).collect() } fn validate(user_id: i64, asset_id: &str, upload_id: Option<&str>) -> Result<(), StorageError> { if user_id <= 0 || asset_id.is_empty() || asset_id.len() > 256 { return Err(StorageError::Other("invalid asset owner or id".into())); } if upload_id.is_some_and(|id| id.is_empty() || id.len() > 128) { return Err(StorageError::Other("invalid upload id".into())); } Ok(()) } fn upload_from_row(row: &Row<'_>) -> rusqlite::Result<(UploadStatus, String)> { Ok(( UploadStatus { upload_id: row.get(0)?, asset_id: row.get(2)?, size: row.get(4)?, sha256: row.get(5)?, mime_type: row.get(6)?, offset: row.get(7)?, state: row.get(8)?, revision: row.get(9)?, created_at: row.get(10)?, updated_at: row.get(11)?, committed_at: row.get(12)?, }, row.get(3)?, )) } fn metadata_from_row(row: &Row<'_>) -> rusqlite::Result { Ok(UserAssetMetadata { id: row.get(0)?, user_id: row.get(1)?, asset_id: row.get(2)?, size: row.get(4)?, sha256: row.get(5)?, mime_type: row.get(6)?, revision: row.get(7)?, created_at: row.get(8)?, updated_at: row.get(9)?, }) } fn load_upload( user_id: i64, upload_id: &str, ) -> Result, StorageError> { db::with_db(|conn| { Ok(conn.query_row( "SELECT upload_id, user_id, asset_id, file_name, size, sha256, mime_type, received, state, revision, created_at, updated_at, committed_at FROM user_asset_uploads WHERE user_id = ?1 AND upload_id = ?2", params![user_id, upload_id], upload_from_row, ).optional()?) }) } fn reconcile_offset( user_id: i64, mut status: UploadStatus, file_name: &str, ) -> Result { if status.state != "uploading" { return Ok(status); } let actual = fs::metadata(asset_dir()?.join(file_name))?.len() as i64; if actual > status.size { return Err(StorageError::Other( "upload file exceeds declared size".into(), )); } if actual != status.offset { let now = sync::now_millis(); db::with_db(|conn| { conn.execute( "UPDATE user_asset_uploads SET received = ?3, updated_at = ?4 WHERE user_id = ?1 AND upload_id = ?2 AND state = 'uploading'", params![user_id, status.upload_id, actual, now], )?; Ok(()) })?; status.offset = actual; status.updated_at = now; } Ok(status) } pub fn start( user_id: i64, upload_id: &str, asset_id: &str, size: i64, sha256: &[u8], mime_type: &str, ) -> Result { validate(user_id, asset_id, Some(upload_id))?; if size < 0 || sha256.len() != 32 || mime_type.is_empty() || mime_type.len() > 255 { return Err(StorageError::Other("invalid asset metadata".into())); } cleanup_if_due()?; if let Some((status, file_name)) = load_upload(user_id, upload_id)? { if status.asset_id != asset_id || status.size != size || status.sha256 != sha256 || status.mime_type != mime_type { return Err(StorageError::Other( "upload id already has different metadata".into(), )); } return reconcile_offset(user_id, status, &file_name); } let file_name = random_file_name(); let path = asset_dir()?.join(&file_name); OpenOptions::new() .write(true) .create_new(true) .open(&path)? .sync_all()?; let now = sync::now_millis(); let inserted = db::with_db(|conn| { conn.execute( "INSERT INTO user_asset_uploads (upload_id, user_id, asset_id, file_name, size, sha256, mime_type, received, state, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, 0, 'uploading', ?8, ?8)", params![upload_id, user_id, asset_id, file_name, size, sha256, mime_type, now], )?; Ok(()) }); if let Err(error) = inserted { let _ = fs::remove_file(path); return Err(error); } Ok(UploadStatus { upload_id: upload_id.into(), asset_id: asset_id.into(), size, sha256: sha256.into(), mime_type: mime_type.into(), offset: 0, state: "uploading".into(), revision: None, created_at: now, updated_at: now, committed_at: None, }) } pub fn status(user_id: i64, upload_id: &str) -> Result, StorageError> { let Some((status, file_name)) = load_upload(user_id, upload_id)? else { return Ok(None); }; reconcile_offset(user_id, status, &file_name).map(Some) } pub fn append( user_id: i64, upload_id: &str, offset: i64, chunk: &[u8], ) -> Result { if chunk.is_empty() || chunk.len() > MAX_CHUNK_BYTES { return Err(StorageError::Other("invalid asset chunk size".into())); } let Some((status, file_name)) = load_upload(user_id, upload_id)? else { return Err(StorageError::Other("upload not found".into())); }; let status = reconcile_offset(user_id, status, &file_name)?; if status.state != "uploading" || offset != status.offset || offset + chunk.len() as i64 > status.size { return Err(StorageError::Other( "asset chunk offset does not match acknowledged offset".into(), )); } let path = asset_dir()?.join(file_name); let mut file = OpenOptions::new().append(true).open(path)?; file.write_all(chunk)?; file.sync_data()?; let received = offset + chunk.len() as i64; let now = sync::now_millis(); db::with_db(|conn| { conn.execute( "UPDATE user_asset_uploads SET received = ?3, updated_at = ?4 WHERE user_id = ?1 AND upload_id = ?2 AND state = 'uploading'", params![user_id, upload_id, received, now], )?; Ok(()) })?; Ok(UploadStatus { offset: received, updated_at: now, ..status }) } pub fn commit(user_id: i64, upload_id: &str) -> Result<(UserAssetMetadata, bool), StorageError> { let Some((status, file_name)) = load_upload(user_id, upload_id)? else { return Err(StorageError::Other("upload not found".into())); }; if status.state == "committed" { return Ok(( UserAssetMetadata { id: 0, user_id, asset_id: status.asset_id, size: status.size, sha256: status.sha256, mime_type: status.mime_type, revision: status.revision.ok_or_else(|| { StorageError::Other("committed upload has no revision".into()) })?, created_at: status.created_at, updated_at: status.committed_at.unwrap_or(status.updated_at), }, false, )); } let status = reconcile_offset(user_id, status, &file_name)?; if status.offset != status.size { return Err(StorageError::Other("upload is incomplete".into())); } let path = asset_dir()?.join(&file_name); let mut file = File::open(&path)?; let mut hash = Sha256::new(); let mut buffer = [0_u8; 64 * 1024]; loop { let read = file.read(&mut buffer)?; if read == 0 { break; } hash.update(&buffer[..read]); } if hash.finalize().as_slice() != status.sha256 { return Err(StorageError::Other("asset SHA-256 mismatch".into())); } let now = sync::now_millis(); let (metadata, replaced_file) = db::with_immediate_transaction(|tx| { let existing = tx.query_row( "SELECT id, file_name, created_at FROM user_assets WHERE user_id = ?1 AND asset_id = ?2", params![user_id, status.asset_id], |row| Ok((row.get::<_, i64>(0)?, row.get::<_, String>(1)?, row.get::<_, i64>(2)?)), ).optional()?; let id = if let Some(existing) = &existing { existing.0 } else { tx.execute( "INSERT INTO user_assets (user_id, asset_id, file_name, size, sha256, mime_type, revision, created_at, updated_at, published_upload_id, deleted) VALUES (?1, ?2, ?3, ?4, ?5, ?6, 0, ?7, ?7, ?8, 0)", params![user_id, status.asset_id, file_name, status.size, status.sha256, status.mime_type, now, upload_id], )?; tx.last_insert_rowid() }; let revision = sync::record_event( tx, user_id, sync::EntityType::UserAsset, id, sync::Operation::Upsert, )?; let created_at = existing.as_ref().map(|value| value.2).unwrap_or(now); tx.execute( "UPDATE user_assets SET file_name = ?2, size = ?3, sha256 = ?4, mime_type = ?5, revision = ?6, updated_at = ?7, published_upload_id = ?8, deleted = 0 WHERE id = ?1", params![id, file_name, status.size, status.sha256, status.mime_type, revision, now, upload_id], )?; tx.execute( "UPDATE user_asset_uploads SET state = 'committed', revision = ?3, updated_at = ?4, committed_at = ?4 WHERE user_id = ?1 AND upload_id = ?2", params![user_id, upload_id, revision, now], )?; Ok(( UserAssetMetadata { id, user_id, asset_id: status.asset_id.clone(), size: status.size, sha256: status.sha256.clone(), mime_type: status.mime_type.clone(), revision, created_at, updated_at: now, }, existing .map(|value| value.1) .filter(|old| old != &file_name), )) })?; if let Some(old) = replaced_file { let _ = fs::remove_file(asset_dir()?.join(old)); } Ok((metadata, true)) } pub fn abort(user_id: i64, upload_id: &str) -> Result { let upload = load_upload(user_id, upload_id)?; let Some((status, file_name)) = upload else { return Ok(false); }; if status.state == "committed" { return Ok(false); } db::with_db(|conn| { conn.execute("DELETE FROM user_asset_uploads WHERE user_id = ?1 AND upload_id = ?2 AND state = 'uploading'", params![user_id, upload_id])?; Ok(()) })?; match fs::remove_file(asset_dir()?.join(file_name)) { Ok(()) => Ok(true), Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(true), Err(error) => Err(error.into()), } } pub fn list(user_id: i64) -> Result, StorageError> { cleanup_if_due()?; db::with_db(|conn| { let mut statement = conn.prepare("SELECT id, user_id, asset_id, file_name, size, sha256, mime_type, revision, created_at, updated_at FROM user_assets WHERE user_id = ?1 AND deleted = 0 ORDER BY updated_at DESC")?; statement .query_map([user_id], metadata_from_row)? .collect::, _>>() .map_err(StorageError::from) }) } pub fn read_chunk( user_id: i64, asset_id: &str, offset: i64, amount: usize, ) -> Result)>, StorageError> { validate(user_id, asset_id, None)?; if offset < 0 || amount == 0 || amount > MAX_CHUNK_BYTES { return Err(StorageError::Other("invalid asset chunk range".into())); } let value = db::with_db(|conn| { Ok(conn.query_row( "SELECT id, user_id, asset_id, file_name, size, sha256, mime_type, revision, created_at, updated_at FROM user_assets WHERE user_id = ?1 AND asset_id = ?2 AND deleted = 0", params![user_id, asset_id], |row| Ok((metadata_from_row(row)?, row.get::<_, String>(3)?)), ).optional()?) })?; let Some((metadata, file_name)) = value else { return Ok(None); }; if offset > metadata.size { return Err(StorageError::Other("asset offset exceeds size".into())); } let mut file = File::open(asset_dir()?.join(file_name))?; file.seek(SeekFrom::Start(offset as u64))?; let mut bytes = vec![0; amount.min((metadata.size - offset) as usize)]; file.read_exact(&mut bytes)?; Ok(Some((metadata, bytes))) } pub fn delete(user_id: i64, asset_id: &str) -> Result, StorageError> { validate(user_id, asset_id, None)?; let deleted = db::with_immediate_transaction(|tx| { let value = tx.query_row( "SELECT id, user_id, asset_id, file_name, size, sha256, mime_type, revision, created_at, updated_at FROM user_assets WHERE user_id = ?1 AND asset_id = ?2 AND deleted = 0", params![user_id, asset_id], |row| Ok((metadata_from_row(row)?, row.get::<_, String>(3)?)), ).optional()?; let Some((mut metadata, file_name)) = value else { return Ok(None); }; metadata.revision = sync::record_event( tx, user_id, sync::EntityType::UserAsset, metadata.id, sync::Operation::Delete, )?; metadata.updated_at = sync::now_millis(); tx.execute( "UPDATE user_assets SET deleted = 1, revision = ?2, updated_at = ?3 WHERE id = ?1", params![metadata.id, metadata.revision, metadata.updated_at], )?; Ok(Some((metadata, file_name))) })?; let Some((metadata, file_name)) = deleted else { return Ok(None); }; let _ = fs::remove_file(asset_dir()?.join(file_name)); Ok(Some(metadata)) } pub fn purge_user(user_id: i64) -> Result<(), StorageError> { let files = db::with_immediate_transaction(|tx| { let mut files = HashSet::new(); for query in [ "SELECT file_name FROM user_assets WHERE user_id = ?1", "SELECT file_name FROM user_asset_uploads WHERE user_id = ?1", ] { let mut statement = tx.prepare(query)?; for file in statement.query_map([user_id], |row| row.get::<_, String>(0))? { files.insert(file?); } } tx.execute("DELETE FROM user_assets WHERE user_id = ?1", [user_id])?; tx.execute( "DELETE FROM user_asset_uploads WHERE user_id = ?1", [user_id], )?; Ok(files) })?; let directory = asset_dir()?; for file in files { match fs::remove_file(directory.join(file)) { Ok(()) => {} Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} Err(error) => return Err(error.into()), } } Ok(()) } fn cleanup_if_due() -> Result<(), StorageError> { let now = sync::now_millis(); let previous = LAST_CLEANUP.load(Ordering::Relaxed); if now - previous < CLEANUP_INTERVAL_MS || LAST_CLEANUP .compare_exchange(previous, now, Ordering::Relaxed, Ordering::Relaxed) .is_err() { return Ok(()); } let cutoff = now - ABANDONED_UPLOAD_MS; let abandoned = db::with_db(|conn| { let mut statement = conn.prepare("SELECT file_name FROM user_asset_uploads WHERE state = 'uploading' AND updated_at < ?1")?; let files = statement .query_map([cutoff], |row| row.get::<_, String>(0))? .collect::, _>>()?; conn.execute( "DELETE FROM user_asset_uploads WHERE state = 'uploading' AND updated_at < ?1", [cutoff], )?; Ok(files) })?; let dir = asset_dir()?; for file in abandoned { let _ = fs::remove_file(dir.join(file)); } let referenced = db::with_db(|conn| { let mut files = HashSet::new(); for query in [ "SELECT file_name FROM user_assets WHERE deleted = 0", "SELECT file_name FROM user_asset_uploads WHERE state = 'uploading'", ] { let mut statement = conn.prepare(query)?; for file in statement.query_map([], |row| row.get::<_, String>(0))? { files.insert(file?); } } Ok(files) })?; for entry in fs::read_dir(dir)? { let entry = entry?; if entry.file_type()?.is_file() && !referenced.contains(&entry.file_name().to_string_lossy().into_owned()) { let _ = fs::remove_file(entry.path()); } } Ok(()) }