fix(port_mapping): UDP 映射 dest_map 不再无界增长,recv 出错不再杀任务

问题:
1. 每个新 src 创建映射条目与转发任务,任务虽有 60s 空闲超时退出,
   但映射条目从不移除,dest_map 随不同 src 数量无界增长。
2. recv_from 单次出错经 ? 传播直接杀死整个端口映射接收任务。

修复:
- 转发任务退出时经 remove_if_same(same_channel 身份比较)回收
  映射条目,防止误删重建后的新条目。
- recv_from 错误改为 log + 100ms 休眠 + continue。

测试:test_remove_if_same 覆盖旧任务误删新条目的竞态。
This commit is contained in:
lbl
2026-08-21 02:39:24 +08:00
parent 55d4b874bd
commit 90a00215ca
+59 -1
View File
@@ -53,7 +53,16 @@ async fn recv(
let mut buf = vec![0u8; 65536];
let dest_map = Arc::new(Mutex::new(HashMap::<SocketAddr, Sender<Bytes>>::new()));
loop {
let (len, src) = udp_socket.recv_from(&mut buf).await?;
// 单次接收错误不能杀掉整个端口映射任务:记录日志后继续,
// 短暂休眠避免持续性错误造成空转
let (len, src) = match udp_socket.recv_from(&mut buf).await {
Ok(v) => v,
Err(e) => {
log::warn!("udp port mapping recv error: {e:?}");
tokio::time::sleep(Duration::from_millis(100)).await;
continue;
}
};
let bytes = Bytes::copy_from_slice(&buf[..len]);
let tx = {
@@ -67,6 +76,8 @@ async fn recv(
let dst_host = mapping.dst_host.clone();
let dst_port = mapping.dst_port;
let tunnel_client = quic_tunnel_client.clone();
let tx_clone = tx.clone();
let dest_map_clone = dest_map.clone();
task_group.spawn(async move {
if let Err(e) = udp_mapping_handle(
udp_socket,
@@ -81,6 +92,9 @@ async fn recv(
{
log::error!("udp_mapping_handle {e:?},src:{src}");
}
// 任务退出(含 60s 空闲超时)时回收映射条目,
// 否则 dest_map 随不同 src 数量无界增长
remove_if_same(&mut dest_map_clone.lock(), &src, &tx_clone);
});
map.insert(src, tx.clone());
@@ -100,6 +114,22 @@ async fn recv(
}
}
/// 仅当映射仍指向同一个 channel(即仍是本任务的条目)时才移除,
/// 避免条目被回收重建后,旧任务误删新条目
fn remove_if_same(
map: &mut HashMap<SocketAddr, Sender<Bytes>>,
src: &SocketAddr,
tx: &Sender<Bytes>,
) -> bool {
if let Some(cur) = map.get(src)
&& cur.same_channel(tx)
{
map.remove(src);
return true;
}
false
}
async fn udp_mapping_handle(
udp_socket: Arc<UdpSocket>,
src: SocketAddr,
@@ -143,3 +173,31 @@ async fn udp_mapping_handle(
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
/// 空闲回收竞态:条目被重建后,旧任务退出不得误删新条目
#[test]
fn test_remove_if_same() {
let src: SocketAddr = "10.0.0.1:5000".parse().unwrap();
let mut map = HashMap::new();
let (old_tx, _old_rx) = tokio::sync::mpsc::channel::<Bytes>(1);
map.insert(src, old_tx.clone());
// 条目被重建为新 channel
let (new_tx, _new_rx) = tokio::sync::mpsc::channel::<Bytes>(1);
map.insert(src, new_tx);
// 旧任务退出:不允许删除新条目
assert!(!remove_if_same(&mut map, &src, &old_tx));
assert!(map.contains_key(&src));
// 新任务退出:允许删除
let new_tx = map.get(&src).unwrap().clone();
assert!(remove_if_same(&mut map, &src, &new_tx));
assert!(!map.contains_key(&src));
}
}