noq_proto/connection/streams/
state.rs

1use std::{
2    collections::{VecDeque, hash_map},
3    convert::TryFrom,
4    mem,
5};
6
7use rustc_hash::FxHashMap;
8use tracing::{debug, trace};
9
10use super::{
11    PendingStreamsQueue, Recv, Retransmits, Send, SendState, ShouldTransmit, StreamEvent,
12    StreamHalf,
13};
14use crate::{
15    Dir, MAX_STREAM_COUNT, Side, StreamId, TransportError, VarInt,
16    connection::{PacketBuilder, stats::FrameStats},
17    frame::{self, FrameStruct},
18    transport_parameters::TransportParameters,
19};
20
21/// Wrapper around `Recv` that facilitates reusing `Recv` instances
22#[derive(Debug)]
23pub(super) enum StreamRecv {
24    /// A `Recv` that is ready to be opened
25    Free(Box<Recv>),
26    /// A `Recv` that has been opened
27    Open(Box<Recv>),
28}
29
30impl StreamRecv {
31    /// Returns a reference to the inner `Recv` if the stream is open
32    pub(super) fn as_open_recv(&self) -> Option<&Recv> {
33        match self {
34            Self::Open(r) => Some(r),
35            _ => None,
36        }
37    }
38
39    // Returns a mutable reference to the inner `Recv` if the stream is open
40    pub(super) fn as_open_recv_mut(&mut self) -> Option<&mut Recv> {
41        match self {
42            Self::Open(r) => Some(r),
43            _ => None,
44        }
45    }
46
47    // Returns the inner `Recv`
48    pub(super) fn into_inner(self) -> Box<Recv> {
49        match self {
50            Self::Free(r) | Self::Open(r) => r,
51        }
52    }
53
54    // Reinitialize the stream so the inner `Recv` can be reused
55    pub(super) fn free(self, initial_max_data: u64) -> Self {
56        match self {
57            Self::Free(_) => unreachable!("Self::Free on reinit()"),
58            Self::Open(mut recv) => {
59                recv.reinit(initial_max_data);
60                Self::Free(recv)
61            }
62        }
63    }
64}
65
66#[allow(unreachable_pub)] // fuzzing only
67pub struct StreamsState {
68    pub(super) side: Side,
69    // Set of streams that are currently open, or could be immediately opened by the peer
70    pub(super) send: FxHashMap<StreamId, Option<Box<Send>>>,
71    pub(super) recv: FxHashMap<StreamId, Option<StreamRecv>>,
72    pub(super) free_recv: Vec<StreamRecv>,
73    pub(super) next: [u64; 2],
74    /// Maximum number of locally-initiated streams that may be opened over the lifetime of the
75    /// connection so far, per direction
76    pub(super) max: [u64; 2],
77    /// Maximum number of remotely-initiated streams that may be opened over the lifetime of the
78    /// connection so far, per direction
79    pub(super) max_remote: [u64; 2],
80    /// Value of `max_remote` most recently transmitted to the peer in a `MAX_STREAMS` frame
81    sent_max_remote: [u64; 2],
82    /// Number of streams that we've given the peer permission to open and which aren't fully closed
83    pub(super) allocated_remote_count: [u64; 2],
84    /// Size of the desired stream flow control window. May be smaller than `allocated_remote_count`
85    /// due to `set_max_concurrent` calls.
86    max_concurrent_remote_count: [u64; 2],
87    /// Whether `max_concurrent_remote_count` has ever changed
88    flow_control_adjusted: bool,
89    /// Lowest remotely-initiated stream index that haven't actually been opened by the peer
90    pub(super) next_remote: [u64; 2],
91    /// Whether the remote endpoint has opened any streams the application doesn't know about yet,
92    /// per directionality
93    opened: [bool; 2],
94    // Next to report to the application, once opened
95    pub(super) next_reported_remote: [u64; 2],
96    /// Number of outbound streams
97    ///
98    /// This differs from `self.send.len()` in that it does not include streams that the peer is
99    /// permitted to open but which have not yet been opened.
100    pub(super) send_streams: usize,
101    /// Streams with outgoing data queued, sorted by priority
102    pub(super) pending: PendingStreamsQueue,
103
104    events: VecDeque<StreamEvent>,
105    /// Streams blocked on connection-level flow control or stream window space
106    ///
107    /// Streams are only added to this list when a write fails.
108    pub(super) connection_blocked: Vec<StreamId>,
109    /// Connection-level flow control budget dictated by the peer
110    pub(super) max_data: u64,
111    /// The initial receive window
112    receive_window: u64,
113    /// Limit on incoming data, which is transmitted through `MAX_DATA` frames
114    local_max_data: u64,
115    /// The last value of `MAX_DATA` which had been queued for transmission in
116    /// an outgoing `MAX_DATA` frame
117    sent_max_data: VarInt,
118    /// Sum of current offsets of all send streams.
119    pub(super) data_sent: u64,
120    /// Sum of end offsets of all receive streams. Includes gaps, so it's an upper bound.
121    data_recvd: u64,
122    /// Total quantity of unacknowledged outgoing data
123    pub(super) unacked_data: u64,
124    /// Configured upper bound for `unacked_data`.
125    ///
126    /// Note this may be less than `unacked_data` if the user has set a new value.
127    pub(super) send_window: u64,
128    /// Configured upper bound for how much unacked data the peer can send us per stream
129    pub(super) stream_receive_window: u64,
130
131    // Pertinent state from the TransportParameters supplied by the peer
132    initial_max_stream_data_uni: VarInt,
133    initial_max_stream_data_bidi_local: VarInt,
134    initial_max_stream_data_bidi_remote: VarInt,
135
136    /// The shrink to be applied to local_max_data when receive_window is shrunk
137    receive_window_shrink_debt: u64,
138    /// Whether the locally-initiated stream limit has been hit, per direction
139    pub(super) streams_blocked: [bool; 2],
140}
141
142impl StreamsState {
143    #[allow(unreachable_pub)] // fuzzing only
144    pub fn new(
145        side: Side,
146        max_remote_uni: VarInt,
147        max_remote_bi: VarInt,
148        send_window: u64,
149        receive_window: VarInt,
150        stream_receive_window: VarInt,
151    ) -> Self {
152        Self {
153            side,
154            send: FxHashMap::default(),
155            recv: FxHashMap::default(),
156            free_recv: Vec::new(),
157            next: [0, 0],
158            max: [0, 0],
159            max_remote: [max_remote_bi.into(), max_remote_uni.into()],
160            sent_max_remote: [max_remote_bi.into(), max_remote_uni.into()],
161            allocated_remote_count: [max_remote_bi.into(), max_remote_uni.into()],
162            max_concurrent_remote_count: [max_remote_bi.into(), max_remote_uni.into()],
163            flow_control_adjusted: false,
164            next_remote: [0, 0],
165            opened: [false, false],
166            next_reported_remote: [0, 0],
167            send_streams: 0,
168            pending: PendingStreamsQueue::new(),
169            events: VecDeque::new(),
170            connection_blocked: Vec::new(),
171            max_data: 0,
172            receive_window: receive_window.into(),
173            local_max_data: receive_window.into(),
174            sent_max_data: receive_window,
175            data_sent: 0,
176            data_recvd: 0,
177            unacked_data: 0,
178            send_window,
179            stream_receive_window: stream_receive_window.into(),
180            initial_max_stream_data_uni: 0u32.into(),
181            initial_max_stream_data_bidi_local: 0u32.into(),
182            initial_max_stream_data_bidi_remote: 0u32.into(),
183            receive_window_shrink_debt: 0,
184            streams_blocked: [false, false],
185        }
186    }
187
188    pub(crate) fn set_params(&mut self, params: &TransportParameters) {
189        self.initial_max_stream_data_uni = params.initial_max_stream_data_uni;
190        self.initial_max_stream_data_bidi_local = params.initial_max_stream_data_bidi_local;
191        self.initial_max_stream_data_bidi_remote = params.initial_max_stream_data_bidi_remote;
192        self.max[Dir::Bi as usize] = params.initial_max_streams_bidi.into();
193        self.max[Dir::Uni as usize] = params.initial_max_streams_uni.into();
194        self.received_max_data(params.initial_max_data);
195        for (&id, slot) in self.send.iter_mut() {
196            if id.initiator() != self.side
197                && id.dir() == Dir::Bi
198                && let Some(s) = slot.as_mut()
199            {
200                s.max_data = params.initial_max_stream_data_bidi_local.into();
201            }
202        }
203    }
204
205    /// Ensure we have space for at least a full flow control window of remotely-initiated streams
206    /// to be open, and notify the peer if the window has moved
207    fn ensure_remote_streams(&mut self, dir: Dir) {
208        let new_count = self.max_concurrent_remote_count[dir as usize]
209            .saturating_sub(self.allocated_remote_count[dir as usize]);
210        self.allocated_remote_count[dir as usize] += new_count;
211        self.max_remote[dir as usize] += new_count;
212    }
213
214    pub(crate) fn zero_rtt_rejected(&mut self) {
215        // Revert to initial state for outgoing streams
216        for dir in Dir::iter() {
217            for i in 0..self.next[dir as usize] {
218                // We don't bother calling `stream_freed` here because we explicitly reset affected
219                // counters below.
220                let id = StreamId::new(self.side, dir, i);
221                self.send.remove(&id).unwrap();
222                if let Dir::Bi = dir {
223                    self.recv.remove(&id).unwrap();
224                }
225            }
226            self.next[dir as usize] = 0;
227
228            // If 0-RTT was rejected, any flow control frames we sent were lost.
229            if self.flow_control_adjusted {
230                // Conservative approximation of whatever we sent in transport parameters
231                self.sent_max_remote[dir as usize] = 0;
232            }
233        }
234
235        self.pending.clear();
236        self.send_streams = 0;
237        self.data_sent = 0;
238        self.connection_blocked.clear();
239    }
240
241    /// Process incoming stream frame
242    ///
243    /// If successful, returns whether a `MAX_DATA` frame needs to be transmitted
244    pub(crate) fn received(
245        &mut self,
246        frame: frame::Stream,
247        payload_len: usize,
248    ) -> Result<ShouldTransmit, TransportError> {
249        let id = frame.id;
250        self.validate_receive_id(id).inspect_err(|_e| {
251            debug!("received illegal STREAM frame");
252        })?;
253
254        // Create state for this stream if the remote peer created it.
255        let newly_created = self.ensure_remote(id);
256
257        let Some(rs) = self
258            .recv
259            .get_mut(&id)
260            .map(get_or_insert_recv(self.stream_receive_window))
261        else {
262            trace!("dropping frame for closed stream");
263            return Ok(ShouldTransmit(false));
264        };
265
266        if !rs.is_receiving() {
267            trace!("dropping frame for finished stream");
268            return Ok(ShouldTransmit(false));
269        }
270
271        let (new_bytes, closed) =
272            rs.ingest(frame, payload_len, self.data_recvd, self.local_max_data)?;
273        self.data_recvd = self.data_recvd.saturating_add(new_bytes);
274
275        if !rs.stopped {
276            // Newly opened streams are inherently readable; the app discovers them via
277            // `StreamEvent::Opened` and a separate `Readable` would be redundant.
278            if !newly_created {
279                self.events.push_back(StreamEvent::Readable { id });
280            }
281            return Ok(ShouldTransmit(false));
282        }
283
284        // Stopped streams become closed instantly on FIN, so check whether we need to clean up
285        if closed {
286            let rs = self.recv.remove(&id).flatten().unwrap();
287            self.stream_recv_freed(id, rs);
288        }
289
290        // We don't buffer data on stopped streams, so issue flow control credit immediately
291        Ok(self.add_read_credits(new_bytes))
292    }
293
294    /// Process incoming RESET_STREAM frame
295    ///
296    /// If successful, returns whether a `MAX_DATA` frame needs to be transmitted
297    #[allow(unreachable_pub)] // fuzzing only
298    pub fn received_reset(
299        &mut self,
300        frame: frame::ResetStream,
301    ) -> Result<ShouldTransmit, TransportError> {
302        let frame::ResetStream {
303            id,
304            error_code,
305            final_offset,
306        } = frame;
307        self.validate_receive_id(id).inspect_err(|_e| {
308            debug!("received illegal RESET_STREAM frame");
309        })?;
310
311        // Create state for this stream if the remote peer created it.
312        let newly_created = self.ensure_remote(id);
313
314        let Some(rs) = self
315            .recv
316            .get_mut(&id)
317            .map(get_or_insert_recv(self.stream_receive_window))
318        else {
319            trace!("received RESET_STREAM on closed stream");
320            return Ok(ShouldTransmit(false));
321        };
322
323        // State transition
324        if !rs.reset(
325            error_code,
326            final_offset,
327            self.data_recvd,
328            self.local_max_data,
329        )? {
330            // Redundant reset
331            return Ok(ShouldTransmit(false));
332        }
333        let bytes_read = rs.assembler.bytes_read();
334        let stopped = rs.stopped;
335        let end = rs.end;
336        if stopped {
337            // Stopped streams should be disposed immediately on reset
338            let rs = self.recv.remove(&id).flatten().unwrap();
339            self.stream_recv_freed(id, rs);
340        } else if !newly_created {
341            // Newly opened streams are inherently readable; the app discovers them via
342            // `StreamEvent::Opened` and a separate `Readable` would be redundant.
343            self.events.push_back(StreamEvent::Readable { id });
344        }
345
346        // Update connection-level flow control
347        Ok(if bytes_read != final_offset.into_inner() {
348            // bytes_read is always <= end, so this won't underflow.
349            self.data_recvd = self
350                .data_recvd
351                .saturating_add(u64::from(final_offset) - end);
352            self.add_read_credits(u64::from(final_offset) - bytes_read)
353        } else {
354            ShouldTransmit(false)
355        })
356    }
357
358    /// Process incoming `STOP_SENDING` frame
359    #[allow(unreachable_pub)] // fuzzing only
360    pub fn received_stop_sending(&mut self, id: StreamId, error_code: VarInt) {
361        // Create state for this stream if the remote peer created it.
362        self.ensure_remote(id);
363
364        let max_send_data = self.max_send_data(id);
365        let Some(stream) = self
366            .send
367            .get_mut(&id)
368            .map(get_or_insert_send(max_send_data))
369        else {
370            return;
371        };
372
373        if stream.try_stop(error_code) {
374            self.events
375                .push_back(StreamEvent::Stopped { id, error_code });
376        }
377    }
378
379    pub(crate) fn reset_acked(&mut self, id: StreamId) {
380        match self.send.entry(id) {
381            hash_map::Entry::Vacant(_) => {}
382            hash_map::Entry::Occupied(e) => {
383                if let Some(SendState::ResetSent) = e.get().as_ref().map(|s| s.state) {
384                    e.remove_entry();
385                    self.stream_freed(id, StreamHalf::Send);
386                }
387            }
388        }
389    }
390
391    /// Whether any stream data is queued, regardless of control frames
392    pub(crate) fn can_send_stream_data(&self) -> bool {
393        // Reset streams may linger in the pending stream list, but will never produce stream frames
394        self.pending.iter().any(|stream| {
395            self.send
396                .get(&stream.id)
397                .and_then(|s| s.as_ref())
398                .is_some_and(|s| !s.is_reset())
399        })
400    }
401
402    /// Whether MAX_STREAM_DATA frames could be sent for stream `id`
403    pub(crate) fn can_send_flow_control(&self, id: StreamId) -> bool {
404        self.recv
405            .get(&id)
406            .and_then(|s| s.as_ref())
407            .and_then(|s| s.as_open_recv())
408            .is_some_and(|s| s.can_send_flow_control())
409    }
410
411    pub(in crate::connection) fn write_control_frames<'a, 'b>(
412        &mut self,
413        builder: &mut PacketBuilder<'a, 'b>,
414        pending: &mut Retransmits,
415        stats: &mut FrameStats,
416    ) {
417        // RESET_STREAM
418        while builder.frame_space_remaining() > frame::ResetStream::SIZE_BOUND {
419            let Some((id, error_code)) = pending.reset_stream.pop() else {
420                break;
421            };
422            let Some(stream) = self.send.get_mut(&id).and_then(|s| s.as_mut()) else {
423                continue;
424            };
425            let frame = frame::ResetStream {
426                id,
427                error_code,
428                final_offset: VarInt::try_from(stream.offset()).expect("impossibly large offset"),
429            };
430            builder.write_frame(frame, stats);
431        }
432
433        // STOP_SENDING
434        while builder.frame_space_remaining() > frame::StopSending::SIZE_BOUND {
435            let Some(frame) = pending.stop_sending.pop() else {
436                break;
437            };
438            // We may need to transmit STOP_SENDING even for streams whose state we have discarded,
439            // because we are able to discard local state for stopped streams immediately upon
440            // receiving FIN, even if the peer still has arbitrarily large amounts of data to
441            // (re)transmit due to loss or unconventional sending strategy. We could fine-tune this
442            // a little by dropping the frame if we specifically know the stream's been reset by the
443            // peer, but we discard that information as soon as the application consumes it, so it
444            // can't be relied upon regardless.
445            builder.write_frame(frame, stats);
446        }
447
448        // MAX_DATA
449        if pending.max_data && builder.frame_space_remaining() > 9 {
450            pending.max_data = false;
451
452            // `local_max_data` can grow bigger than `VarInt`.
453            // For transmission inside QUIC frames we need to clamp it to the
454            // maximum allowed `VarInt` size.
455            let max = VarInt::try_from(self.local_max_data).unwrap_or(VarInt::MAX);
456
457            if max > self.sent_max_data {
458                // Record that a `MAX_DATA` announcing a certain window was sent. This will
459                // suppress enqueuing further `MAX_DATA` frames unless either the previous
460                // transmission was not acknowledged or the window further increased.
461                self.sent_max_data = max;
462            }
463
464            builder.write_frame(frame::MaxData(max), stats);
465        }
466
467        // MAX_STREAM_DATA
468        while builder.frame_space_remaining() > 17 {
469            let id = match pending.max_stream_data.iter().next() {
470                Some(x) => *x,
471                None => break,
472            };
473            pending.max_stream_data.remove(&id);
474            let Some(rs) = self
475                .recv
476                .get_mut(&id)
477                .and_then(|s| s.as_mut())
478                .and_then(|s| s.as_open_recv_mut())
479            else {
480                continue;
481            };
482            if !rs.can_send_flow_control() {
483                continue;
484            }
485
486            let (max, _) = rs.max_stream_data(self.stream_receive_window);
487            rs.record_sent_max_stream_data(max);
488            builder.write_frame(frame::MaxStreamData { id, offset: max }, stats);
489        }
490
491        // MAX_STREAMS
492        for dir in Dir::iter() {
493            if !pending.max_stream_id[dir as usize] || builder.frame_space_remaining() <= 9 {
494                continue;
495            }
496
497            pending.max_stream_id[dir as usize] = false;
498            self.sent_max_remote[dir as usize] = self.max_remote[dir as usize];
499            let count = self.max_remote[dir as usize];
500            builder.write_frame(frame::MaxStreams { dir, count }, stats);
501        }
502
503        // STREAMS_BLOCKED
504        for dir in Dir::iter() {
505            if self.streams_blocked[dir as usize] {
506                pending.streams_blocked[dir as usize] = true;
507                self.streams_blocked[dir as usize] = false;
508            }
509
510            if !pending.streams_blocked[dir as usize] || builder.frame_space_remaining() <= 9 {
511                continue;
512            }
513
514            pending.streams_blocked[dir as usize] = false;
515            let limit = self.max[dir as usize];
516            builder.write_frame(frame::StreamsBlocked { dir, limit }, stats);
517        }
518    }
519
520    pub(in crate::connection) fn write_stream_frames<'a, 'b>(
521        &mut self,
522        builder: &mut PacketBuilder<'a, 'b>,
523        fair: bool,
524        stats: &mut FrameStats,
525    ) {
526        while builder.frame_space_remaining() > frame::Stream::SIZE_BOUND {
527            // Pop the stream of the highest priority that currently has pending data. If
528            // the stream still has some pending data left after writing, it will be
529            // reinserted, otherwise not
530            let Some(stream) = self.pending.pop() else {
531                break;
532            };
533
534            let id = stream.id;
535
536            let Some(stream) = self.send.get_mut(&id).and_then(|s| s.as_mut()) else {
537                // Stream was reset with pending data and the reset was acknowledged
538                continue;
539            };
540
541            // Reset streams aren't removed from the pending list and still exist while the peer
542            // hasn't acknowledged the reset, but should not generate STREAM frames, so we need to
543            // check for them explicitly.
544            if stream.is_reset() {
545                continue;
546            }
547
548            // Now that we know the `StreamId`, we can better account for how many bytes
549            // are required to encode it.
550            let max_buf_size = builder.frame_space_remaining() - 1 - VarInt::size(id.into());
551            let (offsets, encode_length) = stream.pending.poll_transmit(max_buf_size);
552            let fin = offsets.end == stream.pending.offset()
553                && matches!(stream.state, SendState::DataSent { .. });
554            if fin {
555                stream.fin_pending = false;
556            }
557
558            if stream.is_pending() {
559                // If the stream still has pending data, reinsert it, possibly with an updated priority value
560                // Fairness with other streams is achieved by implementing round-robin scheduling,
561                // so that the other streams will have a chance to write data
562                // before we touch this stream again.
563                if fair {
564                    self.pending.push_pending(id, stream.priority);
565                } else {
566                    self.pending.reinsert_pending(id, stream.priority);
567                }
568            }
569
570            let range = offsets.clone();
571            let meta = frame::StreamMeta { id, offsets, fin };
572            builder.write_frame(meta.encoder(encode_length), stats);
573            stream.pending.get_into(range, builder.buf);
574        }
575    }
576
577    #[cfg(test)]
578    fn write_frames_for_test(&mut self, capacity: usize, fair: bool) -> frame::StreamMetaVec {
579        let buf = &mut Vec::with_capacity(capacity);
580        let mut tbuf = crate::connection::TransmitBuf::new(buf, std::num::NonZeroUsize::MIN, 1_200);
581        tbuf.start_new_datagram_with_size(capacity);
582        let builder = &mut PacketBuilder::simple_data_buf(&mut tbuf);
583        let stats = &mut FrameStats::default();
584        self.write_stream_frames(builder, fair, stats);
585        builder.sent_frames().stream_frames.clone()
586    }
587
588    pub(crate) fn received_ack_of(&mut self, frame: frame::StreamMeta) {
589        let mut entry = match self.send.entry(frame.id) {
590            hash_map::Entry::Vacant(_) => return,
591            hash_map::Entry::Occupied(e) => e,
592        };
593
594        let Some(stream) = entry.get_mut().as_mut() else {
595            // Because we only call this after sending data on this stream,
596            // this closure should be unreachable. If we did somehow screw that up,
597            // then we might hit an underflow below with unpredictable effects down
598            // the line. Best to short-circuit.
599            return;
600        };
601
602        if stream.is_reset() {
603            // We account for outstanding data on reset streams at time of reset
604            return;
605        }
606        let id = frame.id;
607        self.unacked_data -= frame.offsets.end - frame.offsets.start;
608        if !stream.ack(frame) {
609            // The stream is unfinished or may still need retransmits
610            return;
611        }
612
613        entry.remove_entry();
614        self.stream_freed(id, StreamHalf::Send);
615        self.events.push_back(StreamEvent::Finished { id });
616    }
617
618    pub(crate) fn retransmit(&mut self, frame: frame::StreamMeta) {
619        let Some(stream) = self.send.get_mut(&frame.id).and_then(|s| s.as_mut()) else {
620            // Loss of data on a closed stream is a noop
621            return;
622        };
623        if !stream.is_pending() {
624            self.pending.push_pending(frame.id, stream.priority);
625        }
626        stream.fin_pending |= frame.fin;
627        stream.pending.retransmit(frame.offsets);
628    }
629
630    pub(crate) fn retransmit_all_for_0rtt(&mut self) {
631        for dir in Dir::iter() {
632            for index in 0..self.next[dir as usize] {
633                let id = StreamId::new(Side::Client, dir, index);
634                let Some(stream) = self.send.get_mut(&id).and_then(|s| s.as_mut()) else {
635                    continue;
636                };
637                if stream.pending.is_fully_acked() && !stream.fin_pending {
638                    // Stream data can't be acked in 0-RTT, so we must not have sent anything on
639                    // this stream
640                    continue;
641                }
642                if !stream.is_pending() {
643                    self.pending.push_pending(id, stream.priority);
644                }
645                stream.pending.retransmit_all_for_0rtt();
646            }
647        }
648    }
649
650    pub(crate) fn received_max_streams(
651        &mut self,
652        dir: Dir,
653        count: u64,
654    ) -> Result<(), TransportError> {
655        if count > MAX_STREAM_COUNT {
656            return Err(TransportError::FRAME_ENCODING_ERROR(
657                "unrepresentable stream limit",
658            ));
659        }
660
661        let current = &mut self.max[dir as usize];
662        if count > *current {
663            *current = count;
664            self.streams_blocked[dir as usize] = false;
665            self.events.push_back(StreamEvent::Available { dir });
666        }
667
668        Ok(())
669    }
670
671    /// Handle increase to connection-level flow control limit
672    pub(crate) fn received_max_data(&mut self, n: VarInt) {
673        self.max_data = self.max_data.max(n.into());
674    }
675
676    pub(crate) fn received_max_stream_data(
677        &mut self,
678        id: StreamId,
679        offset: u64,
680    ) -> Result<(), TransportError> {
681        if id.initiator() != self.side && id.dir() == Dir::Uni {
682            debug!("got MAX_STREAM_DATA on recv-only {}", id);
683            return Err(TransportError::STREAM_STATE_ERROR(
684                "MAX_STREAM_DATA on recv-only stream",
685            ));
686        }
687
688        // Create state for this stream if the remote peer created it.
689        self.ensure_remote(id);
690
691        let write_limit = self.write_limit();
692        let max_send_data = self.max_send_data(id);
693        if let Some(ss) = self
694            .send
695            .get_mut(&id)
696            .map(get_or_insert_send(max_send_data))
697        {
698            if ss.increase_max_data(offset) {
699                if write_limit > 0 {
700                    self.events.push_back(StreamEvent::Writable { id });
701                } else if !ss.connection_blocked {
702                    // The stream is still blocked on the connection flow control
703                    // window. In order to get unblocked when the window relaxes
704                    // it needs to be in the connection blocked list.
705                    ss.connection_blocked = true;
706                    self.connection_blocked.push(id);
707                }
708            }
709        } else if id.initiator() == self.side && self.is_local_unopened(id) {
710            debug!("got MAX_STREAM_DATA on unopened {}", id);
711            return Err(TransportError::STREAM_STATE_ERROR(
712                "MAX_STREAM_DATA on unopened stream",
713            ));
714        }
715
716        Ok(())
717    }
718
719    /// Returns the maximum amount of data this is allowed to be written on the connection
720    pub(crate) fn write_limit(&self) -> u64 {
721        (self.max_data - self.data_sent)
722            // `send_window` can be set after construction to something *less* than `unacked_data`
723            .min(self.send_window.saturating_sub(self.unacked_data))
724    }
725
726    /// Yield stream events
727    pub(crate) fn poll(&mut self) -> Option<StreamEvent> {
728        if let Some(dir) = Dir::iter().find(|&i| mem::replace(&mut self.opened[i as usize], false))
729        {
730            return Some(StreamEvent::Opened { dir });
731        }
732
733        if self.write_limit() > 0 {
734            while let Some(id) = self.connection_blocked.pop() {
735                let Some(stream) = self.send.get_mut(&id).and_then(|s| s.as_mut()) else {
736                    continue;
737                };
738
739                debug_assert!(stream.connection_blocked);
740                stream.connection_blocked = false;
741
742                // If it's no longer sensible to write to a stream (even to detect an error) then don't
743                // report it.
744                if stream.is_writable() && stream.max_data > stream.offset() {
745                    return Some(StreamEvent::Writable { id });
746                }
747            }
748        }
749
750        self.events.pop_front()
751    }
752
753    /// Queues MAX_STREAM_ID frames in `pending` if needed
754    ///
755    /// Returns whether any frames were queued.
756    pub(crate) fn queue_max_stream_id(&mut self, pending: &mut Retransmits) -> bool {
757        let mut queued = false;
758        for dir in Dir::iter() {
759            let diff = self.max_remote[dir as usize] - self.sent_max_remote[dir as usize];
760            // To reduce traffic, only announce updates if at least 1/8 of the flow control window
761            // has been consumed.
762            if diff > self.max_concurrent_remote_count[dir as usize] / 8 {
763                pending.max_stream_id[dir as usize] = true;
764                queued = true;
765            }
766        }
767        queued
768    }
769
770    /// Check for errors entailed by the peer's use of `id` as a send stream
771    fn validate_receive_id(&mut self, id: StreamId) -> Result<(), TransportError> {
772        if self.side == id.initiator() {
773            match id.dir() {
774                Dir::Uni => {
775                    return Err(TransportError::STREAM_STATE_ERROR(
776                        "illegal operation on send-only stream",
777                    ));
778                }
779                Dir::Bi if id.index() >= self.next[Dir::Bi as usize] => {
780                    return Err(TransportError::STREAM_STATE_ERROR(
781                        "operation on unopened stream",
782                    ));
783                }
784                Dir::Bi => {}
785            };
786        } else {
787            let limit = self.max_remote[id.dir() as usize];
788            if id.index() >= limit {
789                return Err(TransportError::STREAM_LIMIT_ERROR(""));
790            }
791        }
792        Ok(())
793    }
794
795    /// Whether a locally initiated stream has never been open
796    pub(crate) fn is_local_unopened(&self, id: StreamId) -> bool {
797        id.index() >= self.next[id.dir() as usize]
798    }
799
800    pub(crate) fn set_max_concurrent(&mut self, dir: Dir, count: VarInt) {
801        self.flow_control_adjusted = true;
802        self.max_concurrent_remote_count[dir as usize] = count.into();
803        self.ensure_remote_streams(dir);
804    }
805
806    pub(crate) fn max_concurrent(&self, dir: Dir) -> u64 {
807        self.allocated_remote_count[dir as usize]
808    }
809
810    pub(crate) fn set_send_window(&mut self, send_window: u64) {
811        self.send_window = send_window;
812    }
813
814    /// Set the receive_window and returns whether the receive_window has been
815    /// expanded or shrunk: true if expanded, false if shrunk.
816    pub(crate) fn set_receive_window(&mut self, receive_window: VarInt) -> bool {
817        let receive_window = receive_window.into();
818        let mut expanded = false;
819        if receive_window > self.receive_window {
820            self.local_max_data = self
821                .local_max_data
822                .saturating_add(receive_window - self.receive_window);
823            expanded = true;
824        } else {
825            let diff = self.receive_window - receive_window;
826            self.receive_window_shrink_debt = self.receive_window_shrink_debt.saturating_add(diff);
827        }
828        self.receive_window = receive_window;
829        expanded
830    }
831
832    /// Insert `(id, None)` placeholders for a locally-initiated stream into `send` (and `recv`
833    /// for bidi). Called from `Streams::open`; the caller guarantees the id is fresh.
834    pub(super) fn insert_local(&mut self, id: StreamId) {
835        debug_assert_eq!(id.initiator(), self.side);
836        assert!(self.send.insert(id, None).is_none());
837        if id.dir() == Dir::Bi {
838            let recv = self.free_recv.pop();
839            assert!(self.recv.insert(id, recv).is_none());
840        }
841    }
842
843    /// Allocate any new remote streams when a packet arrives for a stream id.
844    /// Any streams above `next_remote` are considered in the default state, avoiding allocations.
845    /// Once we receive a packet for a new remote stream, we advance `next_remote` and allocate actual state.
846    /// If there's a gap, we insert `None` placeholders for the missing streams.
847    /// Returns `true` if `id` was a newly allocated remote stream.
848    fn ensure_remote(&mut self, id: StreamId) -> bool {
849        let dir = id.dir();
850        let dir_idx = dir as usize;
851
852        // If we initiated this stream, nothing to do
853        if id.initiator() == self.side
854            // If this stream is larger than the max allowed, return.
855            // NOTE: STREAM/RESET_STREAM already enforces this, however STOP_SENDING/MAX_STREAM_DATA do not
856            || id.index() >= self.max_remote[dir_idx]
857            // If this stream has already been opened, nothing to do
858            || id.index() < self.next_remote[dir_idx]
859        {
860            return false;
861        }
862
863        // Create all of the streams between the largest opened and this stream.
864        for i in self.next_remote[dir_idx]..=id.index() {
865            let id = StreamId::new(!self.side, dir, i);
866            let recv = self.free_recv.pop();
867            assert!(self.recv.insert(id, recv).is_none());
868            if dir == Dir::Bi {
869                assert!(self.send.insert(id, None).is_none());
870            }
871        }
872        self.next_remote[dir_idx] = id.index() + 1;
873        self.opened[dir_idx] = true;
874        true
875    }
876
877    /// Adds credits to the connection flow control window
878    ///
879    /// Returns whether a `MAX_DATA` frame should be enqueued as soon as possible.
880    /// This will only be the case if the window update would is significant
881    /// enough. As soon as a window update with a `MAX_DATA` frame has been
882    /// queued, the [`Recv::record_sent_max_stream_data`] function should be called to
883    /// suppress sending further updates until the window increases significantly
884    /// again.
885    pub(super) fn add_read_credits(&mut self, credits: u64) -> ShouldTransmit {
886        if credits > self.receive_window_shrink_debt {
887            let net_credits = credits - self.receive_window_shrink_debt;
888            self.local_max_data = self.local_max_data.saturating_add(net_credits);
889            self.receive_window_shrink_debt = 0;
890        } else {
891            self.receive_window_shrink_debt -= credits;
892        }
893
894        if self.local_max_data > VarInt::MAX.into_inner() {
895            return ShouldTransmit(false);
896        }
897
898        // Only announce a window update if it's significant enough
899        // to make it worthwhile sending a MAX_DATA frame.
900        // We use a fraction of the configured connection receive window to make
901        // the decision, to accommodate for connection using bigger windows requiring
902        // less updates.
903        let diff = self.local_max_data - self.sent_max_data.into_inner();
904        ShouldTransmit(diff >= (self.receive_window / 8))
905    }
906
907    /// Update counters for removal of a stream
908    pub(super) fn stream_freed(&mut self, id: StreamId, half: StreamHalf) {
909        if id.initiator() != self.side {
910            let fully_free = id.dir() == Dir::Uni
911                || match half {
912                    StreamHalf::Send => !self.recv.contains_key(&id),
913                    StreamHalf::Recv => !self.send.contains_key(&id),
914                };
915            if fully_free {
916                self.allocated_remote_count[id.dir() as usize] -= 1;
917                self.ensure_remote_streams(id.dir());
918            }
919        }
920        if half == StreamHalf::Send {
921            self.send_streams -= 1;
922        }
923    }
924
925    pub(super) fn stream_recv_freed(&mut self, id: StreamId, recv: StreamRecv) {
926        self.free_recv.push(recv.free(self.stream_receive_window));
927        self.stream_freed(id, StreamHalf::Recv);
928    }
929
930    pub(super) fn max_send_data(&self, id: StreamId) -> VarInt {
931        let remote = self.side != id.initiator();
932        match id.dir() {
933            Dir::Uni => self.initial_max_stream_data_uni,
934            // Remote/local appear reversed here because the transport parameters are named from
935            // the perspective of the peer.
936            Dir::Bi if remote => self.initial_max_stream_data_bidi_local,
937            Dir::Bi => self.initial_max_stream_data_bidi_remote,
938        }
939    }
940}
941
942#[inline]
943pub(super) fn get_or_insert_send(
944    max_data: VarInt,
945) -> impl Fn(&mut Option<Box<Send>>) -> &mut Box<Send> {
946    move |opt| opt.get_or_insert_with(|| Send::new(max_data))
947}
948
949#[inline]
950pub(super) fn get_or_insert_recv(
951    initial_max_data: u64,
952) -> impl FnMut(&mut Option<StreamRecv>) -> &mut Recv {
953    move |opt| {
954        *opt = opt.take().map(|s| match s {
955            StreamRecv::Free(recv) => StreamRecv::Open(recv),
956            s => s,
957        });
958        opt.get_or_insert_with(|| StreamRecv::Open(Recv::new(initial_max_data)))
959            .as_open_recv_mut()
960            .unwrap()
961    }
962}
963
964#[cfg(test)]
965mod tests {
966    use super::*;
967    use crate::{
968        ReadableError, RecvStream, SendStream, TransportErrorCode, WriteError,
969        connection::State as ConnState, connection::Streams,
970    };
971    use bytes::Bytes;
972
973    fn make(side: Side) -> StreamsState {
974        StreamsState::new(
975            side,
976            128u32.into(),
977            128u32.into(),
978            1024 * 1024,
979            (1024 * 1024u32).into(),
980            (1024 * 1024u32).into(),
981        )
982    }
983
984    #[test]
985    fn trivial_flow_control() {
986        let mut client = StreamsState::new(
987            Side::Client,
988            1u32.into(),
989            1u32.into(),
990            1024 * 1024,
991            (1024 * 1024u32).into(),
992            (1024 * 1024u32).into(),
993        );
994        let id = StreamId::new(Side::Server, Dir::Uni, 0);
995        let initial_max = client.local_max_data;
996        const MESSAGE_SIZE: usize = 2048;
997        assert_eq!(
998            client
999                .received(
1000                    frame::Stream {
1001                        id,
1002                        offset: 0,
1003                        fin: true,
1004                        data: Bytes::from_static(&[0; MESSAGE_SIZE]),
1005                    },
1006                    2048
1007                )
1008                .unwrap(),
1009            ShouldTransmit(false)
1010        );
1011        assert_eq!(client.data_recvd, 2048);
1012        assert_eq!(client.local_max_data - initial_max, 0);
1013
1014        let mut pending = Retransmits::default();
1015        let mut recv = RecvStream {
1016            id,
1017            state: &mut client,
1018            pending: &mut pending,
1019        };
1020
1021        let mut chunks = recv.read(true).unwrap();
1022        assert_eq!(
1023            chunks.next(MESSAGE_SIZE).unwrap().unwrap().bytes.len(),
1024            MESSAGE_SIZE
1025        );
1026        assert!(chunks.next(0).unwrap().is_none());
1027        let should_transmit = chunks.finalize();
1028        assert!(should_transmit.0);
1029        assert!(pending.max_stream_id[Dir::Uni as usize]);
1030        assert_eq!(client.local_max_data - initial_max, MESSAGE_SIZE as u64);
1031    }
1032
1033    #[test]
1034    fn reset_flow_control() {
1035        let mut client = make(Side::Client);
1036        let id = StreamId::new(Side::Server, Dir::Uni, 0);
1037        let initial_max = client.local_max_data;
1038        assert_eq!(
1039            client
1040                .received(
1041                    frame::Stream {
1042                        id,
1043                        offset: 0,
1044                        fin: false,
1045                        data: Bytes::from_static(&[0; 2048]),
1046                    },
1047                    2048
1048                )
1049                .unwrap(),
1050            ShouldTransmit(false)
1051        );
1052        assert_eq!(client.data_recvd, 2048);
1053        assert_eq!(client.local_max_data - initial_max, 0);
1054
1055        let mut pending = Retransmits::default();
1056        let mut recv = RecvStream {
1057            id,
1058            state: &mut client,
1059            pending: &mut pending,
1060        };
1061
1062        let mut chunks = recv.read(true).unwrap();
1063        chunks.next(1024).unwrap();
1064        let _ = chunks.finalize();
1065        assert_eq!(client.local_max_data - initial_max, 1024);
1066        assert_eq!(
1067            client
1068                .received_reset(frame::ResetStream {
1069                    id,
1070                    error_code: 0u32.into(),
1071                    final_offset: 4096u32.into(),
1072                })
1073                .unwrap(),
1074            ShouldTransmit(false)
1075        );
1076
1077        assert_eq!(client.data_recvd, 4096);
1078        assert_eq!(client.local_max_data - initial_max, 4096);
1079
1080        // Ensure reading after a reset doesn't issue redundant credit
1081        let mut recv = RecvStream {
1082            id,
1083            state: &mut client,
1084            pending: &mut pending,
1085        };
1086        let mut chunks = recv.read(true).unwrap();
1087        assert_eq!(
1088            chunks.next(1024).unwrap_err(),
1089            crate::ReadError::Reset(0u32.into())
1090        );
1091        let _ = chunks.finalize();
1092        assert_eq!(client.data_recvd, 4096);
1093        assert_eq!(client.local_max_data - initial_max, 4096);
1094    }
1095
1096    #[test]
1097    fn reset_after_empty_frame_flow_control() {
1098        let mut client = make(Side::Client);
1099        let id = StreamId::new(Side::Server, Dir::Uni, 0);
1100        let initial_max = client.local_max_data;
1101        assert_eq!(
1102            client
1103                .received(
1104                    frame::Stream {
1105                        id,
1106                        offset: 4096,
1107                        fin: false,
1108                        data: Bytes::from_static(&[0; 0]),
1109                    },
1110                    0
1111                )
1112                .unwrap(),
1113            ShouldTransmit(false)
1114        );
1115        assert_eq!(client.data_recvd, 4096);
1116        assert_eq!(client.local_max_data - initial_max, 0);
1117        assert_eq!(
1118            client
1119                .received_reset(frame::ResetStream {
1120                    id,
1121                    error_code: 0u32.into(),
1122                    final_offset: 4096u32.into(),
1123                })
1124                .unwrap(),
1125            ShouldTransmit(false)
1126        );
1127        assert_eq!(client.data_recvd, 4096);
1128        assert_eq!(client.local_max_data - initial_max, 4096);
1129    }
1130
1131    #[test]
1132    fn duplicate_reset_flow_control() {
1133        let mut client = make(Side::Client);
1134        let id = StreamId::new(Side::Server, Dir::Uni, 0);
1135        assert_eq!(
1136            client
1137                .received_reset(frame::ResetStream {
1138                    id,
1139                    error_code: 0u32.into(),
1140                    final_offset: 4096u32.into(),
1141                })
1142                .unwrap(),
1143            ShouldTransmit(false)
1144        );
1145        assert_eq!(client.data_recvd, 4096);
1146        assert_eq!(
1147            client
1148                .received_reset(frame::ResetStream {
1149                    id,
1150                    error_code: 0u32.into(),
1151                    final_offset: 4096u32.into(),
1152                })
1153                .unwrap(),
1154            ShouldTransmit(false)
1155        );
1156        assert_eq!(client.data_recvd, 4096);
1157    }
1158
1159    #[test]
1160    fn recv_stopped() {
1161        let mut client = make(Side::Client);
1162        let id = StreamId::new(Side::Server, Dir::Uni, 0);
1163        let initial_max = client.local_max_data;
1164        assert_eq!(
1165            client
1166                .received(
1167                    frame::Stream {
1168                        id,
1169                        offset: 0,
1170                        fin: false,
1171                        data: Bytes::from_static(&[0; 32]),
1172                    },
1173                    32
1174                )
1175                .unwrap(),
1176            ShouldTransmit(false)
1177        );
1178        assert_eq!(client.local_max_data, initial_max);
1179
1180        let mut pending = Retransmits::default();
1181        let mut recv = RecvStream {
1182            id,
1183            state: &mut client,
1184            pending: &mut pending,
1185        };
1186
1187        recv.stop(0u32.into()).unwrap();
1188        assert_eq!(recv.pending.stop_sending.len(), 1);
1189        assert!(!recv.pending.max_data);
1190
1191        assert!(recv.stop(0u32.into()).is_err());
1192        assert_eq!(recv.read(true).err(), Some(ReadableError::ClosedStream));
1193        assert_eq!(recv.read(false).err(), Some(ReadableError::ClosedStream));
1194
1195        assert_eq!(client.local_max_data - initial_max, 32);
1196        assert_eq!(
1197            client
1198                .received(
1199                    frame::Stream {
1200                        id,
1201                        offset: 32,
1202                        fin: true,
1203                        data: Bytes::from_static(&[0; 16]),
1204                    },
1205                    16
1206                )
1207                .unwrap(),
1208            ShouldTransmit(false)
1209        );
1210        assert_eq!(client.local_max_data - initial_max, 48);
1211        assert!(!client.recv.contains_key(&id));
1212    }
1213
1214    #[test]
1215    fn stopped_reset() {
1216        let mut client = make(Side::Client);
1217        let id = StreamId::new(Side::Server, Dir::Uni, 0);
1218        // Server opens stream
1219        assert_eq!(
1220            client
1221                .received(
1222                    frame::Stream {
1223                        id,
1224                        offset: 0,
1225                        fin: false,
1226                        data: Bytes::from_static(&[0; 32])
1227                    },
1228                    32
1229                )
1230                .unwrap(),
1231            ShouldTransmit(false)
1232        );
1233
1234        let mut pending = Retransmits::default();
1235        let mut recv = RecvStream {
1236            id,
1237            state: &mut client,
1238            pending: &mut pending,
1239        };
1240
1241        recv.stop(0u32.into()).unwrap();
1242        assert_eq!(pending.stop_sending.len(), 1);
1243        assert!(!pending.max_data);
1244
1245        // Server complies
1246        let prev_max = client.max_remote[Dir::Uni as usize];
1247        assert_eq!(
1248            client
1249                .received_reset(frame::ResetStream {
1250                    id,
1251                    error_code: 0u32.into(),
1252                    final_offset: 32u32.into(),
1253                })
1254                .unwrap(),
1255            ShouldTransmit(false)
1256        );
1257        assert!(!client.recv.contains_key(&id), "stream state is freed");
1258        assert_eq!(client.max_remote[Dir::Uni as usize], prev_max + 1);
1259    }
1260
1261    #[test]
1262    fn send_stopped() {
1263        let mut server = make(Side::Server);
1264        server.set_params(&TransportParameters {
1265            initial_max_streams_uni: 1u32.into(),
1266            initial_max_data: 42u32.into(),
1267            initial_max_stream_data_uni: 42u32.into(),
1268            ..TransportParameters::default()
1269        });
1270
1271        let (mut pending, state) = (Retransmits::default(), ConnState::established());
1272        let id = Streams {
1273            state: &mut server,
1274            conn_state: &state,
1275        }
1276        .open(Dir::Uni)
1277        .unwrap();
1278
1279        let mut stream = SendStream {
1280            id,
1281            state: &mut server,
1282            pending: &mut pending,
1283            conn_state: &state,
1284        };
1285
1286        let error_code = 0u32.into();
1287        stream.state.received_stop_sending(id, error_code);
1288        assert!(
1289            stream
1290                .state
1291                .events
1292                .contains(&StreamEvent::Stopped { id, error_code })
1293        );
1294        stream.state.events.clear();
1295
1296        assert_eq!(stream.write(&[]), Err(WriteError::Stopped(error_code)));
1297
1298        stream.reset(0u32.into()).unwrap();
1299        assert_eq!(stream.write(&[]), Err(WriteError::ClosedStream));
1300
1301        // A duplicate frame is a no-op
1302        stream.state.received_stop_sending(id, error_code);
1303        assert!(stream.state.events.is_empty());
1304    }
1305
1306    #[test]
1307    fn final_offset_flow_control() {
1308        let mut client = make(Side::Client);
1309        assert_eq!(
1310            client
1311                .received_reset(frame::ResetStream {
1312                    id: StreamId::new(Side::Server, Dir::Uni, 0),
1313                    error_code: 0u32.into(),
1314                    final_offset: VarInt::MAX,
1315                })
1316                .unwrap_err()
1317                .code,
1318            TransportErrorCode::FLOW_CONTROL_ERROR
1319        );
1320    }
1321
1322    #[test]
1323    fn stream_priority() {
1324        let mut server = make(Side::Server);
1325        server.set_params(&TransportParameters {
1326            initial_max_streams_bidi: 3u32.into(),
1327            initial_max_data: 10u32.into(),
1328            initial_max_stream_data_bidi_remote: 10u32.into(),
1329            ..TransportParameters::default()
1330        });
1331
1332        let (mut pending, state) = (Retransmits::default(), ConnState::established());
1333        let mut streams = Streams {
1334            state: &mut server,
1335            conn_state: &state,
1336        };
1337
1338        let id_high = streams.open(Dir::Bi).unwrap();
1339        let id_mid = streams.open(Dir::Bi).unwrap();
1340        let id_low = streams.open(Dir::Bi).unwrap();
1341
1342        let mut mid = SendStream {
1343            id: id_mid,
1344            state: &mut server,
1345            pending: &mut pending,
1346            conn_state: &state,
1347        };
1348        mid.write(b"mid").unwrap();
1349
1350        let mut low = SendStream {
1351            id: id_low,
1352            state: &mut server,
1353            pending: &mut pending,
1354            conn_state: &state,
1355        };
1356        low.set_priority(-1).unwrap();
1357        low.write(b"low").unwrap();
1358
1359        let mut high = SendStream {
1360            id: id_high,
1361            state: &mut server,
1362            pending: &mut pending,
1363            conn_state: &state,
1364        };
1365        high.set_priority(1).unwrap();
1366        high.write(b"high").unwrap();
1367
1368        let meta = server.write_frames_for_test(40, true);
1369        assert_eq!(meta[0].id, id_high);
1370        assert_eq!(meta[1].id, id_mid);
1371        assert_eq!(meta[2].id, id_low);
1372
1373        assert!(!server.can_send_stream_data());
1374        assert_eq!(server.pending.len(), 0);
1375    }
1376
1377    #[test]
1378    fn requeue_stream_priority() {
1379        let mut server = make(Side::Server);
1380        server.set_params(&TransportParameters {
1381            initial_max_streams_bidi: 3u32.into(),
1382            initial_max_data: 1000u32.into(),
1383            initial_max_stream_data_bidi_remote: 1000u32.into(),
1384            ..TransportParameters::default()
1385        });
1386
1387        let (mut pending, state) = (Retransmits::default(), ConnState::established());
1388        let mut streams = Streams {
1389            state: &mut server,
1390            conn_state: &state,
1391        };
1392
1393        let id_high = streams.open(Dir::Bi).unwrap();
1394        let id_mid = streams.open(Dir::Bi).unwrap();
1395
1396        let mut mid = SendStream {
1397            id: id_mid,
1398            state: &mut server,
1399            pending: &mut pending,
1400            conn_state: &state,
1401        };
1402        assert_eq!(mid.write(b"mid").unwrap(), 3);
1403        assert_eq!(server.pending.len(), 1);
1404
1405        let mut high = SendStream {
1406            id: id_high,
1407            state: &mut server,
1408            pending: &mut pending,
1409            conn_state: &state,
1410        };
1411        high.set_priority(1).unwrap();
1412        assert_eq!(high.write(&[0; 200]).unwrap(), 200);
1413        assert_eq!(server.pending.len(), 2);
1414
1415        // Requeue the high priority stream to lowest priority. The initial send
1416        // still uses high priority since it's queued that way. After that it will
1417        // switch to low priority
1418        let mut high = SendStream {
1419            id: id_high,
1420            state: &mut server,
1421            pending: &mut pending,
1422            conn_state: &state,
1423        };
1424        high.set_priority(-1).unwrap();
1425
1426        let meta = server.write_frames_for_test(40, true);
1427        assert_eq!(meta.len(), 1);
1428        assert_eq!(meta[0].id, id_high);
1429
1430        // After requeuing we should end up with 2 priorities - not 3
1431        assert_eq!(server.pending.len(), 2);
1432
1433        // Send the remaining data. The initial mid priority one should go first now
1434        let meta = server.write_frames_for_test(1000 - 40, true);
1435        assert_eq!(meta.len(), 2);
1436        assert_eq!(meta[0].id, id_mid);
1437        assert_eq!(meta[1].id, id_high);
1438
1439        assert!(!server.can_send_stream_data());
1440        assert_eq!(server.pending.len(), 0);
1441    }
1442
1443    #[test]
1444    fn same_stream_priority() {
1445        for fair in [true, false] {
1446            let mut server = make(Side::Server);
1447            server.set_params(&TransportParameters {
1448                initial_max_streams_bidi: 3u32.into(),
1449                initial_max_data: 300u32.into(),
1450                initial_max_stream_data_bidi_remote: 300u32.into(),
1451                ..TransportParameters::default()
1452            });
1453
1454            let (mut pending, state) = (Retransmits::default(), ConnState::established());
1455            let mut streams = Streams {
1456                state: &mut server,
1457                conn_state: &state,
1458            };
1459
1460            // a, b and c all have the same priority
1461            let id_a = streams.open(Dir::Bi).unwrap();
1462            let id_b = streams.open(Dir::Bi).unwrap();
1463            let id_c = streams.open(Dir::Bi).unwrap();
1464
1465            let mut stream_a = SendStream {
1466                id: id_a,
1467                state: &mut server,
1468                pending: &mut pending,
1469                conn_state: &state,
1470            };
1471            stream_a.write(&[b'a'; 100]).unwrap();
1472
1473            let mut stream_b = SendStream {
1474                id: id_b,
1475                state: &mut server,
1476                pending: &mut pending,
1477                conn_state: &state,
1478            };
1479            stream_b.write(&[b'b'; 100]).unwrap();
1480
1481            let mut stream_c = SendStream {
1482                id: id_c,
1483                state: &mut server,
1484                pending: &mut pending,
1485                conn_state: &state,
1486            };
1487            stream_c.write(&[b'c'; 100]).unwrap();
1488
1489            let mut metas = vec![];
1490
1491            // loop until all the streams are written
1492            loop {
1493                let meta = server.write_frames_for_test(40, fair);
1494                if meta.is_empty() {
1495                    break;
1496                }
1497                metas.extend(meta);
1498            }
1499
1500            assert!(!server.can_send_stream_data());
1501            assert_eq!(server.pending.len(), 0);
1502
1503            let stream_ids = metas.iter().map(|m| m.id).collect::<Vec<_>>();
1504            if fair {
1505                // When fairness is enabled, if we run out of buffer space to write out a stream,
1506                // the stream is re-queued after all the streams with the same priority.
1507                assert_eq!(
1508                    stream_ids,
1509                    vec![id_a, id_b, id_c, id_a, id_b, id_c, id_a, id_b, id_c]
1510                );
1511            } else {
1512                // When fairness is disabled the stream is re-queued before all the other streams
1513                // with the same priority.
1514                assert_eq!(
1515                    stream_ids,
1516                    vec![id_a, id_a, id_a, id_b, id_b, id_b, id_c, id_c, id_c]
1517                );
1518            }
1519        }
1520    }
1521
1522    #[test]
1523    fn unfair_priority_bump() {
1524        let mut server = make(Side::Server);
1525        server.set_params(&TransportParameters {
1526            initial_max_streams_bidi: 3u32.into(),
1527            initial_max_data: 300u32.into(),
1528            initial_max_stream_data_bidi_remote: 300u32.into(),
1529            ..TransportParameters::default()
1530        });
1531
1532        let (mut pending, state) = (Retransmits::default(), ConnState::established());
1533        let mut streams = Streams {
1534            state: &mut server,
1535            conn_state: &state,
1536        };
1537
1538        // a, and b have the same priority, c has higher priority
1539        let id_a = streams.open(Dir::Bi).unwrap();
1540        let id_b = streams.open(Dir::Bi).unwrap();
1541        let id_c = streams.open(Dir::Bi).unwrap();
1542
1543        let mut stream_a = SendStream {
1544            id: id_a,
1545            state: &mut server,
1546            pending: &mut pending,
1547            conn_state: &state,
1548        };
1549        stream_a.write(&[b'a'; 100]).unwrap();
1550
1551        let mut stream_b = SendStream {
1552            id: id_b,
1553            state: &mut server,
1554            pending: &mut pending,
1555            conn_state: &state,
1556        };
1557        stream_b.write(&[b'b'; 100]).unwrap();
1558
1559        let mut metas = vec![];
1560
1561        // Write the first chunk of stream_a
1562        let meta = server.write_frames_for_test(40, false);
1563        assert!(!meta.is_empty());
1564        metas.extend(meta);
1565
1566        // Queue stream_c which has higher priority
1567        let mut stream_c = SendStream {
1568            id: id_c,
1569            state: &mut server,
1570            pending: &mut pending,
1571            conn_state: &state,
1572        };
1573        stream_c.set_priority(1).unwrap();
1574        stream_c.write(&[b'b'; 100]).unwrap();
1575
1576        // loop until all the streams are written
1577        loop {
1578            let meta = server.write_frames_for_test(40, false);
1579            if meta.is_empty() {
1580                break;
1581            }
1582            metas.extend(meta);
1583        }
1584
1585        assert!(!server.can_send_stream_data());
1586        assert_eq!(server.pending.len(), 0);
1587
1588        let stream_ids = metas.iter().map(|m| m.id).collect::<Vec<_>>();
1589        assert_eq!(
1590            stream_ids,
1591            // stream_c bumps stream_b but doesn't bump stream_a which had already been partly
1592            // written out
1593            vec![id_a, id_a, id_a, id_c, id_c, id_c, id_b, id_b, id_b]
1594        );
1595    }
1596
1597    #[test]
1598    fn stop_finished() {
1599        let mut client = make(Side::Client);
1600        let id = StreamId::new(Side::Server, Dir::Uni, 0);
1601        // Server finishes stream
1602        let _ = client
1603            .received(
1604                frame::Stream {
1605                    id,
1606                    offset: 0,
1607                    fin: true,
1608                    data: Bytes::from_static(&[0; 32]),
1609                },
1610                32,
1611            )
1612            .unwrap();
1613        let mut pending = Retransmits::default();
1614        let mut stream = RecvStream {
1615            id,
1616            state: &mut client,
1617            pending: &mut pending,
1618        };
1619        stream.stop(0u32.into()).unwrap();
1620        assert!(client.recv.get_mut(&id).is_none(), "stream is freed");
1621    }
1622
1623    // Verify that a stream that's been reset doesn't cause the appearance of pending data
1624    #[test]
1625    fn reset_stream_cannot_send() {
1626        let mut server = make(Side::Server);
1627        server.set_params(&TransportParameters {
1628            initial_max_streams_uni: 1u32.into(),
1629            initial_max_data: 42u32.into(),
1630            initial_max_stream_data_uni: 42u32.into(),
1631            ..TransportParameters::default()
1632        });
1633        let (mut pending, state) = (Retransmits::default(), ConnState::established());
1634        let mut streams = Streams {
1635            state: &mut server,
1636            conn_state: &state,
1637        };
1638
1639        let id = streams.open(Dir::Uni).unwrap();
1640        let mut stream = SendStream {
1641            id,
1642            state: &mut server,
1643            pending: &mut pending,
1644            conn_state: &state,
1645        };
1646        stream.write(b"hello").unwrap();
1647        stream.reset(0u32.into()).unwrap();
1648
1649        assert_eq!(pending.reset_stream, &[(id, 0u32.into())]);
1650        assert!(!server.can_send_stream_data());
1651    }
1652
1653    #[test]
1654    fn stream_limit_fixed() {
1655        let mut client = make(Side::Client);
1656        // Open streams 0-127
1657        assert_eq!(
1658            client.received(
1659                frame::Stream {
1660                    id: StreamId::new(Side::Server, Dir::Uni, 127),
1661                    offset: 0,
1662                    fin: true,
1663                    data: Bytes::from_static(&[]),
1664                },
1665                0
1666            ),
1667            Ok(ShouldTransmit(false))
1668        );
1669        // Try to open stream 128, exceeding limit
1670        assert_eq!(
1671            client
1672                .received(
1673                    frame::Stream {
1674                        id: StreamId::new(Side::Server, Dir::Uni, 128),
1675                        offset: 0,
1676                        fin: true,
1677                        data: Bytes::from_static(&[]),
1678                    },
1679                    0
1680                )
1681                .unwrap_err()
1682                .code,
1683            TransportErrorCode::STREAM_LIMIT_ERROR
1684        );
1685
1686        // Free stream 127
1687        let mut pending = Retransmits::default();
1688        let mut stream = RecvStream {
1689            id: StreamId::new(Side::Server, Dir::Uni, 127),
1690            state: &mut client,
1691            pending: &mut pending,
1692        };
1693        stream.stop(0u32.into()).unwrap();
1694
1695        // Open stream 128
1696        assert_eq!(
1697            client.received(
1698                frame::Stream {
1699                    id: StreamId::new(Side::Server, Dir::Uni, 128),
1700                    offset: 0,
1701                    fin: true,
1702                    data: Bytes::from_static(&[]),
1703                },
1704                0
1705            ),
1706            Ok(ShouldTransmit(false))
1707        );
1708    }
1709
1710    #[test]
1711    fn stream_limit_grows() {
1712        let mut client = make(Side::Client);
1713        // Open streams 0-127
1714        assert_eq!(
1715            client.received(
1716                frame::Stream {
1717                    id: StreamId::new(Side::Server, Dir::Uni, 127),
1718                    offset: 0,
1719                    fin: true,
1720                    data: Bytes::from_static(&[]),
1721                },
1722                0
1723            ),
1724            Ok(ShouldTransmit(false))
1725        );
1726        // Try to open stream 128, exceeding limit
1727        assert_eq!(
1728            client
1729                .received(
1730                    frame::Stream {
1731                        id: StreamId::new(Side::Server, Dir::Uni, 128),
1732                        offset: 0,
1733                        fin: true,
1734                        data: Bytes::from_static(&[]),
1735                    },
1736                    0
1737                )
1738                .unwrap_err()
1739                .code,
1740            TransportErrorCode::STREAM_LIMIT_ERROR
1741        );
1742
1743        // Relax limit by one
1744        client.set_max_concurrent(Dir::Uni, 129u32.into());
1745
1746        // Open stream 128
1747        assert_eq!(
1748            client.received(
1749                frame::Stream {
1750                    id: StreamId::new(Side::Server, Dir::Uni, 128),
1751                    offset: 0,
1752                    fin: true,
1753                    data: Bytes::from_static(&[]),
1754                },
1755                0
1756            ),
1757            Ok(ShouldTransmit(false))
1758        );
1759    }
1760
1761    #[test]
1762    fn stream_limit_shrinks() {
1763        let mut client = make(Side::Client);
1764        // Open streams 0-127
1765        assert_eq!(
1766            client.received(
1767                frame::Stream {
1768                    id: StreamId::new(Side::Server, Dir::Uni, 127),
1769                    offset: 0,
1770                    fin: true,
1771                    data: Bytes::from_static(&[]),
1772                },
1773                0
1774            ),
1775            Ok(ShouldTransmit(false))
1776        );
1777
1778        // Tighten limit by one
1779        client.set_max_concurrent(Dir::Uni, 127u32.into());
1780
1781        // Free stream 127
1782        let mut pending = Retransmits::default();
1783        let mut stream = RecvStream {
1784            id: StreamId::new(Side::Server, Dir::Uni, 127),
1785            state: &mut client,
1786            pending: &mut pending,
1787        };
1788        stream.stop(0u32.into()).unwrap();
1789
1790        // Try to open stream 128, still exceeding limit
1791        assert_eq!(
1792            client
1793                .received(
1794                    frame::Stream {
1795                        id: StreamId::new(Side::Server, Dir::Uni, 128),
1796                        offset: 0,
1797                        fin: true,
1798                        data: Bytes::from_static(&[]),
1799                    },
1800                    0
1801                )
1802                .unwrap_err()
1803                .code,
1804            TransportErrorCode::STREAM_LIMIT_ERROR
1805        );
1806
1807        // Free stream 126
1808        assert_eq!(
1809            client.received_reset(frame::ResetStream {
1810                id: StreamId::new(Side::Server, Dir::Uni, 126),
1811                error_code: 0u32.into(),
1812                final_offset: 0u32.into(),
1813            }),
1814            Ok(ShouldTransmit(false))
1815        );
1816        let mut pending = Retransmits::default();
1817        let mut stream = RecvStream {
1818            id: StreamId::new(Side::Server, Dir::Uni, 126),
1819            state: &mut client,
1820            pending: &mut pending,
1821        };
1822        stream.stop(0u32.into()).unwrap();
1823
1824        // Open stream 128
1825        assert_eq!(
1826            client.received(
1827                frame::Stream {
1828                    id: StreamId::new(Side::Server, Dir::Uni, 128),
1829                    offset: 0,
1830                    fin: true,
1831                    data: Bytes::from_static(&[]),
1832                },
1833                0
1834            ),
1835            Ok(ShouldTransmit(false))
1836        );
1837    }
1838
1839    #[test]
1840    fn remote_stream_capacity() {
1841        let mut client = make(Side::Client);
1842        for _ in 0..2 {
1843            client.set_max_concurrent(Dir::Uni, 200u32.into());
1844            client.set_max_concurrent(Dir::Bi, 201u32.into());
1845            assert_eq!(client.max_remote[Dir::Uni as usize], 200);
1846            assert_eq!(client.max_remote[Dir::Bi as usize], 201);
1847            assert_eq!(client.allocated_remote_count[Dir::Uni as usize], 200);
1848            assert_eq!(client.allocated_remote_count[Dir::Bi as usize], 201);
1849            // Slots are materialized lazily: no remote stream has been touched yet.
1850            assert!(client.recv.is_empty());
1851            assert!(client.send.is_empty());
1852        }
1853    }
1854
1855    #[test]
1856    fn lazy_remote_allocation_starts_empty() {
1857        // `StreamsState::new` must not pre-populate `send`/`recv` with placeholder slots.
1858        let client = StreamsState::new(
1859            Side::Client,
1860            10_000u32.into(),
1861            10_000u32.into(),
1862            1024 * 1024,
1863            (1024 * 1024u32).into(),
1864            (1024 * 1024u32).into(),
1865        );
1866        // No slots allocated until a stream is actually received.
1867        assert!(client.recv.is_empty());
1868        assert!(client.send.is_empty());
1869        assert_eq!(client.recv.capacity(), 0);
1870        assert_eq!(client.send.capacity(), 0);
1871    }
1872
1873    #[test]
1874    fn out_of_order_implicit_open() {
1875        // Receiving idx=5 implicitly opens idx 0..=4. A later frame for idx=3 must be
1876        // processed normally, not dropped as "closed".
1877        const STREAM_5_PAYLOAD: &[u8] = &[0xAA; 8];
1878        const STREAM_3_PAYLOAD: &[u8] = &[0xBB; 4];
1879
1880        let mut client = make(Side::Client);
1881        assert_eq!(
1882            client.received(
1883                frame::Stream {
1884                    id: StreamId::new(Side::Server, Dir::Uni, 5),
1885                    offset: 0,
1886                    fin: true,
1887                    data: Bytes::from_static(STREAM_5_PAYLOAD),
1888                },
1889                STREAM_5_PAYLOAD.len(),
1890            ),
1891            Ok(ShouldTransmit(false))
1892        );
1893        assert_eq!(client.next_remote[Dir::Uni as usize], 6);
1894        assert_eq!(
1895            client.received(
1896                frame::Stream {
1897                    id: StreamId::new(Side::Server, Dir::Uni, 3),
1898                    offset: 0,
1899                    fin: true,
1900                    data: Bytes::from_static(STREAM_3_PAYLOAD),
1901                },
1902                STREAM_3_PAYLOAD.len(),
1903            ),
1904            Ok(ShouldTransmit(false))
1905        );
1906
1907        let id = StreamId::new(Side::Server, Dir::Uni, 3);
1908        let mut pending = Retransmits::default();
1909        let mut recv = RecvStream {
1910            id,
1911            state: &mut client,
1912            pending: &mut pending,
1913        };
1914        let mut chunks = recv.read(true).unwrap();
1915        assert_eq!(
1916            chunks.next(STREAM_3_PAYLOAD.len()).unwrap().unwrap().bytes,
1917            STREAM_3_PAYLOAD
1918        );
1919        let _ = chunks.finalize();
1920    }
1921
1922    #[test]
1923    fn frame_for_closed_stream_is_dropped() {
1924        // After a remote stream is fully freed, a subsequent frame for the same id must be
1925        // dropped — absence from the map unambiguously means "closed" for ids below the
1926        // frontier.
1927        const PAYLOAD: &[u8] = &[0; 4];
1928
1929        let mut client = make(Side::Client);
1930        let id = StreamId::new(Side::Server, Dir::Uni, 0);
1931        assert_eq!(
1932            client.received(
1933                frame::Stream {
1934                    id,
1935                    offset: 0,
1936                    fin: true,
1937                    data: Bytes::from_static(PAYLOAD),
1938                },
1939                PAYLOAD.len(),
1940            ),
1941            Ok(ShouldTransmit(false))
1942        );
1943        // Stop the stream so it's fully freed.
1944        let mut pending = Retransmits::default();
1945        RecvStream {
1946            id,
1947            state: &mut client,
1948            pending: &mut pending,
1949        }
1950        .stop(0u32.into())
1951        .unwrap();
1952        assert!(!client.recv.contains_key(&id));
1953
1954        // A stray retransmit for the freed stream must be dropped without resurrecting state.
1955        assert_eq!(
1956            client.received(
1957                frame::Stream {
1958                    id,
1959                    offset: 0,
1960                    fin: true,
1961                    data: Bytes::from_static(PAYLOAD),
1962                },
1963                PAYLOAD.len(),
1964            ),
1965            Ok(ShouldTransmit(false))
1966        );
1967        assert!(!client.recv.contains_key(&id));
1968    }
1969
1970    #[test]
1971    fn churn_keeps_maps_bounded() {
1972        // Rapidly open + fully close a long sequence of remote streams. The maps must stay
1973        // bounded (only active streams are materialized) even though thousands of ids have
1974        // been used over the connection's lifetime.
1975        const N: u64 = 5_000;
1976
1977        let mut client = make(Side::Client);
1978        for i in 0..N {
1979            let id = StreamId::new(Side::Server, Dir::Uni, i);
1980            assert_eq!(
1981                client.received(
1982                    frame::Stream {
1983                        id,
1984                        offset: 0,
1985                        fin: true,
1986                        data: Bytes::from_static(&[0; 1]),
1987                    },
1988                    1,
1989                ),
1990                Ok(ShouldTransmit(false))
1991            );
1992            let mut pending = Retransmits::default();
1993            let mut recv = RecvStream {
1994                id,
1995                state: &mut client,
1996                pending: &mut pending,
1997            };
1998            let mut chunks = recv.read(true).unwrap();
1999            let _ = chunks.next(1).unwrap();
2000            assert!(chunks.next(1).unwrap().is_none());
2001            let _ = chunks.finalize();
2002        }
2003        // Every stream was fully drained; the map must be empty.
2004        assert_eq!(client.recv.len(), 0);
2005        assert_eq!(client.send.len(), 0);
2006    }
2007
2008    #[test]
2009    fn expand_receive_window() {
2010        let mut server = make(Side::Server);
2011        let new_receive_window = 2 * server.receive_window as u32;
2012        let expanded = server.set_receive_window(new_receive_window.into());
2013        assert!(expanded);
2014        assert_eq!(server.receive_window, new_receive_window as u64);
2015        assert_eq!(server.local_max_data, new_receive_window as u64);
2016        assert_eq!(server.receive_window_shrink_debt, 0);
2017        let prev_local_max_data = server.local_max_data;
2018
2019        // credit, expecting all of them added to local_max_data
2020        let credits = 1024u64;
2021        let should_transmit = server.add_read_credits(credits);
2022        assert_eq!(server.receive_window_shrink_debt, 0);
2023        assert_eq!(server.local_max_data, prev_local_max_data + credits);
2024        assert!(should_transmit.should_transmit());
2025    }
2026
2027    #[test]
2028    fn shrink_receive_window() {
2029        let mut server = make(Side::Server);
2030        let new_receive_window = server.receive_window as u32 / 2;
2031        let prev_local_max_data = server.local_max_data;
2032
2033        // shrink the receive_winbow, local_max_data is not expected to be changed
2034        let shrink_diff = server.receive_window - new_receive_window as u64;
2035        let expanded = server.set_receive_window(new_receive_window.into());
2036        assert!(!expanded);
2037        assert_eq!(server.receive_window, new_receive_window as u64);
2038        assert_eq!(server.local_max_data, prev_local_max_data);
2039        assert_eq!(server.receive_window_shrink_debt, shrink_diff);
2040        let prev_local_max_data = server.local_max_data;
2041
2042        // credit twice, local_max_data does not change as it is absorbed by receive_window_shrink_debt
2043        let credits = 1024u64;
2044        for _ in 0..2 {
2045            let expected_receive_window_shrink_debt = server.receive_window_shrink_debt - credits;
2046            let should_transmit = server.add_read_credits(credits);
2047            assert_eq!(
2048                server.receive_window_shrink_debt,
2049                expected_receive_window_shrink_debt
2050            );
2051            assert_eq!(server.local_max_data, prev_local_max_data);
2052            assert!(!should_transmit.should_transmit());
2053        }
2054
2055        // credit again which exceeds all remaining expected_receive_window_shrink_debt
2056        let credits = 1024 * 512;
2057        let prev_local_max_data = server.local_max_data;
2058        let expected_local_max_data =
2059            server.local_max_data + (credits - server.receive_window_shrink_debt);
2060        let _should_transmit = server.add_read_credits(credits);
2061        assert_eq!(server.receive_window_shrink_debt, 0);
2062        assert_eq!(server.local_max_data, expected_local_max_data);
2063        assert!(server.local_max_data > prev_local_max_data);
2064
2065        // credit again, all should be added to local_max_data
2066        let credits = 1024 * 512;
2067        let expected_local_max_data = server.local_max_data + credits;
2068        let should_transmit = server.add_read_credits(credits);
2069        assert_eq!(server.receive_window_shrink_debt, 0);
2070        assert_eq!(server.local_max_data, expected_local_max_data);
2071        assert!(should_transmit.should_transmit());
2072    }
2073
2074    #[test]
2075    fn expand_send_window() {
2076        let mut server = make(Side::Server);
2077
2078        let initial_send_window = server.send_window;
2079        let larger_send_window = initial_send_window * 2;
2080
2081        // Set `initial_max_data` larger than `send_window` so we're limited by local flow control
2082        server.set_params(&TransportParameters {
2083            initial_max_data: VarInt::MAX,
2084            initial_max_stream_data_uni: VarInt::MAX,
2085            initial_max_streams_uni: VarInt::from_u32(100),
2086            ..TransportParameters::default()
2087        });
2088
2089        assert_eq!(server.write_limit(), initial_send_window);
2090        assert_eq!(server.poll(), None);
2091
2092        let mut retransmits = Retransmits::default();
2093        let conn_state = ConnState::established();
2094
2095        let stream_id = Streams {
2096            state: &mut server,
2097            conn_state: &conn_state,
2098        }
2099        .open(Dir::Uni)
2100        .expect("should be able to open a stream");
2101
2102        let mut stream = SendStream {
2103            id: stream_id,
2104            state: &mut server,
2105            pending: &mut retransmits,
2106            conn_state: &conn_state,
2107        };
2108
2109        // Check that the stream accepts `initial_send_window` bytes
2110        let initial_send_len = initial_send_window as usize;
2111        let data = vec![0xFFu8; initial_send_len];
2112
2113        assert_eq!(stream.write(&data), Ok(initial_send_len));
2114
2115        // Try to write the same data again, observe that it's blocked
2116        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2117
2118        // Check that we get a `Writable` event after increasing the send window
2119        stream.state.set_send_window(larger_send_window);
2120        assert_eq!(
2121            stream.state.poll(),
2122            Some(StreamEvent::Writable { id: stream_id })
2123        );
2124
2125        // Check that the stream accepts the exact same amount of data again
2126        assert_eq!(stream.write(&data), Ok(initial_send_len));
2127        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2128
2129        assert_eq!(stream.state.poll(), None);
2130
2131        // Ack the data
2132        stream.state.received_ack_of(frame::StreamMeta {
2133            id: stream_id,
2134            offsets: 0..larger_send_window,
2135            fin: false,
2136        });
2137
2138        assert_eq!(
2139            stream.state.poll(),
2140            Some(StreamEvent::Writable { id: stream_id })
2141        );
2142
2143        // Check that our full send window is available again
2144        assert_eq!(stream.write(&data), Ok(initial_send_len));
2145        assert_eq!(stream.write(&data), Ok(initial_send_len));
2146        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2147    }
2148
2149    #[test]
2150    fn shrink_send_window() {
2151        let mut server = make(Side::Server);
2152
2153        let initial_send_window = server.send_window;
2154        let smaller_send_window = server.send_window / 2;
2155
2156        // Set `initial_max_data` larger than `send_window` so we're limited by local flow control
2157        server.set_params(&TransportParameters {
2158            initial_max_data: VarInt::MAX,
2159            initial_max_stream_data_uni: VarInt::MAX,
2160            initial_max_streams_uni: VarInt::from_u32(100),
2161            ..TransportParameters::default()
2162        });
2163
2164        assert_eq!(server.write_limit(), initial_send_window);
2165        assert_eq!(server.poll(), None);
2166
2167        let mut retransmits = Retransmits::default();
2168        let conn_state = ConnState::established();
2169
2170        let stream_id = Streams {
2171            state: &mut server,
2172            conn_state: &conn_state,
2173        }
2174        .open(Dir::Uni)
2175        .expect("should be able to open a stream");
2176
2177        let mut stream = SendStream {
2178            id: stream_id,
2179            state: &mut server,
2180            pending: &mut retransmits,
2181            conn_state: &conn_state,
2182        };
2183
2184        let initial_send_len = initial_send_window as usize;
2185
2186        let data = vec![0xFFu8; initial_send_len];
2187
2188        // Assert that the full send window is accepted
2189        assert_eq!(stream.write(&data), Ok(initial_send_len));
2190        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2191
2192        assert_eq!(stream.state.write_limit(), 0);
2193        assert_eq!(stream.state.poll(), None);
2194
2195        // Shrink our send window, assert that it's still not writable
2196        stream.state.set_send_window(smaller_send_window);
2197        assert_eq!(stream.state.write_limit(), 0);
2198        assert_eq!(stream.state.poll(), None);
2199
2200        // Assert that data is still not accepted
2201        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2202
2203        // Ack some data, assert that writes are still not accepted due to outstanding sends
2204        stream.state.received_ack_of(frame::StreamMeta {
2205            id: stream_id,
2206            offsets: 0..smaller_send_window,
2207            fin: false,
2208        });
2209
2210        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2211
2212        // Ack the rest of the data
2213        stream.state.received_ack_of(frame::StreamMeta {
2214            id: stream_id,
2215            offsets: smaller_send_window..initial_send_window,
2216            fin: false,
2217        });
2218
2219        // This should generate a `Writable` event
2220        assert_eq!(
2221            stream.state.poll(),
2222            Some(StreamEvent::Writable { id: stream_id })
2223        );
2224        assert_eq!(stream.state.write_limit(), smaller_send_window);
2225
2226        // Assert that only `smaller_send_window` bytes are accepted
2227        assert_eq!(stream.write(&data), Ok(smaller_send_window as usize));
2228        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2229    }
2230}