pb_mapper_protocol/
lib.rs

1//! Define message protocols and tools for reading and writing
2//! messages
3//!
4//! The reader and writer traits are `async fn` in a public trait, which cannot
5//! state its auto-trait bounds. That is deliberate: these are only ever awaited
6//! on the connection task that owns the stream, never sent across one.
7#![allow(async_fn_in_trait)]
8
9pub mod buffer;
10pub mod command;
11pub mod data;
12pub mod forward;
13mod frame_read;
14pub mod secure;
15use snafu::{ResultExt, ensure};
16use tokio::io::{AsyncReadExt, AsyncWriteExt};
17
18use crate::buffer::{BufferGetter, CommonBuffer, FixedSizeBuffer};
19use pb_mapper_core::checksum::{
20    AesKeyType, get_checksum, get_checksum_for_key, get_msg_header_key, process_checksum_is_ready,
21    valid_checksum, valid_checksum_for_key,
22};
23use pb_mapper_core::codec::{Aes256GcmDeCodec, Aes256GcmEnCodec, Decryptor, Encryptor};
24use pb_mapper_core::error::MsgDatalenExceededSnafu;
25use pb_mapper_core::error::{
26    self, MsgDatalenValidateSnafu, MsgNetworkReadBodySnafu, MsgNetworkReadCheckSumSnafu,
27    MsgNetworkWriteBodySnafu, MsgNetworkWriteCheckSumSnafu, MsgNetworkWriteCodecMsgSnafu,
28    MsgNetworkWriteCodecTagSnafu, MsgNetworkWriteDatalenSnafu, Result,
29};
30
31/// This message protocol contains header and body, and the header
32/// includes checksum, datalen,respectively, u32, u32, where datalen
33/// represents the length of the body, checksum is used to check the
34/// datalen field. This is just the most basic pedestal protocol, in
35/// order to solve the sticky packet problem with TCP streams. We
36/// can build more advanced communication on top of this protocol, for
37/// example, we can use json or other forms of data representation
38///
39/// ```text
40/// ┌─────────────┐
41/// │ u32 checksum│
42/// │ u32 datalen │
43/// └─────────────┘
44/// ┌─────────────┐
45/// │ Body        │
46/// │(Actual Data)│
47/// └─────────────┘
48/// ```
49pub trait MessageReader {
50    async fn read_msg(&mut self) -> Result<&'_ [u8]>;
51}
52
53pub trait MessageWriter {
54    async fn write_msg(&mut self, msg: &[u8]) -> Result<()>;
55
56    /// TODO: Implement this method to fix the encryption zero-copy data corruption bug
57    ///
58    /// This method should be used for encryption scenarios where:
59    /// - Zero-copy performance is needed
60    /// - The caller can provide mutable data
61    /// - No unsafe transmutation is required
62    ///
63    /// Default implementation falls back to the immutable version for compatibility.
64    /// Encryption implementations should override this for true zero-copy operation.
65    async fn write_msg_mut(&mut self, msg: &mut [u8]) -> Result<()> {
66        // Default implementation: delegate to immutable version
67        // This maintains backward compatibility but doesn't solve the corruption issue
68        self.write_msg(msg).await
69    }
70}
71
72/// Maximum plaintext payload size read from local services.
73const MAX_PLAINTEXT_LEN: DataLenType = 8 * 1024 * 1024;
74/// AES-GCM tag length (bytes). Keep in sync with ring's tag length.
75const CODEC_TAG_LEN: DataLenType = 16;
76/// Maximum value of `datalen` to prevent Out of Memory.
77/// For encrypted frames, the tag is appended to the payload.
78const MAX_MSG_LEN: DataLenType = MAX_PLAINTEXT_LEN + CODEC_TAG_LEN;
79
80// Defined in `pb-mapper-core` so that the checksum and error types can name it
81// without depending on this module. Re-exported rather than redeclared: a second
82// `pub type` would be a distinct name for the same width, and the two would read
83// as unrelated at the crate boundary.
84pub use pb_mapper_core::DataLenType;
85
86macro_rules! gen_write_network_with_error {
87    ($func_name:ident, $write_method:ident, $error:expr, $input_type:ty) => {
88        #[inline]
89        async fn $func_name<T: AsyncWriteExt + Unpin>(
90            writer: &mut T,
91            data: $input_type,
92        ) -> Result<()> {
93            writer.$write_method(data).await.context($error)
94        }
95    };
96}
97
98gen_write_network_with_error!(write_checksum, write_u32, MsgNetworkWriteCheckSumSnafu, u32);
99
100gen_write_network_with_error!(write_datalen, write_u32, MsgNetworkWriteDatalenSnafu, u32);
101
102gen_write_network_with_error!(write_msg_body, write_all, MsgNetworkWriteBodySnafu, &[u8]);
103
104gen_write_network_with_error!(
105    write_codec_msg,
106    write_all,
107    MsgNetworkWriteCodecMsgSnafu,
108    &[u8]
109);
110
111gen_write_network_with_error!(
112    write_codec_tag,
113    write_all,
114    MsgNetworkWriteCodecTagSnafu,
115    &[u8]
116);
117
118fn checksum_key_bytes(key: &Option<AesKeyType>) -> Option<&[u8]> {
119    key.as_ref().map(|key| key.as_slice())
120}
121
122#[inline]
123fn checksum_matches(datalen: DataLenType, checksum: u32, key: Option<&[u8]>) -> bool {
124    match key {
125        Some(key) => valid_checksum_for_key(datalen, checksum, key),
126        None => valid_checksum(datalen, checksum),
127    }
128}
129
130#[inline]
131fn checksum_for(len: DataLenType, key: Option<&[u8]>) -> Result<u32> {
132    match key {
133        Some(key) => Ok(get_checksum_for_key(len, key)),
134        None => {
135            if !process_checksum_is_ready() {
136                return Err(error::Error::MsgCodec {
137                    action: "load configured credential",
138                    detail:
139                        "`MSG_HEADER_KEY` is required; no insecure default checksum is available"
140                            .to_string(),
141                });
142            }
143            Ok(get_checksum(len))
144        }
145    }
146}
147
148#[inline]
149async fn set_msg_len<T: AsyncWriteExt + Unpin>(
150    writer: &mut T,
151    len: DataLenType,
152    checksum_key: Option<&[u8]>,
153) -> Result<()> {
154    write_checksum(writer, checksum_for(len, checksum_key)?).await?;
155    write_datalen(writer, len).await
156}
157
158pub struct NormalMessageReader<'a, T: AsyncReadExt + Unpin> {
159    reader: &'a mut T,
160    buffer: CommonBuffer,
161    frame: frame_read::FrameRead<8>,
162    checksum_key: Option<AesKeyType>,
163}
164
165impl<'a, T: AsyncReadExt + Unpin> NormalMessageReader<'a, T> {
166    pub fn new(reader: &'a mut T) -> Self {
167        Self {
168            reader,
169            buffer: CommonBuffer::new(),
170            frame: frame_read::FrameRead::new(),
171            checksum_key: None,
172        }
173    }
174
175    pub fn with_checksum_key(mut self, key: AesKeyType) -> Self {
176        self.checksum_key = Some(key);
177        self
178    }
179
180    async fn read_msg_inner<const ALLOW_EOF: bool>(&mut self) -> Result<Option<&'_ [u8]>> {
181        if ALLOW_EOF {
182            if !self
183                .frame
184                .header_or_eof(self.reader)
185                .await
186                .context(MsgNetworkReadCheckSumSnafu)?
187            {
188                return Ok(None);
189            }
190        } else {
191            self.frame
192                .header(self.reader)
193                .await
194                .context(MsgNetworkReadCheckSumSnafu)?;
195        }
196        let header = self.frame.header;
197        let checksum = u32::from_be_bytes([header[0], header[1], header[2], header[3]]);
198        let datalen = u32::from_be_bytes([header[4], header[5], header[6], header[7]]);
199        ensure!(
200            checksum_matches(datalen, checksum, checksum_key_bytes(&self.checksum_key)),
201            MsgDatalenValidateSnafu { datalen, checksum }
202        );
203        ensure!(
204            datalen <= MAX_MSG_LEN,
205            MsgDatalenExceededSnafu {
206                actual: datalen,
207                max: MAX_MSG_LEN
208            }
209        );
210        self.buffer.fixed_resize(datalen as usize);
211        self.frame
212            .body(self.reader, self.buffer.buffer_mut())
213            .await
214            .context(MsgNetworkReadBodySnafu)?;
215        self.frame.finish();
216        Ok(Some(self.buffer.buffer()))
217    }
218}
219
220impl<'a, T: AsyncReadExt + Unpin> MessageReader for NormalMessageReader<'a, T> {
221    async fn read_msg(&mut self) -> Result<&'_ [u8]> {
222        self.read_msg_inner::<false>()
223            .await
224            .map(|message| message.unwrap_or_default())
225    }
226}
227
228pub struct NormalMessageWriter<'a, T: AsyncWriteExt> {
229    writer: &'a mut T,
230    checksum_key: Option<AesKeyType>,
231}
232
233impl<'a, T: AsyncWriteExt + Unpin> NormalMessageWriter<'a, T> {
234    pub fn new(writer: &'a mut T) -> Self {
235        Self {
236            writer,
237            checksum_key: None,
238        }
239    }
240
241    pub fn with_checksum_key(mut self, key: AesKeyType) -> Self {
242        self.checksum_key = Some(key);
243        self
244    }
245
246    async fn write_msg_inner(&mut self, msg: &[u8]) -> Result<()> {
247        set_msg_len(
248            &mut self.writer,
249            msg.len() as u32,
250            checksum_key_bytes(&self.checksum_key),
251        )
252        .await?;
253
254        write_msg_body(&mut self.writer, msg).await
255    }
256}
257
258impl<'a, T: AsyncWriteExt + Unpin> MessageWriter for NormalMessageWriter<'a, T> {
259    async fn write_msg(&mut self, msg: &[u8]) -> Result<()> {
260        self.write_msg_inner(msg).await
261    }
262}
263
264pub struct CodecMessageReader<'a, T: AsyncReadExt + Unpin, D: Decryptor> {
265    reader: NormalMessageReader<'a, T>,
266    decryptor: D,
267}
268
269impl<'a, T: AsyncReadExt + Unpin, D: Decryptor> CodecMessageReader<'a, T, D> {
270    pub fn new(reader: &'a mut T, decryptor: D) -> Self {
271        Self {
272            reader: NormalMessageReader::new(reader),
273            decryptor,
274        }
275    }
276
277    /// Bind the length checksum to `key` instead of the process credential.
278    /// Isolated relays keep a remote `MSG_HEADER_KEY` while speaking with a
279    /// different local administrator key.
280    pub fn for_session_key(reader: &'a mut T, decryptor: D, key: AesKeyType) -> Self {
281        Self::new(reader, decryptor).with_checksum_key(key)
282    }
283
284    pub fn with_checksum_key(mut self, key: AesKeyType) -> Self {
285        self.reader.checksum_key = Some(key);
286        self
287    }
288}
289
290impl<T: AsyncReadExt + Unpin, D: Decryptor> CodecMessageReader<'_, T, D> {
291    /// Byte-stream framing permits FIN between records. Control/datagram readers remain strict.
292    pub(crate) async fn read_stream(&mut self) -> Result<&'_ [u8]> {
293        let Some(message) = self.reader.read_msg_inner::<true>().await? else {
294            return Ok(&[]);
295        };
296        let n = message.len();
297        self.decrypt(n)
298    }
299    fn decrypt(&mut self, n: usize) -> Result<&'_ [u8]> {
300        self.decryptor
301            .decrypt(&mut self.reader.buffer.buffer_mut()[..n])
302            .map(|bytes| &*bytes)
303            .map_err(|e| error::Error::MsgCodec {
304                action: "decrypt",
305                detail: format!("got {e} when we read msg"),
306            })
307    }
308}
309impl<T: AsyncReadExt + Unpin, D: Decryptor> MessageReader for CodecMessageReader<'_, T, D> {
310    async fn read_msg(&mut self) -> Result<&'_ [u8]> {
311        let n = self.reader.read_msg().await?.len();
312        self.decrypt(n)
313    }
314}
315
316/// NOTE: We copy input data before encryption to avoid mutating shared buffers.
317/// This trades some performance for correctness until a zero-copy mutable API is added.
318pub struct CodecMessageWriter<'a, T: AsyncWriteExt + Unpin, E: Encryptor> {
319    writer: &'a mut T,
320    encryptor: E,
321    /// `None` uses the process `MSG_HEADER_KEY` hash. Isolated relays must set
322    /// this to the session key so continuation frames stay decryptable.
323    checksum_key: Option<AesKeyType>,
324}
325
326impl<'a, T: AsyncWriteExt + Unpin, E: Encryptor> CodecMessageWriter<'a, T, E> {
327    pub fn new(writer: &'a mut T, encryptor: E) -> Self {
328        Self {
329            writer,
330            encryptor,
331            checksum_key: None,
332        }
333    }
334
335    pub fn for_session_key(writer: &'a mut T, encryptor: E, key: AesKeyType) -> Self {
336        Self::new(writer, encryptor).with_checksum_key(key)
337    }
338
339    pub fn with_checksum_key(mut self, key: AesKeyType) -> Self {
340        self.checksum_key = Some(key);
341        self
342    }
343
344    pub async fn shutdown(&mut self) -> std::io::Result<()> {
345        self.writer.shutdown().await
346    }
347}
348
349impl<'a, T: AsyncWriteExt + Unpin, E: Encryptor> MessageWriter for CodecMessageWriter<'a, T, E> {
350    async fn write_msg(&mut self, msg: &[u8]) -> Result<()> {
351        let mut buf = msg.to_vec();
352        let tag = self
353            .encryptor
354            .encrypt(&mut buf)
355            .map_err(|e| error::Error::MsgCodec {
356                action: "encrypt",
357                detail: format!("got {e} when we read msg"),
358            })?;
359        let msg_len = (buf.len() + tag.as_ref().len()) as DataLenType;
360
361        set_msg_len(self.writer, msg_len, checksum_key_bytes(&self.checksum_key)).await?;
362        write_codec_msg(self.writer, &buf).await?;
363        write_codec_tag(self.writer, tag.as_ref()).await
364    }
365}
366
367#[inline]
368pub fn get_header_msg_reader<T: AsyncReadExt + Unpin>(
369    reader: &mut T,
370) -> Result<CodecMessageReader<'_, T, Aes256GcmDeCodec>> {
371    Ok(CodecMessageReader::new(reader, get_default_decodec()?))
372}
373
374#[inline]
375pub fn get_header_msg_writer<T: AsyncWriteExt + Unpin>(
376    writer: &mut T,
377) -> Result<CodecMessageWriter<'_, T, Aes256GcmEnCodec>> {
378    Ok(CodecMessageWriter::new(writer, get_default_encodec()?))
379}
380
381#[inline]
382pub fn get_default_encodec() -> Result<Aes256GcmEnCodec> {
383    let key = get_msg_header_key().map_err(|detail| error::Error::MsgCodec {
384        action: "load configured credential",
385        detail,
386    })?;
387    Aes256GcmEnCodec::try_new(&key).map_err(|e| error::Error::MsgCodec {
388        action: "create default encodec",
389        detail: format!("{e}"),
390    })
391}
392
393#[inline]
394pub fn get_default_decodec() -> Result<Aes256GcmDeCodec> {
395    let key = get_msg_header_key().map_err(|detail| error::Error::MsgCodec {
396        action: "load configured credential",
397        detail,
398    })?;
399    Aes256GcmDeCodec::try_new(&key).map_err(|e| error::Error::MsgCodec {
400        action: "create default decodec",
401        detail: format!("{e}"),
402    })
403}
404
405#[inline]
406pub fn get_encodec(key: &[u8]) -> Result<Aes256GcmEnCodec> {
407    Aes256GcmEnCodec::try_new(key).map_err(|e| error::Error::MsgCodec {
408        action: "create encodec",
409        detail: format!("{e}"),
410    })
411}
412
413#[inline]
414pub fn get_decodec(key: &[u8]) -> Result<Aes256GcmDeCodec> {
415    Aes256GcmDeCodec::try_new(key).map_err(|e| error::Error::MsgCodec {
416        action: "create decodec",
417        detail: format!("{e}"),
418    })
419}
420
421pub mod direct;
422
423#[cfg(feature = "relay-quic")]
424pub mod quic;
425pub mod transport;
pb_mapper_protocol/lib.rs

pb_mapper_protocol/
lib.rs
