Skip to main content

wasmtime_wasi/p2/
ip_name_lookup.rs

1use crate::p2::SocketError;
2use crate::p2::bindings::sockets::ip_name_lookup::{Host, HostResolveAddressStream};
3use crate::p2::bindings::sockets::network::{ErrorCode, IpAddress, Network};
4use crate::runtime::poll_now;
5use crate::sockets::ip_name_lookup::resolve_addresses;
6use crate::sockets::{MaybeSpawned, WasiSocketsCtxView};
7use std::net::IpAddr;
8use std::vec;
9use wasmtime::Result;
10use wasmtime::component::Resource;
11use wasmtime_wasi_io::poll::{DynPollable, Pollable, subscribe};
12
13pub struct ResolveAddressStream(MaybeSpawned<Result<vec::IntoIter<IpAddr>, ErrorCode>>);
14
15impl Host for WasiSocketsCtxView<'_> {
16    fn resolve_addresses(
17        &mut self,
18        network: Resource<Network>,
19        name: String,
20    ) -> Result<Resource<ResolveAddressStream>, SocketError> {
21        // The network resource itself represents the capability to use this
22        // method, so we need to check its validity. Other than that, we have no
23        // use for it.
24        _ = self.table.get(&network)?;
25
26        let fut = resolve_addresses(&self.ctx, name);
27        let stream = ResolveAddressStream(MaybeSpawned::poll_or_spawn(async move {
28            Ok(fut.await?.into_iter())
29        }));
30
31        // Attempt to surface errors immediately.
32        if let MaybeSpawned::Ready(Err(err)) = &stream.0 {
33            return Err((*err).into());
34        }
35        Ok(self.table.push(stream)?)
36    }
37}
38
39impl HostResolveAddressStream for WasiSocketsCtxView<'_> {
40    fn resolve_next_address(
41        &mut self,
42        resource: Resource<ResolveAddressStream>,
43    ) -> Result<Option<IpAddress>, SocketError> {
44        let stream: &mut ResolveAddressStream = self.table.get_mut(&resource)?;
45        let Some(result) = poll_now(|cx| stream.0.poll_ready(cx)) else {
46            return Err(ErrorCode::WouldBlock.into());
47        };
48
49        match result {
50            Ok(iter) => Ok(iter.next().map(|addr| addr.into())),
51            Err(err) => Err((*err).into()),
52        }
53    }
54
55    fn subscribe(
56        &mut self,
57        resource: Resource<ResolveAddressStream>,
58    ) -> Result<Resource<DynPollable>> {
59        subscribe(self.table, resource)
60    }
61
62    async fn drop(&mut self, resource: Resource<ResolveAddressStream>) -> Result<()> {
63        let stream = self.table.delete(resource)?;
64        if let MaybeSpawned::Pending(fut) = stream.0 {
65            fut.cancel().await;
66        }
67        Ok(())
68    }
69}
70
71#[async_trait::async_trait]
72impl Pollable for ResolveAddressStream {
73    async fn ready(&mut self) {
74        std::future::poll_fn(|cx| self.0.poll_ready(cx).map(|_| ())).await
75    }
76}
77
78mod sync {
79    use super::ResolveAddressStream;
80    use crate::p2::SocketError;
81    use crate::p2::bindings::sockets::network::{IpAddress, Network};
82    use crate::p2::bindings::sync::sockets::ip_name_lookup::{Host, HostResolveAddressStream};
83    use crate::runtime::in_tokio;
84    use crate::sockets::WasiSocketsCtxView;
85    use wasmtime::Result;
86    use wasmtime::component::Resource;
87    use wasmtime_wasi_io::poll::DynPollable;
88
89    impl Host for WasiSocketsCtxView<'_> {
90        fn resolve_addresses(
91            &mut self,
92            network: Resource<Network>,
93            name: String,
94        ) -> Result<Resource<ResolveAddressStream>, SocketError> {
95            <Self as super::Host>::resolve_addresses(self, network, name)
96        }
97    }
98
99    impl HostResolveAddressStream for WasiSocketsCtxView<'_> {
100        fn resolve_next_address(
101            &mut self,
102            resource: Resource<ResolveAddressStream>,
103        ) -> Result<Option<IpAddress>, SocketError> {
104            <Self as super::HostResolveAddressStream>::resolve_next_address(self, resource)
105        }
106
107        fn subscribe(
108            &mut self,
109            resource: Resource<ResolveAddressStream>,
110        ) -> Result<Resource<DynPollable>> {
111            <Self as super::HostResolveAddressStream>::subscribe(self, resource)
112        }
113
114        fn drop(&mut self, resource: Resource<ResolveAddressStream>) -> Result<()> {
115            in_tokio(<Self as super::HostResolveAddressStream>::drop(
116                self, resource,
117            ))
118        }
119    }
120}