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

leyline/core/
multipart.rs
