1
//! Internal: Declare the Reader type for tor-bytes
2

            
3
use tor_error::{bad_api_usage, into_internal};
4

            
5
use crate::{Error, Readable, Result};
6
use std::num::NonZeroUsize;
7

            
8
/// A type for reading messages from a slice of bytes.
9
///
10
/// Unlike io::Read, this object has a simpler error type, and is designed
11
/// for in-memory parsing only.
12
///
13
/// The methods in [`Reader`] should never panic, with one exception:
14
/// the `extract` and `extract_n` methods will panic if the underlying
15
/// [`Readable`] object's `take_from` method panics.
16
///
17
/// # Examples
18
///
19
/// You can use a Reader to extract information byte-by-byte:
20
///
21
/// ```
22
/// use tor_bytes::{Reader,Result};
23
/// let msg = [ 0x00, 0x01, 0x23, 0x45, 0x22, 0x00, 0x00, 0x00 ];
24
/// let mut b = Reader::from_slice(&msg[..]);
25
/// // Multi-byte values are always big-endian.
26
/// assert_eq!(b.take_u32()?, 0x12345);
27
/// assert_eq!(b.take_u8()?, 0x22);
28
///
29
/// // You can check on the length of the message...
30
/// assert_eq!(b.total_len(), 8);
31
/// assert_eq!(b.consumed(), 5);
32
/// assert_eq!(b.remaining(), 3);
33
/// // then skip over a some bytes...
34
/// b.advance(3)?;
35
/// // ... and check that the message is really exhausted.
36
/// b.should_be_exhausted()?;
37
/// # Result::Ok(())
38
/// ```
39
///
40
/// You can also use a Reader to extract objects that implement Readable.
41
/// ```
42
/// use tor_bytes::{Reader,Result,Readable};
43
/// use std::net::Ipv4Addr;
44
/// let msg = [ 0x00, 0x04, 0x7f, 0x00, 0x00, 0x01];
45
/// let mut b = Reader::from_slice(&msg[..]);
46
///
47
/// let tp: u16 = b.extract()?;
48
/// let ip: Ipv4Addr = b.extract()?;
49
/// assert_eq!(tp, 4);
50
/// assert_eq!(ip, Ipv4Addr::LOCALHOST);
51
/// # Result::Ok(())
52
/// ```
53
pub struct Reader<'a> {
54
    /// The underlying slice that we're reading from
55
    b: &'a [u8],
56
    /// The next position in the slice that we intend to read from.
57
    off: usize,
58
    /// What to do if we run out of data - IOW are we reading a possibly incomplete message
59
    completeness: Completeness,
60
}
61

            
62
/// Whether we're supposed to have the complete message, or not
63
///
64
/// IOW are we reading a possibly incomplete message?
65
///
66
/// Affects the error return if we run out of data
67
/// ([`Reader::incomplete_error`]).
68
#[derive(Copy, Clone, Debug)]
69
enum Completeness {
70
    /// We might not have the whole message, and that is expected
71
    ///
72
    /// Throw [`Error::Incomplete`]
73
    PossiblyIncomplete,
74
    /// We ought to have the whole message
75
    ///
76
    /// Throw [`Error::MissingData']
77
    SupposedlyComplete,
78
}
79

            
80
impl<'a> Reader<'a> {
81
    /// Construct a new Reader from a slice of bytes.
82
    ///
83
    /// In tests, prefer [`Reader::from_slice_for_test`].
84
203612
    pub fn from_slice(slice: &'a [u8]) -> Self {
85
203612
        Reader {
86
203612
            b: slice,
87
203612
            off: 0,
88
203612
            completeness: Completeness::SupposedlyComplete,
89
203612
        }
90
203612
    }
91
    /// Construct a new Reader from a slice of bytes which may not be complete.
92
    ///
93
    /// This can be used to try to deserialise a message received from a protocol stream,
94
    /// if we don't know how much data we needed to buffer.
95
    ///
96
    /// [`Readable`] methods, [`extract`](Reader::extract), and so on,
97
    /// will return [`Error::Incomplete`] if the message is incomplete,
98
    /// and reading more would help.
99
    ///
100
    /// (This is achieved via [`incomplete_error`](Reader::incomplete_error.)
101
    ///
102
    /// # Warning about denial of service through excessive memory use
103
    ///
104
    /// It is hazardous to use this approach unless the buffer size is limited,
105
    /// since the sender could send an apparently-very-large message.
106
    ///
107
    /// # Warning about sub-readers
108
    ///
109
    /// If you are constructing other readers from data extracted from this one,
110
    /// make sure to use [`Reader::from_slice`] instead of this method!
111
    /// This method is only for the outermost reader.
112
    ///
113
    /// Failure to follow this warning may result in malformed messages
114
    /// being incorrectly reported as `Incomplete`.
115
    //
116
    // TODO this name is quite clumsy!
117
74788
    pub fn from_possibly_incomplete_slice(slice: &'a [u8]) -> Self {
118
74788
        Reader {
119
74788
            b: slice,
120
74788
            off: 0,
121
74788
            completeness: Completeness::PossiblyIncomplete,
122
74788
        }
123
74788
    }
124
    /// Construct a new Reader from a slice of bytes, in tests
125
    ///
126
    /// This is equivalent to [`Reader::from_possibly_incomplete_slice`].
127
    /// It should be used in test cases, because that gives more precise
128
    /// testing of the generation of incomplete data errors.
129
6036
    pub fn from_slice_for_test(slice: &'a [u8]) -> Self {
130
6036
        Self::from_possibly_incomplete_slice(slice)
131
6036
    }
132
    /// Construct a new Reader from a 'Bytes' object.
133
2722
    pub fn from_bytes(b: &'a bytes::Bytes) -> Self {
134
2722
        Self::from_slice(b.as_ref())
135
2722
    }
136
    /// Return the total length of the slice in this reader, including
137
    /// consumed bytes and remaining bytes.
138
18
    pub fn total_len(&self) -> usize {
139
18
        self.b.len()
140
18
    }
141
    /// Return the total number of bytes in this reader that have not
142
    /// yet been read.
143
2825894
    pub fn remaining(&self) -> usize {
144
2825894
        self.b.len() - self.off
145
2825894
    }
146
    /// Consume this reader, and return a slice containing the remaining
147
    /// bytes from its slice that it did not consume.
148
1234
    pub fn into_rest(self) -> &'a [u8] {
149
1234
        &self.b[self.off..]
150
1234
    }
