tonic/codec/
buffer.rs

1use bytes::buf::UninitSlice;
2use bytes::{Buf, BufMut, 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
48impl<'a> EncodeBuf<'a> {
49    pub(crate) fn new(buf: &'a mut BytesMut) -> Self {
50        EncodeBuf { buf }
51    }
52}
53
54impl 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]
61    pub fn reserve(&mut self, additional: usize) {
62        self.buf.reserve(additional);
63    }
64}
65
66unsafe impl BufMut for EncodeBuf<'_> {
67    #[inline]
68    fn remaining_mut(&self) -> usize {
69        self.buf.remaining_mut()
70    }
71
72    #[inline]
73    unsafe fn advance_mut(&mut self, cnt: usize) {
74        self.buf.advance_mut(cnt)
75    }
76
77    #[inline]
78    fn chunk_mut(&mut self) -> &mut UninitSlice {
79        self.buf.chunk_mut()
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
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]
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}
tonic/codec/buffer.rs

tonic/codec/
buffer.rs
