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