diff --git a/easytier/src/tunnel/fake_tcp/mod.rs b/easytier/src/tunnel/fake_tcp/mod.rs index 0f27d84f..9cd052f8 100644 --- a/easytier/src/tunnel/fake_tcp/mod.rs +++ b/easytier/src/tunnel/fake_tcp/mod.rs @@ -12,7 +12,7 @@ use std::{ sync::Arc, task::{Context as TaskContext, Poll}, }; -use tokio::{io::AsyncReadExt, net::TcpStream, sync::Mutex}; +use tokio::{io::AsyncReadExt, net::TcpStream}; use crate::tunnel::{ FromUrl, IpVersion, SinkError, SinkItem, StreamItem, Tunnel, TunnelConnector, TunnelError, @@ -85,7 +85,7 @@ pub struct FakeTcpTunnelListener { addr: url::Url, os_listener: Option, // interface_name -> fake tcp stack - stack_map: DashMap>>, + stack_map: DashMap>, // a cache from ip addr to interface name ip_to_ifname: IpToIfNameCache, } @@ -148,7 +148,7 @@ impl FakeTcpTunnelListener { async fn get_stack( &self, accept_result: &AcceptResult, - ) -> Result>, TunnelError> { + ) -> Result, TunnelError> { let local_socket_addr = accept_result.local_addr; let interface_name = &accept_result.interface_name; @@ -158,29 +158,38 @@ impl FakeTcpTunnelListener { IpAddr::V6(ip) => (None, Some(ip)), }; - let ret = match self.stack_map.entry(interface_name.to_string()) { - dashmap::Entry::Occupied(entry) => entry.get().clone(), - dashmap::Entry::Vacant(entry) => { - let tun = - create_tun_off_runtime(interface_name.to_string(), None, local_socket_addr) - .await?; - tracing::info!( - ?local_socket_addr, - "create new stack with interface_name: {:?}", - interface_name - ); - let stack = Arc::new(Mutex::new(stack::Stack::new( - tun, - local_ip.unwrap_or(Ipv4Addr::UNSPECIFIED), - local_ip6, - accept_result.mac, - ))); - entry.insert(stack.clone()); - stack - } - }; + if let Some(entry) = self.stack_map.get(interface_name) { + let stack = entry.clone(); + drop(entry); - Ok(ret) + if !stack.is_closed() { + return Ok(stack); + } + + tracing::warn!( + interface_name, + "fake_tcp stack reader_task finished, recreating stack" + ); + self.stack_map.remove(interface_name); + } + + let tun = + create_tun_off_runtime(interface_name.to_string(), None, local_socket_addr).await?; + tracing::info!( + ?local_socket_addr, + "create new stack with interface_name: {:?}", + interface_name + ); + let stack = Arc::new(stack::Stack::new( + tun, + local_ip.unwrap_or(Ipv4Addr::UNSPECIFIED), + local_ip6, + accept_result.mac, + )); + self.stack_map + .insert(interface_name.to_string(), stack.clone()); + + Ok(stack) } } @@ -215,19 +224,29 @@ impl TunnelListener for FakeTcpTunnelListener { let os_listener = tokio::net::TcpListener::bind(bind_addr).await?; tracing::info!(port, "FakeTcpTunnelListener listening"); self.os_listener = Some(os_listener); - // self.stack.lock().await.listen(port); Ok(()) } async fn accept(&mut self) -> Result, TunnelError> { tracing::debug!("FakeTcpTunnelListener waiting for accept"); - let res = self.do_accept().await?; - let stack = self.get_stack(&res).await?; - let socket = stack - .lock() - .await - .alloc_established_socket(res.local_addr, res.remote_addr, stack::State::Established) - .await; + let (res, stack, socket) = loop { + let res = self.do_accept().await?; + let stack = self.get_stack(&res).await?; + let socket = stack.try_alloc_established_socket( + res.local_addr, + res.remote_addr, + stack::State::Established, + ); + let Some(socket) = socket else { + tracing::warn!( + interface_name = res.interface_name, + "fake_tcp stack closed while accepting connection, dropping accepted socket" + ); + self.stack_map.remove(&res.interface_name); + continue; + }; + break (res, stack, socket); + }; tracing::info!( ?res, @@ -236,7 +255,7 @@ impl TunnelListener for FakeTcpTunnelListener { ); let info = TunnelInfo { - tunnel_type: get_faketcp_tunnel_type_str(stack.lock().await.driver_type()), + tunnel_type: get_faketcp_tunnel_type_str(stack.driver_type()), local_addr: Some(self.local_url().into()), remote_addr: Some( crate::tunnel::build_url_from_socket_addr( @@ -354,12 +373,14 @@ impl TunnelConnector for FakeTcpTunnelConnector { let tun = create_tun_off_runtime(interface_name.clone(), Some(remote_addr), local_addr).await?; let local_ip = local_ip.unwrap_or("0.0.0.0".parse().unwrap()); - let mut stack = stack::Stack::new(tun, local_ip, local_ip6, mac); + let stack = stack::Stack::new(tun, local_ip, local_ip6, mac); let driver_type = stack.driver_type(); let socket = stack - .alloc_established_socket(local_addr, remote_addr, stack::State::SynSent) - .await; + .try_alloc_established_socket(local_addr, remote_addr, stack::State::SynSent) + .ok_or(TunnelError::InternalError( + "FakeTCP stack closed while allocating socket".into(), + ))?; let os_stream = os_socket.connect(remote_addr).await?; diff --git a/easytier/src/tunnel/fake_tcp/stack.rs b/easytier/src/tunnel/fake_tcp/stack.rs index 6866c926..a7f1b779 100644 --- a/easytier/src/tunnel/fake_tcp/stack.rs +++ b/easytier/src/tunnel/fake_tcp/stack.rs @@ -54,7 +54,7 @@ use std::sync::{ use tokio::sync::broadcast; use tokio::time; use tokio_util::task::AbortOnDropHandle; -use tracing::{info, trace, warn}; +use tracing::{error, info, trace, warn}; const TIMEOUT: time::Duration = time::Duration::from_secs(1); const RETRIES: usize = 6; @@ -83,13 +83,33 @@ impl AddrTuple { } } +#[derive(Default)] +struct StackState { + tuples: HashMap>, + closed: bool, +} + struct Shared { - tuples: RwLock>>, + state: RwLock, listening: RwLock>, tun: Arc, tuples_purge: broadcast::Sender, } +impl Shared { + fn is_closed(&self) -> bool { + self.state.read().unwrap().closed + } + + fn mark_closed_and_clear_tuples(&self) -> usize { + let mut state = self.state.write().unwrap(); + state.closed = true; + let len = state.tuples.len(); + state.tuples.clear(); + len + } +} + pub struct Stack { shared: Arc, local_ip: Ipv4Addr, @@ -353,7 +373,17 @@ impl Drop for Socket { fn drop(&mut self) { let tuple = AddrTuple::new(self.local_addr, self.remote_addr); // dissociates ourself from the dispatch map - assert!(self.shared.tuples.write().unwrap().remove(&tuple).is_some()); + let (removed, closed) = { + let mut state = self.shared.state.write().unwrap(); + (state.tuples.remove(&tuple).is_some(), state.closed) + }; + if !removed { + if closed { + trace!(?tuple, "Fake TCP tuple already removed after stack closed"); + } else { + warn!(?tuple, "Fake TCP tuple missing while dropping socket"); + } + } // purge cache let _ = self.shared.tuples_purge.send(tuple); @@ -400,7 +430,7 @@ impl Stack { ) -> Stack { let (tuples_purge_tx, _tuples_purge_rx) = broadcast::channel(16); let shared = Arc::new(Shared { - tuples: RwLock::new(HashMap::new()), + state: RwLock::new(StackState::default()), tun: tun.clone(), listening: RwLock::new(HashSet::new()), tuples_purge: tuples_purge_tx.clone(), @@ -426,19 +456,31 @@ impl Stack { self.shared.tun.driver_type() } + pub fn is_closed(&self) -> bool { + self.shared.is_closed() || self.reader_task.is_finished() + } + /// Listens for incoming connections on the given `port`. pub fn listen(&mut self, port: u16) { assert!(self.shared.listening.write().unwrap().insert(port)); } - pub async fn alloc_established_socket( - &mut self, + pub fn try_alloc_established_socket( + &self, local_addr: SocketAddr, remote_addr: SocketAddr, state: State, - ) -> Socket { + ) -> Option { let tuple = AddrTuple::new(local_addr, remote_addr); - let mut tuples = self.shared.tuples.write().unwrap(); + let mut stack_state = self.shared.state.write().unwrap(); + if stack_state.closed || self.reader_task.is_finished() { + stack_state.closed = true; + warn!( + ?tuple, + "fake_tcp stack is closed, refusing to allocate socket" + ); + return None; + } let (sock, incoming) = Socket::new( self.shared.clone(), // self.shared.tun.choose(&mut rng).unwrap().clone(), @@ -450,8 +492,8 @@ impl Stack { Some(0), // Initial ACK state, ); - assert!(tuples.insert(tuple, incoming).is_none()); - sock + assert!(stack_state.tuples.insert(tuple, incoming).is_none()); + Some(sock) } async fn reader_task( @@ -466,7 +508,22 @@ impl Stack { tokio::select! { size = tun.recv(&mut buf) => { - let size = size.unwrap(); + let size = match size { + Ok(size) => size, + Err(e) => { + let shared_tuple_count = shared.mark_closed_and_clear_tuples(); + let cached_tuple_count = tuples.len(); + tuples.clear(); + error!( + ?e, + driver_type = tun.driver_type(), + shared_tuple_count, + cached_tuple_count, + "fake_tcp tun recv failed, reader_task exiting" + ); + break; + } + }; tracing::trace!(len = size, ?buf, "PnetTun received packet"); let buf = buf.split().freeze(); @@ -494,8 +551,8 @@ impl Stack { } else { trace!("Cache miss, checking the shared tuples table for connection"); let sender = { - let tuples = shared.tuples.read().unwrap(); - tuples.get(&tuple).cloned() + let state = shared.state.read().unwrap(); + state.tuples.get(&tuple).cloned() }; if let Some(c) = sender { @@ -532,11 +589,107 @@ impl Stack { } }, tuple = tuples_purge.recv() => { - let tuple = tuple.unwrap(); - tuples.remove(&tuple); - trace!("Removed cached tuple: {:?}", tuple); + match tuple { + Ok(tuple) => { + tuples.remove(&tuple); + trace!("Removed cached tuple: {:?}", tuple); + } + Err(broadcast::error::RecvError::Lagged(skipped)) => { + let cached_tuple_count = tuples.len(); + tuples.clear(); + warn!( + skipped, + cached_tuple_count, + "fake_tcp tuples purge receiver lagged, cleared local cache" + ); + } + Err(broadcast::error::RecvError::Closed) => { + let shared_tuple_count = shared.mark_closed_and_clear_tuples(); + let cached_tuple_count = tuples.len(); + tuples.clear(); + warn!( + shared_tuple_count, + cached_tuple_count, + "fake_tcp tuples purge channel closed, reader_task exiting" + ); + break; + } + } } } } } } + +#[cfg(test)] +mod tests { + use super::*; + use std::io; + use tokio::{ + sync::Notify, + time::{Duration, timeout}, + }; + + #[derive(Default)] + struct FailingTun { + fail: Notify, + } + + impl FailingTun { + fn fail(&self) { + self.fail.notify_one(); + } + } + + #[async_trait::async_trait] + impl Tun for FailingTun { + async fn recv(&self, _packet: &mut BytesMut) -> Result { + self.fail.notified().await; + Err(io::Error::new(io::ErrorKind::BrokenPipe, "test tun closed")) + } + + fn try_send(&self, _packet: &Bytes) -> Result<(), io::Error> { + Ok(()) + } + + fn driver_type(&self) -> &'static str { + "test" + } + } + + #[tokio::test] + async fn reader_task_closes_sockets_on_tun_recv_error() { + let tun = Arc::new(FailingTun::default()); + let mut stack = Stack::new(tun.clone(), Ipv4Addr::LOCALHOST, None, None); + let socket = stack + .try_alloc_established_socket( + SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 10_000), + SocketAddr::new(Ipv4Addr::new(192, 0, 2, 1).into(), 20_000), + State::Established, + ) + .expect("socket allocation should succeed before tun failure"); + + tun.fail(); + + let join_result = timeout(Duration::from_secs(1), &mut stack.reader_task) + .await + .expect("reader task should exit after tun recv error"); + assert!(join_result.is_ok()); + assert!(stack.is_closed()); + + let mut buf = BytesMut::new(); + let recv_result = timeout(Duration::from_secs(1), socket.recv(&mut buf)) + .await + .expect("socket recv should not hang after reader task exits"); + assert_eq!(recv_result, None); + + let new_socket = stack.try_alloc_established_socket( + SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 10_001), + SocketAddr::new(Ipv4Addr::new(192, 0, 2, 1).into(), 20_001), + State::Established, + ); + assert!(new_socket.is_none()); + + drop(socket); + } +}