1 use core::ops::Range;
2 
3 use crate::p3::wasi::random;
4 use crate::p3::wasi::sockets::types::{
5     ErrorCode, IpAddress, IpAddressFamily, IpSocketAddress, Ipv4SocketAddress, Ipv6SocketAddress,
6     UdpSocket,
7 };
8 
9 impl IpAddress {
10     pub const IPV4_BROADCAST: IpAddress = IpAddress::Ipv4((255, 255, 255, 255));
11 
12     pub const IPV4_LOOPBACK: IpAddress = IpAddress::Ipv4((127, 0, 0, 1));
13     pub const IPV6_LOOPBACK: IpAddress = IpAddress::Ipv6((0, 0, 0, 0, 0, 0, 0, 1));
14 
15     pub const IPV4_UNSPECIFIED: IpAddress = IpAddress::Ipv4((0, 0, 0, 0));
16     pub const IPV6_UNSPECIFIED: IpAddress = IpAddress::Ipv6((0, 0, 0, 0, 0, 0, 0, 0));
17 
18     pub const IPV4_MAPPED_LOOPBACK: IpAddress =
19         IpAddress::Ipv6((0, 0, 0, 0, 0, 0xFFFF, 0x7F00, 0x0001));
20 
21     pub const fn new_loopback(family: IpAddressFamily) -> IpAddress {
22         match family {
23             IpAddressFamily::Ipv4 => Self::IPV4_LOOPBACK,
24             IpAddressFamily::Ipv6 => Self::IPV6_LOOPBACK,
25         }
26     }
27 
28     pub const fn new_unspecified(family: IpAddressFamily) -> IpAddress {
29         match family {
30             IpAddressFamily::Ipv4 => Self::IPV4_UNSPECIFIED,
31             IpAddressFamily::Ipv6 => Self::IPV6_UNSPECIFIED,
32         }
33     }
34 
35     pub const fn family(&self) -> IpAddressFamily {
36         match self {
37             IpAddress::Ipv4(_) => IpAddressFamily::Ipv4,
38             IpAddress::Ipv6(_) => IpAddressFamily::Ipv6,
39         }
40     }
41 }
42 
43 impl PartialEq for IpAddress {
44     fn eq(&self, other: &Self) -> bool {
45         match (self, other) {
46             (Self::Ipv4(left), Self::Ipv4(right)) => left == right,
47             (Self::Ipv6(left), Self::Ipv6(right)) => left == right,
48             _ => false,
49         }
50     }
51 }
52 
53 impl IpSocketAddress {
54     pub const fn new(ip: IpAddress, port: u16) -> IpSocketAddress {
55         match ip {
56             IpAddress::Ipv4(addr) => IpSocketAddress::Ipv4(Ipv4SocketAddress {
57                 port,
58                 address: addr,
59             }),
60             IpAddress::Ipv6(addr) => IpSocketAddress::Ipv6(Ipv6SocketAddress {
61                 port,
62                 address: addr,
63                 flow_info: 0,
64                 scope_id: 0,
65             }),
66         }
67     }
68 
69     pub const fn ip(&self) -> IpAddress {
70         match self {
71             IpSocketAddress::Ipv4(addr) => IpAddress::Ipv4(addr.address),
72             IpSocketAddress::Ipv6(addr) => IpAddress::Ipv6(addr.address),
73         }
74     }
75 
76     pub const fn port(&self) -> u16 {
77         match self {
78             IpSocketAddress::Ipv4(addr) => addr.port,
79             IpSocketAddress::Ipv6(addr) => addr.port,
80         }
81     }
82 
83     pub const fn family(&self) -> IpAddressFamily {
84         match self {
85             IpSocketAddress::Ipv4(_) => IpAddressFamily::Ipv4,
86             IpSocketAddress::Ipv6(_) => IpAddressFamily::Ipv6,
87         }
88     }
89 }
90 
91 impl PartialEq for Ipv4SocketAddress {
92     fn eq(&self, other: &Self) -> bool {
93         self.port == other.port && self.address == other.address
94     }
95 }
96 
97 impl PartialEq for Ipv6SocketAddress {
98     fn eq(&self, other: &Self) -> bool {
99         self.port == other.port
100             && self.flow_info == other.flow_info
101             && self.address == other.address
102             && self.scope_id == other.scope_id
103     }
104 }
105 
106 impl PartialEq for IpSocketAddress {
107     fn eq(&self, other: &Self) -> bool {
108         match (self, other) {
109             (Self::Ipv4(l0), Self::Ipv4(r0)) => l0 == r0,
110             (Self::Ipv6(l0), Self::Ipv6(r0)) => l0 == r0,
111             _ => false,
112         }
113     }
114 }
115 
116 fn generate_random_u16(range: Range<u16>) -> u16 {
117     let start = range.start as u64;
118     let end = range.end as u64;
119     let port = start + (random::random::get_random_u64() % (end - start));
120     port as u16
121 }
122 
123 /// Execute the inner function with a randomly generated port.
124 /// To prevent random failures, we make a few attempts before giving up.
125 pub fn attempt_random_port<F>(
126     local_address: IpAddress,
127     mut f: F,
128 ) -> Result<IpSocketAddress, ErrorCode>
129 where
130     F: FnMut(IpSocketAddress) -> Result<(), ErrorCode>,
131 {
132     const MAX_ATTEMPTS: u32 = 10;
133     let mut i = 0;
134     loop {
135         i += 1;
136 
137         let port: u16 = generate_random_u16(1024..u16::MAX);
138         let sock_addr = IpSocketAddress::new(local_address, port);
139 
140         match f(sock_addr) {
141             Ok(_) => return Ok(sock_addr),
142             Err(e) if i >= MAX_ATTEMPTS => return Err(e),
143             // Try again if the port is already taken. This can sometimes show up as `AccessDenied` on Windows.
144             Err(ErrorCode::AddressInUse | ErrorCode::AccessDenied) => {}
145             Err(e) => return Err(e),
146         }
147     }
148 }
149 
150 impl UdpSocket {
151     pub fn bind_unspecified(&self) -> Result<(), ErrorCode> {
152         let ip = IpAddress::new_unspecified(self.get_address_family());
153         let port = 0;
154 
155         self.bind(IpSocketAddress::new(ip, port))
156     }
157 }
158