rustango 0.60.7

//! Multipart file uploads. Joins axum's multipart extractor to the
//! [`crate::storage::Storage`] trait, and handles size limits,
//! extension checks and filename sanitizing for you.
//!
//! ## Security
//!
//! Set `max_bytes` and `allowed_extensions` on every upload route. With
//! no extension list any type is accepted except those a browser runs
//! script in ([`ACTIVE_EXTENSIONS`]); list one to accept it. The
//! client's `Content-Type` is never trusted: it is stored for reference
//! only, so sniff or re-encode the bytes yourself if the type matters.
//! Filenames are stripped to a basename, so a client-sent path such as
//! `../../etc/passwd` cannot escape the prefix.
//!
//! ## Quick start
//!
//! ```ignore
//! use rustango::uploads::{UploadConfig, save_uploads};
//! use rustango::storage::BoxedStorage;
//! use axum::extract::{Multipart, State};
//! use axum::Json;
//!
//! async fn upload(
//!     State(store): State<BoxedStorage>,
//!     mp: Multipart,
//! ) -> Json<Vec<String>> {
//!     let cfg = UploadConfig::new("uploads/")
//!         .max_bytes(5 * 1024 * 1024)
//!         .allowed_extensions(&["png", "jpg", "jpeg", "webp"]);
//!     let saved = save_uploads(mp, &cfg, &store).await.unwrap();
//!     Json(saved.into_iter().map(|s| s.key).collect())
//! }
//! ```

use std::collections::HashSet;
use std::path::Path;

use axum::extract::Multipart;

use crate::storage::{BoxedStorage, StorageError};

#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum UploadError {
    #[error("multipart parse error: {0}")]
    Parse(String),
    #[error("file too large: {actual} bytes (max {max})")]
    TooLarge { actual: usize, max: usize },
    #[error("file extension `{0}` not allowed")]
    BadExtension(String),
    #[error("filename missing in multipart field")]
    MissingFilename,
    #[error("too many files (max {max})")]
    TooManyFiles { max: usize },
    #[error("storage error: {0}")]
    Storage(#[from] StorageError),
}

impl From<axum::extract::multipart::MultipartError> for UploadError {
    fn from(e: axum::extract::multipart::MultipartError) -> Self {
        Self::Parse(e.to_string())
    }
}

/// Extensions refused when `allowed_extensions` is empty: served from
/// the app origin, they run script with the viewer's session (#1849).
pub const ACTIVE_EXTENSIONS: &[&str] = &[
    "html", "htm", "xhtml", "xht", "shtml", "svg", "svgz", "xml", "xsl", "xslt", "js", "mjs",
];

/// Per-upload configuration. Cheap to clone.
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct UploadConfig {
    /// Storage prefix (e.g. `"avatars/"` or `"uploads/2026/05/"`).
    /// The original filename is appended after sanitization.
    pub prefix: String,
    /// Hard cap on a single file's size. Default: 10 MiB.
    pub max_bytes: usize,
    /// Extensions to accept (lowercase, no leading dot). An empty set
    /// accepts anything but [`ACTIVE_EXTENSIONS`].
    pub allowed_extensions: HashSet<String>,
    /// Most files one request may carry. Default: 20.
    pub max_files: usize,
    /// How many non-file fields to skip while iterating.
    pub skip_fields: usize,
    /// Prepend a timestamp to the key so two uploads with the same name
    /// do not overwrite each other. Default: true.
    pub randomize_filename: bool,
}

impl UploadConfig {
    pub fn new(prefix: impl Into<String>) -> Self {
        Self {
            prefix: ensure_trailing_slash(prefix.into()),
            max_bytes: 10 * 1024 * 1024,
            allowed_extensions: HashSet::new(),
            max_files: 20,
            skip_fields: 0,
            randomize_filename: true,
        }
    }

    #[must_use]
    pub fn max_bytes(mut self, n: usize) -> Self {
        self.max_bytes = n;
        self
    }

