Skip to main content

gproxy_protocol/codec/
body.rs

1use bytes::Bytes;
2use futures_util::StreamExt;
3
4use super::{CodecError, CodecErrorKind, CodecErrorStage, CodecFuture, CodecLimits};
5use crate::connection::HttpBody;
6
7/// Consumes one HTTP body under an aggregate byte limit.
8///
9/// A stream is consumed incrementally and its transport source is retained on
10/// failure. This helper is intentionally the only operation here that buffers
11/// an entire body; callers must supply a finite `max_body_bytes`.
12pub fn read_http_body(
13    body: HttpBody,
14    limits: CodecLimits,
15) -> CodecFuture<'static, Result<Bytes, CodecError>> {
16    Box::pin(async move {
17        match body {
18            HttpBody::Bytes(bytes) => {
19                if bytes.len() as u64 > limits.max_body_bytes {
20                    return Err(CodecError::new(
21                        CodecErrorKind::Limit,
22                        CodecErrorStage::Body,
23                        "HTTP body exceeds limit",
24                    ));
25                }
26                Ok(bytes)
27            }
28            HttpBody::Stream(mut stream) => {
29                let mut output = Vec::new();
30                while let Some(chunk) = stream.next().await {
31                    let chunk = chunk.map_err(|source| match source.downcast::<CodecError>() {
32                        Ok(error) => *error,
33                        Err(source) => CodecError::with_source(
34                            CodecErrorKind::Transport,
35                            CodecErrorStage::Body,
36                            "HTTP body stream failed",
37                            source,
38                        ),
39                    })?;
40                    let next = output.len().saturating_add(chunk.len());
41                    if next as u64 > limits.max_body_bytes {
42                        return Err(CodecError::new(
43                            CodecErrorKind::Limit,
44                            CodecErrorStage::Body,
45                            "HTTP body exceeds limit",
46                        ));
47                    }
48                    output.extend_from_slice(&chunk);
49                }
50                Ok(Bytes::from(output))
51            }
52        }
53    })
54}