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_rt::async_support::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 
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             fut.await.unwrap()
51         },
52     );
53 
54     // Another client
55     join!(
56         async {
57             let client = TcpSocket::create(family).unwrap();
58             client.connect(addr).await.unwrap();
59             let (mut data_tx, data_rx) = wit_stream::new();
60             join!(
61                 async {
62                     client.send(data_rx).await.unwrap();
63                 },
64                 async {
65                     let remaining = data_tx.write_all(second_message.into()).await;
66                     assert!(remaining.is_empty());
67                     drop(data_tx);
68                 }
69             );
70         },
71         async {
72             let sock = accept.next().await.unwrap();
73             let (mut data_rx, fut) = sock.receive();
74             let (result, data) = data_rx.read(Vec::with_capacity(100)).await;
75             assert_eq!(result, StreamResult::Complete(second_message.len()));
76 
77             // Check that we sent and received our message!
78             assert_eq!(data, second_message); // Not guaranteed to work but should work in practice.
79             fut.await.unwrap()
80         }
81     );
82 }
83 
84 impl test_programs::p3::exports::wasi::cli::run::Guest for Component {
85     async fn run() -> Result<(), ()> {
86         test_tcp_sample_application(
87             IpAddressFamily::Ipv4,
88             IpSocketAddress::Ipv4(Ipv4SocketAddress {
89                 port: 0,                 // use any free port
90                 address: (127, 0, 0, 1), // localhost
91             }),
92         )
93         .await;
94         test_tcp_sample_application(
95             IpAddressFamily::Ipv6,
96             IpSocketAddress::Ipv6(Ipv6SocketAddress {
97                 port: 0,                           // use any free port
98                 address: (0, 0, 0, 0, 0, 0, 0, 1), // localhost
99                 flow_info: 0,
100                 scope_id: 0,
101             }),
102         )
103         .await;
104         Ok(())
105     }
106 }
107 
108 fn main() {}
109