Add 'TempFile::open()' to stream its data.

Resolves #2296.
This commit is contained in:
Sergio Benitez 2023-04-05 12:45:48 -07:00
parent 80b7755317
commit 89534129de
1 changed files with 52 additions and 9 deletions

View File

@ -10,7 +10,7 @@ use crate::fs::FileName;
use tokio::task; use tokio::task;
use tokio::fs::{self, File}; use tokio::fs::{self, File};
use tokio::io::AsyncWriteExt; use tokio::io::{AsyncWriteExt, AsyncBufRead, BufReader};
use tempfile::{NamedTempFile, TempPath}; use tempfile::{NamedTempFile, TempPath};
use either::Either; use either::Either;
@ -110,7 +110,7 @@ pub enum TempFile<'v> {
}, },
#[doc(hidden)] #[doc(hidden)]
Buffered { Buffered {
content: &'v str, content: &'v [u8],
} }
} }
@ -160,7 +160,7 @@ impl<'v> TempFile<'v> {
/// ///
/// Ok(()) /// Ok(())
/// } /// }
/// # let file = TempFile::Buffered { content: "hi".into() }; /// # let file = TempFile::Buffered { content: "hi".as_bytes() };
/// # rocket::async_test(handle(file)).unwrap(); /// # rocket::async_test(handle(file)).unwrap();
/// ``` /// ```
pub async fn persist_to<P>(&mut self, path: P) -> io::Result<()> pub async fn persist_to<P>(&mut self, path: P) -> io::Result<()>
@ -190,7 +190,7 @@ impl<'v> TempFile<'v> {
} }
TempFile::Buffered { content } => { TempFile::Buffered { content } => {
let mut file = File::create(&new_path).await?; let mut file = File::create(&new_path).await?;
file.write_all(content.as_bytes()).await?; file.write_all(content).await?;
*self = TempFile::File { *self = TempFile::File {
file_name: None, file_name: None,
content_type: None, content_type: None,
@ -231,7 +231,7 @@ impl<'v> TempFile<'v> {
/// ///
/// Ok(()) /// Ok(())
/// } /// }
/// # let file = TempFile::Buffered { content: "hi".into() }; /// # let file = TempFile::Buffered { content: "hi".as_bytes() };
/// # rocket::async_test(handle(file)).unwrap(); /// # rocket::async_test(handle(file)).unwrap();
/// ``` /// ```
pub async fn copy_to<P>(&mut self, path: P) -> io::Result<()> pub async fn copy_to<P>(&mut self, path: P) -> io::Result<()>
@ -258,7 +258,7 @@ impl<'v> TempFile<'v> {
TempFile::Buffered { content } => { TempFile::Buffered { content } => {
let path = path.as_ref(); let path = path.as_ref();
let mut file = File::create(path).await?; let mut file = File::create(path).await?;
file.write_all(content.as_bytes()).await?; file.write_all(content).await?;
*self = TempFile::File { *self = TempFile::File {
file_name: None, file_name: None,
content_type: None, content_type: None,
@ -296,7 +296,7 @@ impl<'v> TempFile<'v> {
/// ///
/// Ok(()) /// Ok(())
/// } /// }
/// # let file = TempFile::Buffered { content: "hi".into() }; /// # let file = TempFile::Buffered { content: "hi".as_bytes() };
/// # rocket::async_test(handle(file)).unwrap(); /// # rocket::async_test(handle(file)).unwrap();
/// ``` /// ```
pub async fn move_copy_to<P>(&mut self, path: P) -> io::Result<()> pub async fn move_copy_to<P>(&mut self, path: P) -> io::Result<()>
@ -313,6 +313,49 @@ impl<'v> TempFile<'v> {
Ok(()) Ok(())
} }
/// Open the file for reading, returning an `async` stream of the file.
///
/// This method should be used sparingly. `TempFile` is intended to be used
/// when the incoming data is destined to be stored on disk. If the incoming
/// data is intended to be streamed elsewhere, prefer to implement a custom
/// form guard via [`FromFormField`] that directly streams the incoming data
/// to the ultimate destination.
///
/// # Example
///
/// ```rust
/// # #[macro_use] extern crate rocket;
/// use rocket::fs::TempFile;
/// use rocket::tokio::io;
///
/// #[post("/", data = "<file>")]
/// async fn handle(file: TempFile<'_>) -> std::io::Result<()> {
/// let mut stream = file.open().await?;
/// io::copy(&mut stream, &mut io::stdout()).await?;
/// Ok(())
/// }
/// # let file = TempFile::Buffered { content: "hi".as_bytes() };
/// # rocket::async_test(handle(file)).unwrap();
/// ```
pub async fn open(&self) -> io::Result<impl AsyncBufRead + '_> {
use tokio_util::either::Either;
match self {
TempFile::File { path, .. } => {
let path = match path {
either::Either::Left(p) => p.as_ref(),
either::Either::Right(p) => p.as_path(),
};
let reader = BufReader::new(File::open(path).await?);
Ok(Either::Left(reader))
},
TempFile::Buffered { content } => {
Ok(Either::Right(*content))
},
}
}
/// Returns the size, in bytes, of the file. /// Returns the size, in bytes, of the file.
/// ///
/// This method does not perform any system calls. /// This method does not perform any system calls.
@ -353,7 +396,7 @@ impl<'v> TempFile<'v> {
/// ///
/// Ok(()) /// Ok(())
/// } /// }
/// # let file = TempFile::Buffered { content: "hi".into() }; /// # let file = TempFile::Buffered { content: "hi".as_bytes() };
/// # rocket::async_test(handle(file)).unwrap(); /// # rocket::async_test(handle(file)).unwrap();
/// ``` /// ```
pub fn path(&self) -> Option<&Path> { pub fn path(&self) -> Option<&Path> {
@ -474,7 +517,7 @@ impl<'v> TempFile<'v> {
impl<'v> FromFormField<'v> for Capped<TempFile<'v>> { impl<'v> FromFormField<'v> for Capped<TempFile<'v>> {
fn from_value(field: ValueField<'v>) -> Result<Self, Errors<'v>> { fn from_value(field: ValueField<'v>) -> Result<Self, Errors<'v>> {
let n = N { written: field.value.len() as u64, complete: true }; let n = N { written: field.value.len() as u64, complete: true };
Ok(Capped::new(TempFile::Buffered { content: field.value }, n)) Ok(Capped::new(TempFile::Buffered { content: field.value.as_bytes() }, n))
} }
async fn from_data( async fn from_data(