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