poem/web/
tempfile.rs

1use std::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6use tokio::{
7    fs::File,
8    io::{AsyncRead, AsyncSeekExt, ReadBuf, SeekFrom},
9};
10
11use crate::{error::ReadBodyError, FromRequest, Request, RequestBody, Result};
12
13/// An extractor that extracts the body and writes the contents to a temporary
14/// file.
15///
16/// # Errors
17///
18/// - [`ReadBodyError`]
19#[cfg_attr(docsrs, doc(cfg(feature = "tempfile")))]
20pub struct TempFile(File);
21
22impl TempFile {
23    async fn internal_from_request(body: &mut RequestBody) -> Result<Self, ReadBodyError> {
24        let body = body.take()?;
25        let mut reader = body.into_async_read();
26        let mut file = tokio::fs::File::from_std(::libtempfile::tempfile()?);
27        tokio::io::copy(&mut reader, &mut file).await?;
28        file.seek(SeekFrom::Start(0)).await?;
29        Ok(Self(file))
30    }
31}
32
33impl<'a> FromRequest<'a> for TempFile {
34    async fn from_request(_req: &'a Request, body: &mut RequestBody) -> Result<Self> {
35        Self::internal_from_request(body).await.map_err(Into::into)
36    }
37}
38
39impl AsyncRead for TempFile {
40    fn poll_read(
41        mut self: Pin<&mut Self>,
42        cx: &mut Context<'_>,
43        buf: &mut ReadBuf<'_>,
44    ) -> Poll<std::io::Result<()>> {
45        Pin::new(&mut self.0).poll_read(cx, buf)
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use tokio::io::AsyncReadExt;
52
53    use super::*;
54    use crate::{handler, test::TestClient};
55
56    #[tokio::test]
57    async fn test_tempfile_extractor() {
58        #[handler(internal)]
59        async fn index(mut file: TempFile) {
60            let mut s = String::new();
61            file.read_to_string(&mut s).await.unwrap();
62            assert_eq!(s, "abcdef");
63        }
64
65        let cli = TestClient::new(index);
66        cli.get("/")
67            .body("abcdef")
68            .send()
69            .await
70            .assert_status_is_ok();
71    }
72}
poem/web/tempfile.rs

poem/web/
tempfile.rs