151
    /// Return the total number of bytes in this reader that have
152
    /// already been read.
153
90536
    pub fn consumed(&self) -> usize {
154
90536
        self.off
155
90536
    }
156
    /// Skip `n` bytes from the reader.
157
    ///
158
    /// Returns Ok on success.  Throws MissingData or Incomplete if there were
159
    /// not enough bytes to skip.
160
1217894
    pub fn advance(&mut self, n: usize) -> Result<()> {
161
1217894
        self.peek(n)?;
162
1217892
        self.off += n;
163
1217892
        Ok(())
164
1217894
    }
165
    /// Check whether this reader is exhausted (out of bytes).
166
    ///
167
    /// Return Ok if it is, and Err(Error::ExtraneousBytes)
168
    /// if there were extra bytes.
169
92446
    pub fn should_be_exhausted(&self) -> Result<()> {
170
92446
        if self.remaining() != 0 {
171
140
            return Err(Error::ExtraneousBytes);
172
92306
        }
173
92306
        Ok(())
174
92446
    }
175
    /// Truncate this reader, so that no more than `n` bytes remain.
176
    ///
177
    /// Fewer than `n` bytes may remain if there were not enough bytes
178
    /// to begin with.
179
98196
    pub fn truncate(&mut self, n: usize) {
180
98196
        if n < self.remaining() {
181
11018
            self.b = &self.b[..self.off + n];
182
87178
        }
183
98196
    }
184
    /// Try to return a slice of `n` bytes from this reader without
185
    /// consuming them.
186
    ///
187
    /// On success, returns Ok(slice).  If there are fewer than n
188
    /// bytes, Throws MissingData or Incomplete if there were
189
    /// not enough bytes to skip.
190
2309730
    pub fn peek(&self, n: usize) -> Result<&'a [u8]> {
191
2309730
        if let Some(deficit) = n
192
2309730
            .checked_sub(self.remaining())
193
2317060
            .and_then(|d| d.try_into().ok())
194
        {
195
53474
            return Err(self.incomplete_error(deficit));
196
2256256
        }
197
2256256

            
198
2256256
        Ok(&self.b[self.off..(n + self.off)])
199
2309730
    }
200
    /// Try to consume and return a slice of `n` bytes from this reader.
201
    ///
202
    /// On success, returns Ok(Slice).  If there are fewer than n
203
    /// bytes, Throws MissingData or Incomplete.
204
    ///
205
    /// # Example
206
    /// ```
207
    /// use tor_bytes::{Reader,Result};
208
    /// let m = b"Hello World";
209
    /// let mut b = Reader::from_slice(m);
210
    /// assert_eq!(b.take(5)?, b"Hello");
211
    /// assert_eq!(b.take_u8()?, 0x20);
212
    /// assert_eq!(b.take(5)?, b"World");
213
    /// b.should_be_exhausted()?;
214
    /// # Result::Ok(())
215
    /// ```
216
1056062
    pub fn take(&mut self, n: usize) -> Result<&'a [u8]> {
