diff --git a/qbase/src/time.rs b/qbase/src/time.rs index 11377f4a..6392ce68 100644 --- a/qbase/src/time.rs +++ b/qbase/src/time.rs @@ -152,6 +152,8 @@ impl IdleTimer { // Checks health of the path and // determines whether a heartbeat packet needs to be sent. pub fn health(&mut self) -> Result, TimeOut> { + // TODO: 考虑连接还没关闭,闲置路径的 defer_idle_timeout 不能生效 + // 得统一考虑成连接级的 defer_idle_timeout if let Some(t) = self.last_effective_comm { let elapsed = t.elapsed(); if elapsed > self.idle_config.defer_idle_timeout() { diff --git a/qbase/src/util.rs b/qbase/src/util.rs index 48a3cda9..7bcb2032 100644 --- a/qbase/src/util.rs +++ b/qbase/src/util.rs @@ -14,4 +14,4 @@ mod unique_id; pub use unique_id::{UniqueId, UniqueIdGenerator}; mod wakers; -pub use wakers::{WakerVec, Wakers}; +pub use wakers::{WakerGroup, Wakers}; diff --git a/qbase/src/util/wakers.rs b/qbase/src/util/wakers.rs index 4d1d000e..b2a5f896 100644 --- a/qbase/src/util/wakers.rs +++ b/qbase/src/util/wakers.rs @@ -8,30 +8,30 @@ use std::{ use smallvec::SmallVec; #[derive(Debug, Clone)] -pub struct WakerVec { +pub struct WakerGroup { wakers: SmallVec<[Waker; N]>, } -impl Default for WakerVec { +impl Default for WakerGroup { fn default() -> Self { Self::new() } } -impl WakerVec { +impl WakerGroup { pub const fn new() -> Self { Self { wakers: SmallVec::new_const(), } } - pub fn register(&mut self, waker: &Waker) { + pub fn add(&mut self, waker: &Waker) { if !self.wakers.iter().any(|w| w.will_wake(waker)) { self.wakers.push(waker.clone()); } } - pub fn unregister(&mut self, waker: &Waker) { + pub fn remove(&mut self, waker: &Waker) { self.wakers .retain(|registered| !registered.will_wake(waker)); } @@ -43,7 +43,7 @@ impl WakerVec { } } -impl Drop for WakerVec { +impl Drop for WakerGroup { fn drop(&mut self) { self.wake_all(); } @@ -51,7 +51,7 @@ impl Drop for WakerVec { #[derive(Debug)] pub struct Wakers { - wakers: Mutex>, + inner: Mutex>, } impl Wake for Wakers { @@ -73,24 +73,30 @@ impl Default for Wakers { impl Wakers { pub const fn new() -> Self { Self { - wakers: Mutex::new(WakerVec::new()), + inner: Mutex::new(WakerGroup::new()), } } - fn lock(&self) -> MutexGuard<'_, WakerVec> { - self.wakers.lock().expect("Wakers mutex poisoned") + fn lock_guard(&self) -> MutexGuard<'_, WakerGroup> { + self.inner.lock().expect("Wakers mutex poisoned") } - pub fn register(&self, waker: &Waker) { - self.lock().register(waker) + pub fn add(&self, waker: &Waker) { + self.lock_guard().add(waker) } - pub fn unregister(&self, waker: &Waker) { - self.lock().unregister(waker) + pub fn remove(&self, waker: &Waker) { + self.lock_guard().remove(waker) + } + + pub fn together_with(self: &Arc, waker: &Waker) -> Waker { + let mut guard = self.lock_guard(); + guard.add(waker); + Waker::from(self.clone()) } pub fn wake_all(&self) { - { mem::replace(&mut *self.lock(), WakerVec::new()) }.wake_all() + { mem::replace(&mut *self.lock_guard(), WakerGroup::new()) }.wake_all() } pub fn to_waker(self: &Arc) -> Waker { @@ -102,10 +108,10 @@ impl Wakers { cx: &mut Context<'_>, poll: impl FnOnce(&mut Context<'_>) -> Poll, ) -> Poll { - self.register(cx.waker()); + self.add(cx.waker()); let result = poll(&mut Context::from_waker(&self.to_waker())); if result.is_ready() { - self.unregister(cx.waker()); + self.remove(cx.waker()); } result } diff --git a/qinterface/src/io/handy.rs b/qinterface/src/io/handy.rs index 3bd2eb7a..79140f6e 100644 --- a/qinterface/src/io/handy.rs +++ b/qinterface/src/io/handy.rs @@ -84,7 +84,7 @@ pub mod qudp { } } - fn usc(&self) -> io::Result<&qudp::UdpSocket> { + fn socket(&self) -> io::Result<&qudp::UdpSocket> { self.io .as_ref() .map_err(|e| io::Error::from(e.clone())) @@ -98,7 +98,7 @@ pub mod qudp { } fn bound_addr(&self) -> io::Result { - self.usc()?.local_addr() + self.socket()?.local_addr() } fn max_segments(&self) -> io::Result { @@ -115,11 +115,17 @@ pub mod qudp { pkts: &[io::IoSlice], route: Route, ) -> Poll> { - let io = self.usc()?; - self.send_wakers.combine_with(cx, |cx| { - debug_assert_eq!(route.ecn(), None); - io.poll_send(cx, pkts, &route.line) - }) + let io = self.socket()?; + let waker = cx.waker(); + let waker_group = self.send_wakers.together_with(waker); + let cx = &mut Context::from_waker(&waker_group); + + debug_assert_eq!(route.ecn(), None); + let result = io.poll_send(cx, pkts, &route); + if result.is_ready() { + self.send_wakers.remove(waker); + } + result } fn poll_recv( @@ -128,7 +134,7 @@ pub mod qudp { pkts: &mut [BytesMut], route: &mut [Route], ) -> Poll> { - let io = self.usc()?; + let io = self.socket()?; self.recv_wakers.combine_with(cx, |cx| { let len = route.len().min(pkts.len()); let mut rcvd_lines = Vec::with_capacity(len); @@ -152,7 +158,7 @@ pub mod qudp { } fn poll_close(&mut self, _cx: &mut Context) -> Poll> { - self.usc()?; + self.socket()?; self.send_wakers.wake_all(); self.recv_wakers.wake_all(); self.io = Ok(Err(Closed(()))); diff --git a/qprotocol/Cargo.toml b/qprotocol/Cargo.toml index 156b27fc..dcf38ec1 100644 --- a/qprotocol/Cargo.toml +++ b/qprotocol/Cargo.toml @@ -12,9 +12,13 @@ categories.workspace = true [dependencies] bytes = { workspace = true } nom = { workspace = true } +dashmap = { workspace = true } qbase = { workspace = true } +qinterface = { workspace = true } +qudp = { workspace = true } rand = { workspace = true } thiserror = { workspace = true } +tokio = { workspace = true } [features] # Enable shorter TTL only for tests (especially integration tests in other crates). diff --git a/qprotocol/src/io.rs b/qprotocol/src/io.rs index 8b137891..fb2ab596 100644 --- a/qprotocol/src/io.rs +++ b/qprotocol/src/io.rs @@ -1 +1,231 @@ +use std::{ + collections::HashMap, + io::{self, IoSlice}, + net::SocketAddr, + sync::{Arc, atomic::AtomicUsize}, + task::{Context, Poll}, +}; +use dashmap::DashMap; +use qbase::{ + ArcReceiving, + net::{ + NetFeature, + addr::EndpointAddr, + route::{Link, Route}, + }, +}; +use qinterface::bind_uri::BindUri; +use qudp::ext::UdpSocket; +use tokio::sync::{SetOnce, mpsc}; + +use crate::stun; + +/// A unique handle identifying a managed UDP socket inside the [`Dock`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct Fd(usize); + +static FD_COUNTER: AtomicUsize = AtomicUsize::new(0); + +impl Fd { + fn next() -> Self { + Fd(FD_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed)) + } +} + +/// STUN agent handle – the concrete keepalive logic lives in `stun/agent.rs`. +/// Wraps an inner agent behind `Arc` so that the keepalive task and the Dock +/// can share ownership. +pub struct Agent { + // TODO: fill in keepalive / probe state + pub stun_addr: SocketAddr, +} + +pub type ArcAgent = Arc; + +/// A packet that should be forwarded (relayed). +pub struct ForwardPacket { + pub payload: bytes::Bytes, + pub link: Link, +} + +/// Aggregated external addresses learned from STUN for the resident sockets. +#[derive(Default)] +pub struct AddressBook { + /// bind_uri -> list of discovered endpoint addresses + entries: HashMap>, +} + +/// The harbour where all managed UDP sockets are docked. +/// +/// It owns every UDP socket, their STUN agents, NAT probing results, +/// transaction routing, and protocol-level channels. +pub struct Dock { + // ── resident listeners ────────────────────────────────────────────── + /// BindUri → Fd mapping for sockets that **must** stay open. + resident: DashMap, + + // ── sockets ───────────────────────────────────────────────────────── + /// All managed sockets (both long-lived residents and ephemeral ones). + sockets: DashMap, + + // ── STUN agents ───────────────────────────────────────────────────── + /// Per-socket STUN agents with their associated server address. + /// `Vec` because the number of agents per socket is typically tiny. + agents: DashMap>, + /// Pending STUN BindingRequest→BindingResponse transactions. + transactions: DashMap>, + // ── NAT feature ───────────────────────────────────────────────────── + /// Per-socket NAT type, resolved asynchronously by the NAT probe task. + nat_features: DashMap>, + /// Externally-visible addresses discovered by agents. + endpoint_addrs: DashMap>>>, + + // ── global address book ───────────────────────────────────────────── + /// Aggregated endpoint addresses **only** for resident sockets. + address_book: AddressBook, + + // ── QUIC routing ──────────────────────────────────────────────────── + // router: HashMap, // TODO: wire up QUIC CID router + + // ── channels ──────────────────────────────────────────────────────── + /// Incoming STUN BindingRequests (for the built-in STUN server). + stun_tx: mpsc::UnboundedSender, + stun_rx: mpsc::UnboundedReceiver, + + /// Packets that need to be forwarded / relayed. + forward_tx: mpsc::UnboundedSender, + forward_rx: mpsc::UnboundedReceiver, +} + +impl Dock { + /// Create an empty `Dock`. + pub fn new() -> Self { + let (stun_tx, stun_rx) = mpsc::unbounded_channel(); + let (forward_tx, forward_rx) = mpsc::unbounded_channel(); + Self { + resident: DashMap::new(), + sockets: DashMap::new(), + agents: DashMap::new(), + endpoint_addrs: DashMap::new(), + transactions: DashMap::new(), + nat_features: DashMap::new(), + stun_tx, + stun_rx, + forward_tx, + forward_rx, + address_book: AddressBook::default(), + } + } + + /// Bind a new UDP socket for the given URI and return its [`Fd`]. + pub fn bind(&self, uri: BindUri) -> io::Result { + let addr: SocketAddr = (&uri) + .try_into() + .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, format!("{e}")))?; + let socket = UdpSocket::bind(addr); + let fd = Fd::next(); + + self.sockets.insert(fd, socket); + self.resident.insert(uri, fd); + self.nat_features.insert(fd, SetOnce::const_new()); + self.endpoint_addrs.insert(fd, HashMap::new()); + Ok(fd) + } + + /// Close and remove the socket identified by `fd`. + pub fn close(&self, fd: Fd) { + if let Some((_, mut socket)) = self.sockets.remove(&fd) { + socket.close(); + } + self.agents.remove(&fd); + self.nat_features.remove(&fd); + self.endpoint_addrs.remove(&fd); + // Remove from resident map if present. + self.resident.retain(|_, v| *v != fd); + } + + // ── send ──────────────────────────────────────────────────────────── + + /// Send datagrams through the socket identified by `fd`. + pub fn poll_send( + &self, + fd: Fd, + cx: &mut Context<'_>, + pkts: &[IoSlice<'_>], + route: Route, + ) -> Poll> { + let socket = self + .sockets + .get(&fd) + .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "unknown fd"))?; + socket.poll_send(cx, pkts, route) + } + + // ── STUN agent management ─────────────────────────────────────────── + + /// Register a STUN agent for the given socket and agent server address. + pub fn add_agent(&self, fd: Fd, agent_addr: SocketAddr) -> io::Result<()> { + if !self.sockets.contains_key(&fd) { + return Err(io::Error::new(io::ErrorKind::NotFound, "unknown fd")); + } + let agent = Arc::new(Agent { + stun_addr: agent_addr, + }); + let mut agents = self.agents.entry(fd).or_default(); + if agents.iter().any(|(a, _)| *a == agent_addr) { + return Err(io::Error::new( + io::ErrorKind::AlreadyExists, + "agent already registered", + )); + } + agents.push((agent_addr, agent)); + Ok(()) + } + + /// Remove a previously registered STUN agent. + pub fn del_agent(&self, fd: Fd, agent_addr: SocketAddr) { + if let Some(mut agents) = self.agents.get_mut(&fd) { + agents.retain(|(a, _)| *a != agent_addr); + } + if let Some(mut ep_addrs) = self.endpoint_addrs.get_mut(&fd) { + ep_addrs.remove(&agent_addr); + } + } + + // ── queries ───────────────────────────────────────────────────────── + + /// Return the locally bound address of the socket. + pub fn bind_addr(&self, fd: Fd) -> io::Result { + self.sockets + .get(&fd) + .map(|s| s.local_addr()) + .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "unknown fd"))? + } + + /// Asynchronously wait for the NAT type probing result of the socket. + pub async fn nat_feature(&self, fd: Fd) -> io::Result { + let once = self + .nat_features + .get(&fd) + .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "unknown fd"))?; + Ok(*once.wait().await) + } + + /// Asynchronously wait for the external (STUN-mapped) endpoint address. + pub async fn endpoint_addr(&self, fd: Fd, agent_addr: SocketAddr) -> io::Result { + let once = { + let ep_addrs = self + .endpoint_addrs + .get(&fd) + .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "unknown fd"))?; + + ep_addrs + .get(&agent_addr) + .cloned() + .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "unknown agent"))? + }; + + Ok(*once.wait().await) + } +} diff --git a/qprotocol/src/lib.rs b/qprotocol/src/lib.rs index 7df7211c..29f97adf 100644 --- a/qprotocol/src/lib.rs +++ b/qprotocol/src/lib.rs @@ -1,5 +1,6 @@ pub mod dns; pub mod forward; -pub mod io; +// pub mod io; +pub mod packet; pub mod quic; pub mod stun; diff --git a/qprotocol/src/packet.rs b/qprotocol/src/packet.rs new file mode 100644 index 00000000..20dfc7b8 --- /dev/null +++ b/qprotocol/src/packet.rs @@ -0,0 +1,232 @@ +use bytes::BufMut; +use qbase::net::{ + Family, + addr::{EndpointAddr, WriteEndpointAddr, be_endpoint_addr}, + route::Pathway, +}; + +const STUN_HEADER_MASK: u8 = 0b1111_1110; +const STUN_HEADER_BITS: u8 = 0b1100_0010; + +const FORWARD_HEADER_MASK: u8 = 0b1110_0000; +const FORWARD_VERSION_MASK: u8 = 0b1111_0000; +const FORWARD_HEADER_BITS: u8 = 0b0110_0000; +const FORWARD_BIT: u8 = 0b0000_1000; +const FORWARD_FAMILY_BIT: u8 = 0b0000_0100; +const FORWARD_SRC_TYPE_BIT: u8 = 0b0000_0010; +const FORWARD_DST_TYPE_BIT: u8 = 0b0000_0001; + +#[derive(PartialEq, Eq, Debug)] +pub enum HeaderType { + Stun(u8), // 最后 bit + Forward(u8), // 最后 5bit +} + +// Stun Packet { +// Header Form (1) = 1, +// Fixed Bit (1) = 1, +// Stun Hdr (6), // Request 0b000010 #Response 0b000011 +// Version (32) = 0, +// DDIL(8) = 0, // 伪装0长度的目标连接ID +// SDIL(8) = 0, // 伪装0长度的源连接ID +// Ver(16), // 2个字节,表示我们自定义的版本号,方便未来升级 +// ... Stun payload +// } +#[derive(Clone, Copy)] +pub struct StunHeader { + version: u16, +} + +impl StunHeader { + pub fn new(version: u16) -> Self { + Self { version } + } + + pub fn encoding_size() -> usize { + 1 + 4 + 4 + } +} + +pub fn be_stun_header(input: &[u8]) -> nom::IResult<&[u8], StunHeader> { + let (remain, version) = nom::number::streaming::be_u16(input)?; + Ok((remain, StunHeader { version })) +} + +pub trait WriteStunHeader { + fn put_stun_header(&mut self, stun_header: &StunHeader); +} + +impl WriteStunHeader for T { + fn put_stun_header(&mut self, stun_header: &StunHeader) { + self.put_u8(STUN_HEADER_BITS); + self.put_u32(0); + self.put_u8(0); + self.put_u8(0); + self.put_u16(stun_header.version); + } +} + +// Forward Packet { +// Header Form (1) = 0, +// Fixed Bit (1) = 1, +// Spin Bit (1) = 1, // 1表示带有转发包头 +// Remain (5), // 使其等于真正QUIC包第一字节的后5bit,飘忽不定,伪装够深 +// Version (4), +// Forward (1) = 1, +// Family (1), // 0表示IPv4,1表示IPv6 +// Src type(1), // 0表示直连,1表示带agent +// Dst type(1), // 0表示直连,1表示带agent +// Src endpoint, // 根据src type,是Endpoint::Agent还是Direct +// Dst endpoint, // 根据dst type,是Endpoint::Agent还是Direct +// ... Real Quic Packet +// } +#[derive(Debug, Clone, Copy)] +pub struct ForwardHeader { + remian: u8, // 后 5bits + version: u8, // 前 4bits + pathway: Pathway, +} + +impl ForwardHeader { + pub fn encoding_size(pathway: &Pathway) -> usize { + if matches!(pathway.remote(), EndpointAddr::Direct { .. }) { + return 0; + } + 1 + 1 + pathway.local().encoding_size() + pathway.remote().encoding_size() + } + + pub fn pathway(&self) -> Pathway { + self.pathway + } + + pub fn new(version: u8, pathway: &Pathway, buffer: &[u8]) -> Self { + let remian = buffer[0] & 0b0001_1111; + Self { + remian, + version, + pathway: *pathway, + } + } +} + +pub trait WriteForwardHeader { + fn put_forward_header(&mut self, forward_header: &ForwardHeader); +} + +impl WriteForwardHeader for T { + fn put_forward_header(&mut self, forward_header: &ForwardHeader) { + self.put_u8(FORWARD_HEADER_BITS | forward_header.remian); + let mut flag = (forward_header.version << 4) | FORWARD_BIT; + + if forward_header.pathway.local().ip().is_ipv6() { + flag |= FORWARD_FAMILY_BIT; + } + if matches!(forward_header.pathway.local(), EndpointAddr::Agent { .. }) { + flag |= FORWARD_SRC_TYPE_BIT; + } + if matches!(forward_header.pathway.remote(), EndpointAddr::Agent { .. }) { + flag |= FORWARD_DST_TYPE_BIT; + } + self.put_u8(flag); + self.put_endpoint_addr(forward_header.pathway.local()); + self.put_endpoint_addr(forward_header.pathway.remote()); + } +} + +pub fn be_forward_header(input: &[u8]) -> nom::IResult<&[u8], ForwardHeader> { + let (remain, first) = nom::number::streaming::be_u8(input)?; + let version = (first & FORWARD_VERSION_MASK) >> 4; + let flag = first & !FORWARD_VERSION_MASK; + let family = match flag & FORWARD_FAMILY_BIT { + 0 => Family::V4, + _ => Family::V6, + }; + + let src_ep_typ = flag & FORWARD_SRC_TYPE_BIT; + let dst_ep_typ = flag & FORWARD_DST_TYPE_BIT; + let (remain, src) = be_endpoint_addr(remain, src_ep_typ, family)?; + let (remain, dst) = be_endpoint_addr(remain, dst_ep_typ, family)?; + let pathway = Pathway::new(src, dst); + Ok(( + remain, + ForwardHeader { + remian: first, + version, + pathway, + }, + )) +} + +#[derive(Clone, Copy)] +pub enum Header { + Stun(StunHeader), + Forward(ForwardHeader), +} + +pub fn be_header_type(input: &[u8]) -> nom::IResult<&[u8], HeaderType> { + let (remain, first) = nom::number::streaming::be_u8(input)?; + if first & STUN_HEADER_MASK == STUN_HEADER_BITS { + let (remain, version) = nom::number::streaming::be_u32(remain)?; + if version == 0 { + let (remain, _) = nom::number::streaming::be_u8(remain)?; + let (remain, _) = nom::number::streaming::be_u8(remain)?; + return Ok((remain, HeaderType::Stun(first & 1))); + } + } else if first & FORWARD_HEADER_MASK == FORWARD_HEADER_BITS { + return Ok((remain, HeaderType::Forward(first & 0b0001_1111))); + } + Err(nom::Err::Error(nom::error::make_error( + input, + nom::error::ErrorKind::Alt, + ))) +} + +pub fn be_header(input: &[u8]) -> nom::IResult<&[u8], Header> { + let (remain, ty) = be_header_type(input)?; + match ty { + HeaderType::Stun(_ty) => { + let (remain, stun_hdr) = be_stun_header(remain)?; + Ok((remain, Header::Stun(stun_hdr))) + } + HeaderType::Forward(_ty) => { + let (remain, forward_hdr) = be_forward_header(remain)?; + Ok((remain, Header::Forward(forward_hdr))) + } + } +} + +pub trait WriteHeader { + fn put_header(&mut self, header: &Header); +} + +impl WriteHeader for T { + fn put_header(&mut self, header: &Header) { + match header { + Header::Stun(stun_header) => { + self.put_stun_header(stun_header); + } + Header::Forward(forward_header) => { + self.put_forward_header(forward_header); + } + } + } +} + +#[cfg(test)] +mod tests { + use bytes::BytesMut; + + use super::*; + + #[test] + fn test_stun_header() { + let stun_hdr = StunHeader::new(0); + let mut buf = BytesMut::with_capacity(StunHeader::encoding_size()); + buf.put_stun_header(&stun_hdr); + let (remain, hdr) = be_header_type(&buf[..]).unwrap(); + assert_eq!(hdr, HeaderType::Stun(0)); + let (remain, stun_hdr) = be_stun_header(remain).unwrap(); + assert_eq!(stun_hdr.version, 0); + assert_eq!(remain.len(), 0) + } +} diff --git a/qprotocol/src/stun.rs b/qprotocol/src/stun.rs index d0e87a0d..fdd7e095 100644 --- a/qprotocol/src/stun.rs +++ b/qprotocol/src/stun.rs @@ -1 +1,3 @@ pub mod msg; + +pub use msg::{Request, Response, TransactionId}; diff --git a/qtraversal/src/future.rs b/qtraversal/src/future.rs index 18f02efe..28e1a0d1 100644 --- a/qtraversal/src/future.rs +++ b/qtraversal/src/future.rs @@ -5,11 +5,11 @@ use std::{ task::{Context, Poll}, }; -use qbase::util::WakerVec; +use qbase::util::WakerGroup; #[derive(Debug)] enum FutureState { - Demand(WakerVec), + Demand(WakerGroup), Ready(T), } @@ -137,7 +137,7 @@ impl Future { let mut state = self.state(); match state.deref_mut() { FutureState::Demand(wakers) => { - wakers.register(cx.waker()); + wakers.add(cx.waker()); Poll::Pending } FutureState::Ready(..) => Poll::Ready(ReadyFuture(state)), @@ -167,7 +167,7 @@ impl Future { let mut state = self.state(); *state = match state.deref_mut() { FutureState::Demand(wakers) => FutureState::Demand(mem::take(wakers)), - FutureState::Ready(_) => FutureState::Demand(WakerVec::default()), + FutureState::Ready(_) => FutureState::Demand(WakerGroup::default()), }; } } diff --git a/qtraversal/src/lib.rs b/qtraversal/src/lib.rs index 0c7f2e77..eb1550f3 100644 --- a/qtraversal/src/lib.rs +++ b/qtraversal/src/lib.rs @@ -1,10 +1,6 @@ -use qbase::net::addr::EndpointAddr; - pub mod addr; mod future; pub mod nat; pub mod packet; pub mod punch; pub mod route; - -pub type PathWay = qbase::net::route::Pathway; diff --git a/qtraversal/src/packet.rs b/qtraversal/src/packet.rs index 12a6afd0..20dfc7b8 100644 --- a/qtraversal/src/packet.rs +++ b/qtraversal/src/packet.rs @@ -2,10 +2,9 @@ use bytes::BufMut; use qbase::net::{ Family, addr::{EndpointAddr, WriteEndpointAddr, be_endpoint_addr}, + route::Pathway, }; -use crate::PathWay; - const STUN_HEADER_MASK: u8 = 0b1111_1110; const STUN_HEADER_BITS: u8 = 0b1100_0010; @@ -85,22 +84,22 @@ impl WriteStunHeader for T { pub struct ForwardHeader { remian: u8, // 后 5bits version: u8, // 前 4bits - pathway: PathWay, + pathway: Pathway, } impl ForwardHeader { - pub fn encoding_size(pathway: &PathWay) -> usize { + pub fn encoding_size(pathway: &Pathway) -> usize { if matches!(pathway.remote(), EndpointAddr::Direct { .. }) { return 0; } 1 + 1 + pathway.local().encoding_size() + pathway.remote().encoding_size() } - pub fn pathway(&self) -> PathWay { + pub fn pathway(&self) -> Pathway { self.pathway } - pub fn new(version: u8, pathway: &PathWay, buffer: &[u8]) -> Self { + pub fn new(version: u8, pathway: &Pathway, buffer: &[u8]) -> Self { let remian = buffer[0] & 0b0001_1111; Self { remian, @@ -147,7 +146,7 @@ pub fn be_forward_header(input: &[u8]) -> nom::IResult<&[u8], ForwardHeader> { let dst_ep_typ = flag & FORWARD_DST_TYPE_BIT; let (remain, src) = be_endpoint_addr(remain, src_ep_typ, family)?; let (remain, dst) = be_endpoint_addr(remain, dst_ep_typ, family)?; - let pathway = PathWay::new(src, dst); + let pathway = Pathway::new(src, dst); Ok(( remain, ForwardHeader { diff --git a/qtraversal/src/punch/puncher.rs b/qtraversal/src/punch/puncher.rs index 58891f8b..b7a2a961 100644 --- a/qtraversal/src/punch/puncher.rs +++ b/qtraversal/src/punch/puncher.rs @@ -18,7 +18,7 @@ use qbase::{ net::{ NatType, addr::EndpointAddr, - route::{Line, Link, Route}, + route::{Line, Link, Pathway, Route}, tx::Signals, }, packet::{ @@ -42,7 +42,6 @@ use tokio::{task::AbortHandle, time::timeout}; use tracing::Instrument as _; use crate::{ - PathWay, addr::AddressBook, nat::{ client::{StunClientComponent, StunClientsComponent}, @@ -93,9 +92,9 @@ fn build_validated_way( remote: EndpointAddr, local_addr: SocketAddr, remote_addr: SocketAddr, -) -> Result<(BindUri, Link, PathWay), InvalidWay> { +) -> Result<(BindUri, Link, Pathway), InvalidWay> { let link = Link::new(local_addr, remote_addr); - let pathway = PathWay::new(local, remote); + let pathway = Pathway::new(local, remote); let way = (bind.clone(), pathway, link); validate_outbound_candidate(&way)?; Ok((way.0, way.2, way.1)) @@ -903,7 +902,7 @@ where &self, endpoint: EndpointAddr, source: qresolve::Source, - ) -> io::Result> { + ) -> io::Result> { let local_endpoints = { let mut address_book = self.0.address_book.lock().unwrap(); address_book.add_peer_endpoint(endpoint, source.clone())?; @@ -991,7 +990,7 @@ where fn recv_punch_me_now( &self, - pathway: PathWay, + pathway: Pathway, punch_me_now_frame: PunchMeNowFrame, ) -> io::Result<()> { let punch_id = punch_me_now_frame.punch_id().flip(); @@ -1655,7 +1654,7 @@ where local: &EndpointAddr, remote: &EndpointAddr, source: &qresolve::Source, - ) -> Result<(BindUri, Link, PathWay), ResolvePathError> { + ) -> Result<(BindUri, Link, Pathway), ResolvePathError> { if let qresolve::Source::Mdns { nic, family } = source { let matches_iface = bind .as_iface_bind_uri() @@ -1701,7 +1700,7 @@ where } } -impl ReceiveFrame<(BindUri, PathWay, Link, ReliableFrame)> for ArcPuncher +impl ReceiveFrame<(BindUri, Pathway, Link, ReliableFrame)> for ArcPuncher where TX: SendFrame + Send + Sync + Clone + 'static, PH: ProductHeader + Send + Sync + 'static, @@ -1714,7 +1713,7 @@ where fn recv_frame( &self, - (_bind, pathway, link, frame): (BindUri, PathWay, Link, ReliableFrame), + (_bind, pathway, link, frame): (BindUri, Pathway, Link, ReliableFrame), ) -> Result { tracing::debug!(target: "punch", %pathway, %link, frame = ?frame, "received reliable punch frame"); match frame { @@ -1748,7 +1747,7 @@ where } } -impl ReceiveFrame<(BindUri, PathWay, Link, PunchHelloFrame)> for ArcPuncher +impl ReceiveFrame<(BindUri, Pathway, Link, PunchHelloFrame)> for ArcPuncher where TX: SendFrame + Send + Sync + Clone + 'static, PH: ProductHeader + Send + Sync + 'static, @@ -1761,7 +1760,7 @@ where fn recv_frame( &self, - (_bind, pathway, link, frame): (BindUri, PathWay, Link, PunchHelloFrame), + (_bind, pathway, link, frame): (BindUri, Pathway, Link, PunchHelloFrame), ) -> Result { tracing::debug!(target: "punch", %pathway, %link, frame = ?frame, "received punch hello frame"); let punch_id = frame.punch_id().flip(); @@ -1871,7 +1870,7 @@ mod tests { let (_, link, pathway) = build_validated_way(&bind, local, remote, local_addr, remote_addr) .expect("matching loopback scope must be retained"); assert_eq!(link, Link::new(local_addr, remote_addr)); - assert_eq!(pathway, PathWay::new(local, remote)); + assert_eq!(pathway, Pathway::new(local, remote)); } #[test] @@ -1886,7 +1885,7 @@ mod tests { .expect("wildcard source selection belongs to IO"); assert_eq!(link, Link::new(local_addr, remote_addr)); - assert_eq!(pathway, PathWay::new(local, remote)); + assert_eq!(pathway, Pathway::new(local, remote)); } #[test] diff --git a/qtraversal/src/route.rs b/qtraversal/src/route.rs index f4ca8cef..921c4dc7 100644 --- a/qtraversal/src/route.rs +++ b/qtraversal/src/route.rs @@ -14,7 +14,7 @@ use bytes::BytesMut; use qbase::{ net::{ addr::EndpointAddr, - route::{Line, Link, Route}, + route::{Line, Link, Pathway, Route}, }, util::ArcAsyncDeque, }; @@ -30,10 +30,9 @@ use smallvec::SmallVec; use tokio_util::task::AbortOnDropHandle; use tracing::Instrument as _; -pub type ArcRecvQueue = ArcAsyncDeque<(BytesMut, PathWay, Link)>; +pub type ArcRecvQueue = ArcAsyncDeque<(BytesMut, Pathway, Link)>; use crate::{ - PathWay, nat::{ client::{StunClients, StunClientsComponent}, router::{StunRouter, StunRouterComponent}, diff --git a/qudp/examples/receive.rs b/qudp/examples/receive.rs index aadb21af..9f34baae 100644 --- a/qudp/examples/receive.rs +++ b/qudp/examples/receive.rs @@ -1,5 +1,7 @@ +use bytes::BytesMut; use clap::Parser; -use qudp::UdpSocket; +use qbase::net::route::Line; +use qudp::{BATCH_SIZE, UdpSocket}; #[derive(Parser, Debug)] #[command(version, about, long_about = None)] @@ -18,16 +20,24 @@ async fn main() { let addr = args.bind.parse().unwrap(); let socket = UdpSocket::bind(addr).expect("failed to create socket"); - let mut receiver = socket.receiver(); + let mut iovecs: Vec = (0..BATCH_SIZE) + .map(|_| { + let mut buf = BytesMut::with_capacity(1500); + buf.resize(1500, 0); + buf + }) + .collect(); + let mut lines: Vec = (0..BATCH_SIZE).map(|_| Line::default()).collect(); + loop { - match receiver.recv().await { + match socket.receive(&mut iovecs, &mut lines).await { Ok(n) => { tracing::info!( target: "qudp", packets = n, - dst = %receiver.lines[0].dst, - src = %receiver.lines[0].src, - len = receiver.lines[0].seg_size, + dst = %lines[0].dst, + src = %lines[0].src, + len = lines[0].seg_size, "received packets" ); } diff --git a/qudp/src/ext.rs b/qudp/src/ext.rs new file mode 100644 index 00000000..81155c3b --- /dev/null +++ b/qudp/src/ext.rs @@ -0,0 +1,95 @@ +use std::{ + io::{self, IoSliceMut}, + net::SocketAddr, + sync::Arc, + task::{Context, Poll, ready}, +}; + +use bytes::BytesMut; +use qbase::{ + net::route::{Line, Link, Pathway, Route}, + util::Wakers, +}; + +pub struct UdpSocket { + send_wakers: Arc>, + io: io::Result, +} + +impl UdpSocket { + pub fn bind(addr: SocketAddr) -> Self { + UdpSocket { + send_wakers: Arc::new(Wakers::new()), + io: super::UdpSocket::bind(addr), + } + } + + fn socket(&self) -> io::Result<&super::UdpSocket> { + self.io + .as_ref() + .map_err(|e| io::Error::new(e.kind(), e.to_string())) + } + + pub fn local_addr(&self) -> io::Result { + self.socket()?.local_addr() + } + + pub fn max_segments(&self) -> io::Result { + Ok(super::BATCH_SIZE) + } + + pub fn max_segment_size(&self) -> io::Result { + Ok(1500) + } + + pub fn poll_send( + &self, + cx: &mut Context, + pkts: &[io::IoSlice], + route: Route, + ) -> Poll> { + let io = self.socket()?; + let waker = cx.waker(); + let waker_group = self.send_wakers.together_with(waker); + let cx = &mut Context::from_waker(&waker_group); + + debug_assert_eq!(route.ecn(), None); + let result = io.poll_send(cx, pkts, &route); + if result.is_ready() { + self.send_wakers.remove(waker); + } + result + } + + pub fn poll_recv( + &self, + cx: &mut Context, + pkts: &mut [BytesMut], + route: &mut [Route], + ) -> Poll> { + let io = self.socket()?; + let dst = io.local_addr()?; + let len = route.len().min(pkts.len()); + let mut rcvd_lines = Vec::with_capacity(len); + rcvd_lines.resize_with(route.len(), Line::default); + let mut bufs = pkts[..len] + .iter_mut() + .map(|p| IoSliceMut::new(p.as_mut())) + .collect::>(); + debug_assert_eq!(rcvd_lines.len(), bufs.len()); + let nrcvd = ready!(io.poll_recv(cx, &mut bufs, &mut rcvd_lines))?; + + for (idx, mut line) in rcvd_lines.into_iter().take(nrcvd).enumerate() { + let pathway = Pathway::new(line.link.src.into(), dst.into()); + line.link = Link::new(line.src, io.local_addr()?).flip(); + route[idx] = Route::new(pathway.flip(), line); + } + + Poll::Ready(Ok(nrcvd)) + } + + pub fn close(&mut self) { + self.io = Err(io::Error::new(io::ErrorKind::NotFound, "socket closed")); + self.send_wakers.wake_all(); + } +} diff --git a/qudp/src/lib.rs b/qudp/src/lib.rs index 1b711ba1..f71959da 100644 --- a/qudp/src/lib.rs +++ b/qudp/src/lib.rs @@ -25,6 +25,8 @@ cfg_if::cfg_if! { } } +pub mod ext; + #[derive(Debug)] pub struct UdpSocket { io: tokio::net::UdpSocket, @@ -191,25 +193,19 @@ impl UdpSocket { } } - pub fn receiver(&self) -> Receiver<'_> { - Receiver { + pub fn receive<'a>(&'a self, iovecs: &'a mut [BytesMut], lines: &'a mut [Line]) -> Receive<'a> { + Receive { socket: self, - iovecs: (0..BATCH_SIZE) - .map(|_| { - let mut buf = BytesMut::with_capacity(1500); - buf.resize(1500, 0); - buf - }) - .collect::>(), - lines: (0..BATCH_SIZE).map(|_| Line::default()).collect::>(), + iovecs, + lines, } } } pub struct Send<'a> { - pub socket: &'a UdpSocket, - pub iovecs: &'a [IoSlice<'a>], - pub line: Line, + socket: &'a UdpSocket, + iovecs: &'a [IoSlice<'a>], + line: Line, } impl Future for Send<'_> { @@ -221,27 +217,23 @@ impl Future for Send<'_> { } } -pub struct Receiver<'u> { - pub socket: &'u UdpSocket, - pub iovecs: Vec, - pub lines: Vec, +pub struct Receive<'a> { + socket: &'a UdpSocket, + iovecs: &'a mut [BytesMut], + lines: &'a mut [Line], } -impl Receiver<'_> { - #[inline] - pub fn poll_recv(&mut self, cx: &mut Context) -> Poll> { - let mut bufs = self +impl Future for Receive<'_> { + type Output = io::Result; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + let mut bufs = this .iovecs .iter_mut() .map(|b| IoSliceMut::new(b)) .collect::>(); - - self.socket.poll_recv(cx, &mut bufs, &mut self.lines) - } - - #[inline] - pub async fn recv(&mut self) -> io::Result { - core::future::poll_fn(|cx| self.poll_recv(cx)).await + this.socket.poll_recv(cx, &mut bufs, this.lines) } } @@ -257,23 +249,28 @@ mod tests { #[tokio::test(flavor = "current_thread")] async fn ipv6_wildcard_socket_does_not_receive_ipv4_packets() -> io::Result<()> { - let socket = UdpSocket::bind(SocketAddr::V6(SocketAddrV6::new( + let socket4 = UdpSocket::bind(SocketAddr::V6(SocketAddrV6::new( Ipv6Addr::UNSPECIFIED, 0, 0, 0, )))?; - let port = socket.local_addr()?.port(); + let port = socket4.local_addr()?.port(); - let sender = + let socket6 = std::net::UdpSocket::bind(SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)))?; - sender.send_to( + socket6.send_to( b"ping", SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port)), )?; - let mut receiver = socket.receiver(); - let result = tokio::time::timeout(Duration::from_millis(200), receiver.recv()).await; + let mut iovecs = [BytesMut::with_capacity(1500); 1]; + let mut lines = [Line::default(); 1]; + let result = tokio::time::timeout( + Duration::from_millis(200), + socket4.receive(&mut iovecs, &mut lines), + ) + .await; assert!( result.is_err(), "unexpected ipv4 datagram arrived on ipv6 wildcard socket: {result:?}"