cloudpub_common/transport/
mod.rs

1use crate::config::TransportConfig;
2use crate::constants::{DEFAULT_KEEPALIVE_INTERVAL, DEFAULT_KEEPALIVE_SECS};
3use crate::utils::to_socket_addr;
4use anyhow::{Context, Result};
5use async_trait::async_trait;
6use std::fmt::{Debug, Display};
7#[cfg(unix)]
8use std::os::fd::RawFd;
9use std::time::Duration;
10use tokio::io::{AsyncRead, AsyncWrite};
11use tracing::error;
12
13use crate::protocol::message::Message as ProtocolMessage;
14
15#[async_trait]
16pub trait ProtobufStream {
17    async fn recv_message(&mut self) -> anyhow::Result<Option<ProtocolMessage>>;
18    async fn send_message(&mut self, msg: &ProtocolMessage) -> anyhow::Result<()>;
19    async fn close(&mut self) -> anyhow::Result<()>;
20}
21
22#[cfg(unix)]
23use anyhow::bail;
24
25mod tcp;
26pub use tcp::{Listener, NamedSocketAddr, SocketAddr, Stream, TcpTransport};
27
28mod websocket;
29pub use websocket::{WebsocketStream, WebsocketTransport};
30
31#[cfg(feature = "rustls")]
32pub mod rustls;
33#[cfg(feature = "rustls")]
34use rustls as tls;
35#[cfg(feature = "rustls")]
36pub use tls::TlsTransport;
37
38#[derive(Clone)]
39pub struct AddrMaybeCached {
40    pub addr: String,
41    pub socket_addr: Option<NamedSocketAddr>,
42}
43
44impl AddrMaybeCached {
45    pub fn new(addr: &str) -> AddrMaybeCached {
46        AddrMaybeCached {
47            addr: addr.to_string(),
48            socket_addr: None,
49        }
50    }
51
52    pub async fn resolve(&mut self) -> Result<()> {
53        match to_socket_addr(&self.addr).await {
54            Ok(s) => {
55                self.socket_addr = Some(NamedSocketAddr::Inet(s));
56                Ok(())
57            }
58            Err(e) => Err(e),
59        }
60    }
61}
62
63impl Display for AddrMaybeCached {
64    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65        match self.socket_addr.as_ref() {
66            Some(s) => f.write_fmt(format_args!("{}", s)),
67            None => f.write_str(&self.addr),
68        }
69    }
70}
71
72/// Specify a transport layer, like TCP, TLS
73#[async_trait]
74pub trait Transport: Debug + Send + Sync {
75    type Acceptor: Send + Sync;
76    type RawStream: Send + Sync;
77    type Stream: 'static + AsyncRead + AsyncWrite + ProtobufStream + Unpin + Send + Sync + Debug;
78
79    fn new(config: &TransportConfig) -> Result<Self>
80    where
81        Self: Sized;
82    /// Get the stream id, which is used to identify the transport layer
83    #[cfg(unix)]
84    fn as_raw_fd(conn: &Self::Stream) -> RawFd;
85    /// Provide the transport with socket options, which can be handled at the need of the transport
86    fn hint(conn: &Self::Stream, opts: SocketOpts);
87    async fn bind(&self, addr: NamedSocketAddr) -> Result<Self::Acceptor>;
88    /// accept must be cancel safe
89    async fn accept(&self, a: &Self::Acceptor) -> Result<(Self::RawStream, SocketAddr)>;
90    async fn handshake(&self, conn: Self::RawStream) -> Result<Self::Stream>;
91    async fn connect(&self, addr: &AddrMaybeCached) -> Result<Self::Stream>;
92
93    fn get_header(&self, _name: &str) -> Option<String> {
94        None
95    }
96}
97
98#[derive(Debug, Clone, Copy)]
99pub struct Keepalive {
100    // tcp_keepalive_time if the underlying protocol is TCP
101    pub keepalive_secs: u64,
102    // tcp_keepalive_intvl if the underlying protocol is TCP
103    pub keepalive_interval: u64,
104}
105
106#[derive(Debug, Clone, Copy)]
107pub struct SocketOpts {
108    // None means do not change
109    pub nodelay: Option<bool>,
110    // keepalive must be Some or None at the same time, or the behavior will be platform-dependent
111    pub keepalive: Option<Keepalive>,
112    // SO_PRIORITY
113    pub priority: Option<u8>,
114}
115
116impl Default for Keepalive {
117    fn default() -> Self {
118        Keepalive {
119            keepalive_secs: DEFAULT_KEEPALIVE_SECS,
120            keepalive_interval: DEFAULT_KEEPALIVE_INTERVAL,
121        }
122    }
123}
124
125impl SocketOpts {
126    /// Socket options for the control channel
127    pub fn for_control_channel() -> SocketOpts {
128        SocketOpts {
129            nodelay: Some(true), // Always set nodelay for the control channel
130            keepalive: Some(Keepalive::default()),
131            priority: Some(0), // Set high priority for the control channel
132        }
133    }
134
135    pub fn for_data_channel() -> SocketOpts {
136        SocketOpts {
137            nodelay: Some(true), // Always set nodelay for the data channel
138            keepalive: Some(Keepalive::default()),
139            priority: Some(0),
140        }
141    }
142}
143
144/// Bound a dial with a deadline. `secs` of 0 means no deadline.
145///
146/// The steps this guards - TCP connect, TLS handshake, WebSocket upgrade -
147/// exchange no protocol messages, so `MESSAGE_TIMEOUT_SECS` never applies to
148/// them, and a peer that accepts the connection and then stops responding
149/// would otherwise park the caller indefinitely.
150pub(crate) async fn with_connect_deadline<F, T>(secs: u64, what: &str, fut: F) -> Result<T>
151where
152    F: std::future::Future<Output = Result<T>>,
153{
154    if secs == 0 {
155        return fut.await;
156    }
157    match tokio::time::timeout(Duration::from_secs(secs), fut).await {
158        Ok(result) => result,
159        Err(_) => Err(anyhow::anyhow!("{} timed out after {} seconds", what, secs)),
160    }
161}
162
163#[cfg(unix)]
164pub fn set_reuse(s: &dyn std::os::fd::AsRawFd) -> Result<()> {
165    use libc;
166    use std::{io, mem};
167    unsafe {
168        let optval: libc::c_int = 1;
169        let ret = libc::setsockopt(
170            s.as_raw_fd(),
171            libc::SOL_SOCKET,
172            libc::SO_REUSEPORT | libc::SO_REUSEADDR,
173            &optval as *const _ as *const libc::c_void,
174            mem::size_of_val(&optval) as libc::socklen_t,
175        );
176        if ret != 0 {
177            bail!("Set sock option failed: {:?}", io::Error::last_os_error());
178        }
179    }
180    Ok(())
181}
182
183#[cfg(target_os = "linux")]
184pub fn set_low_latency(s: &dyn std::os::fd::AsRawFd) -> Result<()> {
185    use libc;
186    use std::{io, mem};
187
188    unsafe {
189        let fd = s.as_raw_fd();
190        // 1. TCP_NODELAY - Disable Nagle's algorithm for immediate packet sending
191        let nodelay: libc::c_int = 1;
192        let ret = libc::setsockopt(
193            fd,
194            libc::IPPROTO_TCP,
195            libc::TCP_NODELAY,
196            &nodelay as *const _ as *const libc::c_void,
197            mem::size_of_val(&nodelay) as libc::socklen_t,
198        );
199        if ret != 0 {
200            bail!(
201                "Failed to set TCP_NODELAY: {:?}",
202                io::Error::last_os_error()
203            );
204        }
205
206        // 2. TCP_QUICKACK - Enable quick ACK mode to reduce ACK delay
207        let quickack: libc::c_int = 1;
208        let ret = libc::setsockopt(
209            fd,
210            libc::IPPROTO_TCP,
211            libc::TCP_QUICKACK,
212            &quickack as *const _ as *const libc::c_void,
213            mem::size_of_val(&quickack) as libc::socklen_t,
214        );
215        if ret != 0 {
216            bail!(
217                "Failed to set TCP_QUICKACK: {:?}",
218                io::Error::last_os_error()
219            );
220        }
221    }
222    Ok(())
223}
224
225// Set socket priority: 0 - lowest (default), 7 - higest
226#[cfg(target_os = "linux")]
227pub fn set_priority(s: &dyn std::os::fd::AsRawFd, priority: libc::c_int) -> Result<()> {
228    use libc;
229    use std::{io, mem};
230
231    unsafe {
232        let fd = s.as_raw_fd();
233
234        // 1. SO_PRIORITY - Set high priority for the socket
235        let ret = libc::setsockopt(
236            fd,
237            libc::SOL_SOCKET,
238            libc::SO_PRIORITY,
239            &priority as *const _ as *const libc::c_void,
240            mem::size_of_val(&priority) as libc::socklen_t,
241        );
242        if ret != 0 {
243            bail!(
244                "Failed to set SO_PRIORITY: {:?}",
245                io::Error::last_os_error()
246            );
247        }
248    }
249
250    Ok(())
251}
252
253impl SocketOpts {
254    pub fn apply(&self, conn: &Stream) {
255        if let Some(v) = self.keepalive {
256            let keepalive_duration = Duration::from_secs(v.keepalive_secs);
257            let keepalive_interval = Duration::from_secs(v.keepalive_interval);
258            if let Err(e) = tcp::try_set_tcp_keepalive(conn, keepalive_duration, keepalive_interval)
259                .with_context(|| "Failed to set keepalive")
260            {
261                error!("{:#}", e);
262            }
263        }
264
265        match conn {
266            Stream::Tcp(conn) => {
267                #[cfg(unix)]
268                if let Err(e) = set_reuse(conn) {
269                    error!("{:#}", e);
270                }
271                if let Some(nodelay) = self.nodelay {
272                    #[cfg(not(target_os = "linux"))]
273                    if let Err(e) = conn
274                        .set_nodelay(nodelay)
275                        .with_context(|| "Failed to set nodelay")
276                    {
277                        error!("{:#}", e);
278                    }
279                    #[cfg(target_os = "linux")]
280                    if nodelay {
281                        if let Err(e) = set_low_latency(conn) {
282                            error!("Failed to set low latency options: {:#}", e);
283                        }
284                    }
285                }
286                #[cfg(target_os = "linux")]
287                if let Some(priority) = self.priority {
288                    if let Err(e) = set_priority(conn, priority as libc::c_int) {
289                        error!("Failed to set socket priority: {:#}", e);
290                    }
291                }
292            }
293            #[cfg(unix)]
294            Stream::Unix(_conn) =>
295            {
296                #[cfg(target_os = "linux")]
297                if let Some(priority) = self.priority {
298                    if let Err(e) = set_priority(_conn, priority as libc::c_int) {
299                        error!("Failed to set socket priority: {:#}", e);
300                    }
301                }
302            }
303        }
304    }
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310    use crate::config::{TcpConfig, TransportType, WebsocketConfig};
311    use crate::constants::{DEFAULT_KEEPALIVE_INTERVAL, DEFAULT_KEEPALIVE_SECS};
312    use std::time::Instant;
313    use tokio::net::TcpListener;
314
315    const TEST_CONNECT_TIMEOUT_SECS: u64 = 1;
316    /// Well above the connect timeout, so a hung connect() fails the test
317    /// instead of hanging the whole suite.
318    const TEST_GIVE_UP: Duration = Duration::from_secs(15);
319
320    /// A peer that completes the TCP handshake and then never says anything.
321    /// This is what a dead path behind a stateful middlebox looks like: the
322    /// socket is ESTABLISHED, so neither TLS nor the WebSocket upgrade ever
323    /// gets a reply, and nothing below us times out on its own.
324    async fn silent_peer() -> (String, tokio::task::JoinHandle<()>) {
325        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
326        let addr = listener.local_addr().unwrap().to_string();
327        let handle = tokio::spawn(async move {
328            let _held = listener.accept().await;
329            std::future::pending::<()>().await;
330        });
331        (addr, handle)
332    }
333
334    fn transport_config(tls: bool, connect_timeout_secs: u64) -> TransportConfig {
335        TransportConfig {
336            transport_type: TransportType::Websocket,
337            tcp: TcpConfig {
338                connect_timeout_secs,
339                ..Default::default()
340            },
341            tls: Some(Default::default()),
342            websocket: Some(WebsocketConfig { tls }),
343        }
344    }
345
346    async fn resolved(addr: &str) -> AddrMaybeCached {
347        let mut remote = AddrMaybeCached::new(addr);
348        remote.resolve().await.unwrap();
349        remote
350    }
351
352    #[tokio::test]
353    async fn websocket_connect_times_out_on_silent_peer() {
354        let (addr, _peer) = silent_peer().await;
355        let transport =
356            WebsocketTransport::new(&transport_config(false, TEST_CONNECT_TIMEOUT_SECS)).unwrap();
357        let remote = resolved(&addr).await;
358
359        let started = Instant::now();
360        let result = tokio::time::timeout(TEST_GIVE_UP, transport.connect(&remote))
361            .await
362            .expect("connect() hung on the WebSocket upgrade instead of timing out");
363
364        assert!(
365            result.is_err(),
366            "connect() must fail against a peer that never completes the upgrade"
367        );
368        assert!(
369            started.elapsed() < TEST_GIVE_UP,
370            "connect() took {:?}, expected roughly {}s",
371            started.elapsed(),
372            TEST_CONNECT_TIMEOUT_SECS
373        );
374    }
375
376    #[cfg(feature = "rustls")]
377    #[tokio::test]
378    async fn tls_connect_times_out_on_silent_peer() {
379        let (addr, _peer) = silent_peer().await;
380        let transport =
381            TlsTransport::new(&transport_config(true, TEST_CONNECT_TIMEOUT_SECS)).unwrap();
382        let remote = resolved(&addr).await;
383
384        let started = Instant::now();
385        let result = tokio::time::timeout(TEST_GIVE_UP, transport.connect(&remote))
386            .await
387            .expect("connect() hung on the TLS handshake instead of timing out");
388
389        assert!(
390            result.is_err(),
391            "connect() must fail against a peer that never answers the ClientHello"
392        );
393        assert!(
394            started.elapsed() < TEST_GIVE_UP,
395            "connect() took {:?}, expected roughly {}s",
396            started.elapsed(),
397            TEST_CONNECT_TIMEOUT_SECS
398        );
399    }
400
401    #[tokio::test]
402    async fn zero_connect_timeout_disables_the_deadline() {
403        let (addr, _peer) = silent_peer().await;
404        let transport = WebsocketTransport::new(&transport_config(false, 0)).unwrap();
405        let remote = resolved(&addr).await;
406
407        assert!(
408            tokio::time::timeout(Duration::from_secs(2), transport.connect(&remote))
409                .await
410                .is_err(),
411            "a connect timeout of 0 must mean 'wait forever'"
412        );
413    }
414
415    /// The control channel carries no traffic between heartbeats, so a
416    /// half-open socket is only detected if the kernel probes it.
417    #[tokio::test]
418    async fn control_channel_socket_opts_enable_tcp_keepalive() {
419        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
420        let addr = listener.local_addr().unwrap();
421        let (client, _accepted) = tokio::join!(
422            async { tokio::net::TcpStream::connect(addr).await.unwrap() },
423            async { listener.accept().await.unwrap() }
424        );
425
426        let stream = Stream::Tcp(client);
427        SocketOpts::for_control_channel().apply(&stream);
428
429        let Stream::Tcp(ref tcp) = stream else {
430            unreachable!("bound a TCP listener")
431        };
432        let sock = socket2::SockRef::from(tcp);
433        assert!(sock.keepalive().unwrap(), "SO_KEEPALIVE must be set");
434        assert_eq!(
435            sock.tcp_keepalive_time().unwrap(),
436            Duration::from_secs(DEFAULT_KEEPALIVE_SECS)
437        );
438        assert_eq!(
439            sock.tcp_keepalive_interval().unwrap(),
440            Duration::from_secs(DEFAULT_KEEPALIVE_INTERVAL)
441        );
442    }
443
444    #[tokio::test]
445    async fn data_channel_socket_opts_enable_tcp_keepalive() {
446        assert!(SocketOpts::for_data_channel().keepalive.is_some());
447    }
448}
cloudpub_common/transport/mod.rs

cloudpub_common/transport/
mod.rs
