Skip to main content

tor_bytes/
reader.rs

1//! Internal: Declare the Reader type for tor-bytes
2
3use tor_error::{bad_api_usage, into_internal};
4
5use crate::{
6    Error::{self},
7    Readable, Result,
8};
9use std::num::NonZeroUsize;
10
11/// A type for reading messages from a slice of bytes.
12///
13/// Unlike io::Read, this object has a simpler error type, and is designed
14/// for in-memory parsing only.
15///
16/// The methods in [`Reader`] should never panic, with one exception:
17/// the `extract` and `extract_n` methods will panic if the underlying
18/// [`Readable`] object's `take_from` method panics.
19///
20/// # Examples
21///
22/// You can use a Reader to extract information byte-by-byte:
23///
24/// ```
25/// use tor_bytes::{Reader,Result};
26/// let msg = [ 0x00, 0x01, 0x23, 0x45, 0x22, 0x00, 0x00, 0x00 ];
27/// let mut b = Reader::from_slice(&msg[..]);
28/// // Multi-byte values are always big-endian.
29/// assert_eq!(b.take_u32()?, 0x12345);
30/// assert_eq!(b.take_u8()?, 0x22);
31///
32/// // You can check on the length of the message...
33/// assert_eq!(b.total_len(), 8);
34/// assert_eq!(b.consumed(), 5);
35/// assert_eq!(b.remaining(), 3);
36/// // then skip over a some bytes...
37/// b.advance(3)?;
38/// // ... and check that the message is really exhausted.
39/// b.should_be_exhausted()?;
40/// # Result::Ok(())
41/// ```
42///
43/// You can also use a Reader to extract objects that implement Readable.
44/// ```
45/// use tor_bytes::{Reader,Result,Readable};
46/// use std::net::Ipv4Addr;
47/// let msg = [ 0x00, 0x04, 0x7f, 0x00, 0x00, 0x01];
48/// let mut b = Reader::from_slice(&msg[..]);
49///
50/// let tp: u16 = b.extract()?;
51/// let ip: Ipv4Addr = b.extract()?;
52/// assert_eq!(tp, 4);
53/// assert_eq!(ip, Ipv4Addr::LOCALHOST);
54/// # Result::Ok(())
55/// ```
56pub struct Reader<'a> {
57    /// The underlying slice that we're reading from
58    b: &'a [u8],
59    /// The next position in the slice that we intend to read from.
60    off: usize,
61    /// What to do if we run out of data - IOW are we reading a possibly incomplete message
62    completeness: Completeness,
63}
64
65/// Whether we're supposed to have the complete message, or not
66///
67/// IOW are we reading a possibly incomplete message?
68///
69/// Affects the error return if we run out of data
70/// ([`Reader::incomplete_error`]).
71#[derive(Copy, Clone, Debug)]
72enum Completeness {
73    /// We might not have the whole message, and that is expected
74    ///
75    /// Throw [`Error::Incomplete`]
76    PossiblyIncomplete,
77    /// We ought to have the whole message
78    ///
79    /// Throw [`Error::MissingData`]
80    SupposedlyComplete,
81}
82
83impl<'a> Reader<'a> {
84    /// Construct a new Reader from a slice of bytes.
85    ///
86    /// In tests, prefer [`Reader::from_slice_for_test`].
87    pub fn from_slice(slice: &'a [u8]) -> Self {
88        Reader {
89            b: slice,
90            off: 0,
91            completeness: Completeness::SupposedlyComplete,
92        }
93    }
94    /// Construct a new Reader from a slice of bytes which may not be complete.
95    ///
96    /// This can be used to try to deserialise a message received from a protocol stream,
97    /// if we don't know how much data we needed to buffer.
98    ///
99    /// [`Readable`] methods, [`extract`](Reader::extract), and so on,
100    /// will return [`Error::Incomplete`] if the message is incomplete,
101    /// and reading more would help.
102    ///
103    /// (This is achieved via [`incomplete_error`](Reader::incomplete_error.)
104    ///
105    /// # Warning about denial of service through excessive memory use
106    ///
107    /// It is hazardous to use this approach unless the buffer size is limited,
108    /// since the sender could send an apparently-very-large message.
109    ///
110    /// # Warning about sub-readers
111    ///
112    /// If you are constructing other readers from data extracted from this one,
113    /// make sure to use [`Reader::from_slice`] instead of this method!
114    /// This method is only for the outermost reader.
115    ///
116    /// Failure to follow this warning may result in malformed messages
117    /// being incorrectly reported as `Incomplete`.
118    //
119    // TODO this name is quite clumsy!
120    pub fn from_possibly_incomplete_slice(slice: &'a [u8]) -> Self {
121        Reader {
122            b: slice,
123            off: 0,
124            completeness: Completeness::PossiblyIncomplete,
125        }
126    }
127    /// Construct a new Reader from a slice of bytes, in tests
128    ///
129    /// This is equivalent to [`Reader::from_possibly_incomplete_slice`].
130    /// It should be used in test cases, because that gives more precise
131    /// testing of the generation of incomplete data errors.
132    pub fn from_slice_for_test(slice: &'a [u8]) -> Self {
133        Self::from_possibly_incomplete_slice(slice)
134    }
135    /// Construct a new Reader from a 'Bytes' object.
136    pub fn from_bytes(b: &'a bytes::Bytes) -> Self {
137        Self::from_slice(b.as_ref())
138    }
139    /// Return the total length of the slice in this reader, including
140    /// consumed bytes and remaining bytes.
141    pub fn total_len(&self) -> usize {
142        self.b.len()
143    }
144    /// Return the total number of bytes in this reader that have not
145    /// yet been read.
146    pub fn remaining(&self) -> usize {
147        self.b.len() - self.off
148    }
149    /// Consume this reader, and return a slice containing the remaining
150    /// bytes from its slice that it did not consume.
151    pub fn into_rest(self) -> &'a [u8] {
152        &self.b[self.off..]
153    }
154    /// Return the total number of bytes in this reader that have
155    /// already been read.
156    pub fn consumed(&self) -> usize {
157        self.off
158    }
159    /// Skip `n` bytes from the reader.
160    ///
161    /// Returns Ok on success.  Throws MissingData or Incomplete if there were
162    /// not enough bytes to skip.
163    pub fn advance(&mut self, n: usize) -> Result<()> {
164        self.peek(n)?;
165        self.off += n;
166        Ok(())
167    }
168    /// Check whether this reader is exhausted (out of bytes).
169    ///
170    /// Return Ok if it is, and Err(Error::ExtraneousBytes)
171    /// if there were extra bytes.
172    pub fn should_be_exhausted(&self) -> Result<()> {
173        if self.remaining() != 0 {
174            return Err(Error::ExtraneousBytes);
175        }
176        Ok(())
177    }
178    /// Truncate this reader, so that no more than `n` bytes remain.
179    ///
180    /// Fewer than `n` bytes may remain if there were not enough bytes
181    /// to begin with.
182    pub fn truncate(&mut self, n: usize) {
183        if n < self.remaining() {
184            self.b = &self.b[..self.off + n];
185        }
186    }
187    /// Try to return a slice of `n` bytes from this reader without
188    /// consuming them.
189    ///
190    /// On success, returns Ok(slice).  If there are fewer than n
191    /// bytes, Throws MissingData or Incomplete if there were
192    /// not enough bytes to skip.
193    pub fn peek(&self, n: usize) -> Result<&'a [u8]> {
194        if let Some(deficit) = n
195            .checked_sub(self.remaining())
196            .and_then(|d| d.try_into().ok())
197        {
198            return Err(self.incomplete_error(deficit));
199        }
200
201        Ok(&self.b[self.off..(n + self.off)])
202    }
203    /// Try to consume and return a slice of `n` bytes from this reader.
204    ///
205    /// On success, returns Ok(Slice).  If there are fewer than n
206    /// bytes, Throws MissingData or Incomplete.
207    ///
208    /// # Example
209    /// ```
210    /// use tor_bytes::{Reader,Result};
211    /// let m = b"Hello World";
212    /// let mut b = Reader::from_slice(m);
213    /// assert_eq!(b.take(5)?, b"Hello");
214    /// assert_eq!(b.take_u8()?, 0x20);
215    /// assert_eq!(b.take(5)?, b"World");
216    /// b.should_be_exhausted()?;
217    /// # Result::Ok(())
218    /// ```
219    pub fn take(&mut self, n: usize) -> Result<&'a [u8]> {
220        let b = self.peek(n)?;
221        self.advance(n)?;
222        Ok(b)
223    }
224    /// Try to fill a provided buffer with bytes consumed from this reader.
225    ///
226    /// On success, the buffer will be filled with data from the
227    /// reader, the reader will advance by the length of the buffer,
228    /// and we'll return Ok(()).  On failure the buffer will be
229    /// unchanged.
230    ///
231    /// # Example
232    /// ```
233    /// use tor_bytes::Reader;
234    /// let m = b"Hello world";
235    /// let mut v1 = vec![0; 5];
236    /// let mut v2 = vec![0; 5];
237    /// let mut b = Reader::from_slice(m);
238    /// b.take_into(&mut v1[..])?;
239    /// assert_eq!(b.take_u8()?, b' ');
240    /// b.take_into(&mut v2[..])?;
241    /// assert_eq!(&v1[..], b"Hello");
242    /// assert_eq!(&v2[..], b"world");
243    /// b.should_be_exhausted()?;
244    /// # tor_bytes::Result::Ok(())
245    /// ```
246    pub fn take_into(&mut self, buf: &mut [u8]) -> Result<()> {
247        let n = buf.len();
248        let b = self.take(n)?;
249        buf.copy_from_slice(b);
250        Ok(())
251    }
252    /// Try to consume and return a u8 from this reader.
253    pub fn take_u8(&mut self) -> Result<u8> {
254        let b = self.take(1)?;
255        Ok(b[0])
256    }
257    /// Try to consume and return a big-endian u16 from this reader.
258    pub fn take_u16(&mut self) -> Result<u16> {
259        let b: [u8; 2] = self.extract()?;
260        let r = u16::from_be_bytes(b);
261        Ok(r)
262    }
263    /// Try to consume and return a big-endian u32 from this reader.
264    pub fn take_u32(&mut self) -> Result<u32> {
265        let b: [u8; 4] = self.extract()?;
266        let r = u32::from_be_bytes(b);
267        Ok(r)
268    }
269    /// Try to consume and return a big-endian u64 from this reader.
270    pub fn take_u64(&mut self) -> Result<u64> {
271        let b: [u8; 8] = self.extract()?;
272        let r = u64::from_be_bytes(b);
273        Ok(r)
274    }
275    /// Try to consume and return a big-endian u128 from this reader.
276    pub fn take_u128(&mut self) -> Result<u128> {
277        let b: [u8; 16] = self.extract()?;
278        let r = u128::from_be_bytes(b);
279        Ok(r)
280    }
281    /// Try to consume and return bytes from this buffer until we
282    /// encounter a terminating byte equal to `term`.
283    ///
284    /// On success, returns Ok(Slice), where the slice does not
285    /// include the terminating byte.  Throws MissingData or Incomplete
286    /// if we do not find the terminating bytes.
287    ///
288    /// Advances the reader to the point immediately after the terminating
289    /// byte.
290    ///
291    /// # Example
292    /// ```
293    /// use tor_bytes::{Reader,Result};
294    /// let m = b"Hello\0wrld";
295    /// let mut b = Reader::from_slice(m);
296    /// assert_eq!(b.take_until(0)?, b"Hello");
297    /// assert_eq!(b.into_rest(), b"wrld");
298    /// # Result::Ok(())
299    /// ```
300    pub fn take_until(&mut self, term: u8) -> Result<&'a [u8]> {
301        self.take_until_with_limit(term, usize::MAX)
302        // Since take_until_with_limit used if max_len > self.remaining()
303        // So when max_len is usize::MAX, it will always be greater than self.remaining(),
304        // and thus we will always return the error from incomplete_error
305    }
306    /// Try to consume and return bytes from this buffer until we
307    /// encounter a terminating byte equal to `term`.
308    /// Or we have read `max_len` bytes without finding the terminator.
309    /// The maximum value that will be returned is max_len - 1,
310    /// because we need to remove the terminator.
311    /// If we didn't find the terminator within the first `max_len` bytes,
312    /// we will return an error.
313    ///     
314    /// # Example
315    /// ```rust
316    /// use tor_bytes::{Reader,Result};
317    /// let m = b"Hello\0wrld";
318    /// let mut b = Reader::from_slice(m);
319    /// assert_eq!(b.take_until_with_limit(0, 10)?, b"Hello");
320    /// assert_eq!(b.into_rest(), b"wrld");
321    /// # Result::Ok(())
322    /// ```
323    pub fn take_until_with_limit(&mut self, term: u8, max_len: usize) -> Result<&'a [u8]> {
324        let limit: usize = std::cmp::min(max_len, self.remaining());
325        let pos = match self.b[self.off..limit + self.off]
326            .iter()
327            .position(|b| *b == term)
328        {
329            Some(p) => p,
330            None => {
331                if max_len > self.remaining() {
332                    // since user asked for more than we have, we should return Incomplete, not LimitExceeded
333                    return Err(self.incomplete_error(
334                        //
335                        1.try_into().expect("1 == 0"),
336                    ));
337                } else {
338                    return Err(Error::LimitExceeded {
339                        limit,
340                        terminator: term,
341                    });
342                }
343            }
344        };
345
346        let result: &[u8] = self.take(pos)?;
347        self.advance(1)?;
348        Ok(result)
349    }
350    /// Consume and return all the remaining bytes, but do not consume the reader
351    ///
352    /// This can be useful if you need to possibly read either fixed-length data,
353    /// or variable length data eating the rest of the `Reader`.
354    ///
355    /// The `Reader` will be left devoid of further bytes.
356    /// Consider using `into_rest()` instead.
357    pub fn take_rest(&mut self) -> &'a [u8] {
358        self.take(self.remaining())
359            .expect("taking remaining failed")
360    }
361
362    /// Consume and return all but the last `n` remaining bytes.
363    ///
364    /// Gives `Error::MissingData` if there are fewer than `n` remaining bytes.
365    ///
366    /// It is invalid to call this method on a `Reader` constructed with
367    /// [`Reader::from_possibly_incomplete_slice`].  (If we don't know where the
368    /// data actually ends, we can't take all but the last `n` bytes.)
369    /// Such calls cause an internal error.
370    ///
371    /// # Example
372    /// ```
373    /// use tor_bytes::{Reader,Result};
374    /// let m = b"Hello World";
375    /// let mut b = Reader::from_slice(m);
376    /// assert_eq!(b.take_all_but(2)?, b"Hello Wor");
377    /// assert_eq!(b.into_rest(), b"ld");
378    /// # Result::Ok(())
379    /// ```
380    pub fn take_all_but(&mut self, n: usize) -> Result<&'a [u8]> {
381        match self.completeness {
382            Completeness::PossiblyIncomplete => {
383                return Err(Error::Bug(bad_api_usage!(
384                    "Called take_all_but on a PossiblyIncomplete reader."
385                )));
386            }
387            Completeness::SupposedlyComplete => {}
388        }
389
390        let n_to_take = self.remaining().checked_sub(n).ok_or(Error::MissingData)?;
391
392        let result = self
393            .take(n_to_take)
394            .map_err(into_internal!("Subtraction misled us somehow"))?;
395        debug_assert_eq!(self.remaining(), n);
396        Ok(result)
397    }
398
399    /// Try to decode and remove a Readable from this reader, using its
400    /// take_from() method.
401    ///
402    /// On failure, consumes nothing.
403    pub fn extract<E: Readable>(&mut self) -> Result<E> {
404        let off_orig = self.off;
405        let result = E::take_from(self);
406        if result.is_err() {
407            // We encountered an error; we should rewind.
408            self.off = off_orig;
409        }
410        result
411    }
412
413    /// Try to decode and remove `n` Readables from this reader, using the
414    /// Readable's take_from() method.
415    ///
416    /// On failure, consumes nothing.
417    pub fn extract_n<E: Readable>(&mut self, n: usize) -> Result<Vec<E>> {
418        // This `min` will help us defend against a pathological case where an
419        // attacker tells us that there are BIGNUM elements forthcoming, and our
420        // attempt to allocate `Vec::with_capacity(BIGNUM)` makes us panic.
421        //
422        // The `min` can be incorrect if E is somehow encodable in zero bytes
423        // (!?), but that will only cause our initial allocation to be too
424        // small.
425        //
426        // In practice, callers should always check that `n` is reasonable
427        // before calling this function, and protocol designers should not
428        // provide e.g. 32-bit counters for object types of which we should
429        // never allocate u32::MAX.
430        let n_alloc = std::cmp::min(n, self.remaining());
431        let mut result = Vec::with_capacity(n_alloc);
432        let off_orig = self.off;
433        for _ in 0..n {
434            match E::take_from(self) {
435                Ok(item) => result.push(item),
436                Err(e) => {
437                    // Encountered an error; we should rewind.
438                    self.off = off_orig;
439                    return Err(e);
440                }
441            }
442        }
443        Ok(result)
444    }
445
446    /// Decode something with a `u8` length field
447    ///
448    /// Prefer to use this function, rather than ad-hoc `take_u8`
449    /// and subsequent manual length checks.
450    /// Using this facility eliminates the need to separately keep track of the lengths.
451    ///
452    /// `read_nested` consumes a length field,
453    /// and provides the closure `f` with an inner `Reader` that
454    /// contains precisely that many bytes -
455    /// the bytes which follow the length field in the original reader.
456    /// If the closure is successful, `read_nested` checks that that inner reader is exhausted,
457    /// i.e. that the inner contents had the same length as was specified.
458    ///
459    /// The closure should read whatever is inside the nested structure
460    /// from the nested reader.
461    /// It may well want to use `take_rest`, to consume all of the counted bytes.
462    ///
463    /// On failure, the amount consumed is not specified.
464    pub fn read_nested_u8len<F, T>(&mut self, f: F) -> Result<T>
465    where
466        F: FnOnce(&mut Reader) -> Result<T>,
467    {
468        read_nested_generic::<u8, _, _>(self, f)
469    }
470
471    /// Start decoding something with a u16 length field
472    pub fn read_nested_u16len<F, T>(&mut self, f: F) -> Result<T>
473    where
474        F: FnOnce(&mut Reader) -> Result<T>,
475    {
476        read_nested_generic::<u16, _, _>(self, f)
477    }
478
479    /// Start decoding something with a u32 length field
480    pub fn read_nested_u32len<F, T>(&mut self, f: F) -> Result<T>
481    where
482        F: FnOnce(&mut Reader) -> Result<T>,
483    {
484        read_nested_generic::<u32, _, _>(self, f)
485    }
486
487    /// Return a cursor object describing the current position of this Reader
488    /// within its underlying byte stream.
489    ///
490    /// The resulting [`Cursor`] can be used with `range`, but nothing else.
491    ///
492    /// Note that having to use a `Cursor` is typically an anti-pattern: it
493    /// tends to indicate that whatever you're parsing could probably have a
494    /// better design that would better separate data from metadata.
495    /// Unfortunately, there are a few places like that in the Tor  protocols.
496    //
497    // TODO: This could instead be a function that takes a closure, passes a
498    // reader to that closure, and returns the closure's output along with
499    // whatever the reader consumed.
500    pub fn cursor(&self) -> Cursor<'a> {
501        Cursor {
502            pos: self.off,
503            _phantom: std::marker::PhantomData,
504        }
505    }
506
507    /// Return the slice of bytes between the start cursor (inclusive) and end
508    /// cursor (exclusive).
509    ///
510    /// If the cursors are not in order, return an empty slice.
511    ///
512    /// This function is guaranteed not to panic if the inputs were generated
513    /// from a different Reader, but if so the byte slice that it returns will
514    /// not be meaningful.
515    pub fn range(&self, start: Cursor<'a>, end: Cursor<'a>) -> &'a [u8] {
516        if start.pos <= end.pos && end.pos <= self.b.len() {
517            &self.b[start.pos..end.pos]
518        } else {
519            &self.b[..0]
520        }
521    }
522
523    /// Returns the error that should be returned if we ran out of data
524    ///
525    /// For a usual `Reader` this is [`Error::MissingData`].
526    /// For a reader from
527    /// [`Reader::from_possibly_incomplete_slice`]
528    /// it's [`Error::Incomplete`].
529    pub fn incomplete_error(&self, deficit: NonZeroUsize) -> Error {
530        use Completeness as C;
531        use Error as E;
532        match self.completeness {
533            C::PossiblyIncomplete => E::Incomplete {
534                deficit: deficit.into(),
535            },
536            C::SupposedlyComplete => E::MissingData,
537        }
538    }
539}
540
541/// A reference to a position within a [`Reader`].
542#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
543pub struct Cursor<'a> {
544    /// The underlying position within the reader.
545    pos: usize,
546    /// Used so that we can restrict the cursor to the lifetime of the
547    /// underlying byte slice.
548    _phantom: std::marker::PhantomData<&'a [u8]>,
549}
550
551/// Implementation of `read_nested_*` -- generic
552fn read_nested_generic<L, F, T>(b: &mut Reader, f: F) -> Result<T>
553where
554    F: FnOnce(&mut Reader) -> Result<T>,
555    L: Readable + Copy + Sized + TryInto<usize>,
556{
557    let length: L = b.extract()?;
558    let length: usize = length.try_into().map_err(|_| Error::BadLengthValue)?;
559    let slice = b.take(length)?;
560    let mut inner = Reader::from_slice(slice);
561    let out = f(&mut inner)?;
562    inner.should_be_exhausted()?;
563    Ok(out)
564}
565
566#[cfg(test)]
567mod tests {
568    #![allow(clippy::unwrap_used)]
569    use super::*;
570    #[test]
571    fn bytecursor_read_ok() {
572        let bytes = b"On a mountain halfway between Reno and Rome";
573        let mut bc = Reader::from_slice(&bytes[..]);
574
575        assert_eq!(bc.consumed(), 0);
576        assert_eq!(bc.remaining(), 43);
577        assert_eq!(bc.total_len(), 43);
578
579        assert_eq!(bc.take(3).unwrap(), &b"On "[..]);
580        assert_eq!(bc.consumed(), 3);
581
582        assert_eq!(bc.take_u16().unwrap(), 0x6120);
583        assert_eq!(bc.take_u8().unwrap(), 0x6d);
584        assert_eq!(bc.take_u64().unwrap(), 0x6f756e7461696e20);
585        assert_eq!(bc.take_u32().unwrap(), 0x68616c66);
586        assert_eq!(bc.consumed(), 18);
587        assert_eq!(bc.remaining(), 25);
588        assert_eq!(bc.total_len(), 43);
589
590        assert_eq!(bc.peek(7).unwrap(), &b"way bet"[..]);
591        assert_eq!(bc.consumed(), 18); // no change
592        assert_eq!(bc.remaining(), 25); // no change
593        assert_eq!(bc.total_len(), 43); // no change
594
595        assert_eq!(bc.peek(7).unwrap(), &b"way bet"[..]);
596        assert_eq!(bc.consumed(), 18); // no change this time either.
597
598        bc.advance(12).unwrap();
599        assert_eq!(bc.consumed(), 30);
600        assert_eq!(bc.remaining(), 13);
601
602        let rem = bc.into_rest();
603        assert_eq!(rem, &b"Reno and Rome"[..]);
604
605        // now let's try consuming right up to the end.
606        let mut bc = Reader::from_slice(&bytes[..]);
607        bc.advance(22).unwrap();
608        assert_eq!(bc.remaining(), 21);
609        let rem = bc.take(21).unwrap();
610        assert_eq!(rem, &b"between Reno and Rome"[..]);
611        assert_eq!(bc.consumed(), 43);
612        assert_eq!(bc.remaining(), 0);
613
614        // We can still take a zero-length slice.
615        assert_eq!(bc.take(0).unwrap(), &b""[..]);
616    }
617
618    #[test]
619    fn read_u128() {
620        let bytes = bytes::Bytes::from(&b"irreproducibility?"[..]); // 18 bytes
621        let mut b = Reader::from_bytes(&bytes);
622
623        assert_eq!(b.take_u8().unwrap(), b'i');
624        assert_eq!(b.take_u128().unwrap(), 0x72726570726f6475636962696c697479);
625        assert_eq!(b.remaining(), 1);
626    }
627
628    #[test]
629    fn bytecursor_read_missing() {
630        let bytes = b"1234567";
631        let mut bc = Reader::from_slice_for_test(&bytes[..]);
632
633        assert_eq!(bc.consumed(), 0);
634        assert_eq!(bc.remaining(), 7);
635        assert_eq!(bc.total_len(), 7);
636
637        assert_eq!(bc.take_u64(), Err(Error::new_incomplete_for_test(1)));
638        assert_eq!(bc.take(8), Err(Error::new_incomplete_for_test(1)));
639        assert_eq!(bc.peek(8), Err(Error::new_incomplete_for_test(1)));
640
641        assert_eq!(bc.consumed(), 0);
642        assert_eq!(bc.remaining(), 7);
643        assert_eq!(bc.total_len(), 7);
644
645        assert_eq!(bc.take_u32().unwrap(), 0x31323334); // get 4 bytes. 3 left.
646        assert_eq!(bc.take_u32(), Err(Error::new_incomplete_for_test(1)));
647
648        assert_eq!(bc.consumed(), 4);
649        assert_eq!(bc.remaining(), 3);
650        assert_eq!(bc.total_len(), 7);
651
652        assert_eq!(bc.take_u16().unwrap(), 0x3536); // get 2 bytes. 1 left.
653        assert_eq!(bc.take_u16(), Err(Error::new_incomplete_for_test(1)));
654
655        assert_eq!(bc.consumed(), 6);
656        assert_eq!(bc.remaining(), 1);
657        assert_eq!(bc.total_len(), 7);
658
659        assert_eq!(bc.take_u8().unwrap(), 0x37); // get 1 byte. 0 left.
660        assert_eq!(bc.take_u8(), Err(Error::new_incomplete_for_test(1)));
661
662        assert_eq!(bc.consumed(), 7);
663        assert_eq!(bc.remaining(), 0);
664        assert_eq!(bc.total_len(), 7);
665    }
666
667    #[test]
668    fn advance_too_far() {
669        let bytes = b"12345";
670        let mut b = Reader::from_slice_for_test(&bytes[..]);
671        assert_eq!(b.remaining(), 5);
672        assert_eq!(b.advance(16), Err(Error::new_incomplete_for_test(11)));
673        assert_eq!(b.remaining(), 5);
674        assert_eq!(b.advance(5), Ok(()));
675        assert_eq!(b.remaining(), 0);
676    }
677
678    #[test]
679    fn truncate() {
680        let bytes = b"Hello universe!!!1!";
681        let mut b = Reader::from_slice_for_test(&bytes[..]);
682
683        assert_eq!(b.take(5).unwrap(), &b"Hello"[..]);
684        assert_eq!(b.remaining(), 14);
685        assert_eq!(b.consumed(), 5);
686        b.truncate(9);
687        assert_eq!(b.remaining(), 9);
688        assert_eq!(b.consumed(), 5);
689        assert_eq!(b.take_u8().unwrap(), 0x20);
690        assert_eq!(b.into_rest(), &b"universe"[..]);
691    }
692
693    #[test]
694    fn exhaust() {
695        let b = Reader::from_slice_for_test(&b""[..]);
696        assert_eq!(b.should_be_exhausted(), Ok(()));
697
698        let mut b = Reader::from_slice_for_test(&b"outis"[..]);
699        assert_eq!(b.should_be_exhausted(), Err(Error::ExtraneousBytes));
700        b.take(4).unwrap();
701        assert_eq!(b.should_be_exhausted(), Err(Error::ExtraneousBytes));
702        b.take(1).unwrap();
703        assert_eq!(b.should_be_exhausted(), Ok(()));
704    }
705
706    #[test]
707    fn take_rest() {
708        let mut b = Reader::from_slice_for_test(b"si vales valeo");
709        assert_eq!(b.take(3).unwrap(), b"si ");
710        assert_eq!(b.take_rest(), b"vales valeo");
711        assert_eq!(b.take_rest(), b"");
712    }
713
714    #[test]
715    fn take_until() {
716        let mut b = Reader::from_slice_for_test(&b"si vales valeo"[..]);
717        assert_eq!(b.take_until(b' ').unwrap(), &b"si"[..]);
718        assert_eq!(b.take_until(b' ').unwrap(), &b"vales"[..]);
719        assert_eq!(b.take_until(b' '), Err(Error::new_incomplete_for_test(1)));
720    }
721
722    #[test]
723    fn take_until_with_limit() {
724        let mut b = Reader::from_slice_for_test(&b"si vales valeo"[..]);
725        assert_eq!(b.take_until_with_limit(b' ', 10).unwrap(), &b"si"[..]);
726        assert_eq!(b.take_until_with_limit(b' ', 10).unwrap(), &b"vales"[..]);
727        assert_eq!(
728            b.take_until_with_limit(b' ', 100),
729            Err(Error::new_incomplete_for_test(1)),
730        );
731        let mut b = Reader::from_slice_for_test(&b"Hello\0World"[..]);
732        assert_eq!(
733            b.take_until_with_limit(b'\0', 1),
734            Err(Error::LimitExceeded {
735                limit: 1,
736                terminator: b'\0'
737            })
738        );
739
740        // Test the case where the terminator is exactly at the max_len
741        let mut b = Reader::from_slice_for_test(&b"si vales valeo"[..]);
742        assert_eq!(b.take_until_with_limit(b' ', 3).unwrap(), &b"si"[..]);
743        // Test the case where the terminator is exactly one more than the limit
744        assert_eq!(
745            b.take_until_with_limit(b' ', 5),
746            Err(Error::LimitExceeded {
747                limit: 5,
748                terminator: b' '
749            })
750        );
751    }
752
753    #[test]
754    fn truncate_badly() {
755        let mut b = Reader::from_slice_for_test(&b"abcdefg"[..]);
756        b.truncate(1000);
757        assert_eq!(b.total_len(), 7);
758        assert_eq!(b.remaining(), 7);
759    }
760
761    #[test]
762    fn nested_good() {
763        let mut b = Reader::from_slice_for_test(b"abc\0\0\x04defghijkl");
764        assert_eq!(b.take(3).unwrap(), b"abc");
765
766        b.read_nested_u16len(|s| {
767            assert!(s.should_be_exhausted().is_ok());
768            Ok(())
769        })
770        .unwrap();
771
772        b.read_nested_u8len(|s| {
773            assert_eq!(s.take(4).unwrap(), b"defg");
774            assert!(s.should_be_exhausted().is_ok());
775            Ok(())
776        })
777        .unwrap();
778
779        assert_eq!(b.take(2).unwrap(), b"hi");
780    }
781
782    #[test]
783    fn nested_bad() {
784        let mut b = Reader::from_slice_for_test(b"................");
785        assert_eq!(
786            read_nested_generic::<u128, _, ()>(&mut b, |_| panic!())
787                .err()
788                .unwrap(),
789            Error::BadLengthValue
790        );
791
792        let mut b = Reader::from_slice_for_test(b"................");
793        assert_eq!(
794            b.read_nested_u32len::<_, ()>(|_| panic!()).err().unwrap(),
795            Error::new_incomplete_for_test(774778414 - (16 - 4))
796        );
797    }
798
799    #[test]
800    fn nested_inner_bad() {
801        let mut b = Reader::from_slice_for_test(&[1, 66]);
802        assert_eq!(
803            b.read_nested_u8len(|b| b.take_u32()),
804            Err(Error::MissingData),
805        );
806    }
807
808    #[test]
809    fn incomplete_slice() {
810        // Test specifically the from_possibly_incomplete_slice constructor -
811        // ie, deliberately don't use Reader::from_slice_for_test.
812        let mut b = Reader::from_possibly_incomplete_slice(&[]);
813        assert_eq!(b.take_u32(), Err(Error::new_incomplete_for_test(4)));
814    }
815
816    #[test]
817    fn extract() {
818        // For example purposes, declare a length-then-bytes string type.
819        #[derive(Debug)]
820        struct LenEnc(Vec<u8>);
821        impl Readable for LenEnc {
822            fn take_from(b: &mut Reader<'_>) -> Result<Self> {
823                let length = b.take_u8()?;
824                let content = b.take(length as usize)?.into();
825                Ok(LenEnc(content))
826            }
827        }
828
829        let bytes = b"\x04this\x02is\x09sometimes\x01a\x06string!";
830        let mut b = Reader::from_slice_for_test(&bytes[..]);
831
832        let le: LenEnc = b.extract().unwrap();
833        assert_eq!(&le.0[..], &b"this"[..]);
834
835        let les: Vec<LenEnc> = b.extract_n(4).unwrap();
836        assert_eq!(&les[3].0[..], &b"string"[..]);
837
838        assert_eq!(b.remaining(), 1);
839
840        // Make sure that we don't advance on a failing extract().
841        let le: Result<LenEnc> = b.extract();
842        assert_eq!(le.unwrap_err(), Error::new_incomplete_for_test(33));
843        assert_eq!(b.remaining(), 1);
844
845        // Make sure that we don't advance on a failing extract_n()
846        let mut b = Reader::from_slice_for_test(&bytes[..]);
847        assert_eq!(b.remaining(), 28);
848        let les: Result<Vec<LenEnc>> = b.extract_n(10);
849        assert_eq!(les.unwrap_err(), Error::new_incomplete_for_test(33));
850        assert_eq!(b.remaining(), 28);
851    }
852
853    #[test]
854    fn cursor() -> Result<()> {
855        let alphabet = b"abcdefghijklmnopqrstuvwxyz";
856        let mut b = Reader::from_slice_for_test(&alphabet[..]);
857
858        let c1 = b.cursor();
859        let _ = b.take_u16()?;
860        let c2 = b.cursor();
861        let c2b = b.cursor();
862        b.advance(7)?;
863        let c3 = b.cursor();
864
865        assert_eq!(b.range(c1, c2), &b"ab"[..]);
866        assert_eq!(b.range(c2, c3), &b"cdefghi"[..]);
867        assert_eq!(b.range(c1, c3), &b"abcdefghi"[..]);
868        assert_eq!(b.range(c1, c1), &b""[..]);
869        assert_eq!(b.range(c3, c1), &b""[..]);
870        assert_eq!(c2, c2b);
871        assert!(c1 < c2);
872        assert!(c2 < c3);
873
874        Ok(())
875    }
876
877    #[test]
878    fn take_all_but() -> Result<()> {
879        let message = b"byte manipulation for fun and (non)-profit";
880
881        // Case 1: Successful, complete reader
882        // (Can't use from_slice_for_test here: that's a possibly-incomplete reader.)
883        let mut b = Reader::from_slice(message);
884        assert_eq!(b.take_all_but(6)?, b"byte manipulation for fun and (non)-");
885        assert_eq!(b.into_rest(), b"profit");
886
887        // Case 1b: Successful, take nothing, complete reader.
888        let mut b = Reader::from_slice(message);
889        assert_eq!(b.take_all_but(message.len())?, b"");
890        assert_eq!(b.into_rest(), message);
891
892        // Case 1c: Successful, take everything, complete reader.
893        let mut b = Reader::from_slice(message);
894        assert_eq!(b.take_all_but(0)?, message);
895        assert_eq!(b.into_rest(), b"");
896
897        // Case 2: Unsuccessful, complete reader
898        let mut b = Reader::from_slice(message);
899        assert!(matches!(
900            b.take_all_but(message.len() + 1),
901            Err(Error::MissingData)
902        ));
903
904        // Case 3: Anything, incomplete reader.
905        let mut b = Reader::from_possibly_incomplete_slice(message);
906        assert!(matches!(b.take_all_but(6), Err(Error::Bug(_))));
907
908        Ok(())
909    }
910}