salvo_core 0.94.0

//! Form parse module.
use std::ffi::OsStr;
use std::path::{Path, PathBuf};

use base64::engine::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use bytes::{Bytes, BytesMut};
use futures_util::stream::{Stream, TryStreamExt};
use mime::Mime;
use multimap::MultiMap;
use multra::{Field, Multipart};
use rand::TryRng;
use rand::rngs::SysRng;
use tempfile::Builder;
use tokio::fs::File;
use tokio::io::AsyncWriteExt;

use crate::http::ParseError;
use crate::http::header::{CONTENT_TYPE, HeaderMap};

/// The extracted text fields and uploaded files from a `multipart/form-data` request.
#[derive(Debug)]
#[non_exhaustive]
pub struct FormData {
    /// Name-value pairs for plain text fields. Technically, these are form data parts with no
    /// filename specified in the part's `Content-Disposition`.
    pub fields: MultiMap<String, String>,
    /// Name-value pairs for temporary files. Technically, these are form data parts with a
    /// filename specified in the part's `Content-Disposition`.
    pub files: MultiMap<String, FilePart>,
}

impl FormData {
    /// Creates a new `FormData`.
    #[inline]
    #[must_use]
    pub fn new() -> Self {
        Self {
            fields: MultiMap::new(),
            files: MultiMap::new(),
        }
    }

    /// Parse MIME `multipart/*` information from a stream as a `FormData`.
    pub(crate) async fn read<S, O, E>(headers: &HeaderMap, body: S) -> Result<Self, ParseError>
    where
        S: Stream<Item = Result<O, E>> + Send,
        O: Into<Bytes> + 'static,
        E: Into<Box<dyn std::error::Error + Send + Sync>>,
    {
        let content_type: Option<Mime> = headers
            .get(CONTENT_TYPE)
            .and_then(|h| h.to_str().ok())
            .and_then(|v| v.parse().ok());
        match content_type {
            Some(content_type) if content_type.subtype() == mime::WWW_FORM_URLENCODED => {
                futures_util::pin_mut!(body);
                let mut data = BytesMut::new();
                while let Some(chunk) = body.try_next().await.map_err(|e| {
                    let err = e.into();
                    if err.is::<http_body_util::LengthLimitError>() {
                        ParseError::PayloadTooLarge
                    } else {
                        ParseError::other(err)
                    }
                })? {
                    data.extend_from_slice(&chunk.into());
                }
                let mut form_data = Self::new();
                form_data.fields = form_urlencoded::parse(&data).into_owned().collect();
                Ok(form_data)
            }
            Some(content_type) if content_type.type_() == mime::MULTIPART => {
                let mut form_data = Self::new();
                let Some(boundary) = headers
                    .get(CONTENT_TYPE)
                    .and_then(|ct| ct.to_str().ok())
                    .and_then(|ct| multra::parse_boundary(ct).ok())
                else {
                    return Err(ParseError::InvalidContentType);
                };
                let mut multipart = Multipart::new(body, boundary);
                while let Some(mut field) = multipart.next_field().await.map_err(|e| {
                    // Check if the multra error contains a LengthLimitError
                    let mut source = std::error::Error::source(&e);
                    while let Some(err) = source {
                        if err.is::<http_body_util::LengthLimitError>() {
                            return ParseError::PayloadTooLarge;
                        }
                        source = std::error::Error::source(err);
                    }
                    ParseError::Multer(e)
                })? {
                    if let Some(name) = field.name().map(|s| s.to_owned()) {
                        if field.file_name().is_some() {
                            form_data
                                .files
                                .insert(name, FilePart::create(&mut field).await?);
                        } else {
                            form_data.fields.insert(name, field.text().await?);
                        }
                    }
                }
                Ok(form_data)
            }
            _ => Err(ParseError::InvalidContentType),
        }
    }
}
impl Default for FormData {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}
/// A file that is to be inserted into a `multipart/*` or alternatively an uploaded file that
/// was received as part of `multipart/*` parsing.
#[derive(Clone, Debug)]
pub struct FilePart {
    name: Option<String>,
    /// The headers of the part
    headers: HeaderMap,
    /// A temporary file containing the file content
    path: PathBuf,
    /// Optionally, the size of the file.  This is filled when multiparts are parsed, but is
    /// not necessary when they are generated.
    size: u64,
    // The temporary directory the upload was put into, saved for the Drop trait
    temp_dir: Option<PathBuf>,
}
impl FilePart {
    /// Get file name.
    #[inline]
    #[must_use]
    pub fn name(&self) -> Option<&str> {
        self.name.as_deref()
    }
    /// Get file name mutable reference.
    #[inline]
    pub fn name_mut(&mut self) -> Option<&mut String> {
        self.name.as_mut()
    }
    /// Get headers.
    #[inline]
    #[must_use]
    pub fn headers(&self) -> &HeaderMap {
        &self.headers
    }
    /// Get headers mutable reference.
    pub fn headers_mut(&mut self) -> &mut HeaderMap {
        &mut self.headers
    }
    /// Get content type.
    #[inline]
    pub fn content_type(&self) -> Option<Mime> {
        self.headers
            .get(CONTENT_TYPE)
            .and_then(|h| h.to_str().ok())
            .and_then(|v| v.parse().ok())
    }
    /// Get file path.
    #[inline]
    #[must_use]
    pub fn path(&self) -> &PathBuf {
        &self.path
    }
    /// Get file size.
    #[inline]
    #[must_use]
    pub fn size(&self) -> u64 {
        self.size
    }
    /// If you do not want the file on disk to be deleted when Self drops, call this
    /// function.  It will become your responsibility to clean up.
    #[inline]
    pub fn do_not_delete_on_drop(&mut self) {
        self.temp_dir = None;
    }

