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

            
3
use crate::{Error, Readable, Result};
4
use arrayref::array_ref;
5

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

            
58
impl<'a> Reader<'a> {
59
    /// Construct a new Reader from a slice of bytes.
60
78324
    pub fn from_slice(slice: &'a [u8]) -> Self {
61
78324
        Reader { b: slice, off: 0 }
62
78324
    }
63
    /// Construct a new Reader from a 'Bytes' object.
64
987
    pub fn from_bytes(b: &'a bytes::Bytes) -> Self {
65
987
        Self::from_slice(b.as_ref())
66
987
    }
67
    /// Return the total length of the slice in this reader, including
68
    /// consumed bytes and remaining bytes.
69
9
    pub fn total_len(&self) -> usize {
70
9
        self.b.len()
71
9
    }
72
    /// Return the total number of bytes in this reader that have not
73
    /// yet been read.
74
998255
    pub fn remaining(&self) -> usize {
75
998255
        self.b.len() - self.off
76
998255
    }
77
    /// Consume this reader, and return a slice containing the remaining
78
    /// bytes from its slice that it did not consume.
79
60
    pub fn into_rest(self) -> &'a [u8] {
80
60
        &self.b[self.off..]
81
60
    }
82
    /// Return the total number of bytes in this reader that have
83
    /// already been read.
84
1464
    pub fn consumed(&self) -> usize {
85
1464
        self.off
86
1464
    }
87
    /// Skip `n` bytes from the reader.
88
    ///
89
    /// Returns Ok on success.  Returns Err(Error::Truncated) if there were
90
    /// not enough bytes to skip.
91
462230
    pub fn advance(&mut self, n: usize) -> Result<()> {
92
462230
        if n > self.remaining() {
93
1
            return Err(Error::Truncated);
94
462229
        }
95
462229
        self.off += n;
96
462229
        Ok(())
97
462230
    }
98
    /// Check whether this reader is exhausted (out of bytes).
99
    ///
100
    /// Return Ok if it is, and Err(Error::ExtraneousBytes)
101
    /// if there were extra bytes.
102
1144
    pub fn should_be_exhausted(&self) -> Result<()> {
103
1144
        if self.remaining() != 0 {
104
2
            return Err(Error::ExtraneousBytes);
105
1142
        }
106
1142
        Ok(())
107
1144
    }
108
    /// Truncate this reader, so that no more than `n` bytes remain.
109
    ///
110
    /// Fewer than `n` bytes may remain if there were not enough bytes
111
    /// to begin with.
112
71980
    pub fn truncate(&mut self, n: usize) {
113
71980
        if n < self.remaining() {
114
2350
            self.b = &self.b[..self.off + n];
115
69630
        }
116
71980
    }
117
    /// Try to return a slice of `n` bytes from this reader without
118
    /// consuming them.
119
    ///
120
    /// On success, returns Ok(slice).  If there are fewer than n
121
    /// bytes, returns Err(Error::Truncated).
122
316742
    pub fn peek(&self, n: usize) -> Result<&'a [u8]> {
123
316742
        if self.remaining() < n {
124
96
            return Err(Error::Truncated);
125
316646
        }
126
316646

            
127
316646
        Ok(&self.b[self.off..(n + self.off)])
128
316742
    }
129
    /// Try to consume and return a slice of `n` bytes from this reader.
130
    ///
131
    /// On success, returns Ok(Slice).  If there are fewer than n
132
    /// bytes, returns Err(Error::Truncated).
133
    ///
134
    /// # Example
135
    /// ```
136
    /// use tor_bytes::{Reader,Result};
137
    /// let m = b"Hello World";
138
    /// let mut r = Reader::from_slice(m);
139
    /// assert_eq!(r.take(5)?, b"Hello");
140
    /// assert_eq!(r.take_u8()?, 0x20);
141
    /// assert_eq!(r.take(5)?, b"World");
142
    /// r.should_be_exhausted()?;
143
    /// # Result::Ok(())
144
    /// ```
145
316217
    pub fn take(&mut self, n: usize) -> Result<&'a [u8]> {
146
316217
        let b = self.peek(n)?;
147
316122
        self.advance(n)?;
148
316209
        Ok(b)
149
316304
    }
150
    /// Try to fill a provided buffer with bytes consumed from this reader.
151
    ///
152
    /// On success, the buffer will be filled with data from the
153
    /// reader, the reader will advance by the length of the buffer,
154
    /// and we'll return Ok(()).  On failure the buffer will be
155
    /// unchanged.
156
    ///
157
    /// # Example
158
    /// ```
159
    /// use tor_bytes::Reader;
