• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 // Copyright (C) 2018-2019, Cloudflare, Inc.
2 // All rights reserved.
3 //
4 // Redistribution and use in source and binary forms, with or without
5 // modification, are permitted provided that the following conditions are
6 // met:
7 //
8 //     * Redistributions of source code must retain the above copyright notice,
9 //       this list of conditions and the following disclaimer.
10 //
11 //     * Redistributions in binary form must reproduce the above copyright
12 //       notice, this list of conditions and the following disclaimer in the
13 //       documentation and/or other materials provided with the distribution.
14 //
15 // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
16 // IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
17 // THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
18 // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
19 // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
20 // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
21 // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
22 // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
23 // LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
24 // NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
25 // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26 
27 #![allow(dead_code)]
28 
29 /// Zero-copy abstraction for parsing and constructing network packets.
30 use std::mem;
31 use std::ptr;
32 
33 /// A specialized [`Result`] type for [`OctetsMut`] operations.
34 ///
35 /// [`Result`]: https://doc.rust-lang.org/std/result/enum.Result.html
36 /// [`OctetsMut`]: struct.OctetsMut.html
37 pub type Result<T> = std::result::Result<T, BufferTooShortError>;
38 
39 /// An error indicating that the provided [`OctetsMut`] is not big enough.
40 ///
41 /// [`OctetsMut`]: struct.OctetsMut.html
42 #[derive(Clone, Copy, Debug, PartialEq)]
43 pub struct BufferTooShortError;
44 
45 impl std::fmt::Display for BufferTooShortError {
fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result46     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
47         write!(f, "BufferTooShortError")
48     }
49 }
50 
51 impl std::error::Error for BufferTooShortError {
source(&self) -> Option<&(dyn std::error::Error + 'static)>52     fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
53         None
54     }
55 }
56 
57 macro_rules! peek_u {
58     ($b:expr, $ty:ty, $len:expr) => {{
59         let len = $len;
60         let src = &$b.buf[$b.off..];
61 
62         if src.len() < len {
63             return Err(BufferTooShortError);
64         }
65 
66         let mut out: $ty = 0;
67         unsafe {
68             let dst = &mut out as *mut $ty as *mut u8;
69             let off = (mem::size_of::<$ty>() - len) as isize;
70 
71             ptr::copy_nonoverlapping(src.as_ptr(), dst.offset(off), len);
72         };
73 
74         Ok(<$ty>::from_be(out))
75     }};
76 }
77 
78 macro_rules! get_u {
79     ($b:expr, $ty:ty, $len:expr) => {{
80         let out = peek_u!($b, $ty, $len);
81 
82         $b.off += $len;
83 
84         out
85     }};
86 }
87 
88 macro_rules! put_u {
89     ($b:expr, $ty:ty, $v:expr, $len:expr) => {{
90         let len = $len;
91 
92         if $b.buf.len() < $b.off + len {
93             return Err(BufferTooShortError);
94         }
95 
96         let v = $v;
97 
98         #[allow(clippy::range_plus_one)]
99         let dst = &mut $b.buf[$b.off..($b.off + len)];
100 
101         unsafe {
102             let src = &<$ty>::to_be(v) as *const $ty as *const u8;
103             let off = (mem::size_of::<$ty>() - len) as isize;
104 
105             ptr::copy_nonoverlapping(src.offset(off), dst.as_mut_ptr(), len);
106         }
107 
108         $b.off += $len;
109 
110         Ok(dst)
111     }};
112 }
113 
114 /// A zero-copy immutable byte buffer.
115 ///
116 /// `Octets` wraps an in-memory buffer of bytes and provides utility functions
117 /// for manipulating it. The underlying buffer is provided by the user and is
118 /// not copied when creating an `Octets`. Operations are panic-free and will
119 /// avoid indexing the buffer past its end.
120 ///
121 /// Additionally, an offset (initially set to the start of the buffer) is
122 /// incremented as bytes are read from / written to the buffer, to allow for
123 /// sequential operations.
124 #[derive(Debug, PartialEq)]
125 pub struct Octets<'a> {
126     buf: &'a [u8],
127     off: usize,
128 }
129 
130 impl<'a> Octets<'a> {
131     /// Creates an `Octets` from the given slice, without copying.
132     ///
133     /// Since there's no copy, the input slice needs to be mutable to allow
134     /// modifications.
with_slice(buf: &'a [u8]) -> Self135     pub fn with_slice(buf: &'a [u8]) -> Self {
136         Octets { buf, off: 0 }
137     }
138 
139     /// Reads an unsigned 8-bit integer from the current offset and advances
140     /// the buffer.
get_u8(&mut self) -> Result<u8>141     pub fn get_u8(&mut self) -> Result<u8> {
142         get_u!(self, u8, 1)
143     }
144 
145     /// Reads an unsigned 8-bit integer from the current offset without
146     /// advancing the buffer.
peek_u8(&mut self) -> Result<u8>147     pub fn peek_u8(&mut self) -> Result<u8> {
148         peek_u!(self, u8, 1)
149     }
150 
151     /// Reads an unsigned 16-bit integer in network byte-order from the current
152     /// offset and advances the buffer.
get_u16(&mut self) -> Result<u16>153     pub fn get_u16(&mut self) -> Result<u16> {
154         get_u!(self, u16, 2)
155     }
156 
157     /// Reads an unsigned 24-bit integer in network byte-order from the current
158     /// offset and advances the buffer.
get_u24(&mut self) -> Result<u32>159     pub fn get_u24(&mut self) -> Result<u32> {
160         get_u!(self, u32, 3)
161     }
162 
163     /// Reads an unsigned 32-bit integer in network byte-order from the current
164     /// offset and advances the buffer.
get_u32(&mut self) -> Result<u32>165     pub fn get_u32(&mut self) -> Result<u32> {
166         get_u!(self, u32, 4)
167     }
168 
169     /// Reads an unsigned 64-bit integer in network byte-order from the current
170     /// offset and advances the buffer.
get_u64(&mut self) -> Result<u64>171     pub fn get_u64(&mut self) -> Result<u64> {
172         get_u!(self, u64, 8)
173     }
174 
175     /// Reads an unsigned variable-length integer in network byte-order from
176     /// the current offset and advances the buffer.
get_varint(&mut self) -> Result<u64>177     pub fn get_varint(&mut self) -> Result<u64> {
178         let first = self.peek_u8()?;
179 
180         let len = varint_parse_len(first);
181 
182         if len > self.cap() {
183             return Err(BufferTooShortError);
184         }
185 
186         let mut vec = self.get_bytes(len)?.to_vec();
187 
188         // Mask the 2 most significant bits to remove the encoded length.
189         vec[0] &= 0x3f;
190 
191         let mut b = OctetsMut::with_slice(&mut vec);
192 
193         let out = match len {
194             1 => u64::from(b.get_u8()?),
195             2 => u64::from(b.get_u16()?),
196             4 => u64::from(b.get_u32()?),
197             8 => b.get_u64()?,
198             _ => unreachable!(),
199         };
200 
201         Ok(out)
202     }
203 
204     /// Reads `len` bytes from the current offset without copying and advances
205     /// the buffer.
get_bytes(&mut self, len: usize) -> Result<Octets>206     pub fn get_bytes(&mut self, len: usize) -> Result<Octets> {
207         if self.cap() < len {
208             return Err(BufferTooShortError);
209         }
210 
211         let out = Octets {
212             buf: &self.buf[self.off..self.off + len],
213             off: 0,
214         };
215 
216         self.off += len;
217 
218         Ok(out)
219     }
220 
221     /// Reads `len` bytes from the current offset without copying and advances
222     /// the buffer, where `len` is an unsigned 8-bit integer prefix.
get_bytes_with_u8_length(&mut self) -> Result<Octets>223     pub fn get_bytes_with_u8_length(&mut self) -> Result<Octets> {
224         let len = self.get_u8()?;
225         self.get_bytes(len as usize)
226     }
227 
228     /// Reads `len` bytes from the current offset without copying and advances
229     /// the buffer, where `len` is an unsigned 16-bit integer prefix in network
230     /// byte-order.
get_bytes_with_u16_length(&mut self) -> Result<Octets>231     pub fn get_bytes_with_u16_length(&mut self) -> Result<Octets> {
232         let len = self.get_u16()?;
233         self.get_bytes(len as usize)
234     }
235 
236     /// Reads `len` bytes from the current offset without copying and advances
237     /// the buffer, where `len` is an unsigned variable-length integer prefix
238     /// in network byte-order.
get_bytes_with_varint_length(&mut self) -> Result<Octets>239     pub fn get_bytes_with_varint_length(&mut self) -> Result<Octets> {
240         let len = self.get_varint()?;
241         self.get_bytes(len as usize)
242     }
243 
244     /// Reads `len` bytes from the current offset without copying and without
245     /// advancing the buffer.
peek_bytes(&self, len: usize) -> Result<Octets>246     pub fn peek_bytes(&self, len: usize) -> Result<Octets> {
247         if self.cap() < len {
248             return Err(BufferTooShortError);
249         }
250 
251         let out = Octets {
252             buf: &self.buf[self.off..self.off + len],
253             off: 0,
254         };
255 
256         Ok(out)
257     }
258 
259     /// Returns a slice of `len` elements from the current offset.
slice(&'a self, len: usize) -> Result<&'a [u8]>260     pub fn slice(&'a self, len: usize) -> Result<&'a [u8]> {
261         if len > self.cap() {
262             return Err(BufferTooShortError);
263         }
264 
265         Ok(&self.buf[self.off..self.off + len])
266     }
267 
268     /// Returns a slice of `len` elements from the end of the buffer.
slice_last(&'a self, len: usize) -> Result<&'a [u8]>269     pub fn slice_last(&'a self, len: usize) -> Result<&'a [u8]> {
270         if len > self.cap() {
271             return Err(BufferTooShortError);
272         }
273 
274         let cap = self.cap();
275         Ok(&self.buf[cap - len..])
276     }
277 
278     /// Returns the remaining capacity in the buffer.
cap(&self) -> usize279     pub fn cap(&self) -> usize {
280         self.buf.len() - self.off
281     }
282 
283     /// Returns the total length of the buffer.
len(&self) -> usize284     pub fn len(&self) -> usize {
285         self.buf.len()
286     }
287 
288     /// Returns the current offset of the buffer.
off(&self) -> usize289     pub fn off(&self) -> usize {
290         self.off
291     }
292 
293     /// Returns a reference to the internal buffer.
buf(&self) -> &[u8]294     pub fn buf(&self) -> &[u8] {
295         self.buf
296     }
297 
298     /// Copies the buffer from the current offset into a new `Vec<u8>`.
to_vec(&self) -> Vec<u8>299     pub fn to_vec(&self) -> Vec<u8> {
300         self.as_ref().to_vec()
301     }
302 }
303 
304 impl<'a> AsRef<[u8]> for Octets<'a> {
as_ref(&self) -> &[u8]305     fn as_ref(&self) -> &[u8] {
306         &self.buf[self.off..]
307     }
308 }
309 
310 /// A zero-copy mutable byte buffer.
311 ///
312 /// Like `Octets` but mutable.
313 #[derive(Debug, PartialEq)]
314 pub struct OctetsMut<'a> {
315     buf: &'a mut [u8],
316     off: usize,
317 }
318 
319 impl<'a> OctetsMut<'a> {
320     /// Creates an `OctetsMut` from the given slice, without copying.
321     ///
322     /// Since there's no copy, the input slice needs to be mutable to allow
323     /// modifications.
with_slice(buf: &'a mut [u8]) -> Self324     pub fn with_slice(buf: &'a mut [u8]) -> Self {
325         OctetsMut { buf, off: 0 }
326     }
327 
328     /// Reads an unsigned 8-bit integer from the current offset and advances
329     /// the buffer.
get_u8(&mut self) -> Result<u8>330     pub fn get_u8(&mut self) -> Result<u8> {
331         get_u!(self, u8, 1)
332     }
333 
334     /// Reads an unsigned 8-bit integer from the current offset without
335     /// advancing the buffer.
peek_u8(&mut self) -> Result<u8>336     pub fn peek_u8(&mut self) -> Result<u8> {
337         peek_u!(self, u8, 1)
338     }
339 
340     /// Writes an unsigned 8-bit integer at the current offset and advances
341     /// the buffer.
put_u8(&mut self, v: u8) -> Result<&mut [u8]>342     pub fn put_u8(&mut self, v: u8) -> Result<&mut [u8]> {
343         put_u!(self, u8, v, 1)
344     }
345 
346     /// Reads an unsigned 16-bit integer in network byte-order from the current
347     /// offset and advances the buffer.
get_u16(&mut self) -> Result<u16>348     pub fn get_u16(&mut self) -> Result<u16> {
349         get_u!(self, u16, 2)
350     }
351 
352     /// Writes an unsigned 16-bit integer in network byte-order at the current
353     /// offset and advances the buffer.
put_u16(&mut self, v: u16) -> Result<&mut [u8]>354     pub fn put_u16(&mut self, v: u16) -> Result<&mut [u8]> {
355         put_u!(self, u16, v, 2)
356     }
357 
358     /// Reads an unsigned 24-bit integer in network byte-order from the current
359     /// offset and advances the buffer.
get_u24(&mut self) -> Result<u32>360     pub fn get_u24(&mut self) -> Result<u32> {
361         get_u!(self, u32, 3)
362     }
363 
364     /// Writes an unsigned 24-bit integer in network byte-order at the current
365     /// offset and advances the buffer.
put_u24(&mut self, v: u32) -> Result<&mut [u8]>366     pub fn put_u24(&mut self, v: u32) -> Result<&mut [u8]> {
367         put_u!(self, u32, v, 3)
368     }
369 
370     /// Reads an unsigned 32-bit integer in network byte-order from the current
371     /// offset and advances the buffer.
get_u32(&mut self) -> Result<u32>372     pub fn get_u32(&mut self) -> Result<u32> {
373         get_u!(self, u32, 4)
374     }
375 
376     /// Writes an unsigned 32-bit integer in network byte-order at the current
377     /// offset and advances the buffer.
put_u32(&mut self, v: u32) -> Result<&mut [u8]>378     pub fn put_u32(&mut self, v: u32) -> Result<&mut [u8]> {
379         put_u!(self, u32, v, 4)
380     }
381 
382     /// Reads an unsigned 64-bit integer in network byte-order from the current
383     /// offset and advances the buffer.
get_u64(&mut self) -> Result<u64>384     pub fn get_u64(&mut self) -> Result<u64> {
385         get_u!(self, u64, 8)
386     }
387 
388     /// Writes an unsigned 64-bit integer in network byte-order at the current
389     /// offset and advances the buffer.
put_u64(&mut self, v: u64) -> Result<&mut [u8]>390     pub fn put_u64(&mut self, v: u64) -> Result<&mut [u8]> {
391         put_u!(self, u64, v, 8)
392     }
393 
394     /// Reads an unsigned variable-length integer in network byte-order from
395     /// the current offset and advances the buffer.
get_varint(&mut self) -> Result<u64>396     pub fn get_varint(&mut self) -> Result<u64> {
397         let first = self.peek_u8()?;
398 
399         let len = varint_parse_len(first);
400 
401         if len > self.cap() {
402             return Err(BufferTooShortError);
403         }
404 
405         let mut vec = self.get_bytes(len)?.to_vec();
406 
407         // Mask the 2 most significant bits to remove the encoded length.
408         vec[0] &= 0x3f;
409 
410         let mut b = OctetsMut::with_slice(&mut vec);
411 
412         let out = match len {
413             1 => u64::from(b.get_u8()?),
414             2 => u64::from(b.get_u16()?),
415             4 => u64::from(b.get_u32()?),
416             8 => b.get_u64()?,
417             _ => unreachable!(),
418         };
419 
420         Ok(out)
421     }
422 
423     /// Writes an unsigned variable-length integer in network byte-order at the
424     /// current offset and advances the buffer.
put_varint(&mut self, v: u64) -> Result<&mut [u8]>425     pub fn put_varint(&mut self, v: u64) -> Result<&mut [u8]> {
426         self.put_varint_with_len(v, varint_len(v))
427     }
428 
429     /// Writes an unsigned variable-length integer of the specified length, in
430     /// network byte-order at the current offset and advances the buffer.
put_varint_with_len( &mut self, v: u64, len: usize, ) -> Result<&mut [u8]>431     pub fn put_varint_with_len(
432         &mut self, v: u64, len: usize,
433     ) -> Result<&mut [u8]> {
434         if self.cap() < len {
435             return Err(BufferTooShortError);
436         }
437 
438         let buf = match len {
439             1 => self.put_u8(v as u8)?,
440 
441             2 => {
442                 let buf = self.put_u16(v as u16)?;
443                 buf[0] |= 0x40;
444                 buf
445             },
446 
447             4 => {
448                 let buf = self.put_u32(v as u32)?;
449                 buf[0] |= 0x80;
450                 buf
451             },
452 
453             8 => {
454                 let buf = self.put_u64(v)?;
455                 buf[0] |= 0xc0;
456                 buf
457             },
458 
459             _ => panic!("value is too large for varint"),
460         };
461 
462         Ok(buf)
463     }
464 
465     /// Reads `len` bytes from the current offset without copying and advances
466     /// the buffer.
get_bytes(&mut self, len: usize) -> Result<Octets>467     pub fn get_bytes(&mut self, len: usize) -> Result<Octets> {
468         if self.cap() < len {
469             return Err(BufferTooShortError);
470         }
471 
472         let out = Octets {
473             buf: &self.buf[self.off..self.off + len],
474             off: 0,
475         };
476 
477         self.off += len;
478 
479         Ok(out)
480     }
481 
482     /// Reads `len` bytes from the current offset without copying and advances
483     /// the buffer.
get_bytes_mut(&mut self, len: usize) -> Result<OctetsMut>484     pub fn get_bytes_mut(&mut self, len: usize) -> Result<OctetsMut> {
485         if self.cap() < len {
486             return Err(BufferTooShortError);
487         }
488 
489         let out = OctetsMut {
490             buf: &mut self.buf[self.off..self.off + len],
491             off: 0,
492         };
493 
494         self.off += len;
495 
496         Ok(out)
497     }
498 
499     /// Reads `len` bytes from the current offset without copying and advances
500     /// the buffer, where `len` is an unsigned 8-bit integer prefix.
get_bytes_with_u8_length(&mut self) -> Result<Octets>501     pub fn get_bytes_with_u8_length(&mut self) -> Result<Octets> {
502         let len = self.get_u8()?;
503         self.get_bytes(len as usize)
504     }
505 
506     /// Reads `len` bytes from the current offset without copying and advances
507     /// the buffer, where `len` is an unsigned 16-bit integer prefix in network
508     /// byte-order.
get_bytes_with_u16_length(&mut self) -> Result<Octets>509     pub fn get_bytes_with_u16_length(&mut self) -> Result<Octets> {
510         let len = self.get_u16()?;
511         self.get_bytes(len as usize)
512     }
513 
514     /// Reads `len` bytes from the current offset without copying and advances
515     /// the buffer, where `len` is an unsigned variable-length integer prefix
516     /// in network byte-order.
get_bytes_with_varint_length(&mut self) -> Result<Octets>517     pub fn get_bytes_with_varint_length(&mut self) -> Result<Octets> {
518         let len = self.get_varint()?;
519         self.get_bytes(len as usize)
520     }
521 
522     /// Reads `len` bytes from the current offset without copying and without
523     /// advancing the buffer.
peek_bytes(&mut self, len: usize) -> Result<Octets>524     pub fn peek_bytes(&mut self, len: usize) -> Result<Octets> {
525         if self.cap() < len {
526             return Err(BufferTooShortError);
527         }
528 
529         let out = Octets {
530             buf: &self.buf[self.off..self.off + len],
531             off: 0,
532         };
533 
534         Ok(out)
535     }
536 
537     /// Reads `len` bytes from the current offset without copying and without
538     /// advancing the buffer.
peek_bytes_mut(&mut self, len: usize) -> Result<OctetsMut>539     pub fn peek_bytes_mut(&mut self, len: usize) -> Result<OctetsMut> {
540         if self.cap() < len {
541             return Err(BufferTooShortError);
542         }
543 
544         let out = OctetsMut {
545             buf: &mut self.buf[self.off..self.off + len],
546             off: 0,
547         };
548 
549         Ok(out)
550     }
551 
552     /// Writes `len` bytes from the current offset without copying and advances
553     /// the buffer.
put_bytes(&mut self, v: &[u8]) -> Result<()>554     pub fn put_bytes(&mut self, v: &[u8]) -> Result<()> {
555         let len = v.len();
556 
557         if self.cap() < len {
558             return Err(BufferTooShortError);
559         }
560 
561         if len == 0 {
562             return Ok(());
563         }
564 
565         self.as_mut()[..len].copy_from_slice(v);
566 
567         self.off += len;
568 
569         Ok(())
570     }
571 
572     /// Splits the buffer in two at the given absolute offset.
split_at(&mut self, off: usize) -> Result<(OctetsMut, OctetsMut)>573     pub fn split_at(&mut self, off: usize) -> Result<(OctetsMut, OctetsMut)> {
574         if self.len() < off {
575             return Err(BufferTooShortError);
576         }
577 
578         let (left, right) = self.buf.split_at_mut(off);
579 
580         let first = OctetsMut { buf: left, off: 0 };
581 
582         let last = OctetsMut { buf: right, off: 0 };
583 
584         Ok((first, last))
585     }
586 
587     /// Returns a slice of `len` elements from the current offset.
slice(&'a mut self, len: usize) -> Result<&'a mut [u8]>588     pub fn slice(&'a mut self, len: usize) -> Result<&'a mut [u8]> {
589         if len > self.cap() {
590             return Err(BufferTooShortError);
591         }
592 
593         Ok(&mut self.buf[self.off..self.off + len])
594     }
595 
596     /// Returns a slice of `len` elements from the end of the buffer.
slice_last(&'a mut self, len: usize) -> Result<&'a mut [u8]>597     pub fn slice_last(&'a mut self, len: usize) -> Result<&'a mut [u8]> {
598         if len > self.cap() {
599             return Err(BufferTooShortError);
600         }
601 
602         let cap = self.cap();
603         Ok(&mut self.buf[cap - len..])
604     }
605 
606     /// Returns the remaining capacity in the buffer.
cap(&self) -> usize607     pub fn cap(&self) -> usize {
608         self.buf.len() - self.off
609     }
610 
611     /// Returns the total length of the buffer.
len(&self) -> usize612     pub fn len(&self) -> usize {
613         self.buf.len()
614     }
615 
616     /// Returns the current offset of the buffer.
off(&self) -> usize617     pub fn off(&self) -> usize {
618         self.off
619     }
620 
621     /// Returns a reference to the internal buffer.
buf(&self) -> &[u8]622     pub fn buf(&self) -> &[u8] {
623         self.buf
624     }
625 
626     /// Copies the buffer from the current offset into a new `Vec<u8>`.
to_vec(&self) -> Vec<u8>627     pub fn to_vec(&self) -> Vec<u8> {
628         self.as_ref().to_vec()
629     }
630 }
631 
632 impl<'a> AsRef<[u8]> for OctetsMut<'a> {
as_ref(&self) -> &[u8]633     fn as_ref(&self) -> &[u8] {
634         &self.buf[self.off..]
635     }
636 }
637 
638 impl<'a> AsMut<[u8]> for OctetsMut<'a> {
as_mut(&mut self) -> &mut [u8]639     fn as_mut(&mut self) -> &mut [u8] {
640         &mut self.buf[self.off..]
641     }
642 }
643 
644 /// Returns how many bytes it would take to encode `v` as a variable-length
645 /// integer.
varint_len(v: u64) -> usize646 pub fn varint_len(v: u64) -> usize {
647     if v <= 63 {
648         1
649     } else if v <= 16383 {
650         2
651     } else if v <= 1_073_741_823 {
652         4
653     } else if v <= 4_611_686_018_427_387_903 {
654         8
655     } else {
656         unreachable!()
657     }
658 }
659 
660 /// Returns how long the variable-length integer is, given its first byte.
varint_parse_len(first: u8) -> usize661 pub fn varint_parse_len(first: u8) -> usize {
662     match first >> 6 {
663         0 => 1,
664         1 => 2,
665         2 => 4,
666         3 => 8,
667         _ => unreachable!(),
668     }
669 }
670 
671 #[cfg(test)]
672 mod tests {
673     use super::*;
674 
675     #[test]
get_u()676     fn get_u() {
677         let d = [
678             1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
679         ];
680 
681         let mut b = Octets::with_slice(&d);
682         assert_eq!(b.cap(), 18);
683         assert_eq!(b.off(), 0);
684 
685         assert_eq!(b.get_u8().unwrap(), 1);
686         assert_eq!(b.cap(), 17);
687         assert_eq!(b.off(), 1);
688 
689         assert_eq!(b.get_u16().unwrap(), 0x203);
690         assert_eq!(b.cap(), 15);
691         assert_eq!(b.off(), 3);
692 
693         assert_eq!(b.get_u24().unwrap(), 0x40506);
694         assert_eq!(b.cap(), 12);
695         assert_eq!(b.off(), 6);
696 
697         assert_eq!(b.get_u32().unwrap(), 0x0708090a);
698         assert_eq!(b.cap(), 8);
699         assert_eq!(b.off(), 10);
700 
701         assert_eq!(b.get_u64().unwrap(), 0x0b0c0d0e0f101112);
702         assert_eq!(b.cap(), 0);
703         assert_eq!(b.off(), 18);
704 
705         assert!(b.get_u8().is_err());
706         assert!(b.get_u16().is_err());
707         assert!(b.get_u24().is_err());
708         assert!(b.get_u32().is_err());
709         assert!(b.get_u64().is_err());
710     }
711 
712     #[test]
get_u_mut()713     fn get_u_mut() {
714         let mut d = [
715             1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
716         ];
717 
718         let mut b = OctetsMut::with_slice(&mut d);
719         assert_eq!(b.cap(), 18);
720         assert_eq!(b.off(), 0);
721 
722         assert_eq!(b.get_u8().unwrap(), 1);
723         assert_eq!(b.cap(), 17);
724         assert_eq!(b.off(), 1);
725 
726         assert_eq!(b.get_u16().unwrap(), 0x203);
727         assert_eq!(b.cap(), 15);
728         assert_eq!(b.off(), 3);
729 
730         assert_eq!(b.get_u24().unwrap(), 0x40506);
731         assert_eq!(b.cap(), 12);
732         assert_eq!(b.off(), 6);
733 
734         assert_eq!(b.get_u32().unwrap(), 0x0708090a);
735         assert_eq!(b.cap(), 8);
736         assert_eq!(b.off(), 10);
737 
738         assert_eq!(b.get_u64().unwrap(), 0x0b0c0d0e0f101112);
739         assert_eq!(b.cap(), 0);
740         assert_eq!(b.off(), 18);
741 
742         assert!(b.get_u8().is_err());
743         assert!(b.get_u16().is_err());
744         assert!(b.get_u24().is_err());
745         assert!(b.get_u32().is_err());
746         assert!(b.get_u64().is_err());
747     }
748 
749     #[test]
peek_u()750     fn peek_u() {
751         let d = [1, 2];
752 
753         let mut b = Octets::with_slice(&d);
754         assert_eq!(b.cap(), 2);
755         assert_eq!(b.off(), 0);
756 
757         assert_eq!(b.peek_u8().unwrap(), 1);
758         assert_eq!(b.cap(), 2);
759         assert_eq!(b.off(), 0);
760 
761         assert_eq!(b.peek_u8().unwrap(), 1);
762         assert_eq!(b.cap(), 2);
763         assert_eq!(b.off(), 0);
764 
765         b.get_u16().unwrap();
766 
767         assert!(b.peek_u8().is_err());
768     }
769 
770     #[test]
peek_u_mut()771     fn peek_u_mut() {
772         let mut d = [1, 2];
773 
774         let mut b = OctetsMut::with_slice(&mut d);
775         assert_eq!(b.cap(), 2);
776         assert_eq!(b.off(), 0);
777 
778         assert_eq!(b.peek_u8().unwrap(), 1);
779         assert_eq!(b.cap(), 2);
780         assert_eq!(b.off(), 0);
781 
782         assert_eq!(b.peek_u8().unwrap(), 1);
783         assert_eq!(b.cap(), 2);
784         assert_eq!(b.off(), 0);
785 
786         b.get_u16().unwrap();
787 
788         assert!(b.peek_u8().is_err());
789     }
790 
791     #[test]
get_bytes()792     fn get_bytes() {
793         let d = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
794         let mut b = Octets::with_slice(&d);
795         assert_eq!(b.cap(), 10);
796         assert_eq!(b.off(), 0);
797 
798         assert_eq!(b.get_bytes(5).unwrap().as_ref(), [1, 2, 3, 4, 5]);
799         assert_eq!(b.cap(), 5);
800         assert_eq!(b.off(), 5);
801 
802         assert_eq!(b.get_bytes(3).unwrap().as_ref(), [6, 7, 8]);
803         assert_eq!(b.cap(), 2);
804         assert_eq!(b.off(), 8);
805 
806         assert!(b.get_bytes(3).is_err());
807         assert_eq!(b.cap(), 2);
808         assert_eq!(b.off(), 8);
809 
810         assert_eq!(b.get_bytes(2).unwrap().as_ref(), [9, 10]);
811         assert_eq!(b.cap(), 0);
812         assert_eq!(b.off(), 10);
813 
814         assert!(b.get_bytes(2).is_err());
815     }
816 
817     #[test]
get_bytes_mut()818     fn get_bytes_mut() {
819         let mut d = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
820         let mut b = OctetsMut::with_slice(&mut d);
821         assert_eq!(b.cap(), 10);
822         assert_eq!(b.off(), 0);
823 
824         assert_eq!(b.get_bytes(5).unwrap().as_ref(), [1, 2, 3, 4, 5]);
825         assert_eq!(b.cap(), 5);
826         assert_eq!(b.off(), 5);
827 
828         assert_eq!(b.get_bytes(3).unwrap().as_ref(), [6, 7, 8]);
829         assert_eq!(b.cap(), 2);
830         assert_eq!(b.off(), 8);
831 
832         assert!(b.get_bytes(3).is_err());
833         assert_eq!(b.cap(), 2);
834         assert_eq!(b.off(), 8);
835 
836         assert_eq!(b.get_bytes(2).unwrap().as_ref(), [9, 10]);
837         assert_eq!(b.cap(), 0);
838         assert_eq!(b.off(), 10);
839 
840         assert!(b.get_bytes(2).is_err());
841     }
842 
843     #[test]
peek_bytes()844     fn peek_bytes() {
845         let d = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
846         let mut b = Octets::with_slice(&d);
847         assert_eq!(b.cap(), 10);
848         assert_eq!(b.off(), 0);
849 
850         assert_eq!(b.peek_bytes(5).unwrap().as_ref(), [1, 2, 3, 4, 5]);
851         assert_eq!(b.cap(), 10);
852         assert_eq!(b.off(), 0);
853 
854         assert_eq!(b.peek_bytes(5).unwrap().as_ref(), [1, 2, 3, 4, 5]);
855         assert_eq!(b.cap(), 10);
856         assert_eq!(b.off(), 0);
857 
858         b.get_bytes(5).unwrap();
859     }
860 
861     #[test]
peek_bytes_mut()862     fn peek_bytes_mut() {
863         let mut d = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
864         let mut b = OctetsMut::with_slice(&mut d);
865         assert_eq!(b.cap(), 10);
866         assert_eq!(b.off(), 0);
867 
868         assert_eq!(b.peek_bytes(5).unwrap().as_ref(), [1, 2, 3, 4, 5]);
869         assert_eq!(b.cap(), 10);
870         assert_eq!(b.off(), 0);
871 
872         assert_eq!(b.peek_bytes(5).unwrap().as_ref(), [1, 2, 3, 4, 5]);
873         assert_eq!(b.cap(), 10);
874         assert_eq!(b.off(), 0);
875 
876         b.get_bytes(5).unwrap();
877     }
878 
879     #[test]
get_varint()880     fn get_varint() {
881         let d = [0xc2, 0x19, 0x7c, 0x5e, 0xff, 0x14, 0xe8, 0x8c];
882         let mut b = Octets::with_slice(&d);
883         assert_eq!(b.get_varint().unwrap(), 151288809941952652);
884         assert_eq!(b.cap(), 0);
885         assert_eq!(b.off(), 8);
886 
887         let d = [0x9d, 0x7f, 0x3e, 0x7d];
888         let mut b = Octets::with_slice(&d);
889         assert_eq!(b.get_varint().unwrap(), 494878333);
890         assert_eq!(b.cap(), 0);
891         assert_eq!(b.off(), 4);
892 
893         let d = [0x7b, 0xbd];
894         let mut b = Octets::with_slice(&d);
895         assert_eq!(b.get_varint().unwrap(), 15293);
896         assert_eq!(b.cap(), 0);
897         assert_eq!(b.off(), 2);
898 
899         let d = [0x40, 0x25];
900         let mut b = Octets::with_slice(&d);
901         assert_eq!(b.get_varint().unwrap(), 37);
902         assert_eq!(b.cap(), 0);
903         assert_eq!(b.off(), 2);
904 
905         let d = [0x25];
906         let mut b = Octets::with_slice(&d);
907         assert_eq!(b.get_varint().unwrap(), 37);
908         assert_eq!(b.cap(), 0);
909         assert_eq!(b.off(), 1);
910     }
911 
912     #[test]
get_varint_mut()913     fn get_varint_mut() {
914         let mut d = [0xc2, 0x19, 0x7c, 0x5e, 0xff, 0x14, 0xe8, 0x8c];
915         let mut b = OctetsMut::with_slice(&mut d);
916         assert_eq!(b.get_varint().unwrap(), 151288809941952652);
917         assert_eq!(b.cap(), 0);
918         assert_eq!(b.off(), 8);
919 
920         let mut d = [0x9d, 0x7f, 0x3e, 0x7d];
921         let mut b = OctetsMut::with_slice(&mut d);
922         assert_eq!(b.get_varint().unwrap(), 494878333);
923         assert_eq!(b.cap(), 0);
924         assert_eq!(b.off(), 4);
925 
926         let mut d = [0x7b, 0xbd];
927         let mut b = OctetsMut::with_slice(&mut d);
928         assert_eq!(b.get_varint().unwrap(), 15293);
929         assert_eq!(b.cap(), 0);
930         assert_eq!(b.off(), 2);
931 
932         let mut d = [0x40, 0x25];
933         let mut b = OctetsMut::with_slice(&mut d);
934         assert_eq!(b.get_varint().unwrap(), 37);
935         assert_eq!(b.cap(), 0);
936         assert_eq!(b.off(), 2);
937 
938         let mut d = [0x25];
939         let mut b = OctetsMut::with_slice(&mut d);
940         assert_eq!(b.get_varint().unwrap(), 37);
941         assert_eq!(b.cap(), 0);
942         assert_eq!(b.off(), 1);
943     }
944 
945     #[test]
put_varint()946     fn put_varint() {
947         let mut d = [0; 8];
948         {
949             let mut b = OctetsMut::with_slice(&mut d);
950             assert!(b.put_varint(151288809941952652).is_ok());
951             assert_eq!(b.cap(), 0);
952             assert_eq!(b.off(), 8);
953         }
954         let exp = [0xc2, 0x19, 0x7c, 0x5e, 0xff, 0x14, 0xe8, 0x8c];
955         assert_eq!(&d, &exp);
956 
957         let mut d = [0; 4];
958         {
959             let mut b = OctetsMut::with_slice(&mut d);
960             assert!(b.put_varint(494878333).is_ok());
961             assert_eq!(b.cap(), 0);
962             assert_eq!(b.off(), 4);
963         }
964         let exp = [0x9d, 0x7f, 0x3e, 0x7d];
965         assert_eq!(&d, &exp);
966 
967         let mut d = [0; 2];
968         {
969             let mut b = OctetsMut::with_slice(&mut d);
970             assert!(b.put_varint(15293).is_ok());
971             assert_eq!(b.cap(), 0);
972             assert_eq!(b.off(), 2);
973         }
974         let exp = [0x7b, 0xbd];
975         assert_eq!(&d, &exp);
976 
977         let mut d = [0; 1];
978         {
979             let mut b = OctetsMut::with_slice(&mut d);
980             assert!(b.put_varint(37).is_ok());
981             assert_eq!(b.cap(), 0);
982             assert_eq!(b.off(), 1);
983         }
984         let exp = [0x25];
985         assert_eq!(&d, &exp);
986 
987         let mut d = [0; 3];
988         {
989             let mut b = OctetsMut::with_slice(&mut d);
990             assert!(b.put_varint(151288809941952652).is_err());
991             assert_eq!(b.cap(), 3);
992             assert_eq!(b.off(), 0);
993         }
994         let exp = [0; 3];
995         assert_eq!(&d, &exp);
996     }
997 
998     #[test]
999     #[should_panic]
varint_too_large()1000     fn varint_too_large() {
1001         let mut d = [0; 3];
1002         let mut b = OctetsMut::with_slice(&mut d);
1003         assert!(b.put_varint(std::u64::MAX).is_err());
1004     }
1005 
1006     #[test]
put_u()1007     fn put_u() {
1008         let mut d = [0; 18];
1009 
1010         {
1011             let mut b = OctetsMut::with_slice(&mut d);
1012             assert_eq!(b.cap(), 18);
1013             assert_eq!(b.off(), 0);
1014 
1015             assert!(b.put_u8(1).is_ok());
1016             assert_eq!(b.cap(), 17);
1017             assert_eq!(b.off(), 1);
1018 
1019             assert!(b.put_u16(0x203).is_ok());
1020             assert_eq!(b.cap(), 15);
1021             assert_eq!(b.off(), 3);
1022 
1023             assert!(b.put_u24(0x40506).is_ok());
1024             assert_eq!(b.cap(), 12);
1025             assert_eq!(b.off(), 6);
1026 
1027             assert!(b.put_u32(0x0708090a).is_ok());
1028             assert_eq!(b.cap(), 8);
1029             assert_eq!(b.off(), 10);
1030 
1031             assert!(b.put_u64(0x0b0c0d0e0f101112).is_ok());
1032             assert_eq!(b.cap(), 0);
1033             assert_eq!(b.off(), 18);
1034 
1035             assert!(b.put_u8(1).is_err());
1036         }
1037 
1038         let exp = [
1039             1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
1040         ];
1041         assert_eq!(&d, &exp);
1042     }
1043 
1044     #[test]
put_bytes()1045     fn put_bytes() {
1046         let mut d = [0; 5];
1047 
1048         {
1049             let mut b = OctetsMut::with_slice(&mut d);
1050             assert_eq!(b.cap(), 5);
1051             assert_eq!(b.off(), 0);
1052 
1053             let p = [0x0a, 0x0b, 0x0c, 0x0d, 0x0e];
1054             assert!(b.put_bytes(&p).is_ok());
1055             assert_eq!(b.cap(), 0);
1056             assert_eq!(b.off(), 5);
1057 
1058             assert!(b.put_u8(1).is_err());
1059         }
1060 
1061         let exp = [0xa, 0xb, 0xc, 0xd, 0xe];
1062         assert_eq!(&d, &exp);
1063     }
1064 
1065     #[test]
split()1066     fn split() {
1067         let mut d = b"helloworld".to_vec();
1068 
1069         let mut b = OctetsMut::with_slice(&mut d);
1070         assert_eq!(b.cap(), 10);
1071         assert_eq!(b.off(), 0);
1072         assert_eq!(b.as_ref(), b"helloworld");
1073 
1074         assert!(b.get_bytes(5).is_ok());
1075         assert_eq!(b.cap(), 5);
1076         assert_eq!(b.off(), 5);
1077         assert_eq!(b.as_ref(), b"world");
1078 
1079         let off = b.off();
1080 
1081         let (first, last) = b.split_at(off).unwrap();
1082         assert_eq!(first.cap(), 5);
1083         assert_eq!(first.off(), 0);
1084         assert_eq!(first.as_ref(), b"hello");
1085 
1086         assert_eq!(last.cap(), 5);
1087         assert_eq!(last.off(), 0);
1088         assert_eq!(last.as_ref(), b"world");
1089     }
1090 
1091     #[test]
split_at()1092     fn split_at() {
1093         let mut d = b"helloworld".to_vec();
1094 
1095         {
1096             let mut b = OctetsMut::with_slice(&mut d);
1097             let (first, second) = b.split_at(5).unwrap();
1098 
1099             let mut exp1 = b"hello".to_vec();
1100             assert_eq!(first.as_ref(), &mut exp1[..]);
1101 
1102             let mut exp2 = b"world".to_vec();
1103             assert_eq!(second.as_ref(), &mut exp2[..]);
1104         }
1105 
1106         {
1107             let mut b = OctetsMut::with_slice(&mut d);
1108             let (first, second) = b.split_at(10).unwrap();
1109 
1110             let mut exp1 = b"helloworld".to_vec();
1111             assert_eq!(first.as_ref(), &mut exp1[..]);
1112 
1113             let mut exp2 = b"".to_vec();
1114             assert_eq!(second.as_ref(), &mut exp2[..]);
1115         }
1116 
1117         {
1118             let mut b = OctetsMut::with_slice(&mut d);
1119             let (first, second) = b.split_at(9).unwrap();
1120 
1121             let mut exp1 = b"helloworl".to_vec();
1122             assert_eq!(first.as_ref(), &mut exp1[..]);
1123 
1124             let mut exp2 = b"d".to_vec();
1125             assert_eq!(second.as_ref(), &mut exp2[..]);
1126         }
1127 
1128         {
1129             let mut b = OctetsMut::with_slice(&mut d);
1130             assert!(b.split_at(11).is_err());
1131         }
1132     }
1133 
1134     #[test]
slice()1135     fn slice() {
1136         let d = b"helloworld".to_vec();
1137 
1138         {
1139             let b = Octets::with_slice(&d);
1140             let exp = b"hello".to_vec();
1141             assert_eq!(b.slice(5), Ok(&exp[..]));
1142         }
1143 
1144         {
1145             let b = Octets::with_slice(&d);
1146             let exp = b"".to_vec();
1147             assert_eq!(b.slice(0), Ok(&exp[..]));
1148         }
1149 
1150         {
1151             let mut b = Octets::with_slice(&d);
1152             b.get_bytes(5).unwrap();
1153 
1154             let exp = b"world".to_vec();
1155             assert_eq!(b.slice(5), Ok(&exp[..]));
1156         }
1157 
1158         {
1159             let b = Octets::with_slice(&d);
1160             assert!(b.slice(11).is_err());
1161         }
1162     }
1163 
1164     #[test]
slice_mut()1165     fn slice_mut() {
1166         let mut d = b"helloworld".to_vec();
1167 
1168         {
1169             let mut b = OctetsMut::with_slice(&mut d);
1170             let mut exp = b"hello".to_vec();
1171             assert_eq!(b.slice(5), Ok(&mut exp[..]));
1172         }
1173 
1174         {
1175             let mut b = OctetsMut::with_slice(&mut d);
1176             let mut exp = b"".to_vec();
1177             assert_eq!(b.slice(0), Ok(&mut exp[..]));
1178         }
1179 
1180         {
1181             let mut b = OctetsMut::with_slice(&mut d);
1182             b.get_bytes(5).unwrap();
1183 
1184             let mut exp = b"world".to_vec();
1185             assert_eq!(b.slice(5), Ok(&mut exp[..]));
1186         }
1187 
1188         {
1189             let mut b = OctetsMut::with_slice(&mut d);
1190             assert!(b.slice(11).is_err());
1191         }
1192     }
1193 
1194     #[test]
slice_last()1195     fn slice_last() {
1196         let d = b"helloworld".to_vec();
1197 
1198         {
1199             let b = Octets::with_slice(&d);
1200             let exp = b"orld".to_vec();
1201             assert_eq!(b.slice_last(4), Ok(&exp[..]));
1202         }
1203 
1204         {
1205             let b = Octets::with_slice(&d);
1206             let exp = b"d".to_vec();
1207             assert_eq!(b.slice_last(1), Ok(&exp[..]));
1208         }
1209 
1210         {
1211             let b = Octets::with_slice(&d);
1212             let exp = b"".to_vec();
1213             assert_eq!(b.slice_last(0), Ok(&exp[..]));
1214         }
1215 
1216         {
1217             let b = Octets::with_slice(&d);
1218             let exp = b"helloworld".to_vec();
1219             assert_eq!(b.slice_last(10), Ok(&exp[..]));
1220         }
1221 
1222         {
1223             let b = Octets::with_slice(&d);
1224             assert!(b.slice_last(11).is_err());
1225         }
1226     }
1227 
1228     #[test]
slice_last_mut()1229     fn slice_last_mut() {
1230         let mut d = b"helloworld".to_vec();
1231 
1232         {
1233             let mut b = OctetsMut::with_slice(&mut d);
1234             let mut exp = b"orld".to_vec();
1235             assert_eq!(b.slice_last(4), Ok(&mut exp[..]));
1236         }
1237 
1238         {
1239             let mut b = OctetsMut::with_slice(&mut d);
1240             let mut exp = b"d".to_vec();
1241             assert_eq!(b.slice_last(1), Ok(&mut exp[..]));
1242         }
1243 
1244         {
1245             let mut b = OctetsMut::with_slice(&mut d);
1246             let mut exp = b"".to_vec();
1247             assert_eq!(b.slice_last(0), Ok(&mut exp[..]));
1248         }
1249 
1250         {
1251             let mut b = OctetsMut::with_slice(&mut d);
1252             let mut exp = b"helloworld".to_vec();
1253             assert_eq!(b.slice_last(10), Ok(&mut exp[..]));
1254         }
1255 
1256         {
1257             let mut b = OctetsMut::with_slice(&mut d);
1258             assert!(b.slice_last(11).is_err());
1259         }
1260     }
1261 }
1262