mirror of
https://github.com/jmcorgan/fips.git
synced 2026-07-30 19:46:15 +00:00
Implement non-blocking transport connect for connection-oriented transports
TCP (and future Tor) transports previously established connections synchronously inside send(), blocking the node's RX event loop during TCP handshake. This is particularly problematic for Tor where SOCKS5 circuit establishment can take 30-120 seconds. Add a non-blocking connect path: - ConnectionState enum in transport layer (None/Connecting/Connected/Failed) - connect_async() on TcpTransport spawns background TCP connect task - connection_state_sync() polls task completion, promotes to pool - TransportHandle gains connect() and connection_state() dispatch methods - Node tracks PendingConnect entries for connection-oriented transports - initiate_connection() defers handshake for connection-oriented transports - start_handshake() extracted as separate method for deferred invocation - poll_pending_connects() in tick handler polls and completes handshakes - Failed connects trigger retry via schedule_retry() Connectionless transports (UDP, Ethernet) are unchanged — connect() is a no-op and connection_state() always returns Connected. The existing connect-on-send path in send_async() is preserved as fallback for reconnection after connection drops. 811 tests pass (6 new), clippy clean.
This commit is contained in:
Johnathan Corgan
@@ -107,6 +107,7 @@ impl Node {
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as u64)
|
||||
.unwrap_or(0);
|
||||
self.poll_pending_connects().await;
|
||||
self.resend_pending_handshakes(now_ms).await;
|
||||
self.resend_pending_rekeys(now_ms).await;
|
||||
self.resend_pending_session_handshakes(now_ms).await;
|
||||
|
||||
+154
-10
@@ -3,7 +3,7 @@
|
||||
use super::{Node, NodeError, NodeState};
|
||||
use crate::peer::PeerConnection;
|
||||
use crate::protocol::{Disconnect, DisconnectReason};
|
||||
use crate::transport::{packet_channel, Link, LinkDirection, TransportAddr, TransportId};
|
||||
use crate::transport::{packet_channel, Link, LinkDirection, LinkId, TransportAddr, TransportId};
|
||||
use crate::upper::tun::{run_tun_reader, shutdown_tun_interface, TunDevice, TunState};
|
||||
use crate::node::wire::build_msg1;
|
||||
use crate::{NodeAddr, PeerIdentity};
|
||||
@@ -143,9 +143,14 @@ impl Node {
|
||||
|
||||
/// Initiate a connection to a peer on a specific transport and address.
|
||||
///
|
||||
/// Allocates a link, starts the Noise IK handshake, sends msg1, and
|
||||
/// registers the connection for msg2 dispatch. Used by both static peer
|
||||
/// config and transport discovery auto-connect paths.
|
||||
/// For connectionless transports (UDP, Ethernet): allocates a link, starts
|
||||
/// the Noise IK handshake, sends msg1, and registers the connection for
|
||||
/// msg2 dispatch.
|
||||
///
|
||||
/// For connection-oriented transports (TCP, Tor): allocates a link and
|
||||
/// starts a non-blocking transport connect. The handshake is deferred
|
||||
/// until the transport connection is established — the tick handler
|
||||
/// polls `connection_state()` and initiates the handshake when ready.
|
||||
pub(super) async fn initiate_connection(
|
||||
&mut self,
|
||||
transport_id: TransportId,
|
||||
@@ -154,15 +159,14 @@ impl Node {
|
||||
) -> Result<(), NodeError> {
|
||||
let peer_node_addr = *peer_identity.node_addr();
|
||||
|
||||
let is_connection_oriented = self.transports.get(&transport_id)
|
||||
.map(|t| t.transport_type().connection_oriented)
|
||||
.unwrap_or(false);
|
||||
|
||||
// Allocate link ID and create link
|
||||
let link_id = self.allocate_link_id();
|
||||
|
||||
// Use Link::new() for connection-oriented transports (Connecting state)
|
||||
// and Link::connectionless() for connectionless transports (Connected state)
|
||||
let link = if self.transports.get(&transport_id)
|
||||
.map(|t| t.transport_type().connection_oriented)
|
||||
.unwrap_or(false)
|
||||
{
|
||||
let link = if is_connection_oriented {
|
||||
Link::new(
|
||||
link_id,
|
||||
transport_id,
|
||||
@@ -186,6 +190,53 @@ impl Node {
|
||||
self.addr_to_link
|
||||
.insert((transport_id, remote_addr.clone()), link_id);
|
||||
|
||||
if is_connection_oriented {
|
||||
// Connection-oriented: start non-blocking connect, defer handshake
|
||||
if let Some(transport) = self.transports.get(&transport_id) {
|
||||
match transport.connect(&remote_addr).await {
|
||||
Ok(()) => {
|
||||
debug!(
|
||||
peer = %self.peer_display_name(&peer_node_addr),
|
||||
transport_id = %transport_id,
|
||||
remote_addr = %remote_addr,
|
||||
link_id = %link_id,
|
||||
"Transport connect initiated (non-blocking)"
|
||||
);
|
||||
self.pending_connects.push(super::PendingConnect {
|
||||
link_id,
|
||||
transport_id,
|
||||
remote_addr,
|
||||
peer_identity,
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
// Clean up link
|
||||
self.links.remove(&link_id);
|
||||
self.addr_to_link.remove(&(transport_id, remote_addr));
|
||||
return Err(NodeError::TransportError(e.to_string()));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
} else {
|
||||
// Connectionless: proceed with immediate handshake
|
||||
self.start_handshake(link_id, transport_id, remote_addr, peer_identity).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Start the Noise handshake on a link and send msg1.
|
||||
///
|
||||
/// Called immediately for connectionless transports, or after the
|
||||
/// transport connection is established for connection-oriented transports.
|
||||
pub(super) async fn start_handshake(
|
||||
&mut self,
|
||||
link_id: LinkId,
|
||||
transport_id: TransportId,
|
||||
remote_addr: TransportAddr,
|
||||
peer_identity: PeerIdentity,
|
||||
) -> Result<(), NodeError> {
|
||||
let peer_node_addr = *peer_identity.node_addr();
|
||||
|
||||
// Create connection in handshake phase (outbound knows expected identity)
|
||||
let current_time_ms = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
@@ -336,6 +387,99 @@ impl Node {
|
||||
}
|
||||
}
|
||||
|
||||
/// Poll pending transport connects and initiate handshakes for ready ones.
|
||||
///
|
||||
/// Called from the tick handler. For each pending connect, queries the
|
||||
/// transport's connection state. When a connection is established,
|
||||
/// marks the link as Connected and starts the Noise handshake.
|
||||
/// Failed connections are cleaned up and scheduled for retry.
|
||||
pub(super) async fn poll_pending_connects(&mut self) {
|
||||
if self.pending_connects.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let mut completed = Vec::new();
|
||||
|
||||
for (i, pending) in self.pending_connects.iter().enumerate() {
|
||||
let state = if let Some(transport) = self.transports.get(&pending.transport_id) {
|
||||
transport.connection_state(&pending.remote_addr)
|
||||
} else {
|
||||
crate::transport::ConnectionState::Failed("transport removed".into())
|
||||
};
|
||||
|
||||
match state {
|
||||
crate::transport::ConnectionState::Connected => {
|
||||
completed.push((i, true, None));
|
||||
}
|
||||
crate::transport::ConnectionState::Failed(reason) => {
|
||||
completed.push((i, false, Some(reason)));
|
||||
}
|
||||
crate::transport::ConnectionState::Connecting => {
|
||||
// Still in progress, check on next tick
|
||||
}
|
||||
crate::transport::ConnectionState::None => {
|
||||
// Shouldn't happen — treat as failure
|
||||
completed.push((i, false, Some("no connection attempt found".into())));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Process completions in reverse order to preserve indices
|
||||
for (i, success, reason) in completed.into_iter().rev() {
|
||||
let pending = self.pending_connects.remove(i);
|
||||
|
||||
if success {
|
||||
// Mark link as Connected
|
||||
if let Some(link) = self.links.get_mut(&pending.link_id) {
|
||||
link.set_connected();
|
||||
}
|
||||
|
||||
debug!(
|
||||
peer = %self.peer_display_name(pending.peer_identity.node_addr()),
|
||||
transport_id = %pending.transport_id,
|
||||
remote_addr = %pending.remote_addr,
|
||||
link_id = %pending.link_id,
|
||||
"Transport connected, starting handshake"
|
||||
);
|
||||
|
||||
// Start the handshake now that the transport is connected
|
||||
if let Err(e) = self.start_handshake(
|
||||
pending.link_id,
|
||||
pending.transport_id,
|
||||
pending.remote_addr.clone(),
|
||||
pending.peer_identity,
|
||||
).await {
|
||||
warn!(
|
||||
link_id = %pending.link_id,
|
||||
error = %e,
|
||||
"Failed to start handshake after transport connect"
|
||||
);
|
||||
// Clean up link on handshake failure
|
||||
self.remove_link(&pending.link_id);
|
||||
}
|
||||
} else {
|
||||
let reason = reason.unwrap_or_default();
|
||||
warn!(
|
||||
peer = %self.peer_display_name(pending.peer_identity.node_addr()),
|
||||
transport_id = %pending.transport_id,
|
||||
remote_addr = %pending.remote_addr,
|
||||
link_id = %pending.link_id,
|
||||
reason = %reason,
|
||||
"Transport connect failed"
|
||||
);
|
||||
|
||||
// Clean up link and schedule retry
|
||||
self.remove_link(&pending.link_id);
|
||||
self.links.remove(&pending.link_id);
|
||||
let now_ms = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as u64)
|
||||
.unwrap_or(0);
|
||||
self.schedule_retry(*pending.peer_identity.node_addr(), now_ms, false);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// === State Transitions ===
|
||||
|
||||
/// Start the node.
|
||||
|
||||
@@ -119,6 +119,9 @@ pub enum NodeError {
|
||||
|
||||
#[error("handshake failed: {0}")]
|
||||
HandshakeFailed(String),
|
||||
|
||||
#[error("transport error: {0}")]
|
||||
TransportError(String),
|
||||
}
|
||||
|
||||
/// Node operational state.
|
||||
@@ -208,6 +211,22 @@ struct TransportDropState {
|
||||
dropping: bool,
|
||||
}
|
||||
|
||||
/// State for a link waiting for transport-level connection establishment.
|
||||
///
|
||||
/// For connection-oriented transports (TCP, Tor), the transport connect runs
|
||||
/// asynchronously. This struct holds the data needed to complete the handshake
|
||||
/// once the connection is ready.
|
||||
struct PendingConnect {
|
||||
/// The link that was created for this connection.
|
||||
link_id: LinkId,
|
||||
/// Which transport is being used.
|
||||
transport_id: TransportId,
|
||||
/// The remote address being connected to.
|
||||
remote_addr: TransportAddr,
|
||||
/// The peer identity (for handshake initiation).
|
||||
peer_identity: PeerIdentity,
|
||||
}
|
||||
|
||||
/// A running FIPS node instance.
|
||||
///
|
||||
/// This is the top-level container holding all node state.
|
||||
@@ -363,6 +382,13 @@ pub struct Node {
|
||||
/// Rate limiter for source-side CoordsRequired/PathBroken responses.
|
||||
coords_response_rate_limiter: RoutingErrorRateLimiter,
|
||||
|
||||
// === Pending Transport Connects ===
|
||||
/// Links waiting for transport-level connection establishment before
|
||||
/// sending handshake msg1. For connection-oriented transports (TCP, Tor),
|
||||
/// the transport connect runs in the background; the tick handler polls
|
||||
/// connection_state() and initiates the handshake when connected.
|
||||
pending_connects: Vec<PendingConnect>,
|
||||
|
||||
// === Connection Retry ===
|
||||
/// Retry state for peers whose outbound connections have failed.
|
||||
/// Keyed by NodeAddr. Entries are created when a handshake times out
|
||||
@@ -499,6 +525,7 @@ impl Node {
|
||||
coords_response_rate_limiter: RoutingErrorRateLimiter::with_interval(
|
||||
std::time::Duration::from_millis(coords_response_interval_ms),
|
||||
),
|
||||
pending_connects: Vec::new(),
|
||||
retry_pending: HashMap::new(),
|
||||
last_parent_reeval: None,
|
||||
last_congestion_log: None,
|
||||
@@ -601,6 +628,7 @@ impl Node {
|
||||
coords_response_rate_limiter: RoutingErrorRateLimiter::with_interval(
|
||||
std::time::Duration::from_millis(coords_response_interval_ms),
|
||||
),
|
||||
pending_connects: Vec::new(),
|
||||
retry_pending: HashMap::new(),
|
||||
last_parent_reeval: None,
|
||||
last_congestion_log: None,
|
||||
|
||||
@@ -791,6 +791,26 @@ pub trait Transport {
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Connection State (for non-blocking connect)
|
||||
// ============================================================================
|
||||
|
||||
/// State of a transport-level connection attempt.
|
||||
///
|
||||
/// Used by connection-oriented transports (TCP, Tor) to report the progress
|
||||
/// of a background connection attempt initiated by `connect()`.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum ConnectionState {
|
||||
/// No connection attempt in progress for this address.
|
||||
None,
|
||||
/// Connection attempt is in progress (background task running).
|
||||
Connecting,
|
||||
/// Connection is established and ready for send().
|
||||
Connected,
|
||||
/// Connection attempt failed with the given error message.
|
||||
Failed(String),
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Transport Congestion
|
||||
// ============================================================================
|
||||
@@ -968,6 +988,36 @@ impl TransportHandle {
|
||||
}
|
||||
}
|
||||
|
||||
/// Initiate a non-blocking connection to a remote address.
|
||||
///
|
||||
/// For connection-oriented transports (TCP, Tor), spawns a background
|
||||
/// task to establish the connection. For connectionless transports
|
||||
/// (UDP, Ethernet), this is a no-op that returns Ok immediately.
|
||||
///
|
||||
/// Poll `connection_state()` to check when the connection is ready.
|
||||
pub async fn connect(&self, addr: &TransportAddr) -> Result<(), TransportError> {
|
||||
match self {
|
||||
TransportHandle::Udp(_) => Ok(()), // connectionless
|
||||
#[cfg(target_os = "linux")]
|
||||
TransportHandle::Ethernet(_) => Ok(()), // connectionless
|
||||
TransportHandle::Tcp(t) => t.connect_async(addr).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Query the state of a connection attempt to a remote address.
|
||||
///
|
||||
/// For connectionless transports, always returns `ConnectionState::Connected`
|
||||
/// (they are always "connected"). For connection-oriented transports, returns
|
||||
/// the current state of the background connection attempt.
|
||||
pub fn connection_state(&self, addr: &TransportAddr) -> ConnectionState {
|
||||
match self {
|
||||
TransportHandle::Udp(_) => ConnectionState::Connected,
|
||||
#[cfg(target_os = "linux")]
|
||||
TransportHandle::Ethernet(_) => ConnectionState::Connected,
|
||||
TransportHandle::Tcp(t) => t.connection_state_sync(addr),
|
||||
}
|
||||
}
|
||||
|
||||
/// Close a specific connection on this transport.
|
||||
///
|
||||
/// No-op for connectionless transports. For TCP, removes the
|
||||
|
||||
+395
-4
@@ -26,13 +26,14 @@ pub mod stats;
|
||||
pub mod stream;
|
||||
|
||||
use super::{
|
||||
DiscoveredPeer, PacketTx, ReceivedPacket, Transport, TransportAddr, TransportError,
|
||||
TransportId, TransportState, TransportType,
|
||||
ConnectionState, DiscoveredPeer, PacketTx, ReceivedPacket, Transport, TransportAddr,
|
||||
TransportError, TransportId, TransportState, TransportType,
|
||||
};
|
||||
use crate::config::TcpConfig;
|
||||
use stats::TcpStats;
|
||||
use stream::read_fmp_packet;
|
||||
|
||||
use futures::FutureExt;
|
||||
use socket2::TcpKeepalive;
|
||||
use std::collections::HashMap;
|
||||
use std::net::SocketAddr;
|
||||
@@ -67,6 +68,18 @@ struct TcpConnection {
|
||||
/// Shared connection pool.
|
||||
type ConnectionPool = Arc<Mutex<HashMap<TransportAddr, TcpConnection>>>;
|
||||
|
||||
/// A pending background connection attempt.
|
||||
///
|
||||
/// Holds the JoinHandle for a spawned TCP connect task. The task
|
||||
/// produces a configured `TcpStream` and MSS-derived MTU on success.
|
||||
struct ConnectingEntry {
|
||||
/// Background task performing TCP connect + socket configuration.
|
||||
task: JoinHandle<Result<(TcpStream, u16), TransportError>>,
|
||||
}
|
||||
|
||||
/// Map of addresses with background connection attempts in progress.
|
||||
type ConnectingPool = Arc<Mutex<HashMap<TransportAddr, ConnectingEntry>>>;
|
||||
|
||||
// ============================================================================
|
||||
// TCP Transport
|
||||
// ============================================================================
|
||||
@@ -85,8 +98,10 @@ pub struct TcpTransport {
|
||||
config: TcpConfig,
|
||||
/// Current state.
|
||||
state: TransportState,
|
||||
/// Connection pool: addr -> per-connection state.
|
||||
/// Connection pool: addr -> established connections.
|
||||
pool: ConnectionPool,
|
||||
/// Pending connection attempts: addr -> background connect task.
|
||||
connecting: ConnectingPool,
|
||||
/// Channel for delivering received packets to Node.
|
||||
packet_tx: PacketTx,
|
||||
/// Accept loop task handle (if listener bound).
|
||||
@@ -111,6 +126,7 @@ impl TcpTransport {
|
||||
config,
|
||||
state: TransportState::Configured,
|
||||
pool: Arc::new(Mutex::new(HashMap::new())),
|
||||
connecting: Arc::new(Mutex::new(HashMap::new())),
|
||||
packet_tx,
|
||||
accept_task: None,
|
||||
local_addr: None,
|
||||
@@ -212,7 +228,19 @@ impl TcpTransport {
|
||||
let _ = task.await;
|
||||
}
|
||||
|
||||
// Close all connections
|
||||
// Abort pending connection attempts
|
||||
let mut connecting = self.connecting.lock().await;
|
||||
for (addr, entry) in connecting.drain() {
|
||||
entry.task.abort();
|
||||
debug!(
|
||||
transport_id = %self.transport_id,
|
||||
remote_addr = %addr,
|
||||
"TCP connect aborted (transport stopping)"
|
||||
);
|
||||
}
|
||||
drop(connecting);
|
||||
|
||||
// Close all established connections
|
||||
let mut pool = self.pool.lock().await;
|
||||
for (addr, conn) in pool.drain() {
|
||||
conn.recv_task.abort();
|
||||
@@ -396,6 +424,214 @@ impl TcpTransport {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Initiate a non-blocking connection to a remote address.
|
||||
///
|
||||
/// Spawns a background task that performs TCP connect with timeout,
|
||||
/// configures socket options, and reads MSS. The connection becomes
|
||||
/// available for `send_async()` once the task completes successfully.
|
||||
///
|
||||
/// Poll `connection_state_sync()` to check progress.
|
||||
pub async fn connect_async(&self, addr: &TransportAddr) -> Result<(), TransportError> {
|
||||
if !self.state.is_operational() {
|
||||
return Err(TransportError::NotStarted);
|
||||
}
|
||||
|
||||
// Already established?
|
||||
{
|
||||
let pool = self.pool.lock().await;
|
||||
if pool.contains_key(addr) {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
// Already connecting?
|
||||
{
|
||||
let connecting = self.connecting.lock().await;
|
||||
if connecting.contains_key(addr) {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
let socket_addr = parse_socket_addr(addr)?;
|
||||
let timeout_ms = self.config.connect_timeout_ms();
|
||||
let config = self.config.clone();
|
||||
let transport_id = self.transport_id;
|
||||
let remote_addr = addr.clone();
|
||||
|
||||
debug!(
|
||||
transport_id = %transport_id,
|
||||
remote_addr = %remote_addr,
|
||||
timeout_ms,
|
||||
"Initiating background TCP connect"
|
||||
);
|
||||
|
||||
let task = tokio::spawn(async move {
|
||||
// Connect with timeout
|
||||
let stream = match tokio::time::timeout(
|
||||
Duration::from_millis(timeout_ms),
|
||||
TcpStream::connect(socket_addr),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(stream)) => stream,
|
||||
Ok(Err(e)) => {
|
||||
debug!(
|
||||
transport_id = %transport_id,
|
||||
remote_addr = %remote_addr,
|
||||
error = %e,
|
||||
"Background TCP connect refused"
|
||||
);
|
||||
return Err(TransportError::ConnectionRefused);
|
||||
}
|
||||
Err(_) => {
|
||||
debug!(
|
||||
transport_id = %transport_id,
|
||||
remote_addr = %remote_addr,
|
||||
"Background TCP connect timed out"
|
||||
);
|
||||
return Err(TransportError::Timeout);
|
||||
}
|
||||
};
|
||||
|
||||
// Configure socket options via socket2
|
||||
let std_stream = stream.into_std()
|
||||
.map_err(|e| TransportError::StartFailed(format!("into_std: {}", e)))?;
|
||||
configure_socket(&std_stream, &config)?;
|
||||
|
||||
// Read TCP_MAXSEG for per-connection MTU
|
||||
let mss_mtu = read_mss_mtu(&std_stream, config.mtu());
|
||||
|
||||
// Convert back to tokio
|
||||
let stream = TcpStream::from_std(std_stream)
|
||||
.map_err(|e| TransportError::StartFailed(format!("from_std: {}", e)))?;
|
||||
|
||||
Ok((stream, mss_mtu))
|
||||
});
|
||||
|
||||
let mut connecting = self.connecting.lock().await;
|
||||
connecting.insert(addr.clone(), ConnectingEntry { task });
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Query the state of a connection to a remote address.
|
||||
///
|
||||
/// Checks both established and connecting pools. If a background
|
||||
/// connect task has completed, promotes it to the established pool
|
||||
/// (spawning a receive loop) or reports the failure.
|
||||
///
|
||||
/// This method is synchronous but uses `try_lock` internally.
|
||||
/// Returns `ConnectionState::Connecting` if locks can't be acquired.
|
||||
pub fn connection_state_sync(&self, addr: &TransportAddr) -> ConnectionState {
|
||||
// Check established pool first
|
||||
if let Ok(pool) = self.pool.try_lock() {
|
||||
if pool.contains_key(addr) {
|
||||
return ConnectionState::Connected;
|
||||
}
|
||||
} else {
|
||||
return ConnectionState::Connecting; // can't tell, assume still going
|
||||
}
|
||||
|
||||
// Check connecting pool
|
||||
let mut connecting = match self.connecting.try_lock() {
|
||||
Ok(c) => c,
|
||||
Err(_) => return ConnectionState::Connecting,
|
||||
};
|
||||
|
||||
let entry = match connecting.get_mut(addr) {
|
||||
Some(e) => e,
|
||||
None => return ConnectionState::None,
|
||||
};
|
||||
|
||||
// Check if the background task has completed
|
||||
if !entry.task.is_finished() {
|
||||
return ConnectionState::Connecting;
|
||||
}
|
||||
|
||||
// Task is done — take the result and remove from connecting pool.
|
||||
// We need to poll the finished task. Since it's finished, we use
|
||||
// now_or_never to get the result without blocking.
|
||||
let addr_clone = addr.clone();
|
||||
let task = connecting.remove(&addr_clone).unwrap().task;
|
||||
|
||||
// Use futures::FutureExt::now_or_never or block_on for the finished task.
|
||||
// Since the task is finished, we can safely poll it.
|
||||
match task.now_or_never() {
|
||||
Some(Ok(Ok((stream, mss_mtu)))) => {
|
||||
// Promote to established pool
|
||||
self.promote_connection(addr, stream, mss_mtu);
|
||||
ConnectionState::Connected
|
||||
}
|
||||
Some(Ok(Err(e))) => {
|
||||
ConnectionState::Failed(format!("{}", e))
|
||||
}
|
||||
Some(Err(e)) => {
|
||||
// JoinError (panic or cancel)
|
||||
ConnectionState::Failed(format!("task failed: {}", e))
|
||||
}
|
||||
None => {
|
||||
// Shouldn't happen since is_finished() was true
|
||||
ConnectionState::Connecting
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Promote a completed background connection to the established pool.
|
||||
///
|
||||
/// Splits the stream, spawns a receive loop, and inserts into the pool.
|
||||
/// Called from `connection_state_sync()` when a background task completes.
|
||||
fn promote_connection(&self, addr: &TransportAddr, stream: TcpStream, mss_mtu: u16) {
|
||||
let (read_half, write_half) = stream.into_split();
|
||||
let writer = Arc::new(Mutex::new(write_half));
|
||||
|
||||
let transport_id = self.transport_id;
|
||||
let packet_tx = self.packet_tx.clone();
|
||||
let pool = self.pool.clone();
|
||||
let recv_stats = self.stats.clone();
|
||||
let remote_addr = addr.clone();
|
||||
|
||||
let recv_task = tokio::spawn(async move {
|
||||
tcp_receive_loop(
|
||||
read_half,
|
||||
transport_id,
|
||||
remote_addr.clone(),
|
||||
packet_tx,
|
||||
pool,
|
||||
mss_mtu,
|
||||
recv_stats,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
|
||||
let conn = TcpConnection {
|
||||
writer,
|
||||
recv_task,
|
||||
mtu: mss_mtu,
|
||||
established_at: Instant::now(),
|
||||
};
|
||||
|
||||
// Use try_lock since we're in a sync context and the pool
|
||||
// should be available (connection_state_sync already checked it)
|
||||
if let Ok(mut pool) = self.pool.try_lock() {
|
||||
pool.insert(addr.clone(), conn);
|
||||
self.stats.record_connection_established();
|
||||
debug!(
|
||||
transport_id = %self.transport_id,
|
||||
remote_addr = %addr,
|
||||
mtu = mss_mtu,
|
||||
"TCP connection established (background connect)"
|
||||
);
|
||||
} else {
|
||||
// Pool locked — abort the recv task, connection will be retried
|
||||
conn.recv_task.abort();
|
||||
warn!(
|
||||
transport_id = %self.transport_id,
|
||||
remote_addr = %addr,
|
||||
"Failed to promote connection (pool locked)"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Transport for TcpTransport {
|
||||
@@ -1120,4 +1356,159 @@ mod tests {
|
||||
t1.stop_async().await.unwrap();
|
||||
t2.stop_async().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_connect_async_success() {
|
||||
let (tx1, mut rx1) = packet_channel(100);
|
||||
let (tx2, _rx2) = packet_channel(100);
|
||||
|
||||
let mut t1 = TcpTransport::new(TransportId::new(1), None, make_outbound_config(), tx1);
|
||||
let mut t2 = TcpTransport::new(TransportId::new(2), None, make_config(), tx2);
|
||||
|
||||
t1.start_async().await.unwrap();
|
||||
t2.start_async().await.unwrap();
|
||||
|
||||
let addr2 = t2.local_addr().unwrap();
|
||||
let remote = TransportAddr::from_string(&addr2.to_string());
|
||||
|
||||
// State should be None before connect
|
||||
assert_eq!(t1.connection_state_sync(&remote), ConnectionState::None);
|
||||
|
||||
// Initiate non-blocking connect
|
||||
t1.connect_async(&remote).await.unwrap();
|
||||
|
||||
// Wait for the background connect to complete
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
|
||||
// Poll state — should be Connected now
|
||||
let state = t1.connection_state_sync(&remote);
|
||||
assert_eq!(state, ConnectionState::Connected);
|
||||
|
||||
// Now send should work (connection already established)
|
||||
let mut msg1 = vec![0xAA; 114];
|
||||
msg1[0] = 0x01;
|
||||
msg1[1] = 0x00;
|
||||
msg1[2..4].copy_from_slice(&110u16.to_le_bytes());
|
||||
|
||||
t1.send_async(&remote, &msg1).await.unwrap();
|
||||
|
||||
let packet = timeout(Duration::from_secs(2), rx1.recv())
|
||||
.await;
|
||||
// We receive on rx1 but that's the wrong receiver — t2's rx gets the packet
|
||||
// Just verify send didn't error
|
||||
drop(packet);
|
||||
|
||||
t1.stop_async().await.unwrap();
|
||||
t2.stop_async().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_connect_async_timeout() {
|
||||
let (tx, _rx) = packet_channel(100);
|
||||
let config = TcpConfig {
|
||||
bind_addr: None,
|
||||
connect_timeout_ms: Some(100), // Very short timeout
|
||||
..Default::default()
|
||||
};
|
||||
let mut transport = TcpTransport::new(TransportId::new(1), None, config, tx);
|
||||
transport.start_async().await.unwrap();
|
||||
|
||||
let remote = TransportAddr::from_string("192.0.2.1:2121");
|
||||
transport.connect_async(&remote).await.unwrap();
|
||||
|
||||
// Wait for timeout
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
|
||||
let state = transport.connection_state_sync(&remote);
|
||||
assert!(matches!(state, ConnectionState::Failed(_)));
|
||||
|
||||
transport.stop_async().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_connect_async_not_started() {
|
||||
let (tx, _rx) = packet_channel(100);
|
||||
let transport = TcpTransport::new(TransportId::new(1), None, make_config(), tx);
|
||||
|
||||
let result = transport
|
||||
.connect_async(&TransportAddr::from_string("127.0.0.1:9999"))
|
||||
.await;
|
||||
|
||||
assert!(matches!(result, Err(TransportError::NotStarted)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_connect_async_already_connected() {
|
||||
let (tx1, _rx1) = packet_channel(100);
|
||||
let (tx2, _rx2) = packet_channel(100);
|
||||
|
||||
let mut t1 = TcpTransport::new(TransportId::new(1), None, make_outbound_config(), tx1);
|
||||
let mut t2 = TcpTransport::new(TransportId::new(2), None, make_config(), tx2);
|
||||
|
||||
t1.start_async().await.unwrap();
|
||||
t2.start_async().await.unwrap();
|
||||
|
||||
let addr2 = t2.local_addr().unwrap();
|
||||
let remote = TransportAddr::from_string(&addr2.to_string());
|
||||
|
||||
// Connect first time
|
||||
t1.connect_async(&remote).await.unwrap();
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
assert_eq!(t1.connection_state_sync(&remote), ConnectionState::Connected);
|
||||
|
||||
// Second connect should be a no-op (already connected)
|
||||
t1.connect_async(&remote).await.unwrap();
|
||||
|
||||
t1.stop_async().await.unwrap();
|
||||
t2.stop_async().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_connect_async_then_send_recv() {
|
||||
let (tx1, _rx1) = packet_channel(100);
|
||||
let (tx2, mut rx2) = packet_channel(100);
|
||||
|
||||
let mut t1 = TcpTransport::new(TransportId::new(1), None, make_outbound_config(), tx1);
|
||||
let mut t2 = TcpTransport::new(TransportId::new(2), None, make_config(), tx2);
|
||||
|
||||
t1.start_async().await.unwrap();
|
||||
t2.start_async().await.unwrap();
|
||||
|
||||
let addr2 = t2.local_addr().unwrap();
|
||||
let remote = TransportAddr::from_string(&addr2.to_string());
|
||||
|
||||
// Connect first, then send
|
||||
t1.connect_async(&remote).await.unwrap();
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
assert_eq!(t1.connection_state_sync(&remote), ConnectionState::Connected);
|
||||
|
||||
// Build valid FMP msg1 frame
|
||||
let mut msg1 = vec![0xAA; 114];
|
||||
msg1[0] = 0x01;
|
||||
msg1[1] = 0x00;
|
||||
msg1[2..4].copy_from_slice(&110u16.to_le_bytes());
|
||||
|
||||
// Send using the pre-established connection
|
||||
t1.send_async(&remote, &msg1).await.unwrap();
|
||||
|
||||
let packet = timeout(Duration::from_secs(2), rx2.recv())
|
||||
.await
|
||||
.expect("timeout")
|
||||
.expect("channel closed");
|
||||
assert_eq!(packet.data, msg1);
|
||||
|
||||
t1.stop_async().await.unwrap();
|
||||
t2.stop_async().await.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_connection_state_none_for_unknown() {
|
||||
let (tx, _rx) = packet_channel(100);
|
||||
let transport = TcpTransport::new(TransportId::new(1), None, make_config(), tx);
|
||||
|
||||
let state = transport.connection_state_sync(
|
||||
&TransportAddr::from_string("unknown:1234"),
|
||||
);
|
||||
assert_eq!(state, ConnectionState::None);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user