diff --git a/migrations/0006_add_pending_attachments.sql b/migrations/0006_add_pending_attachments.sql new file mode 100644 index 0000000..3744f2d --- /dev/null +++ b/migrations/0006_add_pending_attachments.sql @@ -0,0 +1,16 @@ +-- Track attachments that are still uploading; moved to attachments on success. +CREATE TABLE IF NOT EXISTS attachments_pending ( + id TEXT PRIMARY KEY NOT NULL, + cipher_id TEXT NOT NULL, + file_name TEXT NOT NULL, + file_size INTEGER NOT NULL, + akey TEXT, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + organization_id TEXT, + FOREIGN KEY (cipher_id) REFERENCES ciphers(id) ON DELETE CASCADE +); + +CREATE INDEX IF NOT EXISTS idx_attachments_pending_cipher ON attachments_pending(cipher_id); +CREATE INDEX IF NOT EXISTS idx_attachments_pending_created_at ON attachments_pending(created_at); + diff --git a/sql/schema.sql b/sql/schema.sql index 17001dd..ec318ad 100644 --- a/sql/schema.sql +++ b/sql/schema.sql @@ -51,6 +51,21 @@ CREATE TABLE IF NOT EXISTS attachments ( ); CREATE INDEX IF NOT EXISTS idx_attachments_cipher ON attachments(cipher_id); +-- Pending attachments table for in-flight uploads +CREATE TABLE IF NOT EXISTS attachments_pending ( + id TEXT PRIMARY KEY NOT NULL, + cipher_id TEXT NOT NULL, + file_name TEXT NOT NULL, + file_size INTEGER NOT NULL, + akey TEXT, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + organization_id TEXT, + FOREIGN KEY (cipher_id) REFERENCES ciphers(id) ON DELETE CASCADE +); +CREATE INDEX IF NOT EXISTS idx_attachments_pending_cipher ON attachments_pending(cipher_id); +CREATE INDEX IF NOT EXISTS idx_attachments_pending_created_at ON attachments_pending(created_at); + -- TwoFactor table for two-factor authentication -- Types: 0=Authenticator(TOTP), 1=Email, 5=Remember, 8=RecoveryCode CREATE TABLE IF NOT EXISTS twofactor ( diff --git a/src/entry.js b/src/entry.js index 8ef096a..b91efd6 100644 --- a/src/entry.js +++ b/src/entry.js @@ -147,16 +147,32 @@ function getTotalLimitBytes(env) { // Get user's current attachment usage async function getUserAttachmentUsage(db, userId, excludeAttachmentId) { const query = excludeAttachmentId - ? `SELECT COALESCE(SUM(a.file_size), 0) as total - FROM attachments a - JOIN ciphers c ON c.id = a.cipher_id - WHERE c.user_id = ?1 AND a.id != ?2` - : `SELECT COALESCE(SUM(a.file_size), 0) as total - FROM attachments a - JOIN ciphers c ON c.id = a.cipher_id - WHERE c.user_id = ?1`; + ? `SELECT COALESCE(SUM(file_size), 0) as total FROM ( + SELECT a.file_size as file_size + FROM attachments a + JOIN ciphers c ON c.id = a.cipher_id + WHERE c.user_id = ?1 AND a.id != ?2 + UNION ALL + SELECT p.file_size as file_size + FROM attachments_pending p + JOIN ciphers c2 ON c2.id = p.cipher_id + WHERE c2.user_id = ?1 AND p.id != ?2 + ) AS files` + : `SELECT COALESCE(SUM(file_size), 0) as total FROM ( + SELECT a.file_size as file_size + FROM attachments a + JOIN ciphers c ON c.id = a.cipher_id + WHERE c.user_id = ?1 + UNION ALL + SELECT p.file_size as file_size + FROM attachments_pending p + JOIN ciphers c2 ON c2.id = p.cipher_id + WHERE c2.user_id = ?1 + ) AS files`; - const bindings = excludeAttachmentId ? [userId, excludeAttachmentId] : [userId]; + const bindings = excludeAttachmentId + ? [userId, excludeAttachmentId] + : [userId]; const result = await db .prepare(query) @@ -259,20 +275,20 @@ async function handleAzureUpload(request, env, cipherId, attachmentId, token) { ); } - // Fetch attachment record - const attachment = await db - .prepare("SELECT * FROM attachments WHERE id = ?1") + // Fetch pending attachment record + const pending = await db + .prepare("SELECT * FROM attachments_pending WHERE id = ?1") .bind(attachmentId) .first(); - if (!attachment) { - return new Response(JSON.stringify({ error: "Attachment not found" }), { + if (!pending) { + return new Response(JSON.stringify({ error: "Attachment not found or already uploaded" }), { status: 404, headers: { "Content-Type": "application/json" }, }); } - if (attachment.cipher_id !== cipherId) { + if (pending.cipher_id !== cipherId) { return new Response( JSON.stringify({ error: "Attachment does not belong to cipher" }), { status: 400, headers: { "Content-Type": "application/json" } } @@ -350,39 +366,29 @@ async function handleAzureUpload(request, env, cipherId, attachmentId, token) { ); } - // Re-enforce limits with actual uploaded size (should be same, but be safe) - try { - await enforceLimits(db, env, userId, uploadedSize, attachmentId); - } catch (err) { - try { - await bucket.delete(r2Key); - } catch { - // Ignore cleanup errors - } - return new Response(JSON.stringify({ error: err.message }), { - status: 400, - headers: { "Content-Type": "application/json" }, - }); - } - - // Update attachment record in database + // Finalize upload: move pending -> attachments and touch revision timestamps const now = nowString(); - await db - .prepare("UPDATE attachments SET file_size = ?1, updated_at = ?2 WHERE id = ?3") - .bind(uploadedSize, now, attachmentId) - .run(); - - // Touch cipher updated_at - await db - .prepare("UPDATE ciphers SET updated_at = ?1 WHERE id = ?2") - .bind(now, cipherId) - .run(); - - // Touch user updated_at - await db - .prepare("UPDATE users SET updated_at = ?1 WHERE id = ?2") - .bind(now, userId) - .run(); + await db.batch([ + db + .prepare( + "INSERT INTO attachments (id, cipher_id, file_name, file_size, akey, created_at, updated_at, organization_id) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)" + ) + .bind( + attachmentId, + cipherId, + pending.file_name, + uploadedSize, + pending.akey, + pending.created_at || now, + now, + pending.organization_id || null + ), + db + .prepare("DELETE FROM attachments_pending WHERE id = ?1") + .bind(attachmentId), + db.prepare("UPDATE ciphers SET updated_at = ?1 WHERE id = ?2").bind(now, cipherId), + db.prepare("UPDATE users SET updated_at = ?1 WHERE id = ?2").bind(now, userId), + ]); return new Response(null, { status: 201 }); } diff --git a/src/handlers/attachments.rs b/src/handlers/attachments.rs index ea7559e..59a450a 100644 --- a/src/handlers/attachments.rs +++ b/src/handlers/attachments.rs @@ -146,7 +146,7 @@ pub async fn create_attachment_v2( query!( &db, - "INSERT INTO attachments (id, cipher_id, file_name, file_size, akey, created_at, updated_at, organization_id) + "INSERT INTO attachments_pending (id, cipher_id, file_name, file_size, akey, created_at, updated_at, organization_id) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?6, ?7)", attachment_id, cipher.id, @@ -165,8 +165,26 @@ pub async fn create_attachment_v2( let mut cipher_response: Cipher = cipher.into(); hydrate_cipher_attachments(&db, &env, &mut cipher_response).await?; - touch_cipher_updated_at(&db, &cipher_id).await?; - db::touch_user_updated_at(&db, &claims.sub).await?; + // add pending attachment to response + let pending_attachment = AttachmentDB { + id: attachment_id.clone(), + cipher_id: cipher_id.clone(), + file_name, + file_size: declared_size, + akey: Some(key), + created_at: now.clone(), + updated_at: now, + organization_id: cipher_response.organization_id.clone(), + }; + + let pending_response = pending_attachment.to_response(None); + match &mut cipher_response.attachments { + Some(list) => list.push(pending_response), + None => cipher_response.attachments = Some(vec![pending_response]), + } + + // no need to touch cipher updated_at and user updated_at here + // it will be touched in after upload Ok(Json(AttachmentUploadResponse { object: "attachment-fileUpload".to_string(), @@ -190,8 +208,8 @@ pub async fn upload_attachment_v2_data( let _cipher = ensure_cipher_for_user(&db, &cipher_id, &claims.sub).await?; - let mut attachment = fetch_attachment(&db, &attachment_id).await?; - if attachment.cipher_id != cipher_id { + let mut pending = fetch_pending_attachment(&db, &attachment_id).await?; + if pending.cipher_id != cipher_id { return Err(AppError::BadRequest( "Attachment does not belong to cipher".to_string(), )); @@ -202,45 +220,63 @@ pub async fn upload_attachment_v2_data( let actual_size = file_bytes.len() as i64; // Validate actual size against declared value deviation - if let Err(e) = validate_size_within_declared(&attachment, actual_size) { - query!(&db, "DELETE FROM attachments WHERE id = ?1", attachment.id) - .map_err(|_| AppError::Database)? - .run() - .await?; + if let Err(e) = validate_size_within_declared(&pending, actual_size) { + query!( + &db, + "DELETE FROM attachments_pending WHERE id = ?1", + pending.id + ) + .map_err(|_| AppError::Database)? + .run() + .await?; return Err(e); } // Validate capacity limits (replace with actual size) - enforce_limits(&db, &env, &claims.sub, actual_size, Some(&attachment.id)).await?; + enforce_limits(&db, &env, &claims.sub, actual_size, Some(&pending.id)).await?; // Need a key - if attachment.akey.is_none() && key_override.is_none() { + if pending.akey.is_none() && key_override.is_none() { return Err(AppError::BadRequest( "No attachment key provided".to_string(), )); } if let Some(k) = key_override { - attachment.akey = Some(k); + pending.akey = Some(k); } // Save to R2 upload_to_r2( &bucket, - &attachment.r2_key(), + &pending.r2_key(), content_type, file_bytes.to_vec(), ) .await?; - // Update metadata + // Finalize: move pending -> attachments and touch timestamps let now = now_string(); query!( &db, - "UPDATE attachments SET file_size = ?1, akey = COALESCE(?2, akey), updated_at = ?3 WHERE id = ?4", + "INSERT INTO attachments (id, cipher_id, file_name, file_size, akey, created_at, updated_at, organization_id) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", + pending.id, + pending.cipher_id, + pending.file_name, actual_size, - attachment.akey, + pending.akey, + pending.created_at, now, - attachment.id, + pending.organization_id, + ) + .map_err(|_| AppError::Database)? + .run() + .await?; + + query!( + &db, + "DELETE FROM attachments_pending WHERE id = ?1", + pending.id ) .map_err(|_| AppError::Database)? .run() @@ -312,7 +348,7 @@ pub async fn upload_attachment_legacy( touch_cipher_updated_at(&db, &cipher_id).await?; db::touch_user_updated_at(&db, &claims.sub).await?; - // 返回最新的 Cipher（含附件） + // reload cipher to return fresh updated_at and attachments state let mut cipher_response: Cipher = cipher.into(); hydrate_cipher_attachments(&db, &env, &mut cipher_response).await?; @@ -571,6 +607,18 @@ async fn fetch_attachment(db: &D1Database, attachment_id: &str) -> Result Result { + db.prepare("SELECT * FROM attachments_pending WHERE id = ?1") + .bind(&[attachment_id.into()])? + .first(None) + .await + .map_err(|_| AppError::Database)? + .ok_or_else(|| AppError::NotFound("Attachment not found".to_string())) +} + async fn load_attachment_map_json( db: &D1Database, json_body: &str, @@ -842,23 +890,40 @@ async fn user_attachment_usage( user_id: &str, exclude_attachment: Option<&str>, ) -> Result { - let query_str = if exclude_attachment.is_some() { - "SELECT COALESCE(SUM(a.file_size), 0) as total - FROM attachments a - JOIN ciphers c ON c.id = a.cipher_id - WHERE c.user_id = ?1 AND a.id != ?2" + let (query_str, bindings): (String, Vec) = if let Some(id) = exclude_attachment { + ( + "SELECT COALESCE(SUM(file_size), 0) as total FROM ( + SELECT a.file_size AS file_size + FROM attachments a + JOIN ciphers c ON c.id = a.cipher_id + WHERE c.user_id = ?1 AND a.id != ?2 + UNION ALL + SELECT p.file_size AS file_size + FROM attachments_pending p + JOIN ciphers c2 ON c2.id = p.cipher_id + WHERE c2.user_id = ?1 AND p.id != ?2 + ) AS files" + .to_string(), + vec![JsValue::from_str(user_id), JsValue::from_str(id)], + ) } else { - "SELECT COALESCE(SUM(a.file_size), 0) as total - FROM attachments a - JOIN ciphers c ON c.id = a.cipher_id - WHERE c.user_id = ?1" + ( + "SELECT COALESCE(SUM(file_size), 0) as total FROM ( + SELECT a.file_size AS file_size + FROM attachments a + JOIN ciphers c ON c.id = a.cipher_id + WHERE c.user_id = ?1 + UNION ALL + SELECT p.file_size AS file_size + FROM attachments_pending p + JOIN ciphers c2 ON c2.id = p.cipher_id + WHERE c2.user_id = ?1 + ) AS files" + .to_string(), + vec![JsValue::from_str(user_id)], + ) }; - let mut bindings: Vec = vec![JsValue::from_str(user_id)]; - if let Some(id) = exclude_attachment { - bindings.push(JsValue::from_str(id)); - } - let row: Option = db .prepare(query_str) .bind(&bindings)? diff --git a/src/handlers/purge.rs b/src/handlers/purge.rs index ce5b126..3d56a85 100644 --- a/src/handlers/purge.rs +++ b/src/handlers/purge.rs @@ -11,6 +11,8 @@ use worker::{query, D1Database, Env}; /// Default number of days to keep soft-deleted items before purging const DEFAULT_PURGE_DAYS: i64 = 30; +/// Retain pending attachments for at most this many days before cleanup +const PENDING_RETENTION_DAYS: i64 = 1; /// Get the purge threshold days from environment variable or use default fn get_purge_days(env: &Env) -> i64 { @@ -20,6 +22,43 @@ fn get_purge_days(env: &Env) -> i64 { .unwrap_or(DEFAULT_PURGE_DAYS) } +/// Purge pending attachments older than the configured retention window. +pub async fn purge_stale_pending_attachments(env: &Env) -> Result { + let db: D1Database = env.d1("vault1")?; + let now = Utc::now(); + let pending_cutoff = now - Duration::days(PENDING_RETENTION_DAYS); + let pending_cutoff_str = pending_cutoff.format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string(); + + let pending_count_result = query!( + &db, + "SELECT COUNT(*) as count FROM attachments_pending WHERE created_at < ?1", + pending_cutoff_str + )? + .first::(None) + .await?; + + let pending_count = pending_count_result.map(|r| r.count).unwrap_or(0); + + if pending_count > 0 { + query!( + &db, + "DELETE FROM attachments_pending WHERE created_at < ?1", + pending_cutoff_str + )? + .run() + .await?; + log::info!( + "Purged {} pending attachment(s) older than {} day(s)", + pending_count, + PENDING_RETENTION_DAYS + ); + } else { + log::info!("No pending attachments to purge"); + } + + Ok(pending_count) +} + /// Purge soft-deleted ciphers that are older than the configured threshold. /// /// This function: diff --git a/src/lib.rs b/src/lib.rs index b9e89de..9ddf828 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -66,6 +66,19 @@ pub async fn scheduled(_event: ScheduledEvent, env: Env, _ctx: ScheduleContext) console_error_panic_hook::set_once(); let _ = console_log::init_with_level(log::Level::Debug); + log::info!("Scheduled task triggered: purging stale pending attachments"); + match handlers::purge::purge_stale_pending_attachments(&env).await { + Ok(count) => { + log::info!( + "Pending attachment purge completed: {} record(s) removed", + count + ); + } + Err(e) => { + log::error!("Pending attachment purge failed: {:?}", e); + } + } + log::info!("Scheduled task triggered: purging soft-deleted ciphers"); match handlers::purge::purge_deleted_ciphers(&env).await {