pb_mapper_protocol/
transport.rs

1//! Relay carriers. Dispatch once before forwarding; control I/O may use the enum.
2use serde::{Deserialize, Serialize};
3use std::{
4    io,
5    net::SocketAddr,
6    pin::Pin,
7    task::{Context, Poll},
8};
9use tokio::{
10    io::{AsyncRead, AsyncWrite, ReadBuf},
11    net::TcpStream,
12};
13
14/// Budget for setting a carrier up: QUIC discovery over TCP and the
15/// handshake after it, on the client; admission of it, on the relay.
16pub const CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
17
18#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq, Serialize, Deserialize)]
19#[serde(rename_all = "snake_case")]
20pub enum RelayTransport {
21    #[default]
22    Tcp,
23    Quic,
24}
25
26#[derive(Debug)]
27pub enum RelayStream {
28    Tcp(TcpStream),
29    #[cfg(feature = "relay-quic")]
30    Quic(Box<crate::quic::QuicStream>),
31}
32
33impl From<TcpStream> for RelayStream {
34    fn from(stream: TcpStream) -> Self {
35        Self::Tcp(stream)
36    }
37}
38
39impl RelayStream {
40    pub fn relay_addr(&self) -> io::Result<SocketAddr> {
41        match self {
42            Self::Tcp(stream) => stream.peer_addr(),
43            #[cfg(feature = "relay-quic")]
44            Self::Quic(stream) => Ok(stream.relay_addr()),
45        }
46    }
47    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
48        match self {
49            Self::Tcp(stream) => stream.peer_addr(),
50            #[cfg(feature = "relay-quic")]
51            Self::Quic(stream) => Ok(stream.peer_addr()),
52        }
53    }
54    pub fn configure(&self, keep_alive: bool) -> io::Result<()> {
55        match self {
56            Self::Tcp(stream) => {
57                let keep_alive_result = if keep_alive {
58                    uni_stream::stream::set_tcp_keep_alive(stream)
59                } else {
60                    Ok(())
61                };
62                // Keep Nagle disabled even if the platform rejects keep-alive.
63                stream.set_nodelay(true)?;
64                keep_alive_result
65            }
66            #[cfg(feature = "relay-quic")]
67            Self::Quic(_) => Ok(()),
68        }
69    }
70    pub fn into_split(self) -> (RelayReadHalf, RelayWriteHalf) {
71        match self {
72            Self::Tcp(stream) => {
73                let (read, write) = stream.into_split();
74                (RelayReadHalf::Tcp(read), RelayWriteHalf::Tcp(write))
75            }
76            #[cfg(feature = "relay-quic")]
77            Self::Quic(stream) => {
78                let (read, write) = stream.into_split();
79                (RelayReadHalf::Quic(read), RelayWriteHalf::Quic(write))
80            }
81        }
82    }
83}
84
85pub enum RelayReadHalf {
86    Tcp(tokio::net::tcp::OwnedReadHalf),
87    #[cfg(feature = "relay-quic")]
88    Quic(crate::quic::QuicReadHalf),
89}
90pub enum RelayWriteHalf {
91    Tcp(tokio::net::tcp::OwnedWriteHalf),
92    #[cfg(feature = "relay-quic")]
93    Quic(crate::quic::QuicWriteHalf),
94}
95
96macro_rules! read_impl {
97    ($ty:ty) => {
98        impl AsyncRead for $ty {
99            fn poll_read(
100                mut self: Pin<&mut Self>,
101                cx: &mut Context<'_>,
102                buf: &mut ReadBuf<'_>,
103            ) -> Poll<io::Result<()>> {
104                match &mut *self {
105                    Self::Tcp(stream) => Pin::new(stream).poll_read(cx, buf),
106                    #[cfg(feature = "relay-quic")]
107                    Self::Quic(stream) => Pin::new(stream).poll_read(cx, buf),
108                }
109            }
110        }
111    };
112}
113macro_rules! write_impl {
114    ($ty:ty) => {
115        impl AsyncWrite for $ty {
116            fn poll_write(
117                mut self: Pin<&mut Self>,
118                cx: &mut Context<'_>,
119                buf: &[u8],
120            ) -> Poll<io::Result<usize>> {
121                match &mut *self {
122                    Self::Tcp(stream) => Pin::new(stream).poll_write(cx, buf),
123                    #[cfg(feature = "relay-quic")]
124                    Self::Quic(stream) => Pin::new(stream).poll_write(cx, buf),
125                }
126            }
127            fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
128                match &mut *self {
129                    Self::Tcp(stream) => Pin::new(stream).poll_flush(cx),
130                    #[cfg(feature = "relay-quic")]
131                    Self::Quic(stream) => Pin::new(stream).poll_flush(cx),
132                }
133            }
134            fn poll_shutdown(
135                mut self: Pin<&mut Self>,
136                cx: &mut Context<'_>,
137            ) -> Poll<io::Result<()>> {
138                match &mut *self {
139                    Self::Tcp(stream) => Pin::new(stream).poll_shutdown(cx),
140                    #[cfg(feature = "relay-quic")]
141                    Self::Quic(stream) => Pin::new(stream).poll_shutdown(cx),
142                }
143            }
144        }
145    };
146}
147read_impl!(RelayStream);
148read_impl!(RelayReadHalf);
149write_impl!(RelayStream);
150write_impl!(RelayWriteHalf);
151
152/// Monomorphize the data path for the selected carrier. No per-I/O enum match.
153#[cfg(feature = "relay-quic")]
154#[macro_export]
155macro_rules! with_relay_stream {
156    ($stream:expr, $name:ident, $body:expr) => {
157        match $stream {
158            $crate::transport::RelayStream::Tcp($name) => $body,
159            $crate::transport::RelayStream::Quic($name) => $body,
160        }
161    };
162}
163
164/// TCP-only builds erase carrier dispatch entirely, including in downstream crates.
165#[cfg(not(feature = "relay-quic"))]
166#[macro_export]
167macro_rules! with_relay_stream {
168    ($stream:expr, $name:ident, $body:expr) => {
169        match $stream {
170            $crate::transport::RelayStream::Tcp($name) => $body,
171        }
172    };
173}
pb_mapper_protocol/transport.rs

pb_mapper_protocol/
transport.rs
