net_io/
session.rs

1//! Sans-IO sessions: a protocol engine reacts to bytes and returns what to do next.
2
3use std::collections::VecDeque;
4use std::fmt;
5use std::sync::Arc;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::time::{Duration, Instant};
8
9use crate::buf::{Buffer, BufferPool};
10use crate::completer::Completer;
11use crate::error::{CloseReason, LimitKind, NetError, RefusalCause, TimeoutKind};
12use crate::gate::KillHandle;
13use crate::id::WorkerId;
14use crate::info::{ConnectionInfo, GateTag};
15use crate::limits::SizeClass;
16use crate::link::Link;
17
18/// A sans-IO protocol engine bound to one connection.
19///
20/// net-io reads, calls the session with the bytes, and does what the returned
21/// [`Step`] says. Sessions are called only on the worker that owns the
22/// connection, never concurrently.
23///
24/// # Examples
25///
26/// ```no_run
27/// use net_io::{Input, Session, SessionContext, Step};
28///
29/// struct Echo;
30/// impl Session for Echo {
31///     fn on_data(&mut self, cx: &mut SessionContext<'_>, input: &mut Input) -> Step {
32///         // Copy as much as the output budget takes; the rest stays in `input`,
33///         // and net-io calls on_data again after the flush because the input changed.
34///         let taken = cx.output().push_partial(input.bytes());
35///         input.consume(taken);
36///         Step::Write
37///     }
38/// }
39///
40/// let server = net_io::serve_sessions("127.0.0.1:7002", |_info: &net_io::ConnectionInfo| Echo)?;
41/// server.wait_for_stop();
42/// # Ok::<(), net_io::NetError>(())
43/// ```
44pub trait Session: Send + 'static {
45    /// Once, after admission and the handshake.
46    fn on_open(&mut self, cx: &mut SessionContext<'_>) -> Step {
47        let _ = cx;
48        Step::Read
49    }
50
51    /// Bytes arrived; `input` holds everything not yet consumed.
52    fn on_data(&mut self, cx: &mut SessionContext<'_>, input: &mut Input) -> Step;
53
54    /// After a `Step::Write` flushed all pending output.
55    fn on_writable(&mut self, cx: &mut SessionContext<'_>) -> Step {
56        let _ = cx;
57        Step::Read
58    }
59
60    /// Another thread posted a completion through this connection's [`Completer`].
61    fn on_wake(&mut self, cx: &mut SessionContext<'_>, token: u32, value: u64) -> Step {
62        let _ = (cx, token, value);
63        Step::Abort
64    }
65
66    /// A deadline passed. `Close` or `Abort` closes with the timeout's own reason
67    /// and close style; anything else re-arms with the same limits.
68    fn on_timeout(&mut self, cx: &mut SessionContext<'_>, which: TimeoutKind) -> Step {
69        let _ = (cx, which);
70        Step::Close
71    }
72
73    /// The peer sent FIN; the session may keep writing (half-close).
74    fn on_peer_closed(&mut self, cx: &mut SessionContext<'_>) -> Step {
75        let _ = cx;
76        Step::Close
77    }
78
79    /// The gate denied a write and no byte of the current response has left:
80    /// write the protocol's refusal for `code` into `output`.
81    fn on_denied(&mut self, code: u16, output: &mut Output) {
82        let _ = (code, output);
83    }
84
85    /// Exactly once, with the one close reason.
86    fn on_close(&mut self, info: &ConnectionInfo, reason: CloseReason) {
87        let _ = (info, reason);
88    }
89}
90
91/// Creates one session per admitted connection. Any `Fn(&ConnectionInfo) -> S` is one.
92///
93/// # Examples
94///
95/// ```
96/// use net_io::{ConnectionInfo, Input, RefusalCause, Session, SessionContext, SessionFactory, Step};
97///
98/// struct Discard;
99/// impl Session for Discard {
100///     fn on_data(&mut self, _cx: &mut SessionContext<'_>, input: &mut Input) -> Step {
101///         let n = input.bytes().len();
102///         input.consume(n);
103///         Step::Read
104///     }
105/// }
106///
107/// struct OnlyLoopback;
108/// impl SessionFactory for OnlyLoopback {
109///     type Session = Discard;
110///     fn open(&self, info: &ConnectionInfo) -> Result<Discard, RefusalCause> {
111///         match info.client_addr() {
112///             Some(addr) if addr.ip().is_loopback() => Ok(Discard),
113///             _ => Err(RefusalCause::Session),
114///         }
115///     }
116/// }
117///
118/// let info = ConnectionInfo::for_testing("10.0.0.1:1".parse().unwrap());
119/// assert!(OnlyLoopback.open(&info).is_err());
120/// ```
121pub trait SessionFactory: Send + Sync + 'static {
122    /// The session type.
123    type Session: Session;
124
125    /// A session for this connection.
126    ///
127    /// # Errors
128    ///
129    /// Refusing closes the connection with `CloseReason::Refused(cause)`.
130    fn open(&self, info: &ConnectionInfo) -> Result<Self::Session, RefusalCause>;
131}
132
133impl<F, S> SessionFactory for F
134where
135    F: Fn(&ConnectionInfo) -> S + Send + Sync + 'static,
136    S: Session,
137{
138    type Session = S;
139    fn open(&self, info: &ConnectionInfo) -> Result<S, RefusalCause> {
140        Ok(self(info))
141    }
142}
143
144/// What the session wants next.
145///
146/// ```
147/// assert_ne!(net_io::Step::Close, net_io::Step::Abort);
148/// ```
149#[non_exhaustive]
150#[derive(Clone, Copy, Debug, PartialEq, Eq)]
151pub enum Step {
152    /// Wait for more input. Output pushed in this call is not sent: return
153    /// `Write` (or `Close`) to flush it first.
154    Read,
155    /// Flush pending output, then `on_writable`.
156    Write,
157    /// No reads, no idle timer, until a completion or the park deadline.
158    Park {
159        /// The deadline; `None` uses the park timeout.
160        until: Option<Instant>,
161    },
162    /// Graceful: flush, FIN, bounded linger.
163    Close,
164    /// RST now.
165    Abort,
166}
167
168/// Which deadlines run.
169///
170/// ```
171/// assert_eq!(net_io::Phase::default(), net_io::Phase::Idle);
172/// ```
173#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
174pub enum Phase {
175    /// Between exchanges: the idle timeout.
176    #[default]
177    Idle,
178    /// Receiving an exchange: the read timeout, the inbound rate, the exchange deadline.
179    Receiving,
180    /// Sending: the write timeout, the outbound rate, the exchange deadline.
181    Sending,
182}
183
184/// Why a swap was refused.
185///
186/// ```
187/// assert_ne!(net_io::SwapError::NotEmpty, net_io::SwapError::Busy);
188/// ```
189#[non_exhaustive]
190#[derive(Clone, Copy, Debug, PartialEq, Eq)]
191pub enum SwapError {
192    /// The buffers are of different size classes.
193    ClassMismatch,
194    /// The incoming buffer has data.
195    NotEmpty,
196    /// A driver operation holds the buffer.
197    Busy,
198}
199
200impl fmt::Display for SwapError {
201    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
202        f.write_str(match self {
203            SwapError::ClassMismatch => "the buffers are of different size classes",
204            SwapError::NotEmpty => "the incoming buffer is not empty",
205            SwapError::Busy => "a driver operation holds the buffer",
206        })
207    }
208}
209
210impl std::error::Error for SwapError {}
211
212/// The connection's read buffer: a zero-copy view of what was read and not yet consumed.
213///
214/// ```
215/// let mut input = net_io::Input::from_bytes(b"GET / HTTP/1.1\r\n");
216/// assert_eq!(input.bytes(), b"GET / HTTP/1.1\r\n");
217/// input.consume(4);
218/// assert_eq!(input.bytes(), b"/ HTTP/1.1\r\n");
219/// ```
220#[derive(Debug)]
221pub struct Input {
222    pub(crate) buf: Buffer,
223    pub(crate) start: usize,
224    pub(crate) changed: bool,
225}
226
227impl Input {
228    pub(crate) fn new(buf: Buffer) -> Input {
229        Input {
230            buf,
231            start: 0,
232            changed: false,
233        }
234    }
235
236    /// An input holding `bytes`, for tests of sessions (allocates).
237    #[must_use]
238    pub fn from_bytes(bytes: &[u8]) -> Input {
239        let class = (bytes.len() as u32)
240            .max(1024)
241            .next_power_of_two()
242            .max(SizeClass::SMALL.bytes());
243        let class = SizeClass::new(class.min(1 << 20)).unwrap_or(SizeClass::LARGE);
244        let mut pool = BufferPool::standalone(u64::from(class.bytes()));
245        let mut input = Input::new(match pool.take(class) {
246            Some(b) => b,
247            None => return Input::new(Buffer::detached(class)),
248        });
249        let n = bytes.len().min(input.buf.capacity());
250        input.buf.spare()[..n].copy_from_slice(&bytes[..n]);
251        input.buf.set_len(n);
252        input
253    }
254
255    /// The unconsumed bytes.
256    #[must_use]
257    pub fn bytes(&self) -> &[u8] {
258        &self.buf.as_slice()[self.start..]
259    }
260
261    /// Consumes `n` bytes (saturating at the length).
262    pub fn consume(&mut self, n: usize) {
263        let n = n.min(self.buf.len() - self.start);
264        self.start += n;
265        if n > 0 {
266            self.changed = true;
267        }
268        if self.start == self.buf.len() {
269            self.buf.clear();
270            self.start = 0;
271        }
272    }
273
274    /// Whether the buffer is full: the session must consume or grow it.
275    #[must_use]
276    pub fn is_full(&self) -> bool {
277        self.buf.len() == self.buf.capacity() && self.start == 0
278    }
279
280    /// The buffer's size class.
281    #[must_use]
282    pub fn class(&self) -> SizeClass {
283        self.buf.class()
284    }
285
286    /// Exchanges the read buffer with an empty `other` of the same class in O(1):
287    /// `other` now holds the unconsumed bytes, and the connection reads into the
288    /// buffer it was given.
289    ///
290    /// # Errors
291    ///
292    /// [`SwapError::ClassMismatch`], [`SwapError::NotEmpty`].
293    ///
294    /// ```
295    /// use net_io::{BufferPool, Input};
296    /// let mut input = Input::from_bytes(b"body");
297    /// let mut pool = BufferPool::standalone(1 << 20);
298    /// let mut spare = pool.take(input.class()).unwrap();
299    /// input.swap(&mut spare).unwrap();
300    /// assert_eq!(spare.as_slice(), b"body");
301    /// assert!(input.bytes().is_empty());
302    /// ```
303    pub fn swap(&mut self, other: &mut Buffer) -> Result<(), SwapError> {
304        if other.class() != self.buf.class() {
305            return Err(SwapError::ClassMismatch);
306        }
307        if !other.is_empty() {
308            return Err(SwapError::NotEmpty);
309        }
310        if self.start > 0 {
311            self.buf.compact(self.start);
312            self.start = 0;
313        }
314        std::mem::swap(&mut self.buf, other);
315        self.changed = true;
316        Ok(())
317    }
318
319    /// Room to read into, compacting first when the free tail is under a quarter.
320    pub(crate) fn spare_for_read(&mut self) -> &mut [u8] {
321        let free = self.buf.capacity() - self.buf.len();
322        if self.start > 0 && free < self.buf.capacity() / 4 {
323            self.buf.compact(self.start);
324            self.start = 0;
325        }
326        self.buf.spare()
327    }
328
329    pub(crate) fn filled(&mut self, n: usize) {
330        let len = self.buf.len() + n;
331        self.buf.set_len(len);
332        self.changed = true;
333    }
334}
335
336/// The output budget was reached; return `Step::Write` and continue after the flush.
337///
338/// ```
339/// assert_eq!(net_io::OutputFull.to_string(), "the connection's output budget is full");
340/// ```
341#[derive(Clone, Copy, Debug, PartialEq, Eq)]
342pub struct OutputFull;
343
344impl fmt::Display for OutputFull {
345    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
346        f.write_str("the connection's output budget is full")
347    }
348}
349
350impl std::error::Error for OutputFull {}
351
352#[derive(Debug)]
353pub(crate) enum Chunk {
354    Pooled(Buffer, usize),
355    Static(&'static [u8], usize),
356}
357
358impl Chunk {
359    pub(crate) fn remaining(&self) -> &[u8] {
360        match self {
361            Chunk::Pooled(b, off) => &b.as_slice()[*off..],
362            Chunk::Static(s, off) => &s[*off..],
363        }
364    }
365}
366
367/// What the session wants written: a bounded queue of chunks.
368///
369/// ```
370/// let mut out = net_io::Output::standalone(64 * 1024);
371/// out.push(b"hello ").unwrap();
372/// out.push_static(b"world").unwrap();
373/// assert_eq!(out.pending(), 11);
374/// ```
375#[derive(Debug)]
376pub struct Output {
377    pub(crate) chunks: VecDeque<Chunk>,
378    free: Vec<Buffer>,
379    budget: Arc<AtomicU64>,
380    limit: u32,
381    used: u32,
382    wipe: bool,
383    pub(crate) response_open: bool,
384    pub(crate) response_tag: Option<GateTag>,
385    pub(crate) bytes_since_begin: u64,
386}
387
388const OUTPUT_CLASS: SizeClass = SizeClass::MEDIUM;
389const MAX_CHUNKS: usize = 64;
390
391impl Output {
392    pub(crate) fn new(budget: Arc<AtomicU64>, limit: u32, wipe: bool) -> Output {
393        Output {
394            chunks: VecDeque::with_capacity(MAX_CHUNKS),
395            free: Vec::new(),
396            budget,
397            limit,
398            used: 0,
399            wipe,
400            response_open: false,
401            response_tag: None,
402            bytes_since_begin: 0,
403        }
404    }
405
406    /// An output with its own budget, for tests of sessions (allocates).
407    #[must_use]
408    pub fn standalone(limit: u32) -> Output {
409        Output::new(Arc::new(AtomicU64::new(u64::from(limit))), limit, true)
410    }
411
412    fn take_buffer(&mut self) -> Option<Buffer> {
413        if let Some(b) = self.free.pop() {
414            return Some(b);
415        }
416        let bytes = u64::from(OUTPUT_CLASS.bytes());
417        let mut left = self.budget.load(Ordering::Relaxed);
418        loop {
419            if left < bytes {
420                return None;
421            }
422            match self.budget.compare_exchange_weak(
423                left,
424                left - bytes,
425                Ordering::AcqRel,
426                Ordering::Relaxed,
427            ) {
428                Ok(_) => return Some(Buffer::detached(OUTPUT_CLASS)),
429                Err(actual) => left = actual,
430            }
431        }
432    }
433
434    fn room(&self) -> usize {
435        let tail = match self.chunks.back() {
436            Some(Chunk::Pooled(b, _)) => b.capacity() - b.len(),
437            _ => 0,
438        };
439        let more = (self.limit.saturating_sub(self.used) / OUTPUT_CLASS.bytes()) as usize
440            * OUTPUT_CLASS.bytes() as usize;
441        tail + more
442    }
443
444    /// Copies as much of `bytes` as the budget allows; returns how many were taken.
445    pub fn push_partial(&mut self, mut bytes: &[u8]) -> usize {
446        let mut taken = 0;
447        while !bytes.is_empty() {
448            if let Some(Chunk::Pooled(b, _)) = self.chunks.back_mut() {
449                let spare = b.spare();
450                if !spare.is_empty() {
451                    let n = spare.len().min(bytes.len());
452                    spare[..n].copy_from_slice(&bytes[..n]);
453                    let len = b.len() + n;
454                    b.set_len(len);
455                    bytes = &bytes[n..];
456                    taken += n;
457                    continue;
458                }
459            }
460            if self.chunks.len() >= MAX_CHUNKS
461                || self.used.saturating_add(OUTPUT_CLASS.bytes()) > self.limit
462            {
463                break;
464            }
465            match self.take_buffer() {
466                Some(b) => {
467                    self.used += OUTPUT_CLASS.bytes();
468                    self.chunks.push_back(Chunk::Pooled(b, 0));
469                }
470                None => break,
471            }
472        }
473        taken
474    }
475
476    /// Copies all of `bytes`, or nothing.
477    ///
478    /// # Errors
479    ///
480    /// [`OutputFull`] when the budget cannot take all of it.
481    pub fn push(&mut self, bytes: &[u8]) -> Result<(), OutputFull> {
482        if bytes.len() > self.room() {
483            return Err(OutputFull);
484        }
485        let n = self.push_partial(bytes);
486        if n == bytes.len() {
487            Ok(())
488        } else {
489            Err(OutputFull)
490        }
491    }
492
493    /// Moves an owned buffer in (no copy).
494    ///
495    /// # Errors
496    ///
497    /// [`OutputFull`] when the budget or the chunk limit is reached.
498    pub fn push_buffer(&mut self, buffer: Buffer) -> Result<(), OutputFull> {
499        let size = u32::try_from(buffer.capacity()).unwrap_or(u32::MAX);
500        if self.chunks.len() >= MAX_CHUNKS || self.used.saturating_add(size) > self.limit {
501            return Err(OutputFull);
502        }
503        self.used += size;
504        self.chunks.push_back(Chunk::Pooled(buffer, 0));
505        Ok(())
506    }
507
508    /// Queues a static slice (no copy, not counted against the budget).
509    ///
510    /// # Errors
511    ///
512    /// [`OutputFull`] when the chunk limit is reached.
513    pub fn push_static(&mut self, bytes: &'static [u8]) -> Result<(), OutputFull> {
514        if self.chunks.len() >= MAX_CHUNKS {
515            return Err(OutputFull);
516        }
517        if !bytes.is_empty() {
518            self.chunks.push_back(Chunk::Static(bytes, 0));
519        }
520        Ok(())
521    }
522
523    /// Marks the start of a response (for the gate's denial boundary).
524    pub fn begin_response(&mut self) {
525        self.response_open = true;
526        self.response_tag = None;
527        self.bytes_since_begin = 0;
528    }
529
530    /// Marks the start of a response and tags it for the gate (`WriteCheck::response`).
531    pub fn begin_tagged_response(&mut self, tag: GateTag) {
532        self.begin_response();
533        self.response_tag = Some(tag);
534    }
535
536    /// Marks the end of the response; clears the tag.
537    pub fn end_response(&mut self) {
538        self.response_open = false;
539        self.response_tag = None;
540    }
541
542    /// Bytes queued and not yet sent.
543    #[must_use]
544    pub fn pending(&self) -> usize {
545        self.chunks.iter().map(|c| c.remaining().len()).sum()
546    }
547
548    /// Whether nothing is queued.
549    #[must_use]
550    pub fn is_empty(&self) -> bool {
551        self.chunks.is_empty()
552    }
553
554    /// Overwrites every pending chunk net-io owns with zeroes and drops it.
555    pub fn wipe_all(&mut self) {
556        while let Some(chunk) = self.chunks.pop_front() {
557            if let Chunk::Pooled(mut b, _) = chunk {
558                b.wipe();
559                self.release_chunk(b);
560            }
561        }
562        self.used = 0;
563    }
564
565    fn release_chunk(&mut self, mut b: Buffer) {
566        self.used = self
567            .used
568            .saturating_sub(u32::try_from(b.capacity()).unwrap_or(u32::MAX));
569        if self.wipe {
570            b.wipe();
571        } else {
572            b.clear();
573        }
574        if b.class() == OUTPUT_CLASS {
575            self.free.push(b);
576        } else {
577            // A buffer moved in with push_buffer goes back to the shared budget.
578            self.budget
579                .fetch_add(b.capacity() as u64, Ordering::Release);
580        }
581    }
582
583    /// The bytes one send offers: at most the chunks one vectored call takes.
584    #[cfg(feature = "reactor")]
585    pub(crate) fn pending_in_send(&self) -> usize {
586        self.chunks
587            .iter()
588            .take(crate::sys::MAX_SLICES)
589            .map(|c| c.remaining().len())
590            .sum()
591    }
592
593    /// The chunks as slices for one send; returns how many.
594    #[cfg(feature = "reactor")]
595    pub(crate) fn fill_slices<'a>(&'a self, out: &mut [std::io::IoSlice<'a>]) -> usize {
596        let mut n = 0;
597        for chunk in self.chunks.iter().take(out.len()) {
598            out[n] = std::io::IoSlice::new(chunk.remaining());
599            n += 1;
600        }
601        n
602    }
603
604    /// Marks `n` bytes as sent, releasing finished chunks.
605    pub(crate) fn advance(&mut self, mut n: usize) {
606        while let Some(front) = self.chunks.front_mut() {
607            let rem = front.remaining().len();
608            if n < rem {
609                match front {
610                    Chunk::Pooled(_, off) | Chunk::Static(_, off) => *off += n,
611                }
612                return;
613            }
614            n -= rem;
615            if let Some(Chunk::Pooled(b, _)) = self.chunks.pop_front() {
616                self.release_chunk(b);
617            }
618        }
619    }
620
621    /// Resets for the next connection.
622    pub(crate) fn reset(&mut self) {
623        self.wipe_all();
624        self.response_open = false;
625        self.response_tag = None;
626        self.bytes_since_begin = 0;
627    }
628}
629
630/// Everything a session may touch, for one call.
631pub struct SessionContext<'a> {
632    pub(crate) link: &'a mut Link,
633    pub(crate) output: &'a mut Output,
634    pub(crate) pool: &'a mut BufferPool,
635    pub(crate) completer: &'a Completer,
636    pub(crate) draining: bool,
637    pub(crate) pressure: f32,
638    pub(crate) worker: WorkerId,
639    pub(crate) grow_to: Option<SizeClass>,
640}
641
642impl SessionContext<'_> {
643    /// The connection's facts.
644    #[must_use]
645    pub fn info(&self) -> &ConnectionInfo {
646        &self.link.info
647    }
648
649    /// The output queue.
650    pub fn output(&mut self) -> &mut Output {
651        self.output
652    }
653
654    /// The worker's clock.
655    #[must_use]
656    pub fn now(&self) -> Instant {
657        self.link.now()
658    }
659
660    /// Chooses which deadlines run.
661    pub fn set_phase(&mut self, phase: crate::session::Phase) {
662        self.link.set_phase(phase);
663    }
664
665    /// The deadline for the current exchange (`None` clears it).
666    pub fn set_exchange_deadline(&mut self, at: Option<Instant>) {
667        self.link.set_exchange_deadline(at);
668    }
669
670    /// Overrides this connection's idle timeout from its next wait for an
671    /// exchange on; `None` restores the listener's limit. For example, a
672    /// protocol shortens keep-alive while the worker pool is nearly full.
673    /// Values below the timer tick are raised to one tick.
674    ///
675    /// ```no_run
676    /// use std::time::Duration;
677    /// use net_io::{Input, Session, SessionContext, Step};
678    ///
679    /// struct KeepAlive;
680    /// impl Session for KeepAlive {
681    ///     fn on_data(&mut self, cx: &mut SessionContext<'_>, input: &mut Input) -> Step {
682    ///         let n = input.bytes().len();
683    ///         input.consume(n);
684    ///         let short = (cx.pressure() > 0.75).then_some(Duration::from_secs(2));
685    ///         cx.set_idle_timeout(short);
686    ///         Step::Write
687    ///     }
688    /// }
689    /// ```
690    pub fn set_idle_timeout(&mut self, timeout: Option<Duration>) {
691        self.link.set_idle_timeout(timeout);
692    }
693
694    /// Whether drain started: finish the current exchange, then close.
695    #[must_use]
696    pub fn draining(&self) -> bool {
697        self.draining
698    }
699
700    /// 0.0 to 1.0: busy workers (blocking driver); 1.0 while the control line sheds.
701    #[must_use]
702    pub fn pressure(&self) -> f32 {
703        self.pressure
704    }
705
706    /// A handle other threads can use to abort this connection.
707    #[must_use]
708    pub fn kill_handle(&self) -> KillHandle {
709        KillHandle {
710            slot: self.link.slot.clone(),
711            generation: self.link.generation,
712            id: self.link.info.id,
713        }
714    }
715
716    /// This worker's completion queue.
717    #[must_use]
718    pub fn completer(&self) -> Completer {
719        self.completer.clone()
720    }
721
722    /// The worker that runs this connection.
723    #[must_use]
724    pub fn worker(&self) -> WorkerId {
725        self.worker
726    }
727
728    /// This worker's buffer pool.
729    pub fn buffers(&mut self) -> &mut BufferPool {
730        self.pool
731    }
732
733    /// Returns a buffer (for example one taken by `Input::swap`); it is wiped.
734    pub fn release(&mut self, buffer: Buffer) {
735        self.pool.put(buffer);
736    }
737
738    /// Moves the unconsumed input into a buffer of `class` after this call returns.
739    ///
740    /// # Errors
741    ///
742    /// `Limit { which: ReadBuffer }` above the maximum read class.
743    pub fn grow_input(&mut self, class: SizeClass) -> Result<(), NetError> {
744        if class > self.link.limits.read_class_max {
745            return Err(NetError::Limit {
746                which: LimitKind::ReadBuffer,
747            });
748        }
749        self.grow_to = Some(class);
750        Ok(())
751    }
752}
753
754impl fmt::Debug for SessionContext<'_> {
755    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
756        f.debug_struct("SessionContext")
757            .field("id", &self.link.info.id)
758            .field("draining", &self.draining)
759            .finish()
760    }
761}
762
763#[cfg(test)]
764mod tests {
765    use super::*;
766
767    #[test]
768    fn test_input_consume_saturates_and_resets() {
769        let mut i = Input::from_bytes(b"abc");
770        i.consume(10);
771        assert!(i.bytes().is_empty());
772        assert_eq!(i.start, 0);
773    }
774
775    #[test]
776    fn test_input_swap_rules() {
777        let mut i = Input::from_bytes(b"xyz");
778        let mut pool = BufferPool::standalone(1 << 22);
779        let mut wrong = pool.take(SizeClass::LARGE).unwrap();
780        assert_eq!(i.swap(&mut wrong), Err(SwapError::ClassMismatch));
781        let mut full = pool.take(i.class()).unwrap();
782        full.set_len(1);
783        assert_eq!(i.swap(&mut full), Err(SwapError::NotEmpty));
784        let mut empty = pool.take(i.class()).unwrap();
785        i.consume(1);
786        i.swap(&mut empty).unwrap();
787        assert_eq!(empty.as_slice(), b"yz");
788    }
789
790    #[test]
791    fn test_output_push_all_or_nothing_and_partial() {
792        let mut o = Output::standalone(OUTPUT_CLASS.bytes());
793        let big = vec![7u8; OUTPUT_CLASS.bytes() as usize + 1];
794        assert_eq!(o.push(&big), Err(OutputFull));
795        assert!(o.is_empty());
796        assert_eq!(o.push_partial(&big), OUTPUT_CLASS.bytes() as usize);
797        assert_eq!(o.pending(), OUTPUT_CLASS.bytes() as usize);
798    }
799
800    #[test]
801    fn test_output_advance_releases_chunks() {
802        let mut o = Output::standalone(1 << 20);
803        o.push(b"hello").unwrap();
804        o.push_static(b" world").unwrap();
805        o.advance(7);
806        assert_eq!(o.pending(), 4);
807        assert_eq!(o.chunks.len(), 1);
808        assert_eq!(o.used, 0);
809        o.advance(4);
810        assert!(o.is_empty());
811    }
812
813    #[test]
814    fn test_output_response_markers() {
815        let mut o = Output::standalone(1 << 20);
816        o.begin_tagged_response(GateTag(3));
817        assert_eq!(o.response_tag, Some(GateTag(3)));
818        o.end_response();
819        assert_eq!(o.response_tag, None);
820    }
821
822    /// Property: whatever mix of pushes is accepted and however the sends split,
823    /// the bytes that leave are the accepted bytes in order, and the budget returns.
824    #[test]
825    fn prop_output_stream_matches_accepted_pushes() {
826        use crate::testing::Gen;
827        static TEXT: [u8; 4096] = {
828            let mut t = [0u8; 4096];
829            let mut i = 0;
830            while i < t.len() {
831                t[i] = (i % 251) as u8;
832                i += 1;
833            }
834            t
835        };
836        let mut g = Gen::new(0x0047_0041);
837        for round in 0..if cfg!(miri) { 3 } else { 300 } {
838            let mut o = Output::standalone(64 * 1024);
839            let mut expected = Vec::new();
840            let mut sent = Vec::new();
841            for _ in 0..200 {
842                match g.below(4) {
843                    0 => {
844                        let n = g.below(40_000) as usize;
845                        let data: Vec<u8> = (0..n).map(|_| g.next_u64() as u8).collect();
846                        let took = o.push_partial(&data);
847                        expected.extend_from_slice(&data[..took]);
848                    }
849                    1 => {
850                        let n = g.below(3000) as usize;
851                        let data: Vec<u8> = (0..n).map(|_| g.next_u64() as u8).collect();
852                        if o.push(&data).is_ok() {
853                            expected.extend_from_slice(&data);
854                        }
855                    }
856                    2 => {
857                        let start = g.below(TEXT.len() as u64) as usize;
858                        let piece = &TEXT[start..];
859                        if o.push_static(piece).is_ok() {
860                            expected.extend_from_slice(piece);
861                        }
862                    }
863                    _ => {
864                        // A send of a random size, taken from the front chunks as a vectored write would.
865                        let mut want = g.below(70_000) as usize;
866                        let mut taken = 0;
867                        for chunk in o.chunks.iter() {
868                            if want == 0 {
869                                break;
870                            }
871                            let r = chunk.remaining();
872                            let n = r.len().min(want);
873                            sent.extend_from_slice(&r[..n]);
874                            want -= n;
875                            taken += n;
876                        }
877                        o.advance(taken);
878                    }
879                }
880            }
881            while !o.is_empty() {
882                let n = o.chunks.front().map_or(0, |c| c.remaining().len());
883                if let Some(c) = o.chunks.front() {
884                    sent.extend_from_slice(c.remaining());
885                }
886                o.advance(n);
887            }
888            assert!(sent == expected, "round {round}: stream differs");
889            assert_eq!(o.used, 0, "round {round}: budget not returned");
890            assert_eq!(o.pending(), 0);
891        }
892    }
893}
net_io/session.rs

net_io/
session.rs
