tonic/codec/
buffer.rs

1/*
2 *
3 * Copyright 2025 gRPC authors.
4 *
5 * Permission is hereby granted, free of charge, to any person obtaining a copy
6 * of this software and associated documentation files (the "Software"), to
7 * deal in the Software without restriction, including without limitation the
8 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
9 * sell copies of the Software, and to permit persons to whom the Software is
10 * furnished to do so, subject to the following conditions:
11 *
12 * The above copyright notice and this permission notice shall be included in
13 * all copies or substantial portions of the Software.
14 *
15 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
20 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
21 * IN THE SOFTWARE.
22 *
23 */
24
25use bytes::buf::UninitSlice;
26use bytes::{Buf, BufMut, Bytes, BytesMut};
27
28/// A specialized buffer to decode gRPC messages from.
29#[derive(Debug)]
30pub struct DecodeBuf<'a> {
31    buf: &'a mut BytesMut,
32    len: usize,
33}
34
35/// A specialized buffer to encode gRPC messages into.
36#[derive(Debug)]
37pub struct EncodeBuf<'a> {
38    buf: &'a mut BytesMut,
39}
40
41impl<'a> DecodeBuf<'a> {
42    pub(crate) fn new(buf: &'a mut BytesMut, len: usize) -> Self {
43        DecodeBuf { buf, len }
44    }
45}
46
47impl Buf for DecodeBuf<'_> {
48    #[inline]
49    fn remaining(&self) -> usize {
50        self.len
51    }
52
53    #[inline]
54    fn chunk(&self) -> &[u8] {
55        let ret = self.buf.chunk();
56
57        if ret.len() > self.len {
58            &ret[..self.len]
59        } else {
60            ret
61        }
62    }
63
64    #[inline]
65    fn advance(&mut self, cnt: usize) {
66        assert!(cnt <= self.len);
67        self.buf.advance(cnt);
68        self.len -= cnt;
69    }
70
71    #[inline]
72    fn copy_to_bytes(&mut self, len: usize) -> Bytes {
73        assert!(len <= self.len);
74        self.len -= len;
75        self.buf.copy_to_bytes(len)
76    }
77}
78
79impl<'a> EncodeBuf<'a> {
80    pub(crate) fn new(buf: &'a mut BytesMut) -> Self {
81        EncodeBuf { buf }
82    }
83}
84
85impl EncodeBuf<'_> {
86    /// Reserves capacity for at least `additional` more bytes to be inserted
87    /// into the buffer.
88    ///
89    /// More than `additional` bytes may be reserved in order to avoid frequent
90    /// reallocations. A call to `reserve` may result in an allocation.
91    #[inline]
92    pub fn reserve(&mut self, additional: usize) {
93        self.buf.reserve(additional);
94    }
95}
96
97unsafe impl BufMut for EncodeBuf<'_> {
98    #[inline]
99    fn remaining_mut(&self) -> usize {
100        self.buf.remaining_mut()
101    }
102
103    #[inline]
104    unsafe fn advance_mut(&mut self, cnt: usize) {
105        unsafe { self.buf.advance_mut(cnt) }
106    }
107
108    #[inline]
109    fn chunk_mut(&mut self) -> &mut UninitSlice {
110        self.buf.chunk_mut()
111    }
112
113    #[inline]
114    fn put<T: Buf>(&mut self, src: T)
115    where
116        Self: Sized,
117    {
118        self.buf.put(src)
119    }
120
121    #[inline]
122    fn put_slice(&mut self, src: &[u8]) {
123        self.buf.put_slice(src)
124    }
125
126    #[inline]
127    fn put_bytes(&mut self, val: u8, cnt: usize) {
128        self.buf.put_bytes(val, cnt);
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn decode_buf() {
138        let mut payload = BytesMut::with_capacity(100);
139        payload.put(&vec![0u8; 50][..]);
140        let mut buf = DecodeBuf::new(&mut payload, 20);
141
142        assert_eq!(buf.len, 20);
143        assert_eq!(buf.remaining(), 20);
144        assert_eq!(buf.chunk().len(), 20);
145
146        buf.advance(10);
147        assert_eq!(buf.remaining(), 10);
148
149        let mut out = [0; 5];
150        buf.copy_to_slice(&mut out);
151        assert_eq!(buf.remaining(), 5);
152        assert_eq!(buf.chunk().len(), 5);
153
154        assert_eq!(buf.copy_to_bytes(5).len(), 5);
155        assert!(!buf.has_remaining());
156    }
157
158    #[test]
159    fn encode_buf() {
160        let mut bytes = BytesMut::with_capacity(100);
161        let mut buf = EncodeBuf::new(&mut bytes);
162
163        let initial = buf.remaining_mut();
164        unsafe { buf.advance_mut(20) };
165        assert_eq!(buf.remaining_mut(), initial - 20);
166
167        buf.put_u8(b'a');
168        assert_eq!(buf.remaining_mut(), initial - 20 - 1);
169    }
170}
tonic/codec/buffer.rs

tonic/codec/
buffer.rs
