danvics/warden-worker
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

feat: implement pending attachments feature in upload handling

Browse source
This commit is contained in:
qaz741wsd856 2025-12-11 03:22:10 +00:00
parent 5e04837b5d
commit 023b010504
6 changed files with 234 additions and 80 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
migrations/0006_add_pending_attachments.sql Normal file
View file

@ -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);

15
sql/schema.sql
View file

@ -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 (

100
src/entry.js
View file

@ -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 });
}

131
src/handlers/attachments.rs
View file

@ -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<Attach
.ok_or_else(|| AppError::NotFound("Attachment not found".to_string()))
}
async fn fetch_pending_attachment(
db: &D1Database,
attachment_id: &str,
) -> Result<AttachmentDB, AppError> {
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<i64, AppError> {
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<JsValue>) = 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<JsValue> = vec![JsValue::from_str(user_id)];
if let Some(id) = exclude_attachment {
bindings.push(JsValue::from_str(id));
}
let row: Option<Value> = db
.prepare(query_str)
.bind(&bindings)?

39
src/handlers/purge.rs
View file

@ -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<u32, worker::Error> {
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::<CountResult>(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:

13
src/lib.rs
View file

@ -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 {