use std::{
collections::HashMap,
net::SocketAddr,
pin::Pin,
sync::{Arc, OnceLock},
task::{Context, Poll},
time::Duration,
};
use anyhow::{anyhow, bail, Context as _, Result};
use nfrp_common::{
config::{ClientConfig, ProxyConfig},
frp::{
kcp::{self, KcpStream},
msg::{constant_time_eq, NatHoleClient, NatHoleResp, NatHoleVisitor},
},
p2p::{self, Packet, Role, Transport, SERVER_NAME},
};
use quinn::{Connection, Endpoint, EndpointConfig, RecvStream, SendStream, TokioRuntime};
use tokio::{
io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf},
net::UdpSocket,
time::Instant,
};
use tracing::{debug, info, warn};
const PUNCH_TIMEOUT: Duration = Duration::from_secs(8);
const PUNCH_INTERVAL: Duration = Duration::from_millis(50);
const HELLO_INTERVAL: Duration = Duration::from_millis(200);
const TOKEN_LEN: usize = 64;
const KEEP_ALIVE: Duration = Duration::from_secs(5);
const IDLE_TIMEOUT: Duration = Duration::from_secs(30);
const PREDICT_SAMPLES: usize = 3;
const CANDIDATE_STAGGER: Duration = Duration::from_millis(60);
const KCP_PUNCH_ROUNDS: u32 = 6;
pub async fn connect_as_visitor(
server: &SocketAddr,
sid: &str,
secret_key: &str,
transport: Transport,
window: u16,
) -> Result<P2PStream> {
let sock = UdpSocket::bind(any_addr(server))
.await
.context("绑定 P2P 本地端口失败")?;
let rd = rendezvous(&sock, server, Role::Visitor, sid, transport, window).await?;
debug!(peer = ?rd.addrs, %sid, transport = ?rd.transport, "牵线完成,visitor 开始直连");
let raw = sock.into_std().context("转换 UDP socket 失败")?;
let stream = match rd.transport {
Transport::Quic => {
let endpoint =
Endpoint::new(EndpointConfig::default(), None, raw, Arc::new(TokioRuntime))
.context("创建 QUIC endpoint 失败")?;
let conn = quic_connect_any(&endpoint, &rd.candidates, sid).await?;
let (send, recv) = conn.open_bi().await.map_err(quic_err("打开双向流"))?;
let mut s = P2PStream::Quic(QuicStream {
send,
recv,
_endpoint: endpoint,
});
greet_client(&mut s, sid, secret_key).await?;
s
}
Transport::Kcp => {
let sock = UdpSocket::from_std(raw).context("包装 KCP socket 失败")?;
let mut s = P2PStream::Kcp(kcp_connect_any(sock, &rd.candidates, sid).await?);
greet_client(&mut s, sid, secret_key).await?;
s
}
};
info!(peer = ?rd.addrs, %sid, transport = ?rd.transport, "xtcp P2P 直连已建立(数据不再经过服务端)");
Ok(stream)
}
pub async fn accept_as_provider(
server: &SocketAddr,
sid: &str,
secret_key: &str,
window: u16,
) -> Result<P2PStream> {
let sock = UdpSocket::bind(any_addr(server))
.await
.context("绑定 P2P 本地端口失败")?;
let rd = rendezvous(&sock, server, Role::Provider, sid, Transport::Quic, window).await?;
debug!(peer = ?rd.addrs, %sid, transport = ?rd.transport, "牵线完成,provider 开始打洞并等待入站");
let expect_ip = rd.addrs.first().map(|a| a.ip());
let raw = sock.into_std().context("转换 UDP socket 失败")?;
let stream = match rd.transport {
Transport::Quic => {
let puncher = UdpSocket::from_std(raw.try_clone().context("克隆 UDP socket 失败")?)
.context("包装打洞 socket 失败")?;
let endpoint = Endpoint::new(
EndpointConfig::default(),
Some(server_config()?),
raw,
Arc::new(TokioRuntime),
)
.context("创建 QUIC endpoint 失败")?;
let deadline = Instant::now() + PUNCH_TIMEOUT;
let punching = tokio::spawn(punch_loop(puncher, rd.candidates.clone(), deadline));
let accepted = quic_accept(&endpoint, expect_ip, deadline).await;
punching.abort();
let conn = accepted?;
let (send, recv) = conn
.accept_bi()
.await
.map_err(quic_err("等待 QUIC 双向流"))?;
let mut s = P2PStream::Quic(QuicStream {
send,
recv,
_endpoint: endpoint,
});
greet_server(&mut s, sid, secret_key).await?;
s
}
Transport::Kcp => {
let sock = UdpSocket::from_std(raw).context("包装 KCP socket 失败")?;
let s = kcp_accept(sock, &rd.candidates, expect_ip, sid, PUNCH_TIMEOUT).await?;
let mut s = P2PStream::Kcp(s);
greet_server(&mut s, sid, secret_key).await?;
s
}
};
info!(peer = ?rd.addrs, %sid, transport = ?rd.transport, "xtcp P2P 直连已建立(provider 侧)");
Ok(stream)
}
fn quic_err(what: &str) -> impl Fn(quinn::ConnectionError) -> anyhow::Error + '_ {
move |e| anyhow!("{what}失败:{e}")
}
pub struct PunchBus {
req_tx: tokio::sync::mpsc::UnboundedSender<NatHoleVisitor>,
waiters: std::sync::Mutex<HashMap<String, tokio::sync::oneshot::Sender<NatHoleResp>>>,
}
impl PunchBus {
fn new() -> (
Arc<Self>,
tokio::sync::mpsc::UnboundedReceiver<NatHoleVisitor>,
) {
let (req_tx, req_rx) = tokio::sync::mpsc::unbounded_channel();
(
Arc::new(Self {
req_tx,
waiters: std::sync::Mutex::new(HashMap::new()),
}),
req_rx,
)
}
pub fn request(
&self,
m: NatHoleVisitor,
) -> Result<tokio::sync::oneshot::Receiver<NatHoleResp>> {
let (tx, rx) = tokio::sync::oneshot::channel();
let txn = m.transaction_id.clone();
self.lock().insert(txn.clone(), tx);
if let Err(e) = self.req_tx.send(m) {
self.lock().remove(&txn);
return Err(anyhow!("控制连接已关闭,无法发起打洞:{e}"));
}
Ok(rx)
}
pub fn deliver(&self, resp: NatHoleResp) {
let waiter = self.lock().remove(&resp.transaction_id);
if let Some(tx) = waiter {
let _ = tx.send(resp);
} else {
debug!(txn = %resp.transaction_id, "收到无人等待的 NatHoleResp,丢弃");
}
}
pub fn clear(&self) {
self.lock().clear();
}
fn lock(
&self,
) -> std::sync::MutexGuard<'_, HashMap<String, tokio::sync::oneshot::Sender<NatHoleResp>>> {
self.waiters.lock().unwrap_or_else(|e| e.into_inner())
}
}
pub struct P2PRoute {
pub server_udp: SocketAddr,
pub bus: Arc<PunchBus>,
pub transport: Transport,
pub predict_window: u16,
}
pub async fn setup(
cfg: &ClientConfig,
) -> Option<(
Arc<P2PRoute>,
tokio::sync::mpsc::UnboundedReceiver<NatHoleVisitor>,
)> {
let port = cfg.p2p_port?;
if !cfg.p2p_enable {
return None;
}
let addr = nfrp_common::util::resolve_addr(&format!("{}:{}", cfg.server_addr, port))
.await
.ok()?;
let (bus, rx) = PunchBus::new();
Some((
Arc::new(P2PRoute {
server_udp: addr,
bus,
transport: parse_transport(&cfg.xtcp_transport),
predict_window: if cfg.xtcp_port_predict {
cfg.xtcp_predict_window.max(1)
} else {
0
},
}),
rx,
))
}
fn parse_transport(s: &str) -> Transport {
match s.trim().to_ascii_lowercase().as_str() {
"kcp" => Transport::Kcp,
_ => Transport::Quic,
}
}
pub async fn serve_as_provider(
cfg: Arc<ClientConfig>,
proxy: ProxyConfig,
m: NatHoleClient,
) -> Result<()> {
let Some(port) = cfg.p2p_port else {
bail!("收到打洞请求但未配置 p2p_port,忽略");
};
let server_udp = nfrp_common::util::resolve_addr(&format!("{}:{}", cfg.server_addr, port))
.await
.with_context(|| format!("解析牵线地址 {}:{} 失败", cfg.server_addr, port))?;
let local_addr = nfrp_common::util::resolve_addr(&proxy.local_addr)
.await
.with_context(|| format!("解析内网地址 {} 失败", proxy.local_addr))?;
let window = if cfg.xtcp_port_predict {
cfg.xtcp_predict_window.max(1)
} else {
0
};
let stream = accept_as_provider(&server_udp, &m.sid, &proxy.secret_key, window).await?;
debug!(proxy = %m.proxy_name, local = %local_addr, "P2P 已建立,连接内网服务");
let mut local = tokio::net::TcpStream::connect(local_addr)
.await
.with_context(|| format!("连接内网服务 {local_addr} 失败"))?;
local.set_nodelay(true).ok();
let mut stream = stream;
match nfrp_common::util::relay_between(&mut local, &mut stream).await {
Ok((up, down)) => debug!(proxy = %m.proxy_name, "P2P 转发结束:上行 {up}B / 下行 {down}B"),
Err(e) => debug!(proxy = %m.proxy_name, "P2P 转发中断:{e}"),
}
Ok(())
}
pub async fn try_punch_as_visitor(
route: &P2PRoute,
run_id: &str,
target: &str,
secret_key: &str,
) -> Result<P2PStream> {
let txn = nfrp_common::util::new_run_id();
let ts = nfrp_common::util::now_unix_secs() as i64;
let req = NatHoleVisitor {
transaction_id: txn.clone(),
proxy_name: target.to_string(),
sign_key: nfrp_common::frp::msg::auth_key(secret_key, ts),
timestamp: ts,
protocol: match route.transport {
Transport::Quic => "quic".to_string(),
Transport::Kcp => "kcp".to_string(),
},
..Default::default()
};
let _ = run_id;
let rx = route.bus.request(req)?;
let resp = tokio::time::timeout(PUNCH_TIMEOUT, rx)
.await
.map_err(|_| anyhow!("等待服务端牵线响应超时"))?
.map_err(|_| anyhow!("打洞请求被丢弃(控制连接断开了?)"))?;
if !resp.error.is_empty() {
bail!("服务端拒绝打洞:{}", resp.error);
}
if resp.sid.is_empty() {
bail!("服务端未下发 sid(可能未启用 P2P)");
}
connect_as_visitor(
&route.server_udp,
&resp.sid,
secret_key,
route.transport,
route.predict_window,
)
.await
}
struct Rendezvous {
addrs: Vec<SocketAddr>,
candidates: Vec<SocketAddr>,
transport: Transport,
}
async fn rendezvous(
sock: &UdpSocket,
server: &SocketAddr,
role: Role,
sid: &str,
transport: Transport,
window: u16,
) -> Result<Rendezvous> {
let hello =
p2p::encode_hello(role, sid, transport).context("sid 长度不是 32,无法编码 HELLO")?;
let mut probes: Vec<UdpSocket> = Vec::new();
if window > 0 {
for _ in 1..PREDICT_SAMPLES {
match UdpSocket::bind(any_addr(server)).await {
Ok(s) => probes.push(s),
Err(e) => {
debug!("开采样 socket 失败,端口预测样本会少一个:{e}");
break;
}
}
}
}
let deadline = Instant::now() + PUNCH_TIMEOUT;
let mut buf = [0u8; 512];
let mut ticker = tokio::time::interval(HELLO_INTERVAL);
ticker.tick().await;
loop {
let remain = deadline.saturating_duration_since(Instant::now());
if remain.is_zero() {
bail!("等待服务端牵线下发对端地址超时({PUNCH_TIMEOUT:?})");
}
sock.send_to(&hello, server)
.await
.context("发送 HELLO 失败")?;
for p in &probes {
let _ = p.send_to(&hello, server).await;
}
let got = tokio::time::timeout(remain, sock.recv_from(&mut buf)).await;
match got {
Ok(Ok((n, _from))) => {
match p2p::decode(&buf[..n]) {
Some(Packet::Peers {
role: r,
sid: s,
transport: t,
addrs,
}) if s == sid && r == role.peer() => {
return Ok(build_rendezvous(addrs, t, window));
}
Some(Packet::Peer {
role: r,
sid: s,
transport: t,
addr,
}) if s == sid && r == role.peer() => {
return Ok(build_rendezvous(vec![addr], t, window));
}
other => debug!(?other, "忽略非预期的牵线报文"),
}
}
Ok(Err(e)) => return Err(e).context("读取牵线响应失败"),
Err(_) => bail!("等待服务端牵线下发对端地址超时({PUNCH_TIMEOUT:?})"),
}
tokio::time::sleep(HELLO_INTERVAL).await;
}
}
fn build_rendezvous(addrs: Vec<SocketAddr>, transport: Transport, window: u16) -> Rendezvous {
let candidates = if window > 0 {
p2p::expand_candidates(&addrs, window)
} else {
addrs.clone()
};
Rendezvous {
addrs,
candidates,
transport,
}
}
fn any_addr(remote: &SocketAddr) -> SocketAddr {
if remote.is_ipv4() {
"0.0.0.0:0"
} else {
"[::]:0"
}
.parse()
.expect("硬编码的任意地址必然合法")
}
async fn punch_loop(sock: UdpSocket, candidates: Vec<SocketAddr>, deadline: Instant) {
while Instant::now() < deadline {
for c in &candidates {
if sock.send_to(p2p::PUNCH_MAGIC, c).await.is_err() {
return;
}
}
tokio::time::sleep(PUNCH_INTERVAL).await;
}
}
async fn quic_connect_any(
endpoint: &Endpoint,
candidates: &[SocketAddr],
sid: &str,
) -> Result<Connection> {
if candidates.is_empty() {
bail!("牵线没有给出任何候选地址");
}
let deadline = Instant::now() + PUNCH_TIMEOUT;
let cfg = client_config()?;
if candidates.len() == 1 {
let connecting = endpoint
.connect_with(cfg, candidates[0], SERVER_NAME)
.map_err(|e| anyhow!("发起 QUIC 连接失败:{e}"))?;
return tokio::time::timeout_at(deadline, connecting)
.await
.map_err(|_| anyhow!("QUIC 握手超时(打洞未成功)"))?
.map_err(|e| anyhow!("QUIC 握手失败:{e}"));
}
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
for (i, addr) in candidates.iter().enumerate() {
let ep = endpoint.clone();
let cfg = cfg.clone();
let addr = *addr;
let tx = tx.clone();
tokio::spawn(async move {
tokio::time::sleep(CANDIDATE_STAGGER * i as u32).await;
let r = match ep.connect_with(cfg, addr, SERVER_NAME) {
Ok(c) => c.await.map_err(|e| anyhow!("QUIC 握手失败({addr}):{e}")),
Err(e) => Err(anyhow!("发起 QUIC 连接失败({addr}):{e}")),
};
let _ = tx.send((addr, r));
});
}
drop(tx);
let mut last_err = None;
let mut pending = candidates.len();
while pending > 0 {
match tokio::time::timeout_at(deadline, rx.recv()).await {
Ok(Some((addr, Ok(conn)))) => {
debug!(%addr, %sid, "候选端口握手成功");
return Ok(conn);
}
Ok(Some((_addr, Err(e)))) => {
last_err = Some(e);
pending -= 1;
}
Ok(None) => break,
Err(_) => break,
}
}
Err(last_err.unwrap_or_else(|| anyhow!("QUIC 握手超时(打洞未成功)")))
}
async fn quic_accept(
endpoint: &Endpoint,
expect_ip: Option<std::net::IpAddr>,
deadline: Instant,
) -> Result<Connection> {
let incoming = tokio::time::timeout_at(deadline, endpoint.accept())
.await
.map_err(|_| anyhow!("等待入站 P2P 连接超时(打洞未成功)"))?
.ok_or_else(|| anyhow!("QUIC endpoint 已关闭"))?;
let conn = incoming
.accept()
.map_err(|e| anyhow!("接受 QUIC 连接失败:{e}"))?
.await
.map_err(|e| anyhow!("QUIC 握手失败:{e}"))?;
let remote = conn.remote_address();
if let Some(ip) = expect_ip {
if remote.ip() != ip {
conn.close(0u32.into(), b"unexpected peer");
bail!("拒绝来自 {remote} 的 P2P 连接(牵线下发的 IP 是 {ip})");
}
}
Ok(conn)
}
async fn kcp_connect_any(
sock: UdpSocket,
candidates: &[SocketAddr],
sid: &str,
) -> Result<KcpStream> {
let conv = kcp::conv_from_sid(sid);
if candidates.is_empty() {
bail!("牵线没有给出任何候选地址");
}
for _ in 0..KCP_PUNCH_ROUNDS {
for c in candidates {
let _ = sock.send_to(p2p::PUNCH_MAGIC, c).await;
}
tokio::time::sleep(PUNCH_INTERVAL).await;
}
debug!(first = %candidates[0], total = candidates.len(), %sid, "KCP 开始建链(未锁定前会轮流试候选)");
Ok(KcpStream::spawn_candidates(
std::sync::Arc::new(sock),
candidates,
conv,
None,
))
}
async fn kcp_accept(
sock: UdpSocket,
candidates: &[SocketAddr],
expect_ip: Option<std::net::IpAddr>,
sid: &str,
timeout: Duration,
) -> Result<KcpStream> {
let conv = kcp::conv_from_sid(sid);
let deadline = Instant::now() + timeout;
let mut buf = vec![0u8; 2048];
loop {
let remain = deadline.saturating_duration_since(Instant::now());
if remain.is_zero() {
bail!("等待入站 P2P 连接超时(打洞未成功)");
}
for c in candidates {
let _ = sock.send_to(p2p::PUNCH_MAGIC, c).await;
}
match tokio::time::timeout(remain.min(PUNCH_INTERVAL * 4), sock.recv_from(&mut buf)).await {
Ok(Ok((n, from))) => {
if let Some(ip) = expect_ip {
if from.ip() != ip {
debug!(%from, "忽略陌生来源的 P2P 报文");
continue;
}
}
debug!(%from, %sid, "KCP 收到首个数据报,开始建链");
return Ok(KcpStream::spawn(
std::sync::Arc::new(sock),
from,
conv,
Some(buf[..n].to_vec()),
));
}
Ok(Err(e)) => return Err(e).context("读取 P2P 数据报失败"),
Err(_) => continue,
}
}
}
async fn greet_client<S>(s: &mut S, sid: &str, secret_key: &str) -> Result<()>
where
S: AsyncRead + AsyncWrite + Unpin,
{
s.write_all(p2p::handshake_token(secret_key, sid).as_bytes())
.await?;
let expect = p2p::handshake_ack(secret_key, sid);
let mut got = vec![0u8; expect.len()];
tokio::time::timeout(PUNCH_TIMEOUT, s.read_exact(&mut got))
.await
.map_err(|_| anyhow!("等待 P2P 握手响应超时"))??;
let got = String::from_utf8(got).map_err(|_| anyhow!("P2P 握手响应不是 UTF-8"))?;
if !constant_time_eq(&got, &expect) {
bail!("对端未能证明持有 secret_key(P2P 握手应答不匹配)");
}
Ok(())
}
async fn greet_server<S>(s: &mut S, sid: &str, secret_key: &str) -> Result<()>
where
S: AsyncRead + AsyncWrite + Unpin,
{
let mut got = vec![0u8; TOKEN_LEN];
tokio::time::timeout(PUNCH_TIMEOUT, s.read_exact(&mut got))
.await
.map_err(|_| anyhow!("等待 P2P 口令超时"))??;
let got = String::from_utf8(got).map_err(|_| anyhow!("P2P 口令不是 UTF-8"))?;
let expect = p2p::handshake_token(secret_key, sid);
if !constant_time_eq(&got, &expect) {
warn!(%sid, "P2P 握手口令不匹配,拒绝该连接");
let _ = s.write_all(b"BAD").await;
let _ = s.shutdown().await;
bail!("P2P 握手口令不匹配(对方不是本次 xtcp 会话的 peer)");
}
s.write_all(p2p::handshake_ack(secret_key, sid).as_bytes())
.await?;
Ok(())
}
static SERVER_CFG: OnceLock<quinn::ServerConfig> = OnceLock::new();
fn server_config() -> Result<quinn::ServerConfig> {
if let Some(c) = SERVER_CFG.get() {
return Ok(c.clone());
}
let key = rcgen::generate_simple_self_signed(vec![SERVER_NAME.to_string()])
.context("生成 P2P 自签证书失败")?;
let cert_der = key.cert.der().clone();
let key_der = key.key_pair.serialize_der();
let mut tls = rustls::ServerConfig::builder_with_provider(common_crypto_provider())
.with_protocol_versions(&[&rustls::version::TLS13])
.map_err(|e| anyhow!("TLS 版本配置失败:{e}"))?
.with_no_client_auth()
.with_single_cert(
vec![cert_der],
rustls::pki_types::PrivateKeyDer::try_from(key_der)
.map_err(|e| anyhow!("转换 P2P 私钥失败:{e}"))?,
)
.context("构造 P2P 服务端 TLS 配置失败")?;
tls.alpn_protocols = vec![p2p::ALPN.to_vec()];
let quic_cfg = quinn::crypto::rustls::QuicServerConfig::try_from(tls)
.map_err(|e| anyhow!("构造 QUIC 服务端配置失败:{e}"))?;
let mut cfg = quinn::ServerConfig::with_crypto(Arc::new(quic_cfg));
cfg.transport_config(transport_config());
let _ = SERVER_CFG.set(cfg.clone());
Ok(cfg)
}
fn client_config() -> Result<quinn::ClientConfig> {
use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
use rustls::pki_types::{CertificateDer, ServerName, UnixTime};
use rustls::{DigitallySignedStruct, SignatureScheme};
#[derive(Debug)]
struct SkipVerify;
impl ServerCertVerifier for SkipVerify {
fn verify_server_cert(
&self,
_end_entder: &CertificateDer<'_>,
_intermediates: &[CertificateDer<'_>],
_server_name: &ServerName<'_>,
_ocsp: &[u8],
_now: UnixTime,
) -> Result<ServerCertVerified, rustls::Error> {
Ok(ServerCertVerified::assertion())
}
fn verify_tls12_signature(
&self,
_m: &[u8],
_c: &CertificateDer<'_>,
_d: &DigitallySignedStruct,
) -> Result<HandshakeSignatureValid, rustls::Error> {
Ok(HandshakeSignatureValid::assertion())
}
fn verify_tls13_signature(
&self,
_m: &[u8],
_c: &CertificateDer<'_>,
_d: &DigitallySignedStruct,
) -> Result<HandshakeSignatureValid, rustls::Error> {
Ok(HandshakeSignatureValid::assertion())
}
fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
rustls::crypto::ring::default_provider()
.signature_verification_algorithms
.supported_schemes()
}
}
let mut tls = rustls::ClientConfig::builder_with_provider(common_crypto_provider())
.with_protocol_versions(&[&rustls::version::TLS13])
.map_err(|e| anyhow!("TLS 版本配置失败:{e}"))?
.dangerous()
.with_custom_certificate_verifier(Arc::new(SkipVerify))
.with_no_client_auth();
tls.alpn_protocols = vec![p2p::ALPN.to_vec()];
let quic_cfg = quinn::crypto::rustls::QuicClientConfig::try_from(tls)
.map_err(|e| anyhow!("构造 QUIC 客户端配置失败:{e}"))?;
let mut cfg = quinn::ClientConfig::new(Arc::new(quic_cfg));
cfg.transport_config(transport_config());
Ok(cfg)
}
fn common_crypto_provider() -> Arc<rustls::crypto::CryptoProvider> {
nfrp_common::frp::tls::crypto_provider()
}
fn transport_config() -> Arc<quinn::TransportConfig> {
let mut t = quinn::TransportConfig::default();
t.keep_alive_interval(Some(KEEP_ALIVE));
let _ = t.max_idle_timeout(Some(IDLE_TIMEOUT.try_into().expect("超时溢出")));
Arc::new(t)
}
pub(crate) struct QuicStream {
send: SendStream,
recv: RecvStream,
_endpoint: Endpoint,
}
impl AsyncRead for QuicStream {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.recv).poll_read(cx, buf)
}
}
impl AsyncWrite for QuicStream {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
Pin::new(&mut self.send)
.poll_write(cx, buf)
.map_err(std::io::Error::from)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.send).poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.send)
.poll_shutdown(cx)
.map_err(std::io::Error::from)
}
}
pub enum P2PStream {
Quic(QuicStream),
Kcp(KcpStream),
}
impl AsyncRead for P2PStream {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
match self.get_mut() {
Self::Quic(s) => Pin::new(s).poll_read(cx, buf),
Self::Kcp(s) => Pin::new(s).poll_read(cx, buf),
}
}
}
impl AsyncWrite for P2PStream {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
match self.get_mut() {
Self::Quic(s) => Pin::new(s).poll_write(cx, buf),
Self::Kcp(s) => Pin::new(s).poll_write(cx, buf),
}
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
match self.get_mut() {
Self::Quic(s) => Pin::new(s).poll_flush(cx),
Self::Kcp(s) => Pin::new(s).poll_flush(cx),
}
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
match self.get_mut() {
Self::Quic(s) => Pin::new(s).poll_shutdown(cx),
Self::Kcp(s) => Pin::new(s).poll_shutdown(cx),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
#[test]
fn any_addr_follows_remote_family() {
let v4: SocketAddr = "1.2.3.4:7000".parse().unwrap();
let v6: SocketAddr = "[2001:db8::1]:7000".parse().unwrap();
assert!(any_addr(&v4).is_ipv4());
assert!(any_addr(&v6).is_ipv6());
assert_eq!(any_addr(&v4).port(), 0, "端口交给内核分配");
}
#[test]
fn transport_names_are_lenient() {
assert_eq!(parse_transport("kcp"), Transport::Kcp);
assert_eq!(parse_transport("KCP"), Transport::Kcp);
assert_eq!(parse_transport("quic"), Transport::Quic);
assert_eq!(parse_transport("kcpp"), Transport::Quic);
assert_eq!(parse_transport(""), Transport::Quic);
}
#[test]
fn server_and_client_config_both_build() {
let _ = server_config().expect("服务端配置");
let _ = server_config().expect("服务端配置(缓存命中)");
let _ = client_config().expect("客户端配置");
}
#[test]
fn keepalive_is_shorter_than_idle_timeout() {
assert!(
KEEP_ALIVE < IDLE_TIMEOUT,
"心跳必须比空闲超时短,否则永远等不到心跳"
);
assert!(
KEEP_ALIVE <= Duration::from_secs(30),
"心跳间隔必须小于常见 NAT 的 UDP 老化时间,否则洞会关掉"
);
let _ = transport_config();
}
#[test]
fn candidates_always_include_the_observed_address() {
let rd = build_rendezvous(
vec!["203.0.113.9:45001".parse().unwrap()],
Transport::Quic,
8,
);
assert!(rd
.candidates
.contains(&"203.0.113.9:45001".parse().unwrap()));
assert!(rd.candidates.len() > 1, "开了预测就该多出候选端口");
}
#[test]
fn zero_window_falls_back_to_single_address() {
let rd = build_rendezvous(
vec!["203.0.113.9:45001".parse().unwrap()],
Transport::Quic,
0,
);
assert_eq!(rd.candidates.len(), 1);
}
#[tokio::test]
async fn handshake_over_real_quic() {
let scfg = server_config().unwrap();
let srv_sock = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
let srv_addr = srv_sock.local_addr().unwrap();
let server = Endpoint::new(
EndpointConfig::default(),
Some(scfg),
srv_sock,
Arc::new(TokioRuntime),
)
.unwrap();
let cli_sock = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
let client = Endpoint::new(
EndpointConfig::default(),
None,
cli_sock,
Arc::new(TokioRuntime),
)
.unwrap();
let server2 = server.clone();
let sid = "0123456789abcdef0123456789abcdef";
let sk = "top-secret";
let (done_tx, done_rx) = tokio::sync::oneshot::channel::<()>();
let srv_task = tokio::spawn(async move {
let incoming = server.accept().await.expect("accept");
let conn = incoming.accept().unwrap().await.expect("handshake");
let (send, recv) = conn.accept_bi().await.expect("accept_bi");
let mut s = P2PStream::Quic(QuicStream {
send,
recv,
_endpoint: server2,
});
let r = greet_server(&mut s, sid, sk).await;
let _ = done_rx.await;
r
});
let conn = client
.connect_with(client_config().unwrap(), srv_addr, SERVER_NAME)
.unwrap()
.await
.expect("连接成功");
let (send, recv) = conn.open_bi().await.expect("open_bi");
let mut s = P2PStream::Quic(QuicStream {
send,
recv,
_endpoint: client,
});
greet_client(&mut s, sid, sk).await.expect("握手成功");
s.write_all(b"ping").await.unwrap();
drop(done_tx);
srv_task.await.unwrap().expect("服务端握手");
}
async fn mini_rendezvous(sock: UdpSocket) {
let mut buf = [0u8; 512];
let mut seen: HashMap<Role, (Vec<SocketAddr>, Transport)> = HashMap::new();
loop {
let Ok((n, peer)) = sock.recv_from(&mut buf).await else {
break;
};
let Some(Packet::Hello {
role,
sid,
transport,
}) = p2p::decode(&buf[..n])
else {
continue;
};
let (addrs, t) = seen.entry(role).or_insert((Vec::new(), transport));
if !addrs.contains(&peer) {
addrs.push(peer);
}
*t = transport;
let Some((other_addrs, _)) = seen.get(&role.peer()) else {
continue;
};
let chosen = seen
.get(&Role::Visitor)
.map(|(_, t)| *t)
.unwrap_or(transport);
let mine = seen[&role].0.clone();
let other = other_addrs.clone();
if let Some(pkt) = p2p::encode_peers(role.peer(), &sid, chosen, &other) {
let _ = sock.send_to(&pkt, peer).await;
}
if let Some(pkt) = p2p::encode_peers(role, &sid, chosen, &mine) {
for a in &other {
let _ = sock.send_to(&pkt, a).await;
}
}
return;
}
}
#[tokio::test]
async fn full_p2p_over_rendezvous() {
let rd = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let rd_addr = rd.local_addr().unwrap();
tokio::spawn(mini_rendezvous(rd));
let sid = "aaaabbbbccccddddeeeeffff00001111";
let sk = "p2p-secret";
let provider = tokio::spawn({
let server = rd_addr;
async move { accept_as_provider(&server, sid, sk, 4).await }
});
let mut visitor = connect_as_visitor(&rd_addr, sid, sk, Transport::Quic, 4)
.await
.expect("visitor 侧 P2P 建立失败");
let mut provider_stream = provider.await.unwrap().expect("provider 侧 P2P 建立失败");
visitor.write_all(b"ping").await.unwrap();
let mut buf = [0u8; 4];
provider_stream.read_exact(&mut buf).await.unwrap();
assert_eq!(&buf, b"ping");
provider_stream.write_all(b"pong").await.unwrap();
visitor.read_exact(&mut buf).await.unwrap();
assert_eq!(&buf, b"pong", "P2P 链路必须双向可通");
}
#[tokio::test]
async fn full_p2p_over_rendezvous_with_kcp() {
let rd = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let rd_addr = rd.local_addr().unwrap();
tokio::spawn(mini_rendezvous(rd));
let sid = "11112222333344445555666677778888";
let sk = "kcp-secret";
let provider = tokio::spawn({
let server = rd_addr;
async move { accept_as_provider(&server, sid, sk, 4).await }
});
let mut visitor = connect_as_visitor(&rd_addr, sid, sk, Transport::Kcp, 4)
.await
.expect("visitor 侧 KCP P2P 建立失败");
let mut provider_stream = provider
.await
.unwrap()
.expect("provider 侧 KCP P2P 建立失败");
assert!(
matches!(provider_stream, P2PStream::Kcp(_)),
"provider 也必须走 KCP"
);
visitor.write_all(b"ping").await.unwrap();
let mut buf = [0u8; 4];
provider_stream.read_exact(&mut buf).await.unwrap();
assert_eq!(&buf, b"ping");
provider_stream.write_all(b"pong").await.unwrap();
visitor.read_exact(&mut buf).await.unwrap();
assert_eq!(&buf, b"pong", "KCP P2P 链路必须双向可通");
}
#[tokio::test]
async fn wrong_token_is_rejected() {
let scfg = server_config().unwrap();
let srv_sock = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
let srv_addr = srv_sock.local_addr().unwrap();
let server = Endpoint::new(
EndpointConfig::default(),
Some(scfg),
srv_sock,
Arc::new(TokioRuntime),
)
.unwrap();
let cli_sock = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
let client = Endpoint::new(
EndpointConfig::default(),
None,
cli_sock,
Arc::new(TokioRuntime),
)
.unwrap();
let sid = "0123456789abcdef0123456789abcdef";
let server2 = server.clone();
let srv_task = tokio::spawn(async move {
let incoming = server.accept().await.expect("accept");
let conn = incoming.accept().unwrap().await.expect("handshake");
let (send, recv) = conn.accept_bi().await.expect("accept_bi");
let mut s = P2PStream::Quic(QuicStream {
send,
recv,
_endpoint: server2,
});
greet_server(&mut s, sid, "right-key").await
});
let conn = client
.connect_with(client_config().unwrap(), srv_addr, SERVER_NAME)
.unwrap()
.await
.expect("连接成功");
let (send, recv) = conn.open_bi().await.expect("open_bi");
let mut s = P2PStream::Quic(QuicStream {
send,
recv,
_endpoint: client,
});
let r = greet_client(&mut s, sid, "wrong-key").await;
assert!(r.is_err(), "口令不对却成功了:等于 NAT 外谁都能进来");
let srv = srv_task.await.unwrap();
assert!(srv.is_err());
}
#[tokio::test]
async fn fake_provider_without_secret_is_rejected() {
let sid = "deadbeefdeadbeefdeadbeefdeadbeef";
let sk = "real-secret";
let fake_replies: Vec<Vec<u8>> = vec![
b"RTNL-OK".to_vec(),
vec![b'A'; 64],
nfrp_common::p2p::handshake_ack("wrong-key", sid).into_bytes(),
];
for reply in fake_replies {
let reply_desc = String::from_utf8_lossy(&reply).to_string();
let scfg = server_config().unwrap();
let srv_sock = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
let srv_addr = srv_sock.local_addr().unwrap();
let server = Endpoint::new(
EndpointConfig::default(),
Some(scfg),
srv_sock,
Arc::new(TokioRuntime),
)
.unwrap();
let cli_sock = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
let client = Endpoint::new(
EndpointConfig::default(),
None,
cli_sock,
Arc::new(TokioRuntime),
)
.unwrap();
let (done_tx, done_rx) = tokio::sync::oneshot::channel::<()>();
let server2 = server.clone();
let fake = tokio::spawn(async move {
let incoming = server.accept().await.expect("accept");
let conn = incoming.accept().unwrap().await.expect("handshake");
let (mut send, mut recv) = conn.accept_bi().await.expect("accept_bi");
let mut tok = [0u8; 64];
recv.read_exact(&mut tok)
.await
.expect("假 provider 收到口令");
send.write_all(&reply).await.expect("假 provider 回应答");
let _ = send.finish();
let _ = done_rx.await;
drop(server2);
});
let conn = client
.connect_with(client_config().unwrap(), srv_addr, SERVER_NAME)
.unwrap()
.await
.expect("连接成功");
let (send, recv) = conn.open_bi().await.expect("open_bi");
let mut s = P2PStream::Quic(QuicStream {
send,
recv,
_endpoint: client,
});
let r = greet_client(&mut s, sid, sk).await;
assert!(
r.is_err(),
"假 provider 用 {reply_desc:?} 应答却被接受了:B-1 劫持链没堵住"
);
let _ = done_tx.send(());
fake.await.unwrap();
}
}
}