    /// Create a new temporary FilePart (when created this way, the file will be
    /// deleted once the FilePart object goes out of scope).
    ///
    /// # Security
    ///
    /// This streams the entire field to a temporary file with **no size limit
    /// of its own** — it writes until the field ends. When the body is read for
    /// the first time through
    /// [`Request::form_data`](crate::http::Request::form_data) it is bounded by
    /// the request's secure-max-size limit (configurable via
    /// [`Request::set_secure_max_size`](crate::http::Request::set_secure_max_size),
    /// the [`set_global_secure_max_size`](crate::http::request::set_global_secure_max_size)
    /// default, or the `SecureMaxSize` middleware).
    ///
    /// That guarantee only covers the **first** read of the body: if the body
    /// was already consumed and cached by an earlier `payload`/`payload_with_max_size`
    /// call (possibly under a larger limit or none), `form_data` reuses the
    /// cached payload without re-applying the limit. In that case, or when
    /// calling `create` directly on an otherwise unbounded `Field`, a client can
    /// fill the temporary directory's disk (denial of service), so make sure the
    /// request body is bounded before the first read.
    pub async fn create(field: &mut Field<'_>) -> Result<Self, ParseError> {
        // Setup a file to capture the contents.
        // Map the `JoinError` to a `ParseError` instead of panicking, so a
        // runtime issue (e.g. shutdown) on this request path doesn't abort.
        let mut path =
            tokio::task::spawn_blocking(|| Builder::new().prefix("salvo_http_multipart").tempdir())
                .await
                .map_err(ParseError::other)??
                .keep();
        let temp_dir = Some(path.clone());
        let name = field.file_name().map(|s| {
            // Sanitize filename by removing invalid characters
            s.chars()
                .filter(|c| {
                    !matches!(
                        c,
                        '/' | '\\' | '\0' | '<' | '>' | ':' | '"' | '|' | '?' | '*'
                    )
                })
                .collect::<String>()
        });
        path.push(format!(
            "{}.{}",
            text_nonce(),
            name.as_deref()
                .and_then(|name| { Path::new(name).extension().and_then(OsStr::to_str) })
                .unwrap_or("unknown")
        ));
        let mut file = File::create(&path).await?;
        let mut size = 0;
        while let Some(chunk) = field.chunk().await? {
            size += chunk.len() as u64;
            file.write_all(&chunk).await?;
        }
        file.sync_all().await?;
        Ok(Self {
            name,
            headers: field.headers().to_owned(),
            path,
            size,
            temp_dir,
        })
    }
}

fn cleanup_temporary_upload(path: &Path, temp_dir: &Path) {
    // Log warnings if cleanup fails to help identify potential disk space issues.
    if let Err(e) = std::fs::remove_file(path) {
        // Only log if the file still exists (ENOENT is expected if already cleaned up).
        if e.kind() != std::io::ErrorKind::NotFound {
            tracing::warn!(
                error = %e,
                path = %path.display(),
                "failed to remove temporary upload file"
            );
        }
    }
    if let Err(e) = std::fs::remove_dir(temp_dir) {
        // Only log if directory still exists and is not empty.
        if e.kind() != std::io::ErrorKind::NotFound {
            tracing::debug!(
                error = %e,
                path = %temp_dir.display(),
                "failed to remove temporary upload directory"
            );
        }
    }
}

impl Drop for FilePart {
    fn drop(&mut self) {
        if let Some(temp_dir) = self.temp_dir.take() {
            let path = self.path.clone();
            if tokio::runtime::Handle::try_current().is_ok() {
                tokio::task::spawn_blocking(move || cleanup_temporary_upload(&path, &temp_dir));
            } else {
                cleanup_temporary_upload(&path, &temp_dir);
            }
        }
    }
}

// Port from https://github.com/mikedilger/textnonce/blob/master/src/lib.rs
fn text_nonce() -> String {
    const BYTE_LEN: usize = 24;
    let mut raw = [0u8; BYTE_LEN];

    // First 12 bytes are derived from a monotonically advancing time source so
    // that nonce values issued within the same process tend to differ even on
    // RNG failure; the trailing 12 bytes are pure CSPRNG output. If the RNG
    // fails (extremely rare), we fall back to a wider time window so we still
    // emit a valid nonce instead of panicking.
    if let Ok(now) = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH) {
        raw[..4].copy_from_slice(&now.subsec_nanos().to_le_bytes());
        raw[4..12].copy_from_slice(&now.as_secs().to_le_bytes());
    }
    if SysRng.try_fill_bytes(&mut raw[12..]).is_err() {
        let micros = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_micros() as u64)
            .unwrap_or_default();
        raw[12..20].copy_from_slice(&micros.to_le_bytes());
        raw[20..].copy_from_slice(&micros.rotate_left(17).to_le_bytes()[..4]);
    }

    URL_SAFE_NO_PAD.encode(raw)
}

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

    #[test]
    fn file_part_drop_removes_temp_files_without_tokio_runtime() {
        let temp_dir = Builder::new()
            .prefix("salvo_http_multipart_drop_test")
            .tempdir()
            .expect("create temp dir")
            .keep();
        let path = temp_dir.join("upload.tmp");
        std::fs::write(&path, b"data").expect("write temp upload");

        {
            let _file_part = FilePart {
                name: None,
                headers: HeaderMap::new(),
                path: path.clone(),
                size: 4,
                temp_dir: Some(temp_dir.clone()),
            };
        }

        assert!(!path.exists());
        assert!(!temp_dir.exists());
    }
}