    #[must_use]
    pub fn allowed_extensions(mut self, exts: &[&str]) -> Self {
        self.allowed_extensions = exts.iter().map(|e| e.to_ascii_lowercase()).collect();
        self
    }

    #[must_use]
    pub fn max_files(mut self, n: usize) -> Self {
        self.max_files = n;
        self
    }

    #[must_use]
    pub fn randomize_filename(mut self, on: bool) -> Self {
        self.randomize_filename = on;
        self
    }
}

/// Result of saving one uploaded file.
#[derive(Debug, Clone)]
pub struct SavedUpload {
    /// Storage key. Pass to `storage.url(&saved.key)` for a public URL.
    pub key: String,
    /// Original filename from the multipart field.
    pub original_filename: String,
    /// Content type the client claimed. Clients lie; do not trust it
    /// for access or rendering decisions.
    pub content_type: Option<String>,
    pub size_bytes: usize,
}

/// Save every file field in `mp` to `storage` and return one
/// [`SavedUpload`] per file. Fields without a `filename` are skipped.
///
/// The first error stops the loop and deletes the files already saved,
/// so a refused request leaves nothing behind (#1905).
///
/// # Errors
/// See [`UploadError`].
pub async fn save_uploads(
    mp: Multipart,
    cfg: &UploadConfig,
    storage: &BoxedStorage,
) -> Result<Vec<SavedUpload>, UploadError> {
    let mut out: Vec<SavedUpload> = Vec::new();
    let result = save_each(mp, cfg, storage, &mut out).await;
    if result.is_err() {
        for saved in &out {
            let _ = storage.delete(&saved.key).await;
        }
    }
    result.map(|()| out)
}

async fn save_each(
    mut mp: Multipart,
    cfg: &UploadConfig,
    storage: &BoxedStorage,
    out: &mut Vec<SavedUpload>,
) -> Result<(), UploadError> {
    let mut skipped = 0;
    while let Some(mut field) = mp.next_field().await? {
        let Some(filename) = field.file_name().map(str::to_owned) else {
            // Not a file field.
            if skipped < cfg.skip_fields {
                skipped += 1;
            }
            continue;
        };
        if out.len() >= cfg.max_files {
            return Err(UploadError::TooManyFiles { max: cfg.max_files });
        }
        let content_type = field.content_type().map(str::to_owned);

        // Read chunk by chunk and stop at the first chunk that crosses
        // `max_bytes`. Buffering the whole field first would let a huge
        // body eat memory even when the cap is small.
        let mut bytes: Vec<u8> = Vec::new();
        while let Some(chunk) = field.chunk().await? {
            if bytes.len().saturating_add(chunk.len()) > cfg.max_bytes {
                return Err(UploadError::TooLarge {
                    actual: bytes.len() + chunk.len(),
                    max: cfg.max_bytes,
                });
            }
            bytes.extend_from_slice(&chunk);
        }

        let ext = lowercase_ext(&filename);
        let allowed = match (&ext, cfg.allowed_extensions.is_empty()) {
            (Some(e), true) => !ACTIVE_EXTENSIONS.contains(&e.as_str()),
            (None, true) => true,
            (Some(e), false) => cfg.allowed_extensions.contains(e),
            (None, false) => false,
        };
        if !allowed {
            return Err(UploadError::BadExtension(
                ext.unwrap_or_else(|| "<none>".into()),
            ));
        }

        let safe_name = sanitize_filename(&filename);
        let key_filename = if cfg.randomize_filename {
            randomize(&safe_name)
        } else {
            safe_name
        };
        let key = format!("{}{key_filename}", cfg.prefix);

        let size = bytes.len();
        storage.save(&key, &bytes).await?;
        out.push(SavedUpload {
            key,
            original_filename: filename,
            content_type,
            size_bytes: size,
        });
    }
    Ok(())
}

// ---- pure helpers ----

fn ensure_trailing_slash(mut s: String) -> String {
    if !s.is_empty() && !s.ends_with('/') {
        s.push('/');
    }
    s
}

