From 1cc71c8d6219cb1efdefa489637244323d60965e Mon Sep 17 00:00:00 2001 From: lbl <1791778603@qq.com> Date: Thu, 20 Aug 2026 23:20:06 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20TCP=20=E5=8F=8C=E5=90=91=E8=BD=AC?= =?UTF-8?q?=E5=8F=91=E5=8D=8A=E5=85=B3=E9=97=AD=E6=97=B6=E4=B8=8D=E5=86=8D?= =?UTF-8?q?=E6=88=AA=E6=96=AD=E5=AF=B9=E7=AB=AF=E6=95=B0=E6=8D=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 问题：tcp_port_mapping::stream_copy 与 tcp_nat::stream_nat 用 tokio::select! 组合两个 io::copy，任一方向先结束就 drop 另一方向。 客户端半关闭（shutdown 写方向）后，服务端的响应数据会被静默截断。 修复：抽出 copy_bidirectional_split，用 tokio::io::join 合并分离的 读写半部后走 copy_bidirectional——任一方向 EOF 时对另一端 shutdown 并继续转发剩余方向，直到双向完成。 测试：test_half_close_no_truncation 用 duplex 流验证半关闭后响应 完整送达（旧实现此处读到空数据）。 --- vnt-core/src/nat/internal_nat/tcp_nat.rs | 15 +++-- vnt-core/src/port_mapping/tcp_port_mapping.rs | 67 +++++++++++++++++-- 2 files changed, 69 insertions(+), 13 deletions(-) diff --git a/vnt-core/src/nat/internal_nat/tcp_nat.rs b/vnt-core/src/nat/internal_nat/tcp_nat.rs index d7c5329..3ec7be8 100644 --- a/vnt-core/src/nat/internal_nat/tcp_nat.rs +++ b/vnt-core/src/nat/internal_nat/tcp_nat.rs @@ -78,8 +78,8 @@ async fn stream_task( } pub(crate) async fn stream_nat( - mut recv_stream: R, - mut send_stream: W, + recv_stream: R, + send_stream: W, addr: A, ) -> anyhow::Result<()> where @@ -89,10 +89,11 @@ where let mut tokio_stream = TcpStream::connect(&addr) .await .with_context(|| format!("error connecting to {:?}", addr))?; - let (mut tcp_r, mut tcp_w) = tokio_stream.split(); - tokio::select! { - _ = tokio::io::copy(&mut recv_stream, &mut tcp_w) => {}, - _ = tokio::io::copy(&mut tcp_r, &mut send_stream) => {}, - } + crate::port_mapping::tcp_port_mapping::copy_bidirectional_split( + &mut tokio_stream, + recv_stream, + send_stream, + ) + .await?; Ok(()) } diff --git a/vnt-core/src/port_mapping/tcp_port_mapping.rs b/vnt-core/src/port_mapping/tcp_port_mapping.rs index 518a16e..f48f2ea 100644 --- a/vnt-core/src/port_mapping/tcp_port_mapping.rs +++ b/vnt-core/src/port_mapping/tcp_port_mapping.rs @@ -7,8 +7,26 @@ use crate::utils::task_control::TaskGroup; use anyhow::Context; use pnet_packet::ip::IpNextHeaderProtocols; use std::net::{Ipv4Addr, SocketAddr}; +use tokio::io::{AsyncRead, AsyncWrite}; use tokio::net::{TcpListener, TcpStream}; +/// 双向转发：任一方向 EOF 时对另一端执行 shutdown 并继续转发剩余方向， +/// 直到两个方向都完成。避免 select! 下任一方向先结束就 drop 另一方向 +/// 造成的半关闭截断（如对端半关闭后响应数据丢失）。 +pub(crate) async fn copy_bidirectional_split( + stream: &mut T, + reader: R, + writer: W, +) -> anyhow::Result<(u64, u64)> +where + T: AsyncRead + AsyncWrite + Unpin, + R: AsyncRead + Unpin, + W: AsyncWrite + Unpin, +{ + let mut other = tokio::io::join(reader, writer); + Ok(tokio::io::copy_bidirectional(stream, &mut other).await?) +} + pub async fn start( task_group: &TaskGroup, list: &Vec, @@ -64,7 +82,7 @@ async fn stream_copy( dst_port: u16, quic_tunnel_client: QuicTunnelClient, ) -> anyhow::Result<()> { - let (mut send_stream, mut recv_stream) = quic_tunnel_client.open_bi(target_ip).await?; + let (mut send_stream, recv_stream) = quic_tunnel_client.open_bi(target_ip).await?; let handshake = QuicProxyHandshake { handshake: Some(quic_proxy_handshake::Handshake::TcpPortMapping( PortProxyHandshake { @@ -76,10 +94,47 @@ async fn stream_copy( )), }; send_handshake(&mut send_stream, handshake).await?; - let (mut tcp_r, mut tcp_w) = tcp_stream.split(); - tokio::select! { - _ = tokio::io::copy(&mut recv_stream, &mut tcp_w) => {}, - _ = tokio::io::copy(&mut tcp_r, &mut send_stream) => {}, - } + copy_bidirectional_split(&mut tcp_stream, recv_stream, send_stream).await?; Ok(()) } + + +#[cfg(test)] +mod tests { + use super::*; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + /// 半关闭场景：客户端发完请求后 shutdown 写方向， + /// 服务端的响应必须完整送达，不能被截断。 + #[tokio::test] + async fn test_half_close_no_truncation() { + let (mut client, mut relay_tcp) = tokio::io::duplex(64); + let (relay_tunnel, mut server) = tokio::io::duplex(64); + let (relay_tunnel_r, relay_tunnel_w) = tokio::io::split(relay_tunnel); + + let relay = tokio::spawn(async move { + copy_bidirectional_split(&mut relay_tcp, relay_tunnel_r, relay_tunnel_w) + .await + .unwrap(); + }); + + // 客户端发请求后立即半关闭写方向 + client.write_all(b"ping").await.unwrap(); + client.shutdown().await.unwrap(); + + // 服务端读到完整请求（读到 EOF 前数据不能丢） + let mut buf = [0u8; 4]; + server.read_exact(&mut buf).await.unwrap(); + assert_eq!(&buf, b"ping"); + // 服务端回响应并关闭 + server.write_all(b"pong").await.unwrap(); + server.shutdown().await.unwrap(); + + // 客户端必须能读到完整响应（旧实现此处会被截断为空） + let mut out = Vec::new(); + client.read_to_end(&mut out).await.unwrap(); + assert_eq!(out, b"pong"); + + relay.await.unwrap(); + } +}