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

@@ -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",

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"] }

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);
+  },
+};
+

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<String>,
-        Option<String>,
-        Option<String>,
-    ),
-    AppError,
-> {
-    let mut file_bytes: Option<axum::body::Bytes> = None;
+) -> Result<(Bytes, Option<String>, Option<String>, Option<String>), AppError> {
+    let mut file_bytes: Option<Bytes> = None;
     let mut content_type: Option<String> = None;
     let mut key: Option<String> = None;
     let mut file_name: Option<String> = 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<String, AppError> {
+    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<String, AppError> {
     Ok(env.secret("JWT_SECRET")?.to_string())
 }

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，附件上传需要更大的上限

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

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