217
1056062
        let b = self.peek(n)?;
218
1006468
        self.advance(n)?;
219
1006468
        Ok(b)
220
1056062
    }
221
    /// Try to fill a provided buffer with bytes consumed from this reader.
222
    ///
223
    /// On success, the buffer will be filled with data from the
224
    /// reader, the reader will advance by the length of the buffer,
225
    /// and we'll return Ok(()).  On failure the buffer will be
226
    /// unchanged.
227
    ///
228
    /// # Example
229
    /// ```
230
    /// use tor_bytes::Reader;
231
    /// let m = b"Hello world";
232
    /// let mut v1 = vec![0; 5];
233
    /// let mut v2 = vec![0; 5];
234
    /// let mut b = Reader::from_slice(m);
235
    /// b.take_into(&mut v1[..])?;
236
    /// assert_eq!(b.take_u8()?, b' ');
237
    /// b.take_into(&mut v2[..])?;
238
    /// assert_eq!(&v1[..], b"Hello");
239
    /// assert_eq!(&v2[..], b"world");
240
    /// b.should_be_exhausted()?;
241
    /// # tor_bytes::Result::Ok(())
242
    /// ```
243
432052
    pub fn take_into(&mut self, buf: &mut [u8]) -> Result<()> {
244
432052
        let n = buf.len();
245
432052
        let b = self.take(n)?;
246
419120
        buf.copy_from_slice(b);
247
419120
        Ok(())
248
432052
    }
249
    /// Try to consume and return a u8 from this reader.
250
433878
    pub fn take_u8(&mut self) -> Result<u8> {
251
433878
        let b = self.take(1)?;
252
403888
        Ok(b[0])
253
433878
    }
254
    /// Try to consume and return a big-endian u16 from this reader.
255
243450
    pub fn take_u16(&mut self) -> Result<u16> {
256
243450
        let b: [u8; 2] = self.extract()?;
257
239640
        let r = u16::from_be_bytes(b);
258
239640
        Ok(r)
259
243450
    }
260
    /// Try to consume and return a big-endian u32 from this reader.
261
51218
    pub fn take_u32(&mut self) -> Result<u32> {
262
51218
        let b: [u8; 4] = self.extract()?;
263
45704
        let r = u32::from_be_bytes(b);
264
45704
        Ok(r)
265
51218
    }
266
    /// Try to consume and return a big-endian u64 from this reader.
267
71270
    pub fn take_u64(&mut self) -> Result<u64> {
268
71270
        let b: [u8; 8] = self.extract()?;
269
71268
        let r = u64::from_be_bytes(b);
270
71268
        Ok(r)
271
71270
    }
272
    /// Try to consume and return a big-endian u128 from this reader.
273
5174
    pub fn take_u128(&mut self) -> Result<u128> {
274
5174
        let b: [u8; 16] = self.extract()?;
275
1706
        let r = u128::from_be_bytes(b);
276
1706
        Ok(r)
277
5174
    }
278
    /// Try to consume and return bytes from this buffer until we
279
    /// encounter a terminating byte equal to `term`.
280
    ///
281
    /// On success, returns Ok(Slice), where the slice does not
282
    /// include the terminating byte.  Throws MissingData or Incomplete
283
    /// if we do not find the terminating bytes.
284
    ///
285
    /// Advances the reader to the point immediately after the terminating
286
    /// byte.
287
    ///
288
    /// # Example
289
    /// ```
290
    /// use tor_bytes::{Reader,Result};
291
    /// let m = b"Hello\0wrld";
292
    /// let mut b = Reader::from_slice(m);
293
    /// assert_eq!(b.take_until(0)?, b"Hello");
294
    /// assert_eq!(b.into_rest(), b"wrld");
295
    /// # Result::Ok(())
296
    /// ```
297
20338
    pub fn take_until(&mut self, term: u8) -> Result<&'a [u8]> {
298
12176
        let pos =
299
20338
            self.b[self.off..]
300
20338
                .iter()
301
130958
                .position(|b| *b == term)
302
20338
                .ok_or(self.incomplete_error(
303
20338
                    //
304
20338
                    1.try_into().expect("1 == 0"),
305
20338
                ))?;
306
12176
        let result = self.take(pos)?;
307
12176
        self.advance(1)?;
308
12176
        Ok(result)
309
20338
    }
310
    /// Consume and return all the remaining bytes, but do not consume the reader
311
    ///
312
    /// This can be useful if you need to possibly read either fixed-length data,
313
    /// or variable length data eating the rest of the `Reader`.
314
    ///
315
    /// The `Reader` will be left devoid of further bytes.
316
    /// Consider using `into_rest()` instead.
317
23328
    pub fn take_rest(&mut self) -> &'a [u8] {
