wasmtime_wasi/p2/
udp.rs

1use futures::TryFutureExt;
2
3use crate::{
4    p2::bindings::sockets::network::ErrorCode,
5    runtime::poll_now,
6    sockets::{MaybeSpawned, UdpSocket as P3Socket},
7};
8use std::{
9    net::SocketAddr,
10    sync::{Arc, Mutex},
11};
12
13/// A UDP socket + associated p2 bookkeeping.
14pub struct UdpSocket {
15    pub(crate) inner: Arc<Mutex<P3Socket>>,
16    pub(crate) in_progress_operation: Option<AsyncOperation>,
17}
18impl UdpSocket {
19    pub(crate) fn new(inner: P3Socket) -> Self {
20        Self {
21            inner: Arc::new(Mutex::new(inner)),
22            in_progress_operation: None,
23        }
24    }
25    pub(crate) fn get_mut(&mut self) -> Option<&mut P3Socket> {
26        Arc::get_mut(&mut self.inner)?.get_mut().ok()
27    }
28    pub(crate) fn lock(&self) -> std::sync::MutexGuard<'_, P3Socket> {
29        self.inner.lock().expect("other thread panicked")
30    }
31}
32
33#[derive(Clone, Copy, Debug, PartialEq, Eq)]
34pub(crate) enum AsyncOperation {
35    Bind,
36}
37
38pub struct IncomingDatagramStream {
39    pub(crate) inner: Arc<Mutex<P3Socket>>,
40    pub(crate) connected_addr: Option<SocketAddr>,
41    pub(crate) current_recv: Option<MaybeSpawned<Result<(Vec<u8>, SocketAddr), ErrorCode>>>,
42}
43impl IncomingDatagramStream {
44    pub(crate) fn new(inner: Arc<Mutex<P3Socket>>) -> Self {
45        let connected_addr = inner.lock().unwrap().remote_address().ok();
46        Self {
47            inner,
48            connected_addr,
49            current_recv: None,
50        }
51    }
52    pub(crate) fn poll_recv_ready(
53        &mut self,
54        cx: &mut std::task::Context<'_>,
55    ) -> std::task::Poll<()> {
56        if self.current_recv.is_none() {
57            let connected_addr = self.connected_addr;
58            let inner = self.inner.clone();
59            let recv = MaybeSpawned::poll_or_spawn(async move {
60                loop {
61                    let fut = inner.lock().unwrap().recv();
62                    let (data, addr) = fut.await?;
63
64                    // Only process the packet if it matches the expected remote
65                    // address (if connected). Under normal circumstances, the
66                    // OS should already do this filtering for us. However,
67                    // nothing in POSIX guarantees this behavior, especially
68                    // after (re)connecting a socket with already queued packets
69                    // from a different peer. Case in point: on Linux the
70                    // filtering happens when the packet is received from the
71                    // network, *not* when the packet is delivered to the
72                    // application as part of `recvfrom`.
73                    if let Some(connected_addr) = connected_addr
74                        && connected_addr != addr
75                    {
76                        continue;
77                    }
78
79                    return Ok((data, addr));
80                }
81            });
82            self.current_recv = Some(recv);
83        }
84
85        self.current_recv
86            .as_mut()
87            .unwrap()
88            .poll_ready(cx)
89            .map(|_| ())
90    }
91
92    pub(crate) fn try_recv(&mut self) -> Result<(Vec<u8>, SocketAddr), ErrorCode> {
93        if poll_now(|cx| self.poll_recv_ready(cx)).is_none() {
94            return Err(ErrorCode::WouldBlock);
95        }
96
97        self.current_recv.take().unwrap().unwrap_ready()
98    }
99
100    pub(crate) async fn finish(mut self) {
101        let Some(MaybeSpawned::Pending(recv)) = self.current_recv.take() else {
102            return;
103        };
104        recv.cancel().await;
105    }
106}
107
108pub struct OutgoingDatagramStream {
109    pub(crate) inner: Arc<Mutex<P3Socket>>,
110    /// Number of datagrams permitted by most recent `check-send` call.
111    pub(crate) check_send_permit_count: usize,
112    pub(crate) prev_send: Option<MaybeSpawned<Result<(), ErrorCode>>>,
113}
114impl OutgoingDatagramStream {
115    pub(crate) fn new(inner: Arc<Mutex<P3Socket>>) -> Self {
116        Self {
117            inner,
118            check_send_permit_count: 0,
119            prev_send: None,
120        }
121    }
122    pub(crate) fn poll_send_ready(
123        &mut self,
124        cx: &mut std::task::Context<'_>,
125    ) -> std::task::Poll<()> {
126        match &mut self.prev_send {
127            Some(send) => send.poll_ready(cx).map(|_| ()),
128            None => std::task::Poll::Ready(()),
129        }
130    }
131
132    pub(crate) fn try_send(
133        &mut self,
134        data: Vec<u8>,
135        addr: Option<std::net::SocketAddr>,
136    ) -> Result<(), ErrorCode> {
137        if let Some(send) = &mut self.prev_send {
138            if poll_now(|cx| send.poll_ready(cx)).is_none() {
139                return Err(ErrorCode::WouldBlock);
140            }
141
142            let result = self.prev_send.take().unwrap().unwrap_ready();
143            if let Err(e) = result {
144                return Err(e);
145            }
146        }
147
148        debug_assert!(self.prev_send.is_none());
149
150        let mut send = MaybeSpawned::poll_or_spawn(
151            self.inner
152                .lock()
153                .unwrap()
154                .send(data, addr)
155                .map_err(|e| e.into()),
156        );
157        if poll_now(|cx| send.poll_ready(cx)).is_some() {
158            send.unwrap_ready()
159        } else {
160            self.prev_send = Some(send);
161            Ok(())
162        }
163    }
164}
wasmtime_wasi/p2/udp.rs

wasmtime_wasi/p2/
udp.rs
