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

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

            
5
use crate::{
6
    Error::{self},
7
    Readable, Result,
8
};
9
use 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
/// ```
56
pub 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)]
72
enum 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

            
83
impl<'a> Reader<'a> {
84
    /// Construct a new Reader from a slice of bytes.
85
    ///
86
    /// In tests, prefer [`Reader::from_slice_for_test`].
87
391520
    pub fn from_slice(slice: &'a [u8]) -> Self {
88
391520
        Reader {
89
391520
            b: slice,
90
391520
            off: 0,
91
391520
            completeness: Completeness::SupposedlyComplete,
92
391520
        }
93
391520
    }
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
85337
    pub fn from_possibly_incomplete_slice(slice: &'a [u8]) -> Self {
121
85337
        Reader {
122
85337
            b: slice,
123
85337
            off: 0,
124
85337
            completeness: Completeness::PossiblyIncomplete,
125
85337
        }
126
85337
    }
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
6658
    pub fn from_slice_for_test(slice: &'a [u8]) -> Self {
133
6658
        Self::from_possibly_incomplete_slice(slice)
134
6658
    }
135
    /// Construct a new Reader from a 'Bytes' object.
136
7127
    pub fn from_bytes(b: &'a bytes::Bytes) -> Self {
137
7127
        Self::from_slice(b.as_ref())
138
7127
    }
139
    /// Return the total length of the slice in this reader, including
140
    /// consumed bytes and remaining bytes.
141
18
    pub fn total_len(&self) -> usize {
142
18
        self.b.len()
143
18
    }
144
    /// Return the total number of bytes in this reader that have not
145
    /// yet been read.
146
5262320
    pub fn remaining(&self) -> usize {
147
5262320
        self.b.len() - self.off
148
5262320
    }
149
    /// Consume this reader, and return a slice containing the remaining
150
    /// bytes from its slice that it did not consume.
151
1960
    pub fn into_rest(self) -> &'a [u8] {
152
1960
        &self.b[self.off..]
153
1960
    }
154
    /// Return the total number of bytes in this reader that have
155
    /// already been read.
156
121303
    pub fn consumed(&self) -> usize {
157
121303
        self.off
158
121303
    }
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
2272437
    pub fn advance(&mut self, n: usize) -> Result<()> {
164
2272437
        self.peek(n)?;
165
2272435
        self.off += n;
166
2272435
        Ok(())
167
2272437
    }
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
144859
    pub fn should_be_exhausted(&self) -> Result<()> {
173
144859
        if self.remaining() != 0 {
174
154
            return Err(Error::ExtraneousBytes);
175
144705
        }
176
144705
        Ok(())
177
144859
    }
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
210454
    pub fn truncate(&mut self, n: usize) {
183
210454
        if n < self.remaining() {
184
66302
            self.b = &self.b[..self.off + n];
185
144152
        }
186
210454
    }
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
4200257
    pub fn peek(&self, n: usize) -> Result<&'a [u8]> {
194
4200257
        if let Some(deficit) = n
195
4200257
            .checked_sub(self.remaining())
196
4212006
            .and_then(|d| d.try_into().ok())
197
        {
198
61751
            return Err(self.incomplete_error(deficit));
199
4138506
        }
200

            
201
4138506
        Ok(&self.b[self.off..(n + self.off)])
202
4200257
    }
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
1885514
    pub fn take(&mut self, n: usize) -> Result<&'a [u8]> {
220
1885514
        let b = self.peek(n)?;
221
1828042
        self.advance(n)?;
222
1828042
        Ok(b)
223
1885514
    }
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
808698
    pub fn take_into(&mut self, buf: &mut [u8]) -> Result<()> {
247
808698
        let n = buf.len();
248
808698
        let b = self.take(n)?;
249
791736
        buf.copy_from_slice(b);
250
791736
        Ok(())
251
808698
    }
252
    /// Try to consume and return a u8 from this reader.
253
725063
    pub fn take_u8(&mut self) -> Result<u8> {
254
725063
        let b = self.take(1)?;
255
691911
        Ok(b[0])
256
725063
    }
257
    /// Try to consume and return a big-endian u16 from this reader.
258
499510
    pub fn take_u16(&mut self) -> Result<u16> {
259
499510
        let b: [u8; 2] = self.extract()?;
260
495308
        let r = u16::from_be_bytes(b);
261
495308
        Ok(r)
262
499510
    }
263
    /// Try to consume and return a big-endian u32 from this reader.
264
84614
    pub fn take_u32(&mut self) -> Result<u32> {
265
84614
        let b: [u8; 4] = self.extract()?;
266
79508
        let r = u32::from_be_bytes(b);
267
79508
        Ok(r)
268
84614
    }
269
    /// Try to consume and return a big-endian u64 from this reader.
270
87456
    pub fn take_u64(&mut self) -> Result<u64> {
271
87456
        let b: [u8; 8] = self.extract()?;
272
87454
        let r = u64::from_be_bytes(b);
273
87454
        Ok(r)
274
87456
    }
275
    /// Try to consume and return a big-endian u128 from this reader.
276
12456
    pub fn take_u128(&mut self) -> Result<u128> {
277
12456
        let b: [u8; 16] = self.extract()?;
278
4956
        let r = u128::from_be_bytes(b);
279
4956
        Ok(r)
280
12456
    }
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
27231
    pub fn take_until(&mut self, term: u8) -> Result<&'a [u8]> {
301
27231
        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
27231
    }
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
27243
    pub fn take_until_with_limit(&mut self, term: u8, max_len: usize) -> Result<&'a [u8]> {
324
27243
        let limit: usize = std::cmp::min(max_len, self.remaining());
