diff --git a/src/websocket.rs b/src/websocket.rs index 2af824fb1..dc3deedf5 100644 --- a/src/websocket.rs +++ b/src/websocket.rs @@ -18,7 +18,7 @@ use std::{ net::SocketAddr, pin::Pin, sync::Arc, - task::{ready, Context, Poll}, + task::{Context, Poll}, time::Duration, }; use tokio::{ @@ -47,7 +47,10 @@ impl AsyncRead for FlushBeforeRead { cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { - ready!(Pin::new(&mut self.0).poll_flush(cx))?; + // A blocked writer must not prevent reading data that lets the peer make progress. + if let Poll::Ready(Err(err)) = Pin::new(&mut self.0).poll_flush(cx) { + return Poll::Ready(Err(err)); + } Pin::new(&mut self.0).poll_read(cx, buf) } } @@ -525,6 +528,9 @@ pub fn check_ws(endpoint: &str) -> String { format!("{}://{}", protocol, address) } +#[cfg(test)] +mod proxy_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/src/websocket/proxy_tests.rs b/src/websocket/proxy_tests.rs new file mode 100644 index 000000000..06db80965 --- /dev/null +++ b/src/websocket/proxy_tests.rs @@ -0,0 +1,149 @@ +use super::FlushBeforeRead; +use crate::{proxy::Proxy, tcp::DynTcpStream, tls::TlsType}; +use futures::{SinkExt, StreamExt}; +use rustls_pki_types::{pem::PemObject, CertificateDer, PrivateKeyDer}; +use std::{sync::Arc, time::Duration}; +use tokio::{ + io::{duplex, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufStream}, + time::timeout, +}; +use tokio_native_tls::native_tls; +use tokio_rustls::{rustls, TlsAcceptor}; +use tokio_socks::IntoTargetAddr; +use tokio_tungstenite::{accept_async, client_async_tls_with_config, Connector, WebSocketStream}; + +fn acceptor() -> TlsAcceptor { + let cert = + CertificateDer::from_pem_slice(include_bytes!("../../tests/data/websocket-cert.pem")) + .unwrap(); + let key = PrivateKeyDer::from_pem_slice(include_bytes!("../../tests/data/websocket-key.pem")) + .unwrap(); + let config = rustls::ServerConfig::builder_with_provider(Arc::new( + rustls::crypto::ring::default_provider(), + )) + .with_safe_default_protocol_versions() + .unwrap() + .with_no_client_auth() + .with_single_cert(vec![cert], key) + .unwrap(); + TlsAcceptor::from(Arc::new(config)) +} + +async fn exchange( + ws: WebSocketStream, + sent: u8, + received: u8, +) { + let (mut writer, mut reader) = ws.split(); + let send = async { + for sequence in 0..16 { + writer + .send(vec![sent + sequence; 256 * 1024].into()) + .await + .unwrap(); + } + }; + let receive = async { + for sequence in 0..16 { + let message = reader.next().await.unwrap().unwrap(); + assert!(message.is_binary()); + assert_eq!( + message.into_data().as_ref(), + vec![received + sequence; 256 * 1024] + ); + } + }; + tokio::join!(send, receive); +} + +#[tokio::test] +async fn native_tls_wss_bidirectional_transfer_through_proxy_buffers() { + for outer_tls in [TlsType::Plain, TlsType::Rustls, TlsType::NativeTls] { + // A 4 KiB pipe forces backpressure while both sides send 4 MiB. + let (client, server) = duplex(4096); + let server = async { + let mut stream = DynTcpStream(Box::new(server)); + if !matches!(outer_tls, TlsType::Plain) { + stream = DynTcpStream(Box::new(acceptor().accept(stream).await.unwrap())); + } + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + assert!(headers.len() < 4096); + headers.push(stream.read_u8().await.unwrap()); + } + assert!(headers.starts_with(b"CONNECT target.invalid:443 HTTP/1.1\r\n")); + stream + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await + .unwrap(); + stream.flush().await.unwrap(); + let stream = acceptor().accept(stream).await.unwrap(); + exchange(accept_async(stream).await.unwrap(), 64, 0).await; + }; + let client = async { + let proxy = Proxy::new("https://proxy.invalid:4443", 5_000).unwrap(); + let target = "target.invalid:443".into_target_addr().unwrap(); + let stream = match outer_tls { + TlsType::Plain => { + DynTcpStream(Box::new(proxy.http_connect(client, &target).await.unwrap())) + } + TlsType::Rustls => DynTcpStream(Box::new( + proxy + .https_connect_rustls(client, &target, true) + .await + .unwrap(), + )), + TlsType::NativeTls => DynTcpStream(Box::new( + proxy + .https_connect_nativetls(client, &target, true) + .await + .unwrap(), + )), + }; + let connector = native_tls::TlsConnector::builder() + .danger_accept_invalid_certs(true) + .build() + .unwrap(); + let (ws, _) = client_async_tls_with_config( + "wss://target.invalid/ws/relay", + FlushBeforeRead(stream), + None, + Some(Connector::NativeTls(connector)), + ) + .await + .unwrap(); + exchange(ws, 0, 64).await; + }; + timeout(Duration::from_secs(20), async { + tokio::join!(client, server) + }) + .await + .unwrap(); + } +} + +#[tokio::test] +async fn read_progresses_while_proxy_flush_is_blocked() { + let (client, mut peer) = duplex(64); + let mut client = FlushBeforeRead(DynTcpStream(Box::new(BufStream::new(client)))); + let payload = [0x5a; 4096]; + client.write_all(&payload).await.unwrap(); + peer.write_all(b"pong").await.unwrap(); + + let mut reply = [0; 4]; + timeout(Duration::from_secs(1), client.read_exact(&mut reply)) + .await + .expect("a blocked flush must not prevent receiving queued data") + .unwrap(); + assert_eq!(&reply, b"pong"); + + let mut received = [0; 4096]; + timeout(Duration::from_secs(1), async { + let (flushed, read) = tokio::join!(client.flush(), peer.read_exact(&mut received)); + flushed.unwrap(); + read.unwrap(); + }) + .await + .unwrap(); + assert_eq!(received, payload); +}