mirror of
https://github.com/EasyTier/EasyTier.git
synced 2026-08-06 20:49:46 +00:00
fix(faketcp): handle closed tun reader without panic (#2308)
Handle TUN receive errors by marking the fake TCP stack closed and clearing registered sockets instead of panicking. Refuse new sockets on closed stacks and let listeners recreate stacks when the reader task exits.
This commit is contained in:
@@ -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<tokio::net::TcpListener>,
|
||||
// interface_name -> fake tcp stack
|
||||
stack_map: DashMap<String, Arc<Mutex<stack::Stack>>>,
|
||||
stack_map: DashMap<String, Arc<stack::Stack>>,
|
||||
// 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<Arc<Mutex<stack::Stack>>, TunnelError> {
|
||||
) -> Result<Arc<stack::Stack>, 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<Box<dyn Tunnel>, 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?;
|
||||
|
||||
|
||||
@@ -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<AddrTuple, flume::Sender<Bytes>>,
|
||||
closed: bool,
|
||||
}
|
||||
|
||||
struct Shared {
|
||||
tuples: RwLock<HashMap<AddrTuple, flume::Sender<Bytes>>>,
|
||||
state: RwLock<StackState>,
|
||||
listening: RwLock<HashSet<u16>>,
|
||||
tun: Arc<dyn Tun>,
|
||||
tuples_purge: broadcast::Sender<AddrTuple>,
|
||||
}
|
||||
|
||||
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<Shared>,
|
||||
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<Socket> {
|
||||
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<usize, io::Error> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user