wasmtime_wasi/p2/
udp.rs

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

wasmtime_wasi/p2/
udp.rs
