diff --git a/quinn-proto/src/connection/streams/mod.rs b/quinn-proto/src/connection/streams/mod.rs index 635b5a3f3..49f216d20 100644 --- a/quinn-proto/src/connection/streams/mod.rs +++ b/quinn-proto/src/connection/streams/mod.rs @@ -1,6 +1,6 @@ use std::{ cell::RefCell, - collections::{BinaryHeap, VecDeque}, + collections::{hash_map, BinaryHeap, VecDeque}, }; use bytes::Bytes; @@ -114,10 +114,11 @@ impl<'a> RecvStream<'a> { /// Discards unread data and notifies the peer to stop transmitting. Once stopped, further /// attempts to operate on a stream will yield `UnknownStream` errors. pub fn stop(&mut self, error_code: VarInt) -> Result<(), UnknownStream> { - let stream = match self.state.recv.get_mut(&self.id) { - Some(s) => s, - None => return Err(UnknownStream { _private: () }), + let mut entry = match self.state.recv.entry(self.id) { + hash_map::Entry::Occupied(s) => s, + hash_map::Entry::Vacant(_) => return Err(UnknownStream { _private: () }), }; + let stream = entry.get_mut(); let (read_credits, stop_sending) = stream.stop()?; if stop_sending.should_transmit() { @@ -127,6 +128,14 @@ impl<'a> RecvStream<'a> { }); } + // We need to keep stopped streams around until they're finished or reset so we can update + // connection-level flow control to account for discarded data. Otherwise, we can discard + // state immediately. + if !stream.receiving_unknown_size() { + entry.remove(); + self.state.stream_freed(self.id, StreamHalf::Recv); + } + if self.state.add_read_credits(read_credits).should_transmit() { self.pending.max_data = true; } diff --git a/quinn-proto/src/connection/streams/state.rs b/quinn-proto/src/connection/streams/state.rs index bb3cc1559..fd9681b1c 100644 --- a/quinn-proto/src/connection/streams/state.rs +++ b/quinn-proto/src/connection/streams/state.rs @@ -1089,4 +1089,30 @@ mod tests { assert_eq!(meta[1].id, id_mid); assert_eq!(meta[2].id, id_low); } + + #[test] + fn stop_finished() { + let mut client = make(Side::Client); + let id = StreamId::new(Side::Server, Dir::Uni, 0); + // Server finishes stream + let _ = client + .received( + frame::Stream { + id, + offset: 0, + fin: true, + data: Bytes::from_static(&[0; 32]), + }, + 32, + ) + .unwrap(); + let mut pending = Retransmits::default(); + let mut stream = RecvStream { + id, + state: &mut client, + pending: &mut pending, + }; + stream.stop(0u32.into()).unwrap(); + assert!(client.recv.get_mut(&id).is_none(), "stream is freed"); + } }