318
23328
        self.take(self.remaining())
319
23328
            .expect("taking remaining failed")
320
23328
    }
321

            
322
    /// Consume and return all but the last `n` remaining bytes.
323
    ///
324
    /// Gives `Error::MissingData` if there are fewer than `n` remaining bytes.
325
    ///
326
    /// It is invalid to call this method on a `Reader` constructed with
327
    /// [`Reader::from_possibly_incomplete_slice`].  (If we don't know where the
328
    /// data actually ends, we can't take all but the last `n` bytes.)
329
    /// Such calls cause an internal error.
330
    ///
331
    /// # Example
332
    /// ```
333
    /// use tor_bytes::{Reader,Result};
334
    /// let m = b"Hello World";
335
    /// let mut b = Reader::from_slice(m);
336
    /// assert_eq!(b.take_all_but(2)?, b"Hello Wor");
337
    /// assert_eq!(b.into_rest(), b"ld");
338
    /// # Result::Ok(())
339
    /// ```
340
1098
    pub fn take_all_but(&mut self, n: usize) -> Result<&'a [u8]> {
341
1098
        match self.completeness {
342
            Completeness::PossiblyIncomplete => {
343
2
                return Err(Error::Bug(bad_api_usage!(
344
2
                    "Called take_all_but on a PossiblyIncomplete reader."
345
2
                )))
346
            }
347
1096
            Completeness::SupposedlyComplete => {}
348
        }
349

            
350
1096
        let n_to_take = self.remaining().checked_sub(n).ok_or(Error::MissingData)?;
351

            
352
1094
        let result = self
353
1094
            .take(n_to_take)
354
1094
            .map_err(into_internal!("Subtraction misled us somehow"))?;
355
1094
        debug_assert_eq!(self.remaining(), n);
356
1094
        Ok(result)
357
1098
    }
358

            
359
    /// Try to decode and remove a Readable from this reader, using its
360
    /// take_from() method.
361
    ///
362
    /// On failure, consumes nothing.
363
531104
    pub fn extract<E: Readable>(&mut self) -> Result<E> {
364
531104
        let off_orig = self.off;
365
531104
        let result = E::take_from(self);
366
531104
        if result.is_err() {
367
14952
            // We encountered an error; we should rewind.
368
14952
            self.off = off_orig;
369
516298
        }
370
531104
        result
371
531104
    }
372

            
373
    /// Try to decode and remove `n` Readables from this reader, using the
374
    /// Readable's take_from() method.
375
    ///
376
    /// On failure, consumes nothing.
377
5229
    pub fn extract_n<E: Readable>(&mut self, n: usize) -> Result<Vec<E>> {
378
5229
        // This `min` will help us defend against a pathological case where an
379
5229
        // attacker tells us that there are BIGNUM elements forthcoming, and our
380
5229
        // attempt to allocate `Vec::with_capacity(BIGNUM)` makes us panic.
381
5229
        //
382
5229
        // The `min` can be incorrect if E is somehow encodable in zero bytes
383
5229
        // (!?), but that will only cause our initial allocation to be too
384
5229
        // small.
385
5229
        //
386
5229
        // In practice, callers should always check that `n` is reasonable
387
5229
        // before calling this function, and protocol designers should not
388
5229
        // provide e.g. 32-bit counters for object types of which we should
389
5229
        // never allocate u32::MAX.
390
5229
        let n_alloc = std::cmp::min(n, self.remaining());
391
5229
        let mut result = Vec::with_capacity(n_alloc);
392
5229
        let off_orig = self.off;
393
5229
        for _ in 0..n {
394
17840
            match E::take_from(self) {
395
17838
                Ok(item) => result.push(item),
396
2
                Err(e) => {
397
2
                    // Encountered an error; we should rewind.
398
2
                    self.off = off_orig;
399
2
                    return Err(e);
400
                }
401
            }
402
        }
403
5227
        Ok(result)
404
5229
    }
405

            
406
    /// Decode something with a `u8` length field
407
    ///
408
    /// Prefer to use this function, rather than ad-hoc `take_u8`
409
    /// and subsequent manual length checks.
410
    /// Using this facility eliminates the need to separately keep track of the lengths.
411
    ///
412
    /// `read_nested` consumes a length field,
413
    /// and provides the closure `f` with an inner `Reader` that
414
    /// contains precisely that many bytes -
415
    /// the bytes which follow the length field in the original reader.
416
    /// If the closure is successful, `read_nested` checks that that inner reader is exhausted,
417
    /// i.e. that the inner contents had the same length as was specified.
418
    ///
419
    /// The closure should read whatever is inside the nested structure
420
    /// from the nested reader.
421
    /// It may well want to use `take_rest`, to consume all of the counted bytes.
