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