1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use crate::preview2::bindings::sockets::ip_name_lookup::{Host, HostResolveAddressStream};
use crate::preview2::bindings::sockets::network::{ErrorCode, IpAddress, IpAddressFamily, Network};
use crate::preview2::poll::{subscribe, Pollable, Subscribe};
use crate::preview2::{spawn_blocking, AbortOnDropJoinHandle, SocketError, WasiView};
use anyhow::Result;
use std::mem;
use std::net::{SocketAddr, ToSocketAddrs};
use std::pin::Pin;
use std::vec;
use wasmtime::component::Resource;

pub enum ResolveAddressStream {
    Waiting(AbortOnDropJoinHandle<Result<Vec<IpAddress>, SocketError>>),
    Done(Result<vec::IntoIter<IpAddress>, SocketError>),
}

#[async_trait::async_trait]
impl<T: WasiView> Host for T {
    fn resolve_addresses(
        &mut self,
        network: Resource<Network>,
        name: String,
        family: Option<IpAddressFamily>,
        include_unavailable: bool,
    ) -> Result<Resource<ResolveAddressStream>, SocketError> {
        let network = self.table().get_resource(&network)?;

        // `Host::parse` serves us two functions:
        // 1. validate the input is not an IP address,
        // 2. convert unicode domains to punycode.
        let name = match url::Host::parse(&name).map_err(|_| ErrorCode::InvalidArgument)? {
            url::Host::Domain(name) => name,
            url::Host::Ipv4(_) => return Err(ErrorCode::InvalidArgument.into()),
            url::Host::Ipv6(_) => return Err(ErrorCode::InvalidArgument.into()),
        };

        if !network.allow_ip_name_lookup {
            return Err(ErrorCode::PermanentResolverFailure.into());
        }

        // ignored for now, should probably have a future PR to actually take
        // this into account. This would require invoking `getaddrinfo` directly
        // rather than using the standard library to do it for us.
        let _ = include_unavailable;

        // For now use the standard library to perform actual resolution through
        // the usage of the `ToSocketAddrs` trait. This blocks the current
        // thread, so use `spawn_blocking`. Finally note that this is only
        // resolving names, not ports, so force the port to be 0.
        let task = spawn_blocking(move || -> Result<Vec<_>, SocketError> {
            let result = (name.as_str(), 0)
                .to_socket_addrs()
                .map_err(|_| ErrorCode::NameUnresolvable)?; // If/when we use `getaddrinfo` directly, map the error properly.
            Ok(result
                .filter_map(|addr| {
                    // In lieu of preventing these addresses from being resolved
                    // in the first place, filter them out here.
                    match addr {
                        SocketAddr::V4(addr) => match family {
                            None | Some(IpAddressFamily::Ipv4) => {
                                let [a, b, c, d] = addr.ip().octets();
                                Some(IpAddress::Ipv4((a, b, c, d)))
                            }
                            Some(IpAddressFamily::Ipv6) => None,
                        },
                        SocketAddr::V6(addr) => match family {
                            None | Some(IpAddressFamily::Ipv6) => {
                                let [a, b, c, d, e, f, g, h] = addr.ip().segments();
                                Some(IpAddress::Ipv6((a, b, c, d, e, f, g, h)))
                            }
                            Some(IpAddressFamily::Ipv4) => None,
                        },
                    }
                })
                .collect())
        });
        let resource = self
            .table_mut()
            .push_resource(ResolveAddressStream::Waiting(task))?;
        Ok(resource)
    }
}

#[async_trait::async_trait]
impl<T: WasiView> HostResolveAddressStream for T {
    fn resolve_next_address(
        &mut self,
        resource: Resource<ResolveAddressStream>,
    ) -> Result<Option<IpAddress>, SocketError> {
        let stream = self.table_mut().get_resource_mut(&resource)?;
        loop {
            match stream {
                ResolveAddressStream::Waiting(future) => {
                    match crate::preview2::poll_noop(Pin::new(future)) {
                        Some(result) => {
                            *stream = ResolveAddressStream::Done(result.map(|v| v.into_iter()));
                        }
                        None => return Err(ErrorCode::WouldBlock.into()),
                    }
                }
                ResolveAddressStream::Done(slot @ Err(_)) => {
                    mem::replace(slot, Ok(Vec::new().into_iter()))?;
                    unreachable!();
                }
                ResolveAddressStream::Done(Ok(iter)) => return Ok(iter.next()),
            }
        }
    }

    fn subscribe(
        &mut self,
        resource: Resource<ResolveAddressStream>,
    ) -> Result<Resource<Pollable>> {
        subscribe(self.table_mut(), resource)
    }

    fn drop(&mut self, resource: Resource<ResolveAddressStream>) -> Result<()> {
        self.table_mut().delete_resource(resource)?;
        Ok(())
    }
}

#[async_trait::async_trait]
impl Subscribe for ResolveAddressStream {
    async fn ready(&mut self) {
        if let ResolveAddressStream::Waiting(future) = self {
            *self = ResolveAddressStream::Done(future.await.map(|v| v.into_iter()));
        }
    }
}