feat(lp-client): add LP session management and UDP data plane support
- Add wrap_data() and session_id() to LpRegistrationClient for LP packet creation after handshake - Add init_lp_session() and close_lp_session() to SpeedtestClient for managing LP sessions - Extract prepare_sphinx_fragments() helper to reduce code duplication between send_data_with_surbs() and send_data_via_lp() - Add send_data_via_lp() for sending Sphinx packets through LP's UDP data plane (port 51264) The LP session is kept alive after TCP handshake closes, allowing subsequent wrap_data() calls for UDP transmission without re-handshaking.
This commit is contained in:
committed by
Jędrzej Stuczyński
parent
1cfa79abc9
commit
b924729992
@@ -1134,6 +1134,90 @@ impl LpRegistrationClient {
|
||||
|
||||
Ok(response_data.to_vec())
|
||||
}
|
||||
|
||||
/// Wrap data in an LP packet for UDP transmission to the data plane (port 51264).
|
||||
///
|
||||
/// This method encrypts the provided data using the established LP session
|
||||
/// and returns serialized LP packet bytes ready to send over UDP.
|
||||
///
|
||||
/// # Prerequisites
|
||||
/// - Handshake must be completed (`perform_handshake()`)
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `data` - Raw application data to wrap (e.g., Sphinx packet bytes)
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Ok(Vec<u8>)` - Serialized LP packet bytes (outer header + encrypted payload)
|
||||
///
|
||||
/// # Wire Format
|
||||
/// The returned bytes are in LP wire format:
|
||||
/// - Outer header (12B): receiver_idx(4) + counter(8) - always cleartext
|
||||
/// - Encrypted payload: proto(1) + reserved(3) + msg_type(4) + content + tag(16)
|
||||
///
|
||||
/// # Usage
|
||||
/// After LP handshake, wrap Sphinx packets before sending to gateway's LP data port:
|
||||
/// ```ignore
|
||||
/// client.perform_handshake().await?;
|
||||
/// let sphinx_bytes = build_sphinx_packet(...);
|
||||
/// let lp_bytes = client.wrap_data(&sphinx_bytes)?;
|
||||
/// socket.send_to(&lp_bytes, gateway_lp_data_address).await?; // UDP:51264
|
||||
/// ```
|
||||
pub fn wrap_data(&mut self, data: &[u8]) -> Result<Vec<u8>> {
|
||||
let state_machine = self.state_machine.as_mut().ok_or_else(|| {
|
||||
LpClientError::Transport("Cannot wrap data: handshake not completed".to_string())
|
||||
})?;
|
||||
|
||||
// Process data through state machine to create LP packet
|
||||
let action = state_machine
|
||||
.process_input(LpInput::SendData(data.to_vec()))
|
||||
.ok_or_else(|| LpClientError::Transport("State machine returned no action".to_string()))?
|
||||
.map_err(|e| LpClientError::Transport(format!("Failed to encrypt data: {}", e)))?;
|
||||
|
||||
let packet = match action {
|
||||
LpAction::SendPacket(packet) => packet,
|
||||
other => {
|
||||
return Err(LpClientError::Transport(format!(
|
||||
"Unexpected action when wrapping data: {:?}",
|
||||
other
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
// Get outer AEAD key for encryption
|
||||
let outer_key = state_machine
|
||||
.session()
|
||||
.ok()
|
||||
.and_then(|s| s.outer_aead_key_for_sending());
|
||||
|
||||
// Serialize the packet with outer AEAD encryption
|
||||
let mut buf = BytesMut::new();
|
||||
serialize_lp_packet(&packet, &mut buf, outer_key.as_ref()).map_err(|e| {
|
||||
LpClientError::Transport(format!("Failed to serialize LP packet: {}", e))
|
||||
})?;
|
||||
|
||||
Ok(buf.to_vec())
|
||||
}
|
||||
|
||||
/// Get the LP session ID (receiver_idx) for this client.
|
||||
///
|
||||
/// This ID is included in the outer header of LP packets and is used by
|
||||
/// the gateway to look up the session for decryption.
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Ok(u32)` - The session ID
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if handshake has not been completed.
|
||||
pub fn session_id(&self) -> Result<u32> {
|
||||
let state_machine = self.state_machine.as_ref().ok_or_else(|| {
|
||||
LpClientError::Transport("Cannot get session ID: handshake not completed".to_string())
|
||||
})?;
|
||||
|
||||
state_machine
|
||||
.session()
|
||||
.map(|s| s.id())
|
||||
.map_err(|e| LpClientError::Transport(format!("Failed to get session: {}", e)))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -30,6 +30,7 @@ use tracing::{debug, info, trace};
|
||||
|
||||
use crate::topology::{GatewayInfo, SpeedtestTopology};
|
||||
use nym_ip_packet_requests::v8::request::IpPacketRequest;
|
||||
use nym_sphinx::forwarding::packet::MixPacket;
|
||||
|
||||
/// Conv ID for KCP - hash of source and destination addresses
|
||||
fn compute_conv_id(local: SocketAddr, remote: SocketAddr) -> u32 {
|
||||
@@ -56,6 +57,24 @@ pub struct SpeedtestClient {
|
||||
kcp_driver: Option<KcpDriver>,
|
||||
/// RNG for packet building
|
||||
rng: ChaCha8Rng,
|
||||
/// LP registration client (kept alive for data framing)
|
||||
lp_client: Option<LpRegistrationClient>,
|
||||
}
|
||||
|
||||
/// Prepared Sphinx packet data ready for sending
|
||||
struct PreparedPackets {
|
||||
/// Message fragments ready for Sphinx wrapping
|
||||
fragments: Vec<nym_sphinx::chunking::fragment::Fragment>,
|
||||
/// Route through mixnet (Mix1, Mix2, Mix3, Gateway)
|
||||
route: Vec<nym_sphinx_types::Node>,
|
||||
/// Sphinx destination (gateway's sphinx key)
|
||||
destination: Destination,
|
||||
/// Zero delays for each hop
|
||||
delays: Vec<Delay>,
|
||||
/// SURB encryption keys for decrypting replies
|
||||
encryption_keys: Vec<SurbEncryptionKey>,
|
||||
/// First hop address for sending
|
||||
first_hop_addr: NymNodeRoutingAddress,
|
||||
}
|
||||
|
||||
impl SpeedtestClient {
|
||||
@@ -73,6 +92,7 @@ impl SpeedtestClient {
|
||||
socket: None,
|
||||
kcp_driver: None,
|
||||
rng,
|
||||
lp_client: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -116,6 +136,79 @@ impl SpeedtestClient {
|
||||
Ok(duration)
|
||||
}
|
||||
|
||||
/// Initialize LP session for data plane communication.
|
||||
///
|
||||
/// Performs LP handshake with the gateway and stores the cryptographic session
|
||||
/// state (encryption keys, counters, session ID) for wrapping UDP data packets.
|
||||
/// The TCP control connection is closed immediately after handshake; only the
|
||||
/// state machine is retained for `wrap_data()` calls.
|
||||
///
|
||||
/// # Data Flow
|
||||
/// After calling this method, use `send_data_via_lp()` to:
|
||||
/// 1. Build Sphinx packets (as usual)
|
||||
/// 2. Wrap them in LP via the stored state machine
|
||||
/// 3. Send to gateway's LP data port (UDP:51264)
|
||||
///
|
||||
/// # Returns
|
||||
/// Handshake duration on success.
|
||||
pub async fn init_lp_session(&mut self) -> Result<Duration> {
|
||||
info!(
|
||||
"Initializing LP session with gateway at {}",
|
||||
self.gateway.lp_address
|
||||
);
|
||||
|
||||
let client_ip = "0.0.0.0".parse().unwrap();
|
||||
|
||||
let mut lp_client = LpRegistrationClient::new_with_default_psk(
|
||||
self.identity_keypair.clone(),
|
||||
self.gateway.identity,
|
||||
self.gateway.lp_address,
|
||||
client_ip,
|
||||
);
|
||||
|
||||
let start = Instant::now();
|
||||
lp_client
|
||||
.perform_handshake()
|
||||
.await
|
||||
.context("LP handshake failed")?;
|
||||
let duration = start.elapsed();
|
||||
|
||||
// Close TCP connection - we only need the state machine for UDP data
|
||||
// (close() only drops stream, preserves state_machine)
|
||||
lp_client.close();
|
||||
|
||||
info!(
|
||||
"LP session established in {:?}, session_id={}",
|
||||
duration,
|
||||
lp_client.session_id().unwrap_or(0)
|
||||
);
|
||||
|
||||
// Store the client (with state machine) for data plane operations
|
||||
self.lp_client = Some(lp_client);
|
||||
|
||||
Ok(duration)
|
||||
}
|
||||
|
||||
/// Check if LP session is established
|
||||
pub fn has_lp_session(&self) -> bool {
|
||||
self.lp_client
|
||||
.as_ref()
|
||||
.map(|c| c.is_handshake_complete())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Close LP session and cleanup resources.
|
||||
///
|
||||
/// This fully destroys the LP session (state machine + TCP stream), unlike
|
||||
/// `LpRegistrationClient::close()` which only drops the TCP stream.
|
||||
/// After calling this, `init_lp_session()` must be called again to re-establish.
|
||||
pub fn close_lp_session(&mut self) {
|
||||
if let Some(mut client) = self.lp_client.take() {
|
||||
client.close();
|
||||
info!("LP session closed");
|
||||
}
|
||||
}
|
||||
|
||||
/// Initialize UDP socket and KCP for data plane
|
||||
pub async fn init_data_channel(&mut self) -> Result<()> {
|
||||
let socket = UdpSocket::bind("0.0.0.0:0")
|
||||
@@ -147,28 +240,33 @@ impl SpeedtestClient {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Send data with SURBs for bidirectional communication
|
||||
/// Prepare Sphinx packet fragments for sending.
|
||||
///
|
||||
/// Returns the SURB encryption keys needed to decrypt replies.
|
||||
/// The `num_surbs` parameter controls how many reply SURBs to attach.
|
||||
pub async fn send_data_with_surbs(
|
||||
/// Common logic for both direct and LP-wrapped sending:
|
||||
/// - Wraps payload in IpPacketRequest + KCP
|
||||
/// - Builds route and destination
|
||||
/// - Creates SURBs if requested
|
||||
/// - Fragments the message
|
||||
///
|
||||
/// Returns prepared data ready for Sphinx wrapping and sending.
|
||||
async fn prepare_sphinx_fragments(
|
||||
&mut self,
|
||||
payload: &[u8],
|
||||
num_surbs: usize,
|
||||
) -> Result<Vec<SurbEncryptionKey>> {
|
||||
) -> Result<PreparedPackets> {
|
||||
if self.socket.is_none() {
|
||||
self.init_data_channel().await?;
|
||||
}
|
||||
|
||||
let driver = self.kcp_driver.as_mut().context("KCP not initialized")?;
|
||||
let socket = self.socket.as_ref().context("socket not initialized")?;
|
||||
|
||||
// Step 1: Wrap payload in IpPacketRequest (DataRequest) and feed to KCP
|
||||
let data_request = IpPacketRequest::new_data_request(Bytes::copy_from_slice(payload));
|
||||
let data_bytes = data_request.to_bytes()
|
||||
let data_bytes = data_request
|
||||
.to_bytes()
|
||||
.context("failed to serialize IpPacketRequest")?;
|
||||
driver.send(&data_bytes);
|
||||
driver.update(10); // Process KCP state machine to produce outgoing packets
|
||||
driver.update(10);
|
||||
|
||||
let outgoing = driver.fetch_outgoing();
|
||||
if outgoing.is_empty() {
|
||||
@@ -208,17 +306,16 @@ impl SpeedtestClient {
|
||||
let surb = ReplySurb::construct(
|
||||
&mut self.rng,
|
||||
&recipient,
|
||||
Duration::from_millis(0), // zero delay for speed testing
|
||||
false, // use_legacy_surb_format
|
||||
Duration::from_millis(0),
|
||||
false,
|
||||
&route_provider,
|
||||
false, // disable_mix_hops
|
||||
false,
|
||||
)
|
||||
.context("failed to construct reply SURB")?;
|
||||
surbs_with_keys.push(surb.with_key_rotation(SphinxKeyRotation::Unknown));
|
||||
}
|
||||
}
|
||||
|
||||
// Extract encryption keys for later decryption
|
||||
let encryption_keys: Vec<SurbEncryptionKey> = surbs_with_keys
|
||||
.iter()
|
||||
.map(|s| s.encryption_key().clone())
|
||||
@@ -228,7 +325,7 @@ impl SpeedtestClient {
|
||||
let nym_message = if num_surbs > 0 {
|
||||
let sender_tag = AnonymousSenderTag::new_random(&mut self.rng);
|
||||
let repliable_message = RepliableMessage::new_data(
|
||||
false, // use_legacy_surb_format
|
||||
false,
|
||||
kcp_buf.to_vec(),
|
||||
sender_tag,
|
||||
surbs_with_keys,
|
||||
@@ -241,7 +338,7 @@ impl SpeedtestClient {
|
||||
let nym_message =
|
||||
nym_message.pad_to_full_packet_lengths(PacketSize::RegularPacket.plaintext_size());
|
||||
|
||||
// Step 6: Fragment and send
|
||||
// Step 6: Fragment
|
||||
let fragments = nym_message
|
||||
.split_into_fragments(&mut self.rng, PacketSize::RegularPacket.plaintext_size());
|
||||
|
||||
@@ -251,41 +348,145 @@ impl SpeedtestClient {
|
||||
fragments.len()
|
||||
);
|
||||
|
||||
let first_hop_addr = NymNodeRoutingAddress::try_from(route[0].address)
|
||||
.context("invalid first hop address")?;
|
||||
|
||||
Ok(PreparedPackets {
|
||||
fragments,
|
||||
route,
|
||||
destination,
|
||||
delays,
|
||||
encryption_keys,
|
||||
first_hop_addr,
|
||||
})
|
||||
}
|
||||
|
||||
/// Send data with SURBs for bidirectional communication (direct to Mix1).
|
||||
///
|
||||
/// Returns the SURB encryption keys needed to decrypt replies.
|
||||
/// The `num_surbs` parameter controls how many reply SURBs to attach.
|
||||
///
|
||||
/// Note: This sends directly to the first mix node. For LP transport,
|
||||
/// use `send_data_via_lp()` instead.
|
||||
pub async fn send_data_with_surbs(
|
||||
&mut self,
|
||||
payload: &[u8],
|
||||
num_surbs: usize,
|
||||
) -> Result<Vec<SurbEncryptionKey>> {
|
||||
let prepared = self.prepare_sphinx_fragments(payload, num_surbs).await?;
|
||||
let socket = self.socket.as_ref().context("socket not initialized")?;
|
||||
|
||||
let mut packet_buf = BytesMut::new();
|
||||
for fragment in fragments {
|
||||
for fragment in prepared.fragments {
|
||||
let nym_packet = NymPacket::sphinx_build(
|
||||
false, // use_legacy_sphinx_format
|
||||
false,
|
||||
PacketSize::RegularPacket.payload_size(),
|
||||
fragment.into_bytes(),
|
||||
&route,
|
||||
&destination,
|
||||
&delays,
|
||||
&prepared.route,
|
||||
&prepared.destination,
|
||||
&prepared.delays,
|
||||
)?;
|
||||
|
||||
let framed = FramedNymPacket::new(
|
||||
nym_packet,
|
||||
PacketType::Mix,
|
||||
SphinxKeyRotation::Unknown,
|
||||
false, // use_legacy_packet_encoding
|
||||
false,
|
||||
);
|
||||
let mut codec = NymCodec;
|
||||
codec.encode(framed, &mut packet_buf)?;
|
||||
}
|
||||
|
||||
// Send to first hop
|
||||
let first_hop_addr: SocketAddr =
|
||||
NymNodeRoutingAddress::try_from(route[0].address)?.into();
|
||||
|
||||
socket.send_to(&packet_buf, first_hop_addr).await?;
|
||||
let first_hop_socket: SocketAddr = prepared.first_hop_addr.into();
|
||||
socket.send_to(&packet_buf, first_hop_socket).await?;
|
||||
info!(
|
||||
"Sent {} bytes (KCP) with {} SURBs ({} packet bytes) to {}",
|
||||
kcp_buf.len(),
|
||||
num_surbs,
|
||||
"Sent {} packet bytes with {} SURBs to {}",
|
||||
packet_buf.len(),
|
||||
first_hop_addr
|
||||
num_surbs,
|
||||
first_hop_socket
|
||||
);
|
||||
|
||||
Ok(encryption_keys)
|
||||
Ok(prepared.encryption_keys)
|
||||
}
|
||||
|
||||
/// Send data via LP data plane (UDP:51264) with SURBs for bidirectional communication.
|
||||
///
|
||||
/// This is the primary method for sending data through the mixnet via the LP transport.
|
||||
/// Requires `init_lp_session()` to be called first to establish the LP cryptographic session.
|
||||
///
|
||||
/// # Data Flow (see gateway/src/node/lp_listener/data_handler.rs)
|
||||
/// ```text
|
||||
/// LP Client → UDP:51264 → LP Data Handler → Mixnet Entry
|
||||
/// LP(Sphinx) decrypt LP forward Sphinx
|
||||
/// ```
|
||||
///
|
||||
/// # Why LP instead of direct to Mix1?
|
||||
/// - Client may be behind NAT/firewall (can't reach Mix1 directly)
|
||||
/// - LP provides authenticated, encrypted session with gateway
|
||||
/// - This is the standard client-to-gateway transport for lewes-protocol
|
||||
///
|
||||
/// Returns SURB encryption keys needed to decrypt replies.
|
||||
pub async fn send_data_via_lp(
|
||||
&mut self,
|
||||
payload: &[u8],
|
||||
num_surbs: usize,
|
||||
) -> Result<Vec<SurbEncryptionKey>> {
|
||||
// Check LP session exists first (before borrowing self mutably)
|
||||
if !self.has_lp_session() {
|
||||
bail!("LP session not initialized - call init_lp_session() first");
|
||||
}
|
||||
|
||||
let prepared = self.prepare_sphinx_fragments(payload, num_surbs).await?;
|
||||
|
||||
// Now get mutable references after prepare_sphinx_fragments is done
|
||||
let lp_client = self.lp_client.as_mut().unwrap(); // safe: checked above
|
||||
let socket = self.socket.as_ref().context("socket not initialized")?;
|
||||
let lp_data_address = self.gateway.lp_data_address;
|
||||
|
||||
let mut total_sent = 0usize;
|
||||
let fragment_count = prepared.fragments.len();
|
||||
|
||||
for fragment in prepared.fragments {
|
||||
let nym_packet = NymPacket::sphinx_build(
|
||||
false,
|
||||
PacketSize::RegularPacket.payload_size(),
|
||||
fragment.into_bytes(),
|
||||
&prepared.route,
|
||||
&prepared.destination,
|
||||
&prepared.delays,
|
||||
)?;
|
||||
|
||||
// Wrap in MixPacket v2: packet_type || key_rotation || next_hop || sphinx_data
|
||||
let mix_packet = MixPacket::new(
|
||||
prepared.first_hop_addr,
|
||||
nym_packet,
|
||||
PacketType::Mix,
|
||||
SphinxKeyRotation::Unknown,
|
||||
);
|
||||
|
||||
let mix_bytes = mix_packet
|
||||
.into_v2_bytes()
|
||||
.context("failed to serialize MixPacket")?;
|
||||
|
||||
// Wrap in LP for UDP data plane
|
||||
let lp_packet = lp_client
|
||||
.wrap_data(&mix_bytes)
|
||||
.context("failed to wrap in LP")?;
|
||||
|
||||
// Send to gateway's LP data port (51264) with timeout
|
||||
tokio::time::timeout(Duration::from_secs(5), socket.send_to(&lp_packet, lp_data_address))
|
||||
.await
|
||||
.context("UDP send timed out")?
|
||||
.context("UDP send failed")?;
|
||||
total_sent += lp_packet.len();
|
||||
}
|
||||
|
||||
info!(
|
||||
"Sent {} bytes via LP ({} fragments, {} SURBs) to {}",
|
||||
total_sent, fragment_count, num_surbs, lp_data_address
|
||||
);
|
||||
|
||||
Ok(prepared.encryption_keys)
|
||||
}
|
||||
|
||||
/// Receive UDP data with timeout
|
||||
|
||||
Reference in New Issue
Block a user