diff --git a/Cargo.lock b/Cargo.lock index 7bb9027f83..f155696d07 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7084,6 +7084,7 @@ dependencies = [ "serde", "thiserror 2.0.17", "tokio", + "tokio-util", "tracing", ] diff --git a/common/nymsphinx/Cargo.toml b/common/nymsphinx/Cargo.toml index 4a472f706a..e5cd7941ae 100644 --- a/common/nymsphinx/Cargo.toml +++ b/common/nymsphinx/Cargo.toml @@ -16,6 +16,7 @@ thiserror = { workspace = true } serde = { workspace = true, features = ["derive"] } bincode = { workspace = true } log = { workspace = true } +tokio-util = { workspace = true, features = ["codec"] } nym-sphinx-acknowledgements = { workspace = true } nym-sphinx-addressing = { workspace = true } diff --git a/common/nymsphinx/src/receiver.rs b/common/nymsphinx/src/receiver.rs index eb83ec43ea..ddd584109b 100644 --- a/common/nymsphinx/src/receiver.rs +++ b/common/nymsphinx/src/receiver.rs @@ -1,6 +1,8 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 +use std::io; + use crate::message::{NymMessage, NymMessageError, PaddedMessage, PlainMessage}; use nym_crypto::aes::cipher::{KeyIvInit, StreamCipher}; use nym_crypto::asymmetric::x25519; @@ -15,10 +17,13 @@ use nym_sphinx_chunking::reconstruction::MessageReconstructor; use nym_sphinx_params::{ PacketEncryptionAlgorithm, PacketHkdfAlgorithm, ReplySurbEncryptionAlgorithm, }; +use serde::{Deserialize, Serialize}; use thiserror::Error; +use tokio_util::bytes::{Buf, BytesMut}; +use tokio_util::codec::{Decoder, Encoder}; // TODO: should this live in this file? -#[derive(Debug)] +#[derive(Debug, Serialize, Deserialize)] pub struct ReconstructedMessage { /// The actual plaintext message that was received. pub message: Vec, @@ -56,6 +61,48 @@ impl From for ReconstructedMessage { } } +pub struct ReconstructedMessageCodec; + +impl Encoder for ReconstructedMessageCodec { + type Error = MessageRecoveryError; + + fn encode(&mut self, item: ReconstructedMessage, buf: &mut BytesMut) -> Result<(), Self::Error> { + let encoded = bincode::serialize(&item).expect("failed to serialize ReconstructedMessage"); + let encoded_len = encoded.len() as u32; + let mut encoded_with_len = encoded_len.to_le_bytes().to_vec(); + encoded_with_len.extend(encoded); + buf.reserve(encoded_with_len.len()); + buf.extend_from_slice(&encoded_with_len); + Ok(()) + } +} + +impl Decoder for ReconstructedMessageCodec { + type Item = ReconstructedMessage; + type Error = MessageRecoveryError; + + fn decode(&mut self, buf: &mut BytesMut) -> Result, Self::Error> { + if buf.len() < 4 { + return Ok(None); + } + + let len = u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]) as usize; + if buf.len() < len + 4 { + return Ok(None); + } + + let decoded = match bincode::deserialize(&buf[4..len]) { + Ok(decoded) => decoded, + Err(_) => return Ok(None) + }; + + buf.advance(len + 4); + + Ok(Some(decoded)) + } +} + + #[derive(Debug, Error)] pub enum MessageRecoveryError { #[error( @@ -75,6 +122,9 @@ pub enum MessageRecoveryError { #[error("Failed to recover message fragment - {0}")] FragmentRecoveryError(#[from] ChunkingError), + + #[error("Failed to recover message fragment - {0}")] + MessageRecoveryError(#[from] io::Error), } pub trait MessageReceiver {