src/kcp_stream.rs
Rama master · réplica de lectura del repositorio original.
use hbb_common::{
anyhow,
bytes::{Bytes, BytesMut},
bytes_codec::BytesCodec,
config, log,
tcp::{DynTcpStream, FramedStream},
tokio::{self, net::UdpSocket, sync::mpsc, sync::oneshot},
tokio_util, ResultType, Stream,
};
use kcp_sys::{
endpoint::{ConnId, KcpEndpoint},
packet_def::{KcpPacket, KcpPacketHeader},
stream,
};
use std::{net::SocketAddr, sync::Arc};
pub struct KcpStream {
endpoint: KcpEndpoint,
conn_id: ConnId,
stop_sender: Option<oneshot::Sender<()>>,
}
const KCP_IO_ERR_LOG_INTERVAL: std::time::Duration = std::time::Duration::from_secs(5);
static KCP_SEND_ERR_LOG: hbb_common::log_throttle::LogThrottle =
hbb_common::log_throttle::LogThrottle::new(KCP_IO_ERR_LOG_INTERVAL);
static KCP_RECV_ERR_LOG: hbb_common::log_throttle::LogThrottle =
hbb_common::log_throttle::LogThrottle::new(KCP_IO_ERR_LOG_INTERVAL);
impl KcpStream {
// Opt in to KCP's built-in congestion window (nc=0) instead of the pure turbo profile
// (nc=1) that has always shipped; see `get_kcp_cc_enabled` for why this is not the default.
// Sender-side only, so no wire negotiation is needed and either peer may run either profile.
// Requires kcp-sys from the `rustdesk-patches` branch, which wires the config factory into
// connection setup (on older revs the factory was stored but never consulted).
fn apply_kcp_config(endpoint: &mut KcpEndpoint) {
if crate::get_kcp_cc_enabled() {
endpoint.set_kcp_config_factory(Box::new(|conv| {
let mut config = kcp_sys::ffi_safe::KcpConfig::new_turbo(conv);
config.nc = Some(0);
config
}));
}
}
/// How long since a valid packet was last received from the peer, or `None` once the
/// connection is gone. Answered by the KCP endpoint's own tasks, not by the session's read
/// loop, so it stays meaningful while that loop is busy sending a large message; and the
/// endpoint pings an idle peer often enough that silence here means the peer, not quiet.
pub fn peer_silent_for(&self) -> Option<std::time::Duration> {
self.endpoint.peer_silent_for(&self.conn_id)
}
fn create_framed(stream: stream::KcpStream, local_addr: Option<SocketAddr>) -> Stream {
Stream::Tcp(FramedStream(
tokio_util::codec::Framed::new(DynTcpStream(Box::new(stream)), BytesCodec::new()),
local_addr.unwrap_or(config::Config::get_any_listen_addr(true)),
None,
0,
))
}
pub async fn accept(
udp_socket: Arc<UdpSocket>,
timeout: std::time::Duration,
init_packet: Option<BytesMut>,
) -> ResultType<(Self, Stream)> {
let mut endpoint = KcpEndpoint::new();
Self::apply_kcp_config(&mut endpoint);
endpoint.run().await;
let (input, output) = (
endpoint.input_sender(),
endpoint
.output_receiver()
.ok_or_else(|| anyhow::anyhow!("Failed to get output receiver"))?,
);
let (stop_sender, stop_receiver) = oneshot::channel();
if let Some(packet) = init_packet {
if packet.len() >= std::mem::size_of::<KcpPacketHeader>() {
input.send(packet.into()).await?;
}
}
Self::kcp_io(udp_socket.clone(), input, output, stop_receiver).await;
let conn_id = tokio::time::timeout(timeout, endpoint.accept()).await??;
if let Some(stream) = stream::KcpStream::new(&endpoint, conn_id) {
Ok((
Self {
endpoint,
conn_id,
stop_sender: Some(stop_sender),
},
Self::create_framed(stream, udp_socket.local_addr().ok()),
))
} else {
Err(anyhow::anyhow!("Failed to create KcpStream"))
}
}
pub async fn connect(
udp_socket: Arc<UdpSocket>,
timeout: std::time::Duration,
) -> ResultType<(Self, Stream)> {
let mut endpoint = KcpEndpoint::new();
Self::apply_kcp_config(&mut endpoint);
endpoint.run().await;
let (input, output) = (
endpoint.input_sender(),
endpoint
.output_receiver()
.ok_or_else(|| anyhow::anyhow!("Failed to get output receiver"))?,
);
let (stop_sender, stop_receiver) = oneshot::channel();
Self::kcp_io(udp_socket.clone(), input, output, stop_receiver).await;
let conn_id = endpoint.connect(timeout, 0, 0, Bytes::new()).await?;
if let Some(stream) = stream::KcpStream::new(&endpoint, conn_id) {
Ok((
Self {
endpoint,
conn_id,
stop_sender: Some(stop_sender),
},
Self::create_framed(stream, udp_socket.local_addr().ok()),
))
} else {
Err(anyhow::anyhow!("Failed to create KcpStream"))
}
}
async fn kcp_io(
udp_socket: Arc<UdpSocket>,
input: mpsc::Sender<KcpPacket>,
mut output: mpsc::Receiver<KcpPacket>,
mut stop_receiver: oneshot::Receiver<()>,
) {
let udp = udp_socket.clone();
tokio::spawn(async move {
let mut buf = vec![0; 1500];
// Socket errors are ICMP unreachable on a connected UDP socket — advisory, and
// routine while a hole forms — so treat them as loss and let KCP's pong timeout reap
// a link that is really dead. One throttle PER DIRECTION: the error is reported once
// and cleared, so send-ok/recv-err alternates and a shared counter never fires.
loop {
tokio::select! {
_ = &mut stop_receiver => {
log::debug!("KCP io loop received stop signal");
break;
}
Some(data) = output.recv() => {
if let Err(e) = udp.send(&data.inner()).await {
if let Some(n) = KCP_SEND_ERR_LOG.due() {
log::debug!("KCP send error x{n} (treated as loss), last: {e}");
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
}
result = udp.recv_from(&mut buf) => {
match result {
Ok((size, _)) => {
if size < std::mem::size_of::<KcpPacketHeader>() {
continue;
}
input
.send(BytesMut::from(&buf[..size]).into())
.await.ok();
}
Err(e) => {
if let Some(n) = KCP_RECV_ERR_LOG.due() {
log::debug!("KCP recv error x{n} (treated as loss), last: {e}");
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
}
}
else => {
log::debug!("KCP endpoint input closed");
break;
}
}
}
});
}
}
impl Drop for KcpStream {
fn drop(&mut self) {
if let Some(sender) = self.stop_sender.take() {
let _ = sender.send(());
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
async fn connected_pair() -> (Arc<UdpSocket>, Arc<UdpSocket>) {
let a = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let b = UdpSocket::bind("127.0.0.1:0").await.unwrap();
a.connect(b.local_addr().unwrap()).await.unwrap();
b.connect(a.local_addr().unwrap()).await.unwrap();
(Arc::new(a), Arc::new(b))
}
async fn establish() -> ((KcpStream, Stream), (KcpStream, Stream)) {
let (a, b) = connected_pair().await;
let (accept_res, connect_res) = tokio::join!(
KcpStream::accept(b, Duration::from_secs(5), None),
KcpStream::connect(a, Duration::from_secs(5))
);
(
connect_res.expect("connect over loopback"),
accept_res.expect("accept over loopback"),
)
}
// The full client path over real loopback sockets: handshake through the kcp_io
// pumps, framed data both ways, then a graceful close. The endpoint guard stays
// alive across the stream drop so the FIN can go out, and the peer's framed
// stream must end (BrokenPipe from the kcp reader) instead of hanging.
#[tokio::test]
async fn test_kcp_stream_loopback_roundtrip_and_close() {
let ((_guard_a, mut stream_a), (_guard_b, mut stream_b)) = establish().await;
stream_a
.send_bytes(Bytes::from_static(b"ping"))
.await
.unwrap();
let got = stream_b.next_timeout(5000).await.unwrap().unwrap();
assert_eq!(&got[..], b"ping");
stream_b
.send_bytes(Bytes::from_static(b"pong"))
.await
.unwrap();
let got = stream_a.next_timeout(5000).await.unwrap().unwrap();
assert_eq!(&got[..], b"pong");
drop(stream_a);
match stream_b.next_timeout(10_000).await {
None | Some(Err(_)) => {}
Some(Ok(data)) => panic!("unexpected data after close: {:?}", data),
}
}
// A writer that queues many frames and closes immediately must not cost the
// reader any of them: every frame arrives intact, in order, before end-of-stream.
// This is the client-side pin for the kcp-sys close-tail-drain semantics, through
// the real BytesCodec framing rustdesk sessions use.
#[tokio::test]
async fn test_kcp_stream_close_delivers_all_frames() {
let ((_guard_a, mut tx), (_guard_b, mut rx)) = establish().await;
const N: usize = 50;
let payload = vec![7u8; 32 * 1024];
for _ in 0..N {
tx.send_bytes(Bytes::from(payload.clone())).await.unwrap();
}
drop(tx);
let mut got = 0usize;
loop {
match rx.next_timeout(10_000).await {
Some(Ok(data)) => {
assert_eq!(data.len(), payload.len(), "frame boundary broken");
assert!(data.iter().all(|&b| b == 7), "frame content corrupted");
got += 1;
}
// BrokenPipe (kcp reader end) or timeout-None both end the stream.
None | Some(Err(_)) => break,
}
}
assert_eq!(got, N, "graceful close lost frames");
}
}