/// Drop directory parts and replace anything outside `[a-zA-Z0-9._-]`
/// with `_`, so the result is safe to use as a storage key. This is
/// what stops a client-supplied path from escaping the upload prefix.
pub fn sanitize_filename(name: &str) -> String {
    // Keep the basename only; clients sometimes send a full path.
    // Split on both separators instead of `Path::file_name`, which
    // treats `\` as a separator only on Windows. Otherwise the same
    // upload gets a different name depending on the server's OS.
    let base = name
        .rsplit(['/', '\\'])
        .next()
        .filter(|s| !s.is_empty())
        .unwrap_or(name);
    let mut out = String::with_capacity(base.len());
    for c in base.chars() {
        if c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-') {
            out.push(c);
        } else {
            out.push('_');
        }
    }
    if out.is_empty() {
        out.push_str("upload");
    }
    out
}

/// The lowercase extension without its dot, or `None`.
fn lowercase_ext(name: &str) -> Option<String> {
    Path::new(name)
        .extension()
        .and_then(|s| s.to_str())
        .map(str::to_ascii_lowercase)
}

/// Build `{uuid}-{name}` so two uploads with the same name do not
/// overwrite each other. A clock prefix collided within a tick (#1905).
fn randomize(name: &str) -> String {
    format!("{}-{name}", uuid::Uuid::new_v4().simple())
}

#[cfg(test)]
mod tests {
    use super::*;

    // -------- sanitize_filename

    #[test]
    fn sanitize_keeps_safe_ascii() {
        assert_eq!(sanitize_filename("photo.jpg"), "photo.jpg");
        assert_eq!(sanitize_filename("My_File-1.PDF"), "My_File-1.PDF");
    }

    #[test]
    fn sanitize_strips_directory_components() {
        // A client-sent path must not escape the prefix.
        assert_eq!(sanitize_filename("/etc/passwd"), "passwd");
        assert_eq!(sanitize_filename("../../etc/passwd"), "passwd");
        // Windows separators too, on every platform.
        assert_eq!(sanitize_filename("C:\\windows\\evil.exe"), "evil.exe");
        assert_eq!(sanitize_filename("C:/Users/me/photo.jpg"), "photo.jpg");
    }

    /// A path that is nothing but separators has no basename to keep.
    #[test]
    fn sanitize_handles_a_path_with_no_filename() {
        assert_eq!(sanitize_filename("/"), "_");
        assert_eq!(sanitize_filename("dir/"), "dir_");
    }

    #[test]
    fn sanitize_replaces_unsafe_chars() {
        assert_eq!(sanitize_filename("a b c.txt"), "a_b_c.txt");
        assert_eq!(sanitize_filename("hello+world.png"), "hello_world.png");
        assert_eq!(sanitize_filename("évil.jpg"), "_vil.jpg");
    }

    #[test]
    fn sanitize_empty_falls_back_to_upload() {
        assert_eq!(sanitize_filename(""), "upload");
    }

    // -------- extension extraction

    #[test]
    fn lowercase_ext_works() {
        assert_eq!(lowercase_ext("a.PNG").as_deref(), Some("png"));
        assert_eq!(lowercase_ext("a.tar.gz").as_deref(), Some("gz"));
        assert_eq!(lowercase_ext("noext"), None);
    }

    // -------- ensure_trailing_slash

    #[test]
    fn ensure_trailing_slash_idempotent() {
        assert_eq!(ensure_trailing_slash("a/".into()), "a/");
        assert_eq!(ensure_trailing_slash("a".into()), "a/");
        assert_eq!(ensure_trailing_slash(String::new()), "");
    }

    // -------- UploadConfig builder

    #[test]
    fn config_defaults() {
        let c = UploadConfig::new("uploads");
        assert_eq!(c.prefix, "uploads/");
        assert_eq!(c.max_bytes, 10 * 1024 * 1024);
        assert!(c.allowed_extensions.is_empty());
        assert!(c.randomize_filename);
    }

