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