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