Upload something!

//! Run with //! //! ```not_rust //! cargo run -p example-stream-to-file //! ``` use axum::{ body::Bytes, extract::{Multipart, Path, Request}, http::StatusCode, response::{Html, Redirect}, routing::{get, post}, BoxError, Router, }; use futures_util::{Stream, TryStreamExt}; use std::{io, pin::pin}; use tokio::{fs::File, io::BufWriter}; use tokio_util::io::StreamReader; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; const UPLOADS_DIRECTORY: &str = "uploads"; #[tokio::main] async fn main() { tracing_subscriber::registry() .with( tracing_subscriber::EnvFilter::try_from_default_env() .unwrap_or_else(|_| format!("{}=debug", env!("CARGO_CRATE_NAME")).into()), ) .with(tracing_subscriber::fmt::layer()) .init(); // save files to a separate directory to not override files in the current directory tokio::fs::create_dir(UPLOADS_DIRECTORY) .await .expect("failed to create `uploads` directory"); let app = Router::new() .route("/", get(show_form).post(accept_form)) .route("/file/{file_name}", post(save_request_body)); let listener = tokio::net::TcpListener::bind("127.0.0.1:3000") .await .unwrap(); tracing::debug!("listening on {}", listener.local_addr().unwrap()); axum::serve(listener, app).await; } // Handler that streams the request body to a file. // // POST'ing to `/file/foo.txt` will create a file called `foo.txt`. async fn save_request_body( Path(file_name): Path, request: Request, ) -> Result<(), (StatusCode, String)> { stream_to_file(&file_name, request.into_body().into_data_stream()).await } // Handler that returns HTML for a multipart form. async fn show_form() -> Html<&'static str> { Html( r#" Upload something! "#, ) } // Handler that accepts a multipart form upload and streams each field to a file. async fn accept_form(mut multipart: Multipart) -> Result { while let Ok(Some(field)) = multipart.next_field().await { let file_name = if let Some(file_name) = field.file_name() { file_name.to_owned() } else { continue; }; stream_to_file(&file_name, field).await?; } Ok(Redirect::to("/")) } // Save a `Stream` to a file async fn stream_to_file(path: &str, stream: S) -> Result<(), (StatusCode, String)> where S: Stream>, E: Into, { if !path_is_valid(path) { return Err((StatusCode::BAD_REQUEST, "Invalid path".to_owned())); } async { // Convert the stream into an `AsyncRead`. let body_with_io_error = stream.map_err(io::Error::other); let mut body_reader = pin!(StreamReader::new(body_with_io_error)); // Create the file. `File` implements `AsyncWrite`. let path = std::path::Path::new(UPLOADS_DIRECTORY).join(path); let mut file = BufWriter::new(File::create(path).await?); // Copy the body into the file. tokio::io::copy(&mut body_reader, &mut file).await?; Ok::<_, io::Error>(()) } .await .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string())) } // to prevent directory traversal attacks we ensure the path consists of exactly one normal // component fn path_is_valid(path: &str) -> bool { let path = std::path::Path::new(path); let mut components = path.components().peekable(); if let Some(first) = components.peek() { if !matches!(first, std::path::Component::Normal(_)) { return false; } } components.count() == 1 }