diff --git a/quinn/src/recv_stream.rs b/quinn/src/recv_stream.rs index 203e84abe..3f581b6ba 100644 --- a/quinn/src/recv_stream.rs +++ b/quinn/src/recv_stream.rs @@ -280,6 +280,9 @@ impl RecvStream { conn.inner.recv_stream(self.stream).stop(error_code)?; conn.wake(); self.all_data_read = true; + // Clean up shared state that might be left over from a cancalled read + // operation, so `drop` doesn't have to + conn.blocked_readers.remove(&self.stream); Ok(()) } diff --git a/quinn/src/tests.rs b/quinn/src/tests.rs index 8ff7aabd9..ffe4507f2 100755 --- a/quinn/src/tests.rs +++ b/quinn/src/tests.rs @@ -10,6 +10,7 @@ use std::{ future::Future, io, net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, UdpSocket}, + pin::pin, str, sync::{ Arc, @@ -27,8 +28,11 @@ use rustls::{ RootCertStore, pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer}, }; -use tokio::runtime::{Builder, Runtime}; use tokio::time::{sleep, timeout}; +use tokio::{ + join, + runtime::{Builder, Runtime}, +}; use tracing::{error_span, info}; use tracing_futures::Instrument as _; use tracing_subscriber::EnvFilter; @@ -1007,6 +1011,52 @@ async fn stream_drop_removes_blocked_reader() { } } +/// Test that dropping a `RecvStream` after cancelling a read and then +/// explicitly `stop`ing it doesn't panic. +#[tokio::test] +async fn recv_stream_cancel_stop_drop() { + let _guard = subscribe(); + let factory = EndpointFactory::new(); + let server = { + let _guard = error_span!("server").entered(); + factory.endpoint() + }; + let server_addr = server.local_addr().unwrap(); + + let client = { + let _guard = error_span!("client").entered(); + factory.endpoint() + }; + let recv_dropped = tokio::sync::SetOnce::new(); + join!( + async { + let conn = server.accept().await.unwrap().await.unwrap(); + let mut recv = conn.accept_uni().await.unwrap(); + // Create a future to read from the stream, poll it once, then immediately drop it + { + let fut = pin!(recv.read_to_end(usize::MAX)); + let mut cx = Context::from_waker(Waker::noop()); + assert!(fut.poll(&mut cx).is_pending()); + } + recv_dropped.set(()).unwrap(); + recv.stop(0u32.into()).unwrap(); + }, + async { + let conn = client + .connect(server_addr, "localhost") + .unwrap() + .await + .unwrap(); + let mut send = conn.open_uni().await.unwrap(); + _ = send.write_all(b"hello").await; + // Don't drop (finish) the send stream until the read has been + // cancelled by the server, ensuring that read_to_end can't complete + // immediately. + recv_dropped.wait().await; + }, + ); +} + #[derive(Default)] struct WakeCounter { wakes: AtomicUsize,