1 use bytes::buf::UninitSlice; 2 use bytes::{Buf, BufMut, Bytes, BytesMut}; 3 4 /// A specialized buffer to decode gRPC messages from. 5 #[derive(Debug)] 6 pub struct DecodeBuf<'a> { 7 buf: &'a mut BytesMut, 8 len: usize, 9 } 10 11 /// A specialized buffer to encode gRPC messages into. 12 #[derive(Debug)] 13 pub struct EncodeBuf<'a> { 14 buf: &'a mut BytesMut, 15 } 16 17 impl<'a> DecodeBuf<'a> { new(buf: &'a mut BytesMut, len: usize) -> Self18 pub(crate) fn new(buf: &'a mut BytesMut, len: usize) -> Self { 19 DecodeBuf { buf, len } 20 } 21 } 22 23 impl Buf for DecodeBuf<'_> { 24 #[inline] remaining(&self) -> usize25 fn remaining(&self) -> usize { 26 self.len 27 } 28 29 #[inline] chunk(&self) -> &[u8]30 fn chunk(&self) -> &[u8] { 31 let ret = self.buf.chunk(); 32 33 if ret.len() > self.len { 34 &ret[..self.len] 35 } else { 36 ret 37 } 38 } 39 40 #[inline] advance(&mut self, cnt: usize)41 fn advance(&mut self, cnt: usize) { 42 assert!(cnt <= self.len); 43 self.buf.advance(cnt); 44 self.len -= cnt; 45 } 46 47 #[inline] copy_to_bytes(&mut self, len: usize) -> Bytes48 fn copy_to_bytes(&mut self, len: usize) -> Bytes { 49 assert!(len <= self.len); 50 self.len -= len; 51 self.buf.copy_to_bytes(len) 52 } 53 } 54 55 impl<'a> EncodeBuf<'a> { new(buf: &'a mut BytesMut) -> Self56 pub(crate) fn new(buf: &'a mut BytesMut) -> Self { 57 EncodeBuf { buf } 58 } 59 } 60 61 impl EncodeBuf<'_> { 62 /// Reserves capacity for at least `additional` more bytes to be inserted 63 /// into the buffer. 64 /// 65 /// More than `additional` bytes may be reserved in order to avoid frequent 66 /// reallocations. A call to `reserve` may result in an allocation. 67 #[inline] reserve(&mut self, additional: usize)68 pub fn reserve(&mut self, additional: usize) { 69 self.buf.reserve(additional); 70 } 71 } 72 73 unsafe impl BufMut for EncodeBuf<'_> { 74 #[inline] remaining_mut(&self) -> usize75 fn remaining_mut(&self) -> usize { 76 self.buf.remaining_mut() 77 } 78 79 #[inline] advance_mut(&mut self, cnt: usize)80 unsafe fn advance_mut(&mut self, cnt: usize) { 81 self.buf.advance_mut(cnt) 82 } 83 84 #[inline] chunk_mut(&mut self) -> &mut UninitSlice85 fn chunk_mut(&mut self) -> &mut UninitSlice { 86 self.buf.chunk_mut() 87 } 88 89 #[inline] put<T: Buf>(&mut self, src: T) where Self: Sized,90 fn put<T: Buf>(&mut self, src: T) 91 where 92 Self: Sized, 93 { 94 self.buf.put(src) 95 } 96 97 #[inline] put_slice(&mut self, src: &[u8])98 fn put_slice(&mut self, src: &[u8]) { 99 self.buf.put_slice(src) 100 } 101 102 #[inline] put_bytes(&mut self, val: u8, cnt: usize)103 fn put_bytes(&mut self, val: u8, cnt: usize) { 104 self.buf.put_bytes(val, cnt); 105 } 106 } 107 108 #[cfg(test)] 109 mod tests { 110 use super::*; 111 112 #[test] decode_buf()113 fn decode_buf() { 114 let mut payload = BytesMut::with_capacity(100); 115 payload.put(&vec![0u8; 50][..]); 116 let mut buf = DecodeBuf::new(&mut payload, 20); 117 118 assert_eq!(buf.len, 20); 119 assert_eq!(buf.remaining(), 20); 120 assert_eq!(buf.chunk().len(), 20); 121 122 buf.advance(10); 123 assert_eq!(buf.remaining(), 10); 124 125 let mut out = [0; 5]; 126 buf.copy_to_slice(&mut out); 127 assert_eq!(buf.remaining(), 5); 128 assert_eq!(buf.chunk().len(), 5); 129 130 assert_eq!(buf.copy_to_bytes(5).len(), 5); 131 assert!(!buf.has_remaining()); 132 } 133 134 #[test] encode_buf()135 fn encode_buf() { 136 let mut bytes = BytesMut::with_capacity(100); 137 let mut buf = EncodeBuf::new(&mut bytes); 138 139 let initial = buf.remaining_mut(); 140 unsafe { buf.advance_mut(20) }; 141 assert_eq!(buf.remaining_mut(), initial - 20); 142 143 buf.put_u8(b'a'); 144 assert_eq!(buf.remaining_mut(), initial - 20 - 1); 145 } 146 } 147