zencodec_io/
window.rs

1//! The driver-owned input buffer and the protocol checks every driver
2//! shares.
3
4use alloc::vec::Vec;
5use core::fmt;
6
7use crate::admission::{InputError, reserve};
8use zencodec::decode::{Input, InputRequest};
9
10/// Bytes the driver holds: `buf[head..end]` starts at absolute `base`.
11///
12/// `buf` is zeroed once as it grows and then reused: reads go straight into
13/// `buf[end..]`, so no read pays for a memset, and live bytes move to the
14/// front only when the tail has no room for the next read.
15#[derive(Debug, Default)]
16pub(crate) struct Window {
17    buf: Vec<u8>,
18    head: usize,
19    end: usize,
20    base: u64,
21    /// Total source length, once known.
22    pub(crate) total: Option<u64>,
23}
24
25impl Window {
26    pub(crate) fn new() -> Self {
27        Self::default()
28    }
29
30    pub(crate) fn base(&self) -> u64 {
31        self.base
32    }
33
34    pub(crate) fn len(&self) -> usize {
35        self.end - self.head
36    }
37
38    pub(crate) fn end(&self) -> u64 {
39        self.base + self.len() as u64
40    }
41
42    pub(crate) fn bytes(&self) -> &[u8] {
43        &self.buf[self.head..self.end]
44    }
45
46    pub(crate) fn capacity(&self) -> usize {
47        self.buf.capacity()
48    }
49
50    pub(crate) fn release(&mut self) {
51        self.buf = Vec::new();
52        self.head = 0;
53        self.end = 0;
54    }
55
56    pub(crate) fn input(&self) -> Input<'_> {
57        Input::new(self.base, self.bytes()).with_total_len_opt(self.total)
58    }
59
60    /// Position at `offset`, keeping the longest held prefix there. A fetched
61    /// range can become the sequential window without copying its allocation.
62    /// Returns whether any bytes were retained; it never fills a hole.
63    pub(crate) fn position(&mut self, offset: u64, fetched: &mut Vec<(u64, Vec<u8>)>) -> bool {
64        let active = offset >= self.base && offset < self.end();
65        let active_end = if active { self.end() } else { offset };
66        if let Some((i, _)) = fetched
67            .iter()
68            .enumerate()
69            .filter(|(_, (start, bytes))| {
70                offset >= *start && offset < start.saturating_add(bytes.len() as u64)
71            })
72            .max_by_key(|(_, (start, bytes))| start.saturating_add(bytes.len() as u64))
73            .filter(|(_, (start, bytes))| start.saturating_add(bytes.len() as u64) > active_end)
74        {
75            let (start, bytes) = fetched.swap_remove(i);
76            self.adopt(start, bytes);
77            self.consume((offset - start) as usize);
78            return true;
79        }
80        if active {
81            self.consume((offset - self.base) as usize);
82            true
83        } else {
84            self.reset_to(offset);
85            false
86        }
87    }
88
89    /// Drop `n` bytes from the front (the decoder took them).
90    pub(crate) fn consume(&mut self, n: usize) {
91        let n = n.min(self.len());
92        self.head += n;
93        self.base += n as u64;
94        if self.head == self.end {
95            self.head = 0;
96            self.end = 0;
97        }
98    }
99
100    /// Forget everything and start at `offset`.
101    pub(crate) fn reset_to(&mut self, offset: u64) {
102        self.head = 0;
103        self.end = 0;
104        self.base = offset;
105    }
106
107    /// Room for `n` more bytes after the held bytes, to read into; follow
108    /// with [`commit`](Self::commit). The contents are stale, not zeroed.
109    pub(crate) fn spare(&mut self, n: usize) -> &mut [u8] {
110        self.try_spare(n, usize::MAX)
111            .expect("input window allocation")
112    }
113
114    pub(crate) fn try_spare(&mut self, n: usize, limit: usize) -> Result<&mut [u8], InputError> {
115        let needed = self.len().checked_add(n).ok_or(InputError::Limit {
116            needed: u64::MAX,
117            limit,
118        })?;
119        if needed > limit {
120            return Err(InputError::Limit {
121                needed: needed as u64,
122                limit,
123            });
124        }
125        if self.buf.len() - self.end < n && self.head > 0 {
126            self.buf.copy_within(self.head..self.end, 0);
127            self.end -= self.head;
128            self.head = 0;
129        }
130        if self.buf.len() - self.end < n {
131            // Grow geometrically so a stream of reads zeroes O(total) bytes.
132            let want = (self.end + n)
133                .max(self.buf.len().saturating_mul(2))
134                .min(limit);
135            reserve(&mut self.buf, want)?;
136            if self.buf.capacity() > limit {
137                let needed = self.buf.capacity() as u64;
138                self.release();
139                return Err(InputError::Limit { needed, limit });
140            }
141            self.buf.resize(want, 0);
142        }
143        Ok(&mut self.buf[self.end..self.end + n])
144    }
145
146    /// Keep `n` of the bytes last handed out by [`spare`](Self::spare).
147    pub(crate) fn commit(&mut self, n: usize) {
148        self.end += n;
149    }
150
151    /// Take ownership of `buf`, holding bytes from `offset`. The old
152    /// contents are dropped; no bytes are copied.
153    #[cfg_attr(not(feature = "tokio"), allow(dead_code))]
154    pub(crate) fn adopt(&mut self, offset: u64, buf: Vec<u8>) {
155        self.end = buf.len();
156        self.buf = buf;
157        self.head = 0;
158        self.base = offset;
159    }
160
161    /// Append bytes received from elsewhere.
162    #[allow(dead_code)]
163    pub(crate) fn extend(&mut self, bytes: &[u8]) {
164        self.spare(bytes.len()).copy_from_slice(bytes);
165        self.commit(bytes.len());
166    }
167}
168
169/// How many bytes to hold before stepping a decoder that asked for
170/// `min_len` contiguous bytes: `min_len` rounded up to a whole multiple of
171/// itself no smaller than `read_size`.
172///
173/// A decoder that consumes `min_len` per step (a strip, a row) then leaves
174/// nothing behind, so the window empties and restarts at its front instead
175/// of moving a leftover tail on every read; one stray header in front of
176/// the strips would otherwise misalign every read for the whole image.
177/// Decoders that ask for little (`min_len` of 1) still get `read_size`.
178pub(crate) fn read_target(min_len: usize, read_size: usize) -> usize {
179    let unit = min_len.max(1);
180    let read_size = read_size.max(unit);
181    read_size.div_ceil(unit).saturating_mul(unit)
182}
183
184/// A decoder broke the [`IncrementalDecode`](zencodec::decode::IncrementalDecode)
185/// protocol. Drivers report this instead of looping.
186#[derive(Debug, Clone, PartialEq, Eq)]
187#[non_exhaustive]
188pub enum ProtocolViolation {
189    /// `Need` with no requests.
190    EmptyRequest,
191    /// A request's offset plus length cannot be represented by the protocol.
192    AddressOverflow(InputRequest),
193    /// The input satisfied one of the decoder's requests, yet it took
194    /// nothing and asked again.
195    NoProgress(InputRequest),
196    /// The decoder asked for bytes past the end it was told about, instead
197    /// of reporting truncated input.
198    NeedPastEnd {
199        /// The request.
200        request: InputRequest,
201        /// The total length the decoder had been given.
202        total: u64,
203    },
204}
205
206impl fmt::Display for ProtocolViolation {
207    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
208        match self {
209            Self::EmptyRequest => write!(f, "decoder returned Need with no requests"),
210            Self::AddressOverflow(r) => write!(f, "decoder request address overflows: {r:?}"),
211            Self::NoProgress(r) => write!(
212                f,
213                "decoder took nothing from an input that satisfied its request {r:?}"
214            ),
215            Self::NeedPastEnd { request, total } => write!(
216                f,
217                "decoder requested {request:?} past the end of a {total}-byte source"
218            ),
219        }
220    }
221}
222
223impl core::error::Error for ProtocolViolation {}
224
225/// Check a `Need` against the input that was just offered: requests must
226/// exist, must not ask past a known end, and an input that satisfied one
227/// must have been consumed from.
228pub(crate) fn check_need(
229    requests: &[InputRequest],
230    offered: &Input<'_>,
231    consumed: usize,
232) -> Result<(), ProtocolViolation> {
233    if requests.is_empty() {
234        return Err(ProtocolViolation::EmptyRequest);
235    }
236    for &request in requests {
237        let span = match request {
238            InputRequest::AtLeast { offset, min_len } => Some((offset, min_len.max(1) as u64)),
239            InputRequest::Range(r) => Some((r.offset, r.length)),
240            _ => None,
241        };
242        if let Some((offset, length)) = span
243            && offset.checked_add(length).is_none()
244        {
245            return Err(ProtocolViolation::AddressOverflow(request));
246        }
247    }
248    if let Some(total) = offered.total_len {
249        for r in requests {
250            let past = match *r {
251                InputRequest::AtLeast { offset, .. } => offset > total,
252                InputRequest::Range(range) => range.end() > total,
253                _ => false,
254            };
255            if past {
256                return Err(ProtocolViolation::NeedPastEnd { request: *r, total });
257            }
258        }
259    }
260    if consumed == 0
261        && let Some(r) = requests.iter().find(|r| r.is_satisfied_by(offered))
262    {
263        return Err(ProtocolViolation::NoProgress(*r));
264    }
265    Ok(())
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271
272    #[test]
273    fn window_consumes_by_head_and_reads_into_spare() {
274        let mut w = Window::new();
275        w.extend(&[1, 2, 3, 4]);
276        w.consume(1);
277        assert_eq!((w.base(), w.bytes()), (1, &[2u8, 3, 4][..]));
278        let s = w.spare(4);
279        s[..2].copy_from_slice(&[5, 6]);
280        w.commit(2);
281        assert_eq!(w.bytes(), &[2, 3, 4, 5, 6]);
282        w.consume(5);
283        assert_eq!((w.base(), w.len()), (6, 0));
284        w.reset_to(100);
285        assert_eq!((w.base(), w.end()), (100, 100));
286    }
287
288    #[test]
289    fn read_target_rounds_to_whole_requests() {
290        assert_eq!(read_target(1, 65536), 65536);
291        assert_eq!(read_target(0, 4096), 4096);
292        assert_eq!(read_target(65536, 65536), 65536);
293        assert_eq!(read_target(20_000, 65536), 80_000);
294        assert_eq!(read_target(100_000, 65536), 100_000);
295    }
296
297    #[test]
298    fn need_checks_catch_spins_and_overreach() {
299        let bytes = [0u8; 8];
300        let full = Input::new(0, &bytes).with_total_len(8);
301        assert_eq!(
302            check_need(&[], &full, 0),
303            Err(ProtocolViolation::EmptyRequest)
304        );
305        let r = InputRequest::at_least(0, 4);
306        assert_eq!(
307            check_need(&[r], &full, 0),
308            Err(ProtocolViolation::NoProgress(r))
309        );
310        assert_eq!(check_need(&[r], &full, 1), Ok(()));
311        let past = InputRequest::range(6, 4);
312        assert!(matches!(
313            check_need(&[past], &full, 0),
314            Err(ProtocolViolation::NeedPastEnd { .. })
315        ));
316        let short = Input::new(0, &bytes[..2]);
317        assert_eq!(check_need(&[r], &short, 0), Ok(()));
318    }
319}
zencodec_io/window.rs

zencodec_io/
window.rs
