wasmtime_wasi/
udp.rs

1use crate::host::network::util;
2use crate::runtime::with_ambient_tokio_runtime;
3use async_trait::async_trait;
4use cap_net_ext::{AddressFamily, Blocking};
5use io_lifetimes::raw::{FromRawSocketlike, IntoRawSocketlike};
6use std::io;
7use std::net::SocketAddr;
8use std::sync::Arc;
9use wasmtime_wasi_io::poll::Pollable;
10
11use super::network::{SocketAddrCheck, SocketAddressFamily};
12
13/// The state of a UDP socket.
14///
15/// This represents the various states a socket can be in during the
16/// activities of binding, and connecting.
17pub(crate) enum UdpState {
18    /// The initial state for a newly-created socket.
19    Default,
20
21    /// Binding started via `start_bind`.
22    BindStarted,
23
24    /// Binding finished via `finish_bind`. The socket has an address but
25    /// is not yet listening for connections.
26    Bound,
27
28    /// The socket is "connected" to a peer address.
29    Connected,
30}
31
32/// A host UDP socket, plus associated bookkeeping.
33///
34/// The inner state is wrapped in an Arc because the same underlying socket is
35/// used for implementing the stream types.
36pub struct UdpSocket {
37    /// The part of a `UdpSocket` which is reference-counted so that we
38    /// can pass it to async tasks.
39    pub(crate) inner: Arc<tokio::net::UdpSocket>,
40
41    /// The current state in the bind/connect progression.
42    pub(crate) udp_state: UdpState,
43
44    /// Socket address family.
45    pub(crate) family: SocketAddressFamily,
46
47    /// The check of allowed addresses
48    pub(crate) socket_addr_check: Option<SocketAddrCheck>,
49}
50
51#[async_trait]
52impl Pollable for UdpSocket {
53    async fn ready(&mut self) {
54        // None of the socket-level operations block natively
55    }
56}
57
58impl UdpSocket {
59    /// Create a new socket in the given family.
60    pub fn new(family: AddressFamily) -> io::Result<Self> {
61        // Create a new host socket and set it to non-blocking, which is needed
62        // by our async implementation.
63        let fd = util::udp_socket(family, Blocking::No)?;
64
65        let socket_address_family = match family {
66            AddressFamily::Ipv4 => SocketAddressFamily::Ipv4,
67            AddressFamily::Ipv6 => {
68                rustix::net::sockopt::set_ipv6_v6only(&fd, true)?;
69                SocketAddressFamily::Ipv6
70            }
71        };
72
73        let socket = Self::setup_tokio_udp_socket(fd)?;
74
75        Ok(UdpSocket {
76            inner: Arc::new(socket),
77            udp_state: UdpState::Default,
78            family: socket_address_family,
79            socket_addr_check: None,
80        })
81    }
82
83    fn setup_tokio_udp_socket(fd: rustix::fd::OwnedFd) -> io::Result<tokio::net::UdpSocket> {
84        let std_socket =
85            unsafe { std::net::UdpSocket::from_raw_socketlike(fd.into_raw_socketlike()) };
86        with_ambient_tokio_runtime(|| tokio::net::UdpSocket::try_from(std_socket))
87    }
88
89    pub fn udp_socket(&self) -> &tokio::net::UdpSocket {
90        &self.inner
91    }
92}
93
94pub struct IncomingDatagramStream {
95    pub(crate) inner: Arc<tokio::net::UdpSocket>,
96
97    /// If this has a value, the stream is "connected".
98    pub(crate) remote_address: Option<SocketAddr>,
99}
100
101pub struct OutgoingDatagramStream {
102    pub(crate) inner: Arc<tokio::net::UdpSocket>,
103
104    /// If this has a value, the stream is "connected".
105    pub(crate) remote_address: Option<SocketAddr>,
106
107    /// Socket address family.
108    pub(crate) family: SocketAddressFamily,
109
110    pub(crate) send_state: SendState,
111
112    /// The check of allowed addresses
113    pub(crate) socket_addr_check: Option<SocketAddrCheck>,
114}
115
116pub(crate) enum SendState {
117    /// Waiting for the API consumer to call `check-send`.
118    Idle,
119
120    /// Ready to send up to x datagrams.
121    Permitted(usize),
122
123    /// Waiting for the OS.
124    Waiting,
125}