1 use test_programs::p3::wasi::sockets::types::{
2     ErrorCode, IpAddress, IpAddressFamily, IpSocketAddress, UdpSocket,
3 };
4 
5 struct Component;
6 
7 test_programs::p3::export!(Component);
8 
9 const SOME_PORT: u16 = 47; // If the tests pass, this will never actually be connected to.
10 
11 fn test_udp_connect_disconnect_reconnect(family: IpAddressFamily) {
12     let unspecified_addr = IpSocketAddress::new(IpAddress::new_unspecified(family), 0);
13     let remote1 = IpSocketAddress::new(IpAddress::new_loopback(family), 4321);
14     let remote2 = IpSocketAddress::new(IpAddress::new_loopback(family), 4320);
15 
16     let client = UdpSocket::create(family).unwrap();
17     client.bind(unspecified_addr).unwrap();
18 
19     assert_eq!(client.disconnect(), Err(ErrorCode::InvalidState));
20     assert_eq!(client.get_remote_address(), Err(ErrorCode::InvalidState));
21 
22     assert_eq!(client.disconnect(), Err(ErrorCode::InvalidState));
23     assert_eq!(client.get_remote_address(), Err(ErrorCode::InvalidState));
24 
25     _ = client.connect(remote1).unwrap();
26     assert_eq!(client.get_remote_address(), Ok(remote1));
27 
28     _ = client.connect(remote1).unwrap();
29     assert_eq!(client.get_remote_address(), Ok(remote1));
30 
31     _ = client.connect(remote2).unwrap();
32     assert_eq!(client.get_remote_address(), Ok(remote2));
33 
34     _ = client.disconnect().unwrap();
35     assert_eq!(client.get_remote_address(), Err(ErrorCode::InvalidState));
36 
37     _ = client.connect(remote1).unwrap();
38     assert_eq!(client.get_remote_address(), Ok(remote1));
39 }
40 
41 /// `0.0.0.0` / `::` is not a valid remote address in WASI.
42 fn test_udp_connect_unspec(family: IpAddressFamily) {
43     let ip = IpAddress::new_unspecified(family);
44     let addr = IpSocketAddress::new(ip, SOME_PORT);
45     let sock = UdpSocket::create(family).unwrap();
46     sock.bind_unspecified().unwrap();
47 
48     assert!(matches!(
49         sock.connect(addr),
50         Err(ErrorCode::InvalidArgument)
51     ));
52 }
53 
54 /// 0 is not a valid remote port.
55 fn test_udp_connect_port_0(family: IpAddressFamily) {
56     let addr = IpSocketAddress::new(IpAddress::new_loopback(family), 0);
57     let sock = UdpSocket::create(family).unwrap();
58     sock.bind_unspecified().unwrap();
59 
60     assert!(matches!(
61         sock.connect(addr),
62         Err(ErrorCode::InvalidArgument)
63     ));
64 }
65 
66 /// Connect should validate the address family.
67 fn test_udp_connect_wrong_family(family: IpAddressFamily) {
68     let wrong_ip = match family {
69         IpAddressFamily::Ipv4 => IpAddress::IPV6_LOOPBACK,
70         IpAddressFamily::Ipv6 => IpAddress::IPV4_LOOPBACK,
71     };
72     let remote_addr = IpSocketAddress::new(wrong_ip, SOME_PORT);
73 
74     let sock = UdpSocket::create(family).unwrap();
75     sock.bind_unspecified().unwrap();
76 
77     assert!(matches!(
78         sock.connect(remote_addr),
79         Err(ErrorCode::InvalidArgument)
80     ));
81 }
82 
83 fn test_udp_connect_dual_stack() {
84     // Set-up:
85     let v4_server = UdpSocket::create(IpAddressFamily::Ipv4).unwrap();
86     v4_server
87         .bind(IpSocketAddress::new(IpAddress::IPV4_LOOPBACK, 0))
88         .unwrap();
89 
90     let v4_server_addr = v4_server.get_local_address().unwrap();
91     let v6_server_addr =
92         IpSocketAddress::new(IpAddress::IPV4_MAPPED_LOOPBACK, v4_server_addr.port());
93 
94     // Tests:
95     let v6_client = UdpSocket::create(IpAddressFamily::Ipv6).unwrap();
96 
97     v6_client.bind_unspecified().unwrap();
98 
99     // Connecting to an IPv4 address on an IPv6 socket should fail:
100     assert!(matches!(
101         v6_client.connect(v4_server_addr),
102         Err(ErrorCode::InvalidArgument)
103     ));
104 
105     // Connecting to an IPv4-mapped-IPv6 address on an IPv6 socket should fail:
106     assert!(matches!(
107         v6_client.connect(v6_server_addr),
108         Err(ErrorCode::InvalidArgument)
109     ));
110 }
111 
112 impl test_programs::p3::exports::wasi::cli::run::Guest for Component {
113     async fn run() -> Result<(), ()> {
114         test_udp_connect_disconnect_reconnect(IpAddressFamily::Ipv4);
115         test_udp_connect_disconnect_reconnect(IpAddressFamily::Ipv6);
116 
117         test_udp_connect_unspec(IpAddressFamily::Ipv4);
118         test_udp_connect_unspec(IpAddressFamily::Ipv6);
119 
120         test_udp_connect_port_0(IpAddressFamily::Ipv4);
121         test_udp_connect_port_0(IpAddressFamily::Ipv6);
122 
123         test_udp_connect_wrong_family(IpAddressFamily::Ipv4);
124         test_udp_connect_wrong_family(IpAddressFamily::Ipv6);
125 
126         test_udp_connect_dual_stack();
127         Ok(())
128     }
129 }
130 
131 fn main() {}
132