160
    /// let m = b"Hello world";
161
    /// let mut v1 = vec![0; 5];
162
    /// let mut v2 = vec![0; 5];
163
    /// let mut r = Reader::from_slice(m);
164
    /// r.take_into(&mut v1[..])?;
165
    /// assert_eq!(r.take_u8()?, b' ');
166
    /// r.take_into(&mut v2[..])?;
167
    /// assert_eq!(&v1[..], b"Hello");
168
    /// assert_eq!(&v2[..], b"world");
169
    /// r.should_be_exhausted()?;
170
    /// # tor_bytes::Result::Ok(())
171
    /// ```
172
639
    pub fn take_into(&mut self, buf: &mut [u8]) -> Result<()> {
173
639
        let n = buf.len();
174
639
        let b = self.take(n)?;
175
639
        buf.copy_from_slice(b);
176
639
        Ok(())
177
639
    }
178
    /// Try to consume and return a u8 from this reader.
179
84349
    pub fn take_u8(&mut self) -> Result<u8> {
180
84349
        let b = self.take(1)?;
181
84319
        Ok(b[0])
182
84349
    }
183
    /// Try to consume and return a big-endian u16 from this reader.
184
147729
    pub fn take_u16(&mut self) -> Result<u16> {
185
147729
        let b = self.take(2)?;
186
147699
        let r = u16::from_be_bytes(*array_ref![b, 0, 2]);
187
147699
        Ok(r)
188
147729
    }
189
    /// Try to consume and return a big-endian u32 from this reader.
190
3629
    pub fn take_u32(&mut self) -> Result<u32> {
191
3629
        let b = self.take(4)?;
192
3628
        let r = u32::from_be_bytes(*array_ref![b, 0, 4]);
193
3628
        Ok(r)
194
3629
    }
195
    /// Try to consume and return a big-endian u64 from this reader.
196
3
    pub fn take_u64(&mut self) -> Result<u64> {
197
3
        let b = self.take(8)?;
198
2
        let r = u64::from_be_bytes(*array_ref![b, 0, 8]);
199
2
        Ok(r)
200
3
    }
201
    /// Try to consume and return a big-endian u128 from this reader.
202
176
    pub fn take_u128(&mut self) -> Result<u128> {
203
176
        let b = self.take(16)?;
204
176
        let r = u128::from_be_bytes(*array_ref![b, 0, 16]);
205
176
        Ok(r)
206
176
    }
207
    /// Try to consume and return bytes from this buffer until we
208
    /// encounter a terminating byte equal to `term`.
209
    ///
210
    /// On success, returns Ok(Slice), where the slice does not
211
    /// include the terminating byte.  Returns Err(Error::Truncated)
212
    /// if we do not find the terminating bytes.
213
    ///
214
    /// Advances the reader to the point immediately after the terminating
215
    /// byte.
216
    ///
217
    /// # Example
218
    /// ```
219
    /// use tor_bytes::{Reader,Result};
220
    /// let m = b"Hello\0wrld";
221
    /// let mut r = Reader::from_slice(m);
222
    /// assert_eq!(r.take_until(0)?, b"Hello");
223
    /// assert_eq!(r.into_rest(), b"wrld");
224
    /// # Result::Ok(())
225
    /// ```
226
931
    pub fn take_until(&mut self, term: u8) -> Result<&'a [u8]> {
227
931
        let pos = self.b[self.off..]
228
931
            .iter()
229
10889
            .position(|b| *b == term)
230
931
            .ok_or(Error::Truncated)?;
231
930
        let result = self.take(pos)?;
232
930
        self.advance(1)?;
233
930
        Ok(result)
234
931
    }
235
    /// Try to decode and remove a Readable from this reader, using its
236
    /// take_from() method.
237
    ///
238
    /// On failure, consumes nothing.
239
5502
    pub fn extract<E: Readable>(&mut self) -> Result<E> {
240
5502
        let off_orig = self.off;
241
5502
        let result = E::take_from(self);
242
5502
        if result.is_err() {
243
126
            // We encountered an error; we should rewind.
244
126
            self.off = off_orig;
245
5376
        }
246
5502
        result
247
5502
    }
248

            
249
    /// Try to decode and remove `n` Readables from this reader, using the
250
    /// Readable's take_from() method.
251
    ///
252
    /// On failure, consumes nothing.
253
134
    pub fn extract_n<E: Readable>(&mut self, n: usize) -> Result<Vec<E>> {
254
134
        let mut result = Vec::new();
255
134
        let off_orig = self.off;
256
134
        for _ in 0..n {
257
274
            match E::take_from(self) {
258
273
                Ok(item) => result.push(item),
259
1
                Err(e) => {
260
1
                    // Encountered an error; we should rewind.
261
1
                    self.off = off_orig;
262
1
                    return Err(e);
263
                }
264
            }
265
        }
266
133
        Ok(result)
267
134
    }
