File Uploads

use std::path::PathBuf;

use axum::{
    Json, Router,
    extract::{DefaultBodyLimit, Multipart, Path},
    http::StatusCode,
    routing::post,
};
use serde::Serialize;
use tokio::io::AsyncWriteExt;

const MAX_COVER_BYTES: usize = 5 * 1024 * 1024;

#[derive(Serialize)]
struct Stored {
    id: String,
    bytes: usize,
    mime: &'static str,
}

/// Content sniffing: trust the bytes, not the client's file name or Content-Type.
fn sniff_image(head: &[u8]) -> Option<&'static str> {
    match head {
        [0xFF, 0xD8, 0xFF, ..] => Some("image/jpeg"),
        [0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A, ..] => Some("image/png"),
        [b'R', b'I', b'F', b'F', _, _, _, _, b'W', b'E', b'B', b'P', ..] => Some("image/webp"),
        _ => None,
    }
}

async fn upload_cover(Path(isbn): Path<String>, mut multipart: Multipart) -> Result<Json<Stored>, (StatusCode, String)> {
    let bad = |msg: String| (StatusCode::BAD_REQUEST, msg);
    let upload_dir = PathBuf::from(std::env::var("UPLOAD_DIR").unwrap_or_else(|_| "uploads".into())); // NOT under static/
    tokio::fs::create_dir_all(&upload_dir).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    while let Some(mut field) = multipart.next_field().await.map_err(|e| (e.status(), e.body_text()))? {
        if field.name() != Some("cover") {
            continue; // ignore unexpected fields
        }
        tracing::debug!(client_name = ?field.file_name(), client_type = ?field.content_type(), "upload started");

        // sniff the first chunk before touching the disk
        let first = field.chunk().await.map_err(|e| (e.status(), e.body_text()))?.ok_or_else(|| bad("empty file".into()))?;
        let mime = sniff_image(&first).ok_or_else(|| bad("only JPEG, PNG or WebP images".into()))?;

        // a server-generated name: no path traversal, no overwriting, no executable extensions
        let id = format!("{isbn}-{}", uuid::Uuid::new_v4().simple());
        let tmp_path = upload_dir.join(format!("{id}.part"));
        let result = async {
            let mut file = tokio::fs::File::create(&tmp_path).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
            file.write_all(&first).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
            let mut written = first.len();
            // stream chunk by chunk: memory use stays at one chunk whatever the file size
            while let Some(chunk) = field.chunk().await.map_err(|e| (e.status(), e.body_text()))? {
                written += chunk.len();
                file.write_all(&chunk).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
            }
            file.flush().await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
            Ok::<_, (StatusCode, String)>(written)
        }
        .await;
        let written = match result {
            Ok(n) => n,
            Err(err) => {
                let _ = tokio::fs::remove_file(&tmp_path).await; // no half-written files left behind
                return Err(err);
            }
        };

        // only a complete, valid upload gets its final name
        let ext = mime.trim_start_matches("image/");
        tokio::fs::rename(&tmp_path, upload_dir.join(format!("{id}.{ext}")))
            .await
            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
        return Ok(Json(Stored { id, bytes: written, mime }));
    }
    Err(bad("missing 'cover' field".into()))
}

#[tokio::main]
async fn main() {
    let app = Router::new().route(
        "/api/v1/books/{isbn}/cover",
        post(upload_cover).layer(DefaultBodyLimit::max(MAX_COVER_BYTES)), // this route only
    );
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

$ printf '\x89PNG\r\n\x1a\n...rest of a png...' > cover.png
$ curl -s -F 'cover=@cover.png' localhost:3000/api/v1/books/9780441013593/cover
{"id":"9780441013593-a7eab4e710f8441fb9e49f3a9cf64f81","bytes":27,"mime":"image/png"}
$ curl -s -F 'cover=@notes.txt;type=image/png' localhost:3000/api/v1/books/9780441013593/cover
only JPEG, PNG or WebP images
$ (printf '\x89PNG\r\n\x1a\n'; head -c 6000000 /dev/zero) > big.png
$ curl -s -w ' %{http_code}\n' -F 'cover=@big.png' localhost:3000/api/v1/books/9780441013593/cover
Request payload is too large 413

use axum::{
    Router,
    extract::Path,
    http::{StatusCode, header},
    response::{IntoResponse, Response},
    routing::get,
};
use axum_extra::response::FileStream;
use tokio_util::io::ReaderStream;

async fn download_export(Path(id): Path<String>) -> Response {
    // only ids we generated: hex characters, so no path tricks are possible
    if id.len() != 32 || !id.bytes().all(|b| b.is_ascii_hexdigit()) {
        return StatusCode::NOT_FOUND.into_response();
    }
    let path = std::path::Path::new("exports").join(format!("{id}.csv"));
    let Ok(file) = tokio::fs::File::open(&path).await else {
        return StatusCode::NOT_FOUND.into_response();
    };
    let size = file.metadata().await.map(|m| m.len()).ok();
    let mut stream = FileStream::new(ReaderStream::new(file)).file_name("bookshelf-export.csv");
    if let Some(size) = size {
        stream = stream.content_size(size);
    }
    let mut response = stream.into_response();
    response.headers_mut().insert(header::CONTENT_TYPE, "text/csv".parse().expect("static value"));
    response.headers_mut().insert(header::X_CONTENT_TYPE_OPTIONS, "nosniff".parse().expect("static value"));
    response
}

#[tokio::main]
async fn main() {
    let app = Router::new().route("/exports/{id}", get(download_export));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}