wasmtime_wasi/p2/
udp.rs

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