wrpc_transport/frame/
codec.rs

1use std::sync::Arc;
2
3use bytes::{Bytes, BytesMut};
4use tracing::{instrument, trace};
5use wasm_tokio::{DEFAULT_MAX_INITIAL_CAPACITY, Leb128DecoderU32, Leb128DecoderU64, Leb128Encoder};
6
7use crate::frame::{Frame, FrameRef, MAX_INITIAL_PATH_CAPACITY};
8
9/// [Frame] decoder
10pub struct Decoder {
11    path: Option<Vec<usize>>,
12    path_cap: usize,
13    data_len: usize,
14    max_depth: u32,
15    max_size: u64,
16}
17
18impl Decoder {
19    /// Construct a new [Frame] decoder
20    #[must_use]
21    pub fn new(max_depth: u32, max_size: u64) -> Self {
22        Self {
23            path: Option::default(),
24            path_cap: 0,
25            data_len: 0,
26            max_depth,
27            max_size,
28        }
29    }
30}
31
32impl Default for Decoder {
33    fn default() -> Self {
34        Self::new(u32::MAX, u32::MAX.into())
35    }
36}
37
38impl tokio_util::codec::Decoder for Decoder {
39    type Item = Frame;
40    type Error = std::io::Error;
41
42    #[instrument(level = "trace", skip_all)]
43    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
44        let path = self.path.take();
45        let mut path = if let Some(path) = path {
46            path
47        } else {
48            trace!("decoding path length");
49            let Some(n) = Leb128DecoderU32.decode(src)? else {
50                return Ok(None);
51            };
52            trace!(n, "decoded path length");
53            if n > self.max_depth {
54                return Err(std::io::Error::new(
55                    std::io::ErrorKind::InvalidInput,
56                    format!(
57                        "path length of `{n}` exceeds maximum of `{}`",
58                        self.max_depth
59                    ),
60                ));
61            }
62            let n = usize::try_from(n)
63                .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
64            self.path_cap = n;
65            Vec::with_capacity(n.min(MAX_INITIAL_PATH_CAPACITY))
66        };
67        while self.path_cap > 0 {
68            trace!(self.path_cap, "decoding path element");
69            let Some(i) = Leb128DecoderU32.decode(src)? else {
70                self.path = Some(path);
71                return Ok(None);
72            };
73            trace!(i, "decoded path element");
74            let i = usize::try_from(i)
75                .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
76            path.push(i);
77            self.path_cap -= 1;
78        }
79        if self.data_len == 0 {
80            trace!("decoding data length");
81            let Some(n) = Leb128DecoderU64.decode(src)? else {
82                self.path = Some(path);
83                return Ok(None);
84            };
85            trace!(n, "decoded data length");
86            if n > self.max_size {
87                return Err(std::io::Error::new(
88                    std::io::ErrorKind::InvalidInput,
89                    format!(
90                        "payload length of `{n}` exceeds maximum of `{}`",
91                        self.max_size
92                    ),
93                ));
94            }
95            let n = usize::try_from(n)
96                .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
97            self.data_len = n;
98            if n == 0 {
99                return Ok(Some(Frame {
100                    path: Arc::from(path),
101                    data: Bytes::default(),
102                }));
103            }
104        }
105        let n = self.data_len.saturating_sub(src.len());
106        if n > 0 {
107            src.reserve(n.min(DEFAULT_MAX_INITIAL_CAPACITY));
108            self.path = Some(path);
109            return Ok(None);
110        }
111        trace!(self.data_len, "decoding data");
112        let data = src.split_to(self.data_len).freeze();
113        self.data_len = 0;
114        Ok(Some(Frame {
115            path: Arc::from(path),
116            data,
117        }))
118    }
119}
120
121/// [Frame] encoder
122pub struct Encoder;
123
124impl tokio_util::codec::Encoder<FrameRef<'_>> for Encoder {
125    type Error = std::io::Error;
126
127    #[instrument(level = "trace", skip_all)]
128    fn encode(
129        &mut self,
130        FrameRef { path, data }: FrameRef<'_>,
131        dst: &mut BytesMut,
132    ) -> Result<(), Self::Error> {
133        let size = data.len();
134        let depth = path.len();
135        dst.reserve(size.saturating_add(depth).saturating_add(5 + 10));
136        let n = u32::try_from(depth)
137            .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
138        trace!(n, "encoding path length");
139        Leb128Encoder.encode(n, dst)?;
140        for p in path {
141            let p = u32::try_from(*p)
142                .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
143            trace!(p, "encoding path element");
144            Leb128Encoder.encode(p, dst)?;
145        }
146        let n = u64::try_from(size)
147            .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
148        trace!(n, "encoding data length");
149        Leb128Encoder.encode(n, dst)?;
150        dst.extend_from_slice(data);
151        Ok(())
152    }
153}
154
155impl tokio_util::codec::Encoder<&Frame> for Encoder {
156    type Error = std::io::Error;
157
158    #[instrument(level = "trace", skip_all)]
159    fn encode(&mut self, frame: &Frame, dst: &mut BytesMut) -> Result<(), Self::Error> {
160        self.encode(FrameRef::from(frame), dst)
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use futures::{SinkExt as _, TryStreamExt as _};
167    use tokio_util::codec::{FramedRead, FramedWrite};
168
169    use super::*;
170
171    #[test_log::test(tokio::test)]
172    async fn codec() -> std::io::Result<()> {
173        let mut tx = FramedWrite::new(vec![], Encoder);
174
175        tx.send(&Frame {
176            path: [0, 1, 2].into(),
177            data: "test".into(),
178        })
179        .await?;
180
181        tx.send(FrameRef {
182            path: &[],
183            data: b"",
184        })
185        .await?;
186
187        tx.send(FrameRef {
188            path: &[0x42],
189            data: "\x7fƒ𐍈Ő".as_bytes(),
190        })
191        .await?;
192
193        let tx = tx.into_inner();
194        assert_eq!(
195            tx,
196            concat!(
197                concat!("\x03", concat!("\0", "\x01", "\x02"), "\x04test"),
198                concat!("\0", "\0"),
199                concat!("\x01", concat!("\x42"), "\x09\x7fƒ𐍈Ő"),
200            )
201            .as_bytes()
202        );
203
204        let mut rx = FramedRead::new(tx.as_slice(), Decoder::default());
205
206        let s = rx.try_next().await?;
207        assert_eq!(
208            s,
209            Some(Frame {
210                path: [0, 1, 2].into(),
211                data: "test".into(),
212            })
213        );
214
215        let s = rx.try_next().await?;
216        assert_eq!(
217            s,
218            Some(Frame {
219                path: [].into(),
220                data: "".into(),
221            })
222        );
223
224        let s = rx.try_next().await?;
225        assert_eq!(
226            s,
227            Some(Frame {
228                path: [0x42].into(),
229                data: "\x7fƒ𐍈Ő".into(),
230            })
231        );
232
233        let s = rx.try_next().await.expect("failed to get EOF");
234        assert_eq!(s, None);
235
236        Ok(())
237    }
238}
wrpc_transport/frame/codec.rs

wrpc_transport/frame/
codec.rs
