pub mod stats;
pub mod stream;
mod tasks;
use tasks::{AcceptConfig, TcpReceiveContext, accept_loop, tcp_receive_loop};
use super::resolve_socket_addrs;
use super::{
ConnectionState, DiscoveredPeer, PacketTx, Transport, TransportAddr, TransportError,
TransportId, TransportState, TransportType,
};
use crate::config::TcpConfig;
use stats::TcpStats;
use futures::FutureExt;
use socket2::TcpKeepalive;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::AsyncWriteExt;
use tokio::net::tcp::OwnedWriteHalf;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use tokio::time::Instant;
use tracing::{debug, info, trace, warn};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Direction {
Inbound,
Outbound,
}
struct TcpConnection {
writer: Arc<Mutex<OwnedWriteHalf>>,
recv_task: JoinHandle<()>,
#[allow(dead_code)]
mtu: u16,
#[allow(dead_code)]
established_at: Instant,
direction: Direction,
}
type ConnectionPool = Arc<Mutex<HashMap<TransportAddr, TcpConnection>>>;
struct ConnectingEntry {
task: JoinHandle<Result<(TcpStream, u16), TransportError>>,
}
type ConnectingPool = Arc<Mutex<HashMap<TransportAddr, ConnectingEntry>>>;
pub struct TcpTransport {
transport_id: TransportId,
name: Option<String>,
config: TcpConfig,
state: TransportState,
pool: ConnectionPool,
connecting: ConnectingPool,
packet_tx: PacketTx,
accept_task: Option<JoinHandle<()>>,
local_addr: Option<SocketAddr>,
stats: Arc<TcpStats>,
}
impl TcpTransport {
pub fn new(
transport_id: TransportId,
name: Option<String>,
config: TcpConfig,
packet_tx: PacketTx,
) -> Self {
Self {
transport_id,
name,
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,
stats: Arc::new(TcpStats::new()),
}
}
pub fn name(&self) -> Option<&str> {
self.name.as_deref()
}
pub fn local_addr(&self) -> Option<SocketAddr> {
self.local_addr
}
pub fn stats(&self) -> &Arc<TcpStats> {
&self.stats
}
pub async fn start_async(&mut self) -> Result<(), TransportError> {
if !self.state.can_start() {
return Err(TransportError::AlreadyStarted);
}
self.state = TransportState::Starting;
if let Some(ref bind_addr) = self.config.bind_addr {
let addr: SocketAddr = bind_addr
.parse()
.map_err(|e| TransportError::StartFailed(format!("invalid bind address: {}", e)))?;
let listener = TcpListener::bind(addr)
.await
.map_err(|e| TransportError::StartFailed(format!("bind failed: {}", e)))?;
self.local_addr = Some(
listener
.local_addr()
.map_err(|e| TransportError::StartFailed(format!("get local addr: {}", e)))?,
);
let transport_id = self.transport_id;
let packet_tx = self.packet_tx.clone();
let pool = self.pool.clone();
let stats = self.stats.clone();
let cfg = AcceptConfig {
mtu: self.config.mtu(),
max_inbound: self.config.max_inbound_connections(),
nodelay: self.config.nodelay(),
keepalive_secs: self.config.keepalive_secs(),
recv_buf: self.config.recv_buf_size(),
send_buf: self.config.send_buf_size(),
first_frame_timeout_ms: self.config.first_frame_timeout_ms(),
};
let accept_task = tokio::spawn(async move {
accept_loop(listener, transport_id, packet_tx, pool, cfg, stats).await;
});
self.accept_task = Some(accept_task);
}
self.state = TransportState::Up;
if let Some(ref name) = self.name {
info!(
name = %name,
local_addr = ?self.local_addr,
mtu = self.config.mtu(),
"TCP transport started"
);
} else {
info!(
local_addr = ?self.local_addr,
mtu = self.config.mtu(),
"TCP transport started"
);
}
Ok(())
}
pub async fn stop_async(&mut self) -> Result<(), TransportError> {
if !self.state.is_operational() {
return Err(TransportError::NotStarted);
}
if let Some(task) = self.accept_task.take() {
task.abort();
let _ = task.await;
}
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);
let mut pool = self.pool.lock().await;
for (addr, conn) in pool.drain() {
conn.recv_task.abort();
let _ = conn.recv_task.await;
match conn.direction {
Direction::Inbound => self.stats.record_pool_inbound_removed(),
Direction::Outbound => self.stats.record_pool_outbound_removed(),
}
debug!(
transport_id = %self.transport_id,
remote_addr = %addr,
direction = ?conn.direction,
"TCP connection closed (transport stopping)"
);
}
drop(pool);
self.local_addr = None;
self.state = TransportState::Down;
info!(
transport_id = %self.transport_id,
"TCP transport stopped"
);
Ok(())
}
pub async fn send_async(
&self,
addr: &TransportAddr,
data: &[u8],
) -> Result<usize, TransportError> {
if !self.state.is_operational() {
return Err(TransportError::NotStarted);
}
let mtu = self.config.mtu() as usize;
if data.len() > mtu {
self.stats.record_mtu_exceeded();
return Err(TransportError::MtuExceeded {
packet_size: data.len(),
mtu: self.config.mtu(),
});
}
let writer = {
let pool = self.pool.lock().await;
pool.get(addr).map(|c| c.writer.clone())
};
let writer = match writer {
Some(w) => w,
None => {
self.connect(addr).await?
}
};
let mut w = writer.lock().await;
match w.write_all(data).await {
Ok(()) => {
self.stats.record_send(data.len());
trace!(
transport_id = %self.transport_id,
remote_addr = %addr,
bytes = data.len(),
"TCP packet sent"
);
Ok(data.len())
}
Err(e) => {
self.stats.record_send_error();
drop(w);
let mut pool = self.pool.lock().await;
if let Some(conn) = pool.remove(addr) {
conn.recv_task.abort();
match conn.direction {
Direction::Inbound => self.stats.record_pool_inbound_removed(),
Direction::Outbound => self.stats.record_pool_outbound_removed(),
}
}
Err(TransportError::SendFailed(format!("{}", e)))
}
}
}
async fn connect(
&self,
addr: &TransportAddr,
) -> Result<Arc<Mutex<OwnedWriteHalf>>, TransportError> {
let socket_addrs = resolve_socket_addrs(addr).await?;
let timeout_ms = self.config.connect_timeout_ms();
let stream = match connect_to_any_addr(&socket_addrs, timeout_ms).await {
Ok(stream) => stream,
Err(error @ TransportError::Timeout) => {
self.stats.record_connect_timeout();
return Err(error);
}
Err(error @ TransportError::ConnectionRefused) => {
self.stats.record_connect_refused();
return Err(error);
}
Err(error) => return Err(error),
};
let std_stream = stream
.into_std()
.map_err(|e| TransportError::StartFailed(format!("into_std: {}", e)))?;
configure_socket(&std_stream, &self.config)?;
let mss_mtu = read_mss_mtu(&std_stream, self.config.mtu());
let stream = TcpStream::from_std(std_stream)
.map_err(|e| TransportError::StartFailed(format!("from_std: {}", e)))?;
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 mtu = mss_mtu;
let recv_task = tokio::spawn(async move {
tcp_receive_loop(
read_half,
TcpReceiveContext {
transport_id,
remote_addr,
packet_tx,
pool,
mtu,
stats: recv_stats,
first_frame_timeout: None,
direction: Direction::Outbound,
},
)
.await;
});
let conn = TcpConnection {
writer: writer.clone(),
recv_task,
mtu: mss_mtu,
established_at: Instant::now(),
direction: Direction::Outbound,
};
let mut pool = self.pool.lock().await;
pool.insert(addr.clone(), conn);
self.stats.record_connection_established();
self.stats.record_pool_outbound_added();
debug!(
transport_id = %self.transport_id,
remote_addr = %addr,
mtu = mss_mtu,
"TCP connection established (connect-on-send)"
);
Ok(writer)
}
pub async fn close_connection_async(&self, addr: &TransportAddr) {
let mut pool = self.pool.lock().await;
if let Some(conn) = pool.remove(addr) {
conn.recv_task.abort();
match conn.direction {
Direction::Inbound => self.stats.record_pool_inbound_removed(),
Direction::Outbound => self.stats.record_pool_outbound_removed(),
}
debug!(
transport_id = %self.transport_id,
remote_addr = %addr,
direction = ?conn.direction,
"TCP connection closed (close_connection)"
);
}
}
pub async fn connect_async(&self, addr: &TransportAddr) -> Result<(), TransportError> {
if !self.state.is_operational() {
return Err(TransportError::NotStarted);
}
{
let pool = self.pool.lock().await;
if pool.contains_key(addr) {
return Ok(());
}
}
{
let connecting = self.connecting.lock().await;
if connecting.contains_key(addr) {
return Ok(());
}
}
let socket_addrs = resolve_socket_addrs(addr).await?;
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 {
let stream = match connect_to_any_addr(&socket_addrs, timeout_ms).await {
Ok(stream) => stream,
Err(error @ TransportError::ConnectionRefused) => {
debug!(
transport_id = %transport_id,
remote_addr = %remote_addr,
error = %error,
"Background TCP connect refused"
);
return Err(error);
}
Err(error @ TransportError::Timeout) => {
debug!(
transport_id = %transport_id,
remote_addr = %remote_addr,
"Background TCP connect timed out"
);
return Err(error);
}
Err(error) => return Err(error),
};
let std_stream = stream
.into_std()
.map_err(|e| TransportError::StartFailed(format!("into_std: {}", e)))?;
configure_socket(&std_stream, &config)?;
let mss_mtu = read_mss_mtu(&std_stream, config.mtu());
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(())
}
pub fn connection_state_sync(&self, addr: &TransportAddr) -> ConnectionState {
if let Ok(pool) = self.pool.try_lock() {
if pool.contains_key(addr) {
return ConnectionState::Connected;
}
} else {
return ConnectionState::Connecting; }
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,
};
if !entry.task.is_finished() {
return ConnectionState::Connecting;
}
let addr_clone = addr.clone();
let task = connecting.remove(&addr_clone).unwrap().task;
match task.now_or_never() {
Some(Ok(Ok((stream, mss_mtu)))) => {
self.promote_connection(addr, stream, mss_mtu);
ConnectionState::Connected
}
Some(Ok(Err(e))) => ConnectionState::Failed(format!("{}", e)),
Some(Err(e)) => {
ConnectionState::Failed(format!("task failed: {}", e))
}
None => {
ConnectionState::Connecting
}
}
}
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,
TcpReceiveContext {
transport_id,
remote_addr,
packet_tx,
pool,
mtu: mss_mtu,
stats: recv_stats,
first_frame_timeout: None,
direction: Direction::Outbound,
},
)
.await;
});
let conn = TcpConnection {
writer,
recv_task,
mtu: mss_mtu,
established_at: Instant::now(),
direction: Direction::Outbound,
};
if let Ok(mut pool) = self.pool.try_lock() {
pool.insert(addr.clone(), conn);
self.stats.record_connection_established();
self.stats.record_pool_outbound_added();
debug!(
transport_id = %self.transport_id,
remote_addr = %addr,
mtu = mss_mtu,
"TCP connection established (background connect)"
);
} else {
conn.recv_task.abort();
warn!(
transport_id = %self.transport_id,
remote_addr = %addr,
"Failed to promote connection (pool locked)"
);
}
}
}
impl Transport for TcpTransport {
fn transport_id(&self) -> TransportId {
self.transport_id
}
fn transport_type(&self) -> &TransportType {
&TransportType::TCP
}
fn state(&self) -> TransportState {
self.state
}
fn mtu(&self) -> u16 {
self.config.mtu()
}
fn link_mtu(&self, _addr: &TransportAddr) -> u16 {
self.config.mtu()
}
fn start(&mut self) -> Result<(), TransportError> {
Err(TransportError::NotSupported(
"use start_async() for TCP transport".into(),
))
}
fn stop(&mut self) -> Result<(), TransportError> {
Err(TransportError::NotSupported(
"use stop_async() for TCP transport".into(),
))
}
fn send(&self, _addr: &TransportAddr, _data: &[u8]) -> Result<(), TransportError> {
Err(TransportError::NotSupported(
"use send_async() for TCP transport".into(),
))
}
fn discover(&self) -> Result<Vec<DiscoveredPeer>, TransportError> {
Ok(Vec::new())
}
fn accept_connections(&self) -> bool {
self.config.bind_addr.is_some()
}
}
async fn connect_to_any_addr(
socket_addrs: &[SocketAddr],
timeout_ms: u64,
) -> Result<TcpStream, TransportError> {
let mut last_error = None;
for socket_addr in socket_addrs {
match tokio::time::timeout(
Duration::from_millis(timeout_ms),
TcpStream::connect(socket_addr),
)
.await
{
Ok(Ok(stream)) => return Ok(stream),
Ok(Err(error)) => {
trace!(
remote_addr = %socket_addr,
error = %error,
"TCP connect candidate failed"
);
last_error = Some(TransportError::ConnectionRefused);
}
Err(_) => {
trace!(
remote_addr = %socket_addr,
timeout_ms,
"TCP connect candidate timed out"
);
last_error = Some(TransportError::Timeout);
}
}
}
Err(last_error
.unwrap_or_else(|| TransportError::InvalidAddress("no TCP addresses to dial".to_string())))
}
fn configure_socket(
stream: &std::net::TcpStream,
config: &TcpConfig,
) -> Result<(), TransportError> {
let socket = socket2::SockRef::from(stream)
.try_clone()
.map_err(|e| TransportError::StartFailed(format!("clone socket: {}", e)))?;
socket
.set_tcp_nodelay(config.nodelay())
.map_err(|e| TransportError::StartFailed(format!("set nodelay: {}", e)))?;
let keepalive_secs = config.keepalive_secs();
if keepalive_secs > 0 {
let keepalive = TcpKeepalive::new().with_time(Duration::from_secs(keepalive_secs));
socket
.set_tcp_keepalive(&keepalive)
.map_err(|e| TransportError::StartFailed(format!("set keepalive: {}", e)))?;
}
socket
.set_recv_buffer_size(config.recv_buf_size())
.map_err(|e| TransportError::StartFailed(format!("set recv buffer: {}", e)))?;
socket
.set_send_buffer_size(config.send_buf_size())
.map_err(|e| TransportError::StartFailed(format!("set send buffer: {}", e)))?;
Ok(())
}
fn configure_accepted_socket(
stream: &std::net::TcpStream,
nodelay: bool,
keepalive_secs: u64,
recv_buf: usize,
send_buf: usize,
) -> Result<(), TransportError> {
let socket = socket2::SockRef::from(stream)
.try_clone()
.map_err(|e| TransportError::StartFailed(format!("clone socket: {}", e)))?;
socket
.set_tcp_nodelay(nodelay)
.map_err(|e| TransportError::StartFailed(format!("set nodelay: {}", e)))?;
if keepalive_secs > 0 {
let keepalive = TcpKeepalive::new().with_time(Duration::from_secs(keepalive_secs));
socket
.set_tcp_keepalive(&keepalive)
.map_err(|e| TransportError::StartFailed(format!("set keepalive: {}", e)))?;
}
socket
.set_recv_buffer_size(recv_buf)
.map_err(|e| TransportError::StartFailed(format!("set recv buffer: {}", e)))?;
socket
.set_send_buffer_size(send_buf)
.map_err(|e| TransportError::StartFailed(format!("set send buffer: {}", e)))?;
Ok(())
}
fn read_mss_mtu(stream: &std::net::TcpStream, default_mtu: u16) -> u16 {
#[cfg(target_os = "linux")]
{
use std::os::unix::io::AsRawFd;
unsafe {
let mut mss: libc::c_int = 0;
let mut len: libc::socklen_t = std::mem::size_of::<libc::c_int>() as libc::socklen_t;
let fd = stream.as_raw_fd();
let ret = libc::getsockopt(
fd,
libc::IPPROTO_TCP,
libc::TCP_MAXSEG,
&mut mss as *mut libc::c_int as *mut libc::c_void,
&mut len,
);
if ret == 0 && mss > 0 {
let mss_mtu = (mss as u32).min(u16::MAX as u32) as u16;
return mss_mtu.min(default_mtu);
}
}
}
#[cfg(not(target_os = "linux"))]
let _ = stream;
default_mtu
}
#[cfg(test)]
mod tests;