    #[test]
    fn config_allowed_extensions_normalize_to_lowercase() {
        let c = UploadConfig::new("u").allowed_extensions(&["PNG", "Jpg"]);
        assert!(c.allowed_extensions.contains("png"));
        assert!(c.allowed_extensions.contains("jpg"));
        assert!(!c.allowed_extensions.contains("PNG"));
    }

    // -------- randomize

    #[test]
    fn randomize_preserves_name_and_prepends_a_hex_id() {
        let r = randomize("photo.jpg");
        assert!(r.ends_with("-photo.jpg"));
        let dash = r.find('-').unwrap();
        assert!(r[..dash].chars().all(|c| c.is_ascii_hexdigit()));
    }

    /// #1905: back-to-back calls in one clock tick must not collide.
    #[test]
    fn randomize_never_repeats_in_a_tight_loop() {
        let keys: HashSet<String> = (0..10_000).map(|_| randomize("x")).collect();
        assert_eq!(keys.len(), 10_000);
    }

    // -------- save_uploads end-to-end via a real multipart body

    #[tokio::test]
    async fn save_uploads_writes_file_to_storage() {
        use crate::storage::InMemoryStorage;
        use axum::body::Body;
        use axum::http::header;
        use axum::http::Request;
        use axum::routing::post;
        use axum::Router;
        use std::sync::Arc as StdArc;
        use tower::ServiceExt;

        // A tiny app so the test runs the real multipart parser.
        let storage: BoxedStorage = StdArc::new(InMemoryStorage::new());
        let storage_for_handler = storage.clone();
        let app = Router::new().route(
            "/upload",
            post(move |mp: Multipart| {
                let storage = storage_for_handler.clone();
                async move {
                    let cfg = UploadConfig::new("uploads/").randomize_filename(false);
                    let saved = save_uploads(mp, &cfg, &storage).await.unwrap();
                    saved
                        .into_iter()
                        .map(|s| s.key)
                        .collect::<Vec<_>>()
                        .join(",")
                }
            }),
        );

        let boundary = "test-boundary";
        let body = format!(
            "--{boundary}\r\n\
             Content-Disposition: form-data; name=\"file\"; filename=\"hello.txt\"\r\n\
             Content-Type: text/plain\r\n\
             \r\n\
             hello world\r\n\
             --{boundary}--\r\n"
        );
        let req = Request::builder()
            .method("POST")
            .uri("/upload")
            .header(
                header::CONTENT_TYPE,
                format!("multipart/form-data; boundary={boundary}"),
            )
            .body(Body::from(body))
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), 200);
        let body_bytes = axum::body::to_bytes(resp.into_body(), 1 << 16)
            .await
            .unwrap();
        let body_str = std::str::from_utf8(&body_bytes).unwrap();
        assert_eq!(body_str, "uploads/hello.txt");

        // The file really landed in storage.
        let stored = storage.load("uploads/hello.txt").await.unwrap();
        assert_eq!(&stored, b"hello world");
    }

    #[tokio::test]
    async fn save_uploads_rejects_oversize_file() {
        use crate::storage::InMemoryStorage;
        use axum::body::Body;
        use axum::http::header;
        use axum::http::Request;
        use axum::routing::post;
        use axum::Router;
        use std::sync::Arc as StdArc;
        use tower::ServiceExt;

        let storage: BoxedStorage = StdArc::new(InMemoryStorage::new());
        let storage_for_handler = storage.clone();
        let app = Router::new().route(
            "/upload",
            post(move |mp: Multipart| {
                let storage = storage_for_handler.clone();
                async move {
                    let cfg = UploadConfig::new("u/").max_bytes(5);
                    match save_uploads(mp, &cfg, &storage).await {
                        Ok(_) => (axum::http::StatusCode::OK, "ok".to_owned()),
                        Err(e) => (axum::http::StatusCode::BAD_REQUEST, e.to_string()),
                    }
                }
            }),
        );

        let boundary = "b";
        // 100 bytes of payload, max is 5.
        let payload = "x".repeat(100);
        let body = format!(
            "--{boundary}\r\n\
             Content-Disposition: form-data; name=\"f\"; filename=\"big.bin\"\r\n\
             Content-Type: application/octet-stream\r\n\
             \r\n\
             {payload}\r\n\
             --{boundary}--\r\n"
        );
        let req = Request::builder()
            .method("POST")
            .uri("/upload")
            .header(
                header::CONTENT_TYPE,
                format!("multipart/form-data; boundary={boundary}"),
            )
            .body(Body::from(body))
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), 400);
        let body_str = std::str::from_utf8(
            &axum::body::to_bytes(resp.into_body(), 1 << 16)
                .await
                .unwrap(),
        )
        .unwrap()
        .to_owned();
        assert!(body_str.contains("file too large"), "got: {body_str}");
    }

    /// An oversize body must error and write nothing. We cannot see the
    /// early abort from outside, but we can check the error surfaces
    /// and storage stays empty.
    #[tokio::test]
    async fn save_uploads_aborts_streaming_on_oversize() {
        use crate::storage::InMemoryStorage;
        use axum::body::Body;
        use axum::http::header;
        use axum::http::Request;
        use axum::routing::post;
        use axum::Router;
        use std::sync::Arc as StdArc;
        use tower::ServiceExt;

        let storage: BoxedStorage = StdArc::new(InMemoryStorage::new());
        let storage_for_handler = storage.clone();
        let app = Router::new().route(
            "/upload",
            post(move |mp: Multipart| {
                let storage = storage_for_handler.clone();
                async move {
                    let cfg = UploadConfig::new("u/").max_bytes(1024); // 1 KiB cap
                    match save_uploads(mp, &cfg, &storage).await {
                        Ok(_) => (axum::http::StatusCode::OK, "ok".to_owned()),
                        Err(e) => (axum::http::StatusCode::BAD_REQUEST, e.to_string()),
                    }
                }
            }),
        );

        let boundary = "b";
        // 50 KiB payload against a 1 KiB cap.
        let payload = "x".repeat(50 * 1024);
        let body = format!(
            "--{boundary}\r\n\
             Content-Disposition: form-data; name=\"f\"; filename=\"big.bin\"\r\n\
             Content-Type: application/octet-stream\r\n\
             \r\n\
             {payload}\r\n\
             --{boundary}--\r\n"
        );
        let req = Request::builder()
            .method("POST")
            .uri("/upload")
            .header(
                header::CONTENT_TYPE,
                format!("multipart/form-data; boundary={boundary}"),
            )
            .body(Body::from(body))
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), 400);
        let body_str = std::str::from_utf8(
            &axum::body::to_bytes(resp.into_body(), 1 << 16)
                .await
                .unwrap(),
        )
        .unwrap()
        .to_owned();
        assert!(body_str.contains("file too large"), "got: {body_str}");
        // Nothing was written.
        assert!(
            storage.load("u/big.bin").await.is_err(),
            "no file should have been saved"
        );
    }

    /// POST one multipart body with `names` as file fields through `cfg`.
    async fn upload(cfg: UploadConfig, names: &[&str]) -> (u16, String, BoxedStorage) {
        use crate::storage::InMemoryStorage;
        use axum::body::Body;
        use axum::http::{header, Request};
        use axum::routing::post;
        use axum::Router;
        use tower::ServiceExt;

        let storage: BoxedStorage = std::sync::Arc::new(InMemoryStorage::new());
        let st = storage.clone();
        let app = Router::new().route(
            "/upload",
            post(move |mp: Multipart| async move {
                match save_uploads(mp, &cfg, &st).await {
                    Ok(v) => (axum::http::StatusCode::OK, v.len().to_string()),
                    Err(e) => (axum::http::StatusCode::BAD_REQUEST, e.to_string()),
                }
            }),
        );
        let mut body = String::new();
        for n in names {
            body.push_str(&format!(
                "--b\r\nContent-Disposition: form-data; name=\"f\"; filename=\"{n}\"\r\n\r\nx\r\n"
            ));
        }
        body.push_str("--b--\r\n");
        let req = Request::builder()
            .method("POST")
            .uri("/upload")
            .header(header::CONTENT_TYPE, "multipart/form-data; boundary=b")
            .body(Body::from(body))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        let status = resp.status().as_u16();
        let bytes = axum::body::to_bytes(resp.into_body(), 1 << 16)
            .await
            .unwrap();
        (
            status,
            String::from_utf8_lossy(&bytes).into_owned(),
            storage,
        )
    }

    /// With no allow-list, types that run script are still refused (#1849).
    #[tokio::test]
    async fn empty_allow_list_refuses_active_types() {
        for name in ["x.html", "x.SVG", "x.xml", "x.js"] {
            let (status, body, _) = upload(UploadConfig::new("u/"), &[name]).await;
            assert_eq!(status, 400, "{name}: {body}");
        }
        let (status, _, _) = upload(UploadConfig::new("u/"), &["x.png", "notes"]).await;
        assert_eq!(status, 200);
        let cfg = UploadConfig::new("u/").allowed_extensions(&["svg"]);
        assert_eq!(upload(cfg, &["x.svg"]).await.0, 200, "listed on purpose");
    }

    #[tokio::test]
    async fn max_files_caps_one_request() {
        let cfg = UploadConfig::new("u/").max_files(2);
        let (status, body, storage) = upload(
            cfg.clone().randomize_filename(false),
            &["a.png", "b.png", "c.png"],
        )
        .await;
        assert_eq!(status, 400);
        assert!(body.contains("too many files"), "{body}");
        for key in ["u/a.png", "u/b.png"] {
            assert!(!storage.exists(key).await.unwrap(), "orphan {key}");
        }
        let (status, body, _) = upload(cfg, &["a.png", "b.png"]).await;
        assert_eq!((status, body.as_str()), (200, "2"));
    }

    /// #1905: a later file's error removes the files saved before it.
    #[tokio::test]
    async fn a_refused_file_removes_the_ones_saved_before_it() {
        let cfg = UploadConfig::new("u/")
            .allowed_extensions(&["png"])
            .randomize_filename(false);
        let (status, body, storage) = upload(cfg, &["a.png", "b.exe"]).await;
        assert_eq!(status, 400, "{body}");
        assert!(!storage.exists("u/a.png").await.unwrap(), "orphan u/a.png");
    }

    #[tokio::test]
    async fn save_uploads_rejects_disallowed_extension() {
        use crate::storage::InMemoryStorage;
        use axum::body::Body;
        use axum::http::header;
        use axum::http::Request;
        use axum::routing::post;
        use axum::Router;
        use std::sync::Arc as StdArc;
        use tower::ServiceExt;

        let storage: BoxedStorage = StdArc::new(InMemoryStorage::new());
        let storage_for_handler = storage.clone();
        let app = Router::new().route(
            "/upload",
            post(move |mp: Multipart| {
                let storage = storage_for_handler.clone();
                async move {
                    let cfg = UploadConfig::new("u/").allowed_extensions(&["png", "jpg"]);
                    match save_uploads(mp, &cfg, &storage).await {
                        Ok(_) => (axum::http::StatusCode::OK, "ok".to_owned()),
                        Err(e) => (axum::http::StatusCode::BAD_REQUEST, e.to_string()),
                    }
                }
            }),
        );

        let boundary = "b";
        let body = format!(
            "--{boundary}\r\n\
             Content-Disposition: form-data; name=\"f\"; filename=\"evil.exe\"\r\n\
             \r\n\
             MZ\r\n\
             --{boundary}--\r\n"
        );
        let req = Request::builder()
            .method("POST")
            .uri("/upload")
            .header(
                header::CONTENT_TYPE,
                format!("multipart/form-data; boundary={boundary}"),
            )
            .body(Body::from(body))
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), 400);
        let body_str = std::str::from_utf8(
            &axum::body::to_bytes(resp.into_body(), 1 << 16)
                .await
                .unwrap(),
        )
        .unwrap()
        .to_owned();
        assert!(body_str.contains("not allowed"), "got: {body_str}");
    }
}