feat: add zero-copy streaming for attachment downloads using js wapper
This commit is contained in:
qaz741wsd856
parent
8902444561
commit
61a3539e23
4 changed files with 156 additions and 129 deletions
149
src/entry.js
149
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);
|
||||
|
|
|
|||
|
|
@ -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<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,
|
||||
|
|
|
|||
32
src/lib.rs
32
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
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
Loading…
Reference in a new issue