leyline/core/
multipart.rs

1use std::io;
2use std::path::Path;
3use std::pin::Pin;
4use std::task::{Context, Poll};
5
6use bytes::Bytes;
7use futures_util::{Stream, TryFutureExt};
8
9use crate::core::body::{Body, BodyKind};
10
11mod debug;
12
13pub struct Part {
14    pub(crate) name: String,
15    pub(crate) body: Body,
16    pub(crate) filename: Option<String>,
17    pub(crate) mime: Option<String>,
18    pub(crate) extra_headers: Vec<(String, String)>,
19}
20
21impl Part {
22    pub fn text(value: impl Into<String>) -> Self {
23        let s = value.into();
24        Self {
25            name: String::new(),
26            body: Body::bytes(Bytes::from(s.into_bytes())),
27            filename: None,
28            mime: Some("text/plain; charset=utf-8".into()),
29            extra_headers: Vec::new(),
30        }
31    }
32
33    pub fn bytes(bytes: impl Into<Bytes>) -> Self {
34        Self {
35            name: String::new(),
36            body: Body::bytes(bytes.into()),
37            filename: None,
38            mime: None,
39            extra_headers: Vec::new(),
40        }
41    }
42
43    pub fn stream<S>(stream: S) -> Self
44    where
45        S: Stream<Item = io::Result<Bytes>> + Send + 'static,
46    {
47        Self {
48            name: String::new(),
49            body: Body::stream(stream, None),
50            filename: None,
51            mime: None,
52            extra_headers: Vec::new(),
53        }
54    }
55
56    pub fn file(path: impl AsRef<Path>) -> io::Result<Self> {
57        let path = path.as_ref().to_path_buf();
58        let metadata = std::fs::metadata(&path)?;
59        let filename = path
60            .file_name()
61            .and_then(|s| s.to_str())
62            .unwrap_or("file")
63            .to_string();
64        let stream = tokio::fs::File::open(path)
65            .map_ok(tokio_util::io::ReaderStream::new)
66            .try_flatten_stream();
67        Ok(Self {
68            name: String::new(),
69            body: Body::stream(stream, Some(metadata.len())),
70            filename: Some(filename),
71            mime: None,
72            extra_headers: Vec::new(),
73        })
74    }
75
76    pub fn filename(mut self, name: impl Into<String>) -> Self {
77        self.filename = Some(name.into());
78        self
79    }
80
81    pub fn mime(mut self, mime: impl Into<String>) -> Self {
82        self.mime = Some(mime.into());
83        self
84    }
85
86    pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
87        self.extra_headers.push((name.into(), value.into()));
88        self
89    }
90}
91
92pub struct Form {
93    pub(crate) parts: Vec<Part>,
94    pub(crate) boundary: String,
95}
96
97impl Default for Form {
98    fn default() -> Self {
99        Self::new()
100    }
101}
102
103impl Form {
104    pub fn new() -> Self {
105        Self {
106            parts: Vec::new(),
107            boundary: random_boundary(),
108        }
109    }
110
111    pub fn boundary(&self) -> &str {
112        &self.boundary
113    }
114
115    pub fn text(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
116        let mut part = Part::text(value);
117        part.name = name.into();
118        self.parts.push(part);
119        self
120    }
121
122    pub fn part(mut self, name: impl Into<String>, mut part: Part) -> Self {
123        part.name = name.into();
124        self.parts.push(part);
125        self
126    }
127
128    pub fn file(mut self, name: impl Into<String>, path: impl AsRef<Path>) -> io::Result<Self> {
129        let mut part = Part::file(path)?;
130        part.name = name.into();
131        self.parts.push(part);
132        Ok(self)
133    }
134
135    pub(crate) fn len_hint(&self) -> Option<u64> {
136        let mut total: u64 = 0;
137        for part in &self.parts {
138            total += 2 + self.boundary.len() as u64 + 2;
139            total += part_header_len(part) as u64;
140            let body_len = part.body.len_hint()?;
141            total += body_len;
142            total += 2;
143        }
144        total += 2 + self.boundary.len() as u64 + 4;
145        Some(total)
146    }
147
148    pub(crate) fn into_stream_body(self) -> Body {
149        if let Some(bytes) = self.buffered() {
150            return Body::bytes(bytes);
151        }
152        let length_hint = self.len_hint();
153        Body::stream(FormStream::new(self), length_hint)
154    }
155
156    fn buffered(&self) -> Option<Bytes> {
157        let mut out = Vec::with_capacity(usize::try_from(self.len_hint()?).ok()?);
158        for part in &self.parts {
159            let body: &[u8] = match &part.body.0 {
160                BodyKind::Empty => &[],
161                BodyKind::Bytes(b) => b,
162                BodyKind::Stream { .. } => return None,
163            };
164            out.extend_from_slice(&render_part_headers(&self.boundary, part).ok()?);
165            out.extend_from_slice(body);
166            out.extend_from_slice(b"\r\n");
167        }
168        out.extend_from_slice(b"--");
169        out.extend_from_slice(self.boundary.as_bytes());
170        out.extend_from_slice(b"--\r\n");
171        Some(Bytes::from(out))
172    }
173
174    pub(crate) fn content_type(&self) -> String {
175        format!("multipart/form-data; boundary={}", self.boundary)
176    }
177}
178
179impl From<Form> for Body {
180    fn from(form: Form) -> Body {
181        form.into_stream_body()
182    }
183}
184
185fn part_header_len(part: &Part) -> usize {
186    let mut n = "Content-Disposition: form-data; name=\"\"\r\n".len() + quoted_len(&part.name);
187    if let Some(fname) = &part.filename {
188        n += "; filename=\"\"".len() + quoted_len(fname);
189    }
190    if let Some(mime) = &part.mime {
191        n += "Content-Type: \r\n".len() + mime.len();
192    }
193    for (k, v) in &part.extra_headers {
194        n += k.len() + 2 + v.len() + 2;
195    }
196    n += 2;
197    n
198}
199
200struct FormStream {
201    boundary: String,
202    parts: std::collections::VecDeque<Part>,
203    state: FormState,
204}
205
206enum FormState {
207    NextPart,
208    InBody(Pin<Box<dyn Stream<Item = io::Result<Bytes>> + Send + 'static>>),
209    PartEnd,
210    Done,
211}
212
213impl FormStream {
214    fn new(form: Form) -> Self {
215        Self {
216            boundary: form.boundary,
217            parts: form.parts.into(),
218            state: FormState::NextPart,
219        }
220    }
221}
222
223fn quoted_len(input: &str) -> usize {
224    input.len() + input.bytes().filter(|b| matches!(b, b'"' | b'\\')).count()
225}
226
227fn escape_quoted(input: &str) -> Option<Vec<u8>> {
228    let mut out = Vec::with_capacity(input.len() + 2);
229    for b in input.bytes() {
230        match b {
231            0x00..=0x1F | 0x7F => return None,
232            b'"' | b'\\' => {
233                out.push(b'\\');
234                out.push(b);
235            }
236            _ => out.push(b),
237        }
238    }
239    Some(out)
240}
241
242fn render_part_headers(boundary: &str, part: &Part) -> io::Result<Bytes> {
243    let mut out = Vec::with_capacity(256);
244    out.extend_from_slice(b"--");
245    out.extend_from_slice(boundary.as_bytes());
246    out.extend_from_slice(b"\r\n");
247
248    out.extend_from_slice(b"Content-Disposition: form-data; name=\"");
249    let escaped_name = escape_quoted(&part.name).ok_or_else(|| {
250        io::Error::new(
251            io::ErrorKind::InvalidInput,
252            "multipart part `name` contains control characters (CRLF/NUL/etc.)",
253        )
254    })?;
255    out.extend_from_slice(&escaped_name);
256    out.push(b'"');
257    if let Some(fname) = &part.filename {
258        let escaped = escape_quoted(fname).ok_or_else(|| {
259            io::Error::new(
260                io::ErrorKind::InvalidInput,
261                "multipart part `filename` contains control characters (CRLF/NUL/etc.)",
262            )
263        })?;
264        out.extend_from_slice(b"; filename=\"");
265        out.extend_from_slice(&escaped);
266        out.push(b'"');
267    }
268    out.extend_from_slice(b"\r\n");
269
270    if let Some(mime) = &part.mime {
271        if mime.bytes().any(|b| matches!(b, 0x00..=0x1F | 0x7F)) {
272            return Err(io::Error::new(
273                io::ErrorKind::InvalidInput,
274                "multipart part `mime` contains control characters",
275            ));
276        }
277        out.extend_from_slice(b"Content-Type: ");
278        out.extend_from_slice(mime.as_bytes());
279        out.extend_from_slice(b"\r\n");
280    }
281
282    for (k, v) in &part.extra_headers {
283        let bad = k.bytes().any(|b| matches!(b, 0x00..=0x1F | 0x7F))
284            || v.bytes().any(|b| matches!(b, 0x00..=0x0A | 0x0D | 0x7F));
285        if bad {
286            return Err(io::Error::new(
287                io::ErrorKind::InvalidInput,
288                "multipart part extra header contains control characters",
289            ));
290        }
291        out.extend_from_slice(k.as_bytes());
292        out.extend_from_slice(b": ");
293        out.extend_from_slice(v.as_bytes());
294        out.extend_from_slice(b"\r\n");
295    }
296
297    out.extend_from_slice(b"\r\n");
298    Ok(Bytes::from(out))
299}
300
301impl Stream for FormStream {
302    type Item = io::Result<Bytes>;
303
304    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
305        loop {
306            // SAFETY: we never move `self.state` variants that hold a `Pin<Box<...>>` out of the structure — we only poll them in place via `.as_mut()`.
307            let this = &mut *self;
308            match &mut this.state {
309                FormState::NextPart => {
310                    let Some(part) = this.parts.pop_front() else {
311                        let mut out = Vec::with_capacity(this.boundary.len() + 6);
312                        out.extend_from_slice(b"--");
313                        out.extend_from_slice(this.boundary.as_bytes());
314                        out.extend_from_slice(b"--\r\n");
315                        this.state = FormState::Done;
316                        return Poll::Ready(Some(Ok(Bytes::from(out))));
317                    };
318                    let headers = match render_part_headers(&this.boundary, &part) {
319                        Ok(h) => h,
320                        Err(e) => {
321                            this.state = FormState::Done;
322                            return Poll::Ready(Some(Err(e)));
323                        }
324                    };
325                    let body = part.body;
326                    let body_stream: Pin<
327                        Box<dyn Stream<Item = io::Result<Bytes>> + Send + 'static>,
328                    > = match body.0 {
329                        BodyKind::Empty => Box::pin(futures_util::stream::empty()),
330                        BodyKind::Bytes(b) => {
331                            Box::pin(futures_util::stream::once(async move { Ok(b) }))
332                        }
333                        BodyKind::Stream { stream, .. } => stream,
334                    };
335                    this.state = FormState::InBody(body_stream);
336                    return Poll::Ready(Some(Ok(headers)));
337                }
338                FormState::InBody(stream) => match stream.as_mut().poll_next(cx) {
339                    Poll::Pending => return Poll::Pending,
340                    Poll::Ready(Some(Ok(b))) => return Poll::Ready(Some(Ok(b))),
341                    Poll::Ready(Some(Err(e))) => {
342                        this.state = FormState::Done;
343                        return Poll::Ready(Some(Err(e)));
344                    }
345                    Poll::Ready(None) => {
346                        this.state = FormState::PartEnd;
347                        continue;
348                    }
349                },
350                FormState::PartEnd => {
351                    this.state = FormState::NextPart;
352                    return Poll::Ready(Some(Ok(Bytes::from_static(b"\r\n"))));
353                }
354                FormState::Done => return Poll::Ready(None),
355            }
356        }
357    }
358}
359
360fn random_boundary() -> String {
361    format!(
362        "----LeylineFormBoundary{}",
363        crate::util::random_hex_token(16)
364    )
365}
366
367#[cfg(test)]
368mod tests;
leyline/core/multipart.rs

leyline/core/
multipart.rs
