use std::sync::Arc; use axum::extract::{Multipart, State}; use axum::http::HeaderMap; use axum::response::IntoResponse; use axum::routing::post; use axum::{Json, Router}; use bytes::BytesMut; use crate::auth::{self, COOKIE_NAME}; use crate::config; use crate::constants; use crate::database; use crate::error::http_error; use crate::state::AppState; use crate::telegram::service::TelegramService; #[derive(Debug, Default)] struct UploadAuthProgress { auth_verified: bool, } #[derive(Debug, PartialEq, Eq)] enum UploadFieldError { FileBeforeAuth, } fn advance_upload_auth_state( mut state: UploadAuthProgress, prechecked_auth: bool, auth_optional: bool, field_name: &str, _field_value: Option<&str>, ) -> Result { if prechecked_auth || auth_optional { state.auth_verified = true; return Ok(state); } if field_name == "key" { state.auth_verified = true; return Ok(state); } if field_name == "file" && !state.auth_verified { return Err(UploadFieldError::FileBeforeAuth); } Ok(state) } /// Sanitize filename: extract basename, limit length, remove dangerous chars. fn sanitize_filename(raw: &str) -> String { let name = std::path::Path::new(raw) .file_name() .and_then(|n| n.to_str()) .unwrap_or("upload"); let clean: String = name .chars() .filter(|c| !c.is_control() && *c != '\0') .collect(); if clean.is_empty() { return "upload".to_string(); } // UTF-8-safe byte-length cap: `clean[..255]` would panic if byte 255 // falls inside a multibyte character (e.g. a Chinese filename). if clean.len() <= 255 { return clean; } let mut cutoff = 0; for (idx, _) in clean.char_indices() { if idx > 255 { break; } cutoff = idx; } clean[..cutoff].to_string() } async fn upload_file( State(state): State>, headers: HeaderMap, mut multipart: Multipart, ) -> Result { let app_settings = config::get_app_settings(&state.settings, &state.db_pool); let bot_token = app_settings .get("BOT_TOKEN") .and_then(|v| v.as_deref()) .unwrap_or(""); let channel_name = app_settings .get("CHANNEL_NAME") .and_then(|v| v.as_deref()) .unwrap_or(""); if bot_token.is_empty() || channel_name.is_empty() { return Err(http_error( axum::http::StatusCode::SERVICE_UNAVAILABLE, "缺少 BOT_TOKEN 或 CHANNEL_NAME,无法上传", "cfg_missing", )); } // Pre-check auth with header-only info (before consuming body) let has_referer = headers.get("referer").is_some(); let cookie_value = headers .get("cookie") .and_then(|v| v.to_str().ok()) .and_then(|cookies| { cookies.split(';').find_map(|c| { let c = c.trim(); c.strip_prefix(&format!("{}=", COOKIE_NAME)) .map(|v| v.to_string()) }) }); let picgo_key = app_settings.get("PICGO_API_KEY").and_then(|v| v.as_deref()); let pass_word = app_settings.get("PASS_WORD").and_then(|v| v.as_deref()); // Upload auth used to derive a sha256 of the password and feed that as // the expected cookie value. That comparison was always a no-op because // session cookies are random tokens (see `auth::generate_session_token`) // that are independent of the password hash. The auth middleware already // validates the session cookie against the stored SESSION_TOKEN before // the request ever reaches this handler, so we just forward the raw // password presence to `ensure_upload_auth` for the referer / picgo-key // branches. The cookie branch inside `ensure_upload_auth` is reachable // only for header-level requests the middleware already allowed. let pass_word_hash_ref = pass_word; let header_key = headers .get("x-api-key") .and_then(|v| v.to_str().ok()) .map(|s| s.to_string()); let auth_optional = picgo_key.map_or(true, |k| k.is_empty()) && pass_word_hash_ref.map_or(true, |p| p.is_empty()); // Pre-check auth using only HEADER-available credentials: the session // cookie (browser login) and/or x-api-key (PicGo / API clients). These // are the only credentials that exist before we consume the multipart // body. Referer is client-controlled and auth.rs ignores it. // // The browser session cookie is the random SESSION_TOKEN from // app_settings, NOT the password. We verify it directly against the // stored token so a logged-in browser can upload without submitting a // form `key` field. `auth_middleware` uses the same comparison. let session_token_owned = app_settings .get("SESSION_TOKEN") .and_then(|v| v.clone()); let cookie_valid = match (cookie_value.as_deref(), session_token_owned.as_deref()) { (Some(c), Some(t)) if !c.is_empty() && !t.is_empty() => { auth::secure_compare(c, t) } _ => false, }; // ensure_upload_auth handles the x-api-key / PicGo path. We pass None // for the cookie because cookie validity is handled via `cookie_valid` // above — that function still compares cookie to the password hash, // which does not match the random session token used in v2.x. let prechecked_auth = cookie_valid || auth::ensure_upload_auth( has_referer, None, picgo_key, pass_word_hash_ref, header_key.as_deref(), ) .is_ok(); // Parse multipart body - stream file chunks to Telegram let mut form_key: Option = None; let mut upload_result: Option> = None; let mut auth_progress = UploadAuthProgress { auth_verified: prechecked_auth || auth_optional, }; let tg_service = TelegramService::new( bot_token.to_string(), channel_name.to_string(), state.http_client.clone(), ); while let Ok(Some(field)) = multipart.next_field().await { let name = field.name().unwrap_or("").to_string(); if name == "key" { let key_text = field.text().await.ok(); if !auth_progress.auth_verified { if let Err((_, msg, code)) = auth::ensure_upload_auth( has_referer, None, picgo_key, pass_word_hash_ref, key_text.as_deref(), ) { return Err(http_error(axum::http::StatusCode::UNAUTHORIZED, msg, code)); } } auth_progress = advance_upload_auth_state( auth_progress, prechecked_auth, auth_optional, &name, key_text.as_deref(), ) .map_err(|_| { http_error( axum::http::StatusCode::UNAUTHORIZED, "upload auth required before file field", "file_before_auth", ) })?; form_key = key_text; } else if name == "file" { auth_progress = advance_upload_auth_state( auth_progress, prechecked_auth, auth_optional, &name, None, ) .map_err(|_| { http_error( axum::http::StatusCode::UNAUTHORIZED, "upload auth required before file field", "file_before_auth", ) })?; let raw_filename = field.file_name().unwrap_or("upload").to_string(); let filename = sanitize_filename(&raw_filename); // Stream the file in chunks to Telegram upload_result = Some( stream_upload_to_telegram(&tg_service, field, &filename, &state.db_pool).await, ); } } // Final auth check with form-level `key`. Only needed when header-level // credentials (cookie / x-api-key) did not already satisfy auth — e.g. // PicGo clients that authenticate by submitting PICGO_API_KEY in the // multipart body instead of as a header. if !prechecked_auth { let final_key = form_key.as_deref(); if let Err((_, msg, code)) = auth::ensure_upload_auth( has_referer, None, picgo_key, pass_word_hash_ref, final_key, ) { return Err(http_error(axum::http::StatusCode::UNAUTHORIZED, msg, code)); } } let short_id = upload_result .ok_or_else(|| http_error(axum::http::StatusCode::BAD_REQUEST, "未提供文件", "no_file"))? .map_err(|e| { tracing::error!("文件上传失败: {}", e); http_error( axum::http::StatusCode::INTERNAL_SERVER_ERROR, "文件上传失败", "upload_failed", ) })?; let download_path = format!("/d/{}", short_id); Ok(Json(serde_json::json!({ "file_id": short_id, "short_id": short_id, "download_path": download_path, "path": download_path, "url": download_path, }))) } /// Stream file upload: reads multipart field in chunks, uploads each chunk to Telegram /// as it reaches TELEGRAM_CHUNK_SIZE. Peak memory is ~1 chunk (~20MB) instead of the full file. async fn stream_upload_to_telegram( tg_service: &TelegramService, mut field: axum::extract::multipart::Field<'_>, filename: &str, db_pool: &database::DbPool, ) -> Result { let chunk_size = constants::TELEGRAM_CHUNK_SIZE; let mut buffer = BytesMut::with_capacity(chunk_size); let mut total_size: usize = 0; let mut chunk_ids: Vec = Vec::new(); let mut first_message_id: Option = None; let mut chunk_num: u32 = 0; // Read field data incrementally while let Ok(Some(bytes)) = field.chunk().await { buffer.extend_from_slice(&bytes); total_size += bytes.len(); // When buffer reaches chunk size, send it while buffer.len() >= chunk_size { chunk_num += 1; let chunk_data = buffer.split_to(chunk_size).freeze().to_vec(); let chunk_name = format!("{}.part{}", filename, chunk_num); let message = tg_service .send_document_raw(chunk_data, &chunk_name, first_message_id) .await?; if first_message_id.is_none() { first_message_id = Some(message.message_id); } let doc = message.document.ok_or("No document in chunk response")?; chunk_ids.push(format!("{}:{}", message.message_id, doc.file_id)); } } // Handle remaining data in buffer if buffer.is_empty() && chunk_ids.is_empty() { return Err("文件为空".into()); } if chunk_ids.is_empty() { // Small file: single upload (no chunks were sent yet) tracing::info!("直接上传文件: {} ({}字节)", filename, total_size); let data = buffer.freeze().to_vec(); let message = tg_service.send_document_raw(data, filename, None).await?; let doc = message.document.ok_or("No document in response")?; let composite_id = format!("{}:{}", message.message_id, doc.file_id); let short_id = database::add_file_metadata(db_pool, filename, &composite_id, total_size as i64) .map_err(|e| e.to_string())?; return Ok(short_id); } // Send remaining buffer as last chunk if !buffer.is_empty() { chunk_num += 1; let chunk_data = buffer.freeze().to_vec(); let chunk_name = format!("{}.part{}", filename, chunk_num); let message = tg_service .send_document_raw(chunk_data, &chunk_name, first_message_id) .await?; let doc = message .document .ok_or("No document in last chunk response")?; chunk_ids.push(format!("{}:{}", message.message_id, doc.file_id)); } // Multi-chunk: create and upload manifest tracing::info!( "分块上传完成: {} ({}MB, {} 块)", filename, total_size / (1024 * 1024), chunk_ids.len() ); let mut manifest = String::from("tgstate-blob\n"); manifest.push_str(filename); manifest.push('\n'); for cid in &chunk_ids { manifest.push_str(cid); manifest.push('\n'); } let manifest_name = format!("{}.manifest", filename); let message = tg_service .send_document_raw(manifest.into_bytes(), &manifest_name, first_message_id) .await?; let doc = message.document.ok_or("No document in manifest response")?; let manifest_composite = format!("{}:{}", message.message_id, doc.file_id); let short_id = database::add_file_metadata(db_pool, filename, &manifest_composite, total_size as i64) .map_err(|e| e.to_string())?; Ok(short_id) } pub fn router() -> Router> { Router::new().route("/api/upload", post(upload_file)) } #[cfg(test)] mod tests { use super::{advance_upload_auth_state, UploadAuthProgress, UploadFieldError}; use crate::config::Settings; use crate::database; use crate::state::AppState; use axum::body::{to_bytes, Body}; use axum::http::{header, Request, StatusCode}; use axum::Router; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; use tower::util::ServiceExt; fn test_state() -> Arc { let unique = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_nanos(); let data_dir = std::env::temp_dir() .join(format!("tgstate-upload-test-{}", unique)) .to_string_lossy() .to_string(); let settings = Settings { bot_token: Some("123456:test-token".into()), channel_name: Some("@test_channel".into()), pass_word: Some("secret".into()), picgo_api_key: None, base_url: "http://127.0.0.1:8000".into(), _mode: "p".into(), _file_route: "/d/".into(), data_dir: data_dir.clone(), }; let db_pool = database::init_db(&data_dir); let tera = tera::Tera::default(); let http_client = reqwest::Client::new(); let app_settings = crate::config::get_app_settings(&settings, &db_pool); Arc::new(AppState::new( settings, tera, http_client, db_pool, app_settings, true, )) } fn multipart_request_with_file_before_key() -> Request { let boundary = "X-BOUNDARY"; let body = format!( "--{b}\r\nContent-Disposition: form-data; name=\"file\"; filename=\"test.txt\"\r\nContent-Type: text/plain\r\n\r\nhello\r\n--{b}\r\nContent-Disposition: form-data; name=\"key\"\r\n\r\nsecret\r\n--{b}--\r\n", b = boundary ); Request::builder() .method("POST") .uri("/api/upload") .header( header::CONTENT_TYPE, format!("multipart/form-data; boundary={}", boundary), ) .body(Body::from(body)) .unwrap() } #[test] fn upload_requires_key_before_file_for_api_requests() { let state = UploadAuthProgress::default(); let result = advance_upload_auth_state(state, false, false, "file", None); assert!(matches!(result, Err(UploadFieldError::FileBeforeAuth))); } #[test] fn upload_accepts_key_before_file_for_api_requests() { let state = UploadAuthProgress::default(); let state = advance_upload_auth_state(state, false, false, "key", Some("secret")).unwrap(); let state = advance_upload_auth_state(state, false, false, "file", None).unwrap(); assert!(state.auth_verified); } #[tokio::test] async fn upload_route_rejects_file_field_before_auth() { let state = test_state(); let app = Router::new() .merge(super::router()) .with_state(state.clone()); let response = app .oneshot(multipart_request_with_file_before_key()) .await .unwrap(); assert_eq!(response.status(), StatusCode::UNAUTHORIZED); let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); let text = String::from_utf8(body.to_vec()).unwrap(); assert!(text.contains("file_before_auth"), "unexpected body: {}", text); let files = database::get_all_files(&state.db_pool).unwrap(); assert!(files.is_empty(), "unexpected files persisted: {:?}", files); } }