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 or receiving.
131            if self
132                .permissions
133                .check(addr, SocketAddrUse::UdpSend)
134                .await
135                .is_err()
136            {
137                self.permissions
138                    .check(addr, SocketAddrUse::UdpReceive)
139                    .await?;
140            }
141        }
142
143        let result = connect(&self.socket, addr);
144        self.update_remote_address();
145        result.map_err(|e| e.into())
146    }
147
148    pub(crate) fn disconnect(&mut self) -> Result<(), ErrorCode> {
149        if !self.is_connected() {
150            return Err(ErrorCode::InvalidState);
151        }
152
153        // On Linux, disconnecting a UDP socket relinquishes its local port
154        // assignment in some cases. If the socket was bound to the wildcard
155        // address, its local address will then read `0.0.0.0:0` or `[::]:0`
156        // which is indistinguishable from an unbound socket. To ensure
157        // `is_bound()` will continue to return `true` after the disconnect, we
158        // manually settle the `is_bound` state here:
159        self.is_bound = true;
160
161        let result = disconnect(&self.socket);
162        self.update_remote_address();
163        result.map_err(|e| e.into())
164    }
165
166    /// Update our internal bookkeeping based on the actual state of the socket.
167    /// This should be called after any operation that may change the remote
168    /// address.
169    fn update_remote_address(&mut self) {
170        self.remote_addr = if let Ok(addr) = self.socket.peer_addr()
171            && addr != unspecified_addr(self.family)
172        {
173            Some(addr)
174        } else {
175            None
176        }
177    }
178
179    pub(crate) fn send(
180        &mut self,
181        data: Vec<u8>,
182        addr: Option<SocketAddr>,
183    ) -> impl Future<Output = Result<(), ErrorCode>> + Send + use<> {
184        let family = self.family;
185        let socket = self.socket.clone();
186        let permissions = self.permissions.clone();
187        let connected_addr = self.remote_address().ok();
188        let is_bound = self.is_bound();
189
190        async move {
191            if data.len() > MAX_DATAGRAM_SIZE {
192                return Err(ErrorCode::DatagramTooLarge);
193            }
194
195            let effective_addr = if let Some(addr) = addr {
196                if !is_valid_remote_address(addr) || !is_valid_address_family(addr.ip(), family) {
197                    return Err(ErrorCode::InvalidArgument);
198                }
199
200                // If the socket is connected, the provided address must match the
201                // connected address.
202                if connected_addr.is_some() && connected_addr != Some(addr) {
203                    return Err(ErrorCode::InvalidArgument);
204                }
205
206                addr
207            } else if let Some(connected_addr) = connected_addr {
208                connected_addr
209            } else {
210                return Err(ErrorCode::InvalidArgument);
211            };
212
213            // Perform all permission checks before doing any syscalls.
214            {
215                if !is_bound {
216                    // If not explicitly bound, the OS will implicitly bind the
217                    // socket to an ephemeral port when sending.
218                    let implicit_bind_addr = unspecified_addr(family);
219                    permissions
220                        .check(implicit_bind_addr, SocketAddrUse::UdpBind)
221                        .await?;
222                }
223
224                permissions
225                    .check(effective_addr, SocketAddrUse::UdpSend)
226                    .await?;
227            }
228
229            if connected_addr == Some(effective_addr) {
230                socket.send(&data).await?;
231            } else {
232                socket.send_to(&data, effective_addr).await?;
233            }
234
235            Ok(())
236        }
237    }
238
239    pub(crate) fn recv(
240        &mut self,
241    ) -> impl Future<Output = Result<(Vec<u8>, SocketAddr), ErrorCode>> + Send + use<> {
242        let socket = self.socket.clone();
243        let permissions = self.permissions.clone();
244        let is_bound = self.is_bound();
245
246        async move {
247            if !is_bound {
248                return Err(ErrorCode::InvalidState);
249            }
250
251            loop {
252                let mut data = vec![0; MAX_DATAGRAM_SIZE];
253                let (len, addr) = socket.recv_from(&mut data).await?;
254                data.truncate(len);
255
256                match permissions.check(addr, SocketAddrUse::UdpReceive).await {
257                    Ok(()) => return Ok((data, addr)),
258                    Err(_) => {
259                        // Not allowed. Drop the packet and poll again.
260                        continue;
261                    }
262                }
263            }
264        }
265    }
266
267    pub(crate) fn address_family(&self) -> SocketAddressFamily {
268        self.family
269    }
270
271    pub(crate) fn unicast_hop_limit(&self) -> Result<u8, ErrorCode> {
272        let n = get_unicast_hop_limit(&self.socket, self.family)?;
273        Ok(n)
274    }
275
276    pub(crate) fn set_unicast_hop_limit(&self, value: u8) -> Result<(), ErrorCode> {
277        set_unicast_hop_limit(&self.socket, self.family, value)?;
278        Ok(())
279    }
280
281    pub(crate) fn receive_buffer_size(&self) -> Result<u64, ErrorCode> {
282        let n = get_receive_buffer_size(&self.socket)?;
283        Ok(n)
284    }
285
286    pub(crate) fn set_receive_buffer_size(&self, value: u64) -> Result<(), ErrorCode> {
287        set_receive_buffer_size(&self.socket, value)?;
288        Ok(())
289    }
290
291    pub(crate) fn send_buffer_size(&self) -> Result<u64, ErrorCode> {
292        let n = get_send_buffer_size(&self.socket)?;
293        Ok(n)
294    }
295
296    pub(crate) fn set_send_buffer_size(&self, value: u64) -> Result<(), ErrorCode> {
297        set_send_buffer_size(&self.socket, value)?;
298        Ok(())
299    }
300}
301
302/// Creates a non-blocking/cloexec UDP socket.
303fn socket(family: SocketAddressFamily) -> std::io::Result<tokio::net::UdpSocket> {
304    // Let the standard library be responsible for handling `WSAStartup`.
305    #[cfg(windows)]
306    static INIT: std::sync::Once = std::sync::Once::new();
307    #[cfg(windows)]
308    INIT.call_once(|| {
309        let _ = std::net::TcpStream::connect(std::net::SocketAddrV4::new(
310            std::net::Ipv4Addr::UNSPECIFIED,
311            0,
312        ));
313    });
314
315    #[cfg(not(any(windows, target_vendor = "apple")))]
316    let flags = rustix::net::SocketFlags::CLOEXEC | rustix::net::SocketFlags::NONBLOCK;
317    #[cfg(any(windows, target_vendor = "apple"))]
318    let flags = rustix::net::SocketFlags::empty();
319
320    let socket = rustix::net::socket_with(
321        match family {
322            SocketAddressFamily::Ipv4 => rustix::net::AddressFamily::INET,
323            SocketAddressFamily::Ipv6 => rustix::net::AddressFamily::INET6,
324        },
325        rustix::net::SocketType::DGRAM,
326        flags,
327        None,
328    )?;
329    #[cfg(target_vendor = "apple")]
330    rustix::io::ioctl_fioclex(&socket)?;
331    #[cfg(any(windows, target_vendor = "apple"))]
332    rustix::io::ioctl_fionbio(&socket, true)?;
333
334    // From the WASI spec:
335    // > On IPv6 sockets, IPV6_V6ONLY is enabled by default and can't
336    // > be configured otherwise.
337    if family == SocketAddressFamily::Ipv6 {
338        rustix::net::sockopt::set_ipv6_v6only(&socket, true)?;
339    }
340
341    Ok(tokio::net::UdpSocket::try_from(std::net::UdpSocket::from(
342        socket,
343    ))?)
344}
345
346fn bind(sockfd: impl AsFd, addr: SocketAddr) -> Result<(), Errno> {
347    rustix::net::bind(sockfd, &addr).map_err(|err| match err {
348        // See: https://learn.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-bind#:~:text=WSAENOBUFS
349        // Windows returns WSAENOBUFS when the ephemeral ports have been exhausted.
350        #[cfg(windows)]
351        Errno::NOBUFS => Errno::ADDRINUSE,
352        // From https://pubs.opengroup.org/onlinepubs/9699919799/functions/bind.html:
353        // > [EAFNOSUPPORT] The specified address is not a valid address for the address family of the specified socket
354        //
355        // The most common reasons for this error should have already
356        // been handled by our own validation. This error mapping is here just
357        // in case there is an edge case we didn't catch.
358        Errno::AFNOSUPPORT => Errno::INVAL,
359        _ => err,
360    })
361}
362
363fn connect(sockfd: impl AsFd, addr: SocketAddr) -> Result<(), Errno> {
364    match rustix::net::connect(sockfd.as_fd(), &addr) {
365        // When connecting a UDP socket, the OS looks up the best route to the
366        // remote address and selects an appropriate outgoing interface.
367        // If the new destination routes through an interface different than the
368        // previously selected interface, most operating systems will
369        // automatically update the socket's local address to match that route.
370        //
371        // Linux however doesn't do that automatically and we manually
372        // dissolve the existing association and then connect again to the
373        // new destination.
374        #[cfg(target_os = "linux")]
375        Err(Errno::INVAL) => {
376            _ = disconnect(sockfd.as_fd());
377            return rustix::net::connect(sockfd.as_fd(), &addr);
378        }
379        // The most common reason for AFNOSUPPORT is an invalid address
380        // family. This should have already been handled by our own
381        // validation. This error mapping is here just in case there is an
382        // edge case we didn't catch.
383        Err(Errno::AFNOSUPPORT) => Err(Errno::INVAL),
384        // EINPROGRESS should only returned by non-blocking TCP sockets,
385        // not UDP sockets.
386        Err(Errno::INPROGRESS) => {
387            debug!("UDP connect returned EINPROGRESS, which should never happen");
388            Ok(())
389        }
390        r => r,
391    }
392}
393
394fn disconnect(sockfd: impl AsFd) -> Result<(), Errno> {
395    match rustix::net::connect_unspec(sockfd) {
396        // BSD platforms return an error even if the UDP socket was disconnected successfully.
397        //
398        // MacOS was kind enough to document this: https://developer.apple.com/library/archive/documentation/System/Conceptual/ManPages_iPhoneOS/man2/connect.2.html
399        // > Datagram sockets may dissolve the association by connecting to an
400        // > invalid address, such as a null address or an address with the address
401        // > family set to AF_UNSPEC (the error EAFNOSUPPORT will be harmlessly
402        // > returned).
403        //
404        // ... except that this appears to be incomplete, because experiments
405        // have shown that MacOS actually returns EINVAL, depending on the
406        // address family of the socket.
407        #[cfg(target_os = "macos")]
408        Err(Errno::INVAL | Errno::AFNOSUPPORT) => Ok(()),
409        r => r,
410    }
411}
wasmtime_wasi/sockets/udp.rs

wasmtime_wasi/sockets/
udp.rs
