Skip to main content

openmoat_proxy/post/
body.rs

1//! The body of an accepted answer, read whole or not at all, at most
2//! [`MAX_BODY`] bytes: an OTLP collector says in it that it took a request
3//! only in part (`partial_success`), which its status does not.
4
5use std::io::Read;
6
7/// The most bytes of an answer's body read; a longer one is not read.
8pub const MAX_BODY: usize = 64 * 1024;
9
10/// How the answer says where its body ends.
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub(super) enum Framing {
13    /// No body (`204`, `304`).
14    Empty,
15    /// `Content-Length`.
16    Length(usize),
17    /// `Transfer-Encoding: chunked`.
18    Chunked,
19    /// Neither: the body ends when the connection closes.
20    Close,
21}
22
23/// The body after the head, `read` holding what came with it: `None` when it
24/// is longer than [`MAX_BODY`], cut off, badly chunked or not read in time.
25pub(super) fn read(stream: &mut impl Read, mut raw: Vec<u8>, framing: Framing) -> Option<Vec<u8>> {
26    // The chunk lines and their line ends take room besides the body.
27    let limit = match framing {
28        Framing::Empty => return Some(Vec::new()),
29        Framing::Length(n) if n > MAX_BODY => return None,
30        Framing::Length(n) => n,
31        Framing::Chunked => 2 * MAX_BODY,
32        Framing::Close => MAX_BODY,
33    };
34    let mut chunk = [0u8; 4096];
35    loop {
36        match framing {
37            Framing::Length(n) if raw.len() >= n => {
38                raw.truncate(n);
39                return Some(raw);
40            }
41            Framing::Chunked => match dechunk(&raw) {
42                Chunks::Done(body) if body.len() <= MAX_BODY => return Some(body),
43                Chunks::Done(_) | Chunks::Bad => return None,
44                Chunks::More => {}
45            },
46            _ => {}
47        }
48        if raw.len() > limit {
49            return None;
50        }
51        let n = stream.read(&mut chunk).ok()?;
52        if n == 0 {
53            // Only a body that ends with the connection may end here.
54            return (framing == Framing::Close).then_some(raw);
55        }
56        raw.extend_from_slice(&chunk[..n]);
57    }
58}
59
60/// What a chunked body read so far holds.
61#[derive(Debug, PartialEq, Eq)]
62pub(super) enum Chunks {
63    /// Every chunk up to the last, empty one.
64    Done(Vec<u8>),
65    /// Not all of it yet.
66    More,
67    /// Not chunked encoding.
68    Bad,
69}
70
71/// Decode `raw` (RFC 9112 ยง7.1); trailers after the last chunk are not read.
72pub(super) fn dechunk(raw: &[u8]) -> Chunks {
73    let mut body = Vec::new();
74    let mut at = 0;
75    loop {
76        let Some(end) = raw[at..].windows(2).position(|w| w == b"\r\n") else {
77            return Chunks::More;
78        };
79        let line = &raw[at..at + end];
80        let size = line.split(|b| *b == b';').next().unwrap_or_default();
81        let Some(size) = std::str::from_utf8(size)
82            .ok()
83            .and_then(|s| usize::from_str_radix(s.trim(), 16).ok())
84        else {
85            return Chunks::Bad;
86        };
87        at += end + 2;
88        if size == 0 {
89            return Chunks::Done(body);
90        }
91        let Some(data_end) = at.checked_add(size) else {
92            return Chunks::Bad;
93        };
94        if raw.len() < data_end.saturating_add(2) {
95            return Chunks::More;
96        }
97        if &raw[data_end..data_end + 2] != b"\r\n" {
98            return Chunks::Bad;
99        }
100        body.extend_from_slice(&raw[at..data_end]);
101        at = data_end + 2;
102    }
103}