diff --git a/blend/message/src/codec.rs b/blend/message/src/codec.rs new file mode 100644 index 000000000..823347e4f --- /dev/null +++ b/blend/message/src/codec.rs @@ -0,0 +1,180 @@ +use lb_blend_proofs::{ + quota::{PROOF_OF_QUOTA_SIZE, ProofOfQuota}, + selection::{PROOF_OF_SELECTION_SIZE, ProofOfSelection}, +}; +use lb_key_management_system_keys::keys::{ + ED25519_PUBLIC_KEY_SIZE, ED25519_SIGNATURE_SIZE, Ed25519PublicKey, Ed25519Signature, +}; + +/// A content error from decoding a wire component. +/// +/// Length is deliberately NOT checked by any [`WireDecode`] implementation: the +/// caller (the network-side size gate that rejects wrongly-sized peer messages +/// up front) guarantees the input holds at least the bytes the component needs. +/// These variants therefore only cover malformed *values*. +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum WireDecodeError { + #[error("Invalid boolean encoding")] + InvalidBool, + #[error("Invalid Ed25519 public key encoding")] + InvalidPublicKey, + #[error("Invalid proof of quota encoding")] + InvalidProofOfQuota, + #[error("Invalid proof of selection encoding")] + InvalidProofOfSelection, + #[error("Unsupported message version")] + UnsupportedVersion, + #[error("Invalid payload type discriminant")] + InvalidPayloadType, +} + +/// Append a message component's fixed-size, prefix-free wire bytes to `out`. +/// +/// Every Blend message component has a size fully determined by the (fixed, +/// network-wide) number of encapsulation layers, so nothing is length-prefixed. +/// The output buffer is allocated by the caller; implementations only append. +pub trait WireEncode { + fn encode_into(&self, out: &mut Vec); +} + +/// Decode a message component from the front of `input`, returning the value +/// and the unconsumed remainder (`(rest, value)`, as in `nom`). +/// +/// Implementations do not check `input`'s length — the caller guarantees it is +/// large enough (see [`WireDecodeError`]). `Context` carries anything the +/// decoder needs that is not on the wire (e.g. the layer count); `()` for +/// self-describing fixed-size components. +pub trait WireDecode: Sized { + type Context; + + fn decode(input: &[u8], context: Self::Context) -> Result<(&[u8], Self), WireDecodeError>; +} + +impl WireEncode for u8 { + fn encode_into(&self, out: &mut Vec) { + out.push(*self); + } +} + +impl WireDecode for u8 { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + Ok((&input[1..], input[0])) + } +} + +impl WireEncode for u16 { + fn encode_into(&self, out: &mut Vec) { + out.extend_from_slice(&self.to_le_bytes()); + } +} + +impl WireDecode for u16 { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (bytes, remaining) = input.split_at(size_of::()); + let value = Self::from_le_bytes(bytes.try_into().expect("split_at guarantees the length")); + Ok((remaining, value)) + } +} + +impl WireEncode for bool { + fn encode_into(&self, out: &mut Vec) { + u8::from(*self).encode_into(out); + } +} + +impl WireDecode for bool { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (remaining, value) = u8::decode(input, ())?; + match value { + 0 => Ok((remaining, false)), + 1 => Ok((remaining, true)), + _ => Err(WireDecodeError::InvalidBool), + } + } +} + +impl WireEncode for Ed25519PublicKey { + fn encode_into(&self, out: &mut Vec) { + out.extend_from_slice(self.as_bytes()); + } +} + +impl WireDecode for Ed25519PublicKey { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (key_bytes, remaining) = input.split_at(ED25519_PUBLIC_KEY_SIZE); + let key_array: [u8; ED25519_PUBLIC_KEY_SIZE] = key_bytes + .try_into() + .expect("split_at guarantees the length"); + let public_key = + Self::from_bytes(&key_array).map_err(|_| WireDecodeError::InvalidPublicKey)?; + Ok((remaining, public_key)) + } +} + +impl WireEncode for ProofOfQuota { + fn encode_into(&self, out: &mut Vec) { + out.extend_from_slice(&<[u8; PROOF_OF_QUOTA_SIZE]>::from(self)); + } +} + +impl WireDecode for ProofOfQuota { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (proof_bytes, remaining) = input.split_at(PROOF_OF_QUOTA_SIZE); + let proof_array: [u8; PROOF_OF_QUOTA_SIZE] = proof_bytes + .try_into() + .expect("split_at guarantees the length"); + let proof = + Self::try_from(proof_array).map_err(|_| WireDecodeError::InvalidProofOfQuota)?; + Ok((remaining, proof)) + } +} + +impl WireEncode for ProofOfSelection { + fn encode_into(&self, out: &mut Vec) { + out.extend_from_slice(&<[u8; PROOF_OF_SELECTION_SIZE]>::from(self)); + } +} + +impl WireDecode for ProofOfSelection { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (proof_bytes, remaining) = input.split_at(PROOF_OF_SELECTION_SIZE); + let proof_array: [u8; PROOF_OF_SELECTION_SIZE] = proof_bytes + .try_into() + .expect("split_at guarantees the length"); + let proof = + Self::try_from(proof_array).map_err(|_| WireDecodeError::InvalidProofOfSelection)?; + Ok((remaining, proof)) + } +} + +impl WireEncode for Ed25519Signature { + fn encode_into(&self, out: &mut Vec) { + out.extend_from_slice(&self.to_bytes()); + } +} + +impl WireDecode for Ed25519Signature { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (sig_bytes, remaining) = input.split_at(ED25519_SIGNATURE_SIZE); + let sig_array: [u8; ED25519_SIGNATURE_SIZE] = sig_bytes + .try_into() + .expect("split_at guarantees the length"); + // `Ed25519Signature::from_bytes` is infallible (any bytes are a valid + // signature value; verification happens elsewhere). + Ok((remaining, Self::from_bytes(&sig_array))) + } +} diff --git a/blend/message/src/encap/encapsulated.rs b/blend/message/src/encap/encapsulated.rs index 111c9f836..859f24eeb 100644 --- a/blend/message/src/encap/encapsulated.rs +++ b/blend/message/src/encap/encapsulated.rs @@ -1,3 +1,5 @@ +use core::num::NonZeroU64; + use derivative::Derivative; use itertools::Itertools as _; use lb_blend_crypto::{ZkHash, cipher::Cipher}; @@ -5,14 +7,15 @@ use lb_blend_proofs::{ quota::{self, VerifiedProofOfQuota}, selection::{self, VerifiedProofOfSelection, inputs::VerifyInputs}, }; -use lb_core::codec::{DeserializeOp as _, SerializeOp as _}; use lb_key_management_system_keys::keys::{ Ed25519PublicKey, Ed25519Signature, SharedKey, UnsecuredEd25519Key, }; use serde::{Deserialize, Serialize}; +use serde_with::serde_as; use crate::{ Error, PayloadType, + codec::{WireDecode, WireDecodeError, WireEncode}, crypto::{domains, key_ext::SharedKeyExt as _}, encap::{ ProofsVerifier, @@ -23,7 +26,9 @@ use crate::{ }, input::EncapsulationInput, message::{ - BlendingHeader, Payload, PublicHeader, payload::PaddedPayloadBody, + BlendingHeader, Payload, PublicHeader, + blending_header::BLENDING_HEADER_ENCODED_SIZE, + payload::{PAYLOAD_ENCODED_SIZE, PaddedPayloadBody}, public_header::VerifiedPublicHeader, }, }; @@ -59,6 +64,43 @@ impl EncapsulatedMessage { (self.public_header, self.encapsulated_part) } + #[cfg(test)] + // Encoding (and sending) of unverified messages should not be done outside of + // tests, so this function is only available in tests. + #[must_use] + pub fn encode(&self) -> Vec { + let expected_encoded_len = + crate::encap::expected_serialized_len(self.encapsulation_layers()); + let mut out = Vec::with_capacity(expected_encoded_len); + self.public_header.encode_into(&mut out); + self.encapsulated_part.encode_into(&mut out); + debug_assert!( + out.len() == expected_encoded_len, + "Message should encode to the expected length but it did not." + ); + out + } + + #[cfg(test)] + fn encapsulation_layers(&self) -> NonZeroU64 { + self.encapsulated_part.encapsulation_layers() + } + + /// Decode a message from the front of `bytes`, returning it and the + /// unconsumed remainder. + /// + /// This does not check `bytes`'s length nor that it is fully consumed — the + /// caller `bytes` is exactly a well-formed `num_layers`-layer message and + /// checks the remainder. + pub fn decode(bytes: &[u8], num_layers: NonZeroU64) -> Result<(&[u8], Self), Error> { + let (remaining, public_header) = PublicHeader::decode(bytes, ())?; + let (remaining, encapsulated_part) = EncapsulatedPart::decode(remaining, num_layers)?; + Ok(( + remaining, + Self::from_components(public_header, encapsulated_part), + )) + } + /// Verify the message public header signature. pub fn verify_header_signature( self, @@ -245,6 +287,33 @@ impl EncapsulatedPart { pub(super) fn sign(&self, key: &UnsecuredEd25519Key) -> Ed25519Signature { key.sign_payload(&signing_body(&self.private_header, &self.payload)) } + + pub(super) fn encapsulation_layers(&self) -> NonZeroU64 { + self.private_header.encapsulation_layers() + } +} + +impl WireEncode for EncapsulatedPart { + fn encode_into(&self, out: &mut Vec) { + self.private_header.encode_into(out); + self.payload.encode_into(out); + } +} + +impl WireDecode for EncapsulatedPart { + type Context = NonZeroU64; + + fn decode(input: &[u8], context: Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (input, private_header) = EncapsulatedPrivateHeader::decode(input, context)?; + let (input, payload) = EncapsulatedPayload::decode(input, ())?; + Ok(( + input, + Self { + private_header, + payload, + }, + )) + } } /// Verify the public header reconstructed when decapsulating all but the very @@ -301,7 +370,7 @@ fn signing_body( // TODO: Consider having `InitializedPrivateHeader` // that just finished the initialization step and doesn't have `decapsulate` method. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] -pub(super) struct EncapsulatedPrivateHeader(Vec); +pub(super) struct EncapsulatedPrivateHeader(Box<[EncapsulatedBlendingHeader]>); impl EncapsulatedPrivateHeader { #[cfg(test)] @@ -351,7 +420,8 @@ impl EncapsulatedPrivateHeader { }); header }) - .collect(), + .collect::>() + .into_boxed_slice(), ) } @@ -492,22 +562,56 @@ impl EncapsulatedPrivateHeader { .iter() .flat_map(EncapsulatedBlendingHeader::iter_bytes) } + + pub(super) fn encapsulation_layers(&self) -> NonZeroU64 { + NonZeroU64::new(self.0.len() as u64) + .expect("An encapsulated part always has at least one blending header.") + } +} + +impl WireEncode for EncapsulatedPrivateHeader { + fn encode_into(&self, out: &mut Vec) { + for layer in &self.0 { + layer.encode_into(out); + } + } +} + +impl WireDecode for EncapsulatedPrivateHeader { + type Context = NonZeroU64; + + fn decode(mut input: &[u8], context: Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let mut layers = Vec::with_capacity(context.get() as usize); + for _ in 0..context.get() { + let (remaining, layer) = EncapsulatedBlendingHeader::decode(input, ())?; + layers.push(layer); + input = remaining; + } + Ok((input, Self(layers.into_boxed_slice()))) + } } /// A blending header encapsulated zero or more times. -// TODO: Consider having `SerializedBlendingHeader` (not encapsulated). +/// +/// Always exactly [`BLENDING_HEADER_ENCODED_SIZE`] bytes (the cipher is +/// length-preserving), so it is a fixed-size array — stored inline, so a whole +/// [`EncapsulatedPrivateHeader`] is one contiguous allocation. +#[serde_as] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] -struct EncapsulatedBlendingHeader(Vec); +struct EncapsulatedBlendingHeader( + #[serde_as(as = "serde_with::Bytes")] [u8; BLENDING_HEADER_ENCODED_SIZE], +); impl EncapsulatedBlendingHeader { /// Build a [`EncapsulatedBlendingHeader`] by serializing a /// [`BlendingHeader`] without any encapsulation. fn initialize(header: &BlendingHeader) -> Self { + let mut bytes = Vec::with_capacity(BLENDING_HEADER_ENCODED_SIZE); + header.encode_into(&mut bytes); Self( - header - .to_bytes() - .expect("BlendingHeader should be able to be serialized") - .to_vec(), + bytes + .try_into() + .expect("A BlendingHeader always encodes to BLENDING_HEADER_ENCODED_SIZE bytes."), ) } @@ -515,17 +619,19 @@ impl EncapsulatedBlendingHeader { /// If there is no encapsulation left, and if the bytes are valid, /// the deserialization will succeed. fn try_deserialize(&self) -> Result { - BlendingHeader::from_bytes(&self.0).map_err(|_| Error::PrivateHeaderDeserializationFailed) + let (_remaining, header) = BlendingHeader::decode(&self.0, ()) + .map_err(|_| Error::PrivateHeaderDeserializationFailed)?; + Ok(header) } /// Add a layer of encapsulation. fn encapsulate(&mut self, cipher: &mut Cipher) { - cipher.encrypt(self.0.as_mut_slice()); + cipher.encrypt(&mut self.0[..]); } /// Remove a layer of encapsulation. fn decapsulate(&mut self, cipher: &mut Cipher) { - cipher.decrypt(self.0.as_mut_slice()); + cipher.decrypt(&mut self.0[..]); } fn iter_bytes(&self) -> impl Iterator + '_ { @@ -533,20 +639,48 @@ impl EncapsulatedBlendingHeader { } } +// The encapsulated leaves already hold their raw ciphered bytes, so encoding is +// the identity and decoding takes a fixed-size slice of the layer/payload size. +// No length checks: the network-side size gate guarantees the input is large +// enough (`split_at`/`try_into` therefore never fail). +impl WireEncode for EncapsulatedBlendingHeader { + fn encode_into(&self, out: &mut Vec) { + out.extend_from_slice(&self.0); + } +} + +impl WireDecode for EncapsulatedBlendingHeader { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (bytes, remaining) = input.split_at(BLENDING_HEADER_ENCODED_SIZE); + Ok(( + remaining, + Self(bytes.try_into().expect("split_at guarantees the length")), + )) + } +} + /// A payload encapsulated zero or more times. -// TODO: Consider having `SerializedPayload` (not encapsulated). +/// +/// Always exactly [`PAYLOAD_ENCODED_SIZE`] bytes; boxed because that is ~34 KiB +/// and must not be stored inline in +/// [`EncapsulatedPart`]/[`EncapsulatedMessage`]. +#[serde_as] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] -struct EncapsulatedPayload(Vec); +struct EncapsulatedPayload(#[serde_as(as = "serde_with::Bytes")] Box<[u8; PAYLOAD_ENCODED_SIZE]>); impl EncapsulatedPayload { /// Build a [`EncapsulatedPayload`] by serializing a [`Payload`] /// without any encapsulation. fn initialize(payload: &Payload) -> Self { + let mut bytes = Vec::with_capacity(PAYLOAD_ENCODED_SIZE); + payload.encode_into(&mut bytes); Self( - payload - .to_bytes() - .expect("Payload should be able to be serialized") - .to_vec(), + bytes + .into_boxed_slice() + .try_into() + .expect("A Payload always encodes to PAYLOAD_ENCODED_SIZE bytes."), ) } @@ -554,18 +688,20 @@ impl EncapsulatedPayload { /// If there is no encapsulation left, and if the bytes are valid, /// the deserialization will succeed. fn try_deserialize(&self) -> Result { - Payload::from_bytes(&self.0).map_err(|_| Error::PayloadDeserializationFailed) + let (_remaining, payload) = + Payload::decode(&self.0[..], ()).map_err(|_| Error::PayloadDeserializationFailed)?; + Ok(payload) } /// Add a layer of encapsulation. fn encapsulate(mut self, cipher: &mut Cipher) -> Self { - cipher.encrypt(self.0.as_mut_slice()); + cipher.encrypt(&mut self.0[..]); self } /// Remove a layer of encapsulation. fn decapsulate(mut self, cipher: &mut Cipher) -> Self { - cipher.decrypt(self.0.as_mut_slice()); + cipher.decrypt(&mut self.0[..]); self } @@ -573,3 +709,23 @@ impl EncapsulatedPayload { self.0.iter().copied() } } + +impl WireEncode for EncapsulatedPayload { + fn encode_into(&self, out: &mut Vec) { + out.extend_from_slice(&self.0[..]); + } +} + +impl WireDecode for EncapsulatedPayload { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (bytes, remaining) = input.split_at(PAYLOAD_ENCODED_SIZE); + let boxed = bytes + .to_vec() + .into_boxed_slice() + .try_into() + .expect("split_at guarantees the length"); + Ok((remaining, Self(boxed))) + } +} diff --git a/blend/message/src/encap/mod.rs b/blend/message/src/encap/mod.rs index 697300800..b701f7f54 100644 --- a/blend/message/src/encap/mod.rs +++ b/blend/message/src/encap/mod.rs @@ -1,10 +1,18 @@ +use core::num::NonZeroU64; + use lb_blend_proofs::{ quota::{ProofOfQuota, VerifiedProofOfQuota}, selection::{ProofOfSelection, VerifiedProofOfSelection, inputs::VerifyInputs}, }; use lb_key_management_system_keys::keys::Ed25519PublicKey; -use crate::crypto::proofs::PoQVerificationInputsMinusSigningKey; +use crate::{ + crypto::proofs::PoQVerificationInputsMinusSigningKey, + message::{ + blending_header::BLENDING_HEADER_ENCODED_SIZE, payload::PAYLOAD_ENCODED_SIZE, + public_header::PUBLIC_HEADER_ENCODED_SIZE, + }, +}; pub mod decapsulated; pub mod encapsulated; @@ -35,3 +43,20 @@ pub trait ProofsVerifier { inputs: &VerifyInputs, ) -> Result; } + +/// The exact serialized size, in bytes, of any well-formed message with +/// `num_layers` encapsulation layers. +/// +/// The wire format is fixed-size, so this is fully determined by the layer +/// count. Used by the network crate to gate received bytes and to size the +/// encode buffer. +#[must_use] +pub fn expected_serialized_len(num_layers: NonZeroU64) -> usize { + let layers_len = (num_layers.get() as usize) + .checked_mul(BLENDING_HEADER_ENCODED_SIZE) + .expect("message encoded length overflow"); + PUBLIC_HEADER_ENCODED_SIZE + .checked_add(layers_len) + .and_then(|len| len.checked_add(PAYLOAD_ENCODED_SIZE)) + .expect("message encoded length overflow") +} diff --git a/blend/message/src/encap/tests.rs b/blend/message/src/encap/tests.rs index 5591f2a2c..57c3ea844 100644 --- a/blend/message/src/encap/tests.rs +++ b/blend/message/src/encap/tests.rs @@ -16,8 +16,10 @@ use crate::{ ProofsVerifier, decapsulated::DecapsulationOutput, encapsulated::{EncapsulatedMessage, EncapsulatedPart}, + expected_serialized_len, validated::{ - EncapsulatedMessageWithVerifiedPublicHeader, RequiredProofOfSelectionVerificationInputs, + EncapsulatedMessageWithVerifiedPublicHeader, EncapsulatedMessageWithVerifiedSignature, + RequiredProofOfSelectionVerificationInputs, }, }, input::EncapsulationInput, @@ -387,6 +389,77 @@ fn decapsulate_empty_private_headers_returns_error() { assert!(matches!(result, Err(Error::EmptyEncapsulationInputs))); } +fn sample_message(num_layers: usize) -> EncapsulatedMessageWithVerifiedPublicHeader { + let (inputs, _) = generate_inputs(num_layers); + EncapsulatedMessageWithVerifiedPublicHeader::try_new( + &inputs, + PayloadType::Data, + b"payload".as_slice().try_into().unwrap(), + ) + .unwrap() +} + +#[test] +fn serialized_size_constants_match_wire_format() { + // The O(1) size gate in `deserialize_from_remote` relies on + // `expected_serialized_len` being exact. Build real, genuinely-encapsulated + // messages of varying layer counts and confirm the constant-derived length + // matches the actual encoded length — this pins every size constant to the + // real wire encoding. + for num_layers in 1..=4u64 { + let message = EncapsulatedMessage::from(sample_message(num_layers as usize)); + + let actual_len = message.encode().len() as u64; + let expected_len = expected_serialized_len(num_layers.try_into().unwrap()) as u64; + + assert_eq!( + expected_len, actual_len, + "expected_serialized_len mismatch for {num_layers} layer(s)" + ); + } +} + +#[test] +fn encode_decode_round_trip() { + // A message encoded to the wire format and decoded back with the expected + // layer count reconstructs the original. + for num_layers in 1..=4u64 { + let message = EncapsulatedMessage::from(sample_message(num_layers as usize)); + + let encoded = message.encode(); + let (remaining, decoded) = + EncapsulatedMessage::decode(&encoded, num_layers.try_into().unwrap()).unwrap(); + + assert!( + remaining.is_empty(), + "leftover bytes for {num_layers} layer(s)" + ); + assert_eq!( + decoded, message, + "round-trip mismatch for {num_layers} layer(s)" + ); + } +} + +#[test] +fn wire_bytes_identical_across_message_types() { + // The send path serializes a verified variant; the receiver decodes an + // `EncapsulatedMessage`. All three must produce byte-identical wire output. + let with_public_header = sample_message(3); + let with_signature: EncapsulatedMessageWithVerifiedSignature = + with_public_header.clone().into(); + let unverified = EncapsulatedMessage::from(with_public_header.clone()); + + let bytes = with_public_header.encode(); + assert_eq!(with_signature.encode(), bytes); + assert_eq!(unverified.encode(), bytes); +} + +// Rejecting a message whose layer count differs from the expected one is now +// the responsibility of the network-side size gate (it compares the received +// length against `EncapsulatedMessage::expected_serialized_len`), covered by +// the `blend-network` tests. `decode` itself assumes a correctly-sized input. + fn generate_inputs(cnt: usize) -> (Vec, Vec) { let recipient_signing_keys = core::iter::repeat_with(UnsecuredEd25519Key::generate_with_blake_rng) diff --git a/blend/message/src/encap/validated.rs b/blend/message/src/encap/validated.rs index a38babaea..ae6c0bb4c 100644 --- a/blend/message/src/encap/validated.rs +++ b/blend/message/src/encap/validated.rs @@ -1,3 +1,5 @@ +use core::num::NonZeroU64; + use derivative::Derivative; use lb_blend_crypto::random_sized_bytes; use lb_blend_proofs::{ @@ -9,11 +11,13 @@ use serde::{Deserialize, Serialize}; use crate::{ Error, MessageIdentifier, PaddedPayloadBody, PayloadType, + codec::WireEncode as _, crypto::key_ext::Ed25519SecretKeyExt as _, encap::{ ProofsVerifier, decapsulated::{DecapsulatedMessage, DecapsulationOutput, PartDecapsulationOutput}, encapsulated::{EncapsulatedMessage, EncapsulatedPart}, + expected_serialized_len, }, input::EncapsulationInput, message::public_header::{PublicHeaderWithVerifiedSignature, VerifiedPublicHeader}, @@ -81,6 +85,25 @@ impl EncapsulatedMessageWithVerifiedSignature { self.public_header_with_verified_signature.id() } + #[must_use] + pub fn encode(&self) -> Vec { + let expected_encoded_len = expected_serialized_len(self.encapsulation_layers()); + let mut out = Vec::with_capacity(expected_encoded_len); + self.public_header_with_verified_signature + .encode_into(&mut out); + self.encapsulated_part.encode_into(&mut out); + debug_assert!( + out.len() == expected_encoded_len, + "Message should encode to the expected length but it did not." + ); + out + } + + #[must_use] + fn encapsulation_layers(&self) -> NonZeroU64 { + self.encapsulated_part.encapsulation_layers() + } + #[cfg(any(feature = "unsafe-test-functions", test))] pub const fn public_header_mut(&mut self) -> &mut PublicHeaderWithVerifiedSignature { &mut self.public_header_with_verified_signature @@ -277,6 +300,24 @@ impl EncapsulatedMessageWithVerifiedPublicHeader { pub const fn public_header_mut(&mut self) -> &mut VerifiedPublicHeader { &mut self.validated_public_header } + + #[must_use] + pub fn encode(&self) -> Vec { + let expected_encoded_len = expected_serialized_len(self.encapsulation_layers()); + let mut out = Vec::with_capacity(expected_encoded_len); + self.validated_public_header.encode_into(&mut out); + self.encapsulated_part.encode_into(&mut out); + debug_assert!( + out.len() == expected_encoded_len, + "Message should encode to the expected length but it did not." + ); + out + } + + #[must_use] + fn encapsulation_layers(&self) -> NonZeroU64 { + self.encapsulated_part.encapsulation_layers() + } } impl From diff --git a/blend/message/src/error.rs b/blend/message/src/error.rs index 60f569d13..aa43b73e0 100644 --- a/blend/message/src/error.rs +++ b/blend/message/src/error.rs @@ -14,6 +14,8 @@ pub enum Error { ProofOfQuotaVerificationFailed(quota::Error), #[error("Encapsulated message deserialization failed")] MessageDeserializationFailed, + #[error(transparent)] + WireDecode(#[from] crate::codec::WireDecodeError), #[error("Payload deserialization failed")] PayloadDeserializationFailed, #[error("Private header deserialization failed")] diff --git a/blend/message/src/lib.rs b/blend/message/src/lib.rs index 91090bf09..6a170c746 100644 --- a/blend/message/src/lib.rs +++ b/blend/message/src/lib.rs @@ -1,3 +1,4 @@ +mod codec; pub mod crypto; pub mod encap; mod error; diff --git a/blend/message/src/message/blending_header.rs b/blend/message/src/message/blending_header.rs index 58de4f68f..2849c5df6 100644 --- a/blend/message/src/message/blending_header.rs +++ b/blend/message/src/message/blending_header.rs @@ -9,7 +9,10 @@ use lb_key_management_system_keys::keys::{ }; use serde::{Deserialize, Serialize}; -use crate::crypto::domains; +use crate::{ + codec::{WireDecode, WireDecodeError, WireEncode}, + crypto::domains, +}; /// A blending header that is fully decapsulated. /// This must be encapsulated when being sent to the blend network. @@ -61,3 +64,49 @@ impl BlendingHeader { fn concat(a: &[u8], b: &[u8]) -> Vec { a.iter().chain(b.iter()).copied().collect::>() } + +/// The exact number of bytes a [`BlendingHeader`] encodes to. Every field is +/// fixed-size, so this is a compile-time constant — which lets the encapsulated +/// (ciphered) form be stored as a `[u8; BLENDING_HEADER_ENCODED_SIZE]`. +pub const BLENDING_HEADER_ENCODED_SIZE: usize = ED25519_PUBLIC_KEY_SIZE + .checked_add(PROOF_OF_QUOTA_SIZE) + .unwrap() + .checked_add(ED25519_SIGNATURE_SIZE) + .unwrap() + .checked_add(PROOF_OF_SELECTION_SIZE) + .unwrap() + .checked_add(size_of::()) + .unwrap(); + +impl WireEncode for BlendingHeader { + fn encode_into(&self, out: &mut Vec) { + self.signing_pubkey.encode_into(out); + self.proof_of_quota.encode_into(out); + self.signature.encode_into(out); + self.proof_of_selection.encode_into(out); + self.is_last.encode_into(out); + } +} + +impl WireDecode for BlendingHeader { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (input, signing_pubkey) = Ed25519PublicKey::decode(input, ())?; + let (input, proof_of_quota) = ProofOfQuota::decode(input, ())?; + let (input, signature) = Ed25519Signature::decode(input, ())?; + let (input, proof_of_selection) = ProofOfSelection::decode(input, ())?; + let (input, is_last) = bool::decode(input, ())?; + + Ok(( + input, + Self { + signing_pubkey, + proof_of_quota, + signature, + proof_of_selection, + is_last, + }, + )) + } +} diff --git a/blend/message/src/message/payload.rs b/blend/message/src/message/payload.rs index ea5a5b648..5fc43ed06 100644 --- a/blend/message/src/message/payload.rs +++ b/blend/message/src/message/payload.rs @@ -1,32 +1,62 @@ use serde::{Deserialize, Serialize}; use serde_with::serde_as; -use crate::Error; +use crate::{ + Error, + codec::{WireDecode, WireDecodeError, WireEncode}, +}; pub const MAX_PAYLOAD_BODY_SIZE: usize = 34 * 1024; -/// A payload header that is fully decapsulated. -/// This must be encapsulated when being sent to the blend network. -#[derive(Clone, Serialize, Deserialize)] -struct PayloadHeader { - payload_type: PayloadType, - body_len: u16, -} - #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[repr(u8)] pub enum PayloadType { Cover = 0x00, Data = 0x01, } +impl TryFrom for PayloadType { + type Error = (); + + fn try_from(value: u8) -> Result { + match value { + 0x00 => Ok(Self::Cover), + 0x01 => Ok(Self::Data), + _ => Err(()), + } + } +} + +impl WireEncode for PayloadType { + fn encode_into(&self, out: &mut Vec) { + (*self as u8).encode_into(out); + } +} + +impl WireDecode for PayloadType { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (remaining, discriminant) = u8::decode(input, ())?; + let payload_type = + Self::try_from(discriminant).map_err(|()| WireDecodeError::InvalidPayloadType)?; + Ok((remaining, payload_type)) + } +} + +/// The decapsulated payload body, padded to a fixed size. +/// +/// `actual_len` is the length of the real (unpadded) content and is the single +/// source of truth for it — the payload no longer stores it a second time. #[serde_as] #[derive(Clone, Serialize, Deserialize)] pub struct PaddedPayloadBody { - /// A body is padded to [`MAX_PAYLOAD_BODY_SIZE`], - /// Box is used to not allocate a big array on the stack. + /// The real content length; `padded[..actual_len]` is the body. + actual_len: u16, + /// A body padded to [`MAX_PAYLOAD_BODY_SIZE`]. `Box` avoids a large stack + /// allocation. #[serde_as(as = "serde_with::Bytes")] padded: Box<[u8; MAX_PAYLOAD_BODY_SIZE]>, - actual_len: u16, } impl TryFrom> for PaddedPayloadBody { @@ -45,7 +75,7 @@ impl TryFrom<&[u8]> for PaddedPayloadBody { return Err(Error::PayloadTooLarge); } - let body_len: u16 = value + let actual_len: u16 = value .len() .try_into() .map_err(|_| Error::InvalidPayloadLength)?; @@ -56,42 +86,60 @@ impl TryFrom<&[u8]> for PaddedPayloadBody { .expect("body must be created with the correct size"); padded[..value.len()].copy_from_slice(value); - Ok(Self { - actual_len: body_len, - padded, - }) + Ok(Self { actual_len, padded }) } } +impl WireEncode for PaddedPayloadBody { + fn encode_into(&self, out: &mut Vec) { + self.actual_len.encode_into(out); + out.extend_from_slice(&self.padded[..]); + } +} + +impl WireDecode for PaddedPayloadBody { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (input, actual_len) = u16::decode(input, ())?; + let (body_bytes, remaining) = input.split_at(MAX_PAYLOAD_BODY_SIZE); + let padded: Box<[u8; MAX_PAYLOAD_BODY_SIZE]> = body_bytes + .to_vec() + .into_boxed_slice() + .try_into() + .expect("split_at guarantees the length"); + Ok((remaining, Self { actual_len, padded })) + } +} + +/// The exact number of bytes a [`Payload`] encodes to: a fixed enum +/// discriminant, the `u16` body length, and the body padded to +/// [`MAX_PAYLOAD_BODY_SIZE`]. Compile-time constant, so the encapsulated +/// (ciphered) form can be stored as a `Box<[u8; PAYLOAD_ENCODED_SIZE]>`. +pub const PAYLOAD_ENCODED_SIZE: usize = + size_of::() + size_of::() + MAX_PAYLOAD_BODY_SIZE; + /// A payload that is fully decapsulated. /// This must be encapsulated when being sent to the blend network. -#[serde_as] #[derive(Clone, Serialize, Deserialize)] pub struct Payload { - header: PayloadHeader, + payload_type: PayloadType, body: PaddedPayloadBody, } impl Payload { - pub const fn new(payload_type: PayloadType, payload_body: PaddedPayloadBody) -> Self { - Self { - header: PayloadHeader { - payload_type, - body_len: payload_body.actual_len, - }, - body: payload_body, - } + pub const fn new(payload_type: PayloadType, body: PaddedPayloadBody) -> Self { + Self { payload_type, body } } pub const fn payload_type(&self) -> PayloadType { - self.header.payload_type + self.payload_type } /// Returns the payload body unpadded. - /// Returns an error if the payload cannot be read up to the length - /// specified in the header + /// Returns an error if the recorded length exceeds the padded buffer. pub fn body(&self) -> Result<&[u8], Error> { - let len = self.header.body_len as usize; + let len = self.body.actual_len as usize; if self.body.padded.len() < len { return Err(Error::InvalidPayloadLength); } @@ -102,3 +150,20 @@ impl Payload { Ok((self.payload_type(), self.body()?.to_vec())) } } + +impl WireEncode for Payload { + fn encode_into(&self, out: &mut Vec) { + self.payload_type.encode_into(out); + self.body.encode_into(out); + } +} + +impl WireDecode for Payload { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (input, payload_type) = PayloadType::decode(input, ())?; + let (input, body) = PaddedPayloadBody::decode(input, ())?; + Ok((input, Self { payload_type, body })) + } +} diff --git a/blend/message/src/message/public_header.rs b/blend/message/src/message/public_header.rs index b4fb9e332..50365b17d 100644 --- a/blend/message/src/message/public_header.rs +++ b/blend/message/src/message/public_header.rs @@ -1,11 +1,22 @@ -use lb_blend_proofs::quota::{self, ProofOfQuota, VerifiedProofOfQuota}; -use lb_key_management_system_keys::keys::{Ed25519PublicKey, Ed25519Signature}; +use lb_blend_proofs::quota::{self, PROOF_OF_QUOTA_SIZE, ProofOfQuota, VerifiedProofOfQuota}; +use lb_key_management_system_keys::keys::{ + ED25519_PUBLIC_KEY_SIZE, ED25519_SIGNATURE_SIZE, Ed25519PublicKey, Ed25519Signature, +}; use serde::{Deserialize, Deserializer, Serialize, de}; -use crate::{Error, MessageIdentifier, encap::ProofsVerifier}; +use crate::{ + Error, MessageIdentifier, + codec::{WireDecode, WireDecodeError, WireEncode}, + encap::ProofsVerifier, +}; const LATEST_BLEND_MESSAGE_VERSION: u8 = 1; +/// The exact number of bytes a [`PublicHeader`] encodes to (a version byte plus +/// fixed-size fields). Compile-time constant. +pub const PUBLIC_HEADER_ENCODED_SIZE: usize = + size_of::() + ED25519_PUBLIC_KEY_SIZE + PROOF_OF_QUOTA_SIZE + ED25519_SIGNATURE_SIZE; + // A public header that is revealed to all nodes. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] pub struct PublicHeader { @@ -102,6 +113,38 @@ impl PublicHeader { } } +impl WireEncode for PublicHeader { + fn encode_into(&self, out: &mut Vec) { + self.version.encode_into(out); + self.signing_pubkey.encode_into(out); + self.proof_of_quota.encode_into(out); + self.signature.encode_into(out); + } +} + +impl WireDecode for PublicHeader { + type Context = (); + + fn decode(input: &[u8], (): Self::Context) -> Result<(&[u8], Self), WireDecodeError> { + let (input, version) = u8::decode(input, ())?; + if version != LATEST_BLEND_MESSAGE_VERSION { + return Err(WireDecodeError::UnsupportedVersion); + } + let (input, signing_pubkey) = Ed25519PublicKey::decode(input, ())?; + let (input, proof_of_quota) = ProofOfQuota::decode(input, ())?; + let (input, signature) = Ed25519Signature::decode(input, ())?; + Ok(( + input, + Self { + version, + signing_pubkey, + proof_of_quota, + signature, + }, + )) + } +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] pub struct PublicHeaderWithVerifiedSignature { version: u8, @@ -175,6 +218,20 @@ impl PublicHeaderWithVerifiedSignature { } } +// The verified public-header variants are never decoded from the wire (a peer's +// bytes always decode into an unverified `PublicHeader`); they only need to +// encode, and all three variants produce identical bytes. Implementing only +// `WireEncode` for them means a verified message can be serialized directly, +// with no conversion/copy through `PublicHeader`. +impl WireEncode for PublicHeaderWithVerifiedSignature { + fn encode_into(&self, out: &mut Vec) { + self.version.encode_into(out); + self.signing_pubkey.encode_into(out); + self.proof_of_quota.encode_into(out); + self.signature.encode_into(out); + } +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] pub struct VerifiedPublicHeader { version: u8, @@ -273,6 +330,15 @@ impl VerifiedPublicHeader { } } +impl WireEncode for VerifiedPublicHeader { + fn encode_into(&self, out: &mut Vec) { + self.version.encode_into(out); + self.signing_pubkey.encode_into(out); + self.proof_of_quota.as_ref().encode_into(out); + self.signature.encode_into(out); + } +} + #[cfg(test)] mod tests { use lb_blend_proofs::quota::VerifiedProofOfQuota; diff --git a/blend/network/src/core/with_core/behaviour/mod.rs b/blend/network/src/core/with_core/behaviour/mod.rs index 4132c9b70..30fd9af9e 100644 --- a/blend/network/src/core/with_core/behaviour/mod.rs +++ b/blend/network/src/core/with_core/behaviour/mod.rs @@ -1,6 +1,6 @@ use core::{ mem::{self}, - num::NonZeroUsize, + num::{NonZeroU64, NonZeroUsize}, }; use std::{ collections::{HashMap, VecDeque, hash_map::Entry}, @@ -59,6 +59,10 @@ pub struct Config { pub peering_degree: RangeInclusive, /// The minimum Blend network size for messages to be relayed between peers. pub minimum_network_size: NonZeroUsize, + /// `ß_c`: the fixed number of encapsulation layers every well-formed Blend + /// message carries. Used to validate the layout of messages received from + /// remote peers before processing them. + pub num_blend_layers: NonZeroU64, } #[derive(Debug, Clone, Copy)] @@ -121,6 +125,9 @@ pub struct Behaviour { protocol_name: StreamProtocol, /// The minimum Blend network size for messages to be relayed between peers. minimum_network_size: NonZeroUsize, + /// `ß_c`: the fixed number of encapsulation layers every well-formed Blend + /// message carries. + num_blend_layers: NonZeroU64, /// States for processing messages from the old epoch /// before the transition period has passed. old_epoch: Option, @@ -239,6 +246,7 @@ impl Behaviour { local_peer_id, protocol_name, minimum_network_size: config.minimum_network_size, + num_blend_layers: config.num_blend_layers, old_epoch: None, } } @@ -266,6 +274,7 @@ impl Behaviour { .collect(), mem::take(&mut self.message_cache), current_epoch_number, + self.num_blend_layers, )); tracing::debug!(target: LOG_TARGET, "Started a new epoch by passing negotiated peers and exchanged message IDs to the old epoch. Now, no negotiated peers in the current epoch."); @@ -933,6 +942,7 @@ impl Behaviour { &mut self.events, &mut self.waker, self.current_epoch_info.1, + self.num_blend_layers, ) { tracing::debug!(target: LOG_TARGET, "Failed to handle message from the current epoch: {receive_error:?}"); let spam_reason = match receive_error { diff --git a/blend/network/src/core/with_core/behaviour/old_epoch.rs b/blend/network/src/core/with_core/behaviour/old_epoch.rs index 06909241b..1c791bd80 100644 --- a/blend/network/src/core/with_core/behaviour/old_epoch.rs +++ b/blend/network/src/core/with_core/behaviour/old_epoch.rs @@ -1,6 +1,7 @@ use std::{ collections::{HashMap, VecDeque, hash_map::Entry}, convert::Infallible, + num::NonZeroU64, task::{Context, Poll, Waker}, }; @@ -38,6 +39,7 @@ pub struct OldEpoch { waker: Option, message_cache: MessageCache, epoch: Epoch, + num_blend_layers: NonZeroU64, } impl OldEpoch { @@ -46,6 +48,7 @@ impl OldEpoch { negotiated_peers: HashMap, message_cache: MessageCache, epoch: Epoch, + num_blend_layers: NonZeroU64, ) -> Self { Self { negotiated_peers, @@ -53,6 +56,7 @@ impl OldEpoch { events: VecDeque::new(), waker: None, epoch, + num_blend_layers, } } @@ -154,6 +158,7 @@ impl OldEpoch { &mut self.events, &mut self.waker, self.epoch, + self.num_blend_layers, ).inspect_err(|receive_error| { tracing::debug!(target: LOG_TARGET, "Failed to handle message from the old epoch: {receive_error:?}. Closing connection with spammy peer."); self.events.push_back(ToSwarm::NotifyHandler { diff --git a/blend/network/src/core/with_core/behaviour/tests/message_handling.rs b/blend/network/src/core/with_core/behaviour/tests/message_handling.rs index 34aa42f66..574d8a0d9 100644 --- a/blend/network/src/core/with_core/behaviour/tests/message_handling.rs +++ b/blend/network/src/core/with_core/behaviour/tests/message_handling.rs @@ -2,6 +2,7 @@ use core::time::Duration; use std::collections::HashSet; use futures::StreamExt as _; +use lb_blend_scheduling::serialize_encapsulated_message_with_verified_public_header; use lb_libp2p::SwarmEvent; use libp2p_swarm_test::SwarmExt as _; use test_log::test; @@ -145,6 +146,63 @@ async fn undeserializable_message_received() { } } +#[test(tokio::test)] +async fn message_with_unexpected_layer_count_disconnects_peer() { + // The listening node expects a single encapsulation layer, but the sender + // delivers a well-formed 3-layer message. The size gate in + // `EncapsulatedMessage::deserialize_from_remote` rejects it up front, and + // the sender is treated exactly like one delivering undeserializable bytes. + let (mut identities, nodes) = new_nodes_with_empty_address(2); + let mut dialing_swarm = TestSwarm::new(&identities.next().unwrap(), |id| { + BehaviourBuilder::new(id).with_membership(&nodes).build() + }); + let mut listening_swarm = TestSwarm::new(&identities.next().unwrap(), |id| { + BehaviourBuilder::new(id) + .with_membership(&nodes) + .with_num_blend_layers(1) + .build() + }); + + listening_swarm.listen().with_memory_addr_external().await; + dialing_swarm + .connect_and_wait_for_upgrade(&mut listening_swarm) + .await; + + let message = TestEncapsulatedMessage::new(b"unexpected_layer_count"); + dialing_swarm + .behaviour_mut() + .force_send_serialized_message_to_current_epoch_peer( + serialize_encapsulated_message_with_verified_public_header(message.as_ref()), + *listening_swarm.local_peer_id(), + ) + .unwrap(); + + let mut events_to_match = 2u8; + loop { + select! { + _ = dialing_swarm.select_next_some() => {} + listening_swarm_event = listening_swarm.select_next_some() => { + match listening_swarm_event { + SwarmEvent::Behaviour(Event::PeerDisconnected(peer_id, peer_state)) => { + assert_eq!(peer_id, *dialing_swarm.local_peer_id()); + assert_eq!(peer_state, NegotiatedPeerState::Spammy(SpamReason::UndeserializableMessage)); + events_to_match -= 1; + } + SwarmEvent::ConnectionClosed { peer_id, endpoint, .. } => { + assert_eq!(peer_id, *dialing_swarm.local_peer_id()); + assert!(endpoint.is_listener()); + events_to_match -= 1; + } + _ => {} + } + } + } + if events_to_match == 0 { + break; + } + } +} + #[test(tokio::test)] async fn duplicate_message_received_from_same_peer() { let (mut identities, nodes) = new_nodes_with_empty_address(2); diff --git a/blend/network/src/core/with_core/behaviour/tests/utils.rs b/blend/network/src/core/with_core/behaviour/tests/utils.rs index ad0679bf6..08827238d 100644 --- a/blend/network/src/core/with_core/behaviour/tests/utils.rs +++ b/blend/network/src/core/with_core/behaviour/tests/utils.rs @@ -1,4 +1,8 @@ -use core::{num::NonZeroUsize, ops::RangeInclusive, time::Duration}; +use core::{ + num::{NonZeroU64, NonZeroUsize}, + ops::RangeInclusive, + time::Duration, +}; use std::{ collections::{HashMap, VecDeque}, iter::repeat_with, @@ -87,6 +91,7 @@ pub struct BehaviourBuilder { provider: Option, peering_degree: Option>, minimum_network_size: Option, + num_blend_layers: Option, } impl BehaviourBuilder { @@ -97,6 +102,7 @@ impl BehaviourBuilder { provider: None, peering_degree: None, minimum_network_size: None, + num_blend_layers: None, } } @@ -124,6 +130,11 @@ impl BehaviourBuilder { self } + pub fn with_num_blend_layers(mut self, num_blend_layers: u64) -> Self { + self.num_blend_layers = Some(num_blend_layers.try_into().unwrap()); + self + } + pub fn build(self) -> Behaviour { Behaviour { negotiated_peers: HashMap::new(), @@ -144,6 +155,9 @@ impl BehaviourBuilder { minimum_network_size: self .minimum_network_size .unwrap_or_else(|| 1usize.try_into().unwrap()), + num_blend_layers: self + .num_blend_layers + .unwrap_or_else(|| 3.try_into().unwrap()), old_epoch: None, message_cache: MessageCache::new(), } diff --git a/blend/network/src/core/with_core/behaviour/utils.rs b/blend/network/src/core/with_core/behaviour/utils.rs index ce93c968f..47eff9cfd 100644 --- a/blend/network/src/core/with_core/behaviour/utils.rs +++ b/blend/network/src/core/with_core/behaviour/utils.rs @@ -1,4 +1,4 @@ -use core::{convert::Infallible, task::Waker}; +use core::{convert::Infallible, num::NonZeroU64, task::Waker}; use std::collections::VecDeque; use either::Either; @@ -76,10 +76,12 @@ pub fn handle_received_serialized_encapsulated_message_and_update_cache( events_queue: &mut VecDeque>>, waker: &mut Option, epoch: Epoch, + num_blend_layers: NonZeroU64, ) -> Result<(), ReceiveError> { // Deserialize the message. - let deserialized_encapsulated_message = deserialize_encapsulated_message(serialized_message) - .map_err(|_| ReceiveError::UndeserializableMessage)?; + let deserialized_encapsulated_message = + deserialize_encapsulated_message(serialized_message, num_blend_layers) + .map_err(|_| ReceiveError::UndeserializableMessage)?; // Add the message to the set of exchanged message identifiers with the sender, // returning `Err` if the message was already sent by this peer previously. diff --git a/blend/network/src/core/with_edge/behaviour/mod.rs b/blend/network/src/core/with_edge/behaviour/mod.rs index ad0cd4858..1dff25dad 100644 --- a/blend/network/src/core/with_edge/behaviour/mod.rs +++ b/blend/network/src/core/with_edge/behaviour/mod.rs @@ -1,4 +1,4 @@ -use core::num::NonZeroUsize; +use core::num::{NonZeroU64, NonZeroUsize}; use std::{ collections::{HashSet, VecDeque}, convert::Infallible, @@ -51,6 +51,10 @@ pub struct Config { pub connection_timeout: Duration, pub max_incoming_connections: usize, pub minimum_network_size: NonZeroUsize, + /// `ß_c`: the fixed number of encapsulation layers every well-formed Blend + /// message carries. Used to validate the layout of messages received from + /// remote peers before processing them. + pub num_blend_layers: NonZeroU64, } /// A [`NetworkBehaviour`]: @@ -67,6 +71,7 @@ pub struct Behaviour { max_incoming_connections: usize, protocol_name: StreamProtocol, minimum_network_size: NonZeroUsize, + num_blend_layers: NonZeroU64, } impl Behaviour { @@ -85,6 +90,7 @@ impl Behaviour { max_incoming_connections: config.max_incoming_connections, protocol_name, minimum_network_size: config.minimum_network_size, + num_blend_layers: config.num_blend_layers, } } @@ -153,7 +159,7 @@ impl Behaviour { fn handle_received_serialized_encapsulated_message(&mut self, serialized_message: &[u8]) { let Ok(deserialized_encapsulated_message) = - deserialize_encapsulated_message(serialized_message) + deserialize_encapsulated_message(serialized_message, self.num_blend_layers) else { tracing::trace!(target: LOG_TARGET, "Failed to deserialize received message. Ignoring..."); return; diff --git a/blend/network/src/core/with_edge/behaviour/tests/message_handling.rs b/blend/network/src/core/with_edge/behaviour/tests/message_handling.rs index 3413078e2..675955d1d 100644 --- a/blend/network/src/core/with_edge/behaviour/tests/message_handling.rs +++ b/blend/network/src/core/with_edge/behaviour/tests/message_handling.rs @@ -50,6 +50,50 @@ async fn receive_valid_message() { } } +#[test(tokio::test)] +async fn reject_message_with_unexpected_layer_count() { + // The behaviour is configured to expect a single encapsulation layer, but + // the shared test helper builds a 3-layer message. The size gate in + // `EncapsulatedMessage::deserialize_from_remote` must reject it up front, + // before the (larger) message is fully parsed or its header processed. + let mut core_swarm = TestSwarm::new_ephemeral(|_| { + BehaviourBuilder::new(PeerId::random()) + .with_num_blend_layers(1) + .build() + }); + let mut edge_swarm = TestSwarm::new_ephemeral(|_| StreamBehaviour::new()); + + core_swarm.listen().with_memory_addr_external().await; + let stream = edge_swarm + .connect_and_upgrade_to_blend(&mut core_swarm) + .await; + let message = TestEncapsulatedMessage::new(b"unexpected_layer_count"); + send_msg( + stream, + serialize_encapsulated_message_with_verified_public_header(message.as_ref()), + ) + .await + .unwrap(); + + loop { + select! { + _ = edge_swarm.select_next_some() => {} + core_swarm_event = core_swarm.select_next_some() => { + match core_swarm_event { + SwarmEvent::Behaviour(Event::Message(_)) => { + panic!("No `Message` event should be generated for a message with an unexpected number of layers."); + } + SwarmEvent::ConnectionClosed { peer_id, .. } => { + assert_eq!(peer_id, *edge_swarm.local_peer_id()); + break; + } + _ => {} + } + } + } + } +} + #[test(tokio::test)] async fn message_timeout() { let mut core_swarm = TestSwarm::new_ephemeral(|_| { diff --git a/blend/network/src/core/with_edge/behaviour/tests/utils.rs b/blend/network/src/core/with_edge/behaviour/tests/utils.rs index e7087e184..9073a72ff 100644 --- a/blend/network/src/core/with_edge/behaviour/tests/utils.rs +++ b/blend/network/src/core/with_edge/behaviour/tests/utils.rs @@ -1,4 +1,7 @@ -use core::{num::NonZeroUsize, time::Duration}; +use core::{ + num::{NonZeroU64, NonZeroUsize}, + time::Duration, +}; use std::collections::{HashSet, VecDeque}; use async_trait::async_trait; @@ -20,6 +23,7 @@ pub struct BehaviourBuilder { max_incoming_connections: Option, timeout: Option, minimum_network_size: Option, + num_blend_layers: Option, } impl BehaviourBuilder { @@ -29,9 +33,15 @@ impl BehaviourBuilder { max_incoming_connections: None, timeout: None, minimum_network_size: None, + num_blend_layers: None, } } + pub fn with_num_blend_layers(mut self, num_blend_layers: u64) -> Self { + self.num_blend_layers = Some(num_blend_layers.try_into().unwrap()); + self + } + pub fn with_max_incoming_connections(mut self, max_incoming_connections: usize) -> Self { self.max_incoming_connections = Some(max_incoming_connections); self @@ -71,6 +81,9 @@ impl BehaviourBuilder { minimum_network_size: self .minimum_network_size .unwrap_or_else(|| 1usize.try_into().unwrap()), + num_blend_layers: self + .num_blend_layers + .unwrap_or_else(|| 3.try_into().unwrap()), } } } diff --git a/blend/scheduling/src/message_blend/crypto/mod.rs b/blend/scheduling/src/message_blend/crypto/mod.rs index e37cb76cd..918ed4944 100644 --- a/blend/scheduling/src/message_blend/crypto/mod.rs +++ b/blend/scheduling/src/message_blend/crypto/mod.rs @@ -10,7 +10,6 @@ use lb_blend_message::{ }, }, }; -use lb_core::codec::{DeserializeOp as _, SerializeOp}; use lb_key_management_system_keys::keys::X25519PrivateKey; pub mod core_and_leader; @@ -39,26 +38,23 @@ pub struct EpochCryptographicProcessorSettings { pub fn serialize_encapsulated_message_with_verified_public_header( message: &EncapsulatedMessageWithVerifiedPublicHeader, ) -> Vec { - serialize_message(message) + message.encode() } #[must_use] pub fn serialize_encapsulated_message_with_verified_signature( message: &EncapsulatedMessageWithVerifiedSignature, ) -> Vec { - serialize_message(message) + message.encode() } -fn serialize_message(message: &Message) -> Vec -where - Message: SerializeOp, -{ - message - .to_bytes() - .expect("Message should be serializable") - .to_vec() -} - -pub fn deserialize_encapsulated_message(message: &[u8]) -> Result { - EncapsulatedMessage::from_bytes(message).map_err(|_| Error::MessageDeserializationFailed) +pub fn deserialize_encapsulated_message( + message: &[u8], + num_blend_layers: NonZeroU64, +) -> Result { + let (remaining, deserialized_message) = EncapsulatedMessage::decode(message, num_blend_layers)?; + if !remaining.is_empty() { + return Err(Error::MessageDeserializationFailed); + } + Ok(deserialized_message) } diff --git a/services/blend/src/core/backends/libp2p/behaviour.rs b/services/blend/src/core/backends/libp2p/behaviour.rs index 725b4ff33..f3cb54d71 100644 --- a/services/blend/src/core/backends/libp2p/behaviour.rs +++ b/services/blend/src/core/backends/libp2p/behaviour.rs @@ -39,11 +39,13 @@ where peering_degree: minimum_core_healthy_peering_degree ..=maximum_core_peering_degree, minimum_network_size: config.minimum_network_size.try_into().unwrap(), + num_blend_layers: config.num_blend_layers, }, with_edge: lb_blend::network::core::with_edge::behaviour::Config { connection_timeout: config.backend.edge_node_connection_timeout, max_incoming_connections: maximum_edge_incoming_connections, minimum_network_size: config.minimum_network_size.try_into().unwrap(), + num_blend_layers: config.num_blend_layers, }, }, observation_window_interval_provider, diff --git a/services/blend/src/core/backends/libp2p/tests/utils.rs b/services/blend/src/core/backends/libp2p/tests/utils.rs index dd5c471d9..611465ce8 100644 --- a/services/blend/src/core/backends/libp2p/tests/utils.rs +++ b/services/blend/src/core/backends/libp2p/tests/utils.rs @@ -198,11 +198,13 @@ impl BlendBehaviourBuilder { with_core: CoreToCoreConfig { peering_degree, minimum_network_size: 1.try_into().unwrap(), + num_blend_layers: 3.try_into().unwrap(), }, with_edge: CoreToEdgeConfig { connection_timeout: Duration::from_secs(1), max_incoming_connections: 300, minimum_network_size: 1.try_into().unwrap(), + num_blend_layers: 3.try_into().unwrap(), }, }, TestObservationWindowProvider {