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        assert!(recv.is_ordered().is_err());
1198
1199        assert_eq!(client.local_max_data - initial_max, 32);
1200        assert_eq!(
1201            client
1202                .received(
1203                    frame::Stream {
1204                        id,
1205                        offset: 32,
1206                        fin: true,
1207                        data: Bytes::from_static(&[0; 16]),
1208                    },
1209                    16
1210                )
1211                .unwrap(),
1212            ShouldTransmit(false)
1213        );
1214        assert_eq!(client.local_max_data - initial_max, 48);
1215        assert!(!client.recv.contains_key(&id));
1216    }
1217
1218    #[test]
1219    fn recv_stream_ordering_mode() {
1220        let mut client = make(Side::Client);
1221        let id = StreamId::new(Side::Server, Dir::Uni, 0);
1222        let _ = client
1223            .received(
1224                frame::Stream {
1225                    id,
1226                    offset: 0,
1227                    fin: false,
1228                    data: Bytes::from_static(b"hello"),
1229                },
1230                5,
1231            )
1232            .unwrap();
1233
1234        let mut pending = Retransmits::default();
1235        let mut recv = RecvStream {
1236            id,
1237            state: &mut client,
1238            pending: &mut pending,
1239        };
1240        assert_eq!(recv.is_ordered(), Ok(true));
1241        let _ = recv.read(false).unwrap().finalize();
1242        assert_eq!(recv.is_ordered(), Ok(false));
1243    }
1244
1245    #[test]
1246    fn stopped_reset() {
1247        let mut client = make(Side::Client);
1248        let id = StreamId::new(Side::Server, Dir::Uni, 0);
1249        // Server opens stream
1250        assert_eq!(
1251            client
1252                .received(
1253                    frame::Stream {
1254                        id,
1255                        offset: 0,
1256                        fin: false,
1257                        data: Bytes::from_static(&[0; 32])
1258                    },
1259                    32
1260                )
1261                .unwrap(),
1262            ShouldTransmit(false)
1263        );
1264
1265        let mut pending = Retransmits::default();
1266        let mut recv = RecvStream {
1267            id,
1268            state: &mut client,
1269            pending: &mut pending,
1270        };
1271
1272        recv.stop(0u32.into()).unwrap();
1273        assert_eq!(pending.stop_sending.len(), 1);
1274        assert!(!pending.max_data);
1275
1276        // Server complies
1277        let prev_max = client.max_remote[Dir::Uni as usize];
1278        assert_eq!(
1279            client
1280                .received_reset(frame::ResetStream {
1281                    id,
1282                    error_code: 0u32.into(),
1283                    final_offset: 32u32.into(),
1284                })
1285                .unwrap(),
1286            ShouldTransmit(false)
1287        );
1288        assert!(!client.recv.contains_key(&id), "stream state is freed");
1289        assert_eq!(client.max_remote[Dir::Uni as usize], prev_max + 1);
1290    }
1291
1292    #[test]
1293    fn send_stopped() {
1294        let mut server = make(Side::Server);
1295        server.set_params(&TransportParameters {
1296            initial_max_streams_uni: 1u32.into(),
1297            initial_max_data: 42u32.into(),
1298            initial_max_stream_data_uni: 42u32.into(),
1299            ..TransportParameters::default()
1300        });
1301
1302        let (mut pending, state) = (Retransmits::default(), ConnState::established());
1303        let id = Streams {
1304            state: &mut server,
1305            conn_state: &state,
1306        }
1307        .open(Dir::Uni)
1308        .unwrap();
1309
1310        let mut stream = SendStream {
1311            id,
1312            state: &mut server,
1313            pending: &mut pending,
1314            conn_state: &state,
1315        };
1316
1317        let error_code = 0u32.into();
1318        stream.state.received_stop_sending(id, error_code);
1319        assert!(
1320            stream
1321                .state
1322                .events
1323                .contains(&StreamEvent::Stopped { id, error_code })
1324        );
1325        stream.state.events.clear();
1326
1327        assert_eq!(stream.write(&[]), Err(WriteError::Stopped(error_code)));
1328
1329        stream.reset(0u32.into()).unwrap();
1330        assert_eq!(stream.write(&[]), Err(WriteError::ClosedStream));
1331
1332        // A duplicate frame is a no-op
1333        stream.state.received_stop_sending(id, error_code);
1334        assert!(stream.state.events.is_empty());
1335    }
1336
1337    #[test]
1338    fn final_offset_flow_control() {
1339        let mut client = make(Side::Client);
1340        assert_eq!(
1341            client
1342                .received_reset(frame::ResetStream {
1343                    id: StreamId::new(Side::Server, Dir::Uni, 0),
1344                    error_code: 0u32.into(),
1345                    final_offset: VarInt::MAX,
1346                })
1347                .unwrap_err()
1348                .code,
1349            TransportErrorCode::FLOW_CONTROL_ERROR
1350        );
1351    }
1352
1353    #[test]
1354    fn stream_priority() {
1355        let mut server = make(Side::Server);
1356        server.set_params(&TransportParameters {
1357            initial_max_streams_bidi: 3u32.into(),
1358            initial_max_data: 10u32.into(),
1359            initial_max_stream_data_bidi_remote: 10u32.into(),
1360            ..TransportParameters::default()
1361        });
1362
1363        let (mut pending, state) = (Retransmits::default(), ConnState::established());
1364        let mut streams = Streams {
1365            state: &mut server,
1366            conn_state: &state,
1367        };
1368
1369        let id_high = streams.open(Dir::Bi).unwrap();
1370        let id_mid = streams.open(Dir::Bi).unwrap();
1371        let id_low = streams.open(Dir::Bi).unwrap();
1372
1373        let mut mid = SendStream {
1374            id: id_mid,
1375            state: &mut server,
1376            pending: &mut pending,
1377            conn_state: &state,
1378        };
1379        mid.write(b"mid").unwrap();
1380
1381        let mut low = SendStream {
1382            id: id_low,
1383            state: &mut server,
1384            pending: &mut pending,
1385            conn_state: &state,
1386        };
1387        low.set_priority(-1).unwrap();
1388        low.write(b"low").unwrap();
1389
1390        let mut high = SendStream {
1391            id: id_high,
1392            state: &mut server,
1393            pending: &mut pending,
1394            conn_state: &state,
1395        };
1396        high.set_priority(1).unwrap();
1397        high.write(b"high").unwrap();
1398
1399        let meta = server.write_frames_for_test(40, true);
1400        assert_eq!(meta[0].id, id_high);
1401        assert_eq!(meta[1].id, id_mid);
1402        assert_eq!(meta[2].id, id_low);
1403
1404        assert!(!server.can_send_stream_data());
1405        assert_eq!(server.pending.len(), 0);
1406    }
1407
1408    #[test]
1409    fn requeue_stream_priority() {
1410        let mut server = make(Side::Server);
1411        server.set_params(&TransportParameters {
1412            initial_max_streams_bidi: 3u32.into(),
1413            initial_max_data: 1000u32.into(),
1414            initial_max_stream_data_bidi_remote: 1000u32.into(),
1415            ..TransportParameters::default()
1416        });
1417
1418        let (mut pending, state) = (Retransmits::default(), ConnState::established());
1419        let mut streams = Streams {
1420            state: &mut server,
1421            conn_state: &state,
1422        };
1423
1424        let id_high = streams.open(Dir::Bi).unwrap();
1425        let id_mid = streams.open(Dir::Bi).unwrap();
1426
1427        let mut mid = SendStream {
1428            id: id_mid,
1429            state: &mut server,
1430            pending: &mut pending,
1431            conn_state: &state,
1432        };
1433        assert_eq!(mid.write(b"mid").unwrap(), 3);
1434        assert_eq!(server.pending.len(), 1);
1435
1436        let mut high = SendStream {
1437            id: id_high,
1438            state: &mut server,
1439            pending: &mut pending,
1440            conn_state: &state,
1441        };
1442        high.set_priority(1).unwrap();
1443        assert_eq!(high.write(&[0; 200]).unwrap(), 200);
1444        assert_eq!(server.pending.len(), 2);
1445
1446        // Requeue the high priority stream to lowest priority. The initial send
1447        // still uses high priority since it's queued that way. After that it will
1448        // switch to low priority
1449        let mut high = SendStream {
1450            id: id_high,
1451            state: &mut server,
1452            pending: &mut pending,
1453            conn_state: &state,
1454        };
1455        high.set_priority(-1).unwrap();
1456
1457        let meta = server.write_frames_for_test(40, true);
1458        assert_eq!(meta.len(), 1);
1459        assert_eq!(meta[0].id, id_high);
1460
1461        // After requeuing we should end up with 2 priorities - not 3
1462        assert_eq!(server.pending.len(), 2);
1463
1464        // Send the remaining data. The initial mid priority one should go first now
1465        let meta = server.write_frames_for_test(1000 - 40, true);
1466        assert_eq!(meta.len(), 2);
1467        assert_eq!(meta[0].id, id_mid);
1468        assert_eq!(meta[1].id, id_high);
1469
1470        assert!(!server.can_send_stream_data());
1471        assert_eq!(server.pending.len(), 0);
1472    }
1473
1474    #[test]
1475    fn same_stream_priority() {
1476        for fair in [true, false] {
1477            let mut server = make(Side::Server);
1478            server.set_params(&TransportParameters {
1479                initial_max_streams_bidi: 3u32.into(),
1480                initial_max_data: 300u32.into(),
1481                initial_max_stream_data_bidi_remote: 300u32.into(),
1482                ..TransportParameters::default()
1483            });
1484
1485            let (mut pending, state) = (Retransmits::default(), ConnState::established());
1486            let mut streams = Streams {
1487                state: &mut server,
1488                conn_state: &state,
1489            };
1490
1491            // a, b and c all have the same priority
1492            let id_a = streams.open(Dir::Bi).unwrap();
1493            let id_b = streams.open(Dir::Bi).unwrap();
1494            let id_c = streams.open(Dir::Bi).unwrap();
1495
1496            let mut stream_a = SendStream {
1497                id: id_a,
1498                state: &mut server,
1499                pending: &mut pending,
1500                conn_state: &state,
1501            };
1502            stream_a.write(&[b'a'; 100]).unwrap();
1503
1504            let mut stream_b = SendStream {
1505                id: id_b,
1506                state: &mut server,
1507                pending: &mut pending,
1508                conn_state: &state,
1509            };
1510            stream_b.write(&[b'b'; 100]).unwrap();
1511
1512            let mut stream_c = SendStream {
1513                id: id_c,
1514                state: &mut server,
1515                pending: &mut pending,
1516                conn_state: &state,
1517            };
1518            stream_c.write(&[b'c'; 100]).unwrap();
1519
1520            let mut metas = vec![];
1521
1522            // loop until all the streams are written
1523            loop {
1524                let meta = server.write_frames_for_test(40, fair);
1525                if meta.is_empty() {
1526                    break;
1527                }
1528                metas.extend(meta);
1529            }
1530
1531            assert!(!server.can_send_stream_data());
1532            assert_eq!(server.pending.len(), 0);
1533
1534            let stream_ids = metas.iter().map(|m| m.id).collect::<Vec<_>>();
1535            if fair {
1536                // When fairness is enabled, if we run out of buffer space to write out a stream,
1537                // the stream is re-queued after all the streams with the same priority.
1538                assert_eq!(
1539                    stream_ids,
1540                    vec![id_a, id_b, id_c, id_a, id_b, id_c, id_a, id_b, id_c]
1541                );
1542            } else {
1543                // When fairness is disabled the stream is re-queued before all the other streams
1544                // with the same priority.
1545                assert_eq!(
1546                    stream_ids,
1547                    vec![id_a, id_a, id_a, id_b, id_b, id_b, id_c, id_c, id_c]
1548                );
1549            }
1550        }
1551    }
1552
1553    #[test]
1554    fn unfair_priority_bump() {
1555        let mut server = make(Side::Server);
1556        server.set_params(&TransportParameters {
1557            initial_max_streams_bidi: 3u32.into(),
1558            initial_max_data: 300u32.into(),
1559            initial_max_stream_data_bidi_remote: 300u32.into(),
1560            ..TransportParameters::default()
1561        });
1562
1563        let (mut pending, state) = (Retransmits::default(), ConnState::established());
1564        let mut streams = Streams {
1565            state: &mut server,
1566            conn_state: &state,
1567        };
1568
1569        // a, and b have the same priority, c has higher priority
1570        let id_a = streams.open(Dir::Bi).unwrap();
1571        let id_b = streams.open(Dir::Bi).unwrap();
1572        let id_c = streams.open(Dir::Bi).unwrap();
1573
1574        let mut stream_a = SendStream {
1575            id: id_a,
1576            state: &mut server,
1577            pending: &mut pending,
1578            conn_state: &state,
1579        };
1580        stream_a.write(&[b'a'; 100]).unwrap();
1581
1582        let mut stream_b = SendStream {
1583            id: id_b,
1584            state: &mut server,
1585            pending: &mut pending,
1586            conn_state: &state,
1587        };
1588        stream_b.write(&[b'b'; 100]).unwrap();
1589
1590        let mut metas = vec![];
1591
1592        // Write the first chunk of stream_a
1593        let meta = server.write_frames_for_test(40, false);
1594        assert!(!meta.is_empty());
1595        metas.extend(meta);
1596
1597        // Queue stream_c which has higher priority
1598        let mut stream_c = SendStream {
1599            id: id_c,
1600            state: &mut server,
1601            pending: &mut pending,
1602            conn_state: &state,
1603        };
1604        stream_c.set_priority(1).unwrap();
1605        stream_c.write(&[b'b'; 100]).unwrap();
1606
1607        // loop until all the streams are written
1608        loop {
1609            let meta = server.write_frames_for_test(40, false);
1610            if meta.is_empty() {
1611                break;
1612            }
1613            metas.extend(meta);
1614        }
1615
1616        assert!(!server.can_send_stream_data());
1617        assert_eq!(server.pending.len(), 0);
1618
1619        let stream_ids = metas.iter().map(|m| m.id).collect::<Vec<_>>();
1620        assert_eq!(
1621            stream_ids,
1622            // stream_c bumps stream_b but doesn't bump stream_a which had already been partly
1623            // written out
1624            vec![id_a, id_a, id_a, id_c, id_c, id_c, id_b, id_b, id_b]
1625        );
1626    }
1627
1628    #[test]
1629    fn stop_finished() {
1630        let mut client = make(Side::Client);
1631        let id = StreamId::new(Side::Server, Dir::Uni, 0);
1632        // Server finishes stream
1633        let _ = client
1634            .received(
1635                frame::Stream {
1636                    id,
1637                    offset: 0,
1638                    fin: true,
1639                    data: Bytes::from_static(&[0; 32]),
1640                },
1641                32,
1642            )
1643            .unwrap();
1644        let mut pending = Retransmits::default();
1645        let mut stream = RecvStream {
1646            id,
1647            state: &mut client,
1648            pending: &mut pending,
1649        };
1650        stream.stop(0u32.into()).unwrap();
1651        assert!(client.recv.get_mut(&id).is_none(), "stream is freed");
1652    }
1653
1654    // Verify that a stream that's been reset doesn't cause the appearance of pending data
1655    #[test]
1656    fn reset_stream_cannot_send() {
1657        let mut server = make(Side::Server);
1658        server.set_params(&TransportParameters {
1659            initial_max_streams_uni: 1u32.into(),
1660            initial_max_data: 42u32.into(),
1661            initial_max_stream_data_uni: 42u32.into(),
1662            ..TransportParameters::default()
1663        });
1664        let (mut pending, state) = (Retransmits::default(), ConnState::established());
1665        let mut streams = Streams {
1666            state: &mut server,
1667            conn_state: &state,
1668        };
1669
1670        let id = streams.open(Dir::Uni).unwrap();
1671        let mut stream = SendStream {
1672            id,
1673            state: &mut server,
1674            pending: &mut pending,
1675            conn_state: &state,
1676        };
1677        stream.write(b"hello").unwrap();
1678        stream.reset(0u32.into()).unwrap();
1679
1680        assert_eq!(pending.reset_stream, &[(id, 0u32.into())]);
1681        assert!(!server.can_send_stream_data());
1682    }
1683
1684    #[test]
1685    fn stream_limit_fixed() {
1686        let mut client = make(Side::Client);
1687        // Open streams 0-127
1688        assert_eq!(
1689            client.received(
1690                frame::Stream {
1691                    id: StreamId::new(Side::Server, Dir::Uni, 127),
1692                    offset: 0,
1693                    fin: true,
1694                    data: Bytes::from_static(&[]),
1695                },
1696                0
1697            ),
1698            Ok(ShouldTransmit(false))
1699        );
1700        // Try to open stream 128, exceeding limit
1701        assert_eq!(
1702            client
1703                .received(
1704                    frame::Stream {
1705                        id: StreamId::new(Side::Server, Dir::Uni, 128),
1706                        offset: 0,
1707                        fin: true,
1708                        data: Bytes::from_static(&[]),
1709                    },
1710                    0
1711                )
1712                .unwrap_err()
1713                .code,
1714            TransportErrorCode::STREAM_LIMIT_ERROR
1715        );
1716
1717        // Free stream 127
1718        let mut pending = Retransmits::default();
1719        let mut stream = RecvStream {
1720            id: StreamId::new(Side::Server, Dir::Uni, 127),
1721            state: &mut client,
1722            pending: &mut pending,
1723        };
1724        stream.stop(0u32.into()).unwrap();
1725
1726        // Open stream 128
1727        assert_eq!(
1728            client.received(
1729                frame::Stream {
1730                    id: StreamId::new(Side::Server, Dir::Uni, 128),
1731                    offset: 0,
1732                    fin: true,
1733                    data: Bytes::from_static(&[]),
1734                },
1735                0
1736            ),
1737            Ok(ShouldTransmit(false))
1738        );
1739    }
1740
1741    #[test]
1742    fn stream_limit_grows() {
1743        let mut client = make(Side::Client);
1744        // Open streams 0-127
1745        assert_eq!(
1746            client.received(
1747                frame::Stream {
1748                    id: StreamId::new(Side::Server, Dir::Uni, 127),
1749                    offset: 0,
1750                    fin: true,
1751                    data: Bytes::from_static(&[]),
1752                },
1753                0
1754            ),
1755            Ok(ShouldTransmit(false))
1756        );
1757        // Try to open stream 128, exceeding limit
1758        assert_eq!(
1759            client
1760                .received(
1761                    frame::Stream {
1762                        id: StreamId::new(Side::Server, Dir::Uni, 128),
1763                        offset: 0,
1764                        fin: true,
1765                        data: Bytes::from_static(&[]),
1766                    },
1767                    0
1768                )
1769                .unwrap_err()
1770                .code,
1771            TransportErrorCode::STREAM_LIMIT_ERROR
1772        );
1773
1774        // Relax limit by one
1775        client.set_max_concurrent(Dir::Uni, 129u32.into());
1776
1777        // Open stream 128
1778        assert_eq!(
1779            client.received(
1780                frame::Stream {
1781                    id: StreamId::new(Side::Server, Dir::Uni, 128),
1782                    offset: 0,
1783                    fin: true,
1784                    data: Bytes::from_static(&[]),
1785                },
1786                0
1787            ),
1788            Ok(ShouldTransmit(false))
1789        );
1790    }
1791
1792    #[test]
1793    fn stream_limit_shrinks() {
1794        let mut client = make(Side::Client);
1795        // Open streams 0-127
1796        assert_eq!(
1797            client.received(
1798                frame::Stream {
1799                    id: StreamId::new(Side::Server, Dir::Uni, 127),
1800                    offset: 0,
1801                    fin: true,
1802                    data: Bytes::from_static(&[]),
1803                },
1804                0
1805            ),
1806            Ok(ShouldTransmit(false))
1807        );
1808
1809        // Tighten limit by one
1810        client.set_max_concurrent(Dir::Uni, 127u32.into());
1811
1812        // Free stream 127
1813        let mut pending = Retransmits::default();
1814        let mut stream = RecvStream {
1815            id: StreamId::new(Side::Server, Dir::Uni, 127),
1816            state: &mut client,
1817            pending: &mut pending,
1818        };
1819        stream.stop(0u32.into()).unwrap();
1820
1821        // Try to open stream 128, still exceeding limit
1822        assert_eq!(
1823            client
1824                .received(
1825                    frame::Stream {
1826                        id: StreamId::new(Side::Server, Dir::Uni, 128),
1827                        offset: 0,
1828                        fin: true,
1829                        data: Bytes::from_static(&[]),
1830                    },
1831                    0
1832                )
1833                .unwrap_err()
1834                .code,
1835            TransportErrorCode::STREAM_LIMIT_ERROR
1836        );
1837
1838        // Free stream 126
1839        assert_eq!(
1840            client.received_reset(frame::ResetStream {
1841                id: StreamId::new(Side::Server, Dir::Uni, 126),
1842                error_code: 0u32.into(),
1843                final_offset: 0u32.into(),
1844            }),
1845            Ok(ShouldTransmit(false))
1846        );
1847        let mut pending = Retransmits::default();
1848        let mut stream = RecvStream {
1849            id: StreamId::new(Side::Server, Dir::Uni, 126),
1850            state: &mut client,
1851            pending: &mut pending,
1852        };
1853        stream.stop(0u32.into()).unwrap();
1854
1855        // Open stream 128
1856        assert_eq!(
1857            client.received(
1858                frame::Stream {
1859                    id: StreamId::new(Side::Server, Dir::Uni, 128),
1860                    offset: 0,
1861                    fin: true,
1862                    data: Bytes::from_static(&[]),
1863                },
1864                0
1865            ),
1866            Ok(ShouldTransmit(false))
1867        );
1868    }
1869
1870    #[test]
1871    fn remote_stream_capacity() {
1872        let mut client = make(Side::Client);
1873        for _ in 0..2 {
1874            client.set_max_concurrent(Dir::Uni, 200u32.into());
1875            client.set_max_concurrent(Dir::Bi, 201u32.into());
1876            assert_eq!(client.max_remote[Dir::Uni as usize], 200);
1877            assert_eq!(client.max_remote[Dir::Bi as usize], 201);
1878            assert_eq!(client.allocated_remote_count[Dir::Uni as usize], 200);
1879            assert_eq!(client.allocated_remote_count[Dir::Bi as usize], 201);
1880            // Slots are materialized lazily: no remote stream has been touched yet.
1881            assert!(client.recv.is_empty());
1882            assert!(client.send.is_empty());
1883        }
1884    }
1885
1886    #[test]
1887    fn lazy_remote_allocation_starts_empty() {
1888        // `StreamsState::new` must not pre-populate `send`/`recv` with placeholder slots.
1889        let client = StreamsState::new(
1890            Side::Client,
1891            10_000u32.into(),
1892            10_000u32.into(),
1893            1024 * 1024,
1894            (1024 * 1024u32).into(),
1895            (1024 * 1024u32).into(),
1896        );
1897        // No slots allocated until a stream is actually received.
1898        assert!(client.recv.is_empty());
1899        assert!(client.send.is_empty());
1900        assert_eq!(client.recv.capacity(), 0);
1901        assert_eq!(client.send.capacity(), 0);
1902    }
1903
1904    #[test]
1905    fn out_of_order_implicit_open() {
1906        // Receiving idx=5 implicitly opens idx 0..=4. A later frame for idx=3 must be
1907        // processed normally, not dropped as "closed".
1908        const STREAM_5_PAYLOAD: &[u8] = &[0xAA; 8];
1909        const STREAM_3_PAYLOAD: &[u8] = &[0xBB; 4];
1910
1911        let mut client = make(Side::Client);
1912        assert_eq!(
1913            client.received(
1914                frame::Stream {
1915                    id: StreamId::new(Side::Server, Dir::Uni, 5),
1916                    offset: 0,
1917                    fin: true,
1918                    data: Bytes::from_static(STREAM_5_PAYLOAD),
1919                },
1920                STREAM_5_PAYLOAD.len(),
1921            ),
1922            Ok(ShouldTransmit(false))
1923        );
1924        assert_eq!(client.next_remote[Dir::Uni as usize], 6);
1925        assert_eq!(
1926            client.received(
1927                frame::Stream {
1928                    id: StreamId::new(Side::Server, Dir::Uni, 3),
1929                    offset: 0,
1930                    fin: true,
1931                    data: Bytes::from_static(STREAM_3_PAYLOAD),
1932                },
1933                STREAM_3_PAYLOAD.len(),
1934            ),
1935            Ok(ShouldTransmit(false))
1936        );
1937
1938        let id = StreamId::new(Side::Server, Dir::Uni, 3);
1939        let mut pending = Retransmits::default();
1940        let mut recv = RecvStream {
1941            id,
1942            state: &mut client,
1943            pending: &mut pending,
1944        };
1945        let mut chunks = recv.read(true).unwrap();
1946        assert_eq!(
1947            chunks.next(STREAM_3_PAYLOAD.len()).unwrap().unwrap().bytes,
1948            STREAM_3_PAYLOAD
1949        );
1950        let _ = chunks.finalize();
1951    }
1952
1953    #[test]
1954    fn frame_for_closed_stream_is_dropped() {
1955        // After a remote stream is fully freed, a subsequent frame for the same id must be
1956        // dropped — absence from the map unambiguously means "closed" for ids below the
1957        // frontier.
1958        const PAYLOAD: &[u8] = &[0; 4];
1959
1960        let mut client = make(Side::Client);
1961        let id = StreamId::new(Side::Server, Dir::Uni, 0);
1962        assert_eq!(
1963            client.received(
1964                frame::Stream {
1965                    id,
1966                    offset: 0,
1967                    fin: true,
1968                    data: Bytes::from_static(PAYLOAD),
1969                },
1970                PAYLOAD.len(),
1971            ),
1972            Ok(ShouldTransmit(false))
1973        );
1974        // Stop the stream so it's fully freed.
1975        let mut pending = Retransmits::default();
1976        RecvStream {
1977            id,
1978            state: &mut client,
1979            pending: &mut pending,
1980        }
1981        .stop(0u32.into())
1982        .unwrap();
1983        assert!(!client.recv.contains_key(&id));
1984
1985        // A stray retransmit for the freed stream must be dropped without resurrecting state.
1986        assert_eq!(
1987            client.received(
1988                frame::Stream {
1989                    id,
1990                    offset: 0,
1991                    fin: true,
1992                    data: Bytes::from_static(PAYLOAD),
1993                },
1994                PAYLOAD.len(),
1995            ),
1996            Ok(ShouldTransmit(false))
1997        );
1998        assert!(!client.recv.contains_key(&id));
1999    }
2000
2001    #[test]
2002    fn churn_keeps_maps_bounded() {
2003        // Rapidly open + fully close a long sequence of remote streams. The maps must stay
2004        // bounded (only active streams are materialized) even though thousands of ids have
2005        // been used over the connection's lifetime.
2006        const N: u64 = 5_000;
2007
2008        let mut client = make(Side::Client);
2009        for i in 0..N {
2010            let id = StreamId::new(Side::Server, Dir::Uni, i);
2011            assert_eq!(
2012                client.received(
2013                    frame::Stream {
2014                        id,
2015                        offset: 0,
2016                        fin: true,
2017                        data: Bytes::from_static(&[0; 1]),
2018                    },
2019                    1,
2020                ),
2021                Ok(ShouldTransmit(false))
2022            );
2023            let mut pending = Retransmits::default();
2024            let mut recv = RecvStream {
2025                id,
2026                state: &mut client,
2027                pending: &mut pending,
2028            };
2029            let mut chunks = recv.read(true).unwrap();
2030            let _ = chunks.next(1).unwrap();
2031            assert!(chunks.next(1).unwrap().is_none());
2032            let _ = chunks.finalize();
2033        }
2034        // Every stream was fully drained; the map must be empty.
2035        assert_eq!(client.recv.len(), 0);
2036        assert_eq!(client.send.len(), 0);
2037    }
2038
2039    #[test]
2040    fn expand_receive_window() {
2041        let mut server = make(Side::Server);
2042        let new_receive_window = 2 * server.receive_window as u32;
2043        let expanded = server.set_receive_window(new_receive_window.into());
2044        assert!(expanded);
2045        assert_eq!(server.receive_window, new_receive_window as u64);
2046        assert_eq!(server.local_max_data, new_receive_window as u64);
2047        assert_eq!(server.receive_window_shrink_debt, 0);
2048        let prev_local_max_data = server.local_max_data;
2049
2050        // credit, expecting all of them added to local_max_data
2051        let credits = 1024u64;
2052        let should_transmit = server.add_read_credits(credits);
2053        assert_eq!(server.receive_window_shrink_debt, 0);
2054        assert_eq!(server.local_max_data, prev_local_max_data + credits);
2055        assert!(should_transmit.should_transmit());
2056    }
2057
2058    #[test]
2059    fn shrink_receive_window() {
2060        let mut server = make(Side::Server);
2061        let new_receive_window = server.receive_window as u32 / 2;
2062        let prev_local_max_data = server.local_max_data;
2063
2064        // shrink the receive_winbow, local_max_data is not expected to be changed
2065        let shrink_diff = server.receive_window - new_receive_window as u64;
2066        let expanded = server.set_receive_window(new_receive_window.into());
2067        assert!(!expanded);
2068        assert_eq!(server.receive_window, new_receive_window as u64);
2069        assert_eq!(server.local_max_data, prev_local_max_data);
2070        assert_eq!(server.receive_window_shrink_debt, shrink_diff);
2071        let prev_local_max_data = server.local_max_data;
2072
2073        // credit twice, local_max_data does not change as it is absorbed by
2074        // receive_window_shrink_debt
2075        let credits = 1024u64;
2076        for _ in 0..2 {
2077            let expected_receive_window_shrink_debt = server.receive_window_shrink_debt - credits;
2078            let should_transmit = server.add_read_credits(credits);
2079            assert_eq!(
2080                server.receive_window_shrink_debt,
2081                expected_receive_window_shrink_debt
2082            );
2083            assert_eq!(server.local_max_data, prev_local_max_data);
2084            assert!(!should_transmit.should_transmit());
2085        }
2086
2087        // credit again which exceeds all remaining expected_receive_window_shrink_debt
2088        let credits = 1024 * 512;
2089        let prev_local_max_data = server.local_max_data;
2090        let expected_local_max_data =
2091            server.local_max_data + (credits - server.receive_window_shrink_debt);
2092        let _should_transmit = server.add_read_credits(credits);
2093        assert_eq!(server.receive_window_shrink_debt, 0);
2094        assert_eq!(server.local_max_data, expected_local_max_data);
2095        assert!(server.local_max_data > prev_local_max_data);
2096
2097        // credit again, all should be added to local_max_data
2098        let credits = 1024 * 512;
2099        let expected_local_max_data = server.local_max_data + credits;
2100        let should_transmit = server.add_read_credits(credits);
2101        assert_eq!(server.receive_window_shrink_debt, 0);
2102        assert_eq!(server.local_max_data, expected_local_max_data);
2103        assert!(should_transmit.should_transmit());
2104    }
2105
2106    #[test]
2107    fn expand_send_window() {
2108        let mut server = make(Side::Server);
2109
2110        let initial_send_window = server.send_window;
2111        let larger_send_window = initial_send_window * 2;
2112
2113        // Set `initial_max_data` larger than `send_window` so we're limited by local flow control
2114        server.set_params(&TransportParameters {
2115            initial_max_data: VarInt::MAX,
2116            initial_max_stream_data_uni: VarInt::MAX,
2117            initial_max_streams_uni: VarInt::from_u32(100),
2118            ..TransportParameters::default()
2119        });
2120
2121        assert_eq!(server.write_limit(), initial_send_window);
2122        assert_eq!(server.poll(), None);
2123
2124        let mut retransmits = Retransmits::default();
2125        let conn_state = ConnState::established();
2126
2127        let stream_id = Streams {
2128            state: &mut server,
2129            conn_state: &conn_state,
2130        }
2131        .open(Dir::Uni)
2132        .expect("should be able to open a stream");
2133
2134        let mut stream = SendStream {
2135            id: stream_id,
2136            state: &mut server,
2137            pending: &mut retransmits,
2138            conn_state: &conn_state,
2139        };
2140
2141        // Check that the stream accepts `initial_send_window` bytes
2142        let initial_send_len = initial_send_window as usize;
2143        let data = vec![0xFFu8; initial_send_len];
2144
2145        assert_eq!(stream.write(&data), Ok(initial_send_len));
2146
2147        // Try to write the same data again, observe that it's blocked
2148        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2149
2150        // Check that we get a `Writable` event after increasing the send window
2151        stream.state.set_send_window(larger_send_window);
2152        assert_eq!(
2153            stream.state.poll(),
2154            Some(StreamEvent::Writable { id: stream_id })
2155        );
2156
2157        // Check that the stream accepts the exact same amount of data again
2158        assert_eq!(stream.write(&data), Ok(initial_send_len));
2159        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2160
2161        assert_eq!(stream.state.poll(), None);
2162
2163        // Ack the data
2164        stream.state.received_ack_of(frame::StreamMeta {
2165            id: stream_id,
2166            offsets: 0..larger_send_window,
2167            fin: false,
2168        });
2169
2170        assert_eq!(
2171            stream.state.poll(),
2172            Some(StreamEvent::Writable { id: stream_id })
2173        );
2174
2175        // Check that our full send window is available again
2176        assert_eq!(stream.write(&data), Ok(initial_send_len));
2177        assert_eq!(stream.write(&data), Ok(initial_send_len));
2178        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2179    }
2180
2181    #[test]
2182    fn shrink_send_window() {
2183        let mut server = make(Side::Server);
2184
2185        let initial_send_window = server.send_window;
2186        let smaller_send_window = server.send_window / 2;
2187
2188        // Set `initial_max_data` larger than `send_window` so we're limited by local flow control
2189        server.set_params(&TransportParameters {
2190            initial_max_data: VarInt::MAX,
2191            initial_max_stream_data_uni: VarInt::MAX,
2192            initial_max_streams_uni: VarInt::from_u32(100),
2193            ..TransportParameters::default()
2194        });
2195
2196        assert_eq!(server.write_limit(), initial_send_window);
2197        assert_eq!(server.poll(), None);
2198
2199        let mut retransmits = Retransmits::default();
2200        let conn_state = ConnState::established();
2201
2202        let stream_id = Streams {
2203            state: &mut server,
2204            conn_state: &conn_state,
2205        }
2206        .open(Dir::Uni)
2207        .expect("should be able to open a stream");
2208
2209        let mut stream = SendStream {
2210            id: stream_id,
2211            state: &mut server,
2212            pending: &mut retransmits,
2213            conn_state: &conn_state,
2214        };
2215
2216        let initial_send_len = initial_send_window as usize;
2217
2218        let data = vec![0xFFu8; initial_send_len];
2219
2220        // Assert that the full send window is accepted
2221        assert_eq!(stream.write(&data), Ok(initial_send_len));
2222        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2223
2224        assert_eq!(stream.state.write_limit(), 0);
2225        assert_eq!(stream.state.poll(), None);
2226
2227        // Shrink our send window, assert that it's still not writable
2228        stream.state.set_send_window(smaller_send_window);
2229        assert_eq!(stream.state.write_limit(), 0);
2230        assert_eq!(stream.state.poll(), None);
2231
2232        // Assert that data is still not accepted
2233        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2234
2235        // Ack some data, assert that writes are still not accepted due to outstanding sends
2236        stream.state.received_ack_of(frame::StreamMeta {
2237            id: stream_id,
2238            offsets: 0..smaller_send_window,
2239            fin: false,
2240        });
2241
2242        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2243
2244        // Ack the rest of the data
2245        stream.state.received_ack_of(frame::StreamMeta {
2246            id: stream_id,
2247            offsets: smaller_send_window..initial_send_window,
2248            fin: false,
2249        });
2250
2251        // This should generate a `Writable` event
2252        assert_eq!(
2253            stream.state.poll(),
2254            Some(StreamEvent::Writable { id: stream_id })
2255        );
2256        assert_eq!(stream.state.write_limit(), smaller_send_window);
2257
2258        // Assert that only `smaller_send_window` bytes are accepted
2259        assert_eq!(stream.write(&data), Ok(smaller_send_window as usize));
2260        assert_eq!(stream.write(&data), Err(WriteError::Blocked));
2261    }
2262}