wasmtime_wasi/sockets/
udp.rs

1use crate::runtime::with_ambient_tokio_runtime;
2use crate::sockets::{
3    ErrorCode, 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, set_receive_buffer_size, set_send_buffer_size, set_unicast_hop_limit,
6    unspecified_addr,
7};
8use rustix::fd::AsFd;
9use rustix::io::Errno;
10use std::net::SocketAddr;
11use std::sync::Arc;
12use tracing::debug;
13
14/// Theoretical maximum byte size of a UDP datagram, the real limit is lower,
15/// but we do not account for e.g. the transport layer here for simplicity.
16/// In practice, datagrams are typically less than 1500 bytes.
17pub(crate) const MAX_DATAGRAM_SIZE: usize = u16::MAX as usize;
18
19/// A host UDP socket, plus associated bookkeeping.
20///
21/// The inner state is wrapped in an Arc because the same underlying socket is
22/// used for implementing the stream types.
23pub struct UdpSocket {
24    socket: Arc<tokio::net::UdpSocket>,
25    family: SocketAddressFamily,
26
27    /// The checks to perform before doing any noteworthy syscall.
28    permissions: SocketAddrCheck,
29
30    /// Cached value of whether the socket is bound. This is cached to avoid
31    /// redundant syscalls in every `send` & `receive`.
32    is_bound: bool,
33
34    /// Cached value of the remote address. This is cached to avoid redundant
35    /// syscalls in every `send`.
36    remote_addr: Option<SocketAddr>,
37}
38
39impl UdpSocket {
40    /// Create a new socket in the given family.
41    pub(crate) async fn new(
42        cx: &WasiSocketsCtx,
43        family: SocketAddressFamily,
44    ) -> Result<Self, ErrorCode> {
45        cx.allowed_network_uses.check_allowed_udp()?;
46
47        let socket = with_ambient_tokio_runtime(|| socket(family))?;
48
49        // Native UDP sockets are immediately writable after creation and
50        // existing guest code out in the wild depends on that. However, due to
51        // the way Tokio is structured internally, a newly created Tokio socket
52        // starts out as "not writable" and a background tokio thread updates
53        // this state asynchronously soon after. Some more info in this thread:
54        // https://github.com/bytecodealliance/wasmtime/issues/12612#issuecomment-3923714174
55        //
56        // To prevent exposing guests to this race condition, we wait for Tokio
57        // to finish its internal setup:
58        socket.writable().await?;
59
60        Ok(Self {
61            socket: Arc::new(socket),
62            is_bound: false,
63            remote_addr: None,
64            permissions: cx.socket_addr_check.clone(),
65            family,
66        })
67    }
68
69    pub(crate) fn is_bound(&mut self) -> bool {
70        // Once bound, a UDP socket can never become unbound again. So we can
71        // skip all work after a previous call has already determined the
72        // socket to be bound.
73        if !self.is_bound {
74            self.is_bound = self
75                .socket
76                .local_addr()
77                .is_ok_and(|addr| addr != unspecified_addr(self.family));
78        }
79        self.is_bound
80    }
81
82    pub(crate) fn local_address(&mut self) -> Result<SocketAddr, ErrorCode> {
83        if !self.is_bound() {
84            return Err(ErrorCode::InvalidState);
85        }
86        self.socket.local_addr().map_err(|e| e.into())
87    }
88
89    pub(crate) fn remote_address(&mut self) -> Result<SocketAddr, ErrorCode> {
90        self.remote_addr.ok_or(ErrorCode::InvalidState)
91    }
92
93    pub(crate) fn is_connected(&mut self) -> bool {
94        self.remote_addr.is_some()
95    }
96
97    pub(crate) async fn bind(&mut self, addr: SocketAddr) -> Result<(), ErrorCode> {
98        if self.is_bound() {
99            return Err(ErrorCode::InvalidState);
100        }
101        if !is_valid_address_family(addr.ip(), self.family) {
102            return Err(ErrorCode::InvalidArgument);
103        }
104
105        self.permissions.check(addr, SocketAddrUse::UdpBind).await?;
106
107        bind(&self.socket, addr)?;
108        Ok(())
109    }
110
111    pub(crate) async fn connect(&mut self, addr: SocketAddr) -> Result<(), ErrorCode> {
112        if !is_valid_address_family(addr.ip(), self.family) || !is_valid_remote_address(addr) {
113            return Err(ErrorCode::InvalidArgument);
114        }
115
116        // Perform all permission checks before doing any syscalls.
117        {
118            if !self.is_bound() {
119                // If not explicitly bound, the OS will implicitly bind the
120                // socket to an ephemeral port when connecting.
121                let implicit_bind_addr = unspecified_addr(self.family);
122                self.permissions
123                    .check(implicit_bind_addr, SocketAddrUse::UdpBind)
124                    .await?;
125            }
126
127            // On UDP sockets, "connecting" is just a local operation that sets the
128            // default remote address for future sends and receives. It does not
129            // actually do any I/O on its own. We'll allow the `connect` call
130            // if the address is permitted for sending and receiving.
131            self.permissions.check(addr, SocketAddrUse::UdpSend).await?;
132            self.permissions
133                .check(addr, SocketAddrUse::UdpReceive)
134                .await?;
135        }
136
137        let result = connect(&self.socket, addr);
138        self.update_remote_address();
139        result.map_err(|e| e.into())
140    }
141
142    pub(crate) fn disconnect(&mut self) -> Result<(), ErrorCode> {
143        if !self.is_connected() {
144            return Err(ErrorCode::InvalidState);
145        }
146
147        // On Linux, disconnecting a UDP socket relinquishes its local port
148        // assignment in some cases. If the socket was bound to the wildcard
149        // address, its local address will then read `0.0.0.0:0` or `[::]:0`
150        // which is indistinguishable from an unbound socket. To ensure
151        // `is_bound()` will continue to return `true` after the disconnect, we
152        // manually settle the `is_bound` state here:
153        self.is_bound = true;
154
155        let result = disconnect(&self.socket);
156        self.update_remote_address();
157        result.map_err(|e| e.into())
158    }
159
160    /// Update our internal bookkeeping based on the actual state of the socket.
161    /// This should be called after any operation that may change the remote
162    /// address.
163    fn update_remote_address(&mut self) {
164        self.remote_addr = if let Ok(addr) = self.socket.peer_addr()
165            && addr != unspecified_addr(self.family)
166        {
167            Some(addr)
168        } else {
169            None
170        }
171    }
172
173    pub(crate) fn send(
174        &mut self,
175        data: Vec<u8>,
176        addr: Option<SocketAddr>,
177    ) -> impl Future<Output = Result<(), ErrorCode>> + Send + use<> {
178        let family = self.family;
179        let socket = self.socket.clone();
180        let permissions = self.permissions.clone();
181        let connected_addr = self.remote_address().ok();
182        let is_bound = self.is_bound();
183
184        async move {
185            if data.len() > MAX_DATAGRAM_SIZE {
186                return Err(ErrorCode::DatagramTooLarge);
187            }
188
189            let effective_addr = if let Some(addr) = addr {
190                if !is_valid_remote_address(addr) || !is_valid_address_family(addr.ip(), family) {
191                    return Err(ErrorCode::InvalidArgument);
192                }
193
194                // If the socket is connected, the provided address must match the
195                // connected address.
196                if connected_addr.is_some() && connected_addr != Some(addr) {
197                    return Err(ErrorCode::InvalidArgument);
198                }
199
200                addr
201            } else if let Some(connected_addr) = connected_addr {
202                connected_addr
203            } else {
204                return Err(ErrorCode::InvalidArgument);
205            };
206
207            // Perform all permission checks before doing any syscalls.
208            {
209                if !is_bound {
210                    // If not explicitly bound, the OS will implicitly bind the
211                    // socket to an ephemeral port when sending.
212                    let implicit_bind_addr = unspecified_addr(family);
213                    permissions
214                        .check(implicit_bind_addr, SocketAddrUse::UdpBind)
215                        .await?;
216                }
217
218                permissions
219                    .check(effective_addr, SocketAddrUse::UdpSend)
220                    .await?;
221            }
222
223            if connected_addr == Some(effective_addr) {
224                socket.send(&data).await?;
225            } else {
226                socket.send_to(&data, effective_addr).await?;
227            }
228
229            Ok(())
230        }
231    }
232
233    pub(crate) fn recv(
234        &mut self,
235    ) -> impl Future<Output = Result<(Vec<u8>, SocketAddr), ErrorCode>> + Send + use<> {
236        let socket = self.socket.clone();
237        let permissions = self.permissions.clone();
238        let is_bound = self.is_bound();
239
240        async move {
241            if !is_bound {
242                return Err(ErrorCode::InvalidState);
243            }
244
245            loop {
246                let mut data = vec![0; MAX_DATAGRAM_SIZE];
247                let (len, addr) = socket.recv_from(&mut data).await?;
248                data.truncate(len);
249
250                match permissions.check(addr, SocketAddrUse::UdpReceive).await {
251                    Ok(()) => return Ok((data, addr)),
252                    Err(_) => {
253                        // Not allowed. Drop the packet and poll again.
254                        continue;
255                    }
256                }
257            }
258        }
259    }
260
261    pub(crate) fn address_family(&self) -> SocketAddressFamily {
262        self.family
263    }
264
265    pub(crate) fn unicast_hop_limit(&self) -> Result<u8, ErrorCode> {
266        let n = get_unicast_hop_limit(&self.socket, self.family)?;
267        Ok(n)
268    }
269
270    pub(crate) fn set_unicast_hop_limit(&self, value: u8) -> Result<(), ErrorCode> {
271        set_unicast_hop_limit(&self.socket, self.family, value)?;
272        Ok(())
273    }
274
275    pub(crate) fn receive_buffer_size(&self) -> Result<u64, ErrorCode> {
276        let n = get_receive_buffer_size(&self.socket)?;
277        Ok(n)
278    }
279
280    pub(crate) fn set_receive_buffer_size(&self, value: u64) -> Result<(), ErrorCode> {
281        set_receive_buffer_size(&self.socket, value)?;
282        Ok(())
283    }
284
285    pub(crate) fn send_buffer_size(&self) -> Result<u64, ErrorCode> {
286        let n = get_send_buffer_size(&self.socket)?;
287        Ok(n)
288    }
289
290    pub(crate) fn set_send_buffer_size(&self, value: u64) -> Result<(), ErrorCode> {
291        set_send_buffer_size(&self.socket, value)?;
292        Ok(())
293    }
294}
295
296/// Creates a non-blocking/cloexec UDP socket.
297fn socket(family: SocketAddressFamily) -> std::io::Result<tokio::net::UdpSocket> {
298    // Let the standard library be responsible for handling `WSAStartup`.
299    #[cfg(windows)]
300    static INIT: std::sync::Once = std::sync::Once::new();
301    #[cfg(windows)]
302    INIT.call_once(|| {
303        let _ = std::net::TcpStream::connect(std::net::SocketAddrV4::new(
304            std::net::Ipv4Addr::UNSPECIFIED,
305            0,
306        ));
307    });
308
309    #[cfg(not(any(windows, target_vendor = "apple")))]
310    let flags = rustix::net::SocketFlags::CLOEXEC | rustix::net::SocketFlags::NONBLOCK;
311    #[cfg(any(windows, target_vendor = "apple"))]
312    let flags = rustix::net::SocketFlags::empty();
313
314    let socket = rustix::net::socket_with(
315        match family {
316            SocketAddressFamily::Ipv4 => rustix::net::AddressFamily::INET,
317            SocketAddressFamily::Ipv6 => rustix::net::AddressFamily::INET6,
318        },
319        rustix::net::SocketType::DGRAM,
320        flags,
321        None,
322    )?;
323    #[cfg(target_vendor = "apple")]
324    rustix::io::ioctl_fioclex(&socket)?;
325    #[cfg(any(windows, target_vendor = "apple"))]
326    rustix::io::ioctl_fionbio(&socket, true)?;
327
328    // From the WASI spec:
329    // > On IPv6 sockets, IPV6_V6ONLY is enabled by default and can't
330    // > be configured otherwise.
331    if family == SocketAddressFamily::Ipv6 {
332        rustix::net::sockopt::set_ipv6_v6only(&socket, true)?;
333    }
334
335    Ok(tokio::net::UdpSocket::try_from(std::net::UdpSocket::from(
336        socket,
337    ))?)
338}
339
340fn bind(sockfd: impl AsFd, addr: SocketAddr) -> Result<(), Errno> {
341    rustix::net::bind(sockfd, &addr).map_err(|err| match err {
342        // See: https://learn.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-bind#:~:text=WSAENOBUFS
343        // Windows returns WSAENOBUFS when the ephemeral ports have been exhausted.
344        #[cfg(windows)]
345        Errno::NOBUFS => Errno::ADDRINUSE,
346        // From https://pubs.opengroup.org/onlinepubs/9699919799/functions/bind.html:
347        // > [EAFNOSUPPORT] The specified address is not a valid address for the address family of the specified socket
348        //
349        // The most common reasons for this error should have already
350        // been handled by our own validation. This error mapping is here just
351        // in case there is an edge case we didn't catch.
352        Errno::AFNOSUPPORT => Errno::INVAL,
353        _ => err,
354    })
355}
356
357fn connect(sockfd: impl AsFd, addr: SocketAddr) -> Result<(), Errno> {
358    match rustix::net::connect(sockfd.as_fd(), &addr) {
359        // When connecting a UDP socket, the OS looks up the best route to the
360        // remote address and selects an appropriate outgoing interface.
361        // If the new destination routes through an interface different than the
362        // previously selected interface, most operating systems will
363        // automatically update the socket's local address to match that route.
364        //
365        // Linux however doesn't do that automatically and we manually
366        // dissolve the existing association and then connect again to the
367        // new destination.
368        #[cfg(target_os = "linux")]
369        Err(Errno::INVAL) => {
370            _ = disconnect(sockfd.as_fd());
371            return rustix::net::connect(sockfd.as_fd(), &addr);
372        }
373        // The most common reason for AFNOSUPPORT is an invalid address
374        // family. This should have already been handled by our own
375        // validation. This error mapping is here just in case there is an
376        // edge case we didn't catch.
377        Err(Errno::AFNOSUPPORT) => Err(Errno::INVAL),
378        // EINPROGRESS should only returned by non-blocking TCP sockets,
379        // not UDP sockets.
380        Err(Errno::INPROGRESS) => {
381            debug!("UDP connect returned EINPROGRESS, which should never happen");
382            Ok(())
383        }
384        r => r,
385    }
386}
387
388fn disconnect(sockfd: impl AsFd) -> Result<(), Errno> {
389    match rustix::net::connect_unspec(sockfd) {
390        // BSD platforms return an error even if the UDP socket was disconnected successfully.
391        //
392        // MacOS was kind enough to document this: https://developer.apple.com/library/archive/documentation/System/Conceptual/ManPages_iPhoneOS/man2/connect.2.html
393        // > Datagram sockets may dissolve the association by connecting to an
394        // > invalid address, such as a null address or an address with the address
395        // > family set to AF_UNSPEC (the error EAFNOSUPPORT will be harmlessly
396        // > returned).
397        //
398        // ... except that this appears to be incomplete, because experiments
399        // have shown that MacOS actually returns EINVAL, depending on the
400        // address family of the socket.
401        #[cfg(target_os = "macos")]
402        Err(Errno::INVAL | Errno::AFNOSUPPORT) => Ok(()),
403        r => r,
404    }
405}
wasmtime_wasi/sockets/udp.rs

wasmtime_wasi/sockets/
udp.rs
