nex_socket/tcp/
sync_impl.rs

1use socket2::{Domain, Protocol, Socket, Type as SockType};
2use std::io;
3use std::net::{SocketAddr, TcpListener, TcpStream};
4use std::time::Duration;
5
6use crate::tcp::TcpConfig;
7
8#[cfg(unix)]
9use std::os::fd::AsRawFd;
10
11#[cfg(unix)]
12use nix::poll::{PollFd, PollFlags, PollTimeout, poll};
13
14/// Low level synchronous TCP socket.
15#[derive(Debug)]
16pub struct TcpSocket {
17    socket: Socket,
18    nonblocking: bool,
19}
20
21impl TcpSocket {
22    /// Build a socket according to `TcpSocketConfig`.
23    pub fn from_config(config: &TcpConfig) -> io::Result<Self> {
24        config.validate()?;
25
26        let socket = Socket::new(
27            config.socket_family.to_domain(),
28            config.socket_type.to_sock_type(),
29            Some(Protocol::TCP),
30        )?;
31
32        socket.set_nonblocking(config.nonblocking)?;
33
34        // Set socket options based on configuration
35        if let Some(flag) = config.reuseaddr {
36            socket.set_reuse_address(flag)?;
37        }
38        #[cfg(any(
39            target_os = "android",
40            target_os = "dragonfly",
41            target_os = "freebsd",
42            target_os = "fuchsia",
43            target_os = "ios",
44            target_os = "linux",
45            target_os = "macos",
46            target_os = "netbsd",
47            target_os = "openbsd",
48            target_os = "tvos",
49            target_os = "visionos",
50            target_os = "watchos"
51        ))]
52        if let Some(flag) = config.reuseport {
53            socket.set_reuse_port(flag)?;
54        }
55        if let Some(flag) = config.nodelay {
56            socket.set_tcp_nodelay(flag)?;
57        }
58        if let Some(dur) = config.linger {
59            socket.set_linger(Some(dur))?;
60        }
61        if let Some(ttl) = config.ttl {
62            socket.set_ttl_v4(ttl)?;
63        }
64        if let Some(hoplimit) = config.hoplimit {
65            socket.set_unicast_hops_v6(hoplimit)?;
66        }
67        if let Some(keepalive) = config.keepalive {
68            socket.set_keepalive(keepalive)?;
69        }
70        if let Some(timeout) = config.read_timeout {
71            socket.set_read_timeout(Some(timeout))?;
72        }
73        if let Some(timeout) = config.write_timeout {
74            socket.set_write_timeout(Some(timeout))?;
75        }
76        if let Some(size) = config.recv_buffer_size {
77            socket.set_recv_buffer_size(size)?;
78        }
79        if let Some(size) = config.send_buffer_size {
80            socket.set_send_buffer_size(size)?;
81        }
82        if let Some(tos) = config.tos {
83            socket.set_tos_v4(tos)?;
84        }
85        crate::apply_tclass_v6(&socket, config.tclass_v6)?;
86        if let Some(only_v6) = config.only_v6 {
87            socket.set_only_v6(only_v6)?;
88        }
89
90        // Linux: optional interface name
91        #[cfg(any(target_os = "linux", target_os = "android", target_os = "fuchsia"))]
92        if let Some(iface) = &config.bind_device {
93            socket.bind_device(Some(iface.as_bytes()))?;
94        }
95
96        // bind to the specified address if provided
97        if let Some(addr) = config.bind_addr {
98            socket.bind(&addr.into())?;
99        }
100
101        Ok(Self {
102            socket,
103            nonblocking: config.nonblocking,
104        })
105    }
106
107    /// Create a socket of arbitrary type (STREAM or RAW).
108    pub fn new(domain: Domain, sock_type: SockType) -> io::Result<Self> {
109        let socket = Socket::new(domain, sock_type, Some(Protocol::TCP))?;
110        socket.set_nonblocking(false)?;
111        Ok(Self {
112            socket,
113            nonblocking: false,
114        })
115    }
116
117    /// Convenience constructor for an IPv4 STREAM socket.
118    pub fn v4_stream() -> io::Result<Self> {
119        Self::new(Domain::IPV4, SockType::STREAM)
120    }
121
122    /// Convenience constructor for an IPv6 STREAM socket.
123    pub fn v6_stream() -> io::Result<Self> {
124        Self::new(Domain::IPV6, SockType::STREAM)
125    }
126
127    /// IPv4 RAW TCP. Requires administrator privileges.
128    pub fn raw_v4() -> io::Result<Self> {
129        Self::new(Domain::IPV4, SockType::RAW)
130    }
131
132    /// IPv6 RAW TCP. Requires administrator privileges.
133    pub fn raw_v6() -> io::Result<Self> {
134        Self::new(Domain::IPV6, SockType::RAW)
135    }
136
137    /// Bind the socket to a specific address.
138    pub fn bind(&self, addr: SocketAddr) -> io::Result<()> {
139        self.socket.bind(&addr.into())
140    }
141
142    /// Connect to a remote address.
143    pub fn connect(&self, addr: SocketAddr) -> io::Result<()> {
144        self.socket.connect(&addr.into())
145    }
146
147    /// Connect to the target address with a timeout and return the connected stream.
148    ///
149    /// The returned `TcpStream` must be used for subsequent I/O.
150    #[cfg(unix)]
151    pub fn connect_timeout(&self, target: SocketAddr, timeout: Duration) -> io::Result<TcpStream> {
152        let socket = self.socket.try_clone()?;
153        socket.set_nonblocking(true)?;
154        let raw_fd = socket.as_raw_fd();
155
156        // Try to connect first
157        match socket.connect(&target.into()) {
158            Ok(_) => { /* succeeded immediately */ }
159            Err(err)
160                if err.kind() == io::ErrorKind::WouldBlock
161                    || err.raw_os_error() == Some(libc::EINPROGRESS) =>
162            {
163                // Continue waiting
164            }
165            Err(e) => return Err(e),
166        }
167
168        // Wait for the connection using poll
169        use std::os::unix::io::BorrowedFd;
170        // SAFETY: `raw_fd` belongs to `socket` and remains valid for this scope;
171        // BorrowedFd does not take ownership.
172        let mut fds = [PollFd::new(
173            unsafe { BorrowedFd::borrow_raw(raw_fd) },
174            PollFlags::POLLOUT,
175        )];
176        let poll_timeout = PollTimeout::try_from(timeout).unwrap_or(PollTimeout::MAX);
177        let n = poll(&mut fds, poll_timeout)?;
178
179        if n == 0 {
180            return Err(io::Error::new(io::ErrorKind::TimedOut, "connect timed out"));
181        }
182
183        // Check the result with `SO_ERROR`
184        let err: i32 = socket
185            .take_error()?
186            .map(|e| e.raw_os_error().unwrap_or(0))
187            .unwrap_or(0);
188        if err != 0 {
189            return Err(io::Error::from_raw_os_error(err));
190        }
191
192        socket.set_nonblocking(self.nonblocking)?;
193
194        match socket.try_clone() {
195            Ok(cloned_socket) => {
196                // Convert the socket into a `std::net::TcpStream`
197                let std_stream: TcpStream = cloned_socket.into();
198                Ok(std_stream)
199            }
200            Err(e) => Err(e),
201        }
202    }
203
204    /// Connect to the target address with a timeout and return the connected stream.
205    ///
206    /// The returned `TcpStream` must be used for subsequent I/O.
207    #[cfg(windows)]
208    pub fn connect_timeout(&self, target: SocketAddr, timeout: Duration) -> io::Result<TcpStream> {
209        use std::mem::size_of;
210        use std::os::windows::io::AsRawSocket;
211        use windows_sys::Win32::Networking::WinSock::{
212            POLLWRNORM, SO_ERROR, SOCKET, SOCKET_ERROR, SOL_SOCKET, WSAPOLLFD, WSAPoll, getsockopt,
213        };
214
215        let socket = self.socket.try_clone()?;
216        socket.set_nonblocking(true)?;
217        let sock = socket.as_raw_socket() as SOCKET;
218
219        // Start connect
220        match socket.connect(&target.into()) {
221            Ok(_) => { /* connection succeeded immediately */ }
222            Err(e) if e.kind() == io::ErrorKind::WouldBlock || e.raw_os_error() == Some(10035) /* WSAEWOULDBLOCK */ => {}
223            Err(e) => return Err(e),
224        }
225
226        // Wait using WSAPoll until writable
227        let mut fds = [WSAPOLLFD {
228            fd: sock,
229            events: POLLWRNORM,
230            revents: 0,
231        }];
232
233        let timeout_ms = timeout.as_millis().clamp(0, i32::MAX as u128) as i32;
234        // SAFETY: `fds` is writable for the supplied element count and remains
235        // live throughout WSAPoll.
236        let result = unsafe { WSAPoll(fds.as_mut_ptr(), fds.len() as u32, timeout_ms) };
237        if result == SOCKET_ERROR {
238            return Err(io::Error::last_os_error());
239        } else if result == 0 {
240            return Err(io::Error::new(io::ErrorKind::TimedOut, "connect timed out"));
241        }
242
243        // Check for errors via `SO_ERROR`
244        let mut so_error: i32 = 0;
245        let mut optlen = size_of::<i32>() as i32;
246        // SAFETY: `so_error` and `optlen` are writable for the duration of
247        // getsockopt and `sock` is open.
248        let ret = unsafe {
249            getsockopt(
250                sock,
251                SOL_SOCKET,
252                SO_ERROR,
253                &mut so_error as *mut _ as *mut _,
254                &mut optlen,
255            )
256        };
257
258        if ret == SOCKET_ERROR || so_error != 0 {
259            return Err(io::Error::from_raw_os_error(so_error));
260        }
261
262        socket.set_nonblocking(self.nonblocking)?;
263
264        let std_stream: TcpStream = socket.into();
265        Ok(std_stream)
266    }
267
268    /// Start listening for incoming connections.
269    pub fn listen(&self, backlog: i32) -> io::Result<()> {
270        self.socket.listen(backlog)
271    }
272
273    /// Accept an incoming connection.
274    pub fn accept(&self) -> io::Result<(TcpStream, SocketAddr)> {
275        let (stream, addr) = self.socket.accept()?;
276        let address = addr.as_socket().ok_or_else(|| {
277            io::Error::new(
278                io::ErrorKind::InvalidData,
279                "accepted peer did not provide an IP socket address",
280            )
281        })?;
282        Ok((stream.into(), address))
283    }
284
285    /// Convert the socket into a `TcpStream`.
286    pub fn to_tcp_stream(self) -> io::Result<TcpStream> {
287        Ok(self.socket.into())
288    }
289
290    /// Convert the socket into a `TcpListener`.
291    pub fn to_tcp_listener(self) -> io::Result<TcpListener> {
292        Ok(self.socket.into())
293    }
294
295    /// Send a raw packet (for RAW TCP use).
296    pub fn send_to(&self, buf: &[u8], target: SocketAddr) -> io::Result<usize> {
297        self.socket.send_to(buf, &target.into())
298    }
299
300    /// Receive a raw packet (for RAW TCP use).
301    pub fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
302        // SAFETY: `MaybeUninit<u8>` is layout-compatible with `u8`, and the
303        // slice preserves the original buffer's length and lifetime.
304        let buf_maybe = unsafe {
305            std::slice::from_raw_parts_mut(
306                buf.as_mut_ptr() as *mut std::mem::MaybeUninit<u8>,
307                buf.len(),
308            )
309        };
310
311        let (n, addr) = self.socket.recv_from(buf_maybe)?;
312        let addr = addr
313            .as_socket()
314            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "invalid address format"))?;
315
316        Ok((n, addr))
317    }
318
319    /// Shutdown the socket.
320    pub fn shutdown(&self, how: std::net::Shutdown) -> io::Result<()> {
321        self.socket.shutdown(how)
322    }
323
324    /// Set the socket to reuse the address.
325    pub fn set_reuseaddr(&self, on: bool) -> io::Result<()> {
326        self.socket.set_reuse_address(on)
327    }
328
329    /// Get the socket address reuse option.
330    pub fn reuseaddr(&self) -> io::Result<bool> {
331        self.socket.reuse_address()
332    }
333
334    /// Set the socket port reuse option where supported.
335    #[cfg(any(
336        target_os = "android",
337        target_os = "dragonfly",
338        target_os = "freebsd",
339        target_os = "fuchsia",
340        target_os = "ios",
341        target_os = "linux",
342        target_os = "macos",
343        target_os = "netbsd",
344        target_os = "openbsd",
345        target_os = "tvos",
346        target_os = "visionos",
347        target_os = "watchos"
348    ))]
349    pub fn set_reuseport(&self, on: bool) -> io::Result<()> {
350        self.socket.set_reuse_port(on)
351    }
352
353    /// Get the socket port reuse option where supported.
354    #[cfg(any(
355        target_os = "android",
356        target_os = "dragonfly",
357        target_os = "freebsd",
358        target_os = "fuchsia",
359        target_os = "ios",
360        target_os = "linux",
361        target_os = "macos",
362        target_os = "netbsd",
363        target_os = "openbsd",
364        target_os = "tvos",
365        target_os = "visionos",
366        target_os = "watchos"
367    ))]
368    pub fn reuseport(&self) -> io::Result<bool> {
369        self.socket.reuse_port()
370    }
371
372    /// Set the socket to not delay packets.
373    pub fn set_nodelay(&self, on: bool) -> io::Result<()> {
374        self.socket.set_tcp_nodelay(on)
375    }
376
377    /// Get the no delay option.
378    pub fn nodelay(&self) -> io::Result<bool> {
379        self.socket.tcp_nodelay()
380    }
381
382    /// Set the linger option for the socket.
383    pub fn set_linger(&self, dur: Option<Duration>) -> io::Result<()> {
384        self.socket.set_linger(dur)
385    }
386
387    /// Set the time-to-live for IPv4 packets.
388    pub fn set_ttl(&self, ttl: u32) -> io::Result<()> {
389        self.socket.set_ttl_v4(ttl)
390    }
391
392    /// Get the time-to-live for IPv4 packets.
393    pub fn ttl(&self) -> io::Result<u32> {
394        self.socket.ttl_v4()
395    }
396
397    /// Set the hop limit for IPv6 packets.
398    pub fn set_hoplimit(&self, hops: u32) -> io::Result<()> {
399        self.socket.set_unicast_hops_v6(hops)
400    }
401
402    /// Get the hop limit for IPv6 packets.
403    pub fn hoplimit(&self) -> io::Result<u32> {
404        self.socket.unicast_hops_v6()
405    }
406
407    /// Set the keepalive option for the socket.
408    pub fn set_keepalive(&self, on: bool) -> io::Result<()> {
409        self.socket.set_keepalive(on)
410    }
411
412    /// Get the keepalive option.
413    pub fn keepalive(&self) -> io::Result<bool> {
414        self.socket.keepalive()
415    }
416
417    /// Set the receive buffer size.
418    pub fn set_recv_buffer_size(&self, size: usize) -> io::Result<()> {
419        self.socket.set_recv_buffer_size(size)
420    }
421
422    /// Get the receive buffer size.
423    pub fn recv_buffer_size(&self) -> io::Result<usize> {
424        self.socket.recv_buffer_size()
425    }
426
427    /// Set the send buffer size.
428    pub fn set_send_buffer_size(&self, size: usize) -> io::Result<()> {
429        self.socket.set_send_buffer_size(size)
430    }
431
432    /// Get the send buffer size.
433    pub fn send_buffer_size(&self) -> io::Result<usize> {
434        self.socket.send_buffer_size()
435    }
436
437    /// Set IPv4 TOS / DSCP.
438    pub fn set_tos(&self, tos: u32) -> io::Result<()> {
439        self.socket.set_tos_v4(tos)
440    }
441
442    /// Get IPv4 TOS / DSCP.
443    pub fn tos(&self) -> io::Result<u32> {
444        self.socket.tos_v4()
445    }
446
447    /// Set IPv6 traffic class where supported.
448    #[cfg(any(
449        target_os = "android",
450        target_os = "dragonfly",
451        target_os = "freebsd",
452        target_os = "fuchsia",
453        target_os = "linux",
454        target_os = "macos",
455        target_os = "netbsd",
456        target_os = "openbsd"
457    ))]
458    pub fn set_tclass_v6(&self, tclass: u32) -> io::Result<()> {
459        self.socket.set_tclass_v6(tclass)
460    }
461
462    /// Get IPv6 traffic class where supported.
463    #[cfg(any(
464        target_os = "android",
465        target_os = "dragonfly",
466        target_os = "freebsd",
467        target_os = "fuchsia",
468        target_os = "linux",
469        target_os = "macos",
470        target_os = "netbsd",
471        target_os = "openbsd"
472    ))]
473    pub fn tclass_v6(&self) -> io::Result<u32> {
474        self.socket.tclass_v6()
475    }
476
477    /// Set whether this socket is IPv6 only.
478    pub fn set_only_v6(&self, only_v6: bool) -> io::Result<()> {
479        self.socket.set_only_v6(only_v6)
480    }
481
482    /// Get whether this socket is IPv6 only.
483    pub fn only_v6(&self) -> io::Result<bool> {
484        self.socket.only_v6()
485    }
486
487    /// Set the bind device for the socket (Linux specific).
488    pub fn set_bind_device(&self, iface: &str) -> io::Result<()> {
489        #[cfg(any(target_os = "linux", target_os = "android", target_os = "fuchsia"))]
490        return self.socket.bind_device(Some(iface.as_bytes()));
491
492        #[cfg(not(any(target_os = "linux", target_os = "android", target_os = "fuchsia")))]
493        {
494            let _ = iface;
495            Err(io::Error::new(
496                io::ErrorKind::Unsupported,
497                "bind_device is not supported on this platform",
498            ))
499        }
500    }
501
502    /// Retrieve the local address of the socket.
503    pub fn local_addr(&self) -> io::Result<SocketAddr> {
504        self.socket
505            .local_addr()?
506            .as_socket()
507            .ok_or_else(|| io::Error::other("failed to retrieve local address"))
508    }
509
510    /// Extract the RAW file descriptor for Unix.
511    #[cfg(unix)]
512    pub fn as_raw_fd(&self) -> std::os::unix::io::RawFd {
513        use std::os::fd::AsRawFd;
514        self.socket.as_raw_fd()
515    }
516
517    /// Extract the RAW socket handle for Windows.
518    #[cfg(windows)]
519    pub fn as_raw_socket(&self) -> std::os::windows::io::RawSocket {
520        use std::os::windows::io::AsRawSocket;
521        self.socket.as_raw_socket()
522    }
523
524    /// Construct from a raw `socket2::Socket`.
525    pub fn from_socket(socket: Socket) -> Self {
526        Self {
527            socket,
528            // `socket2::Socket` does not expose a portable getter for the current
529            // blocking mode, so externally supplied sockets default to blocking
530            // expectations in this synchronous wrapper.
531            nonblocking: false,
532        }
533    }
534
535    /// Borrow the inner `socket2::Socket`.
536    pub fn socket(&self) -> &Socket {
537        &self.socket
538    }
539
540    /// Consume and return the inner `socket2::Socket`.
541    pub fn into_socket(self) -> Socket {
542        self.socket
543    }
544}
545
546#[cfg(test)]
547mod tests {
548    #[cfg(unix)]
549    use super::*;
550    #[cfg(unix)]
551    use libc::{F_GETFL, O_NONBLOCK, fcntl};
552    #[cfg(unix)]
553    use std::net::TcpListener as StdTcpListener;
554
555    #[cfg(unix)]
556    fn socket_is_nonblocking(socket: &Socket) -> bool {
557        // SAFETY: The descriptor belongs to the borrowed live socket and F_GETFL
558        // neither takes ownership nor retains pointers.
559        let flags = unsafe { fcntl(socket.as_raw_fd(), F_GETFL) };
560        assert!(flags >= 0, "F_GETFL failed: {}", io::Error::last_os_error());
561        (flags & O_NONBLOCK) != 0
562    }
563
564    #[cfg(unix)]
565    #[test]
566    fn connect_timeout_does_not_mutate_original_nonblocking_state_after_invalid_input() {
567        let sock = TcpSocket::v4_stream().expect("socket");
568        sock.socket.set_nonblocking(true).expect("set nonblocking");
569
570        let result = sock.connect_timeout("[::1]:80".parse().unwrap(), Duration::from_secs(1));
571        assert!(result.is_err());
572        assert!(socket_is_nonblocking(&sock.socket));
573    }
574
575    #[cfg(unix)]
576    #[test]
577    fn connect_timeout_does_not_mutate_original_blocking_state_after_success() {
578        let listener = StdTcpListener::bind("127.0.0.1:0").expect("listener");
579        let addr = listener.local_addr().expect("local addr");
580        let handle = std::thread::spawn(move || listener.accept().expect("accept"));
581
582        let sock = TcpSocket::v4_stream().expect("socket");
583        sock.socket.set_nonblocking(false).expect("set blocking");
584        let _stream = sock
585            .connect_timeout(addr, Duration::from_secs(1))
586            .expect("connect");
587
588        assert!(!socket_is_nonblocking(&sock.socket));
589        let _ = handle.join();
590    }
591}
nex_socket/tcp/sync_impl.rs

nex_socket/tcp/
sync_impl.rs
