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

leyline/core/
multipart.rs
