gproxy_protocol/codec/
multipart.rs

1use super::{CodecError, CodecErrorKind, CodecErrorStage, CodecFuture, CodecLimits};
2use crate::connection::{
3    ByteStream, HttpBody, Multipart, MultipartPart, PartStream, TransportError,
4};
5use bytes::Bytes;
6use futures_core::Stream;
7use futures_util::{StreamExt, stream};
8use std::{
9    collections::VecDeque,
10    sync::{
11        Arc, Mutex,
12        atomic::{AtomicBool, Ordering},
13    },
14    task::{Context, Poll},
15};
16
17#[cfg(not(target_arch = "wasm32"))]
18type Driver = Arc<Mutex<Pump>>;
19
20#[cfg(target_arch = "wasm32")]
21type Driver = std::rc::Rc<std::cell::RefCell<Pump>>;
22
23#[cfg(not(target_arch = "wasm32"))]
24fn access<T>(driver: &Driver, f: impl FnOnce(&mut Pump) -> T) -> T {
25    f(&mut driver.lock().unwrap())
26}
27
28#[cfg(target_arch = "wasm32")]
29fn access<T>(driver: &Driver, f: impl FnOnce(&mut Pump) -> T) -> T {
30    f(&mut driver.borrow_mut())
31}
32
33fn failure(kind: CodecErrorKind, message: &'static str) -> CodecError {
34    CodecError::new(kind, CodecErrorStage::Body, message)
35}
36
37fn map_error(source: multer::Error) -> CodecError {
38    let kind = match &source {
39        multer::Error::FieldSizeExceeded { .. } | multer::Error::StreamSizeExceeded { .. } => {
40            CodecErrorKind::Limit
41        }
42        multer::Error::StreamReadFailed(_) => CodecErrorKind::Transport,
43        multer::Error::IncompleteStream
44        | multer::Error::IncompleteFieldData { .. }
45        | multer::Error::IncompleteHeaders => CodecErrorKind::UnexpectedEof,
46        _ => CodecErrorKind::Multipart,
47    };
48    CodecError::with_source(
49        kind,
50        CodecErrorStage::Body,
51        "multipart parsing failed",
52        source,
53    )
54}
55
56struct Feed {
57    queue: VecDeque<Bytes>,
58    closed: bool,
59}
60
61struct FeedStream(Arc<Mutex<Feed>>);
62
63impl Stream for FeedStream {
64    type Item = Result<Bytes, TransportError>;
65    fn poll_next(self: std::pin::Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
66        let mut f = self.0.lock().unwrap();
67        if let Some(v) = f.queue.pop_front() {
68            Poll::Ready(Some(Ok(v)))
69        } else if f.closed {
70            Poll::Ready(None)
71        } else {
72            Poll::Pending
73        }
74    }
75}
76
77// Independent framing-size meter: multer parses values/headers; this meter
78// prevents overlong headers from accumulating inside its private buffer.
79#[derive(Clone, Copy)]
80enum MeterMode {
81    Preamble,
82    Headers,
83    Body,
84    Boundary,
85    End,
86}
87
88struct Meter {
89    mode: MeterMode,
90    line: Vec<u8>,
91    marker: Vec<u8>,
92    body_marker: Vec<u8>,
93    tail: Vec<u8>,
94    header_bytes: u64,
95    limits: CodecLimits,
96}
97
98impl Meter {
99    fn new(boundary: &str, limits: CodecLimits) -> Self {
100        Self {
101            mode: MeterMode::Preamble,
102            line: Vec::new(),
103            marker: format!("--{boundary}").into_bytes(),
104            body_marker: format!("\r\n--{boundary}").into_bytes(),
105            tail: Vec::new(),
106            header_bytes: 0,
107            limits,
108        }
109    }
110    fn push(&mut self, bytes: &[u8]) -> Result<(), CodecError> {
111        for &b in bytes {
112            match self.mode {
113                MeterMode::End => {}
114                MeterMode::Body => {
115                    self.tail.push(b);
116                    while !self.body_marker.starts_with(&self.tail) {
117                        self.tail.remove(0);
118                    }
119                    if self.tail.len() == self.body_marker.len() {
120                        self.tail.clear();
121                        self.line = self.marker.clone();
122                        self.mode = MeterMode::Boundary;
123                    }
124                }
125                MeterMode::Preamble | MeterMode::Boundary => {
126                    self.line.push(b);
127                    if self.line.starts_with(&self.marker)
128                        && self.line.get(self.marker.len()..) == Some(b"--".as_slice())
129                    {
130                        self.mode = MeterMode::End;
131                        self.line.clear();
132                        continue;
133                    }
134                    if self.line.len() as u64 > self.limits.max_buffer_bytes {
135                        return Err(failure(
136                            CodecErrorKind::Limit,
137                            "multipart preamble/boundary line exceeds buffer limit",
138                        ));
139                    }
140                    if b == b'\n' {
141                        let line = self
142                            .line
143                            .strip_suffix(b"\n")
144                            .unwrap()
145                            .strip_suffix(b"\r")
146                            .unwrap_or_else(|| self.line.strip_suffix(b"\n").unwrap());
147                        if let Some(suffix) = line.strip_prefix(self.marker.as_slice())
148                            && suffix.iter().all(|b| matches!(b, b' ' | b'\t'))
149                        {
150                            self.mode = MeterMode::Headers;
151                            self.header_bytes = 0;
152                        }
153                        self.line.clear();
154                    }
155                }
156                MeterMode::Headers => {
157                    self.header_bytes += 1;
158                    if self.header_bytes > self.limits.max_buffer_bytes {
159                        return Err(failure(
160                            CodecErrorKind::Limit,
161                            "multipart headers exceed buffer limit",
162                        ));
163                    }
164                    self.line.push(b);
165                    if b == b'\n' {
166                        let line = self.line.strip_suffix(b"\n").unwrap();
167                        let line = line.strip_suffix(b"\r").unwrap_or(line);
168                        if line.len() as u64 > self.limits.max_line_bytes {
169                            return Err(failure(
170                                CodecErrorKind::Limit,
171                                "multipart header line exceeds limit",
172                            ));
173                        }
174                        if line.is_empty() {
175                            self.mode = MeterMode::Body;
176                        }
177                        self.line.clear();
178                    } else if self.line.len() as u64
179                        > self
180                            .limits
181                            .max_line_bytes
182                            .saturating_add(u64::from(b == b'\r'))
183                    {
184                        return Err(failure(
185                            CodecErrorKind::Limit,
186                            "multipart header line exceeds limit",
187                        ));
188                    }
189                }
190            }
191        }
192        Ok(())
193    }
194}
195
196struct Pump {
197    input: Option<ByteStream>,
198    pending: Bytes,
199    feed: Arc<Mutex<Feed>>,
200    seen: u64,
201    limits: CodecLimits,
202    quantum: usize,
203    meter: Meter,
204    error: Option<CodecError>,
205}
206
207impl Pump {
208    fn close(&mut self) {
209        self.input = None;
210        self.pending = Bytes::new();
211        self.feed.lock().unwrap().closed = true;
212    }
213    fn poll_feed(&mut self, cx: &mut Context<'_>) -> Poll<()> {
214        if !self.pending.is_empty() {
215            let n = self.quantum.min(self.pending.len());
216            let bytes = self.pending.split_to(n);
217            if let Err(e) = self.meter.push(&bytes) {
218                self.error = Some(e);
219                self.close();
220            } else {
221                self.feed.lock().unwrap().queue.push_back(bytes);
222            }
223            return Poll::Ready(());
224        }
225        let Some(input) = self.input.as_mut() else {
226            self.feed.lock().unwrap().closed = true;
227            return Poll::Ready(());
228        };
229        match input.as_mut().poll_next(cx) {
230            Poll::Ready(Some(Ok(bytes))) => {
231                let n = self.seen.saturating_add(bytes.len() as u64);
232                if n > self.limits.max_body_bytes {
233                    self.error = Some(failure(
234                        CodecErrorKind::Limit,
235                        "multipart wire body exceeds limit",
236                    ));
237                    self.close();
238                } else {
239                    self.seen = n;
240                    self.pending = bytes;
241                }
242                Poll::Ready(())
243            }
244            Poll::Ready(Some(Err(e))) => {
245                self.error = Some(CodecError::with_source(
246                    CodecErrorKind::Transport,
247                    CodecErrorStage::Body,
248                    "multipart input failed",
249                    e,
250                ));
251                self.close();
252                Poll::Ready(())
253            }
254            Poll::Ready(None) => {
255                self.close();
256                Poll::Ready(())
257            }
258            Poll::Pending => Poll::Pending,
259        }
260    }
261    fn poll_drain(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), CodecError>> {
262        self.pending = Bytes::new();
263        self.feed.lock().unwrap().queue.clear();
264        if let Some(e) = self.error.take() {
265            return Poll::Ready(Err(e));
266        }
267        for _ in 0..32 {
268            let Some(input) = self.input.as_mut() else {
269                return Poll::Ready(Ok(()));
270            };
271            match input.as_mut().poll_next(cx) {
272                Poll::Ready(Some(Ok(bytes))) => {
273                    self.seen = self.seen.saturating_add(bytes.len() as u64);
274                    if self.seen > self.limits.max_body_bytes {
275                        self.close();
276                        return Poll::Ready(Err(failure(
277                            CodecErrorKind::Limit,
278                            "multipart epilogue exceeds body limit",
279                        )));
280                    }
281                }
282                Poll::Ready(Some(Err(e))) => {
283                    self.close();
284                    return Poll::Ready(Err(CodecError::with_source(
285                        CodecErrorKind::Transport,
286                        CodecErrorStage::Body,
287                        "multipart epilogue failed",
288                        e,
289                    )));
290                }
291                Poll::Ready(None) => {
292                    self.close();
293                    return Poll::Ready(Ok(()));
294                }
295                Poll::Pending => return Poll::Pending,
296            }
297        }
298        cx.waker().wake_by_ref();
299        Poll::Pending
300    }
301}
302
303struct FieldStream {
304    field: Option<multer::Field<'static>>,
305    driver: Driver,
306    failed: Arc<AtomicBool>,
307    done: bool,
308}
309
310impl Stream for FieldStream {
311    type Item = Result<Bytes, TransportError>;
312    fn poll_next(
313        mut self: std::pin::Pin<&mut Self>,
314        cx: &mut Context<'_>,
315    ) -> Poll<Option<Self::Item>> {
316        if self.done {
317            return Poll::Ready(None);
318        }
319        for _ in 0..32 {
320            if let Some(error) = access(&self.driver, |p| p.error.take()) {
321                self.done = true;
322                self.field = None;
323                self.failed.store(true, Ordering::Release);
324                return Poll::Ready(Some(Err(Box::new(error))));
325            }
326            let field = self.field.as_mut().expect("active field");
327            match std::pin::Pin::new(field).poll_next(cx) {
328                Poll::Ready(Some(Ok(bytes))) => return Poll::Ready(Some(Ok(bytes))),
329                Poll::Ready(Some(Err(e))) => {
330                    self.done = true;
331                    self.field = None;
332                    self.failed.store(true, Ordering::Release);
333                    return Poll::Ready(Some(Err(Box::new(map_error(e)))));
334                }
335                Poll::Ready(None) => {
336                    self.done = true;
337                    self.field = None;
338                    return Poll::Ready(None);
339                }
340                Poll::Pending => {
341                    if access(&self.driver, |p| p.poll_feed(cx)).is_pending() {
342                        return Poll::Pending;
343                    }
344                }
345            }
346        }
347        cx.waker().wake_by_ref();
348        Poll::Pending
349    }
350}
351
352/// Each part owns a real field stream. Consume/drop it before next_part;
353/// dropping an unread field lets multer discard it under the same size limits.
354/// The caller driving next_part OR the returned body drives the original input.
355/// On wasm only the private queue is Send; the original stream stays !Send.
356pub struct MultipartDecoder {
357    inner: multer::Multipart<'static>,
358    driver: Driver,
359    limits: CodecLimits,
360    parts: usize,
361    failed: Arc<AtomicBool>,
362    done: bool,
363}
364
365impl MultipartDecoder {
366    pub fn new(
367        body: HttpBody,
368        boundary: impl Into<String>,
369        limits: CodecLimits,
370    ) -> Result<Self, CodecError> {
371        let boundary = boundary.into();
372
373        let reserve = boundary.len() + 8;
374        let quantum = usize::try_from(limits.max_buffer_bytes)
375            .unwrap_or(usize::MAX)
376            .checked_sub(reserve)
377            .filter(|q| *q > 0)
378            .ok_or_else(|| {
379                failure(
380                    CodecErrorKind::Limit,
381                    "multipart buffer cannot hold boundary lookbehind",
382                )
383            })?
384            .min(4096);
385        let input: ByteStream = match body {
386            HttpBody::Bytes(b) => Box::pin(stream::once(async move { Ok(b) })),
387            HttpBody::Stream(s) => s,
388        };
389        let feed = Arc::new(Mutex::new(Feed {
390            queue: VecDeque::new(),
391            closed: false,
392        }));
393        let inner = multer::Multipart::with_constraints(
394            FeedStream(feed.clone()),
395            &boundary,
396            multer::Constraints::new().size_limit(
397                multer::SizeLimit::new()
398                    .whole_stream(limits.max_body_bytes)
399                    .per_field(limits.max_part_bytes),
400            ),
401        );
402        let pump = Pump {
403            input: Some(input),
404            pending: Bytes::new(),
405            feed,
406            seen: 0,
407            limits,
408            quantum,
409            meter: Meter::new(&boundary, limits),
410            error: None,
411        };
412        #[cfg(not(target_arch = "wasm32"))]
413        let driver = Arc::new(Mutex::new(pump));
414        #[cfg(target_arch = "wasm32")]
415        let driver = std::rc::Rc::new(std::cell::RefCell::new(pump));
416        Ok(Self {
417            inner,
418            driver,
419            limits,
420            parts: 0,
421            failed: Arc::new(AtomicBool::new(false)),
422            done: false,
423        })
424    }
425    pub fn next_part(&mut self) -> CodecFuture<'_, Result<Option<MultipartPart>, CodecError>> {
426        Box::pin(async move {
427            if self.failed.load(Ordering::Acquire) {
428                return Err(failure(
429                    CodecErrorKind::Multipart,
430                    "multipart decoder terminated",
431                ));
432            }
433            if self.done {
434                return Ok(None);
435            }
436            let driver = self.driver.clone();
437            let mut future = Box::pin(self.inner.next_field());
438            let result = futures_util::future::poll_fn(|cx| {
439                for _ in 0..32 {
440                    if let Some(e) = access(&driver, |p| p.error.take()) {
441                        return Poll::Ready(Err(e));
442                    }
443                    match future.as_mut().poll(cx) {
444                        Poll::Ready(Ok(v)) => return Poll::Ready(Ok(v)),
445                        Poll::Ready(Err(e)) => return Poll::Ready(Err(map_error(e))),
446                        Poll::Pending => {
447                            if access(&driver, |p| p.poll_feed(cx)).is_pending() {
448                                return Poll::Pending;
449                            }
450                        }
451                    }
452                }
453                cx.waker().wake_by_ref();
454                Poll::Pending
455            })
456            .await;
457            let field = match result {
458                Ok(v) => v,
459                Err(e) => {
460                    let concurrent = e.source_error().is_some_and(|s| {
461                        s.downcast_ref::<multer::Error>()
462                            .is_some_and(|e| matches!(e, multer::Error::LockFailure))
463                    });
464                    if !concurrent {
465                        self.failed.store(true, Ordering::Release);
466                        access(&self.driver, |p| p.close());
467                    }
468                    return Err(e);
469                }
470            };
471            let Some(field) = field else {
472                let result =
473                    futures_util::future::poll_fn(|cx| access(&self.driver, |p| p.poll_drain(cx)))
474                        .await;
475                self.done = true;
476                if let Err(e) = result {
477                    self.failed.store(true, Ordering::Release);
478                    return Err(e);
479                }
480                return Ok(None);
481            };
482            if self.parts >= self.limits.max_parts {
483                self.failed.store(true, Ordering::Release);
484                access(&self.driver, |p| p.close());
485                return Err(failure(
486                    CodecErrorKind::Limit,
487                    "multipart part count exceeds limit",
488                ));
489            }
490            self.parts += 1;
491            let headers = field.headers().clone();
492            Ok(Some(MultipartPart {
493                headers,
494                body: HttpBody::Stream(Box::pin(FieldStream {
495                    field: Some(field),
496                    driver: self.driver.clone(),
497                    failed: self.failed.clone(),
498                    done: false,
499                })),
500            }))
501        })
502    }
503}
504
505pub struct MultipartEncoder {
506    boundary: String,
507    parts: PartStream,
508    current: Option<ByteStream>,
509    prefix: Option<Bytes>,
510    closing_sent: bool,
511    failed: bool,
512    limits: CodecLimits,
513    part_bytes: u64,
514    wire_bytes: u64,
515    count: usize,
516}
517
518impl MultipartEncoder {
519    pub fn new(
520        boundary: impl Into<String>,
521        parts: Vec<MultipartPart>,
522        limits: CodecLimits,
523    ) -> Result<Self, CodecError> {
524        if parts.len() > limits.max_parts {
525            return Err(failure(
526                CodecErrorKind::Limit,
527                "multipart part count exceeds limit",
528            ));
529        }
530        Self::from_multipart(
531            boundary,
532            Multipart {
533                parts: Box::pin(stream::iter(parts.into_iter().map(Ok))),
534            },
535            limits,
536        )
537    }
538    pub fn from_multipart(
539        boundary: impl Into<String>,
540        body: Multipart,
541        limits: CodecLimits,
542    ) -> Result<Self, CodecError> {
543        let boundary = boundary.into();
544
545        Ok(Self {
546            boundary,
547            parts: body.parts,
548            current: None,
549            prefix: None,
550            closing_sent: false,
551            failed: false,
552            limits,
553            part_bytes: 0,
554            wire_bytes: 0,
555            count: 0,
556        })
557    }
558    fn account(&mut self, n: usize) -> Result<(), CodecError> {
559        let next = self.wire_bytes.saturating_add(n as u64);
560        if next > self.limits.max_body_bytes {
561            return Err(failure(
562                CodecErrorKind::Limit,
563                "multipart encoded body exceeds limit",
564            ));
565        }
566        self.wire_bytes = next;
567        Ok(())
568    }
569    pub fn next_chunk(&mut self) -> CodecFuture<'_, Result<Option<Bytes>, CodecError>> {
570        Box::pin(async move {
571            if self.failed {
572                return Err(failure(
573                    CodecErrorKind::Multipart,
574                    "multipart encoder terminated",
575                ));
576            }
577            let result = self.next_inner().await;
578            if result.is_err() {
579                self.failed = true;
580                self.current = None;
581                self.prefix = None;
582                self.parts = Box::pin(stream::empty());
583            }
584            result
585        })
586    }
587    async fn next_inner(&mut self) -> Result<Option<Bytes>, CodecError> {
588        loop {
589            if let Some(prefix) = self.prefix.take() {
590                self.account(prefix.len())?;
591                return Ok(Some(prefix));
592            }
593            if let Some(body) = self.current.as_mut() {
594                match body.next().await {
595                    Some(Ok(chunk)) => {
596                        let n = self.part_bytes.saturating_add(chunk.len() as u64);
597                        if n > self.limits.max_part_bytes {
598                            return Err(failure(
599                                CodecErrorKind::Limit,
600                                "multipart encoded part exceeds limit",
601                            ));
602                        }
603                        self.part_bytes = n;
604                        self.account(chunk.len())?;
605                        return Ok(Some(chunk));
606                    }
607                    Some(Err(e)) => {
608                        return Err(CodecError::with_source(
609                            CodecErrorKind::Transport,
610                            CodecErrorStage::Body,
611                            "multipart part stream failed",
612                            e,
613                        ));
614                    }
615                    None => {
616                        self.current = None;
617                        self.prefix = Some(Bytes::from_static(b"\r\n"));
618                        continue;
619                    }
620                }
621            }
622            if self.closing_sent {
623                return Ok(None);
624            }
625            if let Some(part) = self.parts.next().await {
626                let part = part.map_err(|e| {
627                    CodecError::with_source(
628                        CodecErrorKind::Transport,
629                        CodecErrorStage::Body,
630                        "multipart parts stream failed",
631                        e,
632                    )
633                })?;
634                if self.count >= self.limits.max_parts {
635                    return Err(failure(
636                        CodecErrorKind::Limit,
637                        "multipart part count exceeds limit",
638                    ));
639                }
640                self.count += 1;
641                self.part_bytes = 0;
642                self.prefix = Some(encode_prefix(&self.boundary, &part, self.limits)?);
643                self.current = Some(match part.body {
644                    HttpBody::Bytes(b) => Box::pin(stream::once(async move { Ok(b) })),
645                    HttpBody::Stream(s) => s,
646                });
647                continue;
648            }
649            self.closing_sent = true;
650            let end = Bytes::from(format!("--{}--\r\n", self.boundary));
651            self.account(end.len())?;
652            return Ok(Some(end));
653        }
654    }
655}
656
657fn encode_prefix(
658    boundary: &str,
659    part: &MultipartPart,
660    limits: CodecLimits,
661) -> Result<Bytes, CodecError> {
662    let mut length = boundary.len().saturating_add(4).saturating_add(2);
663    for (name, value) in &part.headers {
664        let n = name
665            .as_str()
666            .len()
667            .saturating_add(value.as_bytes().len())
668            .saturating_add(2);
669        if n as u64 > limits.max_line_bytes {
670            return Err(failure(
671                CodecErrorKind::Limit,
672                "multipart header line exceeds limit",
673            ));
674        }
675        length = length.saturating_add(n).saturating_add(2);
676        if length as u64 > limits.max_buffer_bytes {
677            return Err(failure(
678                CodecErrorKind::Limit,
679                "multipart headers exceed buffer limit",
680            ));
681        }
682    }
683    if length as u64 > limits.max_buffer_bytes {
684        return Err(failure(
685            CodecErrorKind::Limit,
686            "multipart prefix exceeds buffer limit",
687        ));
688    }
689    let mut output = Vec::with_capacity(length);
690    output.extend_from_slice(format!("--{boundary}\r\n").as_bytes());
691    for (name, value) in &part.headers {
692        output.extend_from_slice(name.as_str().as_bytes());
693        output.extend_from_slice(b": ");
694        output.extend_from_slice(value.as_bytes());
695        output.extend_from_slice(b"\r\n");
696    }
697    output.extend_from_slice(b"\r\n");
698    Ok(Bytes::from(output))
699}
gproxy_protocol/codec/multipart.rs

gproxy_protocol/codec/
multipart.rs