422
    ///
423
    /// On failure, the amount consumed is not specified.
424
20156
    pub fn read_nested_u8len<F, T>(&mut self, f: F) -> Result<T>
425
20156
    where
426
20156
        F: FnOnce(&mut Reader) -> Result<T>,
427
20156
    {
428
20156
        read_nested_generic::<u8, _, _>(self, f)
429
20156
    }
430

            
431
    /// Start decoding something with a u16 length field
432
607
    pub fn read_nested_u16len<F, T>(&mut self, f: F) -> Result<T>
433
607
    where
434
607
        F: FnOnce(&mut Reader) -> Result<T>,
435
607
    {
436
607
        read_nested_generic::<u16, _, _>(self, f)
437
607
    }
438

            
439
    /// Start decoding something with a u32 length field
440
2
    pub fn read_nested_u32len<F, T>(&mut self, f: F) -> Result<T>
441
2
    where
442
2
        F: FnOnce(&mut Reader) -> Result<T>,
443
2
    {
444
2
        read_nested_generic::<u32, _, _>(self, f)
445
2
    }
446

            
447
    /// Return a cursor object describing the current position of this Reader
448
    /// within its underlying byte stream.
449
    ///
450
    /// The resulting [`Cursor`] can be used with `range`, but nothing else.
451
    ///
452
    /// Note that having to use a `Cursor` is typically an anti-pattern: it
453
    /// tends to indicate that whatever you're parsing could probably have a
454
    /// better design that would better separate data from metadata.
455
    /// Unfortunately, there are a few places like that in the Tor  protocols.
456
    //
457
    // TODO: This could instead be a function that takes a closure, passes a
458
    // reader to that closure, and returns the closure's output along with
459
    // whatever the reader consumed.
460
1368
    pub fn cursor(&self) -> Cursor<'a> {
461
1368
        Cursor {
462
1368
            pos: self.off,
463
1368
            _phantom: std::marker::PhantomData,
464
1368
        }
465
1368
    }
466

            
467
    /// Return the slice of bytes between the start cursor (inclusive) and end
468
    /// cursor (exclusive).
469
    ///
470
    /// If the cursors are not in order, return an empty slice.
471
    ///
472
    /// This function is guaranteed not to panic if the inputs were generated
473
    /// from a different Reader, but if so the byte slice that it returns will
474
    /// not be meaningful.
475
826
    pub fn range(&self, start: Cursor<'a>, end: Cursor<'a>) -> &'a [u8] {
476
826
        if start.pos <= end.pos && end.pos <= self.b.len() {
477
824
            &self.b[start.pos..end.pos]
478
        } else {
479
2
            &self.b[..0]
480
        }
481
826
    }
482

            
483
    /// Returns the error that should be returned if we ran out of data
484
    ///
485
    /// For a usual `Reader` this is [`Error::MissingData`].
486
    /// For a reader from
487
    /// [`Reader::from_possibly_incomplete_slice`]
488
    /// it's [`Error::Incomplete`].
489
73812
    pub fn incomplete_error(&self, deficit: NonZeroUsize) -> Error {
490
        use Completeness as C;
491
        use Error as E;
492
73812
        match self.completeness {
493
68370
            C::PossiblyIncomplete => E::Incomplete {
494
68370
                deficit: deficit.into(),
495
68370
            },
496
5442
            C::SupposedlyComplete => E::MissingData,
497
        }
498
73812
    }
499
}
500

            
501
/// A reference to a position within a [`Reader`].
502
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
503
pub struct Cursor<'a> {
504
    /// The underlying position within the reader.
505
    pos: usize,
506
    /// Used so that we can restrict the cursor to the lifetime of the
507
    /// underlying byte slice.
508
    _phantom: std::marker::PhantomData<&'a [u8]>,
509
}
510

            
511
/// Implementation of `read_nested_*` -- generic
512
20767
fn read_nested_generic<L, F, T>(b: &mut Reader, f: F) -> Result<T>
513
20767
where
514
20767
    F: FnOnce(&mut Reader) -> Result<T>,
515
20767
    L: Readable + Copy + Sized + TryInto<usize>,