268
}
269

            
270
#[cfg(test)]
271
mod tests {
272
    #![allow(clippy::unwrap_used)]
273
    #![allow(clippy::cognitive_complexity)]
274
    use super::*;
275
    #[test]
276
    fn bytecursor_read_ok() {
277
        let bytes = b"On a mountain halfway between Reno and Rome";
278
        let mut bc = Reader::from_slice(&bytes[..]);
279

            
280
        assert_eq!(bc.consumed(), 0);
281
        assert_eq!(bc.remaining(), 43);
282
        assert_eq!(bc.total_len(), 43);
283

            
284
        assert_eq!(bc.take(3).unwrap(), &b"On "[..]);
285
        assert_eq!(bc.consumed(), 3);
286

            
287
        assert_eq!(bc.take_u16().unwrap(), 0x6120);
288
        assert_eq!(bc.take_u8().unwrap(), 0x6d);
289
        assert_eq!(bc.take_u64().unwrap(), 0x6f756e7461696e20);
290
        assert_eq!(bc.take_u32().unwrap(), 0x68616c66);
291
        assert_eq!(bc.consumed(), 18);
292
        assert_eq!(bc.remaining(), 25);
293
        assert_eq!(bc.total_len(), 43);
294

            
295
        assert_eq!(bc.peek(7).unwrap(), &b"way bet"[..]);
296
        assert_eq!(bc.consumed(), 18); // no change
297
        assert_eq!(bc.remaining(), 25); // no change
298
        assert_eq!(bc.total_len(), 43); // no change
299

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

            
303
        bc.advance(12).unwrap();
304
        assert_eq!(bc.consumed(), 30);
305
        assert_eq!(bc.remaining(), 13);
306

            
307
        let rem = bc.into_rest();
308
        assert_eq!(rem, &b"Reno and Rome"[..]);
309

            
310
        // now let's try consuming right up to the end.
311
        let mut bc = Reader::from_slice(&bytes[..]);
312
        bc.advance(22).unwrap();
313
        assert_eq!(bc.remaining(), 21);
314
        let rem = bc.take(21).unwrap();
315
        assert_eq!(rem, &b"between Reno and Rome"[..]);
316
        assert_eq!(bc.consumed(), 43);
317
        assert_eq!(bc.remaining(), 0);
318

            
319
        // We can still take a zero-length slice.
320
        assert_eq!(bc.take(0).unwrap(), &b""[..]);
321
    }
322

            
323
    #[test]
324
    fn read_u128() {
325
        let bytes = bytes::Bytes::from(&b"irreproducibility?"[..]); // 18 bytes
326
        let mut r = Reader::from_bytes(&bytes);
327

            
328
        assert_eq!(r.take_u8().unwrap(), b'i');
329
        assert_eq!(r.take_u128().unwrap(), 0x72726570726f6475636962696c697479);
330
        assert_eq!(r.remaining(), 1);
331
    }
332

            
333
    #[test]
334
    fn bytecursor_read_missing() {
335
        let bytes = b"1234567";
336
        let mut bc = Reader::from_slice(&bytes[..]);
337

            
338
        assert_eq!(bc.consumed(), 0);
339
        assert_eq!(bc.remaining(), 7);
340
        assert_eq!(bc.total_len(), 7);
341

            
342
        assert_eq!(bc.take_u64(), Err(Error::Truncated));
343
        assert_eq!(bc.take(8), Err(Error::Truncated));
344
        assert_eq!(bc.peek(8), Err(Error::Truncated));
345

            
346
        assert_eq!(bc.consumed(), 0);
347
        assert_eq!(bc.remaining(), 7);
348
        assert_eq!(bc.total_len(), 7);
349

            
350
        assert_eq!(bc.take_u32().unwrap(), 0x31323334); // get 4 bytes. 3 left.
351
        assert_eq!(bc.take_u32(), Err(Error::Truncated));
352

            
353
        assert_eq!(bc.consumed(), 4);
354
        assert_eq!(bc.remaining(), 3);
355
        assert_eq!(bc.total_len(), 7);
356

            
357
        assert_eq!(bc.take_u16().unwrap(), 0x3536); // get 2 bytes. 1 left.
358
        assert_eq!(bc.take_u16(), Err(Error::Truncated));
359

            
360
        assert_eq!(bc.consumed(), 6);
361
        assert_eq!(bc.remaining(), 1);
362
        assert_eq!(bc.total_len(), 7);
363

            
364
        assert_eq!(bc.take_u8().unwrap(), 0x37); // get 1 byte. 0 left.
365
        assert_eq!(bc.take_u8(), Err(Error::Truncated));
366

            
367
        assert_eq!(bc.consumed(), 7);
368
        assert_eq!(bc.remaining(), 0);
369
        assert_eq!(bc.total_len(), 7);
370
    }
371

            
372
    #[test]
373
    fn advance_too_far() {
374
        let bytes = b"12345";
375
        let mut r = Reader::from_slice(&bytes[..]);
376
        assert_eq!(r.remaining(), 5);
377
        assert_eq!(r.advance(6), Err(Error::Truncated));
378
        assert_eq!(r.remaining(), 5);
379
        assert_eq!(r.advance(5), Ok(()));
380
        assert_eq!(r.remaining(), 0);
381
    }
382

            
383
    #[test]
384
    fn truncate() {
385
        let bytes = b"Hello universe!!!1!";
386
        let mut r = Reader::from_slice(&bytes[..]);
387

            
388
        assert_eq!(r.take(5).unwrap(), &b"Hello"[..]);
389
        assert_eq!(r.remaining(), 14);
390
        assert_eq!(r.consumed(), 5);
391
        r.truncate(9);
392
        assert_eq!(r.remaining(), 9);
393
        assert_eq!(r.consumed(), 5);
394
        assert_eq!(r.take_u8().unwrap(), 0x20);
395
        assert_eq!(r.into_rest(), &b"universe"[..]);
396
    }
397

            
398
    #[test]
399
    fn exhaust() {
400
        let r = Reader::from_slice(&b""[..]);
401
        assert_eq!(r.should_be_exhausted(), Ok(()));
402

            
403
        let mut r = Reader::from_slice(&b"outis"[..]);
404
        assert_eq!(r.should_be_exhausted(), Err(Error::ExtraneousBytes));
405
        r.take(4).unwrap();
406
        assert_eq!(r.should_be_exhausted(), Err(Error::ExtraneousBytes));
407
        r.take(1).unwrap();
408
        assert_eq!(r.should_be_exhausted(), Ok(()));
409
    }
410

            
411
    #[test]
412
    fn take_until() {
413
        let mut r = Reader::from_slice(&b"si vales valeo"[..]);
414
        assert_eq!(r.take_until(b' ').unwrap(), &b"si"[..]);
415
        assert_eq!(r.take_until(b' ').unwrap(), &b"vales"[..]);
416
        assert_eq!(r.take_until(b' '), Err(Error::Truncated));
417
    }
418

            
419
    #[test]
420
    fn truncate_badly() {
421
        let mut r = Reader::from_slice(&b"abcdefg"[..]);
422
        r.truncate(1000);
423
        assert_eq!(r.total_len(), 7);
424
        assert_eq!(r.remaining(), 7);
425
    }
426

            
427
    #[test]
428
    fn extract() {
429
        // For example purposes, declare a length-then-bytes string type.
430
        #[derive(Debug)]
431
        struct LenEnc(Vec<u8>);
432
        impl Readable for LenEnc {
433
            fn take_from(b: &mut Reader<'_>) -> Result<Self> {
434
                let length = b.take_u8()?;
435
                let content = b.take(length as usize)?.into();
436
                Ok(LenEnc(content))
437
            }
438
        }
439

            
440
        let bytes = b"\x04this\x02is\x09sometimes\x01a\x06string!";
441
        let mut r = Reader::from_slice(&bytes[..]);
442

            
443
        let le: LenEnc = r.extract().unwrap();
444
        assert_eq!(&le.0[..], &b"this"[..]);
445

            
446
        let les: Vec<LenEnc> = r.extract_n(4).unwrap();
447
        assert_eq!(&les[3].0[..], &b"string"[..]);
448

            
449
        assert_eq!(r.remaining(), 1);
450

            
451
        // Make sure that we don't advance on a failing extract().
452
        let le: Result<LenEnc> = r.extract();
453
        assert_eq!(le.unwrap_err(), Error::Truncated);
454
        assert_eq!(r.remaining(), 1);
455

            
456
        // Make sure that we don't advance on a failing extract_n()
457
        let mut r = Reader::from_slice(&bytes[..]);
458
        assert_eq!(r.remaining(), 28);
459
        let les: Result<Vec<LenEnc>> = r.extract_n(10);
460
        assert_eq!(les.unwrap_err(), Error::Truncated);
461
        assert_eq!(r.remaining(), 28);
462
    }
463
}