fix: TCP 双向转发半关闭时不再截断对端数据
问题: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 流验证半关闭后响应 完整送达(旧实现此处读到空数据)。
This commit is contained in:
@@ -78,8 +78,8 @@ async fn stream_task(
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn stream_nat<R, W, A: ToSocketAddrs + Debug>(
|
pub(crate) async fn stream_nat<R, W, A: ToSocketAddrs + Debug>(
|
||||||
mut recv_stream: R,
|
recv_stream: R,
|
||||||
mut send_stream: W,
|
send_stream: W,
|
||||||
addr: A,
|
addr: A,
|
||||||
) -> anyhow::Result<()>
|
) -> anyhow::Result<()>
|
||||||
where
|
where
|
||||||
@@ -89,10 +89,11 @@ where
|
|||||||
let mut tokio_stream = TcpStream::connect(&addr)
|
let mut tokio_stream = TcpStream::connect(&addr)
|
||||||
.await
|
.await
|
||||||
.with_context(|| format!("error connecting to {:?}", addr))?;
|
.with_context(|| format!("error connecting to {:?}", addr))?;
|
||||||
let (mut tcp_r, mut tcp_w) = tokio_stream.split();
|
crate::port_mapping::tcp_port_mapping::copy_bidirectional_split(
|
||||||
tokio::select! {
|
&mut tokio_stream,
|
||||||
_ = tokio::io::copy(&mut recv_stream, &mut tcp_w) => {},
|
recv_stream,
|
||||||
_ = tokio::io::copy(&mut tcp_r, &mut send_stream) => {},
|
send_stream,
|
||||||
}
|
)
|
||||||
|
.await?;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -7,8 +7,26 @@ use crate::utils::task_control::TaskGroup;
|
|||||||
use anyhow::Context;
|
use anyhow::Context;
|
||||||
use pnet_packet::ip::IpNextHeaderProtocols;
|
use pnet_packet::ip::IpNextHeaderProtocols;
|
||||||
use std::net::{Ipv4Addr, SocketAddr};
|
use std::net::{Ipv4Addr, SocketAddr};
|
||||||
|
use tokio::io::{AsyncRead, AsyncWrite};
|
||||||
use tokio::net::{TcpListener, TcpStream};
|
use tokio::net::{TcpListener, TcpStream};
|
||||||
|
|
||||||
|
/// 双向转发:任一方向 EOF 时对另一端执行 shutdown 并继续转发剩余方向,
|
||||||
|
/// 直到两个方向都完成。避免 select! 下任一方向先结束就 drop 另一方向
|
||||||
|
/// 造成的半关闭截断(如对端半关闭后响应数据丢失)。
|
||||||
|
pub(crate) async fn copy_bidirectional_split<T, R, W>(
|
||||||
|
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(
|
pub async fn start(
|
||||||
task_group: &TaskGroup,
|
task_group: &TaskGroup,
|
||||||
list: &Vec<PortMapping>,
|
list: &Vec<PortMapping>,
|
||||||
@@ -64,7 +82,7 @@ async fn stream_copy(
|
|||||||
dst_port: u16,
|
dst_port: u16,
|
||||||
quic_tunnel_client: QuicTunnelClient,
|
quic_tunnel_client: QuicTunnelClient,
|
||||||
) -> anyhow::Result<()> {
|
) -> 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 {
|
let handshake = QuicProxyHandshake {
|
||||||
handshake: Some(quic_proxy_handshake::Handshake::TcpPortMapping(
|
handshake: Some(quic_proxy_handshake::Handshake::TcpPortMapping(
|
||||||
PortProxyHandshake {
|
PortProxyHandshake {
|
||||||
@@ -76,10 +94,47 @@ async fn stream_copy(
|
|||||||
)),
|
)),
|
||||||
};
|
};
|
||||||
send_handshake(&mut send_stream, handshake).await?;
|
send_handshake(&mut send_stream, handshake).await?;
|
||||||
let (mut tcp_r, mut tcp_w) = tcp_stream.split();
|
copy_bidirectional_split(&mut tcp_stream, recv_stream, send_stream).await?;
|
||||||
tokio::select! {
|
|
||||||
_ = tokio::io::copy(&mut recv_stream, &mut tcp_w) => {},
|
|
||||||
_ = tokio::io::copy(&mut tcp_r, &mut send_stream) => {},
|
|
||||||
}
|
|
||||||
Ok(())
|
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();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user