516
20767
{
517
20767
    let length: L = b.extract()?;
518
20765
    let length: usize = length.try_into().map_err(|_| Error::BadLengthValue)?;
519
20763
    let slice = b.take(length)?;
520
20696
    let mut inner = Reader::from_slice(slice);
521
20696
    let out = f(&mut inner)?;
522
20419
    inner.should_be_exhausted()?;
523
20309
    Ok(out)
524
20767
}
525

            
526
#[cfg(test)]
527
mod tests {
528
    #![allow(clippy::unwrap_used)]
529
    #![allow(clippy::cognitive_complexity)]
530
    use super::*;
531
    #[test]
532
    fn bytecursor_read_ok() {
533
        let bytes = b"On a mountain halfway between Reno and Rome";
534
        let mut bc = Reader::from_slice(&bytes[..]);
535

            
536
        assert_eq!(bc.consumed(), 0);
537
        assert_eq!(bc.remaining(), 43);
538
        assert_eq!(bc.total_len(), 43);
539

            
540
        assert_eq!(bc.take(3).unwrap(), &b"On "[..]);
541
        assert_eq!(bc.consumed(), 3);
542

            
543
        assert_eq!(bc.take_u16().unwrap(), 0x6120);
544
        assert_eq!(bc.take_u8().unwrap(), 0x6d);
545
        assert_eq!(bc.take_u64().unwrap(), 0x6f756e7461696e20);
546
        assert_eq!(bc.take_u32().unwrap(), 0x68616c66);
547
        assert_eq!(bc.consumed(), 18);
548
        assert_eq!(bc.remaining(), 25);
549
        assert_eq!(bc.total_len(), 43);
550

            
551
        assert_eq!(bc.peek(7).unwrap(), &b"way bet"[..]);
552
        assert_eq!(bc.consumed(), 18); // no change
553
        assert_eq!(bc.remaining(), 25); // no change
554
        assert_eq!(bc.total_len(), 43); // no change
555

            
556
        assert_eq!(bc.peek(7).unwrap(), &b"way bet"[..]);
557
        assert_eq!(bc.consumed(), 18); // no change this time either.
558

            
559
        bc.advance(12).unwrap();
560
        assert_eq!(bc.consumed(), 30);
561
        assert_eq!(bc.remaining(), 13);
562

            
563
        let rem = bc.into_rest();
564
        assert_eq!(rem, &b"Reno and Rome"[..]);
565

            
566
        // now let's try consuming right up to the end.
567
        let mut bc = Reader::from_slice(&bytes[..]);
568
        bc.advance(22).unwrap();
569
        assert_eq!(bc.remaining(), 21);
570
        let rem = bc.take(21).unwrap();
571
        assert_eq!(rem, &b"between Reno and Rome"[..]);
572
        assert_eq!(bc.consumed(), 43);
573
        assert_eq!(bc.remaining(), 0);
574

            
575
        // We can still take a zero-length slice.
576
        assert_eq!(bc.take(0).unwrap(), &b""[..]);
577
    }
578

            
579
    #[test]
580
    fn read_u128() {
581
        let bytes = bytes::Bytes::from(&b"irreproducibility?"[..]); // 18 bytes
582
        let mut b = Reader::from_bytes(&bytes);
583

            
584
        assert_eq!(b.take_u8().unwrap(), b'i');
585
        assert_eq!(b.take_u128().unwrap(), 0x72726570726f6475636962696c697479);
586
        assert_eq!(b.remaining(), 1);
587
    }
588

            
589
    #[test]
590
    fn bytecursor_read_missing() {
591
        let bytes = b"1234567";
592
        let mut bc = Reader::from_slice_for_test(&bytes[..]);
593

            
594
        assert_eq!(bc.consumed(), 0);
595
        assert_eq!(bc.remaining(), 7);
596
        assert_eq!(bc.total_len(), 7);
597

            
598
        assert_eq!(bc.take_u64(), Err(Error::new_incomplete_for_test(1)));
599
        assert_eq!(bc.take(8), Err(Error::new_incomplete_for_test(1)));
600
        assert_eq!(bc.peek(8), Err(Error::new_incomplete_for_test(1)));
601

            
602
        assert_eq!(bc.consumed(), 0);
603
        assert_eq!(bc.remaining(), 7);
604
        assert_eq!(bc.total_len(), 7);
605

            
606
        assert_eq!(bc.take_u32().unwrap(), 0x31323334); // get 4 bytes. 3 left.
607
        assert_eq!(bc.take_u32(), Err(Error::new_incomplete_for_test(1)));
608

            
609
        assert_eq!(bc.consumed(), 4);
610
        assert_eq!(bc.remaining(), 3);
611
        assert_eq!(bc.total_len(), 7);
612

            
613
        assert_eq!(bc.take_u16().unwrap(), 0x3536); // get 2 bytes. 1 left.
614
        assert_eq!(bc.take_u16(), Err(Error::new_incomplete_for_test(1)));
615

            
616
        assert_eq!(bc.consumed(), 6);
617
        assert_eq!(bc.remaining(), 1);
618
        assert_eq!(bc.total_len(), 7);
619

            
620
        assert_eq!(bc.take_u8().unwrap(), 0x37); // get 1 byte. 0 left.
621
        assert_eq!(bc.take_u8(), Err(Error::new_incomplete_for_test(1)));
622

            
623
        assert_eq!(bc.consumed(), 7);
624
        assert_eq!(bc.remaining(), 0);
625
        assert_eq!(bc.total_len(), 7);
626
    }
627

            
628
    #[test]
629
    fn advance_too_far() {
630
        let bytes = b"12345";
631
        let mut b = Reader::from_slice_for_test(&bytes[..]);
632
        assert_eq!(b.remaining(), 5);
633
        assert_eq!(b.advance(16), Err(Error::new_incomplete_for_test(11)));
634
        assert_eq!(b.remaining(), 5);
635
        assert_eq!(b.advance(5), Ok(()));
636
        assert_eq!(b.remaining(), 0);
637
    }
638

            
639
    #[test]
640
    fn truncate() {
641
        let bytes = b"Hello universe!!!1!";
642
        let mut b = Reader::from_slice_for_test(&bytes[..]);
643

            
644
        assert_eq!(b.take(5).unwrap(), &b"Hello"[..]);
645
        assert_eq!(b.remaining(), 14);
646
        assert_eq!(b.consumed(), 5);
647
        b.truncate(9);
648
        assert_eq!(b.remaining(), 9);
649
        assert_eq!(b.consumed(), 5);
650
        assert_eq!(b.take_u8().unwrap(), 0x20);
651
        assert_eq!(b.into_rest(), &b"universe"[..]);
652
    }
653

            
654
    #[test]
655
    fn exhaust() {
656
        let b = Reader::from_slice_for_test(&b""[..]);
657
        assert_eq!(b.should_be_exhausted(), Ok(()));
658

            
659
        let mut b = Reader::from_slice_for_test(&b"outis"[..]);
660
        assert_eq!(b.should_be_exhausted(), Err(Error::ExtraneousBytes));
661
        b.take(4).unwrap();
662
        assert_eq!(b.should_be_exhausted(), Err(Error::ExtraneousBytes));
663
        b.take(1).unwrap();
664
        assert_eq!(b.should_be_exhausted(), Ok(()));
665
    }
666

            
667
    #[test]
668
    fn take_rest() {
669
        let mut b = Reader::from_slice_for_test(b"si vales valeo");
670
        assert_eq!(b.take(3).unwrap(), b"si ");
671
        assert_eq!(b.take_rest(), b"vales valeo");
672
        assert_eq!(b.take_rest(), b"");
673
    }
674

            
675
    #[test]
676
    fn take_until() {
677
        let mut b = Reader::from_slice_for_test(&b"si vales valeo"[..]);
678
        assert_eq!(b.take_until(b' ').unwrap(), &b"si"[..]);
679
        assert_eq!(b.take_until(b' ').unwrap(), &b"vales"[..]);
680
        assert_eq!(b.take_until(b' '), Err(Error::new_incomplete_for_test(1)));
681
    }
682

            
683
    #[test]
684
    fn truncate_badly() {
685
        let mut b = Reader::from_slice_for_test(&b"abcdefg"[..]);
686
        b.truncate(1000);
687
        assert_eq!(b.total_len(), 7);
688
        assert_eq!(b.remaining(), 7);
689
    }
690

            
691
    #[test]
692
    fn nested_good() {
693
        let mut b = Reader::from_slice_for_test(b"abc\0\0\x04defghijkl");
694
        assert_eq!(b.take(3).unwrap(), b"abc");
695

            
696
        b.read_nested_u16len(|s| {
697
            assert!(s.should_be_exhausted().is_ok());
698
            Ok(())
699
        })
700
        .unwrap();
701

            
702
        b.read_nested_u8len(|s| {
703
            assert_eq!(s.take(4).unwrap(), b"defg");
704
            assert!(s.should_be_exhausted().is_ok());
705
            Ok(())
706
        })
707
        .unwrap();
708

            
709
        assert_eq!(b.take(2).unwrap(), b"hi");
710
    }
711

            
712
    #[test]
713
    fn nested_bad() {
714
        let mut b = Reader::from_slice_for_test(b"................");
715
        assert_eq!(
716
            read_nested_generic::<u128, _, ()>(&mut b, |_| panic!())
717
                .err()
718
                .unwrap(),
719
            Error::BadLengthValue
720
        );
721

            
722
        let mut b = Reader::from_slice_for_test(b"................");
723
        assert_eq!(
724
            b.read_nested_u32len::<_, ()>(|_| panic!()).err().unwrap(),
725
            Error::new_incomplete_for_test(774778414 - (16 - 4))
726
        );
727
    }
728

            
729
    #[test]
730
    fn nested_inner_bad() {
731
        let mut b = Reader::from_slice_for_test(&[1, 66]);
732
        assert_eq!(
733
            b.read_nested_u8len(|b| b.take_u32()),
734
            Err(Error::MissingData),
735
        );
736
    }
737

            
738
    #[test]
739
    fn incomplete_slice() {
740
        // Test specifically the from_possibly_incomplete_slice constructor -
741
        // ie, deliberately don't use Reader::from_slice_for_test.
742
        let mut b = Reader::from_possibly_incomplete_slice(&[]);
743
        assert_eq!(b.take_u32(), Err(Error::new_incomplete_for_test(4)));
744
    }
745

            
746
    #[test]
747
    fn extract() {
748
        // For example purposes, declare a length-then-bytes string type.
749
        #[derive(Debug)]
750
        struct LenEnc(Vec<u8>);
751
        impl Readable for LenEnc {
752
            fn take_from(b: &mut Reader<'_>) -> Result<Self> {
753
                let length = b.take_u8()?;
754
                let content = b.take(length as usize)?.into();
755
                Ok(LenEnc(content))
756
            }
757
        }
758

            
759
        let bytes = b"\x04this\x02is\x09sometimes\x01a\x06string!";
760
        let mut b = Reader::from_slice_for_test(&bytes[..]);
761

            
762
        let le: LenEnc = b.extract().unwrap();
763
        assert_eq!(&le.0[..], &b"this"[..]);
764

            
765
        let les: Vec<LenEnc> = b.extract_n(4).unwrap();
766
        assert_eq!(&les[3].0[..], &b"string"[..]);
767

            
768
        assert_eq!(b.remaining(), 1);
769

            
770
        // Make sure that we don't advance on a failing extract().
771
        let le: Result<LenEnc> = b.extract();
772
        assert_eq!(le.unwrap_err(), Error::new_incomplete_for_test(33));
773
        assert_eq!(b.remaining(), 1);
774

            
775
        // Make sure that we don't advance on a failing extract_n()
776
        let mut b = Reader::from_slice_for_test(&bytes[..]);
777
        assert_eq!(b.remaining(), 28);
778
        let les: Result<Vec<LenEnc>> = b.extract_n(10);
779
        assert_eq!(les.unwrap_err(), Error::new_incomplete_for_test(33));
780
        assert_eq!(b.remaining(), 28);
781
    }
782

            
783
    #[test]
784
    fn cursor() -> Result<()> {
785
        let alphabet = b"abcdefghijklmnopqrstuvwxyz";
786
        let mut b = Reader::from_slice_for_test(&alphabet[..]);
787

            
788
        let c1 = b.cursor();
789
        let _ = b.take_u16()?;
790
        let c2 = b.cursor();
791
        let c2b = b.cursor();
792
        b.advance(7)?;
793
        let c3 = b.cursor();
794

            
795
        assert_eq!(b.range(c1, c2), &b"ab"[..]);
796
        assert_eq!(b.range(c2, c3), &b"cdefghi"[..]);
797
        assert_eq!(b.range(c1, c3), &b"abcdefghi"[..]);
798
        assert_eq!(b.range(c1, c1), &b""[..]);
799
        assert_eq!(b.range(c3, c1), &b""[..]);
800
        assert_eq!(c2, c2b);
801
        assert!(c1 < c2);
802
        assert!(c2 < c3);
803

            
804
        Ok(())
805
    }
806

            
807
    #[test]
808
    fn take_all_but() -> Result<()> {
809
        let message = b"byte manipulation for fun and (non)-profit";
810

            
811
        // Case 1: Successful, complete reader
812
        // (Can't use from_slice_for_test here: that's a possibly-incomplete reader.)
813
        let mut b = Reader::from_slice(message);
814
        assert_eq!(b.take_all_but(6)?, b"byte manipulation for fun and (non)-");
815
        assert_eq!(b.into_rest(), b"profit");
816

            
817
        // Case 1b: Successful, take nothing, complete reader.
818
        let mut b = Reader::from_slice(message);
819
        assert_eq!(b.take_all_but(message.len())?, b"");
820
        assert_eq!(b.into_rest(), message);
821

            
822
        // Case 1c: Successful, take everything, complete reader.
823
        let mut b = Reader::from_slice(message);
824
        assert_eq!(b.take_all_but(0)?, message);
825
        assert_eq!(b.into_rest(), b"");
826

            
827
        // Case 2: Unsuccessful, complete reader
828
        let mut b = Reader::from_slice(message);
829
        assert!(matches!(
830
            b.take_all_but(message.len() + 1),
831
            Err(Error::MissingData)
832
        ));
833

            
834
        // Case 3: Anything, incomplete reader.
835
        let mut b = Reader::from_possibly_incomplete_slice(message);
836
        assert!(matches!(b.take_all_but(6), Err(Error::Bug(_))));
837

            
838
        Ok(())
839
    }
840
}