Skip to main content

linera_rpc/simple/
codec.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{io, mem, ops::DerefMut};
5
6use bytes::{Buf, BufMut, BytesMut};
7use linera_core::node::NodeError;
8use thiserror::Error;
9use tokio_util::codec::{Decoder, Encoder};
10
11use crate::RpcMessage;
12
13/// The size of the frame prefix that contains the payload size.
14#[expect(
15    clippy::cast_possible_truncation,
16    reason = "size_of::<u32>() is always 4"
17)]
18const PREFIX_SIZE: u8 = mem::size_of::<u32>() as u8;
19
20/// An encoder/decoder of [`RpcMessage`]s for the RPC protocol.
21///
22/// The frames are length-delimited by a [`u32`] prefix, and the payload is deserialized by
23/// [`bincode`].
24#[derive(Clone, Copy, Debug)]
25pub struct Codec;
26
27impl Encoder<RpcMessage> for Codec {
28    type Error = Error;
29
30    fn encode(&mut self, message: RpcMessage, buffer: &mut BytesMut) -> Result<(), Self::Error> {
31        let mut frame_buffer = buffer.split_off(buffer.len());
32
33        frame_buffer.put_u32_le(0);
34
35        let mut frame_writer = frame_buffer.writer();
36
37        bincode::serialize_into(&mut frame_writer, &message)
38            .map_err(|error| Error::Serialization(*error))?;
39
40        let mut frame_buffer = frame_writer.into_inner();
41        let frame_size = frame_buffer.len();
42        let payload_size = frame_size - PREFIX_SIZE as usize;
43
44        let mut start_of_frame = frame_buffer.deref_mut();
45
46        start_of_frame.put_u32_le(
47            payload_size
48                .try_into()
49                .map_err(|_| Error::MessageTooBig { size: payload_size })?,
50        );
51
52        buffer.unsplit(frame_buffer);
53
54        Ok(())
55    }
56}
57
58impl Decoder for Codec {
59    type Item = RpcMessage;
60    type Error = Error;
61
62    fn decode(&mut self, buffer: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
63        if buffer.len() < PREFIX_SIZE.into() {
64            return Ok(None);
65        }
66
67        let mut start_of_buffer: &[u8] = &*buffer;
68        let payload_size = start_of_buffer
69            .get_u32_le()
70            .try_into()
71            .expect("u32 should fit in a usize");
72
73        let frame_size = PREFIX_SIZE as usize + payload_size;
74
75        if buffer.len() < frame_size {
76            buffer.reserve(frame_size);
77            return Ok(None);
78        }
79
80        let _prefix = buffer.split_to(PREFIX_SIZE.into());
81        let payload = buffer.split_to(payload_size);
82
83        let message =
84            bincode::deserialize(&payload).map_err(|error| Error::Deserialization(*error))?;
85
86        Ok(Some(message))
87    }
88}
89
90/// Errors that can arise during transmission or reception of [`RpcMessage`]s.
91#[derive(Debug, Error)]
92#[allow(missing_docs)]
93pub enum Error {
94    #[error("I/O error in the underlying transport: {0}")]
95    IoError(#[from] io::Error),
96
97    #[error("Failed to deserialize an incoming message: {0}")]
98    Deserialization(#[source] bincode::ErrorKind),
99
100    #[error("Failed to serialize outgoing message: {0}")]
101    Serialization(#[source] bincode::ErrorKind),
102
103    #[error("RpcMessage is too big to fit in a protocol frame: \
104        message is {size} bytes but can't be larger than {max} bytes.",
105        max = u32::MAX)]
106    MessageTooBig { size: usize },
107}
108
109impl From<Error> for NodeError {
110    fn from(error: Error) -> NodeError {
111        match error {
112            Error::IoError(io_error) => NodeError::ClientIoError {
113                error: format!("{io_error}"),
114            },
115            err => {
116                tracing::error!("Unexpected decoding error: {err}");
117                NodeError::InvalidDecoding
118            }
119        }
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    #![allow(clippy::cast_possible_truncation)]
126
127    use bytes::{BufMut, BytesMut};
128    use linera_core::data_types::ChainInfoQuery;
129    use test_strategy::proptest;
130    use tokio_util::codec::{Decoder, Encoder};
131
132    use super::{Codec, RpcMessage, PREFIX_SIZE};
133
134    /// Test decoding of a frame from a buffer.
135    ///
136    /// The buffer may contain leading or trailing bytes around the frame. The frame contains the
137    /// size of the payload, and the payload is a serialized dummy [`RpcMessage`].
138    ///
139    /// The decoder should produce the exact same message as used as the test input, and it should
140    /// ignore the leading and trailing bytes.
141    #[proptest]
142    fn decodes_frame_ignoring_leading_and_trailing_bytes(
143        leading_bytes: Vec<u8>,
144        message_contents: ChainInfoQuery,
145        trailing_bytes: Vec<u8>,
146    ) {
147        let message = RpcMessage::ChainInfoQuery(Box::new(message_contents));
148        let payload = bincode::serialize(&message).expect("RpcMessage is serializable");
149
150        let mut buffer = BytesMut::with_capacity(
151            leading_bytes.len() + PREFIX_SIZE as usize + payload.len() + trailing_bytes.len(),
152        );
153
154        buffer.extend_from_slice(&leading_bytes);
155
156        let start_of_buffer = buffer.split();
157
158        buffer.put_u32_le(payload.len() as u32);
159        buffer.extend_from_slice(&payload);
160        buffer.extend_from_slice(&trailing_bytes);
161
162        let result = Codec.decode(&mut buffer);
163
164        assert!(result.is_ok());
165        assert_eq!(result.unwrap(), Some(message));
166
167        assert_eq!(&start_of_buffer, &leading_bytes);
168        assert_eq!(&buffer, &trailing_bytes);
169    }
170
171    /// Test encoding a message to buffer.
172    ///
173    /// The buffer may already contain some leading bytes, but the cursor is set to where the frame
174    /// should start.
175    ///
176    /// The encoder should write a prefix with the size of the serialized message, followed by the
177    /// serialized message bytes. It should not touch the leading bytes nor append any trailing
178    /// bytes.
179    #[proptest]
180    fn encodes_at_the_correct_buffer_offset(
181        leading_bytes: Vec<u8>,
182        message_contents: ChainInfoQuery,
183    ) {
184        let message = RpcMessage::ChainInfoQuery(Box::new(message_contents));
185        let serialized_message =
186            bincode::serialize(&message).expect("Serialization should succeed");
187
188        let mut buffer = BytesMut::new();
189
190        buffer.extend_from_slice(&leading_bytes);
191
192        let frame_start = buffer.len();
193        let prefix_end = frame_start + PREFIX_SIZE as usize;
194
195        let result = Codec.encode(message, &mut buffer);
196
197        assert!(matches!(result, Ok(())));
198        assert_eq!(&buffer[..frame_start], &leading_bytes);
199
200        let prefix = u32::from_le_bytes(
201            buffer[frame_start..prefix_end]
202                .try_into()
203                .expect("Incorrect prefix slice indices"),
204        );
205
206        assert_eq!(prefix as usize, serialized_message.len());
207        assert_eq!(
208            buffer.len(),
209            leading_bytes.len() + PREFIX_SIZE as usize + prefix as usize
210        );
211
212        assert_eq!(&buffer[prefix_end..], &serialized_message);
213    }
214}