iota_network_stack/
codec.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{io::Read, marker::PhantomData};
6
7use bytes::{Buf, BufMut};
8use tonic::{
9    Status,
10    codec::{Codec, DecodeBuf, Decoder, EncodeBuf, Encoder},
11};
12
13/// Default upper bound on total decompressed payload size.
14const MAX_DECOMPRESSED_SIZE: u64 = 128 << 20;
15
16/// Decompress a snappy-framed stream, bounding output at
17/// [`MAX_DECOMPRESSED_SIZE`].
18fn decompress_snappy<R: Read>(src: R) -> std::io::Result<Vec<u8>> {
19    decompress_snappy_bounded(src, MAX_DECOMPRESSED_SIZE)
20}
21
22/// Decompress a snappy-framed stream, bounding total output at
23/// `max_allowed` bytes. Exposed separately from [`decompress_snappy`] so
24/// tests can exercise the bound directly.
25fn decompress_snappy_bounded<R: Read>(src: R, max_allowed: u64) -> std::io::Result<Vec<u8>> {
26    let mut snappy_decoder = snap::read::FrameDecoder::new(src).take(max_allowed);
27    let mut bytes = Vec::new();
28    snappy_decoder.read_to_end(&mut bytes)?;
29    Ok(bytes)
30}
31
32#[derive(Debug)]
33pub struct BcsEncoder<T>(PhantomData<T>);
34
35impl<T: serde::Serialize> Encoder for BcsEncoder<T> {
36    type Item = T;
37    type Error = Status;
38
39    fn encode(&mut self, item: Self::Item, buf: &mut EncodeBuf<'_>) -> Result<(), Self::Error> {
40        bcs::serialize_into(&mut buf.writer(), &item).map_err(|e| Status::internal(e.to_string()))
41    }
42}
43
44#[derive(Debug)]
45pub struct BcsDecoder<U>(PhantomData<U>);
46
47impl<U: serde::de::DeserializeOwned> Decoder for BcsDecoder<U> {
48    type Item = U;
49    type Error = Status;
50
51    fn decode(&mut self, buf: &mut DecodeBuf<'_>) -> Result<Option<Self::Item>, Self::Error> {
52        if !buf.has_remaining() {
53            return Ok(None);
54        }
55
56        let chunk = buf.chunk();
57
58        let item: Self::Item =
59            bcs::from_bytes(chunk).map_err(|e| Status::internal(e.to_string()))?;
60        buf.advance(chunk.len());
61
62        Ok(Some(item))
63    }
64}
65
66/// A [`Codec`] that implements `application/grpc+bcs` via the serde library.
67#[derive(Debug, Clone)]
68pub struct BcsCodec<T, U>(PhantomData<(T, U)>);
69
70impl<T, U> Default for BcsCodec<T, U> {
71    fn default() -> Self {
72        Self(PhantomData)
73    }
74}
75
76impl<T, U> Codec for BcsCodec<T, U>
77where
78    T: serde::Serialize + Send + 'static,
79    U: serde::de::DeserializeOwned + Send + 'static,
80{
81    type Encode = T;
82    type Decode = U;
83    type Encoder = BcsEncoder<T>;
84    type Decoder = BcsDecoder<U>;
85
86    fn encoder(&mut self) -> Self::Encoder {
87        BcsEncoder(PhantomData)
88    }
89
90    fn decoder(&mut self) -> Self::Decoder {
91        BcsDecoder(PhantomData)
92    }
93}
94
95#[derive(Debug)]
96pub struct BcsSnappyEncoder<T>(PhantomData<T>);
97
98impl<T: serde::Serialize> Encoder for BcsSnappyEncoder<T> {
99    type Item = T;
100    type Error = Status;
101
102    fn encode(&mut self, item: Self::Item, buf: &mut EncodeBuf<'_>) -> Result<(), Self::Error> {
103        let mut snappy_encoder = snap::write::FrameEncoder::new(buf.writer());
104        bcs::serialize_into(&mut snappy_encoder, &item).map_err(|e| Status::internal(e.to_string()))
105    }
106}
107
108#[derive(Debug)]
109pub struct BcsSnappyDecoder<U>(PhantomData<U>);
110
111impl<U: serde::de::DeserializeOwned> Decoder for BcsSnappyDecoder<U> {
112    type Item = U;
113    type Error = Status;
114
115    fn decode(&mut self, buf: &mut DecodeBuf<'_>) -> Result<Option<Self::Item>, Self::Error> {
116        if !buf.has_remaining() {
117            return Ok(None);
118        }
119        let bytes = decompress_snappy(buf.reader()).map_err(|e| Status::internal(e.to_string()))?;
120        let item =
121            bcs::from_bytes(bytes.as_slice()).map_err(|e| Status::internal(e.to_string()))?;
122        Ok(Some(item))
123    }
124}
125
126/// A [`Codec`] that implements `bcs` encoding/decoding and snappy
127/// compression/decompression via the serde library.
128#[derive(Debug, Clone)]
129pub struct BcsSnappyCodec<T, U>(PhantomData<(T, U)>);
130
131impl<T, U> Default for BcsSnappyCodec<T, U> {
132    fn default() -> Self {
133        Self(PhantomData)
134    }
135}
136
137impl<T, U> Codec for BcsSnappyCodec<T, U>
138where
139    T: serde::Serialize + Send + 'static,
140    U: serde::de::DeserializeOwned + Send + 'static,
141{
142    type Encode = T;
143    type Decode = U;
144    type Encoder = BcsSnappyEncoder<T>;
145    type Decoder = BcsSnappyDecoder<U>;
146
147    fn encoder(&mut self) -> Self::Encoder {
148        BcsSnappyEncoder(PhantomData)
149    }
150
151    fn decoder(&mut self) -> Self::Decoder {
152        BcsSnappyDecoder(PhantomData)
153    }
154}
155
156// Anemo variant of BCS codec using Snappy for compression.
157pub mod anemo {
158    use std::marker::PhantomData;
159
160    use ::anemo::rpc::codec::{Codec, Decoder, Encoder};
161    use bytes::Buf;
162
163    #[derive(Debug)]
164    pub struct BcsSnappyEncoder<T>(PhantomData<T>);
165
166    impl<T: serde::Serialize> Encoder for BcsSnappyEncoder<T> {
167        type Item = T;
168        type Error = bcs::Error;
169
170        fn encode(&mut self, item: Self::Item) -> Result<bytes::Bytes, Self::Error> {
171            let mut buf = Vec::<u8>::new();
172            let mut snappy_encoder = snap::write::FrameEncoder::new(&mut buf);
173            bcs::serialize_into(&mut snappy_encoder, &item)?;
174            drop(snappy_encoder);
175            Ok(buf.into())
176        }
177    }
178
179    #[derive(Debug)]
180    pub struct BcsSnappyDecoder<U>(PhantomData<U>);
181
182    impl<U: serde::de::DeserializeOwned> Decoder for BcsSnappyDecoder<U> {
183        type Item = U;
184        type Error = bcs::Error;
185
186        fn decode(&mut self, buf: bytes::Bytes) -> Result<Self::Item, Self::Error> {
187            let bytes = super::decompress_snappy(buf.reader())?;
188            bcs::from_bytes(bytes.as_slice())
189        }
190    }
191
192    /// A [`Codec`] that implements `bcs` encoding/decoding via the serde
193    /// library.
194    #[derive(Debug, Clone)]
195    pub struct BcsSnappyCodec<T, U>(PhantomData<(T, U)>);
196
197    impl<T, U> Default for BcsSnappyCodec<T, U> {
198        fn default() -> Self {
199            Self(PhantomData)
200        }
201    }
202
203    impl<T, U> Codec for BcsSnappyCodec<T, U>
204    where
205        T: serde::Serialize + Send + 'static,
206        U: serde::de::DeserializeOwned + Send + 'static,
207    {
208        type Encode = T;
209        type Decode = U;
210        type Encoder = BcsSnappyEncoder<T>;
211        type Decoder = BcsSnappyDecoder<U>;
212
213        fn encoder(&mut self) -> Self::Encoder {
214            BcsSnappyEncoder(PhantomData)
215        }
216
217        fn decoder(&mut self) -> Self::Decoder {
218            BcsSnappyDecoder(PhantomData)
219        }
220
221        fn format_name(&self) -> &'static str {
222            "bcs"
223        }
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use ::anemo::rpc::codec::{
230        Codec as AnemoCodec, Decoder as AnemoDecoder, Encoder as AnemoEncoder,
231    };
232
233    use super::*;
234
235    fn snappy_compress(raw: &[u8]) -> Vec<u8> {
236        let mut out = Vec::new();
237        let mut encoder = snap::write::FrameEncoder::new(&mut out);
238        std::io::Write::write_all(&mut encoder, raw).unwrap();
239        drop(encoder);
240        out
241    }
242
243    #[test]
244    fn anemo_roundtrip() {
245        let mut codec: anemo::BcsSnappyCodec<Vec<u64>, Vec<u64>> = anemo::BcsSnappyCodec::default();
246        let value = vec![1u64, 2, 3, 4, 5, 6, 7, 8, 9, 10];
247        let encoded = codec.encoder().encode(value.clone()).unwrap();
248        let decoded = codec.decoder().decode(encoded).unwrap();
249        assert_eq!(decoded, value);
250    }
251
252    #[test]
253    fn bounded_helper_respects_output_limit() {
254        // With `max_allowed` set below the stream's decompressed size,
255        // `decompress_snappy_bounded` returns exactly `max_allowed` bytes.
256        let raw = vec![0u8; 2 * 1024 * 1024];
257        let compressed = snappy_compress(&raw);
258        let limit = 1024u64;
259        let out = decompress_snappy_bounded(&compressed[..], limit).unwrap();
260        assert_eq!(out.len() as u64, limit);
261        assert!((out.len() as u64) < raw.len() as u64);
262    }
263
264    // Tonic-variant round trip via a bespoke HttpBody. Building a real
265    // `DecodeBuf` requires tonic's internal API, so we drive the decoder
266    // through `tonic::codec::Streaming::new_request`, which is the path used
267    // by the gRPC server. This exercises the real `BcsSnappyDecoder::decode`.
268    mod tonic_via_streaming {
269        use std::{
270            pin::Pin,
271            task::{Context, Poll},
272        };
273
274        use bytes::{BufMut, Bytes, BytesMut};
275        use futures::StreamExt;
276        use http_body::{Body as HttpBody, Frame};
277
278        use super::{super::*, snappy_compress};
279
280        /// Minimal HttpBody yielding a single gRPC-framed payload. gRPC length-
281        /// prefixed frames are `[compression:u8][length:u32 BE][payload]`.
282        struct OneFrameBody(Option<Bytes>);
283
284        impl OneFrameBody {
285            fn new(payload: Bytes) -> Self {
286                let mut framed = BytesMut::with_capacity(5 + payload.len());
287                framed.put_u8(0);
288                framed.put_u32(payload.len() as u32);
289                framed.put_slice(&payload);
290                Self(Some(framed.freeze()))
291            }
292        }
293
294        impl HttpBody for OneFrameBody {
295            type Data = Bytes;
296            type Error = std::convert::Infallible;
297
298            fn poll_frame(
299                mut self: Pin<&mut Self>,
300                _cx: &mut Context<'_>,
301            ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
302                Poll::Ready(self.0.take().map(|b| Ok(Frame::data(b))))
303            }
304
305            fn is_end_stream(&self) -> bool {
306                self.0.is_none()
307            }
308        }
309
310        #[tokio::test]
311        async fn tonic_roundtrip() {
312            let mut codec: BcsSnappyCodec<Vec<u64>, Vec<u64>> = BcsSnappyCodec::default();
313            let value = vec![1u64, 2, 3, 4, 5, 6, 7, 8, 9, 10];
314            let raw = bcs::to_bytes(&value).unwrap();
315            let compressed = snappy_compress(&raw);
316            let body = OneFrameBody::new(Bytes::from(compressed));
317            let mut stream =
318                tonic::codec::Streaming::new_request(codec.decoder(), body, None, None);
319            let decoded = stream.next().await.unwrap().unwrap();
320            assert_eq!(decoded, value);
321        }
322    }
323}
iota_network_stack/codec.rs

iota_network_stack/
codec.rs
