• 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 out = match len {
187             1 => u64::from(self.get_u8()?),
188 
189             2 => u64::from(self.get_u16()? & 0x3fff),
190 
191             4 => u64::from(self.get_u32()? & 0x3fffffff),
192 
193             8 => self.get_u64()? & 0x3fffffffffffffff,
194 
195             _ => unreachable!(),
196         };
197 
198         Ok(out)
199     }
200 
201     /// Reads `len` bytes from the current offset without copying and advances
202     /// the buffer.
get_bytes(&mut self, len: usize) -> Result<Octets>203     pub fn get_bytes(&mut self, len: usize) -> Result<Octets> {
204         if self.cap() < len {
205             return Err(BufferTooShortError);
206         }
207 
208         let out = Octets {
209             buf: &self.buf[self.off..self.off + len],
210             off: 0,
211         };
212 
213         self.off += len;
214 
215         Ok(out)
216     }
217 
218     /// Reads `len` bytes from the current offset without copying and advances
219     /// the buffer, where `len` is an unsigned 8-bit integer prefix.
get_bytes_with_u8_length(&mut self) -> Result<Octets>220     pub fn get_bytes_with_u8_length(&mut self) -> Result<Octets> {
221         let len = self.get_u8()?;
222         self.get_bytes(len as usize)
223     }
224 
225     /// Reads `len` bytes from the current offset without copying and advances
226     /// the buffer, where `len` is an unsigned 16-bit integer prefix in network
227     /// byte-order.
get_bytes_with_u16_length(&mut self) -> Result<Octets>228     pub fn get_bytes_with_u16_length(&mut self) -> Result<Octets> {
229         let len = self.get_u16()?;
230         self.get_bytes(len as usize)
231     }
232 
233     /// Reads `len` bytes from the current offset without copying and advances
234     /// the buffer, where `len` is an unsigned variable-length integer prefix
235     /// in network byte-order.
get_bytes_with_varint_length(&mut self) -> Result<Octets>236     pub fn get_bytes_with_varint_length(&mut self) -> Result<Octets> {
237         let len = self.get_varint()?;
238         self.get_bytes(len as usize)
239     }
240 
241     /// Reads `len` bytes from the current offset without copying and without
242     /// advancing the buffer.
peek_bytes(&self, len: usize) -> Result<Octets>243     pub fn peek_bytes(&self, len: usize) -> Result<Octets> {
244         if self.cap() < len {
245             return Err(BufferTooShortError);
246         }
247 
248         let out = Octets {
249             buf: &self.buf[self.off..self.off + len],
250             off: 0,
251         };
252 
253         Ok(out)
254     }
255 
256     /// Returns a slice of `len` elements from the current offset.
slice(&'a self, len: usize) -> Result<&'a [u8]>257     pub fn slice(&'a self, len: usize) -> Result<&'a [u8]> {
258         if len > self.cap() {
259             return Err(BufferTooShortError);
260         }
261 
262         Ok(&self.buf[self.off..self.off + len])
263     }
264 
265     /// Returns a slice of `len` elements from the end of the buffer.
slice_last(&'a self, len: usize) -> Result<&'a [u8]>266     pub fn slice_last(&'a self, len: usize) -> Result<&'a [u8]> {
267         if len > self.cap() {
268             return Err(BufferTooShortError);
269         }
270 
271         let cap = self.cap();
272         Ok(&self.buf[cap - len..])
273     }
274 
275     /// Advances the buffer's offset.
skip(&mut self, skip: usize) -> Result<()>276     pub fn skip(&mut self, skip: usize) -> Result<()> {
277         if skip > self.cap() {
278             return Err(BufferTooShortError);
279         }
280 
281         self.off += skip;
282 
283         Ok(())
284     }
285 
286     /// Returns the remaining capacity in the buffer.
cap(&self) -> usize287     pub fn cap(&self) -> usize {
288         self.buf.len() - self.off
289     }
290 
291     /// Returns the total length of the buffer.
len(&self) -> usize292     pub fn len(&self) -> usize {
293         self.buf.len()
294     }
295 
296     /// Returns the current offset of the buffer.
off(&self) -> usize297     pub fn off(&self) -> usize {
298         self.off
299     }
300 
301     /// Returns a reference to the internal buffer.
buf(&self) -> &[u8]302     pub fn buf(&self) -> &[u8] {
303         self.buf
304     }
305 
306     /// Copies the buffer from the current offset into a new `Vec<u8>`.
to_vec(&self) -> Vec<u8>307     pub fn to_vec(&self) -> Vec<u8> {
308         self.as_ref().to_vec()
309     }
310 }
311 
312 impl<'a> AsRef<[u8]> for Octets<'a> {
as_ref(&self) -> &[u8]313     fn as_ref(&self) -> &[u8] {
314         &self.buf[self.off..]
315     }
316 }
317 
318 /// A zero-copy mutable byte buffer.
319 ///
320 /// Like `Octets` but mutable.
321 #[derive(Debug, PartialEq)]
322 pub struct OctetsMut<'a> {
323     buf: &'a mut [u8],
324     off: usize,
325 }
326 
327 impl<'a> OctetsMut<'a> {
328     /// Creates an `OctetsMut` from the given slice, without copying.
329     ///
330     /// Since there's no copy, the input slice needs to be mutable to allow
331     /// modifications.
with_slice(buf: &'a mut [u8]) -> Self332     pub fn with_slice(buf: &'a mut [u8]) -> Self {
333         OctetsMut { buf, off: 0 }
334     }
335 
336     /// Reads an unsigned 8-bit integer from the current offset and advances
337     /// the buffer.
get_u8(&mut self) -> Result<u8>338     pub fn get_u8(&mut self) -> Result<u8> {
339         get_u!(self, u8, 1)
340     }
341 
342     /// Reads an unsigned 8-bit integer from the current offset without
343     /// advancing the buffer.
peek_u8(&mut self) -> Result<u8>344     pub fn peek_u8(&mut self) -> Result<u8> {
345         peek_u!(self, u8, 1)
346     }
347 
348     /// Writes an unsigned 8-bit integer at the current offset and advances
349     /// the buffer.
put_u8(&mut self, v: u8) -> Result<&mut [u8]>350     pub fn put_u8(&mut self, v: u8) -> Result<&mut [u8]> {
351         put_u!(self, u8, v, 1)
352     }
353 
354     /// Reads an unsigned 16-bit integer in network byte-order from the current
355     /// offset and advances the buffer.
get_u16(&mut self) -> Result<u16>356     pub fn get_u16(&mut self) -> Result<u16> {
357         get_u!(self, u16, 2)
358     }
359 
360     /// Writes an unsigned 16-bit integer in network byte-order at the current
361     /// offset and advances the buffer.
put_u16(&mut self, v: u16) -> Result<&mut [u8]>362     pub fn put_u16(&mut self, v: u16) -> Result<&mut [u8]> {
363         put_u!(self, u16, v, 2)
364     }
365 
366     /// Reads an unsigned 24-bit integer in network byte-order from the current
367     /// offset and advances the buffer.
get_u24(&mut self) -> Result<u32>368     pub fn get_u24(&mut self) -> Result<u32> {
369         get_u!(self, u32, 3)
370     }
371 
372     /// Writes an unsigned 24-bit integer in network byte-order at the current
373     /// offset and advances the buffer.
put_u24(&mut self, v: u32) -> Result<&mut [u8]>374     pub fn put_u24(&mut self, v: u32) -> Result<&mut [u8]> {
375         put_u!(self, u32, v, 3)
376     }
377 
378     /// Reads an unsigned 32-bit integer in network byte-order from the current
379     /// offset and advances the buffer.
get_u32(&mut self) -> Result<u32>380     pub fn get_u32(&mut self) -> Result<u32> {
381         get_u!(self, u32, 4)
382     }
383 
384     /// Writes an unsigned 32-bit integer in network byte-order at the current
385     /// offset and advances the buffer.
put_u32(&mut self, v: u32) -> Result<&mut [u8]>386     pub fn put_u32(&mut self, v: u32) -> Result<&mut [u8]> {
387         put_u!(self, u32, v, 4)
388     }
389 
390     /// Reads an unsigned 64-bit integer in network byte-order from the current
391     /// offset and advances the buffer.
get_u64(&mut self) -> Result<u64>392     pub fn get_u64(&mut self) -> Result<u64> {
393         get_u!(self, u64, 8)
394     }
395 
396     /// Writes an unsigned 64-bit integer in network byte-order at the current
397     /// offset and advances the buffer.
put_u64(&mut self, v: u64) -> Result<&mut [u8]>398     pub fn put_u64(&mut self, v: u64) -> Result<&mut [u8]> {
399         put_u!(self, u64, v, 8)
400     }
401 
402     /// Reads an unsigned variable-length integer in network byte-order from
403     /// the current offset and advances the buffer.
get_varint(&mut self) -> Result<u64>404     pub fn get_varint(&mut self) -> Result<u64> {
405         let first = self.peek_u8()?;
406 
407         let len = varint_parse_len(first);
408 
409         if len > self.cap() {
410             return Err(BufferTooShortError);
411         }
412 
413         let out = match len {
414             1 => u64::from(self.get_u8()?),
415 
416             2 => u64::from(self.get_u16()? & 0x3fff),
417 
418             4 => u64::from(self.get_u32()? & 0x3fffffff),
419 
420             8 => self.get_u64()? & 0x3fffffffffffffff,
421 
422             _ => unreachable!(),
423         };
424 
425         Ok(out)
426     }
427 
428     /// Writes an unsigned variable-length integer in network byte-order at the
429     /// current offset and advances the buffer.
put_varint(&mut self, v: u64) -> Result<&mut [u8]>430     pub fn put_varint(&mut self, v: u64) -> Result<&mut [u8]> {
431         self.put_varint_with_len(v, varint_len(v))
432     }
433 
434     /// Writes an unsigned variable-length integer of the specified length, in
435     /// network byte-order at the current offset and advances the buffer.
put_varint_with_len( &mut self, v: u64, len: usize, ) -> Result<&mut [u8]>436     pub fn put_varint_with_len(
437         &mut self, v: u64, len: usize,
438     ) -> Result<&mut [u8]> {
439         if self.cap() < len {
440             return Err(BufferTooShortError);
441         }
442 
443         let buf = match len {
444             1 => self.put_u8(v as u8)?,
445 
446             2 => {
447                 let buf = self.put_u16(v as u16)?;
448                 buf[0] |= 0x40;
449                 buf
450             },
451 
452             4 => {
453                 let buf = self.put_u32(v as u32)?;
454                 buf[0] |= 0x80;
455                 buf
456             },
457 
458             8 => {
459                 let buf = self.put_u64(v)?;
460                 buf[0] |= 0xc0;
461                 buf
462             },
463 
464             _ => panic!("value is too large for varint"),
465         };
466 
467         Ok(buf)
468     }
469 
470     /// Reads `len` bytes from the current offset without copying and advances
471     /// the buffer.
get_bytes(&mut self, len: usize) -> Result<Octets>472     pub fn get_bytes(&mut self, len: usize) -> Result<Octets> {
473         if self.cap() < len {
474             return Err(BufferTooShortError);
475         }
476 
477         let out = Octets {
478             buf: &self.buf[self.off..self.off + len],
479             off: 0,
480         };
481 
482         self.off += len;
483 
484         Ok(out)
485     }
486 
487     /// Reads `len` bytes from the current offset without copying and advances
488     /// the buffer.
get_bytes_mut(&mut self, len: usize) -> Result<OctetsMut>489     pub fn get_bytes_mut(&mut self, len: usize) -> Result<OctetsMut> {
490         if self.cap() < len {
491             return Err(BufferTooShortError);
492         }
493 
494         let out = OctetsMut {
495             buf: &mut self.buf[self.off..self.off + len],
496             off: 0,
497         };
498 
499         self.off += len;
500 
501         Ok(out)
502     }
503 
504     /// Reads `len` bytes from the current offset without copying and advances
505     /// the buffer, where `len` is an unsigned 8-bit integer prefix.
get_bytes_with_u8_length(&mut self) -> Result<Octets>506     pub fn get_bytes_with_u8_length(&mut self) -> Result<Octets> {
507         let len = self.get_u8()?;
508         self.get_bytes(len as usize)
509     }
510 
511     /// Reads `len` bytes from the current offset without copying and advances
512     /// the buffer, where `len` is an unsigned 16-bit integer prefix in network
513     /// byte-order.
get_bytes_with_u16_length(&mut self) -> Result<Octets>514     pub fn get_bytes_with_u16_length(&mut self) -> Result<Octets> {
515         let len = self.get_u16()?;
516         self.get_bytes(len as usize)
517     }
518 
519     /// Reads `len` bytes from the current offset without copying and advances
520     /// the buffer, where `len` is an unsigned variable-length integer prefix
521     /// in network byte-order.
get_bytes_with_varint_length(&mut self) -> Result<Octets>522     pub fn get_bytes_with_varint_length(&mut self) -> Result<Octets> {
523         let len = self.get_varint()?;
524         self.get_bytes(len as usize)
525     }
526 
527     /// Reads `len` bytes from the current offset without copying and without
528     /// advancing the buffer.
peek_bytes(&mut self, len: usize) -> Result<Octets>529     pub fn peek_bytes(&mut self, len: usize) -> Result<Octets> {
530         if self.cap() < len {
531             return Err(BufferTooShortError);
532         }
533 
534         let out = Octets {
535             buf: &self.buf[self.off..self.off + len],
536             off: 0,
537         };
538 
539         Ok(out)
540     }
541 
542     /// Reads `len` bytes from the current offset without copying and without
543     /// advancing the buffer.
peek_bytes_mut(&mut self, len: usize) -> Result<OctetsMut>544     pub fn peek_bytes_mut(&mut self, len: usize) -> Result<OctetsMut> {
545         if self.cap() < len {
546             return Err(BufferTooShortError);
547         }
548 
549         let out = OctetsMut {
550             buf: &mut self.buf[self.off..self.off + len],
551             off: 0,
552         };
553 
554         Ok(out)
555     }
556 
557     /// Writes `len` bytes from the current offset without copying and advances
558     /// the buffer.
put_bytes(&mut self, v: &[u8]) -> Result<()>559     pub fn put_bytes(&mut self, v: &[u8]) -> Result<()> {
560         let len = v.len();
561 
562         if self.cap() < len {
563             return Err(BufferTooShortError);
564         }
565 
566         if len == 0 {
567             return Ok(());
568         }
569 
570         self.as_mut()[..len].copy_from_slice(v);
571 
572         self.off += len;
573 
574         Ok(())
575     }
576 
577     /// Splits the buffer in two at the given absolute offset.
split_at(&mut self, off: usize) -> Result<(OctetsMut, OctetsMut)>578     pub fn split_at(&mut self, off: usize) -> Result<(OctetsMut, OctetsMut)> {
579         if self.len() < off {
580             return Err(BufferTooShortError);
581         }
582 
583         let (left, right) = self.buf.split_at_mut(off);
584 
585         let first = OctetsMut { buf: left, off: 0 };
586 
587         let last = OctetsMut { buf: right, off: 0 };
588 
589         Ok((first, last))
590     }
591 
592     /// Returns a slice of `len` elements from the current offset.
slice(&'a mut self, len: usize) -> Result<&'a mut [u8]>593     pub fn slice(&'a mut self, len: usize) -> Result<&'a mut [u8]> {
594         if len > self.cap() {
595             return Err(BufferTooShortError);
596         }
597 
598         Ok(&mut self.buf[self.off..self.off + len])
599     }
600 
601     /// Returns a slice of `len` elements from the end of the buffer.
slice_last(&'a mut self, len: usize) -> Result<&'a mut [u8]>602     pub fn slice_last(&'a mut self, len: usize) -> Result<&'a mut [u8]> {
603         if len > self.cap() {
604             return Err(BufferTooShortError);
605         }
606 
607         let cap = self.cap();
608         Ok(&mut self.buf[cap - len..])
609     }
610 
611     /// Advances the buffer's offset.
skip(&mut self, skip: usize) -> Result<()>612     pub fn skip(&mut self, skip: usize) -> Result<()> {
613         if skip > self.cap() {
614             return Err(BufferTooShortError);
615         }
616 
617         self.off += skip;
618 
619         Ok(())
620     }
621 
622     /// Returns the remaining capacity in the buffer.
cap(&self) -> usize623     pub fn cap(&self) -> usize {
624         self.buf.len() - self.off
625     }
626 
627     /// Returns the total length of the buffer.
len(&self) -> usize628     pub fn len(&self) -> usize {
629         self.buf.len()
630     }
631 
632     /// Returns the current offset of the buffer.
off(&self) -> usize633     pub fn off(&self) -> usize {
634         self.off
635     }
636 
637     /// Returns a reference to the internal buffer.
buf(&self) -> &[u8]638     pub fn buf(&self) -> &[u8] {
639         self.buf
640     }
641 
642     /// Copies the buffer from the current offset into a new `Vec<u8>`.
to_vec(&self) -> Vec<u8>643     pub fn to_vec(&self) -> Vec<u8> {
644         self.as_ref().to_vec()
645     }
646 }
647 
648 impl<'a> AsRef<[u8]> for OctetsMut<'a> {
as_ref(&self) -> &[u8]649     fn as_ref(&self) -> &[u8] {
650         &self.buf[self.off..]
651     }
652 }
653 
654 impl<'a> AsMut<[u8]> for OctetsMut<'a> {
as_mut(&mut self) -> &mut [u8]655     fn as_mut(&mut self) -> &mut [u8] {
656         &mut self.buf[self.off..]
657     }
658 }
659 
660 /// Returns how many bytes it would take to encode `v` as a variable-length
661 /// integer.
varint_len(v: u64) -> usize662 pub fn varint_len(v: u64) -> usize {
663     if v <= 63 {
664         1
665     } else if v <= 16383 {
666         2
667     } else if v <= 1_073_741_823 {
668         4
669     } else if v <= 4_611_686_018_427_387_903 {
670         8
671     } else {
672         unreachable!()
673     }
674 }
675 
676 /// Returns how long the variable-length integer is, given its first byte.
varint_parse_len(first: u8) -> usize677 pub fn varint_parse_len(first: u8) -> usize {
678     match first >> 6 {
679         0 => 1,
680         1 => 2,
681         2 => 4,
682         3 => 8,
683         _ => unreachable!(),
684     }
685 }
686 
687 #[cfg(test)]
688 mod tests {
689     use super::*;
690 
691     #[test]
get_u()692     fn get_u() {
693         let d = [
694             1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
695         ];
696 
697         let mut b = Octets::with_slice(&d);
698         assert_eq!(b.cap(), 18);
699         assert_eq!(b.off(), 0);
700 
701         assert_eq!(b.get_u8().unwrap(), 1);
702         assert_eq!(b.cap(), 17);
703         assert_eq!(b.off(), 1);
704 
705         assert_eq!(b.get_u16().unwrap(), 0x203);
706         assert_eq!(b.cap(), 15);
707         assert_eq!(b.off(), 3);
708 
709         assert_eq!(b.get_u24().unwrap(), 0x40506);
710         assert_eq!(b.cap(), 12);
711         assert_eq!(b.off(), 6);
712 
713         assert_eq!(b.get_u32().unwrap(), 0x0708090a);
714         assert_eq!(b.cap(), 8);
715         assert_eq!(b.off(), 10);
716 
717         assert_eq!(b.get_u64().unwrap(), 0x0b0c0d0e0f101112);
718         assert_eq!(b.cap(), 0);
719         assert_eq!(b.off(), 18);
720 
721         assert!(b.get_u8().is_err());
722         assert!(b.get_u16().is_err());
723         assert!(b.get_u24().is_err());
724         assert!(b.get_u32().is_err());
725         assert!(b.get_u64().is_err());
726     }
727 
728     #[test]
get_u_mut()729     fn get_u_mut() {
730         let mut d = [
731             1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
732         ];
733 
734         let mut b = OctetsMut::with_slice(&mut d);
735         assert_eq!(b.cap(), 18);
736         assert_eq!(b.off(), 0);
737 
738         assert_eq!(b.get_u8().unwrap(), 1);
739         assert_eq!(b.cap(), 17);
740         assert_eq!(b.off(), 1);
741 
742         assert_eq!(b.get_u16().unwrap(), 0x203);
743         assert_eq!(b.cap(), 15);
744         assert_eq!(b.off(), 3);
745 
746         assert_eq!(b.get_u24().unwrap(), 0x40506);
747         assert_eq!(b.cap(), 12);
748         assert_eq!(b.off(), 6);
749 
750         assert_eq!(b.get_u32().unwrap(), 0x0708090a);
751         assert_eq!(b.cap(), 8);
752         assert_eq!(b.off(), 10);
753 
754         assert_eq!(b.get_u64().unwrap(), 0x0b0c0d0e0f101112);
755         assert_eq!(b.cap(), 0);
756         assert_eq!(b.off(), 18);
757 
758         assert!(b.get_u8().is_err());
759         assert!(b.get_u16().is_err());
760         assert!(b.get_u24().is_err());
761         assert!(b.get_u32().is_err());
762         assert!(b.get_u64().is_err());
763     }
764 
765     #[test]
peek_u()766     fn peek_u() {
767         let d = [1, 2];
768 
769         let mut b = Octets::with_slice(&d);
770         assert_eq!(b.cap(), 2);
771         assert_eq!(b.off(), 0);
772 
773         assert_eq!(b.peek_u8().unwrap(), 1);
774         assert_eq!(b.cap(), 2);
775         assert_eq!(b.off(), 0);
776 
777         assert_eq!(b.peek_u8().unwrap(), 1);
778         assert_eq!(b.cap(), 2);
779         assert_eq!(b.off(), 0);
780 
781         b.get_u16().unwrap();
782 
783         assert!(b.peek_u8().is_err());
784     }
785 
786     #[test]
peek_u_mut()787     fn peek_u_mut() {
788         let mut d = [1, 2];
789 
790         let mut b = OctetsMut::with_slice(&mut d);
791         assert_eq!(b.cap(), 2);
792         assert_eq!(b.off(), 0);
793 
794         assert_eq!(b.peek_u8().unwrap(), 1);
795         assert_eq!(b.cap(), 2);
796         assert_eq!(b.off(), 0);
797 
798         assert_eq!(b.peek_u8().unwrap(), 1);
799         assert_eq!(b.cap(), 2);
800         assert_eq!(b.off(), 0);
801 
802         b.get_u16().unwrap();
803 
804         assert!(b.peek_u8().is_err());
805     }
806 
807     #[test]
get_bytes()808     fn get_bytes() {
809         let d = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
810         let mut b = Octets::with_slice(&d);
811         assert_eq!(b.cap(), 10);
812         assert_eq!(b.off(), 0);
813 
814         assert_eq!(b.get_bytes(5).unwrap().as_ref(), [1, 2, 3, 4, 5]);
815         assert_eq!(b.cap(), 5);
816         assert_eq!(b.off(), 5);
817 
818         assert_eq!(b.get_bytes(3).unwrap().as_ref(), [6, 7, 8]);
819         assert_eq!(b.cap(), 2);
820         assert_eq!(b.off(), 8);
821 
822         assert!(b.get_bytes(3).is_err());
823         assert_eq!(b.cap(), 2);
824         assert_eq!(b.off(), 8);
825 
826         assert_eq!(b.get_bytes(2).unwrap().as_ref(), [9, 10]);
827         assert_eq!(b.cap(), 0);
828         assert_eq!(b.off(), 10);
829 
830         assert!(b.get_bytes(2).is_err());
831     }
832 
833     #[test]
get_bytes_mut()834     fn get_bytes_mut() {
835         let mut d = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
836         let mut b = OctetsMut::with_slice(&mut d);
837         assert_eq!(b.cap(), 10);
838         assert_eq!(b.off(), 0);
839 
840         assert_eq!(b.get_bytes(5).unwrap().as_ref(), [1, 2, 3, 4, 5]);
841         assert_eq!(b.cap(), 5);
842         assert_eq!(b.off(), 5);
843 
844         assert_eq!(b.get_bytes(3).unwrap().as_ref(), [6, 7, 8]);
845         assert_eq!(b.cap(), 2);
846         assert_eq!(b.off(), 8);
847 
848         assert!(b.get_bytes(3).is_err());
849         assert_eq!(b.cap(), 2);
850         assert_eq!(b.off(), 8);
851 
852         assert_eq!(b.get_bytes(2).unwrap().as_ref(), [9, 10]);
853         assert_eq!(b.cap(), 0);
854         assert_eq!(b.off(), 10);
855 
856         assert!(b.get_bytes(2).is_err());
857     }
858 
859     #[test]
peek_bytes()860     fn peek_bytes() {
861         let d = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
862         let mut b = Octets::with_slice(&d);
863         assert_eq!(b.cap(), 10);
864         assert_eq!(b.off(), 0);
865 
866         assert_eq!(b.peek_bytes(5).unwrap().as_ref(), [1, 2, 3, 4, 5]);
867         assert_eq!(b.cap(), 10);
868         assert_eq!(b.off(), 0);
869 
870         assert_eq!(b.peek_bytes(5).unwrap().as_ref(), [1, 2, 3, 4, 5]);
871         assert_eq!(b.cap(), 10);
872         assert_eq!(b.off(), 0);
873 
874         b.get_bytes(5).unwrap();
875     }
876 
877     #[test]
peek_bytes_mut()878     fn peek_bytes_mut() {
879         let mut d = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
880         let mut b = OctetsMut::with_slice(&mut d);
881         assert_eq!(b.cap(), 10);
882         assert_eq!(b.off(), 0);
883 
884         assert_eq!(b.peek_bytes(5).unwrap().as_ref(), [1, 2, 3, 4, 5]);
885         assert_eq!(b.cap(), 10);
886         assert_eq!(b.off(), 0);
887 
888         assert_eq!(b.peek_bytes(5).unwrap().as_ref(), [1, 2, 3, 4, 5]);
889         assert_eq!(b.cap(), 10);
890         assert_eq!(b.off(), 0);
891 
892         b.get_bytes(5).unwrap();
893     }
894 
895     #[test]
get_varint()896     fn get_varint() {
897         let d = [0xc2, 0x19, 0x7c, 0x5e, 0xff, 0x14, 0xe8, 0x8c];
898         let mut b = Octets::with_slice(&d);
899         assert_eq!(b.get_varint().unwrap(), 151288809941952652);
900         assert_eq!(b.cap(), 0);
901         assert_eq!(b.off(), 8);
902 
903         let d = [0x9d, 0x7f, 0x3e, 0x7d];
904         let mut b = Octets::with_slice(&d);
905         assert_eq!(b.get_varint().unwrap(), 494878333);
906         assert_eq!(b.cap(), 0);
907         assert_eq!(b.off(), 4);
908 
909         let d = [0x7b, 0xbd];
910         let mut b = Octets::with_slice(&d);
911         assert_eq!(b.get_varint().unwrap(), 15293);
912         assert_eq!(b.cap(), 0);
913         assert_eq!(b.off(), 2);
914 
915         let d = [0x40, 0x25];
916         let mut b = Octets::with_slice(&d);
917         assert_eq!(b.get_varint().unwrap(), 37);
918         assert_eq!(b.cap(), 0);
919         assert_eq!(b.off(), 2);
920 
921         let d = [0x25];
922         let mut b = Octets::with_slice(&d);
923         assert_eq!(b.get_varint().unwrap(), 37);
924         assert_eq!(b.cap(), 0);
925         assert_eq!(b.off(), 1);
926     }
927 
928     #[test]
get_varint_mut()929     fn get_varint_mut() {
930         let mut d = [0xc2, 0x19, 0x7c, 0x5e, 0xff, 0x14, 0xe8, 0x8c];
931         let mut b = OctetsMut::with_slice(&mut d);
932         assert_eq!(b.get_varint().unwrap(), 151288809941952652);
933         assert_eq!(b.cap(), 0);
934         assert_eq!(b.off(), 8);
935 
936         let mut d = [0x9d, 0x7f, 0x3e, 0x7d];
937         let mut b = OctetsMut::with_slice(&mut d);
938         assert_eq!(b.get_varint().unwrap(), 494878333);
939         assert_eq!(b.cap(), 0);
940         assert_eq!(b.off(), 4);
941 
942         let mut d = [0x7b, 0xbd];
943         let mut b = OctetsMut::with_slice(&mut d);
944         assert_eq!(b.get_varint().unwrap(), 15293);
945         assert_eq!(b.cap(), 0);
946         assert_eq!(b.off(), 2);
947 
948         let mut d = [0x40, 0x25];
949         let mut b = OctetsMut::with_slice(&mut d);
950         assert_eq!(b.get_varint().unwrap(), 37);
951         assert_eq!(b.cap(), 0);
952         assert_eq!(b.off(), 2);
953 
954         let mut d = [0x25];
955         let mut b = OctetsMut::with_slice(&mut d);
956         assert_eq!(b.get_varint().unwrap(), 37);
957         assert_eq!(b.cap(), 0);
958         assert_eq!(b.off(), 1);
959     }
960 
961     #[test]
put_varint()962     fn put_varint() {
963         let mut d = [0; 8];
964         {
965             let mut b = OctetsMut::with_slice(&mut d);
966             assert!(b.put_varint(151288809941952652).is_ok());
967             assert_eq!(b.cap(), 0);
968             assert_eq!(b.off(), 8);
969         }
970         let exp = [0xc2, 0x19, 0x7c, 0x5e, 0xff, 0x14, 0xe8, 0x8c];
971         assert_eq!(&d, &exp);
972 
973         let mut d = [0; 4];
974         {
975             let mut b = OctetsMut::with_slice(&mut d);
976             assert!(b.put_varint(494878333).is_ok());
977             assert_eq!(b.cap(), 0);
978             assert_eq!(b.off(), 4);
979         }
980         let exp = [0x9d, 0x7f, 0x3e, 0x7d];
981         assert_eq!(&d, &exp);
982 
983         let mut d = [0; 2];
984         {
985             let mut b = OctetsMut::with_slice(&mut d);
986             assert!(b.put_varint(15293).is_ok());
987             assert_eq!(b.cap(), 0);
988             assert_eq!(b.off(), 2);
989         }
990         let exp = [0x7b, 0xbd];
991         assert_eq!(&d, &exp);
992 
993         let mut d = [0; 1];
994         {
995             let mut b = OctetsMut::with_slice(&mut d);
996             assert!(b.put_varint(37).is_ok());
997             assert_eq!(b.cap(), 0);
998             assert_eq!(b.off(), 1);
999         }
1000         let exp = [0x25];
1001         assert_eq!(&d, &exp);
1002 
1003         let mut d = [0; 3];
1004         {
1005             let mut b = OctetsMut::with_slice(&mut d);
1006             assert!(b.put_varint(151288809941952652).is_err());
1007             assert_eq!(b.cap(), 3);
1008             assert_eq!(b.off(), 0);
1009         }
1010         let exp = [0; 3];
1011         assert_eq!(&d, &exp);
1012     }
1013 
1014     #[test]
1015     #[should_panic]
varint_too_large()1016     fn varint_too_large() {
1017         let mut d = [0; 3];
1018         let mut b = OctetsMut::with_slice(&mut d);
1019         assert!(b.put_varint(std::u64::MAX).is_err());
1020     }
1021 
1022     #[test]
put_u()1023     fn put_u() {
1024         let mut d = [0; 18];
1025 
1026         {
1027             let mut b = OctetsMut::with_slice(&mut d);
1028             assert_eq!(b.cap(), 18);
1029             assert_eq!(b.off(), 0);
1030 
1031             assert!(b.put_u8(1).is_ok());
1032             assert_eq!(b.cap(), 17);
1033             assert_eq!(b.off(), 1);
1034 
1035             assert!(b.put_u16(0x203).is_ok());
1036             assert_eq!(b.cap(), 15);
1037             assert_eq!(b.off(), 3);
1038 
1039             assert!(b.put_u24(0x40506).is_ok());
1040             assert_eq!(b.cap(), 12);
1041             assert_eq!(b.off(), 6);
1042 
1043             assert!(b.put_u32(0x0708090a).is_ok());
1044             assert_eq!(b.cap(), 8);
1045             assert_eq!(b.off(), 10);
1046 
1047             assert!(b.put_u64(0x0b0c0d0e0f101112).is_ok());
1048             assert_eq!(b.cap(), 0);
1049             assert_eq!(b.off(), 18);
1050 
1051             assert!(b.put_u8(1).is_err());
1052         }
1053 
1054         let exp = [
1055             1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
1056         ];
1057         assert_eq!(&d, &exp);
1058     }
1059 
1060     #[test]
put_bytes()1061     fn put_bytes() {
1062         let mut d = [0; 5];
1063 
1064         {
1065             let mut b = OctetsMut::with_slice(&mut d);
1066             assert_eq!(b.cap(), 5);
1067             assert_eq!(b.off(), 0);
1068 
1069             let p = [0x0a, 0x0b, 0x0c, 0x0d, 0x0e];
1070             assert!(b.put_bytes(&p).is_ok());
1071             assert_eq!(b.cap(), 0);
1072             assert_eq!(b.off(), 5);
1073 
1074             assert!(b.put_u8(1).is_err());
1075         }
1076 
1077         let exp = [0xa, 0xb, 0xc, 0xd, 0xe];
1078         assert_eq!(&d, &exp);
1079     }
1080 
1081     #[test]
split()1082     fn split() {
1083         let mut d = b"helloworld".to_vec();
1084 
1085         let mut b = OctetsMut::with_slice(&mut d);
1086         assert_eq!(b.cap(), 10);
1087         assert_eq!(b.off(), 0);
1088         assert_eq!(b.as_ref(), b"helloworld");
1089 
1090         assert!(b.get_bytes(5).is_ok());
1091         assert_eq!(b.cap(), 5);
1092         assert_eq!(b.off(), 5);
1093         assert_eq!(b.as_ref(), b"world");
1094 
1095         let off = b.off();
1096 
1097         let (first, last) = b.split_at(off).unwrap();
1098         assert_eq!(first.cap(), 5);
1099         assert_eq!(first.off(), 0);
1100         assert_eq!(first.as_ref(), b"hello");
1101 
1102         assert_eq!(last.cap(), 5);
1103         assert_eq!(last.off(), 0);
1104         assert_eq!(last.as_ref(), b"world");
1105     }
1106 
1107     #[test]
split_at()1108     fn split_at() {
1109         let mut d = b"helloworld".to_vec();
1110 
1111         {
1112             let mut b = OctetsMut::with_slice(&mut d);
1113             let (first, second) = b.split_at(5).unwrap();
1114 
1115             let mut exp1 = b"hello".to_vec();
1116             assert_eq!(first.as_ref(), &mut exp1[..]);
1117 
1118             let mut exp2 = b"world".to_vec();
1119             assert_eq!(second.as_ref(), &mut exp2[..]);
1120         }
1121 
1122         {
1123             let mut b = OctetsMut::with_slice(&mut d);
1124             let (first, second) = b.split_at(10).unwrap();
1125 
1126             let mut exp1 = b"helloworld".to_vec();
1127             assert_eq!(first.as_ref(), &mut exp1[..]);
1128 
1129             let mut exp2 = b"".to_vec();
1130             assert_eq!(second.as_ref(), &mut exp2[..]);
1131         }
1132 
1133         {
1134             let mut b = OctetsMut::with_slice(&mut d);
1135             let (first, second) = b.split_at(9).unwrap();
1136 
1137             let mut exp1 = b"helloworl".to_vec();
1138             assert_eq!(first.as_ref(), &mut exp1[..]);
1139 
1140             let mut exp2 = b"d".to_vec();
1141             assert_eq!(second.as_ref(), &mut exp2[..]);
1142         }
1143 
1144         {
1145             let mut b = OctetsMut::with_slice(&mut d);
1146             assert!(b.split_at(11).is_err());
1147         }
1148     }
1149 
1150     #[test]
slice()1151     fn slice() {
1152         let d = b"helloworld".to_vec();
1153 
1154         {
1155             let b = Octets::with_slice(&d);
1156             let exp = b"hello".to_vec();
1157             assert_eq!(b.slice(5), Ok(&exp[..]));
1158         }
1159 
1160         {
1161             let b = Octets::with_slice(&d);
1162             let exp = b"".to_vec();
1163             assert_eq!(b.slice(0), Ok(&exp[..]));
1164         }
1165 
1166         {
1167             let mut b = Octets::with_slice(&d);
1168             b.get_bytes(5).unwrap();
1169 
1170             let exp = b"world".to_vec();
1171             assert_eq!(b.slice(5), Ok(&exp[..]));
1172         }
1173 
1174         {
1175             let b = Octets::with_slice(&d);
1176             assert!(b.slice(11).is_err());
1177         }
1178     }
1179 
1180     #[test]
slice_mut()1181     fn slice_mut() {
1182         let mut d = b"helloworld".to_vec();
1183 
1184         {
1185             let mut b = OctetsMut::with_slice(&mut d);
1186             let mut exp = b"hello".to_vec();
1187             assert_eq!(b.slice(5), Ok(&mut exp[..]));
1188         }
1189 
1190         {
1191             let mut b = OctetsMut::with_slice(&mut d);
1192             let mut exp = b"".to_vec();
1193             assert_eq!(b.slice(0), Ok(&mut exp[..]));
1194         }
1195 
1196         {
1197             let mut b = OctetsMut::with_slice(&mut d);
1198             b.get_bytes(5).unwrap();
1199 
1200             let mut exp = b"world".to_vec();
1201             assert_eq!(b.slice(5), Ok(&mut exp[..]));
1202         }
1203 
1204         {
1205             let mut b = OctetsMut::with_slice(&mut d);
1206             assert!(b.slice(11).is_err());
1207         }
1208     }
1209 
1210     #[test]
slice_last()1211     fn slice_last() {
1212         let d = b"helloworld".to_vec();
1213 
1214         {
1215             let b = Octets::with_slice(&d);
1216             let exp = b"orld".to_vec();
1217             assert_eq!(b.slice_last(4), Ok(&exp[..]));
1218         }
1219 
1220         {
1221             let b = Octets::with_slice(&d);
1222             let exp = b"d".to_vec();
1223             assert_eq!(b.slice_last(1), Ok(&exp[..]));
1224         }
1225 
1226         {
1227             let b = Octets::with_slice(&d);
1228             let exp = b"".to_vec();
1229             assert_eq!(b.slice_last(0), Ok(&exp[..]));
1230         }
1231 
1232         {
1233             let b = Octets::with_slice(&d);
1234             let exp = b"helloworld".to_vec();
1235             assert_eq!(b.slice_last(10), Ok(&exp[..]));
1236         }
1237 
1238         {
1239             let b = Octets::with_slice(&d);
1240             assert!(b.slice_last(11).is_err());
1241         }
1242     }
1243 
1244     #[test]
slice_last_mut()1245     fn slice_last_mut() {
1246         let mut d = b"helloworld".to_vec();
1247 
1248         {
1249             let mut b = OctetsMut::with_slice(&mut d);
1250             let mut exp = b"orld".to_vec();
1251             assert_eq!(b.slice_last(4), Ok(&mut exp[..]));
1252         }
1253 
1254         {
1255             let mut b = OctetsMut::with_slice(&mut d);
1256             let mut exp = b"d".to_vec();
1257             assert_eq!(b.slice_last(1), Ok(&mut exp[..]));
1258         }
1259 
1260         {
1261             let mut b = OctetsMut::with_slice(&mut d);
1262             let mut exp = b"".to_vec();
1263             assert_eq!(b.slice_last(0), Ok(&mut exp[..]));
1264         }
1265 
1266         {
1267             let mut b = OctetsMut::with_slice(&mut d);
1268             let mut exp = b"helloworld".to_vec();
1269             assert_eq!(b.slice_last(10), Ok(&mut exp[..]));
1270         }
1271 
1272         {
1273             let mut b = OctetsMut::with_slice(&mut d);
1274             assert!(b.slice_last(11).is_err());
1275         }
1276     }
1277 }
1278