Skip to main content

wasmtime_wasi/sockets/
tcp.rs

1use crate::runtime::with_ambient_tokio_runtime;
2use crate::sockets::{
3    ErrorCode, MaybeReady, SocketAddrCheck, SocketAddrUse, SocketAddressFamily, WasiSocketsCtx,
4    get_receive_buffer_size, get_send_buffer_size, get_unicast_hop_limit, is_valid_address_family,
5    is_valid_remote_address, is_valid_unicast_address, set_receive_buffer_size,
6    set_send_buffer_size, set_unicast_hop_limit, unspecified_addr,
7};
8use rustix::fd::AsFd;
9use rustix::io::Errno;
10use rustix::net::sockopt;
11use std::fmt::Debug;
12use std::future::poll_fn;
13use std::mem;
14use std::net::SocketAddr;
15use std::sync::Arc;
16use std::task::{Poll, ready};
17use std::time::Duration;
18
19/// Value taken from rust std library.
20const DEFAULT_BACKLOG: u32 = 128;
21
22const NANOS_PER_SEC: u64 = 1_000_000_000;
23
24/// The state of a TCP socket.
25///
26/// This represents the various states a socket can be in during the
27/// activities of listening, accepting, and connecting.
28enum TcpState {
29    /// The initial state for a newly-created socket.
30    ///
31    /// The socket may be bound to a local address in this state, but doesn't
32    /// have to.
33    ///
34    /// From here a socket can transition to `Listening` or `Connecting`.
35    Default(tokio::net::TcpSocket),
36
37    /// The socket is now listening and waiting for an incoming connection.
38    ///
39    /// Sockets will not leave this state.
40    Listening(Arc<tokio::net::TcpListener>),
41
42    /// An outgoing connection is started.
43    ///
44    /// This is created via the `start_connect` method. The payload is a future
45    /// for the eventual result of the connect.
46    ///
47    /// From here a socket can transition to `Connected` or `Closed`.
48    Connecting(MaybeReady<Result<tokio::net::TcpStream, ErrorCode>>),
49
50    /// A connection has been established.
51    ///
52    /// This is created either via `finish_connect` or for freshly accepted
53    /// sockets from a TCP listener.
54    ///
55    /// A socket will not transition out of this state.
56    Connected {
57        stream: Arc<tokio::net::TcpStream>,
58        receive_taken: bool,
59        send_taken: bool,
60    },
61
62    /// The socket is closed and no more operations can be performed.
63    Closed(ErrorCode),
64}
65impl TcpState {
66    fn connected(stream: tokio::net::TcpStream) -> Self {
67        TcpState::Connected {
68            stream: Arc::new(stream),
69            receive_taken: false,
70            send_taken: false,
71        }
72    }
73    fn take(&mut self) -> Self {
74        mem::replace(self, TcpState::Closed(ErrorCode::Other))
75    }
76}
77impl Debug for TcpState {
78    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        match self {
80            Self::Default(_) => f.debug_tuple("Default").finish(),
81            Self::Listening { .. } => f.debug_tuple("Listening").finish(),
82            Self::Connecting(..) => f.debug_tuple("Connecting").finish(),
83            Self::Connected { .. } => f.debug_tuple("Connected").finish(),
84            Self::Closed(..) => write!(f, "Closed"),
85        }
86    }
87}
88
89/// A host TCP socket, plus associated bookkeeping.
90pub struct TcpSocket {
91    /// The current state in the bind/listen/accept/connect progression.
92    tcp_state: TcpState,
93
94    /// The desired listen queue size.
95    listen_backlog_size: u32,
96
97    family: SocketAddressFamily,
98
99    /// The checks to perform before doing any noteworthy syscall.
100    permissions: SocketAddrCheck,
101
102    /// Persisted socket options to manually apply to newly accepted client
103    /// sockets on platforms that don't inherit socket options from the listener.
104    listener_options: NonInheritedOptions,
105
106    /// Cached value of whether the socket is bound. Various methods use the
107    /// `.is_bound()` method, so we cache it to avoid redundant syscalls.
108    is_bound: bool,
109}
110
111impl TcpSocket {
112    /// Create a new socket in the given family.
113    pub(crate) fn new(
114        ctx: &WasiSocketsCtx,
115        family: SocketAddressFamily,
116    ) -> Result<Self, ErrorCode> {
117        ctx.allowed_network_uses.check_allowed_tcp()?;
118
119        let socket = with_ambient_tokio_runtime(|| socket(family))?;
120
121        Ok(Self {
122            tcp_state: TcpState::Default(socket),
123            listen_backlog_size: DEFAULT_BACKLOG,
124            family,
125            is_bound: false,
126            listener_options: Default::default(),
127            permissions: ctx.socket_addr_check.clone(),
128        })
129    }
130
131    fn as_fd(&self) -> Result<rustix::fd::BorrowedFd<'_>, ErrorCode> {
132        match &self.tcp_state {
133            TcpState::Default(socket) => Ok(socket.as_fd()),
134            TcpState::Connected { stream, .. } => Ok(stream.as_fd()),
135            TcpState::Listening(listener) => Ok(listener.as_fd()),
136            TcpState::Connecting(..) => Err(ErrorCode::InvalidState),
137            TcpState::Closed(err) => Err(*err),
138        }
139    }
140
141    pub(crate) fn is_bound(&mut self) -> bool {
142        // Once bound, a TCP socket can never become unbound again. So we can
143        // skip all work after a previous call has already determined the
144        // socket to be bound.
145        if !self.is_bound {
146            self.is_bound = match &self.tcp_state {
147                TcpState::Default(socket) => socket
148                    .local_addr()
149                    .is_ok_and(|addr| addr != unspecified_addr(self.family)),
150                _ => true,
151            };
152        }
153        self.is_bound
154    }
155
156    pub(crate) async fn bind(&mut self, addr: SocketAddr) -> Result<(), ErrorCode> {
157        if self.is_bound() {
158            return Err(ErrorCode::InvalidState);
159        }
160        let TcpState::Default(sock) = &self.tcp_state else {
161            return Err(ErrorCode::InvalidState);
162        };
163
164        if !is_valid_unicast_address(addr.ip()) || !is_valid_address_family(addr.ip(), self.family)
165        {
166            return Err(ErrorCode::InvalidArgument);
167        }
168
169        self.permissions.check(addr, SocketAddrUse::TcpBind).await?;
170        bind(sock, addr)?;
171        Ok(())
172    }
173
174    pub(crate) fn start_connect(&mut self, addr: SocketAddr) -> Result<(), ErrorCode> {
175        let TcpState::Default(_) = &self.tcp_state else {
176            return Err(ErrorCode::InvalidState);
177        };
178
179        let permissions = self.permissions.clone();
180        let family = self.family;
181        let already_bound = self.is_bound();
182
183        if !is_valid_unicast_address(addr.ip())
184            || !is_valid_remote_address(addr)
185            || !is_valid_address_family(addr.ip(), family)
186        {
187            return Err(ErrorCode::InvalidArgument);
188        };
189
190        let TcpState::Default(sock) = self.tcp_state.take() else {
191            unreachable!();
192        };
193
194        self.tcp_state = TcpState::Connecting(MaybeReady::new(async move {
195            // Perform all checks before doing any syscalls.
196            {
197                if !already_bound {
198                    // If not explicitly bound, the OS will implicitly bind the
199                    // socket to an ephemeral port when connecting. Unlike other
200                    // operations (e.g. `listen`), we will *not* do the implicit
201                    // bind ourselves because that may accelerate port exhaustion.
202                    // For more info, see IP_BIND_ADDRESS_NO_PORT (Linux) or
203                    // SO_REUSE_UNICASTPORT (Windows).
204                    //
205                    // Instead we check the permission to bind, but not perform
206                    // the actual bind:
207                    let implicit = unspecified_addr(family);
208                    permissions.check(implicit, SocketAddrUse::TcpBind).await?;
209                }
210
211                permissions.check(addr, SocketAddrUse::TcpConnect).await?;
212            }
213
214            let stream = sock.connect(addr).await?;
215            Ok(stream)
216        }));
217
218        Ok(())
219    }
220
221    pub(crate) fn poll_finish_connect(
222        &mut self,
223        cx: &mut std::task::Context<'_>,
224    ) -> Poll<Result<(), ErrorCode>> {
225        match &mut self.tcp_state {
226            TcpState::Connecting(connect) => {
227                ready!(with_ambient_tokio_runtime(|| connect.poll_ready(cx)));
228            }
229            TcpState::Connected { .. } => return Poll::Ready(Ok(())),
230            TcpState::Closed(e) => return Poll::Ready(Err(*e)),
231            _ => return Poll::Ready(Err(ErrorCode::InvalidState)),
232        }
233        let TcpState::Connecting(connect) = self.tcp_state.take() else {
234            unreachable!();
235        };
236
237        match connect.unwrap_ready() {
238            Ok(stream) => {
239                self.tcp_state = TcpState::connected(stream);
240                Poll::Ready(Ok(()))
241            }
242            Err(err) => {
243                self.tcp_state = TcpState::Closed(err);
244                Poll::Ready(Err(err))
245            }
246        }
247    }
248
249    pub(crate) async fn listen(&mut self) -> Result<TcpListenStream, ErrorCode> {
250        let already_bound = self.is_bound();
251        let sock = match self.tcp_state.take() {
252            TcpState::Default(sock) => sock,
253            tcp_state => {
254                self.tcp_state = tcp_state;
255                return Err(ErrorCode::InvalidState);
256            }
257        };
258
259        // Perform all checks before doing any syscalls.
260        {
261            if already_bound {
262                self.permissions
263                    .check(sock.local_addr()?, SocketAddrUse::TcpListen)
264                    .await?;
265            } else {
266                let implicit = unspecified_addr(self.family);
267                self.permissions
268                    .check(implicit, SocketAddrUse::TcpBind)
269                    .await?;
270                self.permissions
271                    .check(implicit, SocketAddrUse::TcpListen)
272                    .await?;
273            }
274        }
275
276        // Some platforms automatically perform an implicit bind as part of
277        // the `listen` syscall. However this is not ubiquitous behavior:
278        // - Linux mentions it in their docs [0] that they perform an
279        //   implicit bind. This behavior has been experimentally verified.
280        // - Windows requires a `bind` before `listen`. This is both
281        //   documented [1] and experimentally verified.
282        // - Other platforms (e.g. macOS, FreeBSD) do not explicitly
283        //   document it either way and instead leave it up to the
284        //   individual protocol to decide [2]. However, experiments
285        //   show that MacOS in fact _does_ perform an implicit bind.
286        //
287        // Thus to ensure consistent behavior across all platforms, we
288        // perform the implicit bind ourselves here for unbound sockets.
289        //
290        // [0]: https://man7.org/linux/man-pages/man7/ip.7.html
291        // > An ephemeral port is allocated to a socket in the following
292        // > circumstances: (...) listen(2) is called on a stream socket
293        // > that was not previously bound;
294        //
295        // [1]: https://learn.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-listen
296        // > WSAEINVAL: The socket has not been bound with bind.
297        //
298        // [2]: https://pubs.opengroup.org/onlinepubs/9699919799/functions/listen.html
299        // > EDESTADDRREQ: The socket is not bound to a local address,
300        // > and the protocol does not support listening on an unbound
301        // > socket.
302        if !already_bound {
303            let implicit = unspecified_addr(self.family);
304            bind(&sock, implicit)?;
305        }
306
307        let listener = sock.listen(self.listen_backlog_size).map_err(|err| {
308            match Errno::from_io_error(&err) {
309                // See: https://learn.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-listen#:~:text=WSAEMFILE
310                // According to the docs, `listen` can return EMFILE on Windows.
311                // This is odd, because we're not trying to create a new socket
312                // or file descriptor of any kind. So we rewrite it to less
313                // surprising error code.
314                //
315                // At the time of writing, this behavior has never been experimentally
316                // observed by any of the wasmtime authors, so we're relying fully
317                // on Microsoft's documentation here.
318                #[cfg(windows)]
319                Some(Errno::MFILE) => Errno::NOBUFS.into(),
320
321                _ => err,
322            }
323        })?;
324        let listener = Arc::new(listener);
325        self.tcp_state = TcpState::Listening(listener.clone());
326
327        Ok(TcpListenStream {
328            inner: listener,
329            listener_options: self.listener_options.clone(),
330            family: self.family,
331            permissions: self.permissions.clone(),
332            pending_accept: None,
333        })
334    }
335
336    pub(crate) fn take_send_stream(&mut self) -> Result<TcpSendStream, ErrorCode> {
337        match &mut self.tcp_state {
338            TcpState::Connected {
339                stream, send_taken, ..
340            } if !*send_taken => {
341                *send_taken = true;
342                Ok(TcpSendStream {
343                    inner: stream.clone(),
344                })
345            }
346            TcpState::Closed(err) => Err(*err),
347            _ => Err(ErrorCode::InvalidState),
348        }
349    }
350
351    pub(crate) fn take_receive_stream(&mut self) -> Result<TcpReceiveStream, ErrorCode> {
352        match &mut self.tcp_state {
353            TcpState::Connected {
354                stream,
355                receive_taken,
356                ..
357            } if !*receive_taken => {
358                *receive_taken = true;
359                Ok(TcpReceiveStream {
360                    inner: stream.clone(),
361                })
362            }
363            TcpState::Closed(err) => Err(*err),
364            _ => Err(ErrorCode::InvalidState),
365        }
366    }
367
368    pub(crate) fn local_address(&mut self) -> Result<SocketAddr, ErrorCode> {
369        if !self.is_bound() {
370            return Err(ErrorCode::InvalidState);
371        }
372
373        match &self.tcp_state {
374            TcpState::Default(socket) => Ok(socket.local_addr()?),
375            TcpState::Connecting(_) => Err(ErrorCode::InvalidState),
376            TcpState::Connected { stream, .. } => Ok(stream.local_addr()?),
377            TcpState::Listening(listener) => Ok(listener.local_addr()?),
378            TcpState::Closed(err) => Err(*err),
379        }
380    }
381
382    pub(crate) fn remote_address(&self) -> Result<SocketAddr, ErrorCode> {
383        match &self.tcp_state {
384            TcpState::Connected { stream, .. } => Ok(stream.peer_addr()?),
385            TcpState::Closed(err) => Err(*err),
386            _ => Err(ErrorCode::InvalidState),
387        }
388    }
389
390    pub(crate) fn is_listening(&self) -> bool {
391        matches!(self.tcp_state, TcpState::Listening(_))
392    }
393
394    pub(crate) fn address_family(&self) -> SocketAddressFamily {
395        self.family
396    }
397
398    pub(crate) fn set_listen_backlog_size(&mut self, value: u64) -> Result<(), ErrorCode> {
399        const MIN_BACKLOG: u32 = 1;
400        const MAX_BACKLOG: u32 = i32::MAX as u32; // OS'es will most likely limit it down even further.
401
402        if value == 0 {
403            return Err(ErrorCode::InvalidArgument);
404        }
405        // Silently clamp backlog size. This is OK for us to do, because operating systems do this too.
406        let value = value
407            .try_into()
408            .unwrap_or(MAX_BACKLOG)
409            .clamp(MIN_BACKLOG, MAX_BACKLOG);
410        match &self.tcp_state {
411            TcpState::Default(..) => {
412                // Socket not listening yet. Stash value for first invocation to `listen`.
413                self.listen_backlog_size = value;
414                Ok(())
415            }
416            TcpState::Listening(listener) => {
417                // Try to update the backlog by calling `listen` again.
418                // Not all platforms support this. We'll only update our own value if the OS supports changing the backlog size after the fact.
419                if rustix::net::listen(&listener, value.try_into().unwrap_or(i32::MAX)).is_err() {
420                    return Err(ErrorCode::NotSupported);
421                }
422                self.listen_backlog_size = value;
423                Ok(())
424            }
425            TcpState::Closed(err) => Err(*err),
426            _ => Err(ErrorCode::InvalidState),
427        }
428    }
429
430    pub(crate) fn keep_alive_enabled(&self) -> Result<bool, ErrorCode> {
431        let fd = self.as_fd()?;
432        let v = sockopt::socket_keepalive(fd)?;
433        Ok(v)
434    }
435
436    pub(crate) fn set_keep_alive_enabled(&self, value: bool) -> Result<(), ErrorCode> {
437        let fd = self.as_fd()?;
438        sockopt::set_socket_keepalive(fd, value)?;
439        Ok(())
440    }
441
442    pub(crate) fn keep_alive_idle_time(&self) -> Result<u64, ErrorCode> {
443        let fd = self.as_fd()?;
444        let v = sockopt::tcp_keepidle(fd)?;
445        Ok(v.as_nanos().try_into().unwrap_or(u64::MAX))
446    }
447
448    pub(crate) fn set_keep_alive_idle_time(&mut self, value: u64) -> Result<(), ErrorCode> {
449        if value == 0 {
450            // WIT: "If the provided value is 0, an `invalid-argument` error is returned."
451            return Err(ErrorCode::InvalidArgument);
452        }
453        let fd = self.as_fd()?;
454        let value = clamp_keep_alive_time(value);
455        sockopt::set_tcp_keepidle(fd, Duration::from_nanos(value))?;
456        self.listener_options.set_keep_alive_idle_time(value);
457        Ok(())
458    }
459
460    pub(crate) fn keep_alive_interval(&self) -> Result<u64, ErrorCode> {
461        let fd = self.as_fd()?;
462        let v = sockopt::tcp_keepintvl(fd)?;
463        Ok(v.as_nanos().try_into().unwrap_or(u64::MAX))
464    }
465
466    pub(crate) fn set_keep_alive_interval(&self, value: u64) -> Result<(), ErrorCode> {
467        if value == 0 {
468            // WIT: "If the provided value is 0, an `invalid-argument` error is returned."
469            return Err(ErrorCode::InvalidArgument);
470        }
471        let fd = self.as_fd()?;
472        let value = clamp_keep_alive_time(value);
473        sockopt::set_tcp_keepintvl(fd, Duration::from_nanos(value))?;
474        Ok(())
475    }
476
477    pub(crate) fn keep_alive_count(&self) -> Result<u32, ErrorCode> {
478        let fd = self.as_fd()?;
479        let v = sockopt::tcp_keepcnt(fd)?;
480        Ok(v)
481    }
482
483    pub(crate) fn set_keep_alive_count(&self, value: u32) -> Result<(), ErrorCode> {
484        if value == 0 {
485            // WIT: "If the provided value is 0, an `invalid-argument` error is returned."
486            return Err(ErrorCode::InvalidArgument);
487        }
488        let value = clamp_keep_alive_count(value);
489        let fd = self.as_fd()?;
490        sockopt::set_tcp_keepcnt(fd, value)?;
491        Ok(())
492    }
493
494    pub(crate) fn hop_limit(&self) -> Result<u8, ErrorCode> {
495        let fd = self.as_fd()?;
496        let n = get_unicast_hop_limit(fd, self.family)?;
497        Ok(n)
498    }
499
500    pub(crate) fn set_hop_limit(&mut self, value: u8) -> Result<(), ErrorCode> {
501        {
502            let fd = self.as_fd()?;
503            set_unicast_hop_limit(fd, self.family, value)?;
504        }
505        self.listener_options.set_hop_limit(value);
506        Ok(())
507    }
508
509    pub(crate) fn receive_buffer_size(&self) -> Result<u64, ErrorCode> {
510        let fd = self.as_fd()?;
511        let n = get_receive_buffer_size(fd)?;
512        Ok(n)
513    }
514
515    pub(crate) fn set_receive_buffer_size(&mut self, value: u64) -> Result<(), ErrorCode> {
516        let res = {
517            let fd = self.as_fd()?;
518            set_receive_buffer_size(fd, value)?
519        };
520        self.listener_options.set_receive_buffer_size(res);
521        Ok(())
522    }
523
524    pub(crate) fn send_buffer_size(&self) -> Result<u64, ErrorCode> {
525        let fd = self.as_fd()?;
526        let n = get_send_buffer_size(fd)?;
527        Ok(n)
528    }
529
530    pub(crate) fn set_send_buffer_size(&mut self, value: u64) -> Result<(), ErrorCode> {
531        let res = {
532            let fd = self.as_fd()?;
533            set_send_buffer_size(fd, value)?
534        };
535        self.listener_options.set_send_buffer_size(res);
536        Ok(())
537    }
538}
539
540pub(crate) struct TcpListenStream {
541    inner: Arc<tokio::net::TcpListener>,
542    family: SocketAddressFamily,
543    listener_options: NonInheritedOptions,
544    permissions: SocketAddrCheck,
545    pending_accept: Option<MaybeReady<Result<tokio::net::TcpStream, ErrorCode>>>,
546}
547impl TcpListenStream {
548    pub(crate) fn poll_accept(&mut self, cx: &mut std::task::Context<'_>) -> Poll<TcpSocket> {
549        ready!(self.poll_ready(cx));
550        let result = self.pending_accept.take().unwrap().unwrap_ready();
551        Poll::Ready(TcpSocket {
552            tcp_state: match result {
553                Ok(client) => {
554                    self.listener_options.apply(self.family, &client);
555                    TcpState::connected(client)
556                }
557                Err(err) => TcpState::Closed(err),
558            },
559            listen_backlog_size: DEFAULT_BACKLOG,
560            family: self.family,
561            is_bound: true,
562            listener_options: Default::default(),
563            permissions: self.permissions.clone(),
564        })
565    }
566
567    pub(crate) fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<()> {
568        if self.pending_accept.is_none() {
569            let listener = self.inner.clone();
570            let permissions = self.permissions.clone();
571
572            self.pending_accept = Some(MaybeReady::new(async move {
573                loop {
574                    match accept(&listener).await {
575                        Ok((client, addr)) => {
576                            if permissions
577                                .check(addr, SocketAddrUse::TcpAccept)
578                                .await
579                                .is_ok()
580                            {
581                                return Ok(client);
582                            } else {
583                                reset(client);
584                                continue;
585                            }
586                        }
587                        Err(err) => {
588                            return Err(err.into());
589                        }
590                    }
591                }
592            }));
593        }
594
595        with_ambient_tokio_runtime(|| {
596            self.pending_accept
597                .as_mut()
598                .unwrap()
599                .poll_ready(cx)
600                .map(|_| ())
601        })
602    }
603}
604
605pub(crate) struct TcpSendStream {
606    inner: Arc<tokio::net::TcpStream>,
607}
608impl TcpSendStream {
609    pub(crate) fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<()> {
610        self.inner.poll_write_ready(cx).map(|_| ())
611    }
612
613    pub(crate) fn poll_write(
614        &mut self,
615        cx: &mut std::task::Context<'_>,
616        buf: &[u8],
617    ) -> Poll<Result<usize, ErrorCode>> {
618        loop {
619            return match self.inner.try_write(buf) {
620                Ok(n) => Poll::Ready(Ok(n)),
621                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
622                    match self.inner.poll_write_ready(cx) {
623                        Poll::Ready(Ok(())) => continue,
624                        Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())),
625                        Poll::Pending => Poll::Pending,
626                    }
627                }
628                Err(e) => Poll::Ready(Err(match Errno::from_io_error(&e) {
629                    #[cfg(windows)]
630                    Some(Errno::SHUTDOWN) | Some(Errno::CONNABORTED) => ErrorCode::ConnectionBroken,
631                    #[cfg(not(windows))]
632                    Some(Errno::PIPE) => ErrorCode::ConnectionBroken,
633
634                    _ => e.into(),
635                })),
636            };
637        }
638    }
639    pub(crate) async fn write(&mut self, buf: &[u8]) -> Result<usize, ErrorCode> {
640        poll_fn(|cx| self.poll_write(cx, buf)).await
641    }
642}
643impl Drop for TcpSendStream {
644    fn drop(&mut self) {
645        _ = rustix::net::shutdown(&self.inner, rustix::net::Shutdown::Write);
646    }
647}
648
649pub(crate) struct TcpReceiveStream {
650    inner: Arc<tokio::net::TcpStream>,
651}
652impl TcpReceiveStream {
653    pub(crate) fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<()> {
654        self.inner.poll_read_ready(cx).map(|_| ())
655    }
656
657    pub(crate) fn poll_read(
658        &mut self,
659        cx: &mut std::task::Context<'_>,
660        buf: &mut [u8],
661    ) -> Poll<Result<usize, ErrorCode>> {
662        if buf.is_empty() {
663            return Poll::Ready(Ok(0));
664        }
665        loop {
666            return match self.inner.try_read(buf) {
667                Ok(0) => Poll::Ready(Ok(0)),
668                Ok(n) => Poll::Ready(Ok(n)),
669                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
670                    match self.inner.poll_read_ready(cx) {
671                        Poll::Ready(Ok(())) => continue,
672                        Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())),
673                        Poll::Pending => Poll::Pending,
674                    }
675                }
676                Err(e) => Poll::Ready(Err(e.into())),
677            };
678        }
679    }
680}
681impl Drop for TcpReceiveStream {
682    fn drop(&mut self) {
683        _ = rustix::net::shutdown(&self.inner, rustix::net::Shutdown::Read);
684    }
685}
686
687#[cfg(not(target_os = "macos"))]
688pub use inherits_option::*;
689#[cfg(not(target_os = "macos"))]
690mod inherits_option {
691    use crate::sockets::SocketAddressFamily;
692    use tokio::net::TcpStream;
693
694    #[derive(Default, Clone)]
695    pub struct NonInheritedOptions;
696
697    impl NonInheritedOptions {
698        pub fn set_keep_alive_idle_time(&mut self, _value: u64) {}
699
700        pub fn set_hop_limit(&mut self, _value: u8) {}
701
702        pub fn set_receive_buffer_size(&mut self, _value: usize) {}
703
704        pub fn set_send_buffer_size(&mut self, _value: usize) {}
705
706        pub(crate) fn apply(&self, _family: SocketAddressFamily, _stream: &TcpStream) {}
707    }
708}
709
710#[cfg(target_os = "macos")]
711pub use does_not_inherit_options::*;
712#[cfg(target_os = "macos")]
713mod does_not_inherit_options {
714    use crate::sockets::SocketAddressFamily;
715    use rustix::net::sockopt;
716    use std::sync::Arc;
717    use std::sync::atomic::{AtomicU8, AtomicU64, AtomicUsize, Ordering::Relaxed};
718    use std::time::Duration;
719    use tokio::net::TcpStream;
720
721    // The socket options below are not automatically inherited from the listener
722    // on all platforms. So we keep track of which options have been explicitly
723    // set and manually apply those values to newly accepted clients.
724    #[derive(Default, Clone)]
725    pub struct NonInheritedOptions(Arc<Inner>);
726
727    #[derive(Default)]
728    struct Inner {
729        receive_buffer_size: AtomicUsize,
730        send_buffer_size: AtomicUsize,
731        hop_limit: AtomicU8,
732        keep_alive_idle_time: AtomicU64, // nanoseconds
733    }
734
735    impl NonInheritedOptions {
736        pub fn set_keep_alive_idle_time(&mut self, value: u64) {
737            self.0.keep_alive_idle_time.store(value, Relaxed);
738        }
739
740        pub fn set_hop_limit(&mut self, value: u8) {
741            self.0.hop_limit.store(value, Relaxed);
742        }
743
744        pub fn set_receive_buffer_size(&mut self, value: usize) {
745            self.0.receive_buffer_size.store(value, Relaxed);
746        }
747
748        pub fn set_send_buffer_size(&mut self, value: usize) {
749            self.0.send_buffer_size.store(value, Relaxed);
750        }
751
752        pub(crate) fn apply(&self, family: SocketAddressFamily, stream: &TcpStream) {
753            // Manually inherit socket options from listener. We only have to
754            // do this on platforms that don't already do this automatically
755            // and only if a specific value was explicitly set on the listener.
756
757            let receive_buffer_size = self.0.receive_buffer_size.load(Relaxed);
758            if receive_buffer_size > 0 {
759                // Ignore potential error.
760                _ = sockopt::set_socket_recv_buffer_size(&stream, receive_buffer_size);
761            }
762
763            let send_buffer_size = self.0.send_buffer_size.load(Relaxed);
764            if send_buffer_size > 0 {
765                // Ignore potential error.
766                _ = sockopt::set_socket_send_buffer_size(&stream, send_buffer_size);
767            }
768
769            // For some reason, IP_TTL is inherited, but IPV6_UNICAST_HOPS isn't.
770            if family == SocketAddressFamily::Ipv6 {
771                let hop_limit = self.0.hop_limit.load(Relaxed);
772                if hop_limit > 0 {
773                    // Ignore potential error.
774                    _ = sockopt::set_ipv6_unicast_hops(&stream, Some(hop_limit));
775                }
776            }
777
778            let keep_alive_idle_time = self.0.keep_alive_idle_time.load(Relaxed);
779            if keep_alive_idle_time > 0 {
780                // Ignore potential error.
781                _ = sockopt::set_tcp_keepidle(&stream, Duration::from_nanos(keep_alive_idle_time));
782            }
783        }
784    }
785}
786
787fn socket(family: SocketAddressFamily) -> std::io::Result<tokio::net::TcpSocket> {
788    match family {
789        SocketAddressFamily::Ipv4 => tokio::net::TcpSocket::new_v4(),
790        SocketAddressFamily::Ipv6 => {
791            let socket = tokio::net::TcpSocket::new_v6()?;
792
793            // From the WASI spec:
794            // > On IPv6 sockets, IPV6_V6ONLY is enabled by default and can't
795            // > be configured otherwise.
796            sockopt::set_ipv6_v6only(&socket, true)?;
797            Ok(socket)
798        }
799    }
800}
801
802fn bind(socket: &tokio::net::TcpSocket, local_address: SocketAddr) -> Result<(), ErrorCode> {
803    // From the WASI spec:
804    // > The bind operation shouldn't be affected by the TIME_WAIT state of a
805    // > recently closed socket on the same local address. In practice this
806    // > means that the SO_REUSEADDR socket option should be set implicitly on
807    // > all platforms, except on Windows where this is the default behavior
808    // > and SO_REUSEADDR performs something different.
809    #[cfg(not(windows))]
810    {
811        _ = sockopt::set_socket_reuseaddr(&socket, true);
812    }
813
814    // Perform the OS bind call.
815    socket
816        .bind(local_address)
817        .map_err(|err| match Errno::from_io_error(&err) {
818            // From https://pubs.opengroup.org/onlinepubs/9699919799/functions/bind.html:
819            // > [EAFNOSUPPORT] The specified address is not a valid address for the address family of the specified socket
820            //
821            // The most common reasons for this error should have already
822            // been handled by our own validation.. This error mapping is here
823            // just in case there is an edge case we didn't catch.
824            Some(Errno::AFNOSUPPORT) => ErrorCode::InvalidArgument,
825            // See: https://learn.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-bind#:~:text=WSAENOBUFS
826            // Windows returns WSAENOBUFS when the ephemeral ports have been exhausted.
827            #[cfg(windows)]
828            Some(Errno::NOBUFS) => ErrorCode::AddressInUse,
829            _ => err.into(),
830        })
831}
832
833async fn accept(
834    listener: &tokio::net::TcpListener,
835) -> std::io::Result<(tokio::net::TcpStream, SocketAddr)> {
836    listener
837        .accept()
838        .await
839        .map_err(|err| match Errno::from_io_error(&err) {
840            // From: https://learn.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-accept#:~:text=WSAEINPROGRESS
841            // > WSAEINPROGRESS: A blocking Windows Sockets 1.1 call is in progress,
842            // > or the service provider is still processing a callback function.
843            //
844            // wasi-sockets doesn't have an equivalent to the EINPROGRESS error,
845            // because in POSIX this error is only returned by a non-blocking
846            // `connect` and wasi-sockets has a different solution for that.
847            #[cfg(windows)]
848            Some(Errno::INPROGRESS) => Errno::INTR.into(),
849
850            // Normalize Linux' non-standard behavior.
851            //
852            // From https://man7.org/linux/man-pages/man2/accept.2.html:
853            // > Linux accept() passes already-pending network errors on the
854            // > new socket as an error code from accept(). This behavior
855            // > differs from other BSD socket implementations. (...)
856            #[cfg(target_os = "linux")]
857            Some(
858                Errno::CONNRESET
859                | Errno::NETRESET
860                | Errno::HOSTUNREACH
861                | Errno::HOSTDOWN
862                | Errno::NETDOWN
863                | Errno::NETUNREACH
864                | Errno::PROTO
865                | Errno::NOPROTOOPT
866                | Errno::NONET
867                | Errno::OPNOTSUPP,
868            ) => Errno::CONNABORTED.into(),
869
870            _ => err,
871        })
872}
873
874fn reset(socket: tokio::net::TcpStream) {
875    _ = socket.set_zero_linger();
876    drop(socket);
877}
878
879fn clamp_keep_alive_time(value: u64) -> u64 {
880    // Ensure that the value passed to the actual syscall never gets rounded down to 0.
881    const MIN: u64 = 1 * NANOS_PER_SEC;
882
883    // Cap it at Linux' maximum, which appears to have the lowest limit across our supported platforms.
884    const MAX: u64 = (i16::MAX as u64) * NANOS_PER_SEC;
885
886    value.clamp(MIN, MAX)
887}
888
889fn clamp_keep_alive_count(value: u32) -> u32 {
890    const MIN_CNT: u32 = 1;
891    // Cap it at Linux' maximum, which appears to have the lowest limit across our supported platforms.
892    const MAX_CNT: u32 = i8::MAX as u32;
893
894    value.clamp(MIN_CNT, MAX_CNT)
895}