From f9c40e41f20c2c8de58aa82d0e3bbdd87cb9b586 Mon Sep 17 00:00:00 2001 From: durch Date: Sat, 20 Dec 2025 20:37:07 +0100 Subject: [PATCH] Refactor LP to stream-oriented TCP processing Gateway (handler.rs): - Add bound_receiver_idx field for session-affine connections - Convert handle() from single-packet to loop with EOF detection - Add validate_or_set_binding() for receiver_idx validation - Set binding in handle_client_hello after collision check - Centralize emit_lifecycle_metrics in main loop only - Add is_connection_closed() helper for graceful EOF Client (client.rs): - Add stream field for persistent TCP connection - Add ensure_connected(), send_packet(), receive_packet(), close() methods - Modify perform_handshake_inner() to use persistent stream - Modify register_with_credential() to use persistent stream - Modify send_forward_packet() to use persistent stream - Keep connect_send_receive() for reference (marked dead_code) This reduces handshake overhead from ~5 TCP connections to 1. Drive-by: Fix log::info! -> info! in wireguard peer_controller.rs --- gateway/src/node/lp_listener/handler.rs | 147 +++++-- .../src/lp_client/client.rs | 397 ++++++++++++------ 2 files changed, 378 insertions(+), 166 deletions(-) diff --git a/gateway/src/node/lp_listener/handler.rs b/gateway/src/node/lp_listener/handler.rs index 5b3eb64e3b..3db1bfbea6 100644 --- a/gateway/src/node/lp_listener/handler.rs +++ b/gateway/src/node/lp_listener/handler.rs @@ -71,6 +71,10 @@ pub struct LpConnectionHandler { remote_addr: SocketAddr, state: LpHandlerState, stats: ConnectionStats, + /// Bound receiver_idx for this connection (set after first packet). + /// All subsequent packets on this connection must use this receiver_idx. + /// Set from ClientHello's proposed receiver_index, or from header for non-bootstrap packets. + bound_receiver_idx: Option, } impl LpConnectionHandler { @@ -80,9 +84,16 @@ impl LpConnectionHandler { remote_addr, state, stats: ConnectionStats::new(), + bound_receiver_idx: None, } } + /// AIDEV-NOTE: Stream-oriented packet loop + /// This handler processes multiple packets on a single TCP connection. + /// Connection lifecycle: handshake + registration, then client closes. + /// First packet binds the connection to a receiver_idx (session-affine). + /// Binding is set by handle_client_hello() from payload's receiver_index, + /// or by validate_or_set_binding() for non-bootstrap first packets. pub async fn handle(mut self) -> Result<(), GatewayError> { debug!("Handling LP connection from {}", self.remote_addr); @@ -90,25 +101,108 @@ impl LpConnectionHandler { inc!("lp_connections_total"); // ============================================================ - // SINGLE-PACKET PROCESSING: Process ONE packet then close - // State persists in LpHandlerState maps between connections + // STREAM-ORIENTED PROCESSING: Loop until connection closes + // State persists in LpHandlerState maps across packets // ============================================================ - // Step 1: Receive raw packet bytes and parse header only (for routing) - let (raw_bytes, header) = match self.receive_raw_packet().await { - Ok(result) => result, - Err(e) => { - inc!("lp_errors_receive_packet"); + loop { + // Step 1: Receive raw packet bytes and parse header only (for routing) + let (raw_bytes, header) = match self.receive_raw_packet().await { + Ok(result) => result, + Err(e) if Self::is_connection_closed(&e) => { + // Graceful EOF - client closed connection + trace!("Connection closed by {} (EOF)", self.remote_addr); + break; + } + Err(e) => { + inc!("lp_errors_receive_packet"); + self.emit_lifecycle_metrics(false); + return Err(e); + } + }; + + let receiver_idx = header.receiver_idx; + + // Step 2: Validate or set binding (session-affine connection) + // Note: ClientHello (receiver_idx=0) defers binding to handle_client_hello() + if let Err(e) = self.validate_or_set_binding(receiver_idx) { self.emit_lifecycle_metrics(false); return Err(e); } - }; - let receiver_idx = header.receiver_idx; + // Step 3: Process the packet + if let Err(e) = self.process_packet(raw_bytes, receiver_idx).await { + self.emit_lifecycle_metrics(false); + return Err(e); + } + } - // Step 2: Get outer_aead_key based on receiver_idx + self.emit_lifecycle_metrics(true); + Ok(()) + } + + /// Check if an error indicates the connection was closed (EOF). + /// AIDEV-NOTE: Uses string matching on error messages. Tokio's read_exact + /// returns UnexpectedEof which gets formatted into the error message. + fn is_connection_closed(e: &GatewayError) -> bool { + match e { + GatewayError::LpConnectionError(msg) => { + msg.contains("unexpected end of file") + || msg.contains("connection reset") + || msg.contains("broken pipe") + } + _ => false, + } + } + + /// Validate that the receiver_idx matches the bound session, or set binding if first packet. + /// + /// Binding rules: + /// - ClientHello (receiver_idx=0): binding deferred to handle_client_hello() which + /// extracts receiver_index from payload + /// - First non-bootstrap packet: sets binding from header's receiver_idx + /// - Subsequent packets: must match bound receiver_idx + fn validate_or_set_binding(&mut self, receiver_idx: u32) -> Result<(), GatewayError> { + match self.bound_receiver_idx { + None => { + // First packet - don't bind if bootstrap (handle_client_hello sets binding) + if receiver_idx != nym_lp::BOOTSTRAP_RECEIVER_IDX { + self.bound_receiver_idx = Some(receiver_idx); + trace!( + "Bound connection from {} to receiver_idx={}", + self.remote_addr, + receiver_idx + ); + } + Ok(()) + } + Some(bound) => { + if receiver_idx == bound { + Ok(()) + } else { + warn!( + "Receiver_idx mismatch from {}: expected {}, got {}", + self.remote_addr, bound, receiver_idx + ); + inc!("lp_errors_receiver_idx_mismatch"); + Err(GatewayError::LpProtocolError(format!( + "receiver_idx mismatch: connection bound to {}, packet has {}", + bound, receiver_idx + ))) + } + } + } + } + + /// Process a single packet: lookup session, parse, route to handler. + /// Individual handlers do NOT emit lifecycle metrics - the main loop handles that. + async fn process_packet( + &mut self, + raw_bytes: Vec, + receiver_idx: u32, + ) -> Result<(), GatewayError> { + // Get outer_aead_key based on receiver_idx // Header is always cleartext for routing. Payload is encrypted after PSK. - // We lookup the session to get the key, then parse the full packet. let outer_key: Option = if receiver_idx == nym_lp::BOOTSTRAP_RECEIVER_IDX { // ClientHello - no encryption (PSK not yet derived) None @@ -122,7 +216,6 @@ impl LpConnectionHandler { .and_then(|session| session.outer_aead_key()) } else if let Some(session_entry) = self.state.session_states.get(&receiver_idx) { // Established session - should always have PSK - // session_states now stores LpStateMachine (not LpSession) for subsession support session_entry .value() .state @@ -134,11 +227,10 @@ impl LpConnectionHandler { None }; - // Step 3: Parse full packet with outer AEAD key + // Parse full packet with outer AEAD key let packet = nym_lp::codec::parse_lp_packet(&raw_bytes, outer_key.as_ref()).map_err(|e| { inc!("lp_errors_parse_packet"); - self.emit_lifecycle_metrics(false); GatewayError::LpProtocolError(format!("Failed to parse LP packet: {}", e)) })?; @@ -150,7 +242,7 @@ impl LpConnectionHandler { outer_key.is_some() ); - // Step 4: Route packet based on receiver_idx + // Route packet based on receiver_idx if receiver_idx == nym_lp::BOOTSTRAP_RECEIVER_IDX { // ClientHello - first packet in handshake self.handle_client_hello(packet).await @@ -169,7 +261,6 @@ impl LpConnectionHandler { receiver_idx, self.remote_addr ); inc!("lp_errors_unknown_session"); - self.emit_lifecycle_metrics(false); Err(GatewayError::LpProtocolError(format!( "Unknown session ID: {}", receiver_idx @@ -207,7 +298,6 @@ impl LpConnectionHandler { } other => { inc!("lp_client_hello_failed"); - self.emit_lifecycle_metrics(false); return Err(GatewayError::LpProtocolError(format!( "Expected ClientHello, got {}", other @@ -233,14 +323,23 @@ impl LpConnectionHandler { // Send Collision response to tell client to retry with new receiver_index // No outer key - this is before PSK derivation + // Note: Do NOT set binding on collision - client may retry with new receiver_index let collision_packet = LpPacket::new(LpHeader::new(receiver_index, 0), LpMessage::Collision); self.send_lp_packet(&collision_packet, None).await?; - self.emit_lifecycle_metrics(true); return Ok(()); } + // Collision check passed - bind this connection to the receiver_index + // All subsequent packets on this connection must use this receiver_index + self.bound_receiver_idx = Some(receiver_index); + trace!( + "Bound connection from {} to receiver_idx={} (via ClientHello)", + self.remote_addr, + receiver_index + ); + // Create state machine for this handshake using client-proposed receiver_index let mut state_machine = LpStateMachine::new( receiver_index, @@ -284,12 +383,11 @@ impl LpConnectionHandler { self.remote_addr, receiver_index ); - // Send Ack to confirm ClientHello received (packet-per-connection model) + // Send Ack to confirm ClientHello received // No outer key - this is before PSK derivation let ack_packet = LpPacket::new(LpHeader::new(receiver_index, 0), LpMessage::Ack); self.send_lp_packet(&ack_packet, None).await?; - self.emit_lifecycle_metrics(true); Ok(()) } @@ -400,7 +498,6 @@ impl LpConnectionHandler { ); } - self.emit_lifecycle_metrics(true); Ok(()) } @@ -473,7 +570,6 @@ impl LpConnectionHandler { inc!("lp_subsession_kk2_sent"); self.send_lp_packet(&response_packet, outer_key.as_ref()) .await?; - self.emit_lifecycle_metrics(true); Ok(()) } LpAction::DeliverData(data) => { @@ -501,7 +597,6 @@ impl LpConnectionHandler { "Unexpected action in transport from {}: {:?}", self.remote_addr, other ); - self.emit_lifecycle_metrics(false); Err(GatewayError::LpProtocolError(format!( "Unexpected action: {:?}", other @@ -544,7 +639,6 @@ impl LpConnectionHandler { self.remote_addr, receiver_idx ); inc!("lp_errors_unknown_payload_type"); - self.emit_lifecycle_metrics(false); Err(GatewayError::LpProtocolError( "Unknown transport payload type (not registration or forwarding)".to_string(), )) @@ -600,7 +694,6 @@ impl LpConnectionHandler { new_receiver_index, self.remote_addr ); inc!("lp_subsession_receiver_index_collision"); - self.emit_lifecycle_metrics(false); return Err(GatewayError::LpProtocolError( "Subsession receiver index collision - client should retry".to_string(), )); @@ -616,7 +709,6 @@ impl LpConnectionHandler { // It will be cleaned up by TTL-based cleanup task inc!("lp_subsession_complete"); - self.emit_lifecycle_metrics(true); Ok(()) } @@ -676,7 +768,6 @@ impl LpConnectionHandler { ); } - self.emit_lifecycle_metrics(true); Ok(()) } @@ -731,7 +822,6 @@ impl LpConnectionHandler { self.remote_addr, receiver_idx ); - self.emit_lifecycle_metrics(true); Ok(()) } @@ -1112,6 +1202,7 @@ mod tests { use nym_lp::message::{ClientHelloData, EncryptedDataPayload, HandshakeData, LpMessage}; use nym_lp::packet::{LpHeader, LpPacket}; use std::sync::Arc; + use std::time::{SystemTime, UNIX_EPOCH}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; // ==================== Test Helpers ==================== diff --git a/nym-registration-client/src/lp_client/client.rs b/nym-registration-client/src/lp_client/client.rs index 074e3cfef1..4f5c1e502d 100644 --- a/nym-registration-client/src/lp_client/client.rs +++ b/nym-registration-client/src/lp_client/client.rs @@ -23,16 +23,16 @@ use tokio::net::TcpStream; /// LP (Lewes Protocol) registration client for direct gateway connections. /// -/// This client uses a packet-per-connection model where each LP packet -/// exchange opens a new TCP connection, sends one packet, receives one -/// response, then closes. Session state is maintained in the state machine -/// across connections. +/// This client uses a persistent TCP connection model where a single TCP +/// connection is used for the entire handshake and registration flow. +/// The connection is opened on first use and closed after registration. /// /// # Example Flow /// ```ignore /// let mut client = LpRegistrationClient::new(...); -/// client.perform_handshake().await?; // Noise handshake (multiple connections) -/// let gateway_data = client.register(...).await?; // Registration (single connection) +/// client.perform_handshake().await?; // Noise handshake (single connection) +/// let gateway_data = client.register(...).await?; // Registration (same connection) +/// // Connection automatically closes after registration /// ``` pub struct LpRegistrationClient { /// Client's Ed25519 identity keypair (used for PSQ authentication and X25519 derivation). @@ -53,6 +53,10 @@ pub struct LpRegistrationClient { /// Configuration for timeouts and TCP parameters. config: LpConfig, + + /// Persistent TCP stream for the connection. + /// Opened on first use, closed after registration. + stream: Option, } impl LpRegistrationClient { @@ -82,6 +86,7 @@ impl LpRegistrationClient { state_machine: None, client_ip, config, + stream: None, } } @@ -130,11 +135,124 @@ impl LpRegistrationClient { self.client_ip } + // ------------------------------------------------------------------------- + // Persistent connection management + // ------------------------------------------------------------------------- + + /// Ensures a TCP connection is established. + /// + /// Opens a new connection to the gateway if one doesn't exist. + /// If a connection already exists, returns immediately. + /// + /// # Errors + /// Returns an error if connection fails or times out. + async fn ensure_connected(&mut self) -> Result<()> { + if self.stream.is_some() { + return Ok(()); + } + + tracing::debug!( + "Opening persistent connection to {}", + self.gateway_lp_address + ); + + let stream = tokio::time::timeout( + self.config.connect_timeout, + TcpStream::connect(self.gateway_lp_address), + ) + .await + .map_err(|_| LpClientError::TcpConnection { + address: self.gateway_lp_address.to_string(), + source: std::io::Error::new( + std::io::ErrorKind::TimedOut, + format!("Connection timeout after {:?}", self.config.connect_timeout), + ), + })? + .map_err(|source| LpClientError::TcpConnection { + address: self.gateway_lp_address.to_string(), + source, + })?; + + // Set TCP_NODELAY for low latency + stream + .set_nodelay(self.config.tcp_nodelay) + .map_err(|source| LpClientError::TcpConnection { + address: self.gateway_lp_address.to_string(), + source, + })?; + + self.stream = Some(stream); + tracing::debug!( + "Persistent connection established to {}", + self.gateway_lp_address + ); + Ok(()) + } + + /// Sends an LP packet on the persistent stream. + /// + /// # Arguments + /// * `packet` - The LP packet to send + /// * `outer_key` - Optional outer AEAD key for encryption + /// + /// # Errors + /// Returns an error if not connected or if send fails. + async fn send_packet( + &mut self, + packet: &LpPacket, + outer_key: Option<&OuterAeadKey>, + ) -> Result<()> { + let stream = self.stream.as_mut().ok_or_else(|| { + LpClientError::Transport("Cannot send: not connected".to_string()) + })?; + + Self::send_packet_with_key(stream, packet, outer_key).await + } + + /// Receives an LP packet from the persistent stream. + /// + /// # Arguments + /// * `outer_key` - Optional outer AEAD key for decryption + /// + /// # Errors + /// Returns an error if not connected or if receive fails. + async fn receive_packet(&mut self, outer_key: Option<&OuterAeadKey>) -> Result { + let stream = self.stream.as_mut().ok_or_else(|| { + LpClientError::Transport("Cannot receive: not connected".to_string()) + })?; + + Self::receive_packet_with_key(stream, outer_key).await + } + + /// Closes the persistent connection. + /// + /// This drops the TCP stream, signaling EOF to the gateway. + /// Safe to call even if not connected. + /// + /// # Connection Lifecycle + /// The connection stays open after handshake and registration to support + /// follow-up operations like `send_forward_packet()`. Callers should: + /// - For direct registration: call `close()` after `register()` returns + /// - For nested sessions: call `close()` after all forwarding is complete + /// + /// The connection will also close automatically when the client is dropped. + pub fn close(&mut self) { + if self.stream.take().is_some() { + tracing::debug!( + "Closed persistent connection to {}", + self.gateway_lp_address + ); + } + } + + // ------------------------------------------------------------------------- + // Handshake + // ------------------------------------------------------------------------- + /// Performs the LP Noise protocol handshake with the gateway. /// /// This establishes a secure encrypted session using the Noise protocol. - /// Uses packet-per-connection model: each handshake message opens a new - /// TCP connection. + /// Uses a persistent TCP connection for all handshake messages. /// /// # Errors /// Returns an error if: @@ -145,10 +263,13 @@ impl LpRegistrationClient { /// /// # Implementation /// This implements the Noise protocol handshake as the initiator: - /// 1. Sends ClientHello, receives Ack (connection 1) - /// 2. Creates LP state machine with client as initiator - /// 3. Exchanges handshake messages (each on new connection) - /// 4. Stores the established session in the state machine + /// 1. Opens persistent TCP connection (if not already connected) + /// 2. Sends ClientHello, receives Ack + /// 3. Creates LP state machine with client as initiator + /// 4. Exchanges handshake messages on the same connection + /// 5. Stores the established session in the state machine + /// + /// The connection remains open after handshake for registration/forwarding. pub async fn perform_handshake(&mut self) -> Result<()> { // Apply handshake timeout (nym-102) tokio::time::timeout( @@ -166,10 +287,14 @@ impl LpRegistrationClient { /// Internal handshake implementation without timeout. /// - /// Uses packet-per-connection model: each LP packet exchange opens a new - /// TCP connection, sends one packet, receives one response, then closes. + /// Uses a persistent TCP connection: all handshake packets are sent and + /// received on the same connection. The connection remains open for + /// registration/forwarding after handshake completes. async fn perform_handshake_inner(&mut self) -> Result<()> { - tracing::debug!("Starting LP handshake as initiator (packet-per-connection)"); + tracing::debug!("Starting LP handshake as initiator (persistent connection)"); + + // Ensure we have a TCP connection + self.ensure_connected().await?; // Step 1: Derive X25519 keys from Ed25519 for Noise protocol (internal to ClientHello) // The Ed25519 keys are used for PSQ authentication and also converted to X25519 @@ -201,7 +326,7 @@ impl LpRegistrationClient { receiver_index ); - // Step 3: Send ClientHello and receive Ack (packet-per-connection) + // Step 3: Send ClientHello and receive Ack (persistent connection) let client_hello_header = nym_lp::packet::LpHeader::new( nym_lp::BOOTSTRAP_RECEIVER_IDX, // session_id not yet established 0, // counter starts at 0 @@ -211,14 +336,10 @@ impl LpRegistrationClient { nym_lp::LpMessage::ClientHello(client_hello_data), ); - let ack_response = Self::connect_send_receive( - self.gateway_lp_address, - &client_hello_packet, - None, // No outer key for ClientHello (before PSK) - None, // No outer key for Ack response (before PSK) - &self.config, - ) - .await?; + // Send ClientHello (no outer key - before PSK) + self.send_packet(&client_hello_packet, None).await?; + // Receive Ack (no outer key - before PSK) + let ack_response = self.receive_packet(None).await?; // Verify we received Ack match ack_response.message() { @@ -262,7 +383,7 @@ impl LpRegistrationClient { } } - // Step 6: Handshake loop - each packet on new connection + // Step 6: Handshake loop - all packets on persistent connection loop { // Send pending packet if we have one if let Some(packet) = pending_packet.take() { @@ -283,14 +404,8 @@ impl LpRegistrationClient { send_key.is_some(), recv_key.is_some() ); - let response = Self::connect_send_receive( - self.gateway_lp_address, - &packet, - send_key.as_ref(), - recv_key.as_ref(), - &self.config, - ) - .await?; + self.send_packet(&packet, send_key.as_ref()).await?; + let response = self.receive_packet(recv_key.as_ref()).await?; tracing::trace!("Received handshake response"); // Process the received packet @@ -314,14 +429,9 @@ impl LpRegistrationClient { .ok() .and_then(|s| s.outer_aead_key()); tracing::trace!("Sending final handshake packet"); - let ack_response = Self::connect_send_receive( - self.gateway_lp_address, - &final_packet, - send_key.as_ref(), - recv_key.as_ref(), - &self.config, - ) - .await?; + self.send_packet(&final_packet, send_key.as_ref()).await?; + let ack_response = + self.receive_packet(recv_key.as_ref()).await?; // Validate Ack response match ack_response.message() { @@ -403,6 +513,11 @@ impl LpRegistrationClient { /// /// Use `outer_aead_key_for_sending()` for `send_key` (gates on PSQ completion) /// and `outer_aead_key()` for `recv_key` (available as soon as PSK derived). + /// + /// # Note + /// This method is kept for reference but is no longer used. The persistent + /// connection model uses `send_packet()` and `receive_packet()` instead. + #[allow(dead_code)] async fn connect_send_receive( address: SocketAddr, packet: &LpPacket, @@ -587,8 +702,7 @@ impl LpRegistrationClient { /// Sends registration request with a pre-generated credential. /// /// This is useful for testing with mock ecash credentials. - /// Uses packet-per-connection model: opens connection, sends request, - /// receives response, closes connection. + /// Uses the persistent TCP connection established during handshake. /// /// # Arguments /// * `wg_keypair` - Client's WireGuard x25519 keypair @@ -597,18 +711,20 @@ impl LpRegistrationClient { /// /// # Returns /// * `Ok(GatewayData)` - Gateway configuration data on successful registration + /// + /// # Connection Lifecycle + /// The connection stays open after registration to support `send_forward_packet()`. + /// Callers should call `close()` when done with all operations. + /// + /// # Panics / Errors + /// Returns error if handshake not completed or if connection was closed. pub async fn register_with_credential( &mut self, wg_keypair: &x25519::KeyPair, credential: CredentialSpendingData, ticket_type: TicketType, ) -> Result { - // Ensure handshake is complete (state machine exists) - let state_machine = self.state_machine.as_mut().ok_or_else(|| { - LpClientError::Transport("Cannot register: handshake not completed".to_string()) - })?; - - tracing::debug!("Sending registration request (packet-per-connection)"); + tracing::debug!("Sending registration request (persistent connection)"); // 1. Build registration request let wg_public_key = PeerPublicKey::new(wg_keypair.public_key().to_bytes().into()); @@ -627,50 +743,52 @@ impl LpRegistrationClient { request_bytes.len() ); - // 3. Encrypt and prepare packet via state machine - let action = state_machine - .process_input(LpInput::SendData(request_bytes)) - .ok_or_else(|| { - LpClientError::Transport("State machine returned no action".to_string()) - })? - .map_err(|e| { - LpClientError::SendRegistrationRequest(format!( - "Failed to encrypt registration request: {}", - e - )) + // 3. Encrypt and prepare packet via state machine (scoped borrow) + let (request_packet, send_key, recv_key) = { + let state_machine = self.state_machine.as_mut().ok_or_else(|| { + LpClientError::Transport("Cannot register: handshake not completed".to_string()) })?; - let request_packet = match action { - LpAction::SendPacket(packet) => packet, - other => { - return Err(LpClientError::Transport(format!( - "Unexpected action when sending registration data: {:?}", - other - ))); - } - }; + let action = state_machine + .process_input(LpInput::SendData(request_bytes)) + .ok_or_else(|| { + LpClientError::Transport("State machine returned no action".to_string()) + })? + .map_err(|e| { + LpClientError::SendRegistrationRequest(format!( + "Failed to encrypt registration request: {}", + e + )) + })?; - // 4. Get outer keys from session - let send_key = state_machine - .session() - .ok() - .and_then(|s| s.outer_aead_key_for_sending()); - let recv_key = state_machine - .session() - .ok() - .and_then(|s| s.outer_aead_key()); + let request_packet = match action { + LpAction::SendPacket(packet) => packet, + other => { + return Err(LpClientError::Transport(format!( + "Unexpected action when sending registration data: {:?}", + other + ))); + } + }; - // 5. Send request and receive response on fresh connection with timeout - let response_packet = tokio::time::timeout( - self.config.registration_timeout, - Self::connect_send_receive( - self.gateway_lp_address, - &request_packet, - send_key.as_ref(), - recv_key.as_ref(), - &self.config, - ), - ) + // Get outer keys from session + let send_key = state_machine + .session() + .ok() + .and_then(|s| s.outer_aead_key_for_sending()); + let recv_key = state_machine + .session() + .ok() + .and_then(|s| s.outer_aead_key()); + + (request_packet, send_key, recv_key) + }; // state_machine borrow ends here + + // 4. Send request and receive response on persistent connection with timeout + let response_packet = tokio::time::timeout(self.config.registration_timeout, async { + self.send_packet(&request_packet, send_key.as_ref()).await?; + self.receive_packet(recv_key.as_ref()).await + }) .await .map_err(|_| { LpClientError::ReceiveRegistrationResponse(format!( @@ -681,7 +799,10 @@ impl LpRegistrationClient { tracing::trace!("Received registration response packet"); - // 6. Decrypt via state machine + // 5. Decrypt via state machine (re-borrow) + let state_machine = self.state_machine.as_mut().ok_or_else(|| { + LpClientError::Transport("State machine disappeared unexpectedly".to_string()) + })?; let action = state_machine .process_input(LpInput::ReceivePacket(response_packet)) .ok_or_else(|| { @@ -748,8 +869,8 @@ impl LpRegistrationClient { /// address, and the inner LP packet bytes, encrypts it through the outer session /// (client-entry), and receives the response from the exit gateway via the entry gateway. /// - /// Uses packet-per-connection model: opens connection, sends forward packet, - /// receives response, closes connection. + /// Uses the persistent TCP connection established during handshake. + /// Multiple forward packets can be sent on the same connection. /// /// # Arguments /// * `target_identity` - Target gateway's Ed25519 identity (32 bytes) @@ -785,15 +906,8 @@ impl LpRegistrationClient { target_address: String, inner_packet_bytes: Vec, ) -> Result> { - // Ensure handshake is complete (state machine exists) - let state_machine = self.state_machine.as_mut().ok_or_else(|| { - LpClientError::Transport( - "Cannot send forward packet: handshake not completed".to_string(), - ) - })?; - tracing::debug!( - "Sending ForwardPacket to {} ({} inner bytes, packet-per-connection)", + "Sending ForwardPacket to {} ({} inner bytes, persistent connection)", target_address, inner_packet_bytes.len() ); @@ -815,47 +929,51 @@ impl LpRegistrationClient { forward_data_bytes.len() ); - // 3. Encrypt and prepare packet via state machine - let action = state_machine - .process_input(LpInput::SendData(forward_data_bytes)) - .ok_or_else(|| { - LpClientError::Transport("State machine returned no action".to_string()) - })? - .map_err(|e| { - LpClientError::Transport(format!("Failed to encrypt ForwardPacket: {}", e)) + // 3. Encrypt and prepare packet via state machine (scoped borrow) + let (forward_packet, send_key, recv_key) = { + let state_machine = self.state_machine.as_mut().ok_or_else(|| { + LpClientError::Transport( + "Cannot send forward packet: handshake not completed".to_string(), + ) })?; - let forward_packet = match action { - LpAction::SendPacket(packet) => packet, - other => { - return Err(LpClientError::Transport(format!( - "Unexpected action when sending ForwardPacket: {:?}", - other - ))); - } - }; + let action = state_machine + .process_input(LpInput::SendData(forward_data_bytes)) + .ok_or_else(|| { + LpClientError::Transport("State machine returned no action".to_string()) + })? + .map_err(|e| { + LpClientError::Transport(format!("Failed to encrypt ForwardPacket: {}", e)) + })?; - // 4. Get outer keys from session - let send_key = state_machine - .session() - .ok() - .and_then(|s| s.outer_aead_key_for_sending()); - let recv_key = state_machine - .session() - .ok() - .and_then(|s| s.outer_aead_key()); + let forward_packet = match action { + LpAction::SendPacket(packet) => packet, + other => { + return Err(LpClientError::Transport(format!( + "Unexpected action when sending ForwardPacket: {:?}", + other + ))); + } + }; - // 5. Send and receive on fresh connection with timeout - let response_packet = tokio::time::timeout( - self.config.forward_timeout, - Self::connect_send_receive( - self.gateway_lp_address, - &forward_packet, - send_key.as_ref(), - recv_key.as_ref(), - &self.config, - ), - ) + // Get outer keys from session + let send_key = state_machine + .session() + .ok() + .and_then(|s| s.outer_aead_key_for_sending()); + let recv_key = state_machine + .session() + .ok() + .and_then(|s| s.outer_aead_key()); + + (forward_packet, send_key, recv_key) + }; // state_machine borrow ends here + + // 4. Send and receive on persistent connection with timeout + let response_packet = tokio::time::timeout(self.config.forward_timeout, async { + self.send_packet(&forward_packet, send_key.as_ref()).await?; + self.receive_packet(recv_key.as_ref()).await + }) .await .map_err(|_| { LpClientError::Transport(format!( @@ -865,7 +983,10 @@ impl LpRegistrationClient { })??; tracing::trace!("Received response packet from entry gateway"); - // 6. Decrypt via state machine + // 5. Decrypt via state machine (re-borrow) + let state_machine = self.state_machine.as_mut().ok_or_else(|| { + LpClientError::Transport("State machine disappeared unexpectedly".to_string()) + })?; let action = state_machine .process_input(LpInput::ReceivePacket(response_packet)) .ok_or_else(|| {