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::sockets::instance_network;
5 use crate::wasi::sockets::network::{
6     ErrorCode, IpAddress, IpAddressFamily, IpSocketAddress, Ipv4SocketAddress, Ipv6SocketAddress,
7     Network,
8 };
9 use crate::wasi::sockets::tcp::TcpSocket;
10 use crate::wasi::sockets::udp::{
11     IncomingDatagram, IncomingDatagramStream, OutgoingDatagram, OutgoingDatagramStream, UdpSocket,
12 };
13 use crate::wasi::sockets::{tcp_create_socket, udp_create_socket};
14 use std::ops::Range;
15 
16 const TIMEOUT_NS: u64 = 1_000_000_000;
17 
18 impl Pollable {
19     pub fn wait(&self) {
20         poll::poll_one(self);
21     }
22 
23     pub fn wait_until(&self, timeout: &Pollable) -> Result<(), ErrorCode> {
24         let ready = poll::poll_list(&[self, timeout]);
25         assert!(ready.len() > 0);
26         match ready[0] {
27             0 => Ok(()),
28             1 => Err(ErrorCode::Timeout),
29             _ => unreachable!(),
30         }
31     }
32 }
33 
34 impl OutputStream {
35     pub fn blocking_write_util(&self, mut bytes: &[u8]) -> Result<(), StreamError> {
36         let pollable = self.subscribe();
37 
38         while !bytes.is_empty() {
39             pollable.wait();
40 
41             let permit = self.check_write()?;
42 
43             let len = bytes.len().min(permit as usize);
44             let (chunk, rest) = bytes.split_at(len);
45 
46             self.write(chunk)?;
47 
48             self.blocking_flush()?;
49 
50             bytes = rest;
51         }
52         Ok(())
53     }
54 }
55 
56 impl Network {
57     pub fn default() -> Network {
58         instance_network::instance_network()
59     }
60 }
61 
62 impl TcpSocket {
63     pub fn new(address_family: IpAddressFamily) -> Result<TcpSocket, ErrorCode> {
64         tcp_create_socket::create_tcp_socket(address_family)
65     }
66 
67     pub fn blocking_bind(
68         &self,
69         network: &Network,
70         local_address: IpSocketAddress,
71     ) -> Result<(), ErrorCode> {
72         let sub = self.subscribe();
73 
74         self.start_bind(&network, local_address)?;
75 
76         loop {
77             match self.finish_bind() {
78                 Err(ErrorCode::WouldBlock) => sub.wait(),
79                 result => return result,
80             }
81         }
82     }
83 
84     pub fn blocking_listen(&self) -> Result<(), ErrorCode> {
85         let sub = self.subscribe();
86 
87         self.start_listen()?;
88 
89         loop {
90             match self.finish_listen() {
91                 Err(ErrorCode::WouldBlock) => sub.wait(),
92                 result => return result,
93             }
94         }
95     }
96 
97     pub fn blocking_connect(
98         &self,
99         network: &Network,
100         remote_address: IpSocketAddress,
101     ) -> Result<(InputStream, OutputStream), ErrorCode> {
102         let sub = self.subscribe();
103 
104         self.start_connect(&network, remote_address)?;
105 
106         loop {
107             match self.finish_connect() {
108                 Err(ErrorCode::WouldBlock) => sub.wait(),
109                 result => return result,
110             }
111         }
112     }
113 
114     pub fn blocking_accept(&self) -> Result<(TcpSocket, InputStream, OutputStream), ErrorCode> {
115         let sub = self.subscribe();
116 
117         loop {
118             match self.accept() {
119                 Err(ErrorCode::WouldBlock) => sub.wait(),
120                 result => return result,
121             }
122         }
123     }
124 }
125 
126 impl UdpSocket {
127     pub fn new(address_family: IpAddressFamily) -> Result<UdpSocket, ErrorCode> {
128         udp_create_socket::create_udp_socket(address_family)
129     }
130 
131     pub fn blocking_bind(
132         &self,
133         network: &Network,
134         local_address: IpSocketAddress,
135     ) -> Result<(), ErrorCode> {
136         let sub = self.subscribe();
137 
138         self.start_bind(&network, local_address)?;
139 
140         loop {
141             match self.finish_bind() {
142                 Err(ErrorCode::WouldBlock) => sub.wait(),
143                 result => return result,
144             }
145         }
146     }
147 }
148 
149 impl OutgoingDatagramStream {
150     fn blocking_check_send(&self, timeout: &Pollable) -> Result<u64, ErrorCode> {
151         let sub = self.subscribe();
152 
153         loop {
154             match self.check_send() {
155                 Ok(0) => sub.wait_until(timeout)?,
156                 result => return result,
157             }
158         }
159     }
160 
161     pub fn blocking_send(&self, mut datagrams: &[OutgoingDatagram]) -> Result<(), ErrorCode> {
162         let timeout = monotonic_clock::subscribe(TIMEOUT_NS, false);
163 
164         while !datagrams.is_empty() {
165             let permit = self.blocking_check_send(&timeout)?;
166             let chunk_len = datagrams.len().min(permit as usize);
167             match self.send(&datagrams[..chunk_len]) {
168                 Ok(0) => {}
169                 Ok(packets_sent) => {
170                     let packets_sent = packets_sent as usize;
171                     datagrams = &datagrams[packets_sent..];
172                 }
173                 Err(err) => return Err(err),
174             }
175         }
176 
177         Ok(())
178     }
179 }
180 
181 impl IncomingDatagramStream {
182     pub fn blocking_receive(&self, count: Range<u64>) -> Result<Vec<IncomingDatagram>, ErrorCode> {
183         let timeout = monotonic_clock::subscribe(TIMEOUT_NS, false);
184         let pollable = self.subscribe();
185         let mut datagrams = vec![];
186 
187         loop {
188             match self.receive(count.end - datagrams.len() as u64) {
189                 Ok(mut chunk) => {
190                     datagrams.append(&mut chunk);
191 
192                     if datagrams.len() >= count.start as usize {
193                         return Ok(datagrams);
194                     } else {
195                         pollable.wait_until(&timeout)?;
196                     }
197                 }
198                 Err(err) => return Err(err),
199             }
200         }
201     }
202 }
203 
204 impl IpAddress {
205     pub const IPV4_BROADCAST: IpAddress = IpAddress::Ipv4((255, 255, 255, 255));
206 
207     pub const IPV4_LOOPBACK: IpAddress = IpAddress::Ipv4((127, 0, 0, 1));
208     pub const IPV6_LOOPBACK: IpAddress = IpAddress::Ipv6((0, 0, 0, 0, 0, 0, 0, 1));
209 
210     pub const IPV4_UNSPECIFIED: IpAddress = IpAddress::Ipv4((0, 0, 0, 0));
211     pub const IPV6_UNSPECIFIED: IpAddress = IpAddress::Ipv6((0, 0, 0, 0, 0, 0, 0, 0));
212 
213     pub const IPV4_MAPPED_LOOPBACK: IpAddress =
214         IpAddress::Ipv6((0, 0, 0, 0, 0, 0xFFFF, 0x7F00, 0x0001));
215 
216     pub const fn new_loopback(family: IpAddressFamily) -> IpAddress {
217         match family {
218             IpAddressFamily::Ipv4 => Self::IPV4_LOOPBACK,
219             IpAddressFamily::Ipv6 => Self::IPV6_LOOPBACK,
220         }
221     }
222 
223     pub const fn new_unspecified(family: IpAddressFamily) -> IpAddress {
224         match family {
225             IpAddressFamily::Ipv4 => Self::IPV4_UNSPECIFIED,
226             IpAddressFamily::Ipv6 => Self::IPV6_UNSPECIFIED,
227         }
228     }
229 
230     pub const fn family(&self) -> IpAddressFamily {
231         match self {
232             IpAddress::Ipv4(_) => IpAddressFamily::Ipv4,
233             IpAddress::Ipv6(_) => IpAddressFamily::Ipv6,
234         }
235     }
236 }
237 
238 impl PartialEq for IpAddress {
239     fn eq(&self, other: &Self) -> bool {
240         match (self, other) {
241             (Self::Ipv4(left), Self::Ipv4(right)) => left == right,
242             (Self::Ipv6(left), Self::Ipv6(right)) => left == right,
243             _ => false,
244         }
245     }
246 }
247 
248 impl IpSocketAddress {
249     pub const fn new(ip: IpAddress, port: u16) -> IpSocketAddress {
250         match ip {
251             IpAddress::Ipv4(addr) => IpSocketAddress::Ipv4(Ipv4SocketAddress {
252                 port: port,
253                 address: addr,
254             }),
255             IpAddress::Ipv6(addr) => IpSocketAddress::Ipv6(Ipv6SocketAddress {
256                 port: port,
257                 address: addr,
258                 flow_info: 0,
259                 scope_id: 0,
260             }),
261         }
262     }
263 
264     pub const fn ip(&self) -> IpAddress {
265         match self {
266             IpSocketAddress::Ipv4(addr) => IpAddress::Ipv4(addr.address),
267             IpSocketAddress::Ipv6(addr) => IpAddress::Ipv6(addr.address),
268         }
269     }
270 
271     pub const fn port(&self) -> u16 {
272         match self {
273             IpSocketAddress::Ipv4(addr) => addr.port,
274             IpSocketAddress::Ipv6(addr) => addr.port,
275         }
276     }
277 
278     pub const fn family(&self) -> IpAddressFamily {
279         match self {
280             IpSocketAddress::Ipv4(_) => IpAddressFamily::Ipv4,
281             IpSocketAddress::Ipv6(_) => IpAddressFamily::Ipv6,
282         }
283     }
284 }
285 
286 impl PartialEq for Ipv4SocketAddress {
287     fn eq(&self, other: &Self) -> bool {
288         self.port == other.port && self.address == other.address
289     }
290 }
291 
292 impl PartialEq for Ipv6SocketAddress {
293     fn eq(&self, other: &Self) -> bool {
294         self.port == other.port
295             && self.flow_info == other.flow_info
296             && self.address == other.address
297             && self.scope_id == other.scope_id
298     }
299 }
300 
301 impl PartialEq for IpSocketAddress {
302     fn eq(&self, other: &Self) -> bool {
303         match (self, other) {
304             (Self::Ipv4(l0), Self::Ipv4(r0)) => l0 == r0,
305             (Self::Ipv6(l0), Self::Ipv6(r0)) => l0 == r0,
306             _ => false,
307         }
308     }
309 }
310