1 use bytes::buf::UninitSlice;
2 use bytes::{Buf, BufMut, 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 
48 impl<'a> EncodeBuf<'a> {
new(buf: &'a mut BytesMut) -> Self49     pub(crate) fn new(buf: &'a mut BytesMut) -> Self {
50         EncodeBuf { buf }
51     }
52 }
53 
54 impl EncodeBuf<'_> {
55     /// Reserves capacity for at least `additional` more bytes to be inserted
56     /// into the buffer.
57     ///
58     /// More than `additional` bytes may be reserved in order to avoid frequent
59     /// reallocations. A call to `reserve` may result in an allocation.
60     #[inline]
reserve(&mut self, additional: usize)61     pub fn reserve(&mut self, additional: usize) {
62         self.buf.reserve(additional);
63     }
64 }
65 
66 unsafe impl BufMut for EncodeBuf<'_> {
67     #[inline]
remaining_mut(&self) -> usize68     fn remaining_mut(&self) -> usize {
69         self.buf.remaining_mut()
70     }
71 
72     #[inline]
advance_mut(&mut self, cnt: usize)73     unsafe fn advance_mut(&mut self, cnt: usize) {
74         self.buf.advance_mut(cnt)
75     }
76 
77     #[inline]
chunk_mut(&mut self) -> &mut UninitSlice78     fn chunk_mut(&mut self) -> &mut UninitSlice {
79         self.buf.chunk_mut()
80     }
81 }
82 
83 #[cfg(test)]
84 mod tests {
85     use super::*;
86 
87     #[test]
decode_buf()88     fn decode_buf() {
89         let mut payload = BytesMut::with_capacity(100);
90         payload.put(&vec![0u8; 50][..]);
91         let mut buf = DecodeBuf::new(&mut payload, 20);
92 
93         assert_eq!(buf.len, 20);
94         assert_eq!(buf.remaining(), 20);
95         assert_eq!(buf.chunk().len(), 20);
96 
97         buf.advance(10);
98         assert_eq!(buf.remaining(), 10);
99 
100         let mut out = [0; 5];
101         buf.copy_to_slice(&mut out);
102         assert_eq!(buf.remaining(), 5);
103         assert_eq!(buf.chunk().len(), 5);
104 
105         assert_eq!(buf.copy_to_bytes(5).len(), 5);
106         assert!(!buf.has_remaining());
107     }
108 
109     #[test]
encode_buf()110     fn encode_buf() {
111         let mut bytes = BytesMut::with_capacity(100);
112         let mut buf = EncodeBuf::new(&mut bytes);
113 
114         let initial = buf.remaining_mut();
115         unsafe { buf.advance_mut(20) };
116         assert_eq!(buf.remaining_mut(), initial - 20);
117 
118         buf.put_u8(b'a');
119         assert_eq!(buf.remaining_mut(), initial - 20 - 1);
120     }
121 }
122