base_http/body/
buffered.rs

1// Released under the MIT License.
2// Copyright, 2026, by Samuel Williams.
3
4use super::{BodyFrame, ReadableBody};
5use base_scheduler::{Cancellation, Cancelled};
6
7/// An in-memory, replayable sequence of chunks.
8///
9/// Reads clone each chunk, following Socketry's buffered body. The application
10/// chooses its owner: a `Vec<u8>` yields independent data, whereas an immutable
11/// shared owner can clone cheaply. Closing preserves the chunks for rewinding.
12#[derive(Clone, Debug)]
13pub struct BufferedBody<D = Vec<u8>> {
14    chunks: Chunks<D>,
15    index: usize,
16    length: Option<u64>,
17}
18
19// Keep the common static response allocation-free without changing the public
20// slice view or the replay cursor. Multiple chunks retain their original order.
21#[derive(Clone, Debug)]
22enum Chunks<D> {
23    Empty,
24    One(D),
25    Many(Vec<D>),
26}
27
28impl<D> Chunks<D> {
29    fn as_slice(&self) -> &[D] {
30        match self {
31            Self::Empty => &[],
32            Self::One(chunk) => std::slice::from_ref(chunk),
33            Self::Many(chunks) => chunks,
34        }
35    }
36
37    fn push(&mut self, chunk: D) {
38        match self {
39            Self::Empty => *self = Self::One(chunk),
40            Self::Many(chunks) => chunks.push(chunk),
41            Self::One(_) => {
42                let Self::One(first) = std::mem::replace(self, Self::Empty) else {
43                    unreachable!();
44                };
45                *self = Self::Many(vec![first, chunk]);
46            }
47        }
48    }
49}
50
51impl<D: AsRef<[u8]> + Clone + Send> BufferedBody<D> {
52    pub fn new(chunks: impl IntoIterator<Item = D>) -> Self {
53        let mut iter = chunks.into_iter();
54        let chunks = match iter.next() {
55            None => Chunks::Empty,
56            Some(first) => match iter.next() {
57                None => Chunks::One(first),
58                Some(second) => {
59                    let mut chunks = vec![first, second];
60                    chunks.extend(iter);
61                    Chunks::Many(chunks)
62                }
63            },
64        };
65        let length = chunks.as_slice().iter().try_fold(0u64, |sum, chunk| {
66            u64::try_from(chunk.as_ref().len())
67                .ok()
68                .and_then(|length| sum.checked_add(length))
69        });
70        Self {
71            chunks,
72            index: 0,
73            length,
74        }
75    }
76
77    /// All original chunks, including data already read.
78    pub fn chunks(&self) -> &[D] {
79        self.chunks.as_slice()
80    }
81
82    /// Read immediately available memory without scheduling an operation.
83    pub fn read_chunk(&mut self) -> Option<D> {
84        let chunk = self.chunks.as_slice().get(self.index)?.clone();
85        self.index += 1;
86        Some(chunk)
87    }
88
89    /// A buffered representation with an independent cursor at the beginning.
90    pub fn buffered(&self) -> Self {
91        Self {
92            chunks: self.chunks.clone(),
93            index: 0,
94            length: self.length,
95        }
96    }
97
98    /// Append a chunk, including after earlier chunks have been read or closed.
99    pub fn write(&mut self, chunk: D) {
100        self.length = self
101            .length
102            .zip(u64::try_from(chunk.as_ref().len()).ok())
103            .and_then(|(length, added)| length.checked_add(added));
104        self.chunks.push(chunk);
105    }
106
107    pub fn clear(&mut self) {
108        self.chunks = Chunks::Empty;
109        self.index = 0;
110        self.length = Some(0);
111    }
112}
113
114impl<D> Default for BufferedBody<D> {
115    fn default() -> Self {
116        Self {
117            chunks: Chunks::Empty,
118            index: 0,
119            length: Some(0),
120        }
121    }
122}
123
124impl<D: AsRef<[u8]> + Clone + Send> From<D> for BufferedBody<D> {
125    fn from(chunk: D) -> Self {
126        Self::new([chunk])
127    }
128}
129
130impl<D: AsRef<[u8]> + Clone + Send> ReadableBody for BufferedBody<D> {
131    type Data = D;
132    type Error = Cancelled;
133
134    fn may_have_trailers(&self) -> bool {
135        false
136    }
137
138    fn is_ready(&self) -> bool {
139        true
140    }
141
142    fn is_empty(&self) -> bool {
143        self.index >= self.chunks.as_slice().len()
144    }
145
146    fn length(&self) -> Option<u64> {
147        self.length
148    }
149
150    fn is_rewindable(&self) -> bool {
151        true
152    }
153
154    fn rewind(&mut self) -> bool {
155        self.index = 0;
156        true
157    }
158
159    fn close(&mut self, _error: Option<Self::Error>) {
160        self.index = self.chunks.as_slice().len();
161    }
162
163    async fn read(
164        &mut self,
165        cancellation: &Cancellation,
166    ) -> Result<BodyFrame<Self::Data>, Self::Error> {
167        cancellation.check()?;
168        Ok(match self.read_chunk() {
169            Some(data) => BodyFrame::Data {
170                data,
171                end: self.is_empty(),
172            },
173            None => BodyFrame::End,
174        })
175    }
176}
base_http/body/buffered.rs

base_http/body/
buffered.rs
