tcp/
lib.rs

1//! A TCP implementation with a somewhat-BSD-like socket API. It is written as a
2//! ["sans-I/O"][sans-io] library meaning it doesn't do any networking I/O itself, it just accepts
3//! bytes and provides bytes. A [dependencies](Dependencies) object must be be provided to support
4//! setting timers and getting the current time. The TCP state object should probably be used with a
5//! reference-counting wrapper so that a reference to the state object can be stored in the timer
6//! callbacks.
7//!
8//! [sans-io]: https://sans-io.readthedocs.io
9//!
10//! ```
11//! use std::cell::RefCell;
12//! use std::rc::{Rc, Weak};
13//!
14//! #[derive(Debug)]
15//! struct TcpDependencies {
16//!     // a reference to the tcp state
17//!     state: Weak<RefCell<tcp::TcpState<Self>>>,
18//! }
19//!
20//! impl tcp::Dependencies for TcpDependencies {
21//!     type Instant = std::time::Instant;
22//!     type Duration = std::time::Duration;
23//!
24//!     fn register_timer(
25//!         &self,
26//!         time: Self::Instant,
27//!         f: impl FnOnce(&mut tcp::TcpState<Self>, tcp::TimerRegisteredBy) + Send + Sync + 'static,
28//!     ) {
29//!         let tcp_state = self.state.upgrade().unwrap();
30//!
31//!         // TODO: To register timers, you would likely want to involve an async
32//!         // runtime. A simple example would create a new task for each timer. The
33//!         // task would sleep for some duration and then run the callback.
34//!     }
35//!
36//!     fn current_time(&self) -> Self::Instant {
37//!         std::time::Instant::now()
38//!     }
39//!
40//!     fn fork(&self) -> Self {
41//!         // TODO: the implementation here would depend on the implementation
42//!         // of `register_timer`
43//!         todo!();
44//!     }
45//! }
46//!
47//! // create the TCP state object
48//! let tcp_state = Rc::new_cyclic(|weak| {
49//!     let dependencies = TcpDependencies {
50//!         state: weak.clone(),
51//!     };
52//!     RefCell::new(tcp::TcpState::new(dependencies, tcp::TcpConfig::default()))
53//! });
54//!
55//! let mut tcp_state = tcp_state.borrow_mut();
56//!
57//! // connect to port 80
58//! let dst_addr = "127.0.0.1:80".parse().unwrap();
59//! tcp_state.connect(dst_addr, || {
60//!     // here we would ask the network interface for an unused port (implicit bind),
61//!     // or where we would use the port assigned to a raw IP socket
62//!     let bind_addr = "127.0.0.1:2532".parse().unwrap();
63//!     Ok::<_, ()>((bind_addr, ()))
64//! }).unwrap();
65//!
66//! // get the SYN packet from the connect
67//! let (header, _payload) = tcp_state.pop_packet().unwrap();
68//! assert!(header.flags.contains(tcp::TcpFlags::SYN));
69//! assert_eq!(header.dst(), dst_addr);
70//! ```
71
72// There are three related state types in this crate:
73//
74// - `TcpState` — This is the public-facing type for the TCP state. Its methods take shared or
75//   mutable references. It contains a non-public `TcpStateEnum`.
76// - `TcpStateEnum` — An enum of all individual TCP state types (ex: `ListeningState`,
77//   `EstablishedState`). It implements the `TcpStateTrait` trait, so its methods usually take owned
78//   objects and return owned objects.
79// - `TcpStateTrait` — A trait implemented by each individual TCP state type, as well as the
80//   `TcpStateEnum` enum that encapsulates all individual states. Its methods usually take owned
81//   state objects and return owned `TcpStateEnum` objects.
82
83#![forbid(unsafe_code)]
84
85use std::fmt::Debug;
86use std::io::{Read, Write};
87use std::net::{Ipv4Addr, SocketAddrV4};
88
89use bytes::{Bytes, BytesMut};
90
91pub mod util;
92
93mod buffer;
94mod connection;
95mod seq;
96mod states;
97mod window_scaling;
98
99#[cfg(test)]
100mod tests;
101
102use crate::states::{
103    CloseWaitState, ClosedState, ClosingState, EstablishedState, FinWaitOneState, FinWaitTwoState,
104    InitState, LastAckState, ListenState, RstState, SynReceivedState, SynSentState, TimeWaitState,
105};
106use crate::util::SmallArrayBackedSlice;
107
108/// A collection of methods that allow the TCP state to interact with the external system.
109pub trait Dependencies: Debug + Sized {
110    type Instant: crate::util::time::Instant<Duration = Self::Duration>;
111    type Duration: crate::util::time::Duration;
112
113    /// Register a timer. The callback will be run on the parent [state](TcpState). The callback can
114    /// use the [`TimerRegisteredBy`] argument to know whether the timer was registered by the
115    /// parent state or one of its child states.
116    ///
117    /// If a child state has not yet been accept()ed, it will be owned by a parent state. When a
118    /// child state registers a timer, the timer's callback will run on the parent state and the
119    /// callback will be given the `TimerRegisteredBy::Child` argument so that the callback can
120    /// delegate accordingly.
121    fn register_timer(
122        &self,
123        time: Self::Instant,
124        f: impl FnOnce(&mut TcpState<Self>, TimerRegisteredBy) + Send + Sync + 'static,
125    );
126
127    /// Get the current time.
128    fn current_time(&self) -> Self::Instant;
129
130    /// Create a new `Dependencies` for use by a child state. When a timer is registered by the
131    /// child state using this new object, the timer's callback will be run on the parent's state
132    /// with the `TimerRegisteredBy::Child` argument so that the parent knows to run the callback on
133    /// one of its child states.
134    ///
135    /// When a child state has been accept()ed, it will no longer be owned by the parent state and
136    /// the parent state has no way to access this child state. The child state's `Dependencies`
137    /// should be updated during the [`finalize`](AcceptedTcpState::finalize) call (on the
138    /// [`AcceptedTcpState`] returned from [`accept`](TcpState::accept)) to run callbacks directly
139    /// on this state instead, and the callbacks should be given `TimerRegisteredBy::Parent` (the
140    /// child state has effectively become a parent state). This `Dependencies` object should also
141    /// make sure that all existing timers from before the state was accept()ed are also updated to
142    /// run callbacks directly on the state.
143    fn fork(&self) -> Self;
144}
145
146/// Specifies whether the callback is meant to run on the parent state or a child state.
147///
148/// For example if a child state registers a timer, a value of `TimerRegisteredBy::Child` will be
149/// given to the callback so that it knows to apply the callback to a child state, not the parent
150/// state.
151#[derive(Copy, Clone, Debug, PartialEq, Eq)]
152pub enum TimerRegisteredBy {
153    Parent,
154    Child,
155}
156
157#[enum_dispatch::enum_dispatch]
158trait TcpStateTrait<X>: Debug + Sized
159where
160    X: Dependencies,
161    TcpStateEnum<X>: From<Self>,
162{
163    /// Start closing this socket. It may or may not close immediately depending on what state the
164    /// socket is currently in.
165    fn close(self) -> (TcpStateEnum<X>, Result<(), CloseError>) {
166        (self.into(), Err(CloseError::InvalidState))
167    }
168
169    /// Start closing this socket by sending an RST packet. It may or may not close immediately
170    /// depending on what state the socket is currently in.
171    ///
172    /// TODO:
173    /// RFC 9293: "The side of a connection issuing a reset should enter the TIME-WAIT state, [...]"
174    fn rst_close(self) -> (TcpStateEnum<X>, Result<(), RstCloseError>) {
175        (self.into(), Err(RstCloseError::InvalidState))
176    }
177
178    fn shutdown(self, _how: Shutdown) -> (TcpStateEnum<X>, Result<(), ShutdownError>) {
179        (self.into(), Err(ShutdownError::InvalidState))
180    }
181
182    fn listen<T, E>(
183        self,
184        _backlog: u32,
185        _associate_fn: impl FnOnce() -> Result<T, E>,
186    ) -> (TcpStateEnum<X>, Result<T, ListenError<E>>) {
187        (self.into(), Err(ListenError::InvalidState))
188    }
189
190    fn connect<T, E>(
191        self,
192        _addr: SocketAddrV4,
193        _associate_fn: impl FnOnce() -> Result<(SocketAddrV4, T), E>,
194    ) -> (TcpStateEnum<X>, Result<T, ConnectError<E>>) {
195        (self.into(), Err(ConnectError::InvalidState))
196    }
197
198    /// Accept a new child state from the pending connection queue. The TCP state for the child
199    /// socket is returned. The [`AcceptedTcpState::finalize`] method must be called immediately on
200    /// the returned child state before any code calls into the parent state again, otherwise the
201    /// child may miss some timer events.
202    fn accept(self) -> (TcpStateEnum<X>, Result<AcceptedTcpState<X>, AcceptError>) {
203        (self.into(), Err(AcceptError::InvalidState))
204    }
205
206    fn send(self, _reader: impl Read, _len: usize) -> (TcpStateEnum<X>, Result<usize, SendError>) {
207        (self.into(), Err(SendError::InvalidState))
208    }
209
210    fn recv(self, _writer: impl Write, _len: usize) -> (TcpStateEnum<X>, Result<usize, RecvError>) {
211        (self.into(), Err(RecvError::InvalidState))
212    }
213
214    /// Returns the number of bytes added to the TCP state's receive buffer. This may be
215    /// smaller (ex: duplicate packet) or larger (ex: there is a non-empty reassembly queue)
216    /// than the packet payload length.
217    fn push_packet(
218        self,
219        _header: &TcpHeader,
220        _payload: Payload,
221    ) -> (TcpStateEnum<X>, Result<u32, PushPacketError>) {
222        (self.into(), Err(PushPacketError::InvalidState))
223    }
224
225    fn pop_packet(
226        self,
227    ) -> (
228        TcpStateEnum<X>,
229        Result<(TcpHeader, Payload), PopPacketError>,
230    ) {
231        (self.into(), Err(PopPacketError::InvalidState))
232    }
233
234    fn clear_error(&mut self) -> Option<TcpError>;
235
236    fn poll(&self) -> PollState;
237
238    fn wants_to_send(&self) -> bool;
239
240    fn local_remote_addrs(&self) -> Option<(SocketAddrV4, SocketAddrV4)>;
241}
242
243#[derive(Debug)]
244pub struct TcpState<X: Dependencies>(Option<TcpStateEnum<X>>);
245
246// this exposes many of the methods from `TcpStateTrait`, but not necessarily all of them (for
247// example we don't expose `rst_close()`).
248impl<X: Dependencies> TcpState<X> {
249    pub fn new(deps: X, config: TcpConfig) -> Self {
250        let new_state = InitState::new(deps, config);
251        Self(Some(new_state.into()))
252    }
253
254    #[inline]
255    fn with_state<T>(&mut self, f: impl FnOnce(TcpStateEnum<X>) -> (TcpStateEnum<X>, T)) -> T {
256        // get the current state, pass it to `f`, and then put it back (`f` may actually replace the
257        // state with an entirely different state object)
258        let state = self.0.take().unwrap();
259        let (state, rv) = f(state);
260        self.0 = Some(state);
261
262        rv
263    }
264
265    #[inline]
266    pub fn close(&mut self) -> Result<(), CloseError> {
267        self.with_state(|state| state.close())
268    }
269
270    #[inline]
271    pub fn shutdown(&mut self, how: Shutdown) -> Result<(), ShutdownError> {
272        self.with_state(|state| state.shutdown(how))
273    }
274
275    #[inline]
276    pub fn listen<T, E>(
277        &mut self,
278        backlog: u32,
279        associate_fn: impl FnOnce() -> Result<T, E>,
280    ) -> Result<T, ListenError<E>> {
281        self.with_state(|state| state.listen(backlog, associate_fn))
282    }
283
284    #[inline]
285    pub fn connect<T, E>(
286        &mut self,
287        addr: SocketAddrV4,
288        associate_fn: impl FnOnce() -> Result<(SocketAddrV4, T), E>,
289    ) -> Result<T, ConnectError<E>> {
290        self.with_state(|state| state.connect(addr, associate_fn))
291    }
292
293    #[inline]
294    pub fn accept(&mut self) -> Result<AcceptedTcpState<X>, AcceptError> {
295        self.with_state(|state| state.accept())
296    }
297
298    #[inline]
299    pub fn send(&mut self, reader: impl Read, len: usize) -> Result<usize, SendError> {
300        self.with_state(|state| state.send(reader, len))
301    }
302
303    #[inline]
304    pub fn recv(&mut self, writer: impl Write, len: usize) -> Result<usize, RecvError> {
305        self.with_state(|state| state.recv(writer, len))
306    }
307
308    #[inline]
309    pub fn push_packet(
310        &mut self,
311        header: &TcpHeader,
312        payload: Payload,
313    ) -> Result<u32, PushPacketError> {
314        self.with_state(|state| state.push_packet(header, payload))
315    }
316
317    #[inline]
318    pub fn pop_packet(&mut self) -> Result<(TcpHeader, Payload), PopPacketError> {
319        self.with_state(|state| state.pop_packet())
320    }
321
322    #[inline]
323    pub fn clear_error(&mut self) -> Option<TcpError> {
324        self.0.as_mut().unwrap().clear_error()
325    }
326
327    #[inline]
328    pub fn poll(&self) -> PollState {
329        self.0.as_ref().unwrap().poll()
330    }
331
332    #[inline]
333    pub fn wants_to_send(&self) -> bool {
334        self.0.as_ref().unwrap().wants_to_send()
335    }
336
337    #[inline]
338    pub fn local_remote_addrs(&self) -> Option<(SocketAddrV4, SocketAddrV4)> {
339        self.0.as_ref().unwrap().local_remote_addrs()
340    }
341}
342
343/// A macro that forwards an argument-less method to the inner type.
344///
345/// ```ignore
346/// // forward!(as_init, Option<&InitState<X>>);
347/// #[inline]
348/// pub fn as_init(&self) -> Option<&InitState<X>> {
349///     self.0.as_ref().unwrap().as_init()
350/// }
351/// ```
352#[cfg(test)]
353macro_rules! forward {
354    ($fn_name:ident, $($return_type:tt)*) => {
355        #[inline]
356        pub fn $fn_name(&self) -> $($return_type)* {
357            self.0.as_ref().unwrap().$fn_name()
358        }
359    };
360}
361
362#[cfg(test)]
363impl<X: Dependencies> TcpState<X> {
364    forward!(as_init, Option<&InitState<X>>);
365    forward!(as_listen, Option<&ListenState<X>>);
366    forward!(as_syn_sent, Option<&SynSentState<X>>);
367    forward!(as_syn_received, Option<&SynReceivedState<X>>);
368    forward!(as_established, Option<&EstablishedState<X>>);
369    forward!(as_fin_wait_one, Option<&FinWaitOneState<X>>);
370    forward!(as_fin_wait_two, Option<&FinWaitTwoState<X>>);
371    forward!(as_closing, Option<&ClosingState<X>>);
372    forward!(as_time_wait, Option<&TimeWaitState<X>>);
373    forward!(as_close_wait, Option<&CloseWaitState<X>>);
374    forward!(as_last_ack, Option<&LastAckState<X>>);
375    forward!(as_rst, Option<&RstState<X>>);
376    forward!(as_closed, Option<&ClosedState<X>>);
377}
378
379#[enum_dispatch::enum_dispatch(TcpStateTrait<X>)]
380#[derive(Debug)]
381enum TcpStateEnum<X: Dependencies> {
382    Init(InitState<X>),
383    Listen(ListenState<X>),
384    SynSent(SynSentState<X>),
385    SynReceived(SynReceivedState<X>),
386    Established(EstablishedState<X>),
387    FinWaitOne(FinWaitOneState<X>),
388    FinWaitTwo(FinWaitTwoState<X>),
389    Closing(ClosingState<X>),
390    TimeWait(TimeWaitState<X>),
391    CloseWait(CloseWaitState<X>),
392    LastAck(LastAckState<X>),
393    Rst(RstState<X>),
394    Closed(ClosedState<X>),
395}
396
397/// A macro that creates a method which casts to an inner variant.
398///
399/// ```ignore
400/// // as_impl!(as_init, Init, InitState);
401/// #[inline]
402/// pub fn as_init(&self) -> Option<&InitState<X>> {
403///     match self {
404///         Self::Init(x) => Some(x),
405///         _ => None,
406///     }
407/// }
408/// ```
409#[cfg(test)]
410macro_rules! as_impl {
411    ($fn_name:ident, $variant:ident, $return_type:ident) => {
412        #[inline]
413        pub fn $fn_name(&self) -> Option<&$return_type<X>> {
414            match self {
415                Self::$variant(x) => Some(x),
416                _ => None,
417            }
418        }
419    };
420}
421
422/// Casts to concrete types. This should only be called from unit tests to verify state.
423#[cfg(test)]
424impl<X: Dependencies> TcpStateEnum<X> {
425    as_impl!(as_init, Init, InitState);
426    as_impl!(as_listen, Listen, ListenState);
427    as_impl!(as_syn_sent, SynSent, SynSentState);
428    as_impl!(as_syn_received, SynReceived, SynReceivedState);
429    as_impl!(as_established, Established, EstablishedState);
430    as_impl!(as_fin_wait_one, FinWaitOne, FinWaitOneState);
431    as_impl!(as_fin_wait_two, FinWaitTwo, FinWaitTwoState);
432    as_impl!(as_closing, Closing, ClosingState);
433    as_impl!(as_time_wait, TimeWait, TimeWaitState);
434    as_impl!(as_close_wait, CloseWait, CloseWaitState);
435    as_impl!(as_last_ack, LastAck, LastAckState);
436    as_impl!(as_rst, Rst, RstState);
437    as_impl!(as_closed, Closed, ClosedState);
438}
439
440/// An accept()ed TCP state. This is used to ensure that the caller uses
441/// [`finalize`](Self::finalize) to update the state's `Dependencies` since the state is no longer
442/// owned by the listening socket.
443// we use a wrapper struct around an enum so that public code can't access the inner state object
444pub struct AcceptedTcpState<X: Dependencies>(AcceptedTcpStateInner<X>);
445
446/// An "established" or "close-wait" TCP state can be accept()ed, so we need to be able to store
447/// either state.
448enum AcceptedTcpStateInner<X: Dependencies> {
449    Established(EstablishedState<X>),
450    CloseWait(CloseWaitState<X>),
451}
452
453impl<X: Dependencies> AcceptedTcpState<X> {
454    /// This allows the caller to update the state's `Dependencies`.
455    ///
456    /// This must be called immediately after [`TcpState::accept`], otherwise the accept()ed socket
457    /// may miss some of its timer events.
458    pub fn finalize(mut self, f: impl FnOnce(&mut X)) -> TcpState<X> {
459        let deps = match &mut self.0 {
460            AcceptedTcpStateInner::Established(state) => &mut state.common.deps,
461            AcceptedTcpStateInner::CloseWait(state) => &mut state.common.deps,
462        };
463
464        // allow the caller to update `deps` for this state since this state is changing owners
465        f(deps);
466
467        TcpState(Some(self.0.into()))
468    }
469
470    pub fn local_addr(&self) -> SocketAddrV4 {
471        match &self.0 {
472            AcceptedTcpStateInner::Established(state) => state.connection.local_addr,
473            AcceptedTcpStateInner::CloseWait(state) => state.connection.local_addr,
474        }
475    }
476
477    pub fn remote_addr(&self) -> SocketAddrV4 {
478        match &self.0 {
479            AcceptedTcpStateInner::Established(state) => state.connection.remote_addr,
480            AcceptedTcpStateInner::CloseWait(state) => state.connection.remote_addr,
481        }
482    }
483}
484
485impl<X: Dependencies> TryFrom<TcpStateEnum<X>> for AcceptedTcpState<X> {
486    type Error = TcpStateEnum<X>;
487
488    fn try_from(state: TcpStateEnum<X>) -> Result<Self, Self::Error> {
489        match state {
490            TcpStateEnum::Established(state) => Ok(Self(AcceptedTcpStateInner::Established(state))),
491            TcpStateEnum::CloseWait(state) => Ok(Self(AcceptedTcpStateInner::CloseWait(state))),
492            // return the state back to the caller
493            state => Err(state),
494        }
495    }
496}
497
498impl<X: Dependencies> From<AcceptedTcpStateInner<X>> for TcpStateEnum<X> {
499    fn from(inner: AcceptedTcpStateInner<X>) -> Self {
500        match inner {
501            AcceptedTcpStateInner::Established(state) => state.into(),
502            AcceptedTcpStateInner::CloseWait(state) => state.into(),
503        }
504    }
505}
506
507#[derive(Copy, Clone, Debug, PartialEq, Eq)]
508pub enum Shutdown {
509    Read,
510    Write,
511    Both,
512}
513
514#[derive(Debug)]
515pub enum TcpError {
516    ResetSent,
517    ResetReceived,
518    /// The connection was closed while it was connecting, and no RST was sent or received.
519    ClosedWhileConnecting,
520    TimedOut,
521}
522
523// errors for operations on `TcpStateTrait` objects
524
525#[derive(Debug)]
526pub enum CloseError {
527    InvalidState,
528}
529
530#[derive(Debug)]
531enum RstCloseError {
532    InvalidState,
533}
534
535#[derive(Debug)]
536pub enum ListenError<E> {
537    InvalidState,
538    FailedAssociation(E),
539}
540
541#[derive(Debug)]
542pub enum ConnectError<E> {
543    InvalidState,
544    /// A previous connection attempt is in progress.
545    InProgress,
546    /// A connection has previously been attempted and was either successful or unsuccessful (it may
547    /// or may not have reached the "established" state). The connection may be established, timed
548    /// out, closing, half-closed, closed, etc. This does not include connection attempts that are
549    /// in progress ("syn-sent" or "syn-received" states).
550    AlreadyConnected,
551    /// Is already listening for new connections.
552    IsListening,
553    FailedAssociation(E),
554}
555
556#[derive(Debug)]
557pub enum AcceptError {
558    InvalidState,
559    NothingToAccept,
560}
561
562#[derive(Debug)]
563pub enum ShutdownError {
564    NotConnected,
565    InvalidState,
566}
567
568#[derive(Debug)]
569pub enum SendError {
570    InvalidState,
571    Full,
572    NotConnected,
573    StreamClosed,
574    Io(std::io::Error),
575}
576
577#[derive(Debug)]
578pub enum RecvError {
579    InvalidState,
580    Empty,
581    NotConnected,
582    /// The peer has sent a FIN, so no more data will be received.
583    StreamClosed,
584    Io(std::io::Error),
585}
586
587#[derive(Debug)]
588pub enum PushPacketError {
589    InvalidState,
590}
591
592#[derive(Debug)]
593pub enum PopPacketError {
594    InvalidState,
595    NoPacket,
596}
597
598// segment/packet headers
599
600bitflags::bitflags! {
601    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
602    pub struct PollState: u32 {
603        /// Data can be read.
604        const READABLE = 1 << 0;
605        /// Data can be written.
606        const WRITABLE = 1 << 1;
607        /// There is a pending error that can be read using [`TcpState::clear_error`].
608        const ERROR = 1 << 2;
609        /// The connection has been closed for receiving. This is not mutually exclusive with
610        /// `READABLE` (even if it's closed for receiving, there may still be buffered data to
611        /// read). Some possible causes are:
612        /// - Received a FIN packet.
613        /// - Sent or received a RST packet.
614        /// - TCP was closed.
615        const RECV_CLOSED = 1 << 3;
616        /// The connection has been closed for sending. This should be mutually exclusive with
617        /// `WRITABLE` (there would be no point in writing data if it's closed for sending). Some
618        /// possible causes are:
619        /// - Sent a FIN packet.
620        /// - Sent or received a RST packet.
621        /// - TCP was `shutdown()` for writing.
622        /// - TCP was closed.
623        const SEND_CLOSED = 1 << 4;
624        /// Is listening for new connections.
625        const LISTENING = 1 << 5;
626        /// A listening socket has a new incoming connection that can be accepted.
627        const READY_TO_ACCEPT = 1 << 6;
628        /// Connection is in the process of opening. More specifically this means that it is in
629        /// either the "syn-sent" or "syn-received" states.
630        const CONNECTING = 1 << 7;
631        /// A connection has previously been attempted and was either successful or unsuccessful (it
632        /// may or may not have reached the "established" state). The connection may be established,
633        /// timed out, closing, half-closed, closed, etc. This does not include connection attempts
634        /// that are in progress ("syn-sent" or "syn-received" states).
635        const CONNECTED = 1 << 8;
636        /// TCP is fully closed (in the "closed" state). This may not be set immediately after a
637        /// `close()` call, for example if `close()` was called while in the "established" state,
638        /// and now is in the "fin-wait-1" state. This does not include the initial state (we don't
639        /// consider a new TCP to be "closed").
640        const CLOSED = 1 << 9;
641    }
642}
643
644#[derive(Copy, Clone, Debug)]
645#[non_exhaustive]
646pub struct TcpConfig {
647    pub(crate) window_scaling_enabled: bool,
648}
649
650impl TcpConfig {
651    pub fn window_scaling(&mut self, enable: bool) {
652        self.window_scaling_enabled = enable;
653    }
654}
655
656impl Default for TcpConfig {
657    fn default() -> Self {
658        Self {
659            window_scaling_enabled: true,
660        }
661    }
662}
663
664bitflags::bitflags! {
665    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
666    pub struct TcpFlags: u8 {
667        const FIN = 1 << 0;
668        const SYN = 1 << 1;
669        const RST = 1 << 2;
670        const PSH = 1 << 3;
671        const ACK = 1 << 4;
672        const URG = 1 << 5;
673        const ECE = 1 << 6;
674        const CWR = 1 << 7;
675    }
676}
677
678#[derive(Copy, Clone, Debug)]
679pub struct TcpHeader {
680    pub ip: Ipv4Header,
681    pub flags: TcpFlags,
682    pub src_port: u16,
683    pub dst_port: u16,
684    pub seq: u32,
685    pub ack: u32,
686    pub window_size: u16,
687    pub selective_acks: Option<SmallArrayBackedSlice<4, (u32, u32)>>,
688    pub window_scale: Option<u8>,
689    pub timestamp: Option<u32>,
690    pub timestamp_echo: Option<u32>,
691}
692
693impl TcpHeader {
694    pub fn src(&self) -> SocketAddrV4 {
695        SocketAddrV4::new(self.ip.src, self.src_port)
696    }
697
698    pub fn dst(&self) -> SocketAddrV4 {
699        SocketAddrV4::new(self.ip.dst, self.dst_port)
700    }
701}
702
703#[derive(Copy, Clone, Debug)]
704pub struct Ipv4Header {
705    pub src: Ipv4Addr,
706    pub dst: Ipv4Addr,
707}
708
709/// A packet payload containing a list of [byte](Bytes) chunks.
710///
711/// The sum of the lengths of each chunk must be at most [`u32::MAX`], otherwise operations on the
712/// payload or other code using the payload may panic.
713// TODO: Intuitively this seems like a good place to use a `SmallVec` to optimize the common case
714// where there are a small number of chunks per packet. But I'm leaving this until we can test `Vec`
715// vs `SmallVec` in a benchmark to see if there's any performance improvement in practice.
716#[derive(Clone, Debug, Default)]
717pub struct Payload(pub Vec<Bytes>);
718
719// We don't implement `PartialEq` or `Eq` since it's not clear what equality means. Are payloads
720// equal if they just contain the same bytes, or are they equal only if the chunks are exactly the
721// same? For example is the payload `["hello", "world"]` the same as `["helloworld"]`?
722static_assertions::assert_not_impl_any!(Payload: PartialEq, Eq);
723
724impl Payload {
725    /// Returns the number of bytes in the payload.
726    pub fn len(&self) -> u32 {
727        self.0
728            .iter()
729            // `fold` rather than `sum` so that we always panic on overflow
730            .fold(0usize, |acc, x| acc.checked_add(x.len()).unwrap())
731            .try_into()
732            .unwrap()
733    }
734
735    /// Returns true if the payload has no data (no byte chunks or only empty byte chunks).
736    pub fn is_empty(&self) -> bool {
737        // should be faster than checking `self.len() == 0`
738        self.0.iter().all(|x| x.len() == 0)
739    }
740
741    /// Concatenate the byte chunks into a single byte chunk. Unless the payload is empty or has a
742    /// single chunk, this will allocate a large buffer and copy all of the individual chunks to
743    /// this new buffer.
744    pub fn concat(&self) -> Bytes {
745        let num_bytes = self.len() as usize;
746        let num_chunks = self.0.len();
747
748        if num_bytes == 0 {
749            return Bytes::new();
750        }
751
752        if num_chunks == 1 {
753            // there's only one chunk, so just return a reference to the chunk
754            return self.0[0].clone();
755        }
756
757        let mut bytes = BytesMut::with_capacity(num_bytes);
758
759        for chunk in &self.0 {
760            bytes.extend_from_slice(chunk);
761        }
762
763        debug_assert_eq!(bytes.len(), bytes.capacity());
764
765        bytes.freeze()
766    }
767}
768
769impl From<Bytes> for Payload {
770    fn from(bytes: Bytes) -> Self {
771        Self(vec![bytes])
772    }
773}
774
775impl From<BytesMut> for Payload {
776    fn from(bytes: BytesMut) -> Self {
777        bytes.freeze().into()
778    }
779}