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.
5
6use std::borrow::Cow;
7use std::collections::hash_map::RandomState;
8use std::fmt;
9use std::hash::{BuildHasher, Hasher};
10
11use bytes::{Bytes, BytesMut};
12use futures_util::{future, stream, StreamExt, TryStreamExt};
13use http::HeaderValue;
14use http_body_util::BodyDataStream;
15
16use crate::{Body, BoxError, Error};
17
18/// A `multipart/form-data` form.
19pub struct Form {
20    boundary: String,
21    parts: Vec<(Cow<'static, str>, Part)>,
22}
23
24/// A part of a [`Form`]: its content, file name and media type.
25pub struct Part {
26    content: Content,
27    file_name: Option<Cow<'static, str>>,
28    mime: Option<HeaderValue>,
29}
30
31enum Content {
32    Bytes(Bytes),
33    Stream(Body),
34}
35
36impl Form {
37    /// An empty form with a random boundary.
38    pub fn new() -> Self {
39        Self {
40            boundary: random_boundary(),
41            parts: Vec::new(),
42        }
43    }
44
45    /// The boundary between the parts.
46    pub fn boundary(&self) -> &str {
47        &self.boundary
48    }
49
50    /// Adds a text part.
51    pub fn text<N, V>(self, name: N, value: V) -> Self
52    where
53        N: Into<Cow<'static, str>>,
54        V: Into<Cow<'static, str>>,
55    {
56        self.part(name, Part::text(value))
57    }
58
59    /// Adds a part.
60    pub fn part<N: Into<Cow<'static, str>>>(mut self, name: N, part: Part) -> Self {
61        self.parts.push((name.into(), part));
62        self
63    }
64
65    /// Encodes the form.
66    pub fn into_body(self) -> Body {
67        let closing = Bytes::from(format!("--{}--\r\n", self.boundary));
68        if self
69            .parts
70            .iter()
71            .all(|(_, part)| matches!(part.content, Content::Bytes(_)))
72        {
73            let mut out = BytesMut::new();
74            for (name, part) in &self.parts {
75                out.extend_from_slice(&part.head(&self.boundary, name));
76                if let Content::Bytes(bytes) = &part.content {
77                    out.extend_from_slice(bytes);
78                }
79                out.extend_from_slice(b"\r\n");
80            }
81            out.extend_from_slice(&closing);
82            return Body::from(out.freeze());
83        }
84        let boundary = self.boundary;
85        let chunks = stream::iter(self.parts)
86            .map(move |(name, part)| {
87                let head = stream::once(future::ready(Ok::<_, BoxError>(
88                    part.head(&boundary, &name),
89                )));
90                let content = match part.content {
91                    Content::Bytes(bytes) => stream::once(future::ready(Ok(bytes))).boxed(),
92                    Content::Stream(body) => BodyDataStream::new(body).boxed(),
93                };
94                let tail = stream::once(future::ready(Ok(Bytes::from_static(b"\r\n"))));
95                head.chain(content).chain(tail)
96            })
97            .flatten()
98            .chain(stream::once(future::ready(Ok(closing))));
99        Body::wrap_stream(chunks.into_stream())
100    }
101}
102
103impl Default for Form {
104    fn default() -> Self {
105        Self::new()
106    }
107}
108
109impl fmt::Debug for Form {
110    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
111        f.debug_struct("Form")
112            .field("boundary", &self.boundary)
113            .field("parts", &self.parts)
114            .finish()
115    }
116}
117
118impl Part {
119    /// A text part.
120    pub fn text<V: Into<Cow<'static, str>>>(value: V) -> Self {
121        let bytes = match value.into() {
122            Cow::Borrowed(text) => Bytes::from_static(text.as_bytes()),
123            Cow::Owned(text) => Bytes::from(text),
124        };
125        Self::new(Content::Bytes(bytes))
126    }
127
128    /// A part of bytes.
129    pub fn bytes<V: Into<Cow<'static, [u8]>>>(value: V) -> Self {
130        let bytes = match value.into() {
131            Cow::Borrowed(bytes) => Bytes::from_static(bytes),
132            Cow::Owned(bytes) => Bytes::from(bytes),
133        };
134        Self::new(Content::Bytes(bytes))
135    }
136
137    /// A part read from a body, such as a file stream; a body held in memory stays in memory.
138    pub fn stream<T: Into<Body>>(value: T) -> Self {
139        let body = value.into();
140        match body.as_bytes() {
141            Some(bytes) => Self::new(Content::Bytes(Bytes::copy_from_slice(bytes))),
142            None => Self::new(Content::Stream(body)),
143        }
144    }
145
146    fn new(content: Content) -> Self {
147        Self {
148            content,
149            file_name: None,
150            mime: None,
151        }
152    }
153
154    /// Sets the file name.
155    pub fn file_name<T: Into<Cow<'static, str>>>(mut self, file_name: T) -> Self {
156        self.file_name = Some(file_name.into());
157        self
158    }
159
160    /// Sets the media type (`Content-Type` of the part).
161    pub fn mime_str(mut self, mime: &str) -> Result<Self, Error> {
162        let essence = mime.split(';').next().unwrap_or_default().trim();
163        if essence
164            .split_once('/')
165            .is_none_or(|(kind, sub)| kind.is_empty() || sub.is_empty())
166        {
167            return Err(Error::builder(format!("invalid media type `{mime}`")));
168        }
169        self.mime = Some(HeaderValue::try_from(mime).map_err(Error::builder)?);
170        Ok(self)
171    }
172
173    fn head(&self, boundary: &str, name: &str) -> Bytes {
174        let mut head = format!(
175            "--{boundary}\r\nContent-Disposition: form-data; name=\"{}\"",
176            escape(name)
177        );
178        if let Some(file_name) = &self.file_name {
179            head.push_str(&format!("; filename=\"{}\"", escape(file_name)));
180        }
181        head.push_str("\r\n");
182        if let Some(mime) = &self.mime {
183            head.push_str("Content-Type: ");
184            head.push_str(mime.to_str().unwrap_or("application/octet-stream"));
185            head.push_str("\r\n");
186        }
187        head.push_str("\r\n");
188        Bytes::from(head)
189    }
190}
191
192impl fmt::Debug for Part {
193    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
194        let content = match &self.content {
195            Content::Bytes(bytes) => format!("{} bytes", bytes.len()),
196            Content::Stream(_) => "stream".to_owned(),
197        };
198        f.debug_struct("Part")
199            .field("content", &content)
200            .field("file_name", &self.file_name)
201            .field("mime", &self.mime)
202            .finish()
203    }
204}
205
206/// Escapes a name or file name in a `Content-Disposition` header as browsers do (WHATWG HTML): `"`, CR and LF become
207/// percent escapes.
208fn escape(value: &str) -> Cow<'_, str> {
209    if value.contains(['"', '\r', '\n']) {
210        Cow::Owned(
211            value
212                .replace('"', "%22")
213                .replace('\r', "%0D")
214                .replace('\n', "%0A"),
215        )
216    } else {
217        Cow::Borrowed(value)
218    }
219}
220
221fn random_boundary() -> String {
222    let random = |salt: u64| {
223        let mut hasher = RandomState::new().build_hasher();
224        hasher.write_u64(salt);
225        hasher.finish()
226    };
227    format!(
228        "{:016x}-{:016x}-{:016x}-{:016x}",
229        random(1),
230        random(2),
231        random(3),
232        random(4)
233    )
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239
240    fn form() -> Form {
241        Form {
242            boundary: "B".to_owned(),
243            parts: Vec::new(),
244        }
245    }
246
247    #[tokio::test]
248    async fn parts_in_memory_encode_to_a_cloneable_body() {
249        let body = form()
250            .text("name", "value")
251            .part(
252                "file",
253                Part::bytes(b"abc".to_vec())
254                    .file_name("a.txt")
255                    .mime_str("text/plain")
256                    .unwrap(),
257            )
258            .into_body();
259        assert!(body.try_clone().is_some());
260        let expected = "--B\r\nContent-Disposition: form-data; name=\"name\"\r\n\r\nvalue\r\n\
261                        --B\r\nContent-Disposition: form-data; name=\"file\"; filename=\"a.txt\"\r\n\
262                        Content-Type: text/plain\r\n\r\nabc\r\n--B--\r\n";
263        assert_eq!(body.collect_bytes().await.unwrap(), expected);
264    }
265
266    #[tokio::test]
267    async fn a_streamed_part_encodes_to_a_stream() {
268        let chunks = stream::iter(vec![
269            Ok::<_, std::io::Error>(Bytes::from("a")),
270            Ok(Bytes::from("b")),
271        ]);
272        let body = form()
273            .part(
274                "file",
275                Part::stream(Body::wrap_stream(chunks)).file_name("f"),
276            )
277            .into_body();
278        assert!(body.try_clone().is_none());
279        let expected = "--B\r\nContent-Disposition: form-data; name=\"file\"; filename=\"f\"\r\n\r\nab\r\n--B--\r\n";
280        assert_eq!(body.collect_bytes().await.unwrap(), expected);
281    }
282
283    #[test]
284    fn names_are_escaped_and_media_types_checked() {
285        let part = Part::text("x").file_name("a\"b");
286        assert!(String::from_utf8_lossy(&part.head("B", "n\r\n"))
287            .contains("name=\"n%0D%0A\"; filename=\"a%22b\""));
288        assert!(Part::text("x").mime_str("nonsense").is_err());
289        assert!(Part::text("x")
290            .mime_str("application/json; charset=utf-8")
291            .is_ok());
292    }
293
294    #[test]
295    fn boundaries_differ() {
296        assert_ne!(Form::new().boundary, Form::new().boundary);
297    }
298}
timvw_rust_runtime/multipart.rs

timvw_rust_runtime/
multipart.rs
