feat: add zero-copy streaming for attachment downloads using js wapper

src/entry.js

@@ -1,10 +1,14 @@
 /**
  * JS Wrapper Entry Point for Warden Worker
  *
- * This wrapper intercepts attachment upload requests for zero-copy streaming
+ * This wrapper intercepts attachment upload and download requests for zero-copy streaming
- * uploads to R2. Workers R2 binding can accept request.body directly.
+ * 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);

src/handlers/attachments.rs

@@ -1,21 +1,14 @@
 use std::{collections::HashMap, sync::Arc};
 
-use axum::{
+use axum::{body::Bytes, extract::{Multipart, Path, State}, Extension, Json};
-    body::Bytes,
-    extract::{Multipart, Path, Query, State},
-    http::{header, StatusCode},
-    response::Response,
-    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,
+    query, wasm_bindgen::JsValue, Bucket, D1Database, Env, HttpMetadata,
-    Response as WorkerResponse,
 };
 
 use crate::{
@@ -78,11 +71,6 @@ struct AttachmentDownloadClaims {
 
 #[derive(Debug, Deserialize, Default)]
 #[serde(default)]
-pub struct AttachmentDownloadQuery {
-    pub token: Option<String>,
-}
-
-#[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<Arc<Env>>,
-    Path((cipher_id, attachment_id)): Path<(String, String)>,
-    Query(query): Query<AttachmentDownloadQuery>,
-) -> Result<Response, AppError> {
-    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<String, AppError> {
-    let secret = jwt_secret(env)?;
-    let data = decode::<AttachmentDownloadClaims>(
-        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,

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::<u64>() {
-            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
-    })
-}

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),