base/crates/erp-server/src/handlers/upload.rs

use axum::Extension;
use axum::extract::{FromRef, Multipart, State};
use axum::response::Json;
use erp_core::error::AppError;
use erp_core::types::{ApiResponse, TenantContext};
use serde::Serialize;
use uuid::Uuid;

use crate::state::AppState;

#[derive(Debug, Serialize, utoipa::ToSchema)]
pub struct UploadResp {
    pub url: String,
    pub filename: String,
    pub size: u64,
    pub content_type: String,
}

/// 上传单个文件。
///
/// 接受 multipart/form-data，将文件保存到本地目录，
/// 返回可通过 `/uploads/` 前缀访问的 URL。
#[utoipa::path(
    post,
    path = "/upload",
    request_body(content_type = "multipart/form-data"),
    responses(
        (status = 200, description = "上传成功", body = ApiResponse<UploadResp>),
        (status = 413, description = "文件过大"),
        (status = 400, description = "无文件或不支持的类型"),
    ),
    tag = "文件上传",
)]
pub async fn upload_file<S>(
    State(state): State<AppState>,
    Extension(ctx): Extension<TenantContext>,
    mut multipart: Multipart,
) -> Result<Json<ApiResponse<UploadResp>>, AppError>
where
    AppState: FromRef<S>,
    S: Clone + Send + Sync + 'static,
{
    let max_size = state.config.storage.max_file_size_bytes();
    let upload_dir = &state.config.storage.upload_dir;

    // 确保上传目录存在
    let base_dir = std::path::Path::new(upload_dir);
    let tenant_dir = base_dir.join(ctx.tenant_id.to_string());
    tokio::fs::create_dir_all(&tenant_dir)
        .await
        .map_err(|e| AppError::Internal(format!("创建上传目录失败: {}", e)))?;

    // 读取第一个 field 作为上传文件
    let field = multipart
        .next_field()
        .await
        .map_err(|e| AppError::Validation(format!("读取上传数据失败: {}", e)))?
        .ok_or_else(|| AppError::Validation("未找到上传文件".to_string()))?;

    let content_type = field
        .content_type()
        .unwrap_or("application/octet-stream")
        .to_string();

    // 验证文件类型
    validate_content_type(&content_type)?;

    let original_name = field.name().unwrap_or("file").to_string();

    let data = field
        .bytes()
        .await
        .map_err(|e| AppError::Validation(format!("读取文件数据失败: {}", e)))?;

    if data.len() as u64 > max_size {
        return Err(AppError::Validation(format!(
            "文件大小超过限制（最大 {}）",
            format_size(max_size)
        )));
    }

    // 校验 magic bytes：验证文件实际内容与声明的 Content-Type 一致
    validate_magic_bytes(&content_type, &data)?;

    // 生成唯一文件名，保留原始扩展名
    let ext = std::path::Path::new(&original_name)
        .extension()
        .and_then(|e| e.to_str())
        .unwrap_or("bin");
    let file_id = Uuid::now_v7();
    let filename = format!("{}.{}", file_id, ext);

    let file_path = tenant_dir.join(&filename);
    let data_vec: Vec<u8> = data.to_vec();
    tokio::fs::write(&file_path, &data_vec)
        .await
        .map_err(|e| AppError::Internal(format!("写入文件失败: {}", e)))?;

    let url = format!("/uploads/{}/{}", ctx.tenant_id, filename);

    tracing::info!(
        tenant_id = %ctx.tenant_id,
        filename = %filename,
        size = data_vec.len(),
        content_type = %content_type,
        "文件上传成功"
    );

    Ok(Json(ApiResponse::ok(UploadResp {
        url,
        filename: original_name,
        size: data_vec.len() as u64,
        content_type,
    })))
}

/// 允许的文件类型
fn validate_content_type(content_type: &str) -> Result<(), AppError> {
    const ALLOWED: &[&str] = &[
        "image/jpeg",
        "image/png",
        "image/gif",
        "image/webp",
        "application/pdf",
        "application/msword",
        "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ];

    if !ALLOWED.contains(&content_type) {
        return Err(AppError::Validation(format!(
            "不支持的文件类型: {}",
            content_type
        )));
    }
    Ok(())
}

/// 校验文件 magic bytes（文件签名）与声明的 Content-Type 是否一致。
///
/// 防止攻击者通过修改 Content-Type 头上传恶意文件。
/// 对于 Office 格式等复杂签名，跳过 magic bytes 校验（仅依赖白名单）。
fn validate_magic_bytes(content_type: &str, data: &[u8]) -> Result<(), AppError> {
    // 需要至少几个字节才能校验
    if data.is_empty() {
        return Err(AppError::Validation("文件内容为空".to_string()));
    }

    let signature: &[u8] = match content_type {
        "image/jpeg" => {
            // JPEG: FF D8 FF
            b"\xFF\xD8\xFF"
        }
        "image/png" => {
            // PNG: 89 50 4E 47 0D 0A 1A 0A
            b"\x89PNG\r\n\x1A\n"
        }
        "image/gif" => {
            // GIF: 47 49 46 38 (GIF8)
            b"GIF8"
        }
        "image/webp" => {
            // WebP: RIFF....WEBP (12 bytes)
            // 前 4 字节: 52 49 46 46 (RIFF)
            // 字节 8-11: 57 45 42 50 (WEBP)
            if data.len() < 12 {
                return Err(AppError::Validation(
                    "文件数据不足，无法验证 WebP 格式".to_string(),
                ));
            }
            let riff_ok = &data[0..4] == b"RIFF";
            let webp_ok = &data[8..12] == b"WEBP";
            if riff_ok && webp_ok {
                return Ok(());
            }
            return Err(AppError::Validation(
                "文件内容与声明的类型 (image/webp) 不匹配".to_string(),
            ));
        }
        "application/pdf" => {
            // PDF: 25 50 44 46 (%PDF)
            b"%PDF"
        }
        // Office 格式的 magic bytes 较复杂（OLE2 / ZIP-based OOXML），
        // 仅依赖白名单，跳过 magic bytes 校验
        "application/msword"
        | "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        | "application/vnd.ms-excel"
        | "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" => {
            return Ok(());
        }
        _ => return Ok(()),
    };

    if data.len() < signature.len() {
        return Err(AppError::Validation(
            "文件数据不足，无法验证文件格式".to_string(),
        ));
    }

    if &data[..signature.len()] != signature {
        return Err(AppError::Validation(format!(
            "文件内容与声明的类型 ({}) 不匹配",
            content_type
        )));
    }

    Ok(())
}

fn format_size(bytes: u64) -> String {
    if bytes >= 1024 * 1024 * 1024 {
        format!("{}GB", bytes / (1024 * 1024 * 1024))
    } else if bytes >= 1024 * 1024 {
        format!("{}MB", bytes / (1024 * 1024))
    } else {
        format!("{}KB", bytes / 1024)
    }
}