Skip to main content

wasmtime_wasi/p2/host/
tcp.rs

1use crate::p2::{Pollable, SocketResult, tcp::TcpSocket};
2use crate::p2::{
3    bindings::sockets::{
4        network::{ErrorCode, IpAddressFamily, IpSocketAddress, Network},
5        tcp::{self, ShutdownType},
6    },
7    tcp::AsyncOperation,
8};
9use crate::sockets::{WasiSocketsCtxView, noop_cx};
10use std::net::SocketAddr;
11use std::task::Poll;
12use wasmtime::component::Resource;
13use wasmtime_wasi_io::{
14    poll::DynPollable,
15    streams::{DynInputStream, DynOutputStream},
16};
17
18impl tcp::Host for WasiSocketsCtxView<'_> {}
19
20impl crate::p2::host::tcp::tcp::HostTcpSocket for WasiSocketsCtxView<'_> {
21    async fn start_bind(
22        &mut self,
23        this: Resource<TcpSocket>,
24        network: Resource<Network>,
25        local_address: IpSocketAddress,
26    ) -> SocketResult<()> {
27        // The network resource itself represents the capability to use this
28        // method, so we need to check its validity. Other than that, we have no
29        // use for it.
30        _ = self.table.get(&network)?;
31
32        let local_address: SocketAddr = local_address.into();
33        let socket = self.table.get_mut(&this)?;
34        if socket.in_progress_operation.is_some() {
35            return Err(ErrorCode::ConcurrencyConflict.into());
36        }
37
38        socket.inner.bind(local_address).await?;
39        socket.in_progress_operation = Some(AsyncOperation::Bind);
40        Ok(())
41    }
42
43    fn finish_bind(&mut self, this: Resource<TcpSocket>) -> SocketResult<()> {
44        let socket = self.table.get_mut(&this)?;
45        if socket.in_progress_operation != Some(AsyncOperation::Bind) {
46            return Err(ErrorCode::NotInProgress.into());
47        };
48        socket.in_progress_operation = None;
49        Ok(())
50    }
51
52    fn start_connect(
53        &mut self,
54        this: Resource<TcpSocket>,
55        network: Resource<Network>,
56        remote_address: IpSocketAddress,
57    ) -> SocketResult<()> {
58        // The network resource itself represents the capability to use this
59        // method, so we need to check its validity. Other than that, we have no
60        // use for it.
61        _ = self.table.get(&network)?;
62
63        let remote_address: SocketAddr = remote_address.into();
64        let socket = self.table.get_mut(&this)?;
65        if socket.in_progress_operation.is_some() {
66            return Err(ErrorCode::ConcurrencyConflict.into());
67        }
68
69        socket.inner.start_connect(remote_address)?;
70        socket.in_progress_operation = Some(AsyncOperation::Connect);
71        Ok(())
72    }
73
74    fn finish_connect(
75        &mut self,
76        this: Resource<TcpSocket>,
77    ) -> SocketResult<(Resource<DynInputStream>, Resource<DynOutputStream>)> {
78        let socket = self.table.get_mut(&this)?;
79        if socket.in_progress_operation != Some(AsyncOperation::Connect) {
80            return Err(ErrorCode::NotInProgress.into());
81        };
82
83        let Poll::Ready(result) = socket.inner.poll_finish_connect(&mut noop_cx()) else {
84            return Err(ErrorCode::WouldBlock.into());
85        };
86        socket.in_progress_operation = None;
87
88        if let Err(e) = result {
89            return Err(e.into());
90        }
91
92        let (input, output) = socket.take_streams()?;
93        let input = self.table.push_child(input, &this)?;
94        let output = self.table.push_child(output, &this)?;
95        Ok((input, output))
96    }
97
98    async fn start_listen(&mut self, this: Resource<TcpSocket>) -> SocketResult<()> {
99        let socket = self.table.get_mut(&this)?;
100        if socket.in_progress_operation.is_some() {
101            return Err(ErrorCode::ConcurrencyConflict.into());
102        }
103
104        if !socket.inner.is_bound() {
105            // In WASI 0.2, sockets had to be explicitly bound before listening.
106            return Err(ErrorCode::InvalidState.into());
107        }
108
109        let listener = socket.inner.listen().await?;
110        socket.in_progress_operation = Some(AsyncOperation::Listen);
111        socket.listener = Some(listener);
112        Ok(())
113    }
114
115    fn finish_listen(&mut self, this: Resource<TcpSocket>) -> SocketResult<()> {
116        let socket = self.table.get_mut(&this)?;
117        if socket.in_progress_operation != Some(AsyncOperation::Listen) {
118            return Err(ErrorCode::NotInProgress.into());
119        };
120        socket.in_progress_operation = None;
121        Ok(())
122    }
123
124    fn accept(
125        &mut self,
126        this: Resource<TcpSocket>,
127    ) -> SocketResult<(
128        Resource<TcpSocket>,
129        Resource<DynInputStream>,
130        Resource<DynOutputStream>,
131    )> {
132        let socket = self.table.get_mut(&this)?;
133        let Some(listener) = &mut socket.listener else {
134            return Err(ErrorCode::InvalidState.into());
135        };
136
137        let accepted = match listener.poll_accept(&mut noop_cx()) {
138            Poll::Pending => return Err(ErrorCode::WouldBlock.into()),
139            Poll::Ready(accepted) => accepted,
140        };
141        let mut tcp_socket = TcpSocket::new(accepted);
142        let (input, output) = tcp_socket.take_streams()?;
143
144        let tcp_socket = self.table.push(tcp_socket)?;
145        let input_stream = self.table.push_child(input, &tcp_socket)?;
146        let output_stream = self.table.push_child(output, &tcp_socket)?;
147
148        Ok((tcp_socket, input_stream, output_stream))
149    }
150
151    fn local_address(&mut self, this: Resource<TcpSocket>) -> SocketResult<IpSocketAddress> {
152        let socket = self.table.get_mut(&this)?;
153        Ok(socket.inner.local_address()?.into())
154    }
155
156    fn remote_address(&mut self, this: Resource<TcpSocket>) -> SocketResult<IpSocketAddress> {
157        let socket = self.table.get(&this)?;
158        Ok(socket.inner.remote_address()?.into())
159    }
160
161    fn is_listening(&mut self, this: Resource<TcpSocket>) -> Result<bool, wasmtime::Error> {
162        let socket = self.table.get(&this)?;
163        Ok(socket.inner.is_listening())
164    }
165
166    fn address_family(
167        &mut self,
168        this: Resource<TcpSocket>,
169    ) -> Result<IpAddressFamily, wasmtime::Error> {
170        let socket = self.table.get(&this)?;
171        Ok(socket.inner.address_family().into())
172    }
173
174    fn set_listen_backlog_size(
175        &mut self,
176        this: Resource<TcpSocket>,
177        value: u64,
178    ) -> SocketResult<()> {
179        let socket = self.table.get_mut(&this)?;
180        socket.inner.set_listen_backlog_size(value)?;
181        Ok(())
182    }
183
184    fn keep_alive_enabled(&mut self, this: Resource<TcpSocket>) -> SocketResult<bool> {
185        let socket = self.table.get(&this)?;
186        Ok(socket.inner.keep_alive_enabled()?)
187    }
188
189    fn set_keep_alive_enabled(
190        &mut self,
191        this: Resource<TcpSocket>,
192        value: bool,
193    ) -> SocketResult<()> {
194        let socket = self.table.get(&this)?;
195        socket.inner.set_keep_alive_enabled(value)?;
196        Ok(())
197    }
198
199    fn keep_alive_idle_time(&mut self, this: Resource<TcpSocket>) -> SocketResult<u64> {
200        let socket = self.table.get(&this)?;
201        Ok(socket.inner.keep_alive_idle_time()?)
202    }
203
204    fn set_keep_alive_idle_time(
205        &mut self,
206        this: Resource<TcpSocket>,
207        value: u64,
208    ) -> SocketResult<()> {
209        let socket = self.table.get_mut(&this)?;
210        socket.inner.set_keep_alive_idle_time(value)?;
211        Ok(())
212    }
213
214    fn keep_alive_interval(&mut self, this: Resource<TcpSocket>) -> SocketResult<u64> {
215        let socket = self.table.get(&this)?;
216        Ok(socket.inner.keep_alive_interval()?)
217    }
218
219    fn set_keep_alive_interval(
220        &mut self,
221        this: Resource<TcpSocket>,
222        value: u64,
223    ) -> SocketResult<()> {
224        let socket = self.table.get(&this)?;
225        socket.inner.set_keep_alive_interval(value)?;
226        Ok(())
227    }
228
229    fn keep_alive_count(&mut self, this: Resource<TcpSocket>) -> SocketResult<u32> {
230        let socket = self.table.get(&this)?;
231        Ok(socket.inner.keep_alive_count()?)
232    }
233
234    fn set_keep_alive_count(&mut self, this: Resource<TcpSocket>, value: u32) -> SocketResult<()> {
235        let socket = self.table.get(&this)?;
236        socket.inner.set_keep_alive_count(value)?;
237        Ok(())
238    }
239
240    fn hop_limit(&mut self, this: Resource<TcpSocket>) -> SocketResult<u8> {
241        let socket = self.table.get(&this)?;
242        Ok(socket.inner.hop_limit()?)
243    }
244
245    fn set_hop_limit(&mut self, this: Resource<TcpSocket>, value: u8) -> SocketResult<()> {
246        let socket = self.table.get_mut(&this)?;
247        socket.inner.set_hop_limit(value)?;
248        Ok(())
249    }
250
251    fn receive_buffer_size(&mut self, this: Resource<TcpSocket>) -> SocketResult<u64> {
252        let socket = self.table.get(&this)?;
253        Ok(socket.inner.receive_buffer_size()?)
254    }
255
256    fn set_receive_buffer_size(
257        &mut self,
258        this: Resource<TcpSocket>,
259        value: u64,
260    ) -> SocketResult<()> {
261        let socket = self.table.get_mut(&this)?;
262        socket.inner.set_receive_buffer_size(value)?;
263        Ok(())
264    }
265
266    fn send_buffer_size(&mut self, this: Resource<TcpSocket>) -> SocketResult<u64> {
267        let socket = self.table.get(&this)?;
268        Ok(socket.inner.send_buffer_size()?)
269    }
270
271    fn set_send_buffer_size(&mut self, this: Resource<TcpSocket>, value: u64) -> SocketResult<()> {
272        let socket = self.table.get_mut(&this)?;
273        socket.inner.set_send_buffer_size(value)?;
274        Ok(())
275    }
276
277    fn subscribe(&mut self, this: Resource<TcpSocket>) -> wasmtime::Result<Resource<DynPollable>> {
278        wasmtime_wasi_io::poll::subscribe(self.table, this)
279    }
280
281    fn shutdown(
282        &mut self,
283        this: Resource<TcpSocket>,
284        shutdown_type: ShutdownType,
285    ) -> SocketResult<()> {
286        let socket = self.table.get_mut(&this)?;
287        socket.shutdown(shutdown_type.into())?;
288        Ok(())
289    }
290
291    fn drop(&mut self, this: Resource<TcpSocket>) -> Result<(), wasmtime::Error> {
292        // As in the filesystem implementation, we assume closing a socket
293        // doesn't block.
294        let dropped = self.table.delete(this)?;
295        drop(dropped);
296
297        Ok(())
298    }
299}
300
301#[async_trait::async_trait]
302impl Pollable for TcpSocket {
303    async fn ready(&mut self) {
304        match &self.in_progress_operation {
305            Some(AsyncOperation::Connect) => {
306                _ = std::future::poll_fn(|cx| self.inner.poll_finish_connect(cx)).await;
307            }
308            None if let Some(listener) = &mut self.listener => {
309                std::future::poll_fn(|cx| listener.poll_ready(cx)).await;
310            }
311            _ => {}
312        }
313    }
314}
315
316pub mod sync {
317    use crate::p2::{
318        SocketError,
319        bindings::{
320            sockets::{
321                network::Network,
322                tcp::{self as async_tcp, HostTcpSocket as AsyncHostTcpSocket},
323            },
324            sync::sockets::tcp::{
325                self, Duration, HostTcpSocket, InputStream, IpAddressFamily, IpSocketAddress,
326                OutputStream, Pollable, ShutdownType, TcpSocket,
327            },
328        },
329    };
330    use crate::runtime::in_tokio;
331    use crate::sockets::WasiSocketsCtxView;
332    use wasmtime::component::Resource;
333
334    impl tcp::Host for WasiSocketsCtxView<'_> {}
335
336    impl HostTcpSocket for WasiSocketsCtxView<'_> {
337        fn start_bind(
338            &mut self,
339            self_: Resource<TcpSocket>,
340            network: Resource<Network>,
341            local_address: IpSocketAddress,
342        ) -> Result<(), SocketError> {
343            in_tokio(async {
344                AsyncHostTcpSocket::start_bind(self, self_, network, local_address).await
345            })
346        }
347
348        fn finish_bind(&mut self, self_: Resource<TcpSocket>) -> Result<(), SocketError> {
349            AsyncHostTcpSocket::finish_bind(self, self_)
350        }
351
352        fn start_connect(
353            &mut self,
354            self_: Resource<TcpSocket>,
355            network: Resource<Network>,
356            remote_address: IpSocketAddress,
357        ) -> Result<(), SocketError> {
358            AsyncHostTcpSocket::start_connect(self, self_, network, remote_address)
359        }
360
361        fn finish_connect(
362            &mut self,
363            self_: Resource<TcpSocket>,
364        ) -> Result<(Resource<InputStream>, Resource<OutputStream>), SocketError> {
365            AsyncHostTcpSocket::finish_connect(self, self_)
366        }
367
368        fn start_listen(&mut self, self_: Resource<TcpSocket>) -> Result<(), SocketError> {
369            in_tokio(async { AsyncHostTcpSocket::start_listen(self, self_).await })
370        }
371
372        fn finish_listen(&mut self, self_: Resource<TcpSocket>) -> Result<(), SocketError> {
373            AsyncHostTcpSocket::finish_listen(self, self_)
374        }
375
376        fn accept(
377            &mut self,
378            self_: Resource<TcpSocket>,
379        ) -> Result<
380            (
381                Resource<TcpSocket>,
382                Resource<InputStream>,
383                Resource<OutputStream>,
384            ),
385            SocketError,
386        > {
387            AsyncHostTcpSocket::accept(self, self_)
388        }
389
390        fn local_address(
391            &mut self,
392            self_: Resource<TcpSocket>,
393        ) -> Result<IpSocketAddress, SocketError> {
394            AsyncHostTcpSocket::local_address(self, self_)
395        }
396
397        fn remote_address(
398            &mut self,
399            self_: Resource<TcpSocket>,
400        ) -> Result<IpSocketAddress, SocketError> {
401            AsyncHostTcpSocket::remote_address(self, self_)
402        }
403
404        fn is_listening(&mut self, self_: Resource<TcpSocket>) -> wasmtime::Result<bool> {
405            AsyncHostTcpSocket::is_listening(self, self_)
406        }
407
408        fn address_family(
409            &mut self,
410            self_: Resource<TcpSocket>,
411        ) -> wasmtime::Result<IpAddressFamily> {
412            AsyncHostTcpSocket::address_family(self, self_)
413        }
414
415        fn set_listen_backlog_size(
416            &mut self,
417            self_: Resource<TcpSocket>,
418            value: u64,
419        ) -> Result<(), SocketError> {
420            AsyncHostTcpSocket::set_listen_backlog_size(self, self_, value)
421        }
422
423        fn keep_alive_enabled(&mut self, self_: Resource<TcpSocket>) -> Result<bool, SocketError> {
424            AsyncHostTcpSocket::keep_alive_enabled(self, self_)
425        }
426
427        fn set_keep_alive_enabled(
428            &mut self,
429            self_: Resource<TcpSocket>,
430            value: bool,
431        ) -> Result<(), SocketError> {
432            AsyncHostTcpSocket::set_keep_alive_enabled(self, self_, value)
433        }
434
435        fn keep_alive_idle_time(
436            &mut self,
437            self_: Resource<TcpSocket>,
438        ) -> Result<Duration, SocketError> {
439            AsyncHostTcpSocket::keep_alive_idle_time(self, self_)
440        }
441
442        fn set_keep_alive_idle_time(
443            &mut self,
444            self_: Resource<TcpSocket>,
445            value: Duration,
446        ) -> Result<(), SocketError> {
447            AsyncHostTcpSocket::set_keep_alive_idle_time(self, self_, value)
448        }
449
450        fn keep_alive_interval(
451            &mut self,
452            self_: Resource<TcpSocket>,
453        ) -> Result<Duration, SocketError> {
454            AsyncHostTcpSocket::keep_alive_interval(self, self_)
455        }
456
457        fn set_keep_alive_interval(
458            &mut self,
459            self_: Resource<TcpSocket>,
460            value: Duration,
461        ) -> Result<(), SocketError> {
462            AsyncHostTcpSocket::set_keep_alive_interval(self, self_, value)
463        }
464
465        fn keep_alive_count(&mut self, self_: Resource<TcpSocket>) -> Result<u32, SocketError> {
466            AsyncHostTcpSocket::keep_alive_count(self, self_)
467        }
468
469        fn set_keep_alive_count(
470            &mut self,
471            self_: Resource<TcpSocket>,
472            value: u32,
473        ) -> Result<(), SocketError> {
474            AsyncHostTcpSocket::set_keep_alive_count(self, self_, value)
475        }
476
477        fn hop_limit(&mut self, self_: Resource<TcpSocket>) -> Result<u8, SocketError> {
478            AsyncHostTcpSocket::hop_limit(self, self_)
479        }
480
481        fn set_hop_limit(
482            &mut self,
483            self_: Resource<TcpSocket>,
484            value: u8,
485        ) -> Result<(), SocketError> {
486            AsyncHostTcpSocket::set_hop_limit(self, self_, value)
487        }
488
489        fn receive_buffer_size(&mut self, self_: Resource<TcpSocket>) -> Result<u64, SocketError> {
490            AsyncHostTcpSocket::receive_buffer_size(self, self_)
491        }
492
493        fn set_receive_buffer_size(
494            &mut self,
495            self_: Resource<TcpSocket>,
496            value: u64,
497        ) -> Result<(), SocketError> {
498            AsyncHostTcpSocket::set_receive_buffer_size(self, self_, value)
499        }
500
501        fn send_buffer_size(&mut self, self_: Resource<TcpSocket>) -> Result<u64, SocketError> {
502            AsyncHostTcpSocket::send_buffer_size(self, self_)
503        }
504
505        fn set_send_buffer_size(
506            &mut self,
507            self_: Resource<TcpSocket>,
508            value: u64,
509        ) -> Result<(), SocketError> {
510            AsyncHostTcpSocket::set_send_buffer_size(self, self_, value)
511        }
512
513        fn subscribe(
514            &mut self,
515            self_: Resource<TcpSocket>,
516        ) -> wasmtime::Result<Resource<Pollable>> {
517            AsyncHostTcpSocket::subscribe(self, self_)
518        }
519
520        fn shutdown(
521            &mut self,
522            self_: Resource<TcpSocket>,
523            shutdown_type: ShutdownType,
524        ) -> Result<(), SocketError> {
525            AsyncHostTcpSocket::shutdown(self, self_, shutdown_type.into())
526        }
527
528        fn drop(&mut self, rep: Resource<TcpSocket>) -> wasmtime::Result<()> {
529            AsyncHostTcpSocket::drop(self, rep)
530        }
531    }
532
533    impl From<ShutdownType> for async_tcp::ShutdownType {
534        fn from(other: ShutdownType) -> Self {
535            match other {
536                ShutdownType::Receive => async_tcp::ShutdownType::Receive,
537                ShutdownType::Send => async_tcp::ShutdownType::Send,
538                ShutdownType::Both => async_tcp::ShutdownType::Both,
539            }
540        }
541    }
542}