wasmtime_wasi/p2/
tcp.rs

1use crate::p2::bindings::sockets::network::ErrorCode;
2use crate::p2::{
3    DynInputStream, DynOutputStream, InputStream, OutputStream, Pollable, SocketResult, StreamError,
4};
5use crate::runtime::poll_now;
6use crate::sockets::{
7    MaybeSpawned, TcpListenStream, TcpReceiveStream, TcpSendStream, TcpSocket as P3Socket,
8};
9use std::future::poll_fn;
10use std::mem;
11use std::net::Shutdown;
12use std::sync::Arc;
13use std::sync::Mutex;
14use std::task::{Poll, ready};
15use wasmtime::Result;
16use wasmtime_wasi_io::streams::StreamResult;
17
18/// A TCP socket + associated p2 bookkeeping.
19pub struct TcpSocket {
20    pub(crate) inner: P3Socket,
21    pub(crate) in_progress_operation: Option<AsyncOperation>,
22    pub(crate) listener: Option<TcpListenStream>,
23    reader: Option<TcpReader>,
24    writer: Option<TcpWriter>,
25}
26
27#[derive(Clone, Copy, Debug, PartialEq, Eq)]
28pub(crate) enum AsyncOperation {
29    Bind,
30    Connect,
31    Listen,
32}
33
34impl TcpSocket {
35    pub(crate) fn new(inner: P3Socket) -> Self {
36        Self {
37            inner,
38            in_progress_operation: None,
39            listener: None,
40            reader: None,
41            writer: None,
42        }
43    }
44    pub(crate) fn take_streams(
45        &mut self,
46        max_io_size: usize,
47    ) -> SocketResult<(DynInputStream, DynOutputStream)> {
48        let reader = TcpReader::new(self.inner.take_receive_stream()?, max_io_size);
49        let writer = TcpWriter::new(self.inner.take_send_stream()?, max_io_size);
50        self.reader = Some(reader.clone());
51        self.writer = Some(writer.clone());
52        let input: DynInputStream = Box::new(reader);
53        let output: DynOutputStream = Box::new(writer);
54        Ok((input, output))
55    }
56    pub(crate) fn shutdown(&mut self, how: Shutdown) -> SocketResult<()> {
57        let reader = self.reader.as_mut().ok_or(ErrorCode::InvalidState)?;
58        let writer = self.writer.as_mut().ok_or(ErrorCode::InvalidState)?;
59
60        if let Shutdown::Both | Shutdown::Read = how {
61            reader.state.lock().unwrap().shutdown();
62        }
63
64        if let Shutdown::Both | Shutdown::Write = how {
65            writer.state.lock().unwrap().shutdown();
66        }
67
68        Ok(())
69    }
70}
71
72enum ReadState {
73    Open(TcpReceiveStream),
74    Closed,
75}
76impl ReadState {
77    fn read(&mut self, size: usize) -> StreamResult<bytes::Bytes> {
78        let Self::Open(stream) = self else {
79            return Err(StreamError::Closed);
80        };
81        if size == 0 {
82            return Ok(bytes::Bytes::new());
83        }
84        let mut buf = bytes::BytesMut::zeroed(size);
85        let n = match poll_now(|cx| stream.poll_read(cx, &mut buf)) {
86            None => 0,
87            Some(Ok(0)) => {
88                *self = ReadState::Closed;
89                return Err(StreamError::Closed);
90            }
91            Some(Ok(n)) => n,
92            Some(Err(e)) => {
93                *self = ReadState::Closed;
94                return Err(StreamError::LastOperationFailed(e.into()));
95            }
96        };
97
98        buf.truncate(n);
99        Ok(buf.freeze())
100    }
101
102    fn shutdown(&mut self) {
103        *self = ReadState::Closed;
104    }
105
106    fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<()> {
107        match self {
108            Self::Open(stream) => stream.poll_ready(cx),
109            Self::Closed => Poll::Ready(()),
110        }
111    }
112}
113
114#[derive(Clone)]
115struct TcpReader {
116    state: Arc<Mutex<ReadState>>,
117    max_read_size: usize,
118}
119impl TcpReader {
120    fn new(stream: TcpReceiveStream, max_read_size: usize) -> Self {
121        Self {
122            state: Arc::new(Mutex::new(ReadState::Open(stream))),
123            max_read_size,
124        }
125    }
126}
127
128#[async_trait::async_trait]
129impl InputStream for TcpReader {
130    fn read(&mut self, size: usize) -> StreamResult<bytes::Bytes> {
131        self.state
132            .lock()
133            .unwrap()
134            .read(size.min(self.max_read_size))
135    }
136}
137
138#[async_trait::async_trait]
139impl Pollable for TcpReader {
140    async fn ready(&mut self) {
141        std::future::poll_fn(|cx| self.state.lock().unwrap().poll_ready(cx)).await
142    }
143}
144
145/// A cloneable subset of StreamError
146#[derive(Debug, Clone)]
147enum WriteError {
148    Closed,
149    LastOperationFailed(ErrorCode),
150}
151impl From<WriteError> for StreamError {
152    fn from(err: WriteError) -> Self {
153        match err {
154            WriteError::Closed => StreamError::Closed,
155            WriteError::LastOperationFailed(e) => StreamError::LastOperationFailed(e.into()),
156        }
157    }
158}
159
160enum WriteState {
161    Ready(TcpSendStream, usize),
162    Writing(MaybeSpawned<Result<TcpSendStream, WriteError>>),
163    Closing(MaybeSpawned<Result<(), WriteError>>),
164    Closed(WriteError),
165}
166
167impl WriteState {
168    fn take(&mut self) -> WriteState {
169        mem::replace(self, WriteState::Closed(WriteError::Closed))
170    }
171
172    fn check_write(&mut self, max_write_size: usize) -> StreamResult<usize> {
173        match poll_now(|cx| self.poll_ready(cx)) {
174            None => Ok(0),
175            Some(Ok((_, permit))) => {
176                *permit = max_write_size;
177                Ok(*permit)
178            }
179            Some(Err(e)) => Err(e),
180        }
181    }
182
183    fn write(&mut self, mut bytes: bytes::Bytes) -> StreamResult<()> {
184        if bytes.is_empty() {
185            return match self {
186                WriteState::Closed(e) => Err(e.clone().into()),
187                _ => Ok(()),
188            };
189        }
190
191        let mut stream = match self {
192            WriteState::Ready(_, permit) if bytes.len() <= *permit => {
193                let WriteState::Ready(stream, _) = self.take() else {
194                    unreachable!()
195                };
196                stream
197            }
198            WriteState::Closed(e) => {
199                return Err(e.clone().into());
200            }
201            _ => {
202                return Err(StreamError::Trap(wasmtime::format_err!(
203                    "not permitted to write {} bytes",
204                    bytes.len()
205                )));
206            }
207        };
208
209        *self = WriteState::Writing(MaybeSpawned::poll_or_spawn(async move {
210            while !bytes.is_empty() {
211                match stream.write(&bytes).await {
212                    Ok(n) => {
213                        let _ = bytes.split_to(n);
214                    }
215                    Err(crate::sockets::ErrorCode::ConnectionBroken) => {
216                        return Err(WriteError::Closed);
217                    }
218                    Err(e) => {
219                        return Err(WriteError::LastOperationFailed(e.into()));
220                    }
221                }
222            }
223
224            Ok(stream)
225        }));
226
227        // Attempt to finish the write, surfacing potential errors immediately:
228        match poll_now(|cx| self.poll_ready(cx)) {
229            None | Some(Ok(_)) => Ok(()),
230            Some(Err(e)) => Err(e),
231        }
232    }
233
234    fn flush(&mut self) -> StreamResult<()> {
235        // `flush` is a no-op here. Writes happen on background tasks and will
236        // always be delivered to the OS as soon as possible. There's nothing
237        // for `flush` to do here that will speed up that process.
238        match self {
239            WriteState::Ready(..) | WriteState::Writing(_) | WriteState::Closing(_) => Ok(()),
240            WriteState::Closed(e) => Err(e.clone().into()),
241        }
242    }
243
244    pub(crate) fn shutdown(&mut self) {
245        *self = match self.take() {
246            // No write in progress, immediately drop the inner stream:
247            WriteState::Ready(..) => WriteState::Closed(WriteError::Closed),
248
249            // Schedule the shutdown after the current write has finished:
250            WriteState::Writing(write) => {
251                WriteState::Closing(MaybeSpawned::poll_or_spawn(async move {
252                    _ = write.into_future().await?;
253                    Ok(())
254                }))
255            }
256
257            s => s,
258        };
259    }
260
261    fn poll_ready(
262        &mut self,
263        cx: &mut std::task::Context<'_>,
264    ) -> Poll<StreamResult<(&mut TcpSendStream, &mut usize)>> {
265        match self {
266            WriteState::Writing(write) => {
267                ready!(write.poll_ready(cx));
268                let WriteState::Writing(write) = self.take() else {
269                    unreachable!()
270                };
271                *self = match write.unwrap_ready() {
272                    Ok(stream) => WriteState::Ready(stream, 0),
273                    Err(err) => WriteState::Closed(err),
274                };
275            }
276            WriteState::Closing(close) => {
277                ready!(close.poll_ready(cx));
278                let WriteState::Closing(close) = self.take() else {
279                    unreachable!()
280                };
281                *self = match close.unwrap_ready() {
282                    Ok(()) => WriteState::Closed(WriteError::Closed),
283                    Err(err) => WriteState::Closed(err),
284                };
285            }
286            _ => {}
287        }
288
289        match self {
290            WriteState::Ready(stream, permit) => match stream.poll_ready(cx) {
291                Poll::Ready(()) => Poll::Ready(Ok((stream, permit))),
292                Poll::Pending => Poll::Pending,
293            },
294            WriteState::Writing(..) | WriteState::Closing(..) => Poll::Pending,
295            WriteState::Closed(e) => Poll::Ready(Err(e.clone().into())),
296        }
297    }
298}
299
300#[derive(Clone)]
301struct TcpWriter {
302    state: Arc<Mutex<WriteState>>,
303    max_write_size: usize,
304}
305impl TcpWriter {
306    fn new(stream: TcpSendStream, max_write_size: usize) -> Self {
307        Self {
308            state: Arc::new(Mutex::new(WriteState::Ready(stream, 0))),
309            max_write_size,
310        }
311    }
312}
313
314#[async_trait::async_trait]
315impl OutputStream for TcpWriter {
316    fn write(&mut self, bytes: bytes::Bytes) -> StreamResult<()> {
317        self.state.lock().unwrap().write(bytes)
318    }
319
320    fn flush(&mut self) -> StreamResult<()> {
321        self.state.lock().unwrap().flush()
322    }
323
324    fn check_write(&mut self) -> StreamResult<usize> {
325        self.state.lock().unwrap().check_write(self.max_write_size)
326    }
327
328    async fn cancel(&mut self) {
329        // Wait for background writes to finish in order to prevent silently
330        // dropping data that (from the guest's perspective) was already written.
331        self.ready().await
332    }
333}
334
335#[async_trait::async_trait]
336impl Pollable for TcpWriter {
337    async fn ready(&mut self) {
338        poll_fn(|cx| self.state.lock().unwrap().poll_ready(cx).map(|_| ())).await;
339    }
340}
wasmtime_wasi/p2/tcp.rs

wasmtime_wasi/p2/
tcp.rs
