juiceback/routes/
upload.rs

1//! Handles multipart file uploads. Streams non-gzip data straight to juicehost.
2//! Gzip uploads get buffered, decompressed, then sent.
3
4use axum::{
5    extract::{ConnectInfo, Multipart, State},
6    http::HeaderMap,
7    response::{IntoResponse, Response},
8    Json,
9};
10use bytes::Bytes;
11use serde::Serialize;
12use std::net::SocketAddr;
13use std::sync::Arc;
14use tokio::sync::mpsc;
15use utoipa::ToSchema;
16
17use crate::db::{self, FileRecord};
18use crate::error::AppError;
19use crate::routes::noscript;
20use crate::state::AppState;
21
22/// JSON response returned by a successful upload.
23#[derive(Serialize, ToSchema)]
24pub struct UploadResponse {
25    pub id: String,
26    pub url: String,
27    pub filename: String,
28    pub size_bytes: i64,
29    pub mime_type: String,
30    pub expires_at: i64,
31    pub delete_token: String,
32}
33
34pub(crate) fn sanitize_filename(name: &str) -> String {
35    let name = name.trim();
36    if name.is_empty() {
37        return "upload".to_string();
38    }
39    let mut result = String::with_capacity(name.len().min(255));
40    for c in name.chars() {
41        if c.is_control() || c == '/' || c == '\\' || c == '\0' {
42            continue;
43        }
44        result.push(c);
45        if result.len() > 255 {
46            result.pop();
47            break;
48        }
49    }
50    let result = result.replace("..", "");
51    if result.is_empty() {
52        return "upload".to_string();
53    }
54    result
55}
56
57#[utoipa::path(
58    post,
59    path = "/upload",
60    request_body(content_type = "multipart/form-data", description = "File upload with optional metadata fields (host, upload_mode, ttl_hours)"),
61    responses(
62        (status = 200, description = "File uploaded successfully", body = UploadResponse),
63        (status = 413, description = "File too large"),
64        (status = 429, description = "Rate limit exceeded"),
65    ),
66    tag = "Uploads",
67)]
68#[tracing::instrument(skip_all)]
69pub async fn upload_handler(
70    State(state): State<Arc<AppState>>,
71    ConnectInfo(addr): ConnectInfo<SocketAddr>,
72    headers: HeaderMap,
73    mut multipart: Multipart,
74) -> Result<Response, AppError> {
75    let wants_html = noscript::wants_html(&headers);
76    let is_gzip = headers
77        .get("x-file-encoding")
78        .and_then(|v| v.to_str().ok())
79        == Some("gzip");
80    let mut ttl_seconds = state.config.default_ttl_hours as i64 * 3600;
81
82    let mut file_id = String::new();
83    let mut file_data: Option<Vec<u8>> = None;
84    let mut sanitized_filename = String::new();
85    let mut mime_type = String::new();
86    let mut total_bytes: i64 = 0;
87    let mut selected_host = String::new();
88    let mut selected_upload_mode = String::from("standard");
89
90    let parse_start = std::time::Instant::now();
91    while let Some(mut field) = multipart
92        .next_field()
93        .await
94        .map_err(|e| AppError::InvalidMultipart(e.to_string()))?
95    {
96        let name = field.name().unwrap_or("").to_string();
97        let field_start = std::time::Instant::now();
98
99        if name == "ttl_hours" {
100            let text = field
101                .text()
102                .await
103                .map_err(|e| AppError::InvalidMultipart(e.to_string()))?;
104            if let Ok(val) = text.parse::<f64>() {
105                let max_secs = state.config.max_ttl_hours as i64 * 3600;
106                let min_secs = 3600i64;
107                ttl_seconds = (val * 3600.0).round() as i64;
108                ttl_seconds = ttl_seconds.clamp(min_secs, max_secs);
109            }
110            tracing::debug!("  field ttl_hours took {:?}", field_start.elapsed());
111        } else if name == "upload_mode" {
112            let text = field
113                .text()
114                .await
115                .map_err(|e| AppError::InvalidMultipart(e.to_string()))?;
116            let text = text.trim().to_string();
117            if text == "quic" || text == "standard" {
118                selected_upload_mode = text;
119            }
120            tracing::debug!("  field upload_mode took {:?}", field_start.elapsed());
121        } else if name == "host" {
122            let text = field
123                .text()
124                .await
125                .map_err(|e| AppError::InvalidMultipart(e.to_string()))?;
126            let text = text.trim().trim_end_matches('/').to_string();
127            if !text.is_empty() {
128                selected_host = text;
129            }
130            tracing::debug!("  field host took {:?}", field_start.elapsed());
131        } else if name == "file" {
132            let filename = field.file_name().unwrap_or("upload").to_string();
133            sanitized_filename = sanitize_filename(&filename);
134
135            mime_type = field
136                .content_type()
137                .map(|m| m.to_string())
138                .unwrap_or_else(|| {
139                    mime_guess::from_path(&sanitized_filename)
140                        .first_or_octet_stream()
141                        .to_string()
142                });
143
144            let id = nanoid::nanoid!(8);
145            file_id = id;
146
147            if is_gzip {
148                // Gzip: must buffer to decompress, streaming not possible
149                let read_start = std::time::Instant::now();
150                let mut raw = Vec::new();
151                while let Some(chunk) = field
152                    .chunk()
153                    .await
154                    .map_err(|e| AppError::InvalidMultipart(e.to_string()))?
155                {
156                    raw.extend_from_slice(&chunk);
157                    if raw.len() as i64 > state.config.max_file_size_bytes as i64 {
158                        return Err(AppError::PayloadTooLarge);
159                    }
160                }
161                let read_time = read_start.elapsed();
162                let decomp_start = std::time::Instant::now();
163                use async_compression::tokio::bufread::GzipDecoder;
164                use tokio::io::AsyncReadExt;
165                let mut decoder = GzipDecoder::new(raw.as_slice());
166                let mut decompressed = Vec::new();
167                decoder
168                    .read_to_end(&mut decompressed)
169                    .await
170                    .map_err(|_| AppError::GzipDecodeFailed)?;
171                let decomp_time = decomp_start.elapsed();
172                total_bytes = decompressed.len() as i64;
173                if total_bytes > state.config.max_file_size_bytes as i64 {
174                    return Err(AppError::PayloadTooLarge);
175                }
176                file_data = Some(decompressed);
177                let gzip_bps = if (read_time + decomp_time).as_secs_f64() > 0.0 {
178                    total_bytes as f64 / (read_time + decomp_time).as_secs_f64()
179                } else { 0.0 };
180                tracing::debug!(
181                    "  field file (gzip): read={:?} decompress={:?} bytes={} {:.2} MB/s",
182                    read_time, decomp_time, total_bytes,
183                    gzip_bps / (1024.0 * 1024.0),
184                );
185            } else {
186                // Non-gzip: stream directly to juicehost as chunks arrive
187                let (tx, rx) = mpsc::channel::<Result<Bytes, String>>(64);
188
189                let host_opt = if selected_host.is_empty() { None } else { Some(selected_host.clone()) };
190                let state_push = Arc::clone(&state);
191                let push_id = file_id.clone();
192                let push_filename = sanitized_filename.clone();
193                let push_mime = mime_type.clone();
194
195                let push_mode = selected_upload_mode.clone();
196                let push_handle = tokio::spawn(async move {
197                    crate::juicehost::push_file_streaming(
198                        &state_push,
199                        &push_id,
200                        &push_filename,
201                        &push_mime,
202                        rx,
203                        host_opt,
204                        &push_mode,
205                    )
206                    .await
207                });
208
209                let read_start = std::time::Instant::now();
210                let mut chunk_count: u64 = 0;
211                let mut push_backpressure: std::time::Duration = Default::default();
212                while let Some(chunk) = field
213                    .chunk()
214                    .await
215                    .map_err(|e| AppError::InvalidMultipart(e.to_string()))?
216                {
217                    total_bytes += chunk.len() as i64;
218                    if total_bytes > state.config.max_file_size_bytes as i64 {
219                        return Err(AppError::PayloadTooLarge);
220                    }
221                    chunk_count += 1;
222                    let send_start = std::time::Instant::now();
223                    tx.send(Ok(chunk)).await
224                        .map_err(|_| AppError::JuicehostRejected("streaming channel closed".into()))?;
225                    push_backpressure += send_start.elapsed();
226                }
227                drop(tx);
228                let read_time = read_start.elapsed();
229                let read_bps = if read_time.as_secs_f64() > 0.0 {
230                    total_bytes as f64 / read_time.as_secs_f64()
231                } else { 0.0 };
232
233                // Wait for push to finish
234                let stream_start = std::time::Instant::now();
235                push_handle.await
236                    .map_err(|_| AppError::TaskPanicked("streaming push task panicked".into()))?
237                    .map_err(|e| {
238                        if e.contains("error trying to connect") || e.contains("dns error") || e.contains("connection refused") {
239                            AppError::JuicehostUnreachable(e)
240                        } else if e.contains("[INSUFFICIENT_STORAGE]") {
241                            AppError::InsufficientStorage(e)
242                        } else {
243                            AppError::JuicehostRejected(e)
244                        }
245                    })?;
246                let push_wait = stream_start.elapsed();
247
248                tracing::debug!(
249                    "  stream: id={} bytes={} chunks={} read={:?} ({:.2} MB/s) \
250                     backpressure={:?} push_wait={:?}",
251                    file_id, total_bytes, chunk_count,
252                    read_time, read_bps / (1024.0 * 1024.0),
253                    push_backpressure, push_wait,
254                );
255            }
256            break;
257        }
258    }
259
260    tracing::debug!("multipart parse total took {:?}", parse_start.elapsed());
261
262    let now = chrono::Utc::now().timestamp();
263    let record = FileRecord {
264        id: file_id.clone(),
265        filename: sanitized_filename.clone(),
266        mime_type: mime_type.clone(),
267        size_bytes: total_bytes,
268        storage_path: format!("remote-{}", file_id),
269        delete_token: uuid::Uuid::new_v4().to_string(),
270        uploaded_at: now,
271        expires_at: now + ttl_seconds,
272        uploader_ip: Some(addr.ip().to_string()),
273        storage_host: selected_host.clone(),
274    };
275
276    // For gzip uploads, the data was already buffered, push now.
277    // For streaming uploads, the push already happened concurrently.
278    if let Some(data) = file_data {
279        let host_opt = if selected_host.is_empty() { None } else { Some(selected_host.as_str()) };
280        let push_start = std::time::Instant::now();
281        crate::juicehost::push_file_to_juicehost(
282            &state,
283            &file_id,
284            &sanitized_filename,
285            &mime_type,
286            &data,
287            host_opt,
288        )
289        .await
290        .map_err(|e| {
291            if e.contains("error trying to connect") || e.contains("dns error") || e.contains("connection refused") {
292                AppError::JuicehostUnreachable(e)
293            } else if e.contains("[INSUFFICIENT_STORAGE]") {
294                AppError::InsufficientStorage(e)
295            } else {
296                AppError::JuicehostRejected(e)
297            }
298        })?;
299        tracing::debug!("push_to_juicehost (buffered) took {:?}", push_start.elapsed());
300    }
301
302    // Parallel DB insert, runs concurrently with the push for streaming path
303    let state2 = Arc::clone(&state);
304    let record2 = record.clone();
305
306    let db_start = std::time::Instant::now();
307    tokio::task::spawn_blocking(move || -> Result<(), AppError> {
308        let db = state2.db.get().map_err(|e| {
309            AppError::DbPoolError(e.to_string())
310        })?;
311        db::insert_file(&db, &record2).map_err(AppError::DatabaseError)
312    })
313    .await
314    .map_err(|_| AppError::TaskPanicked("db insert task panicked".into()))?
315    .map_err(|e| e)?;
316    tracing::debug!("db insert took {:?}", db_start.elapsed());
317
318    tracing::info!(
319        "upload: id={} size={} mime={} expires_at={} ip={}",
320        record.id,
321        record.size_bytes,
322        record.mime_type,
323        record.expires_at,
324        addr.ip()
325    );
326
327    let public_url = if record.storage_host.is_empty() {
328        format!("{}/f/{}", state.config.public_base_url, record.id)
329    } else {
330        format!("{}/f/{}", record.storage_host, record.id)
331    };
332
333    if wants_html {
334        return Ok(noscript::upload_confirmation_html(
335            &public_url,
336            &record.filename,
337            &record.delete_token,
338            &record.id,
339        ));
340    }
341
342    Ok(Json(UploadResponse {
343        url: public_url,
344        id: record.id,
345        filename: record.filename,
346        size_bytes: record.size_bytes,
347        mime_type: record.mime_type,
348        expires_at: record.expires_at,
349        delete_token: record.delete_token,
350    })
351    .into_response())
352}