1 use futures::join;
2 use test_programs::p3::wasi::sockets::types::{
3     IpAddressFamily, IpSocketAddress, Ipv4SocketAddress, Ipv6SocketAddress, TcpSocket,
4 };
5 use test_programs::p3::wit_stream;
6 use test_programs::sockets::supports_ipv6;
7 use wit_bindgen::StreamResult;
8 
9 struct Component;
10 
11 test_programs::p3::export!(Component);
12 
13 async fn test_tcp_sample_application(family: IpAddressFamily, bind_address: IpSocketAddress) {
14     let first_message = b"Hello, world!";
15     let second_message = b"Greetings, planet!";
16 
17     let listener = TcpSocket::create(family).unwrap();
18 
19     listener.bind(bind_address).unwrap();
20     listener.set_listen_backlog_size(32).unwrap();
21     let mut accept = listener.listen().unwrap();
22 
23     let addr = listener.get_local_address().unwrap();
24 
25     join!(
26         async {
27             let client = TcpSocket::create(family).unwrap();
28             client.connect(addr).await.unwrap();
29             let (mut data_tx, data_rx) = wit_stream::new();
30             join!(
31                 async {
32                     client.send(data_rx).await.unwrap();
33                 },
34                 async {
35                     let (result, _) = data_tx.write(vec![]).await;
36                     assert_eq!(result, StreamResult::Complete(0));
37                     let remaining = data_tx.write_all(first_message.into()).await;
38                     assert!(remaining.is_empty());
39                     drop(data_tx);
40                 }
41             );
42         },
43         async {
44             let sock = accept.next().await.unwrap();
45             let (mut data_rx, fut) = sock.receive();
46             let (result, data) = data_rx.read(Vec::with_capacity(100)).await;
47             assert_eq!(result, StreamResult::Complete(first_message.len()));
48             // Check that we sent and received our message!
49             assert_eq!(data, first_message); // Not guaranteed to work but should work in practice.
50 
51             let (result, data) = data_rx.read(Vec::with_capacity(1)).await;
52             assert_eq!(result, StreamResult::Dropped);
53             assert_eq!(data, []);
54 
55             fut.await.unwrap();
56         },
57     );
58 
59     // Another client
60     join!(
61         async {
62             let client = TcpSocket::create(family).unwrap();
63             client.connect(addr).await.unwrap();
64             let (mut data_tx, data_rx) = wit_stream::new();
65             join!(
66                 async {
67                     client.send(data_rx).await.unwrap();
68                 },
69                 async {
70                     let remaining = data_tx.write_all(second_message.into()).await;
71                     assert!(remaining.is_empty());
72                     drop(data_tx);
73                 }
74             );
75         },
76         async {
77             let sock = accept.next().await.unwrap();
78             let (mut data_rx, fut) = sock.receive();
79             let (result, data) = data_rx.read(Vec::with_capacity(100)).await;
80             assert_eq!(result, StreamResult::Complete(second_message.len()));
81             // Check that we sent and received our message!
82             assert_eq!(data, second_message); // Not guaranteed to work but should work in practice.
83 
84             let (result, data) = data_rx.read(Vec::with_capacity(1)).await;
85             assert_eq!(result, StreamResult::Dropped);
86             assert_eq!(data, []);
87 
88             fut.await.unwrap();
89         }
90     );
91 }
92 
93 impl test_programs::p3::exports::wasi::cli::run::Guest for Component {
94     async fn run() -> Result<(), ()> {
95         test_tcp_sample_application(
96             IpAddressFamily::Ipv4,
97             IpSocketAddress::Ipv4(Ipv4SocketAddress {
98                 port: 0,                 // use any free port
99                 address: (127, 0, 0, 1), // localhost
100             }),
101         )
102         .await;
103         if supports_ipv6() {
104             test_tcp_sample_application(
105                 IpAddressFamily::Ipv6,
106                 IpSocketAddress::Ipv6(Ipv6SocketAddress {
107                     port: 0,                           // use any free port
108                     address: (0, 0, 0, 0, 0, 0, 0, 1), // localhost
109                     flow_info: 0,
110                     scope_id: 0,
111                 }),
112             )
113             .await;
114         }
115         Ok(())
116     }
117 }
118 
119 fn main() {}
120