1
0
mirror of synced 2025-01-23 22:18:54 +00:00

Mix: Packet algorithms (#859)

* add keyset

* add packet

* add comments

* add comments

* use functions in sphinx
This commit is contained in:
Youngjoon Lee 2024-10-29 17:51:52 +09:00 committed by GitHub
parent 9d52297cdf
commit 0c0aae0712
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 274 additions and 6 deletions

View File

@ -4,4 +4,15 @@ version = "0.1.0"
edition = "2021"
[dependencies]
sha2 = "0.10.8"
serde = { version = "1.0", features = ["derive"] }
sha2 = "0.10"
sphinx-packet = "0.2"
thiserror = "1.0.65"
x25519-dalek = { version = "2.0.1", features = [
"getrandom",
"static_secrets",
"serde",
] }
[dev-dependencies]
nomos-core = { path = "../../nomos-core/chain-defs" }

View File

@ -1,10 +1,14 @@
#[derive(Debug)]
#[derive(thiserror::Error, Debug)]
pub enum Error {
/// Invalid mix message format
#[error("Invalid mix message format")]
InvalidMixMessage,
/// Payload size is too large
#[error("Payload is too large")]
PayloadTooLarge,
/// Unwrapping a message is not allowed
/// (e.g. the message cannot be unwrapped using the private key provided)
#[error("Too many recipients")]
TooManyRecipients,
#[error("Sphinx packet error: {0}")]
SphinxPacketError(#[from] sphinx_packet::Error),
#[error("Unwrapping a message is not allowed to this node")]
/// e.g. the message cannot be unwrapped using the private key provided
MsgUnwrapNotAllowed,
}

View File

@ -1,4 +1,5 @@
mod error;
pub mod packet;
pub use error::Error;
@ -7,6 +8,8 @@ use sha2::{Digest, Sha256};
pub const MSG_SIZE: usize = 2048;
pub const DROP_MESSAGE: [u8; MSG_SIZE] = [0; MSG_SIZE];
// TODO: Remove all the mock below once the actual implementation is integrated to the system.
//
/// A mock implementation of the Sphinx encoding.
///
/// The length of the encoded message is fixed to [`MSG_SIZE`] bytes.

View File

@ -0,0 +1,250 @@
use crate::Error;
use serde::{Deserialize, Serialize};
use sphinx_packet::constants::NODE_ADDRESS_LENGTH;
/// A packet that contains a header and a payload.
/// The header and payload are encrypted for the selected recipients.
/// This packet can be serialized and sent over the network.
#[derive(Debug, Serialize, Deserialize)]
pub struct Packet {
header: Header,
// This crate doesn't limit the payload size.
// Users must choose the appropriate constant size and implement padding.
payload: Vec<u8>,
}
/// The packet header
#[derive(Debug, Serialize, Deserialize)]
struct Header {
/// The ephemeral public key for a recipient to derive the shared secret
/// which can be used to decrypt the header and payload.
ephemeral_public_key: x25519_dalek::PublicKey,
// TODO: Length-preserved layered encryption on RoutingInfo
routing_info: RoutingInfo,
}
#[derive(Debug, Serialize, Deserialize)]
struct RoutingInfo {
// TODO: Change this to `is_final_layer: bool`
// by implementing the length-preserved layered encryption.
// It's not good to expose the info that how many layers remain to the intermediate recipients.
remaining_layers: u8,
// TODO:: Add the following fields
// header_integrity_hamc
// additional data (e.g. incentivization)
}
impl Packet {
pub fn build(
recipient_pubkeys: &[x25519_dalek::PublicKey],
payload: &[u8],
payload_size: usize,
) -> Result<Self, Error> {
// Derive `[sphinx_packet::header::keys::KeyMaterial]` for all recipients.
let ephemeral_privkey = x25519_dalek::StaticSecret::random();
let key_material = Self::derive_key_material(recipient_pubkeys, &ephemeral_privkey);
// Encrypt the payload for all recipients.
let payload_keys = key_material
.routing_keys
.iter()
.map(|key_material| key_material.payload_key)
.collect::<Vec<_>>();
let payload = sphinx_packet::payload::Payload::encapsulate_message(
payload,
&payload_keys,
payload_size,
)?;
Ok(Packet {
header: Header {
ephemeral_public_key: x25519_dalek::PublicKey::from(&ephemeral_privkey),
routing_info: RoutingInfo {
remaining_layers: u8::try_from(recipient_pubkeys.len())
.map_err(|_| Error::TooManyRecipients)?,
},
},
payload: payload.into_bytes(),
})
}
fn derive_key_material(
recipient_pubkeys: &[x25519_dalek::PublicKey],
ephemeral_privkey: &x25519_dalek::StaticSecret,
) -> sphinx_packet::header::keys::KeyMaterial {
// NodeAddress is needed to build [`sphinx_packet::route::Node`]s
// required by [`sphinx_packet::header::keys::KeyMaterial::derive`],
// but it's not actually used inside. So, we can use a dummy address.
let dummy_node_address =
sphinx_packet::route::NodeAddressBytes::from_bytes([0u8; NODE_ADDRESS_LENGTH]);
let route = recipient_pubkeys
.iter()
.map(|pubkey| sphinx_packet::route::Node {
address: dummy_node_address,
pub_key: *pubkey,
})
.collect::<Vec<_>>();
sphinx_packet::header::keys::KeyMaterial::derive(&route, ephemeral_privkey)
}
pub fn unpack(
&self,
private_key: &x25519_dalek::StaticSecret,
) -> Result<UnpackedPacket, Error> {
// Derive the routing keys for the recipient
let routing_keys = sphinx_packet::header::SphinxHeader::compute_routing_keys(
&self.header.ephemeral_public_key,
private_key,
);
// Decrypt one layer of encryption on the payload
let payload = sphinx_packet::payload::Payload::from_bytes(&self.payload)?;
let payload = payload.unwrap(&routing_keys.payload_key)?;
// If this is the last layer of encryption, return the decrypted payload.
if self.header.routing_info.remaining_layers == 1 {
return Ok(UnpackedPacket::FullyUnpacked(payload.recover_plaintext()?));
}
// Derive the new ephemeral public key for the next recipient
let next_ephemeral_pubkey = Self::derive_next_ephemeral_public_key(
&self.header.ephemeral_public_key,
&routing_keys.blinding_factor,
);
Ok(UnpackedPacket::ToForward(Packet {
header: Header {
ephemeral_public_key: next_ephemeral_pubkey,
routing_info: RoutingInfo {
remaining_layers: self.header.routing_info.remaining_layers - 1,
},
},
payload: payload.into_bytes(),
}))
}
/// Derive the next ephemeral public key for the next recipient.
//
// This is a copy of `blind_the_shared_secret` from https://github.com/nymtech/sphinx/blob/344b902df340e0d5af69c5147b05f76f324b8cef/src/header/mod.rs#L234.
// with renaming the function name and the arguments
// because the original function is not exposed to the public.
// This logic is tightly coupled with the [`sphinx_packet::header::keys::KeyMaterial::derive`].
fn derive_next_ephemeral_public_key(
cur_ephemeral_pubkey: &x25519_dalek::PublicKey,
blinding_factor: &x25519_dalek::StaticSecret,
) -> x25519_dalek::PublicKey {
let new_shared_secret = blinding_factor.diffie_hellman(cur_ephemeral_pubkey);
x25519_dalek::PublicKey::from(new_shared_secret.to_bytes())
}
}
pub enum UnpackedPacket {
ToForward(Packet),
FullyUnpacked(Vec<u8>),
}
#[cfg(test)]
mod tests {
use nomos_core::wire;
use super::*;
#[test]
fn unpack() {
// Prepare keys of two recipients
let recipient_privkeys = (0..2)
.map(|_| x25519_dalek::StaticSecret::random())
.collect::<Vec<_>>();
let recipient_pubkeys = recipient_privkeys
.iter()
.map(x25519_dalek::PublicKey::from)
.collect::<Vec<_>>();
// Build a packet
let payload = [10u8; 512];
let packet = Packet::build(&recipient_pubkeys, &payload, 1024).unwrap();
// The 1st recipient unpacks the packet
let packet = match packet.unpack(&recipient_privkeys[0]).unwrap() {
UnpackedPacket::ToForward(packet) => packet,
UnpackedPacket::FullyUnpacked(_) => {
panic!("The unpacked packet should be the ToFoward type");
}
};
// The 2nd recipient unpacks the packet
match packet.unpack(&recipient_privkeys[1]).unwrap() {
UnpackedPacket::ToForward(_) => {
panic!("The unpacked packet should be the FullyUnpacked type");
}
UnpackedPacket::FullyUnpacked(unpacked_payload) => {
// Check if the payload has been decrypted correctly
assert_eq!(unpacked_payload, payload);
}
}
}
#[test]
fn unpack_with_wrong_keys() {
// Build a packet with two public keys
let payload = [10u8; 512];
let packet = Packet::build(
&(0..2)
.map(|_| x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::random()))
.collect::<Vec<_>>(),
&payload,
1024,
)
.unwrap();
// The 1st recipient unpacks the packet with an wrong key
let packet = match packet
.unpack(&x25519_dalek::StaticSecret::random())
.unwrap()
{
UnpackedPacket::ToForward(packet) => packet,
UnpackedPacket::FullyUnpacked(_) => {
panic!("The unpacked packet should be the ToFoward type");
}
};
// The 2nd recipient unpacks the packet with an wrong key
assert!(packet
.unpack(&x25519_dalek::StaticSecret::random())
.is_err());
}
#[test]
fn consistent_size_serialization() {
// Prepare keys of two recipients
let recipient_privkeys = (0..2)
.map(|_| x25519_dalek::StaticSecret::random())
.collect::<Vec<_>>();
let recipient_pubkeys = recipient_privkeys
.iter()
.map(x25519_dalek::PublicKey::from)
.collect::<Vec<_>>();
// Build a packet
let payload = [10u8; 512];
let packet = Packet::build(&recipient_pubkeys, &payload, 1024).unwrap();
// Calculate the expected packet size
let pubkey_size = 32;
let routing_info_size = 1;
let payload_length_enconding_size = 8;
let payload_size = 1024;
let packet_size =
pubkey_size + routing_info_size + payload_length_enconding_size + payload_size;
// The serialized packet size must be the same as the expected size.
assert_eq!(wire::serialize(&packet).unwrap().len(), packet_size);
// The unpacked packet size must be the same as the original packet size.
match packet.unpack(&recipient_privkeys[0]).unwrap() {
UnpackedPacket::ToForward(packet) => {
assert_eq!(wire::serialize(&packet).unwrap().len(), packet_size);
}
UnpackedPacket::FullyUnpacked(_) => {
panic!("The unpacked packet should be the ToFoward type");
}
}
}
}