ullun/upload.rs
1//! Streaming request bodies
2//!
3//! By default a request body is read completely into memory (bounded by
4//! `App::max_body_size`) before the handler runs. That is right for JSON, wrong
5//! for a 2 GB upload. Routes registered with
6//! [`App::stream_body`](crate::app::App::stream_body) skip the buffering: the
7//! handler pulls the body piece by piece with [`Request::body_stream`](crate::request::Request::body_stream) and can
8//! write it to disk while it arrives, with a limit that is enforced as the
9//! bytes come in.
10//!
11//! ```no_run
12//! use ullun::prelude::*;
13//!
14//! # async fn run() {
15//! App::new()
16//! // up to 5 GiB for this prefix; memory use stays at a few chunks
17//! .stream_body("/upload", Some(5 << 30))
18//! .put("/upload/{name}", |mut req: Request| async move {
19//! let name = req.params.get("name")?.to_string();
20//! let mut body = req.body_stream()?;
21//! let bytes = body.save_to(format!("/var/data/{name}.part")).await?;
22//! Ok(Response::json(serde_json::json!({ "stored": bytes })))
23//! })
24//! .run("127.0.0.1:3000")
25//! .await
26//! .unwrap();
27//! # }
28//! ```
29//!
30//! Anything that needs the whole body (`Request::json`, `form`, `multipart`)
31//! sees an empty body on a streaming route; choose per prefix. Validate names
32//! and paths yourself: the stream only moves bytes.
33
34use crate::error::{Error, Result};
35use bytes::Bytes;
36use http_body_util::BodyExt;
37use std::path::Path;
38use tokio::io::{AsyncWrite, AsyncWriteExt};
39
40enum Source {
41 Hyper(hyper::body::Incoming),
42 Memory(Option<Bytes>),
43}
44
45/// A request body that is read in pieces, see the [module documentation](self)
46pub struct BodyStream {
47 source: Source,
48 limit: Option<u64>,
49 read: u64,
50 declared: Option<u64>,
51}
52
53fn too_large(limit: u64) -> Error {
54 Error::custom(413, format!("Request body too large (max {limit} bytes)"))
55}
56
57impl BodyStream {
58 pub(crate) fn from_incoming(body: hyper::body::Incoming, limit: Option<u64>) -> Self {
59 use hyper::body::Body;
60 let hint = body.size_hint();
61 let declared = hint.exact().or_else(|| {
62 // a lower bound above zero is only known when Content-Length was sent
63 (hint.lower() > 0 && hint.upper() == Some(hint.lower())).then(|| hint.lower())
64 });
65 Self {
66 source: Source::Hyper(body),
67 limit,
68 read: 0,
69 declared,
70 }
71 }
72
73 pub(crate) fn from_bytes(body: Bytes, limit: Option<u64>) -> Self {
74 let declared = Some(body.len() as u64);
75 Self {
76 source: Source::Memory(Some(body)),
77 limit,
78 read: 0,
79 declared,
80 }
81 }
82
83 /// Size announced by the client (`Content-Length`), if any. Untrusted; the
84 /// limit is checked against what actually arrives.
85 pub fn content_length(&self) -> Option<u64> {
86 self.declared
87 }
88
89 /// Bytes delivered so far
90 pub fn bytes_read(&self) -> u64 {
91 self.read
92 }
93
94 /// The next piece, or `None` at the end. Fails with `413` as soon as the
95 /// limit is exceeded, and with `400` when the connection breaks mid-body.
96 pub async fn chunk(&mut self) -> Result<Option<Bytes>> {
97 if let (Some(limit), Some(declared)) = (self.limit, self.declared) {
98 if declared > limit {
99 return Err(self.refuse(limit).await);
100 }
101 }
102 let next = match &mut self.source {
103 Source::Memory(slot) => slot.take().filter(|b| !b.is_empty()),
104 Source::Hyper(body) => loop {
105 match body.frame().await {
106 None => break None,
107 Some(Err(_)) => return Err(Error::bad_request("Request body was interrupted")),
108 Some(Ok(frame)) => {
109 if let Ok(data) = frame.into_data() {
110 if !data.is_empty() {
111 break Some(data);
112 }
113 }
114 }
115 }
116 },
117 };
118 if let Some(data) = &next {
119 self.read += data.len() as u64;
120 if let Some(limit) = self.limit {
121 if self.read > limit {
122 return Err(self.refuse(limit).await);
123 }
124 }
125 }
126 Ok(next)
127 }
128
129 /// Fail with 413. Part of the rest is read and thrown away first: closing a
130 /// connection with unread request data makes the OS reset it, and the
131 /// client (notably on Windows) then loses the response.
132 async fn refuse(&mut self, limit: u64) -> Error {
133 if let Source::Hyper(body) = &mut self.source {
134 crate::app::drain(body, self.read).await;
135 }
136 too_large(limit)
137 }
138
139 /// Write the rest of the body to `writer`; returns the bytes copied
140 pub async fn copy_to<W: AsyncWrite + Unpin>(&mut self, writer: &mut W) -> Result<u64> {
141 let mut total = 0;
142 while let Some(chunk) = self.chunk().await? {
143 writer
144 .write_all(&chunk)
145 .await
146 .map_err(|e| Error::internal(format!("Writing the upload failed: {e}")))?;
147 total += chunk.len() as u64;
148 }
149 writer
150 .flush()
151 .await
152 .map_err(|e| Error::internal(format!("Writing the upload failed: {e}")))?;
153 Ok(total)
154 }
155
156 /// Write the rest of the body to a file; returns the bytes written
157 ///
158 /// The file is created or replaced. If the body is cut off, too large or
159 /// the write fails, the partial file is removed. Use a temporary name and
160 /// rename it afterwards when readers must never see half-written data.
161 pub async fn save_to(&mut self, path: impl AsRef<Path>) -> Result<u64> {
162 let path = path.as_ref();
163 let mut file = tokio::fs::File::create(path)
164 .await
165 .map_err(|e| Error::internal(format!("Cannot create the upload file: {e}")))?;
166 match self.copy_to(&mut file).await {
167 Ok(n) => Ok(n),
168 Err(e) => {
169 drop(file);
170 let _ = tokio::fs::remove_file(path).await;
171 Err(e)
172 }
173 }
174 }
175
176 /// Read the rest into memory, still bounded by the limit
177 pub async fn collect(&mut self) -> Result<Bytes> {
178 let mut out = bytes::BytesMut::new();
179 while let Some(chunk) = self.chunk().await? {
180 out.extend_from_slice(&chunk);
181 }
182 Ok(out.freeze())
183 }
184}