325
27243
        let pos = match self.b[self.off..limit + self.off]
326
27243
            .iter()
327
185471
            .position(|b| *b == term)
328
        {
329
18235
            Some(p) => p,
330
            None => {
331
9008
                if max_len > self.remaining() {
332
                    // since user asked for more than we have, we should return Incomplete, not LimitExceeded
333
9004
                    return Err(self.incomplete_error(
334
9004
                        //
335
9004
                        1.try_into().expect("1 == 0"),
336
9004
                    ));
337
                } else {
338
4
                    return Err(Error::LimitExceeded {
339
4
                        limit,
340
4
                        terminator: term,
341
4
                    });
342
                }
343
            }
344
        };
345

            
346
18235
        let result: &[u8] = self.take(pos)?;
347
18235
        self.advance(1)?;
348
18235
        Ok(result)
349
27243
    }
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
47104
    pub fn take_rest(&mut self) -> &'a [u8] {
358
47104
        self.take(self.remaining())
359
47104
            .expect("taking remaining failed")
360
47104
    }
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
1885
    pub fn take_all_but(&mut self, n: usize) -> Result<&'a [u8]> {
381
1885
        match self.completeness {
382
            Completeness::PossiblyIncomplete => {
383
2
                return Err(Error::Bug(bad_api_usage!(
384
2
                    "Called take_all_but on a PossiblyIncomplete reader."
385
2
                )));
386
            }
387
1883
            Completeness::SupposedlyComplete => {}
388
        }
389

            
390
1883
        let n_to_take = self.remaining().checked_sub(n).ok_or(Error::MissingData)?;
391

            
392
1881
        let result = self
393
1881
            .take(n_to_take)
394
1881
            .map_err(into_internal!("Subtraction misled us somehow"))?;
395
1881
        debug_assert_eq!(self.remaining(), n);
396
1881
        Ok(result)
397
1885
    }
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
1040342
    pub fn extract<E: Readable>(&mut self) -> Result<E> {
404
1040342
        let off_orig = self.off;
405
1040342
        let result = E::take_from(self);
406
1040342
        if result.is_err() {
407
19475
            // We encountered an error; we should rewind.
408
19475
            self.off = off_orig;
409
1021083
        }
410
1040342
        result
411
1040342
    }
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
7879
    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
7879
        let n_alloc = std::cmp::min(n, self.remaining());
431
7879
        let mut result = Vec::with_capacity(n_alloc);
432
7879
        let off_orig = self.off;
433
7879
        for _ in 0..n {
434
26606
            match E::take_from(self) {
435
26604
                Ok(item) => result.push(item),
436
2
                Err(e) => {
437
                    // Encountered an error; we should rewind.
438
2
                    self.off = off_orig;
439
2
                    return Err(e);
440
                }
441
            }
442
        }
443
7877
        Ok(result)
444
7879
    }
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
29560
    pub fn read_nested_u8len<F, T>(&mut self, f: F) -> Result<T>
465
29560
    where
466
29560
        F: FnOnce(&mut Reader) -> Result<T>,
467
    {
468
29560
        read_nested_generic::<u8, _, _>(self, f)
469
29560
    }
470

            
471
    /// Start decoding something with a u16 length field
472
695
    pub fn read_nested_u16len<F, T>(&mut self, f: F) -> Result<T>
473
695
    where
474
695
        F: FnOnce(&mut Reader) -> Result<T>,
475
    {
476
695
        read_nested_generic::<u16, _, _>(self, f)
477
695
    }
478

            
479
    /// Start decoding something with a u32 length field
480
2
    pub fn read_nested_u32len<F, T>(&mut self, f: F) -> Result<T>
481
2
    where
482
2
        F: FnOnce(&mut Reader) -> Result<T>,
483
    {
484
2
        read_nested_generic::<u32, _, _>(self, f)
485
2
    }
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
1508
    pub fn cursor(&self) -> Cursor<'a> {
501
1508
        Cursor {
502
1508
            pos: self.off,
503
1508
            _phantom: std::marker::PhantomData,
504
1508
        }
505
1508
    }
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
910
    pub fn range(&self, start: Cursor<'a>, end: Cursor<'a>) -> &'a [u8] {
516
910
        if start.pos <= end.pos && end.pos <= self.b.len() {
517
908
            &self.b[start.pos..end.pos]
518
        } else {
519
2
            &self.b[..0]
520
        }
521
910
    }
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
70755
    pub fn incomplete_error(&self, deficit: NonZeroUsize) -> Error {
530
        use Completeness as C;
531
        use Error as E;
532
70755
        match self.completeness {
533
70078
            C::PossiblyIncomplete => E::Incomplete {
534
70078
                deficit: deficit.into(),
535
70078
            },
536
677
            C::SupposedlyComplete => E::MissingData,
537
        }
538
70755
    }
539
}
540

            
541
/// A reference to a position within a [`Reader`].
542
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
543
pub 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
552
30259
fn read_nested_generic<L, F, T>(b: &mut Reader, f: F) -> Result<T>
553
30259
where
554
30259
    F: FnOnce(&mut Reader) -> Result<T>,
555
30259
    L: Readable + Copy + Sized + TryInto<usize>,
556
{
557
30259
    let length: L = b.extract()?;
558
30257
    let length: usize = length.try_into().map_err(|_| Error::BadLengthValue)?;
559
30255
    let slice = b.take(length)?;
560
30180
    let mut inner = Reader::from_slice(slice);
561
30180
    let out = f(&mut inner)?;
562
29863
    inner.should_be_exhausted()?;
563
29737
    Ok(out)
564
30259
}
565

            
566
#[cfg(test)]
567
mod 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
}