timvw_rust_runtime/
multipart.rs

1//! `multipart/form-data` bodies, with the API of reqwest's `multipart` module that generated clients use.
2//!
3//! A form whose parts are all held in memory encodes to a body held in memory, so a retry layer can send it again.
4//! A form with a streamed part encodes to a stream; when every part can be sent again (bytes, a file
5//! [`Upload`]), a retry makes that stream again, opening the files afresh. When the length of every part is known
6//! (bytes, a resolved file), the form is sent with a `Content-Length` instead of chunked.
7
8use std::borrow::Cow;
9use std::collections::hash_map::RandomState;
10use std::fmt;
11use std::hash::{BuildHasher, Hasher};
12use std::sync::Arc;
13
14use bytes::{Bytes, BytesMut};
15use futures_util::{future, stream, StreamExt};
16use http::HeaderValue;
17use http_body_util::BodyDataStream;
18
19use crate::maybe::box_stream;
20use crate::{Body, BoxError, Error, Upload};
21
22/// A `multipart/form-data` form.
23pub struct Form {
24    boundary: String,
25    parts: Vec<(Cow<'static, str>, Part)>,
26}
27
28/// A part of a [`Form`]: its content, file name and media type.
29pub struct Part {
30    content: Content,
31    file_name: Option<Cow<'static, str>>,
32    mime: Option<HeaderValue>,
33}
34
35enum Content {
36    Bytes(Bytes),
37    Stream(Body),
38}
39
40impl Form {
41    /// An empty form with a random boundary.
42    pub fn new() -> Self {
43        Self {
44            boundary: random_boundary(),
45            parts: Vec::new(),
46        }
47    }
48
49    /// The boundary between the parts.
50    pub fn boundary(&self) -> &str {
51        &self.boundary
52    }
53
54    /// Adds a text part.
55    pub fn text<N, V>(self, name: N, value: V) -> Self
56    where
57        N: Into<Cow<'static, str>>,
58        V: Into<Cow<'static, str>>,
59    {
60        self.part(name, Part::text(value))
61    }
62
63    /// Adds a part.
64    pub fn part<N: Into<Cow<'static, str>>>(mut self, name: N, part: Part) -> Self {
65        self.parts.push((name.into(), part));
66        self
67    }
68
69    /// Encodes the form.
70    ///
71    /// Parts held in memory encode to bytes. Otherwise the form is a stream, with a known length when every part has
72    /// one, and sent again by a retry when every part can be (see [`Body::try_clone`]).
73    pub fn into_body(self) -> Body {
74        let closing = Bytes::from(format!("--{}--\r\n", self.boundary));
75        if self
76            .parts
77            .iter()
78            .all(|(_, part)| matches!(part.content, Content::Bytes(_)))
79        {
80            let mut out = BytesMut::new();
81            for (name, part) in &self.parts {
82                out.extend_from_slice(&part.head(&self.boundary, name));
83                if let Content::Bytes(bytes) = &part.content {
84                    out.extend_from_slice(bytes);
85                }
86                out.extend_from_slice(b"\r\n");
87            }
88            out.extend_from_slice(&closing);
89            return Body::from(out.freeze());
90        }
91        let pieces: Vec<(Bytes, Content)> = self
92            .parts
93            .into_iter()
94            .map(|(name, part)| (part.head(&self.boundary, &name), part.content))
95            .collect();
96        let len = pieces
97            .iter()
98            .try_fold(closing.len() as u64, |len, (head, content)| {
99                Some(len + head.len() as u64 + content.len()? + 2)
100            });
101        if pieces
102            .iter()
103            .all(|(_, content)| content.try_clone().is_some())
104        {
105            let pieces = Arc::new(pieces);
106            return Body::replay(
107                Arc::new(move || {
108                    let fresh = pieces
109                        .iter()
110                        .map(|(head, content)| {
111                            let content = content
112                                .try_clone()
113                                .expect("every part was checked to clone");
114                            (head.clone(), content)
115                        })
116                        .collect();
117                    encode(fresh, closing.clone(), len)
118                }),
119                len,
120            );
121        }
122        encode(pieces, closing, len)
123    }
124}
125
126/// The stream of a form's parts, each a head, its content and a line break, then the closing boundary.
127fn encode(pieces: Vec<(Bytes, Content)>, closing: Bytes, len: Option<u64>) -> Body {
128    let chunks = stream::iter(pieces)
129        .map(|(head, content)| {
130            let head = stream::once(future::ready(Ok::<_, BoxError>(head)));
131            let content = match content {
132                Content::Bytes(bytes) => box_stream(stream::once(future::ready(Ok(bytes)))),
133                Content::Stream(body) => box_stream(BodyDataStream::new(body)),
134            };
135            let tail = stream::once(future::ready(Ok(Bytes::from_static(b"\r\n"))));
136            head.chain(content).chain(tail)
137        })
138        .flatten()
139        .chain(stream::once(future::ready(Ok(closing))));
140    Body::sized_stream(chunks, len)
141}
142
143impl Content {
144    fn len(&self) -> Option<u64> {
145        match self {
146            Content::Bytes(bytes) => Some(bytes.len() as u64),
147            Content::Stream(body) => http_body::Body::size_hint(body).exact(),
148        }
149    }
150
151    fn try_clone(&self) -> Option<Self> {
152        match self {
153            Content::Bytes(bytes) => Some(Content::Bytes(bytes.clone())),
154            Content::Stream(body) => body.try_clone().map(Content::Stream),
155        }
156    }
157
158    fn from_body(body: Body) -> Self {
159        match body.into_bytes() {
160            Ok(bytes) => Content::Bytes(bytes),
161            Err(body) => Content::Stream(body),
162        }
163    }
164}
165
166impl Default for Form {
167    fn default() -> Self {
168        Self::new()
169    }
170}
171
172impl fmt::Debug for Form {
173    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
174        f.debug_struct("Form")
175            .field("boundary", &self.boundary)
176            .field("parts", &self.parts)
177            .finish()
178    }
179}
180
181impl Part {
182    /// A text part.
183    pub fn text<V: Into<Cow<'static, str>>>(value: V) -> Self {
184        let bytes = match value.into() {
185            Cow::Borrowed(text) => Bytes::from_static(text.as_bytes()),
186            Cow::Owned(text) => Bytes::from(text),
187        };
188        Self::new(Content::Bytes(bytes))
189    }
190
191    /// A part of bytes.
192    pub fn bytes<V: Into<Cow<'static, [u8]>>>(value: V) -> Self {
193        let bytes = match value.into() {
194            Cow::Borrowed(bytes) => Bytes::from_static(bytes),
195            Cow::Owned(bytes) => Bytes::from(bytes),
196        };
197        Self::new(Content::Bytes(bytes))
198    }
199
200    /// A part read from a body, such as a file stream; a body held in memory stays in memory.
201    pub fn stream<T: Into<Body>>(value: T) -> Self {
202        Self::new(Content::from_body(value.into()))
203    }
204
205    /// A file part: the content of an [`Upload`], its file name (`file` when it has none) and its media type (none
206    /// when it has none). A file upload is streamed from disk, and opened again when a retry sends the form again.
207    pub fn upload(upload: Upload) -> Self {
208        let file_name = upload.file_name.clone().unwrap_or(Cow::Borrowed("file"));
209        let mime = upload.mime.as_deref().map(|mime| {
210            HeaderValue::try_from(mime)
211                .unwrap_or(HeaderValue::from_static("application/octet-stream"))
212        });
213        let mut part = Self::new(Content::from_body(Body::from(upload)));
214        part.file_name = Some(file_name);
215        part.mime = mime;
216        part
217    }
218
219    fn new(content: Content) -> Self {
220        Self {
221            content,
222            file_name: None,
223            mime: None,
224        }
225    }
226
227    /// Sets the file name.
228    pub fn file_name<T: Into<Cow<'static, str>>>(mut self, file_name: T) -> Self {
229        self.file_name = Some(file_name.into());
230        self
231    }
232
233    /// Sets the media type (`Content-Type` of the part).
234    pub fn mime_str(mut self, mime: &str) -> Result<Self, Error> {
235        let essence = mime.split(';').next().unwrap_or_default().trim();
236        if essence
237            .split_once('/')
238            .is_none_or(|(kind, sub)| kind.is_empty() || sub.is_empty())
239        {
240            return Err(Error::builder(format!("invalid media type `{mime}`")));
241        }
242        self.mime = Some(HeaderValue::try_from(mime).map_err(Error::builder)?);
243        Ok(self)
244    }
245
246    fn head(&self, boundary: &str, name: &str) -> Bytes {
247        let mut head = format!(
248            "--{boundary}\r\nContent-Disposition: form-data; name=\"{}\"",
249            escape(name)
250        );
251        if let Some(file_name) = &self.file_name {
252            head.push_str(&format!("; filename=\"{}\"", escape(file_name)));
253        }
254        head.push_str("\r\n");
255        if let Some(mime) = &self.mime {
256            head.push_str("Content-Type: ");
257            head.push_str(mime.to_str().unwrap_or("application/octet-stream"));
258            head.push_str("\r\n");
259        }
260        head.push_str("\r\n");
261        Bytes::from(head)
262    }
263}
264
265impl fmt::Debug for Part {
266    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
267        let content = match &self.content {
268            Content::Bytes(bytes) => format!("{} bytes", bytes.len()),
269            Content::Stream(_) => "stream".to_owned(),
270        };
271        f.debug_struct("Part")
272            .field("content", &content)
273            .field("file_name", &self.file_name)
274            .field("mime", &self.mime)
275            .finish()
276    }
277}
278
279/// Escapes a name or file name in a `Content-Disposition` header as browsers do (WHATWG HTML): `"`, CR and LF become
280/// percent escapes.
281fn escape(value: &str) -> Cow<'_, str> {
282    if value.contains(['"', '\r', '\n']) {
283        Cow::Owned(
284            value
285                .replace('"', "%22")
286                .replace('\r', "%0D")
287                .replace('\n', "%0A"),
288        )
289    } else {
290        Cow::Borrowed(value)
291    }
292}
293
294/// A random number, from the random keys of the standard library's hasher.
295pub(crate) fn random_u64(salt: u64) -> u64 {
296    let mut hasher = RandomState::new().build_hasher();
297    hasher.write_u64(salt);
298    hasher.finish()
299}
300
301fn random_boundary() -> String {
302    format!(
303        "{:016x}-{:016x}-{:016x}-{:016x}",
304        random_u64(1),
305        random_u64(2),
306        random_u64(3),
307        random_u64(4)
308    )
309}
310
311#[cfg(test)]
312mod tests {
313    use super::*;
314
315    fn form() -> Form {
316        Form {
317            boundary: "B".to_owned(),
318            parts: Vec::new(),
319        }
320    }
321
322    #[tokio::test]
323    async fn parts_in_memory_encode_to_a_cloneable_body() {
324        let body = form()
325            .text("name", "value")
326            .part(
327                "file",
328                Part::bytes(b"abc".to_vec())
329                    .file_name("a.txt")
330                    .mime_str("text/plain")
331                    .unwrap(),
332            )
333            .into_body();
334        assert!(body.try_clone().is_some());
335        let expected = "--B\r\nContent-Disposition: form-data; name=\"name\"\r\n\r\nvalue\r\n\
336                        --B\r\nContent-Disposition: form-data; name=\"file\"; filename=\"a.txt\"\r\n\
337                        Content-Type: text/plain\r\n\r\nabc\r\n--B--\r\n";
338        assert_eq!(body.collect_bytes().await.unwrap(), expected);
339    }
340
341    #[tokio::test]
342    async fn a_streamed_part_encodes_to_a_stream() {
343        let chunks = stream::iter(vec![
344            Ok::<_, std::io::Error>(Bytes::from("a")),
345            Ok(Bytes::from("b")),
346        ]);
347        let body = form()
348            .part(
349                "file",
350                Part::stream(Body::wrap_stream(chunks)).file_name("f"),
351            )
352            .into_body();
353        assert!(body.try_clone().is_none());
354        let expected = "--B\r\nContent-Disposition: form-data; name=\"file\"; filename=\"f\"\r\n\r\nab\r\n--B--\r\n";
355        assert_eq!(body.collect_bytes().await.unwrap(), expected);
356    }
357
358    #[tokio::test]
359    async fn replayable_uploads_encode_to_a_replayable_body_of_known_length() {
360        let counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
361        let made = counter.clone();
362        let upload = Upload::from_fn(move || {
363            made.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
364            stream::iter([Ok::<_, std::io::Error>(Bytes::from("ab"))])
365        })
366        .len(2)
367        .mime("text/plain");
368        let body = form()
369            .text("n", "v")
370            .part("f", Part::upload(upload))
371            .into_body();
372        let expected = "--B\r\nContent-Disposition: form-data; name=\"n\"\r\n\r\nv\r\n\
373                        --B\r\nContent-Disposition: form-data; name=\"f\"; filename=\"file\"\r\n\
374                        Content-Type: text/plain\r\n\r\nab\r\n--B--\r\n";
375        assert_eq!(
376            http_body::Body::size_hint(&body).exact(),
377            Some(expected.len() as u64)
378        );
379        let copy = body.try_clone().expect("every part can be made again");
380        assert_eq!(body.collect_bytes().await.unwrap(), expected);
381        assert_eq!(copy.collect_bytes().await.unwrap(), expected);
382        assert_eq!(counter.load(std::sync::atomic::Ordering::SeqCst), 2);
383    }
384
385    #[test]
386    fn a_part_of_unknown_length_leaves_the_form_length_unknown() {
387        let chunks = stream::iter([Ok::<_, std::io::Error>(Bytes::from("a"))]);
388        let body = form()
389            .text("n", "v")
390            .part("f", Part::upload(Upload::stream(chunks)))
391            .into_body();
392        assert_eq!(http_body::Body::size_hint(&body).exact(), None);
393        assert!(body.try_clone().is_none());
394    }
395
396    #[test]
397    fn names_are_escaped_and_media_types_checked() {
398        let part = Part::text("x").file_name("a\"b");
399        assert!(String::from_utf8_lossy(&part.head("B", "n\r\n"))
400            .contains("name=\"n%0D%0A\"; filename=\"a%22b\""));
401        assert!(Part::text("x").mime_str("nonsense").is_err());
402        assert!(Part::text("x")
403            .mime_str("application/json; charset=utf-8")
404            .is_ok());
405    }
406
407    #[test]
408    fn boundaries_differ() {
409        assert_ne!(Form::new().boundary, Form::new().boundary);
410    }
411}
timvw_rust_runtime/multipart.rs

timvw_rust_runtime/
multipart.rs
