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

net_io/
session.rs
