tonic/codec/
buffer.rs

Help
1use bytes::buf::UninitSlice;
2use bytes::{Buf, BufMut, Bytes, BytesMut};
3
4/// A specialized buffer to decode gRPC messages from.
5#[derive(Debug)]
6pub 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)]
13pub struct EncodeBuf<'a> {
14    buf: &'a mut BytesMut,
15}
16
17impl<'a> DecodeBuf<'a> {
18    pub(crate) fn new(buf: &'a mut BytesMut, len: usize) -> Self {
19        DecodeBuf { buf, len }
20    }
21}
22
23impl Buf for DecodeBuf<'_> {
24    #[inline]
25    fn remaining(&self) -> usize {
26        self.len
27    }
28
29    #[inline]
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]
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]
48    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
55impl<'a> EncodeBuf<'a> {
56    pub(crate) fn new(buf: &'a mut BytesMut) -> Self {
57        EncodeBuf { buf }
58    }
59}
60
61impl 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]
68    pub fn reserve(&mut self, additional: usize) {
69        self.buf.reserve(additional);
70    }
71}
72
73unsafe impl BufMut for EncodeBuf<'_> {
74    #[inline]
75    fn remaining_mut(&self) -> usize {
76        self.buf.remaining_mut()
77    }
78
79    #[inline]
80    unsafe fn advance_mut(&mut self, cnt: usize) {
81        self.buf.advance_mut(cnt)
82    }
83
84    #[inline]
85    fn chunk_mut(&mut self) -> &mut UninitSlice {
86        self.buf.chunk_mut()
87    }
88
89    #[inline]
90    fn put<T: Buf>(&mut self, src: T)
91    where
92        Self: Sized,
93    {
94        self.buf.put(src)
95    }
96
97    #[inline]
98    fn put_slice(&mut self, src: &[u8]) {
99        self.buf.put_slice(src)
100    }
101
102    #[inline]
103    fn put_bytes(&mut self, val: u8, cnt: usize) {
104        self.buf.put_bytes(val, cnt);
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[test]
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]
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}
tonic/codec/buffer.rs

tonic/codec/
buffer.rs
