From 890244456118a60af0a5007397313ded086df0af Mon Sep 17 00:00:00 2001 From: qaz741wsd856 Date: Mon, 8 Dec 2025 11:11:27 +0000 Subject: [PATCH] feat: implement zero-copy streaming for attachment uploads - Added a new entry point in JavaScript to handle zero-copy streaming uploads to R2. - Updated the Rust backend to support the new upload route with JWT validation. - Enhanced the API to generate full URLs for the new azure-upload endpoint. - Introduced new dependencies: `futures-util` and `wasm-streams` for improved async handling. --- Cargo.lock | 2 + Cargo.toml | 2 + src/entry.js | 403 ++++++++++++++++++++++++++++++++++++ src/handlers/attachments.rs | 45 ++-- src/lib.rs | 10 +- src/router.rs | 2 + wrangler.toml | 2 +- 7 files changed, 440 insertions(+), 26 deletions(-) create mode 100644 src/entry.js diff --git a/Cargo.lock b/Cargo.lock index 62fdfe2..4cf1855 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1257,6 +1257,7 @@ dependencies = [ "console_error_panic_hook", "console_log", "constant_time_eq", + "futures-util", "getrandom 0.3.3", "glob-match", "hex", @@ -1273,6 +1274,7 @@ dependencies = [ "uuid", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams", "web-sys", "worker", "worker-macros 0.6.5", diff --git a/Cargo.toml b/Cargo.toml index e9de039..f6591c5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,12 +23,14 @@ js-sys = "0.3" wasm-bindgen-futures = "0.4" web-sys = { version = "0.3", features = ["Crypto", "CryptoKey", "SubtleCrypto", "WorkerGlobalScope", "Pbkdf2Params"] } console_error_panic_hook = "0.1.7" +wasm-streams = "0.4" # Axum and Routing axum = { version = "0.8", default-features = false, features=["json", "macros", "form", "multipart", "query"] } tower-service = "0.3" tower-http = { version = "0.5", features = ["cors"] } bytes = "1" +futures-util = { version = "0.3", default-features = false, features = ["alloc"] } # Data & Serialization serde = { version = "1.0", features = ["derive"] } diff --git a/src/entry.js b/src/entry.js new file mode 100644 index 0000000..6d82a11 --- /dev/null +++ b/src/entry.js @@ -0,0 +1,403 @@ +/** + * JS Wrapper Entry Point for Warden Worker + * + * This wrapper intercepts attachment upload requests for zero-copy streaming + * uploads to R2. Workers R2 binding can accept request.body directly. + * See: https://blog.cloudflare.com/zh-cn/r2-ga/ + * + * All other requests are passed through to the Rust WASM module. + */ + +import RustWorker from "../build/index.js"; + +// JWT validation using Web Crypto API (no external dependencies) +async function verifyJWT(token, secret) { + const encoder = new TextEncoder(); + const parts = token.split("."); + if (parts.length !== 3) { + throw new Error("Invalid token format"); + } + + const [headerB64, payloadB64, signatureB64] = parts; + + // Import the secret key for HMAC-SHA256 + const key = await crypto.subtle.importKey( + "raw", + encoder.encode(secret), + { name: "HMAC", hash: "SHA-256" }, + false, + ["verify"] + ); + + // Decode the signature (base64url to Uint8Array) + const signature = base64UrlDecode(signatureB64); + + // Verify the signature + const data = encoder.encode(`${headerB64}.${payloadB64}`); + const valid = await crypto.subtle.verify("HMAC", key, signature, data); + + if (!valid) { + throw new Error("Invalid token signature"); + } + + // Decode and parse the payload + const payload = JSON.parse( + new TextDecoder().decode(base64UrlDecode(payloadB64)) + ); + + // Check expiration + if (payload.exp && payload.exp < Math.floor(Date.now() / 1000)) { + throw new Error("Token expired"); + } + + return payload; +} + +function base64UrlDecode(str) { + // Convert base64url to base64 + let base64 = str.replace(/-/g, "+").replace(/_/g, "/"); + // Add padding if needed + while (base64.length % 4) { + base64 += "="; + } + const binary = atob(base64); + const bytes = new Uint8Array(binary.length); + for (let i = 0; i < binary.length; i++) { + bytes[i] = binary.charCodeAt(i); + } + return bytes; +} + +// Parse azure-upload route: /api/ciphers/{id}/attachment/{attachment_id}/azure-upload +function parseAzureUploadPath(path) { + const parts = path.replace(/^\//, "").split("/"); + // Expected: ["api", "ciphers", "{cipher_id}", "attachment", "{attachment_id}", "azure-upload"] + if ( + parts.length === 6 && + parts[0] === "api" && + parts[1] === "ciphers" && + parts[3] === "attachment" && + parts[5] === "azure-upload" + ) { + return { cipherId: parts[2], attachmentId: parts[4] }; + } + return null; +} + +// Extract token from query string +function extractTokenFromQuery(url) { + const params = new URL(url).searchParams; + return params.get("token"); +} + +// Generate ISO timestamp string +function nowString() { + return new Date().toISOString().replace("T", " ").replace("Z", ""); +} + +// Helper to get env var with fallback +function getEnvVar(env, name, defaultValue = null) { + try { + const value = env[name]; + if (value && typeof value.toString === "function") { + return value.toString(); + } + return value || defaultValue; + } catch { + return defaultValue; + } +} + +// Get attachment size limits from env +function getAttachmentMaxBytes(env) { + const value = getEnvVar(env, "ATTACHMENT_MAX_BYTES"); + if (!value) return null; + const parsed = parseInt(value, 10); + return isNaN(parsed) ? null : parsed; +} + +function getTotalLimitBytes(env) { + const value = getEnvVar(env, "ATTACHMENT_TOTAL_LIMIT_KB"); + if (!value) return null; + const kb = parseInt(value, 10); + if (isNaN(kb)) return null; + return kb * 1024; +} + +// 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`; + + const bindings = excludeAttachmentId ? [userId, excludeAttachmentId] : [userId]; + + const result = await db + .prepare(query) + .bind(...bindings) + .first(); + + return result?.total || 0; +} + +// Enforce attachment size limits +async function enforceLimits(db, env, userId, newSize, excludeAttachmentId) { + if (newSize < 0) { + throw new Error("Attachment size cannot be negative"); + } + + const maxBytes = getAttachmentMaxBytes(env); + if (maxBytes !== null && newSize > maxBytes) { + throw new Error("Attachment size exceeds limit"); + } + + const limitBytes = getTotalLimitBytes(env); + if (limitBytes !== null) { + const used = await getUserAttachmentUsage(db, userId, excludeAttachmentId); + const newTotal = used + newSize; + if (newTotal > limitBytes) { + throw new Error("Attachment storage limit reached"); + } + } +} + +// Handle azure-upload with zero-copy streaming +async function handleAzureUpload(request, env, cipherId, attachmentId, token) { + // Get R2 bucket + const bucket = env.ATTACHMENTS_BUCKET; + if (!bucket) { + return new Response(JSON.stringify({ error: "Attachments are not enabled" }), { + status: 400, + headers: { "Content-Type": "application/json" }, + }); + } + + // Get D1 database + const db = env.vault1; + if (!db) { + return new Response(JSON.stringify({ error: "Database not available" }), { + status: 500, + headers: { "Content-Type": "application/json" }, + }); + } + + // Validate JWT token + let claims; + try { + const secret = env.JWT_SECRET?.toString?.() || env.JWT_SECRET; + if (!secret) { + throw new Error("JWT_SECRET not configured"); + } + claims = await verifyJWT(token, secret); + } catch (err) { + return new Response(JSON.stringify({ error: `Invalid token: ${err.message}` }), { + status: 401, + headers: { "Content-Type": "application/json" }, + }); + } + + // Validate token claims match the request + if (claims.cipher_id !== cipherId || claims.attachment_id !== attachmentId) { + return new Response(JSON.stringify({ error: "Invalid download token" }), { + status: 401, + headers: { "Content-Type": "application/json" }, + }); + } + + const userId = claims.sub; + + // Verify cipher belongs to user and is not deleted + const cipher = await db + .prepare("SELECT * FROM ciphers WHERE id = ?1 AND user_id = ?2") + .bind(cipherId, userId) + .first(); + + if (!cipher) { + return new Response(JSON.stringify({ error: "Cipher not found" }), { + status: 404, + headers: { "Content-Type": "application/json" }, + }); + } + + if (cipher.organization_id) { + return new Response( + JSON.stringify({ error: "Organization attachments are not supported" }), + { status: 400, headers: { "Content-Type": "application/json" } } + ); + } + + if (cipher.deleted_at) { + return new Response( + JSON.stringify({ error: "Cannot modify attachments for deleted cipher" }), + { status: 400, headers: { "Content-Type": "application/json" } } + ); + } + + // Fetch attachment record + const attachment = await db + .prepare("SELECT * FROM attachments WHERE id = ?1") + .bind(attachmentId) + .first(); + + if (!attachment) { + return new Response(JSON.stringify({ error: "Attachment not found" }), { + status: 404, + headers: { "Content-Type": "application/json" }, + }); + } + + if (attachment.cipher_id !== cipherId) { + return new Response( + JSON.stringify({ error: "Attachment does not belong to cipher" }), + { status: 400, headers: { "Content-Type": "application/json" } } + ); + } + + // Get Content-Length from request headers + const contentLengthHeader = request.headers.get("Content-Length"); + if (!contentLengthHeader) { + return new Response( + JSON.stringify({ error: "Missing Content-Length header" }), + { status: 400, headers: { "Content-Type": "application/json" } } + ); + } + + const contentLength = parseInt(contentLengthHeader, 10); + if (isNaN(contentLength) || contentLength <= 0) { + return new Response( + JSON.stringify({ error: "Invalid Content-Length header" }), + { status: 400, headers: { "Content-Type": "application/json" } } + ); + } + + // Enforce limits before upload + try { + await enforceLimits(db, env, userId, contentLength, attachmentId); + } catch (err) { + return new Response(JSON.stringify({ error: err.message }), { + status: 400, + headers: { "Content-Type": "application/json" }, + }); + } + + // Build R2 key + const r2Key = `${cipherId}/${attachmentId}`; + + // Prepare R2 put options + const putOptions = {}; + const contentType = request.headers.get("Content-Type"); + if (contentType) { + putOptions.httpMetadata = { contentType }; + } + + // Upload to R2 directly using request.body + // Workers R2 binding can accept request.body directly for zero-copy streaming + // See: https://blog.cloudflare.com/zh-cn/r2-ga/ + let r2Object; + try { + r2Object = await bucket.put(r2Key, request.body, putOptions); + } catch (err) { + // Try to clean up on failure + try { + await bucket.delete(r2Key); + } catch { + // Ignore cleanup errors + } + return new Response( + JSON.stringify({ error: `Upload failed: ${err.message}` }), + { status: 500, headers: { "Content-Type": "application/json" } } + ); + } + + const uploadedSize = r2Object.size; + + // Verify uploaded size matches Content-Length + if (uploadedSize !== contentLength) { + try { + await bucket.delete(r2Key); + } catch { + // Ignore cleanup errors + } + return new Response( + JSON.stringify({ error: "Content-Length does not match uploaded size" }), + { status: 400, headers: { "Content-Type": "application/json" } } + ); + } + + // 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 + 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(); + + return new Response(null, { status: 201 }); +} + +// Main fetch handler +export default { + async fetch(request, env, ctx) { + const url = new URL(request.url); + + // Intercept PUT requests to azure-upload endpoint + if (request.method === "PUT") { + const parsed = parseAzureUploadPath(url.pathname); + if (parsed) { + const token = extractTokenFromQuery(request.url); + if (token) { + return handleAzureUpload( + request, + env, + parsed.cipherId, + parsed.attachmentId, + token + ); + } + } + } + + // Pass all other requests to Rust WASM + const worker = new RustWorker(ctx, env); + return worker.fetch(request); + }, + + async scheduled(event, env, ctx) { + // Pass scheduled events to Rust WASM + const worker = new RustWorker(ctx, env); + return worker.scheduled(event); + }, +}; + diff --git a/src/handlers/attachments.rs b/src/handlers/attachments.rs index e57aa0d..6379fe3 100644 --- a/src/handlers/attachments.rs +++ b/src/handlers/attachments.rs @@ -1,6 +1,7 @@ use std::{collections::HashMap, sync::Arc}; use axum::{ + body::Bytes, extract::{Multipart, Path, Query, State}, http::{header, StatusCode}, response::Response, @@ -170,7 +171,7 @@ pub async fn create_attachment_v2( .await?; // Return upload URL pointing to local upload endpoint - let url = upload_url(&cipher_id, &attachment_id); + let url = upload_url(&env, &base_url, &cipher_id, &attachment_id, &claims.sub)?; let mut cipher_response: Cipher = cipher.into(); hydrate_cipher_attachments(&db, &env, &base_url, &mut cipher_response, &claims.sub).await?; @@ -181,7 +182,7 @@ pub async fn create_attachment_v2( object: "attachment-fileUpload".to_string(), attachment_id, url, - file_upload_type: 0, // Direct + file_upload_type: 1, // Direct PUT with token cipher_response, })) } @@ -385,8 +386,9 @@ pub async fn delete_attachment( db::touch_user_updated_at(&db, &claims.sub).await?; // Reload cipher to return fresh updated_at and attachments state - let mut cipher_response: Cipher = - ensure_cipher_for_user(&db, &cipher_id, &claims.sub).await?.into(); + let mut cipher_response: Cipher = ensure_cipher_for_user(&db, &cipher_id, &claims.sub) + .await? + .into(); hydrate_cipher_attachments(&db, &env, &base_url, &mut cipher_response, &claims.sub).await?; Ok(Json(AttachmentDeleteResponse { @@ -536,10 +538,7 @@ fn is_not_found_error(err: &worker::Error) -> bool { msg.contains("NoSuchKey") || msg.contains("404") || msg.contains("NotFound") } -pub(crate) async fn delete_r2_objects( - bucket: &Bucket, - keys: &[String], -) -> Result<(), AppError> { +pub(crate) async fn delete_r2_objects(bucket: &Bucket, keys: &[String]) -> Result<(), AppError> { for key in keys { if let Err(err) = bucket.delete(key).await { if !is_not_found_error(&err) { @@ -627,10 +626,6 @@ pub(crate) async fn list_attachment_keys_for_soft_deleted_before( Ok(map_rows_to_keys(rows)) } -fn upload_url(cipher_id: &str, attachment_id: &str) -> String { - format!("/api/ciphers/{cipher_id}/attachment/{attachment_id}") -} - fn download_url( env: &Env, base_url: &str, @@ -746,16 +741,8 @@ async fn upload_to_r2( async fn read_multipart( multipart: &mut Multipart, -) -> Result< - ( - axum::body::Bytes, - Option, - Option, - Option, - ), - AppError, -> { - let mut file_bytes: Option = None; +) -> Result<(Bytes, Option, Option, Option), AppError> { + let mut file_bytes: Option = None; let mut content_type: Option = None; let mut key: Option = None; let mut file_name: Option = None; @@ -872,6 +859,20 @@ fn validate_download_token( Ok(claims.sub) } +fn upload_url( + env: &Env, + base_url: &str, + cipher_id: &str, + attachment_id: &str, + user_id: &str, +) -> Result { + let token = build_download_token(env, user_id, cipher_id, attachment_id)?; + let normalized_base = base_url.trim_end_matches('/'); + Ok(format!( + "{normalized_base}/api/ciphers/{cipher_id}/attachment/{attachment_id}/azure-upload?token={token}" + )) +} + fn jwt_secret(env: &Env) -> Result { Ok(env.secret("JWT_SECRET")?.to_string()) } diff --git a/src/lib.rs b/src/lib.rs index ff9b683..9e863b0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use axum::{extract::DefaultBodyLimit, Extension}; use tower_http::cors::{Any, CorsLayer}; use tower_service::Service; @@ -25,14 +27,16 @@ pub async fn main( console_error_panic_hook::set_once(); let _ = console_log::init_with_level(log::Level::Debug); - // Extract base URL from request URI - let uri = req.uri(); + // Extract base URL from the incoming request + let uri = req.uri().clone(); let base_url = format!( "{}://{}", uri.scheme_str().unwrap_or("https"), uri.authority().map(|a| a.as_str()).unwrap_or("localhost") ); + let env = Arc::new(env); + // Allow all origins for CORS, which is typical for a public API like Bitwarden's. let cors = CorsLayer::new() .allow_methods(Any) @@ -41,7 +45,7 @@ pub async fn main( let body_limit = attachment_body_limit_bytes(&env); - let mut app = router::api_router(env) + let mut app = router::api_router((*env).clone()) .layer(Extension(BaseUrl(base_url))) .layer(cors) // axum 默认 body 限制为 2MiB，附件上传需要更大的上限 diff --git a/src/router.rs b/src/router.rs index 261c378..9ce118c 100644 --- a/src/router.rs +++ b/src/router.rs @@ -60,6 +60,8 @@ pub fn api_router(env: Env) -> Router { "/api/ciphers/{id}/attachment/v2", post(attachments::create_attachment_v2), ) + // Note: Azure upload route is handled in entry.js for zero-copy body streaming + // PUT /api/ciphers/{id}/attachment/{attachment_id}/azure-upload .route( "/api/ciphers/{id}/attachment", post(attachments::upload_attachment_legacy), diff --git a/wrangler.toml b/wrangler.toml index 7ac9da0..15fe3c0 100644 --- a/wrangler.toml +++ b/wrangler.toml @@ -1,5 +1,5 @@ name = "warden-worker" -main = "build/index.js" +main = "src/entry.js" compatibility_date = "2025-09-19" keep_vars = true