1 use crate::wasi::clocks::monotonic_clock;
2 use crate::wasi::io::poll::{self, Pollable};
3 use crate::wasi::io::streams::{InputStream, OutputStream, StreamError};
4 use crate::wasi::random;
5 use crate::wasi::sockets::instance_network;
6 use crate::wasi::sockets::ip_name_lookup;
7 use crate::wasi::sockets::network::{
8     ErrorCode, IpAddress, IpAddressFamily, IpSocketAddress, Ipv4SocketAddress, Ipv6SocketAddress,
9     Network,
10 };
11 use crate::wasi::sockets::tcp::TcpSocket;
12 use crate::wasi::sockets::udp::{
13     IncomingDatagram, IncomingDatagramStream, OutgoingDatagram, OutgoingDatagramStream, UdpSocket,
14 };
15 use crate::wasi::sockets::{tcp_create_socket, udp_create_socket};
16 use std::ops::Range;
17 
18 const TIMEOUT_NS: u64 = 1_000_000_000;
19 
20 impl Pollable {
21     pub fn block_until(&self, timeout: &Pollable) -> Result<(), ErrorCode> {
22         let ready = poll::poll(&[self, timeout]);
23         assert!(ready.len() > 0);
24         match ready[0] {
25             0 => Ok(()),
26             1 => Err(ErrorCode::Timeout),
27             _ => unreachable!(),
28         }
29     }
30 }
31 
32 impl OutputStream {
33     pub fn blocking_write_util(&self, mut bytes: &[u8]) -> Result<(), StreamError> {
34         let timeout = monotonic_clock::subscribe_duration(TIMEOUT_NS);
35         let pollable = self.subscribe();
36 
37         while !bytes.is_empty() {
38             pollable.block_until(&timeout).expect("write timed out");
39 
40             let permit = self.check_write()?;
41 
42             let len = bytes.len().min(permit as usize);
43             let (chunk, rest) = bytes.split_at(len);
44 
45             self.write(chunk)?;
46 
47             self.blocking_flush()?;
48 
49             bytes = rest;
50         }
51         Ok(())
52     }
53 }
54 
55 impl Network {
56     pub fn default() -> Network {
57         instance_network::instance_network()
58     }
59 
60     pub fn blocking_resolve_addresses(&self, name: &str) -> Result<Vec<IpAddress>, ErrorCode> {
61         let stream = ip_name_lookup::resolve_addresses(&self, name)?;
62 
63         let timeout = monotonic_clock::subscribe_duration(TIMEOUT_NS);
64         let pollable = stream.subscribe();
65 
66         let mut addresses = vec![];
67 
68         loop {
69             match stream.resolve_next_address() {
70                 Ok(Some(addr)) => {
71                     addresses.push(addr);
72                 }
73                 Ok(None) => match addresses[..] {
74                     [] => return Err(ErrorCode::NameUnresolvable),
75                     _ => return Ok(addresses),
76                 },
77                 Err(ErrorCode::WouldBlock) => {
78                     pollable.block_until(&timeout)?;
79                 }
80                 Err(err) => return Err(err),
81             }
82         }
83     }
84 
85     /// Same as `Network::blocking_resolve_addresses` but ignores post validation errors
86     ///
87     /// The ignored error codes signal that the input passed validation
88     /// and a lookup was actually attempted, but failed. These are ignored to
89     /// make the CI tests less flaky.
90     pub fn permissive_blocking_resolve_addresses(
91         &self,
92         name: &str,
93     ) -> Result<Vec<IpAddress>, ErrorCode> {
94         match self.blocking_resolve_addresses(name) {
95             Err(ErrorCode::NameUnresolvable | ErrorCode::TemporaryResolverFailure) => Ok(vec![]),
96             r => r,
97         }
98     }
99 }
100 
101 impl TcpSocket {
102     pub fn new(address_family: IpAddressFamily) -> Result<TcpSocket, ErrorCode> {
103         tcp_create_socket::create_tcp_socket(address_family)
104     }
105 
106     pub fn blocking_bind(
107         &self,
108         network: &Network,
109         local_address: IpSocketAddress,
110     ) -> Result<(), ErrorCode> {
111         let timeout = monotonic_clock::subscribe_duration(TIMEOUT_NS);
112         let sub = self.subscribe();
113 
114         self.start_bind(&network, local_address)?;
115 
116         loop {
117             match self.finish_bind() {
118                 Err(ErrorCode::WouldBlock) => sub.block_until(&timeout)?,
119                 result => return result,
120             }
121         }
122     }
123 
124     pub fn blocking_listen(&self) -> Result<(), ErrorCode> {
125         let timeout = monotonic_clock::subscribe_duration(TIMEOUT_NS);
126         let sub = self.subscribe();
127 
128         self.start_listen()?;
129 
130         loop {
131             match self.finish_listen() {
132                 Err(ErrorCode::WouldBlock) => sub.block_until(&timeout)?,
133                 result => return result,
134             }
135         }
136     }
137 
138     pub fn blocking_connect(
139         &self,
140         network: &Network,
141         remote_address: IpSocketAddress,
142     ) -> Result<(InputStream, OutputStream), ErrorCode> {
143         let timeout = monotonic_clock::subscribe_duration(TIMEOUT_NS);
144         let sub = self.subscribe();
145 
146         self.start_connect(&network, remote_address)?;
147 
148         loop {
149             match self.finish_connect() {
150                 Err(ErrorCode::WouldBlock) => sub.block_until(&timeout)?,
151                 result => return result,
152             }
153         }
154     }
155 
156     pub fn blocking_accept(&self) -> Result<(TcpSocket, InputStream, OutputStream), ErrorCode> {
157         let timeout = monotonic_clock::subscribe_duration(TIMEOUT_NS);
158         let sub = self.subscribe();
159 
160         loop {
161             match self.accept() {
162                 Err(ErrorCode::WouldBlock) => sub.block_until(&timeout)?,
163                 result => return result,
164             }
165         }
166     }
167 }
168 
169 impl UdpSocket {
170     pub fn new(address_family: IpAddressFamily) -> Result<UdpSocket, ErrorCode> {
171         udp_create_socket::create_udp_socket(address_family)
172     }
173 
174     pub fn blocking_bind(
175         &self,
176         network: &Network,
177         local_address: IpSocketAddress,
178     ) -> Result<(), ErrorCode> {
179         let timeout = monotonic_clock::subscribe_duration(TIMEOUT_NS);
180         let sub = self.subscribe();
181 
182         self.start_bind(&network, local_address)?;
183 
184         loop {
185             match self.finish_bind() {
186                 Err(ErrorCode::WouldBlock) => sub.block_until(&timeout)?,
187                 result => return result,
188             }
189         }
190     }
191 
192     pub fn blocking_bind_unspecified(&self, network: &Network) -> Result<(), ErrorCode> {
193         let ip = IpAddress::new_unspecified(self.address_family());
194         let port = 0;
195 
196         self.blocking_bind(network, IpSocketAddress::new(ip, port))
197     }
198 }
199 
200 impl OutgoingDatagramStream {
201     fn blocking_check_send(&self, timeout: &Pollable) -> Result<u64, ErrorCode> {
202         let sub = self.subscribe();
203 
204         loop {
205             match self.check_send() {
206                 Ok(0) => sub.block_until(timeout)?,
207                 result => return result,
208             }
209         }
210     }
211 
212     pub fn blocking_send(&self, mut datagrams: &[OutgoingDatagram]) -> Result<(), ErrorCode> {
213         let timeout = monotonic_clock::subscribe_duration(TIMEOUT_NS);
214 
215         while !datagrams.is_empty() {
216             let permit = self.blocking_check_send(&timeout)?;
217             let chunk_len = datagrams.len().min(permit as usize);
218             match self.send(&datagrams[..chunk_len]) {
219                 Ok(0) => {}
220                 Ok(packets_sent) => {
221                     let packets_sent = packets_sent as usize;
222                     datagrams = &datagrams[packets_sent..];
223                 }
224                 Err(err) => return Err(err),
225             }
226         }
227 
228         Ok(())
229     }
230 }
231 
232 impl IncomingDatagramStream {
233     pub fn blocking_receive(&self, count: Range<u64>) -> Result<Vec<IncomingDatagram>, ErrorCode> {
234         let timeout = monotonic_clock::subscribe_duration(TIMEOUT_NS);
235         let pollable = self.subscribe();
236         let mut datagrams = vec![];
237 
238         loop {
239             match self.receive(count.end - datagrams.len() as u64) {
240                 Ok(mut chunk) => {
241                     datagrams.append(&mut chunk);
242 
243                     if datagrams.len() >= count.start as usize {
244                         return Ok(datagrams);
245                     } else {
246                         pollable.block_until(&timeout)?;
247                     }
248                 }
249                 Err(err) => return Err(err),
250             }
251         }
252     }
253 }
254 
255 impl IpAddress {
256     pub const IPV4_BROADCAST: IpAddress = IpAddress::Ipv4((255, 255, 255, 255));
257 
258     pub const IPV4_LOOPBACK: IpAddress = IpAddress::Ipv4((127, 0, 0, 1));
259     pub const IPV6_LOOPBACK: IpAddress = IpAddress::Ipv6((0, 0, 0, 0, 0, 0, 0, 1));
260 
261     pub const IPV4_UNSPECIFIED: IpAddress = IpAddress::Ipv4((0, 0, 0, 0));
262     pub const IPV6_UNSPECIFIED: IpAddress = IpAddress::Ipv6((0, 0, 0, 0, 0, 0, 0, 0));
263 
264     pub const IPV4_MAPPED_LOOPBACK: IpAddress =
265         IpAddress::Ipv6((0, 0, 0, 0, 0, 0xFFFF, 0x7F00, 0x0001));
266 
267     pub const fn new_loopback(family: IpAddressFamily) -> IpAddress {
268         match family {
269             IpAddressFamily::Ipv4 => Self::IPV4_LOOPBACK,
270             IpAddressFamily::Ipv6 => Self::IPV6_LOOPBACK,
271         }
272     }
273 
274     pub const fn new_unspecified(family: IpAddressFamily) -> IpAddress {
275         match family {
276             IpAddressFamily::Ipv4 => Self::IPV4_UNSPECIFIED,
277             IpAddressFamily::Ipv6 => Self::IPV6_UNSPECIFIED,
278         }
279     }
280 
281     pub const fn family(&self) -> IpAddressFamily {
282         match self {
283             IpAddress::Ipv4(_) => IpAddressFamily::Ipv4,
284             IpAddress::Ipv6(_) => IpAddressFamily::Ipv6,
285         }
286     }
287 }
288 
289 impl PartialEq for IpAddress {
290     fn eq(&self, other: &Self) -> bool {
291         match (self, other) {
292             (Self::Ipv4(left), Self::Ipv4(right)) => left == right,
293             (Self::Ipv6(left), Self::Ipv6(right)) => left == right,
294             _ => false,
295         }
296     }
297 }
298 
299 impl IpSocketAddress {
300     pub const fn new(ip: IpAddress, port: u16) -> IpSocketAddress {
301         match ip {
302             IpAddress::Ipv4(addr) => IpSocketAddress::Ipv4(Ipv4SocketAddress {
303                 port: port,
304                 address: addr,
305             }),
306             IpAddress::Ipv6(addr) => IpSocketAddress::Ipv6(Ipv6SocketAddress {
307                 port: port,
308                 address: addr,
309                 flow_info: 0,
310                 scope_id: 0,
311             }),
312         }
313     }
314 
315     pub const fn ip(&self) -> IpAddress {
316         match self {
317             IpSocketAddress::Ipv4(addr) => IpAddress::Ipv4(addr.address),
318             IpSocketAddress::Ipv6(addr) => IpAddress::Ipv6(addr.address),
319         }
320     }
321 
322     pub const fn port(&self) -> u16 {
323         match self {
324             IpSocketAddress::Ipv4(addr) => addr.port,
325             IpSocketAddress::Ipv6(addr) => addr.port,
326         }
327     }
328 
329     pub const fn family(&self) -> IpAddressFamily {
330         match self {
331             IpSocketAddress::Ipv4(_) => IpAddressFamily::Ipv4,
332             IpSocketAddress::Ipv6(_) => IpAddressFamily::Ipv6,
333         }
334     }
335 }
336 
337 impl PartialEq for Ipv4SocketAddress {
338     fn eq(&self, other: &Self) -> bool {
339         self.port == other.port && self.address == other.address
340     }
341 }
342 
343 impl PartialEq for Ipv6SocketAddress {
344     fn eq(&self, other: &Self) -> bool {
345         self.port == other.port
346             && self.flow_info == other.flow_info
347             && self.address == other.address
348             && self.scope_id == other.scope_id
349     }
350 }
351 
352 impl PartialEq for IpSocketAddress {
353     fn eq(&self, other: &Self) -> bool {
354         match (self, other) {
355             (Self::Ipv4(l0), Self::Ipv4(r0)) => l0 == r0,
356             (Self::Ipv6(l0), Self::Ipv6(r0)) => l0 == r0,
357             _ => false,
358         }
359     }
360 }
361 
362 fn generate_random_u16(range: Range<u16>) -> u16 {
363     let start = range.start as u64;
364     let end = range.end as u64;
365     let port = start + (random::random::get_random_u64() % (end - start));
366     port as u16
367 }
368 
369 /// Execute the inner function with a randomly generated port.
370 /// To prevent random failures, we make a few attempts before giving up.
371 pub fn attempt_random_port<F>(
372     local_address: IpAddress,
373     mut f: F,
374 ) -> Result<IpSocketAddress, ErrorCode>
375 where
376     F: FnMut(IpSocketAddress) -> Result<(), ErrorCode>,
377 {
378     const MAX_ATTEMPTS: u32 = 10;
379     let mut i = 0;
380     loop {
381         i += 1;
382 
383         let port: u16 = generate_random_u16(1024..u16::MAX);
384         let sock_addr = IpSocketAddress::new(local_address, port);
385 
386         match f(sock_addr) {
387             Ok(_) => return Ok(sock_addr),
388             Err(e) if i >= MAX_ATTEMPTS => return Err(e),
389             // Try again if the port is already taken. This can sometimes show up as `AccessDenied` on Windows.
390             Err(ErrorCode::AddressInUse | ErrorCode::AccessDenied) => {}
391             Err(e) => return Err(e),
392         }
393     }
394 }
395