1 use futures::join; 2 use std::pin::pin; 3 use std::task::{Context, Poll, Waker}; 4 use test_programs::p3::wasi::sockets::types::{ 5 IpAddress, IpAddressFamily, IpSocketAddress, TcpSocket, 6 }; 7 use test_programs::p3::wit_stream; 8 use test_programs::sockets::supports_ipv6; 9 use wit_bindgen::StreamResult; 10 11 struct Component; 12 13 test_programs::p3::export!(Component); 14 15 /// InputStream::read should return `StreamError::Closed` after the connection has been shut down by the server. 16 async fn test_tcp_input_stream_should_be_closed_by_remote_shutdown(family: IpAddressFamily) { 17 setup(family, |server, client| async move { 18 drop(server); 19 20 let (mut client_rx, client_fut) = client.receive(); 21 22 // The input stream should immediately signal StreamError::Closed. 23 // Notably, it should _not_ return an empty list (the wasi-io equivalent of EWOULDBLOCK) 24 // See: https://github.com/bytecodealliance/wasmtime/pull/8968 25 26 // Wait for the shutdown signal to reach the client: 27 assert!(client_rx.next().await.is_none()); 28 assert_eq!(client_fut.await, Ok(())); 29 }) 30 .await; 31 } 32 33 /// InputStream::read should return `StreamError::Closed` after the connection has been shut down locally. 34 async fn test_tcp_input_stream_should_be_closed_by_local_shutdown(family: IpAddressFamily) { 35 setup(family, |server, client| async move { 36 let (mut server_tx, server_rx) = wit_stream::new(); 37 join!( 38 async { 39 server.send(server_rx).await.unwrap(); 40 }, 41 async { 42 // On Linux, `recv` continues to work even after `shutdown(sock, SHUT_RD)` 43 // has been called. To properly test that this behavior doesn't happen in 44 // WASI, we make sure there's some data to read by the client: 45 let rest = server_tx.write_all(b"Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.".into()).await; 46 assert!(rest.is_empty()); 47 drop(server_tx); 48 }, 49 ); 50 51 let (client_rx, client_fut) = client.receive(); 52 53 // Shut down socket locally: 54 drop(client_rx); 55 // Wait for the shutdown signal to reach the client: 56 assert_eq!(client_fut.await, Ok(())); 57 }).await; 58 } 59 60 /// StreamWriter should return `StreamError::Closed` after the connection has been locally shut down for sending. 61 async fn test_tcp_output_stream_should_be_closed_by_local_shutdown(family: IpAddressFamily) { 62 setup(family, |_server, client| async move { 63 let (client_tx, client_rx) = wit_stream::new(); 64 join!( 65 async { 66 client.send(client_rx).await.unwrap(); 67 }, 68 async { 69 // TODO: Verify if send on the stream should return an error 70 //assert!(client_tx.send(b"Hi!".into()).await.is_err()); 71 drop(client_tx); 72 } 73 ); 74 }) 75 .await; 76 } 77 78 /// Calling `shutdown` while the StreamWriter is in the middle of a background write should not cause that write to be lost. 79 async fn test_tcp_shutdown_should_not_lose_data(family: IpAddressFamily) { 80 setup(family, |server, client| async move { 81 // Minimize the local send buffer: 82 client.set_send_buffer_size(1024).unwrap(); 83 let small_buffer_size = client.get_send_buffer_size().unwrap(); 84 85 // Create a significantly bigger buffer, so that we can be pretty sure the `write` won't finish immediately: 86 let big_buffer_size = 100 * small_buffer_size; 87 assert!(big_buffer_size > small_buffer_size); 88 let outgoing_data = vec![0; big_buffer_size as usize]; 89 90 // Submit the oversized buffer and immediately initiate the shutdown: 91 let (mut client_tx, client_rx) = wit_stream::new(); 92 join!( 93 async { 94 client.send(client_rx).await.unwrap(); 95 }, 96 async { 97 let ret = client_tx.write_all(outgoing_data.clone()).await; 98 assert!(ret.is_empty()); 99 drop(client_tx); 100 }, 101 async { 102 // The peer should receive _all_ data: 103 let (server_rx, server_fut) = server.receive(); 104 let incoming_data = server_rx.collect().await; 105 assert_eq!( 106 outgoing_data, incoming_data, 107 "Received data should match the sent data" 108 ); 109 server_fut.await.unwrap(); 110 }, 111 ); 112 }) 113 .await; 114 } 115 116 /// Model a situation where there's a continuous stream of data coming into the 117 /// guest from one side and the other side is reading in chunks but also 118 /// cancelling reads occasionally. Should receive the complete stream of data 119 /// into the result. 120 async fn test_tcp_read_cancellation(family: IpAddressFamily) { 121 // Send 2M of data in 256-byte chunks. 122 const CHUNKS: usize = (2 << 20) / 256; 123 let mut data = [0; 256]; 124 for (i, slot) in data.iter_mut().enumerate() { 125 *slot = i as u8; 126 } 127 128 setup(family, |server, client| async move { 129 // Minimize the local send buffer: 130 client.set_send_buffer_size(1024).unwrap(); 131 132 let (mut client_tx, client_rx) = wit_stream::new(); 133 join!( 134 async { 135 client.send(client_rx).await.unwrap(); 136 }, 137 async { 138 for _ in 0..CHUNKS { 139 let ret = client_tx.write_all(data.to_vec()).await; 140 assert!(ret.is_empty()); 141 } 142 drop(client_tx); 143 }, 144 async { 145 let mut buf = Vec::with_capacity(1024); 146 let (mut server_rx, server_fut) = server.receive(); 147 let mut i = 0_usize; 148 let mut consecutive_zero_length_reads = 0; 149 loop { 150 assert!(buf.is_empty()); 151 let (status, b) = { 152 let mut fut = pin!(server_rx.read(buf)); 153 let mut cx = Context::from_waker(Waker::noop()); 154 match fut.as_mut().poll(&mut cx) { 155 Poll::Ready(pair) => pair, 156 Poll::Pending => fut.cancel(), 157 } 158 }; 159 buf = b; 160 match status { 161 StreamResult::Complete(n) => { 162 assert_eq!(buf.len(), n); 163 for slot in buf.iter_mut() { 164 assert_eq!(*slot, i as u8); 165 i = i.wrapping_add(1); 166 } 167 buf.truncate(0); 168 consecutive_zero_length_reads = 0; 169 } 170 StreamResult::Dropped => break, 171 StreamResult::Cancelled => { 172 assert!(consecutive_zero_length_reads < 10); 173 consecutive_zero_length_reads += 1; 174 server_rx.read(Vec::new()).await; 175 } 176 } 177 } 178 assert_eq!(i, CHUNKS * 256); 179 server_fut.await.unwrap(); 180 }, 181 ); 182 }) 183 .await; 184 } 185 186 impl test_programs::p3::exports::wasi::cli::run::Guest for Component { 187 async fn run() -> Result<(), ()> { 188 test_tcp_input_stream_should_be_closed_by_remote_shutdown(IpAddressFamily::Ipv4).await; 189 test_tcp_input_stream_should_be_closed_by_local_shutdown(IpAddressFamily::Ipv4).await; 190 test_tcp_output_stream_should_be_closed_by_local_shutdown(IpAddressFamily::Ipv4).await; 191 test_tcp_shutdown_should_not_lose_data(IpAddressFamily::Ipv4).await; 192 test_tcp_read_cancellation(IpAddressFamily::Ipv4).await; 193 194 if supports_ipv6() { 195 test_tcp_input_stream_should_be_closed_by_remote_shutdown(IpAddressFamily::Ipv6).await; 196 test_tcp_input_stream_should_be_closed_by_local_shutdown(IpAddressFamily::Ipv6).await; 197 test_tcp_output_stream_should_be_closed_by_local_shutdown(IpAddressFamily::Ipv6).await; 198 test_tcp_shutdown_should_not_lose_data(IpAddressFamily::Ipv6).await; 199 } 200 Ok(()) 201 } 202 } 203 204 fn main() {} 205 206 /// Set up a connected pair of sockets 207 async fn setup<Fut: Future<Output = ()>>( 208 family: IpAddressFamily, 209 body: impl FnOnce(TcpSocket, TcpSocket) -> Fut, 210 ) { 211 let bind_address = IpSocketAddress::new(IpAddress::new_loopback(family), 0); 212 let listener = TcpSocket::create(family).unwrap(); 213 listener.bind(bind_address).unwrap(); 214 let mut accept = listener.listen().unwrap(); 215 let bound_address = listener.get_local_address().unwrap(); 216 let client_socket = TcpSocket::create(family).unwrap(); 217 let ((), accepted_socket) = join!( 218 async { 219 client_socket.connect(bound_address).await.unwrap(); 220 }, 221 async { accept.next().await.unwrap() }, 222 ); 223 body(accepted_socket, client_socket).await; 224 } 225