From 61a3539e238a5a3fb81c31d2fc76aa54946870e1 Mon Sep 17 00:00:00 2001 From: qaz741wsd856 Date: Mon, 8 Dec 2025 13:29:55 +0000 Subject: [PATCH] feat: add zero-copy streaming for attachment downloads using js wapper --- src/entry.js | 149 +++++++++++++++++++++++++++++++++++- src/handlers/attachments.rs | 97 +---------------------- src/lib.rs | 32 +------- src/router.rs | 7 +- 4 files changed, 156 insertions(+), 129 deletions(-) diff --git a/src/entry.js b/src/entry.js index 6d82a11..5b4caf6 100644 --- a/src/entry.js +++ b/src/entry.js @@ -1,10 +1,14 @@ /** * 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. + * This wrapper intercepts attachment upload and download requests for zero-copy streaming + * to/from R2. Workers R2 binding can accept request.body directly for uploads, + * and r2Object.body can be passed directly to Response for downloads. * See: https://blog.cloudflare.com/zh-cn/r2-ga/ * + * This avoids CPU time consumption that would occur if the body went through + * the Rust/WASM layer with axum body conversion. + * * All other requests are passed through to the Rust WASM module. */ @@ -84,6 +88,22 @@ function parseAzureUploadPath(path) { return null; } +// Parse download route: /api/ciphers/{id}/attachment/{attachment_id}/download +function parseDownloadPath(path) { + const parts = path.replace(/^\//, "").split("/"); + // Expected: ["api", "ciphers", "{cipher_id}", "attachment", "{attachment_id}", "download"] + if ( + parts.length === 6 && + parts[0] === "api" && + parts[1] === "ciphers" && + parts[3] === "attachment" && + parts[5] === "download" + ) { + return { cipherId: parts[2], attachmentId: parts[4] }; + } + return null; +} + // Extract token from query string function extractTokenFromQuery(url) { const params = new URL(url).searchParams; @@ -367,6 +387,110 @@ async function handleAzureUpload(request, env, cipherId, attachmentId, token) { return new Response(null, { status: 201 }); } +// Handle download with zero-copy streaming +async function handleDownload(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 + 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" }, + }); + } + + // 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" } } + ); + } + + // Build R2 key + const r2Key = `${cipherId}/${attachmentId}`; + + // Get object from R2 + const r2Object = await bucket.get(r2Key); + if (!r2Object) { + return new Response(JSON.stringify({ error: "Attachment not found in storage" }), { + status: 404, + headers: { "Content-Type": "application/json" }, + }); + } + + // Build response headers + const headers = new Headers(); + const contentType = r2Object.httpMetadata?.contentType || "application/octet-stream"; + headers.set("Content-Type", contentType); + headers.set("Content-Length", r2Object.size.toString()); + + // Return response with R2 object body directly - zero-copy streaming! + // The Workers runtime streams the body directly to the client without consuming CPU time + return new Response(r2Object.body, { + status: 200, + headers, + }); +} + // Main fetch handler export default { async fetch(request, env, ctx) { @@ -389,6 +513,27 @@ export default { } } + // Intercept GET requests to download endpoint for zero-copy streaming + if (request.method === "GET") { + const parsed = parseDownloadPath(url.pathname); + if (parsed) { + const token = extractTokenFromQuery(request.url); + if (!token) { + return new Response( + JSON.stringify({ error: "Missing download token" }), + { status: 401, headers: { "Content-Type": "application/json" } } + ); + } + return handleDownload( + request, + env, + parsed.cipherId, + parsed.attachmentId, + token + ); + } + } + // Pass all other requests to Rust WASM const worker = new RustWorker(ctx, env); return worker.fetch(request); diff --git a/src/handlers/attachments.rs b/src/handlers/attachments.rs index 6379fe3..af20102 100644 --- a/src/handlers/attachments.rs +++ b/src/handlers/attachments.rs @@ -1,21 +1,14 @@ use std::{collections::HashMap, sync::Arc}; -use axum::{ - body::Bytes, - extract::{Multipart, Path, Query, State}, - http::{header, StatusCode}, - response::Response, - Extension, Json, -}; +use axum::{body::Bytes, extract::{Multipart, Path, State}, Extension, Json}; use chrono::Utc; -use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation}; +use jsonwebtoken::{encode, EncodingKey, Header}; use log; use serde::{Deserialize, Serialize}; use serde_json::Value; use uuid::Uuid; use worker::{ - query, wasm_bindgen::JsValue, Bucket, D1Database, Env, Headers as WorkerHeaders, HttpMetadata, - Response as WorkerResponse, + query, wasm_bindgen::JsValue, Bucket, D1Database, Env, HttpMetadata, }; use crate::{ @@ -78,11 +71,6 @@ struct AttachmentDownloadClaims { #[derive(Debug, Deserialize, Default)] #[serde(default)] -pub struct AttachmentDownloadQuery { - pub token: Option, -} - -#[derive(Deserialize)] struct AttachmentKeyRow { cipher_id: String, id: String, @@ -414,64 +402,6 @@ pub async fn delete_attachment_post( .await } -/// GET /api/ciphers/{cipher_id}/attachment/{attachment_id}/download -#[worker::send] -pub async fn download_attachment( - State(env): State>, - Path((cipher_id, attachment_id)): Path<(String, String)>, - Query(query): Query, -) -> Result { - let bucket = require_bucket(&env)?; - let db = db::get_db(&env)?; - - let token = query - .token - .ok_or_else(|| AppError::Unauthorized("Missing download token".to_string()))?; - let user_id = validate_download_token(&env, &token, &cipher_id, &attachment_id)?; - - let cipher = ensure_cipher_for_user(&db, &cipher_id, &user_id).await?; - let attachment = fetch_attachment(&db, &attachment_id).await?; - - if attachment.cipher_id != cipher.id { - return Err(AppError::BadRequest( - "Attachment does not belong to cipher".to_string(), - )); - } - - let object = bucket - .get(attachment.r2_key()) - .execute() - .await - .map_err(AppError::Worker)? - .ok_or_else(|| AppError::NotFound("Attachment not found".to_string()))?; - - let content_length = object.size(); - let http_metadata = object.http_metadata(); - - let response_body = object - .body() - .ok_or_else(|| AppError::NotFound("Attachment data not found".to_string()))? - .response_body() - .map_err(AppError::Worker)?; - - let headers = WorkerHeaders::new(); - let content_type = http_metadata - .content_type - .unwrap_or_else(|| "application/octet-stream".to_string()); - headers - .set(header::CONTENT_TYPE.as_str(), &content_type) - .map_err(AppError::Worker)?; - headers - .set(header::CONTENT_LENGTH.as_str(), &content_length.to_string()) - .map_err(AppError::Worker)?; - - let response = WorkerResponse::from_body(response_body) - .map_err(AppError::Worker)? - .with_status(StatusCode::OK.as_u16()) - .with_headers(headers); - - Ok(response.into()) -} /// Attach attachment information to Cipher (used by other handlers) pub async fn hydrate_cipher_attachments( @@ -838,27 +768,6 @@ fn build_download_token( .map_err(AppError::from) } -fn validate_download_token( - env: &Env, - token: &str, - cipher_id: &str, - attachment_id: &str, -) -> Result { - let secret = jwt_secret(env)?; - let data = decode::( - token, - &DecodingKey::from_secret(secret.as_bytes()), - &Validation::default(), - )?; - - let claims = data.claims; - if claims.cipher_id != cipher_id || claims.attachment_id != attachment_id { - return Err(AppError::Unauthorized("Invalid download token".to_string())); - } - - Ok(claims.sub) -} - fn upload_url( env: &Env, base_url: &str, diff --git a/src/lib.rs b/src/lib.rs index 9e863b0..b8a67b7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -43,13 +43,14 @@ pub async fn main( .allow_headers(Any) .allow_origin(Any); - let body_limit = attachment_body_limit_bytes(&env); + // Attachment uploads/downloads are handled in entry.js for zero-copy streaming, + // so we can use a conservative body limit here (5MB) for regular API requests. + const BODY_LIMIT: usize = 5 * 1024 * 1024; let mut app = router::api_router((*env).clone()) .layer(Extension(BaseUrl(base_url))) .layer(cors) - // axum 默认 body 限制为 2MiB，附件上传需要更大的上限 - .layer(DefaultBodyLimit::max(body_limit)); + .layer(DefaultBodyLimit::max(BODY_LIMIT)); Ok(app.call(req).await?) } @@ -77,28 +78,3 @@ pub async fn scheduled(_event: ScheduledEvent, env: Env, _ctx: ScheduleContext) } } -/// Resolve a permissive body size limit, prioritizing ATTACHMENT_MAX_BYTES and falling back to 64MiB. -fn attachment_body_limit_bytes(env: &Env) -> usize { - const DEFAULT_LIMIT: usize = 64 * 1024 * 1024; - - let max_bytes = env.var("ATTACHMENT_MAX_BYTES").ok().and_then(|v| { - let raw = v.to_string(); - match raw.parse::() { - Ok(val) => Some(val), - Err(err) => { - log::error!("Invalid ATTACHMENT_MAX_BYTES '{}': {}", raw, err); - None - } - } - }); - - let limit_u64 = max_bytes.unwrap_or(DEFAULT_LIMIT as u64); - - limit_u64.try_into().unwrap_or_else(|_| { - log::error!( - "Attachment body limit {} did not fit into usize, falling back to default", - limit_u64 - ); - DEFAULT_LIMIT - }) -} diff --git a/src/router.rs b/src/router.rs index 9ce118c..964d864 100644 --- a/src/router.rs +++ b/src/router.rs @@ -60,8 +60,9 @@ 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 + // Note: Azure upload and download routes are handled in entry.js for zero-copy streaming // PUT /api/ciphers/{id}/attachment/{attachment_id}/azure-upload + // GET /api/ciphers/{id}/attachment/{attachment_id}/download?token=... .route( "/api/ciphers/{id}/attachment", post(attachments::upload_attachment_legacy), @@ -78,10 +79,6 @@ pub fn api_router(env: Env) -> Router { "/api/ciphers/{id}/attachment/{attachment_id}", delete(attachments::delete_attachment), ) - .route( - "/api/ciphers/{id}/attachment/{attachment_id}/download", - get(attachments::download_attachment), - ) .route( "/api/ciphers/{id}/attachment/{attachment_id}/delete", post(attachments::delete_attachment_post),