juiceback/routes/
tus.rs

1use axum::{
2    body::Bytes,
3    extract::{ConnectInfo, Path, State},
4    http::{HeaderMap, StatusCode},
5    response::{IntoResponse, Response},
6    routing, Json,
7};
8use base64::Engine;
9use std::net::SocketAddr;
10use std::sync::Arc;
11use tokio::io::AsyncWriteExt;
12
13use crate::db::{self, FileRecord};
14use crate::error::AppError;
15use crate::state::AppState;
16use crate::tus::TusUpload;
17
18fn parse_tus_metadata(header: Option<&axum::http::HeaderValue>) -> Vec<(String, String)> {
19    let raw = match header.and_then(|v| v.to_str().ok()) {
20        Some(v) => v,
21        None => return vec![],
22    };
23    let mut pairs = Vec::new();
24    for pair in raw.split(',') {
25        let pair = pair.trim();
26        if let Some(eq) = pair.find(' ') {
27            let key = pair[..eq].to_string();
28            let value_b64 = pair[eq + 1..].trim();
29            if let Ok(decoded) = base64::engine::general_purpose::STANDARD
30                .decode(value_b64.as_bytes())
31            {
32                if let Ok(value) = String::from_utf8(decoded) {
33                    pairs.push((key, value));
34                }
35            }
36        }
37    }
38    pairs
39}
40
41#[tracing::instrument(skip_all)]
42pub async fn create_upload_handler(
43    State(state): State<Arc<AppState>>,
44    ConnectInfo(addr): ConnectInfo<SocketAddr>,
45    headers: HeaderMap,
46) -> Result<Response, AppError> {
47    let total_length = headers
48        .get("upload-length")
49        .and_then(|v| v.to_str().ok())
50        .and_then(|v| v.parse::<u64>().ok())
51        .ok_or_else(|| AppError::TusMissingLength)?;
52
53    if total_length > state.config.max_file_size_bytes {
54        return Err(AppError::PayloadTooLarge);
55    }
56
57    let metadata = parse_tus_metadata(headers.get("upload-metadata"));
58    let filename = metadata
59        .iter()
60        .find(|(k, _)| k == "filename")
61        .map(|(_, v)| v.clone())
62        .unwrap_or_else(|| "upload".to_string());
63    let mime_type = metadata
64        .iter()
65        .find(|(k, _)| k == "mimetype")
66        .map(|(_, v)| v.clone())
67        .unwrap_or_else(|| "application/octet-stream".to_string());
68    let ttl_hours = metadata
69        .iter()
70        .find(|(k, _)| k == "ttl")
71        .and_then(|(_, v)| v.parse::<f64>().ok())
72        .unwrap_or(state.config.default_ttl_hours as f64)
73        .clamp(1.0, state.config.max_ttl_hours as f64);
74    let storage_host = metadata
75        .iter()
76        .find(|(k, _)| k == "host")
77        .map(|(_, v)| v.clone())
78        .filter(|v| !v.is_empty())
79        .unwrap_or_default();
80    let storage_upload_mode = metadata
81        .iter()
82        .find(|(k, _)| k == "upload_mode")
83        .map(|(_, v)| v.clone())
84        .filter(|v| v == "quic")
85        .unwrap_or_else(|| "standard".to_string());
86
87    let id = nanoid::nanoid!(8);
88
89    // Stage TUS uploads in the system temp directory, never in the permanent files_dir.
90    // The file is pushed to juicehost on completion and the temp file is cleaned up.
91    let storage_path = std::env::temp_dir().join(&id);
92
93    tokio::fs::File::create(&storage_path)
94        .await
95        .map_err(|_| AppError::TusStagingError)?;
96
97    let upload = TusUpload {
98        id: id.clone(),
99        storage_path,
100        offset: 0,
101        total_length,
102        filename,
103        mime_type,
104        ttl_hours,
105        created_at: chrono::Utc::now().timestamp(),
106        delete_token: uuid::Uuid::new_v4().to_string(),
107        uploader_ip: addr.ip().to_string(),
108        storage_host,
109        storage_upload_mode,
110    };
111
112    state.tus.insert(id.clone(), upload);
113
114    let location = format!("/api/tus/{}", id);
115    Ok(Response::builder()
116        .status(StatusCode::CREATED)
117        .header("Location", &location)
118        .header("Tus-Resumable", "1.0.0")
119        .body(axum::body::Body::empty())
120        .unwrap())
121}
122
123pub async fn get_upload_handler(
124    State(state): State<Arc<AppState>>,
125    Path(id): Path<String>,
126) -> Result<Response, AppError> {
127    let upload = state
128        .tus
129        .get(&id)
130        .ok_or(AppError::TusSessionNotFound)?;
131
132    Ok(Response::builder()
133        .status(StatusCode::OK)
134        .header("Upload-Offset", upload.offset.to_string())
135        .header("Upload-Length", upload.total_length.to_string())
136        .header("Tus-Resumable", "1.0.0")
137        .body(axum::body::Body::empty())
138        .unwrap())
139}
140
141#[tracing::instrument(skip_all)]
142pub async fn patch_upload_handler(
143    State(state): State<Arc<AppState>>,
144    Path(id): Path<String>,
145    headers: HeaderMap,
146    body: Bytes,
147) -> Result<Response, AppError> {
148    let upload_offset = headers
149        .get("upload-offset")
150        .and_then(|v| v.to_str().ok())
151        .and_then(|v| v.parse::<u64>().ok())
152        .ok_or_else(|| AppError::TusMissingOffset)?;
153
154    let snapshot = {
155        let upload = state
156            .tus
157            .get(&id)
158            .ok_or(AppError::TusSessionNotFound)?;
159        if upload_offset != upload.offset {
160            return Err(AppError::TusOffsetMismatch);
161        }
162        (upload.storage_path.clone(), upload.offset)
163    };
164
165    let chunk_len = body.len() as u64;
166    let write_start = std::time::Instant::now();
167    {
168        let mut file = tokio::fs::OpenOptions::new()
169            .append(true)
170            .open(&snapshot.0)
171            .await
172            .map_err(AppError::FilesystemError)?;
173        file.write_all(&body)
174            .await
175            .map_err(AppError::FilesystemError)?;
176        file.flush().await.map_err(AppError::FilesystemError)?;
177    }
178    tracing::debug!("tus patch: chunk_write took {:?} ({} bytes)", write_start.elapsed(), chunk_len);
179
180    let new_offset = snapshot.1 + chunk_len;
181    let (is_complete, completed_snapshot) = {
182        let mut upload = state
183            .tus
184            .get_mut(&id)
185            .ok_or(AppError::TusSessionNotFound)?;
186        upload.offset = new_offset;
187        let done = upload.offset >= upload.total_length;
188        if done {
189            (true, Some((
190                upload.id.clone(),
191                upload.filename.clone(),
192                upload.mime_type.clone(),
193                upload.storage_path.clone(),
194                upload.total_length,
195                upload.delete_token.clone(),
196                upload.uploader_ip.clone(),
197                upload.ttl_hours,
198                upload.storage_host.clone(),
199                upload.storage_upload_mode.clone(),
200            )))
201        } else {
202            (false, None)
203        }
204    };
205
206    if is_complete {
207        if let Some((uid, ufilename, umime, upath, usize, utoken, uip, uttl, ustorage_host, ustorage_upload_mode)) = completed_snapshot {
208            state.tus.remove(&id);
209
210            let host_opt = if ustorage_host.is_empty() { None } else { Some(ustorage_host.clone()) };
211            let mode = ustorage_upload_mode.as_str();
212
213            // Push to juicehost before creating DB record
214            let push_start = std::time::Instant::now();
215            if mode == "quic" {
216                let file_data = tokio::fs::read(&upath).await?;
217                let (tx, rx) = tokio::sync::mpsc::channel::<Result<bytes::Bytes, String>>(1);
218                tx.send(Ok(bytes::Bytes::from(file_data))).await.unwrap();
219                drop(tx);
220                crate::juicehost::push_file_streaming(
221                    &state,
222                    &uid,
223                    &ufilename,
224                    &umime,
225                    rx,
226                    host_opt,
227                    mode,
228                )
229                .await
230            } else {
231                crate::juicehost::push_file_path_to_juicehost(
232                    &state,
233                    &uid,
234                    &upath,
235                    &ufilename,
236                    &umime,
237                    host_opt.as_deref(),
238                )
239                .await
240            }
241            .map_err(|e| {
242                if e.contains("error trying to connect") || e.contains("dns error") || e.contains("connection refused") {
243                    AppError::JuicehostUnreachable(e)
244                } else if e.contains("[INSUFFICIENT_STORAGE]") {
245                    AppError::InsufficientStorage(e)
246                } else {
247                    AppError::JuicehostRejected(e)
248                }
249            })?;
250            tracing::debug!("tus complete: push_to_juicehost took {:?}", push_start.elapsed());
251
252            let now = chrono::Utc::now().timestamp();
253            let ttl_seconds = (uttl * 3600.0).round() as i64;
254
255            let record = FileRecord {
256                id: uid.clone(),
257                filename: ufilename,
258                mime_type: umime,
259                size_bytes: usize as i64,
260                storage_path: format!("remote-{}", uid),
261                delete_token: utoken,
262                uploaded_at: now,
263                expires_at: now + ttl_seconds,
264                uploader_ip: Some(uip),
265                storage_host: ustorage_host,
266            };
267
268            let db_start = std::time::Instant::now();
269            let state2 = Arc::clone(&state);
270            let record2 = record.clone();
271            tokio::task::spawn_blocking(move || -> Result<(), AppError> {
272                let db = state2
273                    .db
274                    .get()
275                    .map_err(|e| AppError::DbPoolError(e.to_string()))?;
276                db::insert_file(&db, &record2).map_err(AppError::DatabaseError)
277            })
278            .await
279            .map_err(|_| AppError::TaskPanicked("db insert task panicked".into()))?
280            .map_err(|e| e)?;
281            tracing::debug!("tus complete: db insert took {:?}", db_start.elapsed());
282
283            let public_url = if record.storage_host.is_empty() {
284                format!("{}/f/{}", state.config.public_base_url, uid)
285            } else {
286                format!("{}/f/{}", record.storage_host, uid)
287            };
288
289            return Ok(Json(serde_json::json!({
290                "id": record.id,
291                "url": public_url,
292                "filename": record.filename,
293                "size_bytes": record.size_bytes,
294                "mime_type": record.mime_type,
295                "expires_at": record.expires_at,
296                "delete_token": record.delete_token,
297            }))
298            .into_response());
299        }
300    }
301
302    Ok(Response::builder()
303        .status(StatusCode::NO_CONTENT)
304        .header("Upload-Offset", new_offset.to_string())
305        .header("Tus-Resumable", "1.0.0")
306        .body(axum::body::Body::empty())
307        .unwrap())
308}
309
310pub async fn options_handler() -> Response {
311    Response::builder()
312        .status(StatusCode::NO_CONTENT)
313        .header("Tus-Resumable", "1.0.0")
314        .header("Tus-Version", "1.0.0")
315        .header(
316            "Tus-Extension",
317            "creation,termination",
318        )
319        .body(axum::body::Body::empty())
320        .unwrap()
321}
322
323pub async fn delete_upload_handler(
324    State(state): State<Arc<AppState>>,
325    Path(id): Path<String>,
326) -> Result<Response, AppError> {
327    let upload = state.tus.remove(&id).ok_or(AppError::TusSessionNotFound)?.1;
328    let _ = tokio::fs::remove_file(&upload.storage_path).await;
329    Ok(Response::builder()
330        .status(StatusCode::NO_CONTENT)
331        .header("Tus-Resumable", "1.0.0")
332        .body(axum::body::Body::empty())
333        .unwrap())
334}
335
336pub fn tus_routes() -> axum::Router<Arc<AppState>> {
337    axum::Router::new()
338        .route("/api/tus", routing::post(create_upload_handler).options(options_handler))
339        .route(
340            "/api/tus/:id",
341            routing::get(get_upload_handler)
342                .patch(patch_upload_handler)
343                .delete(delete_upload_handler),
344        )
345}