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:
durch
2025-12-23 22:04:34 +01:00
committed by Jędrzej Stuczyński
parent 1cfa79abc9
commit b924729992
2 changed files with 315 additions and 30 deletions
@@ -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)]
+231 -30
View File
@@ -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